252 lines
10 KiB
Python
252 lines
10 KiB
Python
import motor.motor_asyncio
|
|
import asyncio
|
|
from typing import Optional, Dict, Any, List
|
|
|
|
class MongoHandler:
|
|
"""
|
|
A basic asynchronous handler for MongoDB operations using motor.
|
|
Designed to be used with asyncio.
|
|
"""
|
|
def __init__(self, db_name: str, collection_name: str, mongo_uri: str = "mongodb://localhost:27017/"):
|
|
"""
|
|
Initializes the MongoDB handler.
|
|
|
|
Args:
|
|
db_name (str): The name of the database to connect to.
|
|
collection_name (str): The name of the collection to use.
|
|
mongo_uri (str): The MongoDB connection string URI.
|
|
Defaults to the standard local URI.
|
|
"""
|
|
self.mongo_uri = mongo_uri
|
|
self.db_name = db_name
|
|
self.collection_name = collection_name
|
|
self._client: Optional[motor.motor_asyncio.AsyncIOMotorClient] = None
|
|
self._db: Optional[motor.motor_asyncio.AsyncIOMotorDatabase] = None
|
|
self._collection: Optional[motor.motor_asyncio.AsyncIOMotorCollection] = None
|
|
|
|
async def connect(self):
|
|
"""Establishes an asynchronous connection to MongoDB."""
|
|
if self._client is None:
|
|
try:
|
|
self._client = motor.motor_asyncio.AsyncIOMotorClient(self.mongo_uri)
|
|
# The ismaster command is cheap and does not require auth.
|
|
# It is used to confirm that the client can connect to the deployment.
|
|
await self._client.admin.command('ismaster')
|
|
self._db = self._client[self.db_name]
|
|
self._collection = self._db[self.collection_name]
|
|
print(f"Connected to MongoDB: Database '{self.db_name}', Collection '{self.collection_name}'")
|
|
except Exception as e:
|
|
print(f"Failed to connect to MongoDB at {self.mongo_uri}: {e}")
|
|
self._client = None # Ensure client is None if connection fails
|
|
raise # Re-raise the exception after printing
|
|
|
|
async def close(self):
|
|
"""Closes the MongoDB connection."""
|
|
if self._client:
|
|
self._client.close()
|
|
self._client = None
|
|
self._db = None
|
|
self._collection = None
|
|
print("MongoDB connection closed.")
|
|
|
|
async def __aenter__(self):
|
|
"""Allows using the handler with async with."""
|
|
await self.connect()
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
"""Ensures the connection is closed when exiting async with."""
|
|
await self.close()
|
|
|
|
async def insert_one(self, document: Dict[str, Any]) -> Any:
|
|
"""
|
|
Inserts a single document into the collection.
|
|
|
|
Args:
|
|
document (Dict[str, Any]): The document to insert.
|
|
|
|
Returns:
|
|
Any: The result of the insert operation (InsertOneResult).
|
|
"""
|
|
if self._collection is None:
|
|
raise RuntimeError("MongoDB connection not established. Call connect() first or use async with.")
|
|
print(f"Inserting document into '{self.collection_name}'...")
|
|
result = await self._collection.insert_one(document)
|
|
print(f"Inserted document with ID: {result.inserted_id}")
|
|
return result
|
|
|
|
async def find_one(self, query: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Finds a single document matching the query.
|
|
|
|
Args:
|
|
query (Dict[str, Any]): The query document.
|
|
|
|
Returns:
|
|
Optional[Dict[str, Any]]: The found document, or None if not found.
|
|
"""
|
|
if self._collection is None:
|
|
raise RuntimeError("MongoDB connection not established. Call connect() first or use async with.")
|
|
print(f"Finding one document in '{self.collection_name}' with query: {query}")
|
|
document = await self._collection.find_one(query)
|
|
return document
|
|
|
|
async def find(self, query: Dict[str, Any] = None) -> List[Dict[str, Any]]:
|
|
"""
|
|
Finds multiple documents matching the query.
|
|
|
|
Args:
|
|
query (Dict[str, Any], optional): The query document. Defaults to None (find all).
|
|
|
|
Returns:
|
|
List[Dict[str, Any]]: A list of matching documents.
|
|
"""
|
|
if self._collection is None:
|
|
raise RuntimeError("MongoDB connection not established. Call connect() first or use async with.")
|
|
if query is None:
|
|
query = {}
|
|
print(f"Finding documents in '{self.collection_name}' with query: {query}")
|
|
# Use list comprehension to iterate through the cursor asynchronously
|
|
documents = [doc async for doc in self._collection.find(query)]
|
|
print(f"Found {len(documents)} documents.")
|
|
return documents
|
|
|
|
async def update_one(self, query: Dict[str, Any], update: Dict[str, Any], upsert: bool = False) -> Any:
|
|
"""
|
|
Updates a single document matching the query.
|
|
|
|
Args:
|
|
query (Dict[str, Any]): The query document.
|
|
update (Dict[str, Any]): The update operations to apply.
|
|
upsert (bool): If True, insert a new document if no match is found.
|
|
|
|
Returns:
|
|
Any: The result of the update operation (UpdateResult).
|
|
"""
|
|
if self._collection is None:
|
|
raise RuntimeError("MongoDB connection not established. Call connect() first or use async with.")
|
|
print(f"Updating one document in '{self.collection_name}' with query: {query}, update: {update}")
|
|
result = await self._collection.update_one(query, update, upsert=upsert)
|
|
print(f"Matched {result.matched_count}, Modified {result.modified_count}, Upserted ID: {result.upserted_id}")
|
|
return result
|
|
|
|
async def delete_one(self, query: Dict[str, Any]) -> Any:
|
|
"""
|
|
Deletes a single document matching the query.
|
|
|
|
Args:
|
|
query (Dict[str, Any]): The query document.
|
|
|
|
Returns:
|
|
Any: The result of the delete operation (DeleteResult).
|
|
"""
|
|
if self._collection is None:
|
|
raise RuntimeError("MongoDB connection not established. Call connect() first or use async with.")
|
|
print(f"Deleting one document from '{self.collection_name}' with query: {query}")
|
|
result = await self._collection.delete_one(query)
|
|
print(f"Deleted count: {result.deleted_count}")
|
|
return result
|
|
|
|
async def delete_many(self, query: Dict[str, Any]) -> Any:
|
|
"""
|
|
Deletes multiple documents matching the query.
|
|
|
|
Args:
|
|
query (Dict[str, Any]): The query document.
|
|
|
|
Returns:
|
|
Any: The result of the delete operation (DeleteResult).
|
|
"""
|
|
if self._collection is None:
|
|
raise RuntimeError("MongoDB connection not established. Call connect() first or use async with.")
|
|
print(f"Deleting many documents from '{self.collection_name}' with query: {query}")
|
|
result = await self._collection.delete_many(query)
|
|
print(f"Deleted count: {result.deleted_count}")
|
|
return result
|
|
|
|
|
|
# --- Example Usage ---
|
|
async def example_mongo_usage():
|
|
"""Demonstrates how to use the MongoHandler."""
|
|
# Ensure you have a MongoDB server running, default is localhost:27017
|
|
db_name = "radio_app_db"
|
|
collection_name = "channels"
|
|
|
|
# Using async with ensures the connection is closed automatically
|
|
async with MongoHandler(db_name, collection_name) as mongo:
|
|
# --- Insert Example ---
|
|
print("\n--- Inserting a document ---")
|
|
channel_data = {
|
|
"_id": "channel_3", # You can specify _id or let MongoDB generate one
|
|
"name": "Emergency Services",
|
|
"frequency_khz": 453000,
|
|
"location": "Countywide",
|
|
"avail_on_nodes": ["client-xyz987"],
|
|
"description": "Monitor for emergency broadcasts."
|
|
}
|
|
try:
|
|
insert_result = await mongo.insert_one(channel_data)
|
|
print(f"Insert successful: {insert_result.inserted_id}")
|
|
except Exception as e:
|
|
print(f"Insert failed: {e}")
|
|
|
|
|
|
# --- Find One Example ---
|
|
print("\n--- Finding one document ---")
|
|
query = {"_id": "channel_3"}
|
|
found_channel = await mongo.find_one(query)
|
|
if found_channel:
|
|
print("Found document:", found_channel)
|
|
else:
|
|
print("Document not found.")
|
|
|
|
# --- Find Many Example ---
|
|
print("\n--- Finding all documents ---")
|
|
all_channels = await mongo.find() # Empty query finds all
|
|
print("All documents:", all_channels)
|
|
|
|
# --- Update Example ---
|
|
print("\n--- Updating a document ---")
|
|
update_query = {"_id": "channel_3"}
|
|
update_data = {"$set": {"location": "Statewide", "avail_on_nodes": ["client-xyz987", "client-newnode1"]}}
|
|
update_result = await mongo.update_one(update_query, update_data)
|
|
print(f"Update successful: Matched {update_result.matched_count}, Modified {update_result.modified_count}")
|
|
|
|
print("\n--- Finding the updated document ---")
|
|
updated_channel = await mongo.find_one(update_query)
|
|
print("Updated document:", updated_channel)
|
|
|
|
# --- Delete Example ---
|
|
print("\n--- Deleting a document ---")
|
|
delete_query = {"_id": "channel_3"}
|
|
delete_result = await mongo.delete_one(delete_query)
|
|
print(f"Delete successful: Deleted count {delete_result.deleted_count}")
|
|
|
|
print("\n--- Verifying deletion ---")
|
|
deleted_channel = await mongo.find_one(delete_query)
|
|
if deleted_channel:
|
|
print("Document still found (deletion failed).")
|
|
else:
|
|
print("Document successfully deleted.")
|
|
|
|
# --- Insert another for delete_many example ---
|
|
temp_doc1 = {"_id": "temp_1", "tag": "temp"}
|
|
temp_doc2 = {"_id": "temp_2", "tag": "temp"}
|
|
await mongo.insert_one(temp_doc1)
|
|
await mongo.insert_one(temp_doc2)
|
|
|
|
# --- Delete Many Example ---
|
|
print("\n--- Deleting many documents ---")
|
|
delete_many_query = {"tag": "temp"}
|
|
delete_many_result = await mongo.delete_many(delete_many_query)
|
|
print(f"Delete many successful: Deleted count {delete_many_result.deleted_count}")
|
|
|
|
|
|
# To run the example usage:
|
|
# 1. Ensure you have a MongoDB server running locally on the default port (27017).
|
|
# 2. Save the code as mongodb_handler.py.
|
|
# 3. Run from your terminal: python -m asyncio mongodb_handler.py
|
|
if __name__ == "__main__":
|
|
# Running the example directly requires running within an asyncio loop
|
|
asyncio.run(example_mongo_usage())
|