move into /app dir and add new db handler + wrappers
This commit is contained in:
Logan Cusano
251
app/internal/db_handler.py
Normal file
251
app/internal/db_handler.py
Normal file
@@ -0,0 +1,251 @@
|
||||
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())
|
||||
300
app/internal/db_wrappers.py
Normal file
300
app/internal/db_wrappers.py
Normal file
@@ -0,0 +1,300 @@
|
||||
import os
|
||||
import typing
|
||||
import asyncio
|
||||
from uuid import uuid4
|
||||
from typing import Optional, List, Dict, Any
|
||||
from enum import Enum
|
||||
from internal.db_handler import MongoHandler
|
||||
|
||||
# Init vars
|
||||
DB_NAME = os.getenv("DB_NAME", "default_db")
|
||||
MONGO_URL = os.getenv("MONGO_URL", "mongodb://10.10.202.4:27017/")
|
||||
|
||||
SYSTEM_DB_COLLECTION_NAME = "radio_systems"
|
||||
|
||||
# --- Types
|
||||
class DemodTypes(str, Enum):
|
||||
P25 = "P25"
|
||||
DMR = "DMR"
|
||||
ANALOG = "NBFM"
|
||||
|
||||
class TalkgroupTag:
|
||||
"""Represents a talkgroup tag."""
|
||||
def __init__(self, talkgroup: str, tagDec: int):
|
||||
self.talkgroup = talkgroup
|
||||
self.tagDec = tagDec
|
||||
|
||||
# Add a method to convert to a dictionary, useful for sending as JSON
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {"talkgroup": self.talkgroup, "tagDec": self.tagDec}
|
||||
|
||||
class System:
|
||||
"""
|
||||
A basic data model for a channel/system entry in a radio system.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
_id: str,
|
||||
_type: DemodTypes,
|
||||
name: str,
|
||||
frequency_khz: List[int],
|
||||
location: str,
|
||||
avail_on_nodes: List[str],
|
||||
description: Optional[str] = "",
|
||||
tags: Optional[List[TalkgroupTag]] = None,
|
||||
whitelist: Optional[List[int]] = None):
|
||||
"""
|
||||
Initializes a System object.
|
||||
|
||||
Args:
|
||||
_id: A unique identifier for the entry (e.g., MongoDB ObjectId string).
|
||||
_type: The demodulation type (P25, NBFM, etc.).
|
||||
name: The name of the channel/system.
|
||||
frequency_khz: The frequency in kilohertz.
|
||||
location: The geographical location or coverage area.
|
||||
avail_on_nodes: A list of node identifiers where this is available.
|
||||
description: A brief description.
|
||||
"""
|
||||
self._id: str = _id
|
||||
self.type: DemodTypes = _type
|
||||
self.name: str = name
|
||||
self.frequency_khz: List[int] = frequency_khz
|
||||
self.location: str = location
|
||||
self.avail_on_nodes: List[str] = avail_on_nodes
|
||||
self.description: str = description or ""
|
||||
self.tags: List[TalkgroupTag] = tags or None
|
||||
self.whitelist: List[int] = whitelist or None
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""
|
||||
Provides a developer-friendly string representation of the object.
|
||||
"""
|
||||
# Use self.type.value for string representation of the enum
|
||||
return (f"System(_id='{self._id}', type='{self.type.value}', name='{self.name}', "
|
||||
f"frequency_khz={self.frequency_khz}, location='{self.location}', "
|
||||
f"avail_on_nodes={self.avail_on_nodes}, description='{self.description}',"
|
||||
f" tags='{self.tags}', whitelist='{self.whitelist}')")
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Converts the System object to a dictionary suitable for MongoDB.
|
||||
Converts the DemodTypes enum to its string value.
|
||||
"""
|
||||
return {
|
||||
"_id": self._id,
|
||||
"type": self.type.value, # Store the enum value (string)
|
||||
"name": self.name,
|
||||
"frequency_khz": self.frequency_khz,
|
||||
"location": self.location,
|
||||
"avail_on_nodes": self.avail_on_nodes,
|
||||
"description": self.description,
|
||||
"tags": self.tags,
|
||||
"whitelist": self.whitelist,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "System":
|
||||
"""
|
||||
Creates a System object from a dictionary (e.g., from MongoDB).
|
||||
Converts the 'type' string back to a DemodTypes enum member.
|
||||
"""
|
||||
# Ensure required keys exist and handle potential type mismatches if necessary
|
||||
# Convert the type string back to the DemodTypes enum
|
||||
system_type = DemodTypes(data.get("type")) if data.get("type") else None # Handle missing or invalid type
|
||||
|
||||
if system_type is None:
|
||||
# Handle error: could raise an exception or return None/default
|
||||
# For this example, let's raise an error if type is missing/invalid
|
||||
raise ValueError(f"Invalid or missing 'type' in document data: {data}")
|
||||
|
||||
|
||||
return cls(
|
||||
_id=data.get("_id"),
|
||||
_type=system_type,
|
||||
name=data.get("name", ""), # Provide default empty string if name is missing
|
||||
frequency_khz=data.get("frequency_khz", 0), # Provide default 0 if missing
|
||||
location=data.get("location", ""),
|
||||
avail_on_nodes=data.get("avail_on_nodes", []), # Provide default empty list
|
||||
description=data.get("description", ""),
|
||||
tags=data.get("tags", None),
|
||||
whitelist=data.get("whitelist", None)
|
||||
)
|
||||
|
||||
|
||||
# --- System class ---
|
||||
class SystemDbController():
|
||||
def __init__(self):
|
||||
# Init the handler
|
||||
self.db_h = MongoHandler(DB_NAME, SYSTEM_DB_COLLECTION_NAME, MONGO_URL)
|
||||
|
||||
async def create_system(self, system_data: Dict[str, Any]) -> Optional[System]:
|
||||
"""
|
||||
Creates a new system entry in the database.
|
||||
|
||||
Args:
|
||||
system_data: A dictionary containing the data for the new system.
|
||||
|
||||
Returns:
|
||||
The created System object if successful, None otherwise.
|
||||
"""
|
||||
print("\n--- Creating a document ---")
|
||||
try:
|
||||
# Check if the data to be inserted has an ID
|
||||
if not system_data.get("_id"):
|
||||
system_data['_id'] = uuid4()
|
||||
|
||||
inserted_result = None
|
||||
inserted_id = None
|
||||
async with self.db_h as db:
|
||||
insert_result = await self.db_h.insert_one(system_data)
|
||||
inserted_id = insert_result.inserted_id
|
||||
|
||||
if inserted_id:
|
||||
print(f"Insert successful with ID: {inserted_id}")
|
||||
# Fetch the inserted document to get the complete data including the generated _id
|
||||
query = {"_id": inserted_id}
|
||||
|
||||
inserted_doc = None
|
||||
async with self.db_h as db:
|
||||
inserted_doc = await db.find_one(query)
|
||||
|
||||
if inserted_doc:
|
||||
# Convert the fetched dictionary back to a System object
|
||||
return System.from_dict(inserted_doc)
|
||||
else:
|
||||
print("Insert acknowledged but no ID returned.")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Create failed: {e}")
|
||||
return None
|
||||
|
||||
async def find_system(self, query: Dict[str, Any]) -> Optional[System]:
|
||||
"""
|
||||
Finds a single system entry in the database.
|
||||
|
||||
Args:
|
||||
query: A dictionary representing the query criteria.
|
||||
|
||||
Returns:
|
||||
A System object if found, None otherwise.
|
||||
"""
|
||||
print("\n--- Finding one document ---")
|
||||
try:
|
||||
found_doc = None
|
||||
async with self.db_h as db:
|
||||
found_doc = await db.find_one(query)
|
||||
|
||||
if found_doc:
|
||||
print("Found document (raw dict):", found_doc)
|
||||
# Convert the dictionary result to a System object
|
||||
return System.from_dict(found_doc)
|
||||
else:
|
||||
print("Document not found.")
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"Find failed: {e}")
|
||||
return None
|
||||
|
||||
async def find_systems(self, query: Dict[str, Any]) -> Optional[List[System]]:
|
||||
"""
|
||||
Finds one or more system entries in the database.
|
||||
|
||||
Args:
|
||||
query: A dictionary representing the query criteria.
|
||||
|
||||
Returns:
|
||||
A list of System object(s) if found, None otherwise.
|
||||
"""
|
||||
print("\n--- Finding documents ---")
|
||||
try:
|
||||
found_docs = None
|
||||
async with self.db_h as db:
|
||||
found_docs = await db.find(query)
|
||||
|
||||
if found_docs:
|
||||
print("Found document (raw dict):", found_docs)
|
||||
# Convert the dictionary results to a System object
|
||||
converted_systems = []
|
||||
for doc in found_docs:
|
||||
converted_systems.append(System.from_dict(doc))
|
||||
|
||||
print("YURB", found_docs, converted_systems)
|
||||
|
||||
return converted_systems if len(converted_systems) > 0 else None
|
||||
else:
|
||||
print("Document not found.")
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"Find failed: {e}")
|
||||
return None
|
||||
|
||||
async def find_all_systems(self, query: Dict[str, Any] = {}) -> List[System]:
|
||||
"""
|
||||
Finds multiple system entries in the database.
|
||||
|
||||
Args:
|
||||
query: A dictionary representing the query criteria (default is empty to find all).
|
||||
|
||||
Returns:
|
||||
A list of System objects.
|
||||
"""
|
||||
print("\n--- Finding multiple documents ---")
|
||||
try:
|
||||
found_docs = None
|
||||
async with self.db_h as db:
|
||||
found_docs = await db.find(query)
|
||||
|
||||
if found_docs:
|
||||
print(f"Found {len(found_docs)} documents (raw dicts).")
|
||||
# Convert the list of dictionaries to a list of System objects
|
||||
return [System.from_dict(doc) for doc in found_docs]
|
||||
else:
|
||||
print("No documents found.")
|
||||
return []
|
||||
except Exception as e:
|
||||
print(f"Find all failed: {e}")
|
||||
return []
|
||||
|
||||
async def update_system(self, query: Dict[str, Any], update_data: Dict[str, Any]) -> Optional[int]:
|
||||
"""
|
||||
Updates a single system entry in the database.
|
||||
|
||||
Args:
|
||||
query: A dictionary representing the query criteria to find the document.
|
||||
update_data: A dictionary representing the update operations (e.g., using $set).
|
||||
|
||||
Returns:
|
||||
The number of modified documents if successful, None otherwise.
|
||||
"""
|
||||
print("\n--- Updating a document ---")
|
||||
try:
|
||||
update_result = None
|
||||
async with self.db_h as db:
|
||||
update_result = await db.update_one(query, update_data)
|
||||
|
||||
print(f"Update result: Matched {update_result.matched_count}, Modified {update_result.modified_count}")
|
||||
return update_result.modified_count
|
||||
except Exception as e:
|
||||
print(f"Update failed: {e}")
|
||||
return None
|
||||
|
||||
async def delete_system(self, query: Dict[str, Any]) -> Optional[int]:
|
||||
"""
|
||||
Deletes a single system entry from the database.
|
||||
Args:
|
||||
query: A dictionary representing the query criteria to find the document to delete.
|
||||
Returns:
|
||||
The number of deleted documents if successful, None otherwise.
|
||||
"""
|
||||
print("\n--- Deleting a document ---")
|
||||
try:
|
||||
delete_result = None
|
||||
async with self.db_h as db:
|
||||
delete_result = await self.db_h.delete_one(query)
|
||||
|
||||
print(f"Delete result: Deleted count {delete_result.deleted_count}")
|
||||
return delete_result.deleted_count
|
||||
except Exception as e:
|
||||
print(f"Delete failed: {e}")
|
||||
return None
|
||||
Reference in New Issue
Block a user