mdb-engine 0.1.6__py3-none-any.whl → 0.4.12__py3-none-any.whl

mdb_engine/repositories/unit_of_work.py

@@ -0,0 +1,166 @@
+"""
+Unit of Work Pattern
+
+Manages repository access and provides a clean interface for data operations.
+The UnitOfWork acts as a factory for repositories and manages their lifecycle.
+"""
+
+import logging
+from typing import Any, Generic, TypeVar
+
+from .base import Entity, Repository
+from .mongo import MongoRepository
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=Entity)
+
+
+class UnitOfWork:
+    """
+    Unit of Work for managing repository access.
+
+    Provides a clean interface for accessing repositories through
+    attribute access (e.g., uow.users, uow.orders).
+
+    The UnitOfWork is request-scoped - one instance per HTTP request.
+    Repositories are created lazily and cached for the duration of the request.
+
+    Usage:
+        # In a route handler
+        @app.get("/users/{user_id}")
+        async def get_user(user_id: str, ctx: RequestContext = Depends()):
+            # Access repository through UnitOfWork
+            user = await ctx.uow.users.get(user_id)
+            return user
+
+        # With explicit repository method
+        @app.get("/orders")
+        async def list_orders(ctx: RequestContext = Depends()):
+            repo = ctx.uow.repository("orders", Order)
+            return await repo.find({"status": "pending"})
+
+    Repository Naming Convention:
+        - Attribute access uses collection name: uow.users -> users collection
+        - The entity class defaults to Entity if not registered
+        - Register entity classes for type-safe repositories
+    """
+
+    def __init__(
+        self,
+        db: Any,  # ScopedMongoWrapper - avoid import cycle
+        entity_registry: dict[str, type[Entity]] | None = None,
+    ):
+        """
+        Initialize the Unit of Work.
+
+        Args:
+            db: ScopedMongoWrapper for database access
+            entity_registry: Optional mapping of collection names to entity classes
+        """
+        self._db = db
+        self._repositories: dict[str, Repository] = {}
+        self._entity_registry: dict[str, type[Entity]] = entity_registry or {}
+
+    def register_entity(self, collection_name: str, entity_class: type[Entity]) -> None:
+        """
+        Register an entity class for a collection.
+
+        This enables type-safe repository access.
+
+        Args:
+            collection_name: Name of the collection
+            entity_class: Entity subclass for this collection
+        """
+        self._entity_registry[collection_name] = entity_class
+
+    def repository(
+        self,
+        name: str,
+        entity_class: type[T] | None = None,
+    ) -> Repository[T]:
+        """
+        Get or create a repository for a collection.
+
+        Args:
+            name: Collection name
+            entity_class: Optional entity class override
+
+        Returns:
+            Repository instance for the collection
+        """
+        if name in self._repositories:
+            return self._repositories[name]
+
+        # Determine entity class
+        if entity_class is None:
+            entity_class = self._entity_registry.get(name, Entity)
+
+        # Get collection from db wrapper
+        collection = getattr(self._db, name)
+
+        # Create repository
+        repo = MongoRepository(collection, entity_class)
+        self._repositories[name] = repo
+
+        logger.debug(f"Created repository for '{name}' with entity {entity_class.__name__}")
+        return repo
+
+    def __getattr__(self, name: str) -> Repository:
+        """
+        Access repositories via attribute syntax.
+
+        Example:
+            uow.users  # Returns Repository for 'users' collection
+            uow.orders  # Returns Repository for 'orders' collection
+        """
+        # Prevent recursion on private attributes
+        if name.startswith("_"):
+            raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'")
+
+        return self.repository(name)
+
+    @property
+    def db(self) -> Any:
+        """
+        Direct access to the underlying ScopedMongoWrapper.
+
+        Use this for operations not covered by the Repository interface,
+        like complex aggregations or raw queries.
+
+        Example:
+            # Complex aggregation
+            pipeline = [{"$match": {...}}, {"$group": {...}}]
+            results = await ctx.uow.db.users.aggregate(pipeline).to_list(None)
+        """
+        return self._db
+
+    def dispose(self) -> None:
+        """
+        Dispose of the UnitOfWork and clear cached repositories.
+
+        Called automatically at the end of a request scope.
+        """
+        self._repositories.clear()
+        logger.debug("UnitOfWork disposed")
+
+
+class TypedUnitOfWork(UnitOfWork, Generic[T]):
+    """
+    Generic typed UnitOfWork for better IDE support.
+
+    This is a convenience class that provides type hints for specific
+    repository types.
+
+    Usage:
+        class MyUnitOfWork(TypedUnitOfWork):
+            @property
+            def users(self) -> Repository[User]:
+                return self.repository("users", User)
+
+            @property
+            def orders(self) -> Repository[Order]:
+                return self.repository("orders", Order)
+    """
+
+    pass

mdb_engine/routing/README.md

@@ -286,7 +286,7 @@ async def safe_message_handler(websocket, message):
                 "type": "action_result",
                 "result": result
             })
-    except Exception as e:
+    except (ValueError, TypeError, KeyError) as e:
         logger.error(f"Error handling message: {e}")
         # Send error to client
         manager = await get_websocket_manager("my_app")

mdb_engine/routing/__init__.py

@@ -23,9 +23,7 @@ def _check_websockets_available():
         try:
             from fastapi import WebSocket
 
-            _websockets_module = __import__(
-                ".websockets", fromlist=[""], package=__name__
-            )
+            _websockets_module = __import__(".websockets", fromlist=[""], package=__name__)
             _websockets_available = True
         except (ImportError, AttributeError):
             _websockets_available = False

mdb_engine/routing/websockets.py

@@ -23,9 +23,10 @@ This module is part of MDB_ENGINE - MongoDB Engine.
 import asyncio
 import json
 import logging
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple
+from typing import Any
 
 # Check if FastAPI WebSocket support is available (OPTIONAL dependency)
 try:
@@ -47,8 +48,8 @@ class WebSocketConnection:
 
     websocket: Any
     app_slug: str
-    user_id: Optional[str] = None
-    user_email: Optional[str] = None
+    user_id: str | None = None
+    user_email: str | None = None
     connected_at: datetime = None
 
     def __post_init__(self):
@@ -72,17 +73,15 @@ class WebSocketConnectionManager:
             app_slug: App slug for scoping connections (ensures isolation)
         """
         self.app_slug = app_slug
-        self.active_connections: List[WebSocketConnection] = (
-            []
-        )  # List of connection metadata
+        self.active_connections: list[WebSocketConnection] = []  # List of connection metadata
         self._lock = asyncio.Lock()
         logger.debug(f"Initialized WebSocket manager for app: {app_slug}")
 
     async def connect(
         self,
         websocket: Any,
-        user_id: Optional[str] = None,
-        user_email: Optional[str] = None,
+        user_id: str | None = None,
+        user_email: str | None = None,
     ) -> WebSocketConnection:
         """
         Accept and register a WebSocket connection with metadata.
@@ -97,10 +96,7 @@ class WebSocketConnectionManager:
         """
         # Note: websocket should already be accepted by the endpoint handler
         # This is just for tracking - don't accept again
-        if (
-            hasattr(websocket, "client_state")
-            and websocket.client_state.name != "CONNECTED"
-        ):
+        if hasattr(websocket, "client_state") and websocket.client_state.name != "CONNECTED":
             await websocket.accept()
         connection = WebSocketConnection(
             websocket=websocket,
@@ -130,9 +126,7 @@ class WebSocketConnectionManager:
         async def _disconnect():
             async with self._lock:
                 self.active_connections = [
-                    conn
-                    for conn in self.active_connections
-                    if conn.websocket is not websocket
+                    conn for conn in self.active_connections if conn.websocket is not websocket
                 ]
             logger.info(
                 f"WebSocket disconnected for app '{self.app_slug}'. "
@@ -141,9 +135,7 @@ class WebSocketConnectionManager:
 
         asyncio.create_task(_disconnect())
 
-    async def broadcast(
-        self, message: Dict[str, Any], filter_by_user: Optional[str] = None
-    ) -> int:
+    async def broadcast(self, message: dict[str, Any], filter_by_user: str | None = None) -> int:
         """
         Broadcast a message to all connected clients for this app.
 
@@ -205,7 +197,7 @@ class WebSocketConnectionManager:
 
         return sent_count
 
-    async def send_to_connection(self, websocket: Any, message: Dict[str, Any]) -> None:
+    async def send_to_connection(self, websocket: Any, message: dict[str, Any]) -> None:
         """
         Send a message to a specific WebSocket connection.
 
@@ -241,7 +233,7 @@ class WebSocketConnectionManager:
                 logger.debug(f"Error sending message to specific WebSocket: {e}")
             self.disconnect(websocket)
 
-    def get_connections_by_user(self, user_id: str) -> List[WebSocketConnection]:
+    def get_connections_by_user(self, user_id: str) -> list[WebSocketConnection]:
         """
         Get all connections for a specific user.
 
@@ -271,14 +263,12 @@ class WebSocketConnectionManager:
 
 
 # Global registry of WebSocket managers per app (app-level isolation)
-_websocket_managers: Dict[str, WebSocketConnectionManager] = {}
+_websocket_managers: dict[str, WebSocketConnectionManager] = {}
 _manager_lock = asyncio.Lock()
 
 # Global registry of message handlers per app (for listening to client messages)
 # Note: Registration happens synchronously during app startup, so no lock needed
-_message_handlers: Dict[
-    str, Dict[str, Callable[[Any, Dict[str, Any]], Awaitable[None]]]
-] = {}
+_message_handlers: dict[str, dict[str, Callable[[Any, dict[str, Any]], Awaitable[None]]]] = {}
 
 
 async def get_websocket_manager(app_slug: str) -> WebSocketConnectionManager:
@@ -319,7 +309,7 @@ def get_websocket_manager_sync(app_slug: str) -> WebSocketConnectionManager:
 
 async def authenticate_websocket(
     websocket: Any, app_slug: str, require_auth: bool = True
-) -> Tuple[Optional[str], Optional[str]]:
+) -> tuple[str | None, str | None]:
     """
     Authenticate a WebSocket connection.
 
@@ -367,22 +357,16 @@ async def authenticate_websocket(
             else:
                 logger.info(f"WebSocket query_params is empty for app '{app_slug}'")
         else:
-            logger.warning(
-                f"WebSocket has no query_params attribute for app '{app_slug}'"
-            )
+            logger.warning(f"WebSocket has no query_params attribute for app '{app_slug}'")
 
         # If no token in query, try to get from cookies (if available)
         # Check both ws_token (non-httponly, for JS access) and token (httponly)
         if not token:
             if hasattr(websocket, "cookies"):
-                cookie_token = websocket.cookies.get(
-                    "ws_token"
-                ) or websocket.cookies.get("token")
+                cookie_token = websocket.cookies.get("ws_token") or websocket.cookies.get("token")
                 if cookie_token:
                     token = cookie_token
-                    logger.debug(
-                        f"WebSocket token found in cookies for app '{app_slug}'"
-                    )
+                    logger.debug(f"WebSocket token found in cookies for app '{app_slug}'")
             else:
                 logger.debug(f"WebSocket has no cookies attribute for app '{app_slug}'")
 
@@ -404,6 +388,8 @@ async def authenticate_websocket(
                 return None, None  # Signal auth failure
             return None, None
 
+        import jwt
+
         from ..auth.dependencies import SECRET_KEY
         from ..auth.jwt import decode_jwt_token
 
@@ -412,22 +398,16 @@ async def authenticate_websocket(
             user_id = payload.get("sub") or payload.get("user_id")
             user_email = payload.get("email")
 
-            logger.info(
-                f"WebSocket authenticated successfully for app '{app_slug}': {user_email}"
-            )
+            logger.info(f"WebSocket authenticated successfully for app '{app_slug}': {user_email}")
             return user_id, user_email
-        except Exception as decode_error:
-            logger.error(
-                f"JWT decode error for app '{app_slug}': {decode_error}", exc_info=True
-            )
+        except (jwt.ExpiredSignatureError, jwt.InvalidTokenError) as decode_error:
+            logger.error(f"JWT decode error for app '{app_slug}': {decode_error}", exc_info=True)
             raise
 
     except WebSocketDisconnect:
         raise
     except (ValueError, TypeError, AttributeError, KeyError, RuntimeError) as e:
-        logger.error(
-            f"WebSocket authentication failed for app '{app_slug}': {e}", exc_info=True
-        )
+        logger.error(f"WebSocket authentication failed for app '{app_slug}': {e}", exc_info=True)
         if require_auth:
             # Don't close before accepting - return error info instead
             return None, None  # Signal auth failure
@@ -437,7 +417,7 @@ async def authenticate_websocket(
 def register_message_handler(
     app_slug: str,
     endpoint_name: str,
-    handler: Callable[[Any, Dict[str, Any]], Awaitable[None]],
+    handler: Callable[[Any, dict[str, Any]], Awaitable[None]],
 ) -> None:
     """
     Register a message handler for a WebSocket endpoint.
@@ -469,14 +449,12 @@ def register_message_handler(
     if app_slug not in _message_handlers:
         _message_handlers[app_slug] = {}
     _message_handlers[app_slug][endpoint_name] = handler
-    logger.info(
-        f"Registered message handler for app '{app_slug}', endpoint '{endpoint_name}'"
-    )
+    logger.info(f"Registered message handler for app '{app_slug}', endpoint '{endpoint_name}'")
 
 
 def get_message_handler(
     app_slug: str, endpoint_name: str
-) -> Optional[Callable[[Any, Dict[str, Any]], Awaitable[None]]]:
+) -> Callable[[Any, dict[str, Any]], Awaitable[None]] | None:
     """
     Get a registered message handler for an endpoint.
 
@@ -497,7 +475,7 @@ async def _accept_websocket_connection(websocket: Any, app_slug: str) -> None:
         await websocket.accept()
         print(f"✅ [WEBSOCKET ACCEPTED] App: '{app_slug}'")
         logger.info(f"✅ WebSocket accepted for app '{app_slug}'")
-    except Exception as accept_error:
+    except (RuntimeError, ConnectionError, OSError) as accept_error:
         print(f"❌ [WEBSOCKET ACCEPT FAILED] App: '{app_slug}', Error: {accept_error}")
         logger.error(
             f"❌ Failed to accept WebSocket for app '{app_slug}': {accept_error}",
@@ -511,9 +489,7 @@ async def _authenticate_websocket_connection(
 ) -> tuple:
     """Authenticate WebSocket connection and return (user_id, user_email)."""
     try:
-        user_id, user_email = await authenticate_websocket(
-            websocket, app_slug, require_auth
-        )
+        user_id, user_email = await authenticate_websocket(websocket, app_slug, require_auth)
 
         if require_auth and not user_id:
             logger.warning(
@@ -522,9 +498,7 @@ async def _authenticate_websocket_connection(
             try:
                 await websocket.close(code=1008, reason="Authentication required")
             except (WebSocketDisconnect, RuntimeError, OSError) as e:
-                logger.debug(
-                    f"WebSocket already closed during auth failure cleanup: {e}"
-                )
+                logger.debug(f"WebSocket already closed during auth failure cleanup: {e}")
             raise WebSocketDisconnect(code=1008)
 
         return user_id, user_email
@@ -547,21 +521,19 @@ async def _authenticate_websocket_connection(
             exc_info=True,
         )
         try:
-            await websocket.close(
-                code=1011, reason="Internal server error during authentication"
-            )
-        except (WebSocketDisconnect, RuntimeError, OSError) as e:
-            logger.debug(f"WebSocket already closed during auth error cleanup: {e}")
-        raise WebSocketDisconnect(code=1011)
+            await websocket.close(code=1011, reason="Internal server error during authentication")
+        except (WebSocketDisconnect, RuntimeError, OSError) as close_error:
+            logger.debug(f"WebSocket already closed during auth error cleanup: {close_error}")
+        raise WebSocketDisconnect(code=1011) from None
 
 
 async def _handle_websocket_message(
     websocket: Any,
-    message: Dict[str, Any],
+    message: dict[str, Any],
     manager: Any,
     app_slug: str,
     endpoint_name: str,
-    handler: Optional[Callable],
+    handler: Callable | None,
 ) -> bool:
     """Handle incoming WebSocket message. Returns True if should continue, False if disconnect."""
     if message.get("type") == "websocket.disconnect":
@@ -598,7 +570,7 @@ def create_websocket_endpoint(
     app_slug: str,
     path: str,
     endpoint_name: str,
-    handler: Optional[Callable[[Any, Dict[str, Any]], Awaitable[None]]] = None,
+    handler: Callable[[Any, dict[str, Any]], Awaitable[None]] | None = None,
     require_auth: bool = True,
     ping_interval: int = 30,
 ) -> Callable:
@@ -656,9 +628,7 @@ def create_websocket_endpoint(
             file=sys.stderr,
             flush=True,
         )
-        print(
-            f"🔌 [WEBSOCKET HANDLER CALLED] App: '{app_slug}', Path: {path}", flush=True
-        )
+        print(f"🔌 [WEBSOCKET HANDLER CALLED] App: '{app_slug}', Path: {path}", flush=True)
         logger.info(f"🔌 [WEBSOCKET HANDLER CALLED] App: '{app_slug}', Path: {path}")
         connection = None
         try:
@@ -693,9 +663,7 @@ def create_websocket_endpoint(
             )
 
             # Connect with metadata (websocket already accepted)
-            connection = await manager.connect(
-                websocket, user_id=user_id, user_email=user_email
-            )
+            connection = await manager.connect(websocket, user_id=user_id, user_email=user_email)
 
             # Send initial connection confirmation
             await manager.send_to_connection(
@@ -765,9 +733,7 @@ def create_websocket_endpoint(
             TypeError,
             AttributeError,
         ) as e:
-            logger.error(
-                f"WebSocket connection error for app '{app_slug}': {e}", exc_info=True
-            )
+            logger.error(f"WebSocket connection error for app '{app_slug}': {e}", exc_info=True)
         finally:
             if connection:
                 manager.disconnect(websocket)
@@ -776,7 +742,7 @@ def create_websocket_endpoint(
 
 
 async def broadcast_to_app(
-    app_slug: str, message: Dict[str, Any], user_id: Optional[str] = None
+    app_slug: str, message: dict[str, Any], user_id: str | None = None
 ) -> int:
     """
     Convenience function to broadcast a message to all WebSocket clients for an app.

mdb_engine/utils/__init__.py

@@ -4,4 +4,6 @@ Utility functions and helpers for MDB Engine.
 This module provides utility functions used across the MDB Engine codebase.
 """
 
-__all__ = []
+from .mongo import clean_mongo_doc, clean_mongo_docs
+
+__all__ = ["clean_mongo_doc", "clean_mongo_docs"]

mdb_engine/utils/mongo.py

@@ -0,0 +1,117 @@
+"""
+MongoDB utility functions for MDB Engine.
+
+This module provides utility functions for working with MongoDB documents,
+including JSON serialization helpers.
+"""
+
+from typing import Any
+
+
+def clean_mongo_doc(doc: dict[str, Any] | None) -> dict[str, Any] | None:
+    """
+    Convert MongoDB document to JSON-serializable format.
+
+    Recursively converts MongoDB-specific types to JSON-compatible types:
+    - ObjectId -> str
+    - datetime -> ISO format string
+    - Nested dictionaries and lists are processed recursively
+
+    Args:
+        doc: MongoDB document (dict) or None
+
+    Returns:
+        Cleaned document with all MongoDB types converted, or None if input was None
+
+    Example:
+        ```python
+        from mdb_engine.utils import clean_mongo_doc
+
+        # MongoDB document with ObjectId and datetime
+        doc = {
+            "_id": ObjectId("507f1f77bcf86cd799439011"),
+            "name": "John",
+            "created_at": datetime(2024, 1, 1, 12, 0, 0),
+            "nested": {
+                "id": ObjectId("507f1f77bcf86cd799439012")
+            }
+        }
+
+        # Convert to JSON-serializable format
+        cleaned = clean_mongo_doc(doc)
+        # {
+        #     "_id": "507f1f77bcf86cd799439011",
+        #     "name": "John",
+        #     "created_at": "2024-01-01T12:00:00",
+        #     "nested": {
+        #         "id": "507f1f77bcf86cd799439012"
+        #     }
+        # }
+        ```
+    """
+    from datetime import datetime
+
+    from bson import ObjectId
+
+    if doc is None:
+        return None
+
+    if not isinstance(doc, dict):
+        # If it's not a dict, try to convert it
+        if isinstance(doc, ObjectId):
+            return str(doc)
+        elif isinstance(doc, datetime):
+            return doc.isoformat() if hasattr(doc, "isoformat") else str(doc)
+        else:
+            return doc
+
+    cleaned: dict[str, Any] = {}
+
+    for key, value in doc.items():
+        if isinstance(value, ObjectId):
+            cleaned[key] = str(value)
+        elif isinstance(value, datetime):
+            cleaned[key] = value.isoformat() if hasattr(value, "isoformat") else str(value)
+        elif isinstance(value, dict):
+            cleaned[key] = clean_mongo_doc(value)
+        elif isinstance(value, list):
+            cleaned[key] = [
+                clean_mongo_doc(item)
+                if isinstance(item, dict)
+                else (
+                    str(item)
+                    if isinstance(item, ObjectId)
+                    else (item.isoformat() if isinstance(item, datetime) else item)
+                )
+                for item in value
+            ]
+        else:
+            cleaned[key] = value
+
+    return cleaned
+
+
+def clean_mongo_docs(docs: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """
+    Convert a list of MongoDB documents to JSON-serializable format.
+
+    Convenience function that applies clean_mongo_doc to each document in a list.
+
+    Args:
+        docs: List of MongoDB documents
+
+    Returns:
+        List of cleaned documents
+
+    Example:
+        ```python
+        from mdb_engine.utils import clean_mongo_docs
+
+        # List of MongoDB documents
+        docs = await db.collection.find({}).to_list(length=10)
+
+        # Convert all to JSON-serializable format
+        cleaned = clean_mongo_docs(docs)
+        ```
+    """
+    return [clean_mongo_doc(doc) for doc in docs]