mdb-engine 0.1.6__py3-none-any.whl → 0.2.0__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, Dict, Generic, Optional, Type, 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: Optional[Dict[str, Type[Entity]]] = 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: Optional[Type[T]] = 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

@@ -72,9 +72,7 @@ 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}")
 
@@ -97,10 +95,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 +125,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 +134,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: Optional[str] = None) -> int:
         """
         Broadcast a message to all connected clients for this app.
 
@@ -276,9 +267,7 @@ _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:
@@ -367,22 +356,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 +387,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 +397,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
@@ -469,9 +448,7 @@ 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(
@@ -497,7 +474,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 +488,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 +497,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,12 +520,10 @@ 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(
@@ -656,9 +627,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 +662,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 +732,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)