mdb-engine 0.2.1__py3-none-any.whl → 0.2.4__py3-none-any.whl

mdb_engine/core/types.py

@@ -8,7 +8,7 @@ throughout the codebase.
 This module is part of MDB_ENGINE - MongoDB Engine.
 """
 
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, TypedDict, Union
+from typing import TYPE_CHECKING, Any, Literal, TypedDict
 
 if TYPE_CHECKING:
     from ..memory import Mem0MemoryService
@@ -41,20 +41,20 @@ class IndexDefinitionDict(TypedDict, total=False):
         "partial",
         "hybrid",
     ]
-    keys: Union[Dict[str, Union[int, str]], List[List[Union[str, int]]]]
+    keys: dict[str, int | str] | list[list[str | int]]
     unique: bool
     sparse: bool
     background: bool
     expireAfterSeconds: int
-    partialFilterExpression: Dict[str, Any]
-    weights: Dict[str, int]  # For text indexes
+    partialFilterExpression: dict[str, Any]
+    weights: dict[str, int]  # For text indexes
     default_language: str  # For text indexes
     language_override: str  # For text indexes
     textIndexVersion: int  # For text indexes
     # Vector search specific
-    fields: List[Dict[str, Any]]  # For vectorSearch indexes
+    fields: list[dict[str, Any]]  # For vectorSearch indexes
     # Search index specific
-    mappings: Dict[str, Any]  # For search indexes
+    mappings: dict[str, Any]  # For search indexes
     # Geospatial specific
     bucketSize: float  # For geoHaystack
     # TTL specific
@@ -79,12 +79,12 @@ class AuthAuthorizationDict(TypedDict, total=False):
     model: str
     policies_collection: str
     link_users_roles: bool
-    default_roles: List[str]
+    default_roles: list[str]
     # OSO-specific
-    api_key: Optional[str]
-    url: Optional[str]
-    initial_roles: List[Dict[str, str]]
-    initial_policies: List[Dict[str, str]]
+    api_key: str | None
+    url: str | None
+    initial_roles: list[dict[str, str]]
+    initial_policies: list[dict[str, str]]
 
 
 class AuthPolicyDict(TypedDict, total=False):
@@ -93,12 +93,12 @@ class AuthPolicyDict(TypedDict, total=False):
     required: bool
     provider: Literal["casbin", "oso", "custom"]
     authorization: AuthAuthorizationDict
-    allowed_roles: List[str]
-    allowed_users: List[str]
-    denied_users: List[str]
-    required_permissions: List[str]
+    allowed_roles: list[str]
+    allowed_users: list[str]
+    denied_users: list[str]
+    required_permissions: list[str]
     custom_resource: str
-    custom_actions: List[Literal["access", "read", "write", "admin"]]
+    custom_actions: list[Literal["access", "read", "write", "admin"]]
     allow_anonymous: bool
     owner_can_access: bool
 
@@ -111,7 +111,7 @@ class DemoUserDict(TypedDict, total=False):
     role: str
     auto_create: bool
     link_to_platform: bool
-    extra_data: Dict[str, Any]
+    extra_data: dict[str, Any]
 
 
 class UsersConfigDict(TypedDict, total=False):
@@ -126,7 +126,7 @@ class UsersConfigDict(TypedDict, total=False):
     link_platform_users: bool
     anonymous_user_prefix: str
     user_id_field: str
-    demo_users: List[DemoUserDict]
+    demo_users: list[DemoUserDict]
     auto_link_platform_demo: bool
     demo_user_seed_strategy: Literal["auto", "manual", "disabled"]
     enable_demo_user_access: bool
@@ -312,7 +312,7 @@ class MetricsConfigDict(TypedDict, total=False):
     enabled: bool
     collect_operation_metrics: bool
     collect_performance_metrics: bool
-    custom_metrics: List[str]
+    custom_metrics: list[str]
 
 
 class LoggingConfigDict(TypedDict, total=False):
@@ -341,11 +341,11 @@ class CORSConfigDict(TypedDict, total=False):
     """CORS configuration."""
 
     enabled: bool
-    allow_origins: List[str]
+    allow_origins: list[str]
     allow_credentials: bool
-    allow_methods: List[Literal["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD", "*"]]
-    allow_headers: List[str]
-    expose_headers: List[str]
+    allow_methods: list[Literal["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD", "*"]]
+    allow_headers: list[str]
+    expose_headers: list[str]
     max_age: int
 
 
@@ -371,22 +371,22 @@ class ManifestDict(TypedDict, total=False):
     schema_version: str
     slug: str
     name: str
-    description: Optional[str]
+    description: str | None
     status: Literal["active", "draft", "archived", "inactive"]
     auth_required: bool  # Backward compatibility
-    auth: Optional[AuthConfigDict]
-    token_management: Optional[TokenManagementDict]
-    data_scope: List[str]
-    pip_deps: List[str]
-    managed_indexes: Optional[ManagedIndexesDict]
-    collection_settings: Optional[Dict[str, Dict[str, Any]]]
-    websockets: Optional[WebSocketsDict]
-    embedding_config: Optional[EmbeddingConfigDict]
-    memory_config: Optional[MemoryConfigDict]
-    cors: Optional[CORSConfigDict]
-    observability: Optional[ObservabilityConfigDict]
-    initial_data: Optional[InitialDataDict]
-    developer_id: Optional[str]
+    auth: AuthConfigDict | None
+    token_management: TokenManagementDict | None
+    data_scope: list[str]
+    pip_deps: list[str]
+    managed_indexes: ManagedIndexesDict | None
+    collection_settings: dict[str, dict[str, Any]] | None
+    websockets: WebSocketsDict | None
+    embedding_config: EmbeddingConfigDict | None
+    memory_config: MemoryConfigDict | None
+    cors: CORSConfigDict | None
+    observability: ObservabilityConfigDict | None
+    initial_data: InitialDataDict | None
+    developer_id: str | None
 
 
 # ============================================================================
@@ -398,13 +398,13 @@ class HealthStatusDict(TypedDict, total=False):
     """Health status response."""
 
     status: Literal["healthy", "degraded", "unhealthy"]
-    checks: Dict[str, Dict[str, Any]]
+    checks: dict[str, dict[str, Any]]
     timestamp: str
 
 
 class MetricsDict(TypedDict, total=False):
     """Metrics response."""
 
-    operations: Dict[str, Dict[str, Any]]
-    summary: Dict[str, Any]
+    operations: dict[str, dict[str, Any]]
+    summary: dict[str, Any]
     timestamp: str

mdb_engine/database/abstraction.py

@@ -24,7 +24,8 @@ Usage:
 """
 
 import logging
-from typing import Any, Callable, Dict, List, Optional
+from collections.abc import Callable
+from typing import Any
 
 from ..exceptions import MongoDBEngineError
 from .scoped_wrapper import ScopedMongoWrapper
@@ -88,8 +89,8 @@ class Collection:
         self._collection = scoped_collection
 
     async def find_one(
-        self, filter: Optional[Dict[str, Any]] = None, *args, **kwargs
-    ) -> Optional[Dict[str, Any]]:
+        self, filter: dict[str, Any] | None = None, *args, **kwargs
+    ) -> dict[str, Any] | None:
         """
         Find a single document matching the filter.
 
@@ -120,7 +121,7 @@ class Collection:
                 context={"operation": "find_one"},
             ) from e
 
-    def find(self, filter: Optional[Dict[str, Any]] = None, *args, **kwargs) -> AsyncIOMotorCursor:
+    def find(self, filter: dict[str, Any] | None = None, *args, **kwargs) -> AsyncIOMotorCursor:
         """
         Find documents matching the filter.
 
@@ -146,7 +147,7 @@ class Collection:
         """
         return self._collection.find(filter or {}, *args, **kwargs)
 
-    async def insert_one(self, document: Dict[str, Any], *args, **kwargs) -> InsertOneResult:
+    async def insert_one(self, document: dict[str, Any], *args, **kwargs) -> InsertOneResult:
         """
         Insert a single document.
 
@@ -178,7 +179,7 @@ class Collection:
             ) from e
 
     async def insert_many(
-        self, documents: List[Dict[str, Any]], *args, **kwargs
+        self, documents: list[dict[str, Any]], *args, **kwargs
     ) -> InsertManyResult:
         """
         Insert multiple documents at once.
@@ -213,7 +214,7 @@ class Collection:
             ) from e
 
     async def update_one(
-        self, filter: Dict[str, Any], update: Dict[str, Any], *args, **kwargs
+        self, filter: dict[str, Any], update: dict[str, Any], *args, **kwargs
     ) -> UpdateResult:
         """
         Update a single document matching the filter.
@@ -251,7 +252,7 @@ class Collection:
             ) from e
 
     async def update_many(
-        self, filter: Dict[str, Any], update: Dict[str, Any], *args, **kwargs
+        self, filter: dict[str, Any], update: dict[str, Any], *args, **kwargs
     ) -> UpdateResult:
         """
         Update multiple documents matching the filter.
@@ -288,7 +289,7 @@ class Collection:
             ) from e
 
     async def replace_one(
-        self, filter: Dict[str, Any], replacement: Dict[str, Any], *args, **kwargs
+        self, filter: dict[str, Any], replacement: dict[str, Any], *args, **kwargs
     ) -> UpdateResult:
         """
         Replace a single document matching the filter.
@@ -367,7 +368,7 @@ class Collection:
                 context={"operation": "replace_one"},
             ) from e
 
-    async def delete_one(self, filter: Dict[str, Any], *args, **kwargs) -> DeleteResult:
+    async def delete_one(self, filter: dict[str, Any], *args, **kwargs) -> DeleteResult:
         """
         Delete a single document matching the filter.
 
@@ -400,7 +401,7 @@ class Collection:
             ) from e
 
     async def delete_many(
-        self, filter: Optional[Dict[str, Any]] = None, *args, **kwargs
+        self, filter: dict[str, Any] | None = None, *args, **kwargs
     ) -> DeleteResult:
         """
         Delete multiple documents matching the filter.
@@ -432,9 +433,7 @@ class Collection:
                 context={"operation": "delete_many"},
             ) from e
 
-    async def count_documents(
-        self, filter: Optional[Dict[str, Any]] = None, *args, **kwargs
-    ) -> int:
+    async def count_documents(self, filter: dict[str, Any] | None = None, *args, **kwargs) -> int:
         """
         Count documents matching the filter.
 
@@ -464,7 +463,7 @@ class Collection:
                 context={"operation": "count_documents"},
             ) from e
 
-    def aggregate(self, pipeline: List[Dict[str, Any]], *args, **kwargs) -> AsyncIOMotorCursor:
+    def aggregate(self, pipeline: list[dict[str, Any]], *args, **kwargs) -> AsyncIOMotorCursor:
         """
         Perform aggregation pipeline.
 
@@ -520,7 +519,7 @@ class AppDB:
             raise RuntimeError("ScopedMongoWrapper is not available. Check imports.")
 
         self._wrapper = scoped_wrapper
-        self._collection_cache: Dict[str, Collection] = {}
+        self._collection_cache: dict[str, Collection] = {}
 
     def collection(self, name: str) -> Collection:
         """

mdb_engine/database/connection.py

@@ -23,7 +23,7 @@ Usage:
 import logging
 import os
 import threading
-from typing import Any, Dict, Optional
+from typing import Any
 
 from motor.motor_asyncio import AsyncIOMotorClient
 from pymongo.errors import (
@@ -36,7 +36,7 @@ from pymongo.errors import (
 logger = logging.getLogger(__name__)
 
 # Global singleton instance
-_shared_client: Optional[AsyncIOMotorClient] = None
+_shared_client: AsyncIOMotorClient | None = None
 # Use threading.Lock for cross-thread safety in multi-threaded environments
 # asyncio.Lock isn't sufficient for thread-safe initialization
 _init_lock = threading.Lock()
@@ -44,8 +44,8 @@ _init_lock = threading.Lock()
 
 def get_shared_mongo_client(
     mongo_uri: str,
-    max_pool_size: Optional[int] = None,
-    min_pool_size: Optional[int] = None,
+    max_pool_size: int | None = None,
+    min_pool_size: int | None = None,
     server_selection_timeout_ms: int = 5000,
     max_idle_time_ms: int = 45000,
     retry_writes: bool = True,
@@ -203,8 +203,8 @@ def register_client_for_metrics(client: AsyncIOMotorClient) -> None:
 
 
 async def get_pool_metrics(
-    client: Optional[AsyncIOMotorClient] = None,
-) -> Dict[str, Any]:
+    client: AsyncIOMotorClient | None = None,
+) -> dict[str, Any]:
     """
     Gets connection pool metrics for monitoring.
     Returns information about pool size, active connections, etc.
@@ -247,7 +247,7 @@ async def get_pool_metrics(
     }
 
 
-async def _get_client_pool_metrics(client: AsyncIOMotorClient) -> Dict[str, Any]:
+async def _get_client_pool_metrics(client: AsyncIOMotorClient) -> dict[str, Any]:
     """
     Internal helper to get pool metrics from a specific client.
 
@@ -299,10 +299,32 @@ async def _get_client_pool_metrics(client: AsyncIOMotorClient) -> Dict[str, Any]
 
         try:
             server_status = await client.admin.command("serverStatus")
-            connections = server_status.get("connections", {})
-            current_connections = connections.get("current", 0)
-            available_connections = connections.get("available", 0)
-            total_created = connections.get("totalCreated", 0)
+            if not isinstance(server_status, dict):
+                # Mock or invalid response - skip connection metrics
+                current_connections = None
+                available_connections = None
+                total_created = None
+            else:
+                connections = server_status.get("connections", {})
+                if not isinstance(connections, dict):
+                    # Mock or invalid response - skip connection metrics
+                    current_connections = None
+                    available_connections = None
+                    total_created = None
+                else:
+                    # Get values, ensuring they're numeric (not MagicMocks)
+                    current_raw = connections.get("current", 0)
+                    available_raw = connections.get("available", 0)
+                    total_raw = connections.get("totalCreated", 0)
+
+                    # Only use if actually numeric
+                    current_connections = (
+                        int(current_raw) if isinstance(current_raw, int | float) else None
+                    )
+                    available_connections = (
+                        int(available_raw) if isinstance(available_raw, int | float) else None
+                    )
+                    total_created = int(total_raw) if isinstance(total_raw, int | float) else None
         except (
             OperationFailure,
             ConnectionFailure,
@@ -330,7 +352,13 @@ async def _get_client_pool_metrics(client: AsyncIOMotorClient) -> Dict[str, Any]
             metrics["total_connections_created"] = total_created
 
         # Calculate pool usage if we have max_pool_size and current connections
-        if max_pool_size and current_connections is not None:
+        # Ensure both are numeric (not MagicMock or other types)
+        if (
+            max_pool_size
+            and current_connections is not None
+            and isinstance(max_pool_size, int | float)
+            and isinstance(current_connections, int | float)
+        ):
             usage_percent = (current_connections / max_pool_size) * 100
             metrics["pool_usage_percent"] = round(usage_percent, 2)
             metrics["active_connections"] = current_connections  # Alias for compatibility

mdb_engine/database/query_validator.py

@@ -14,7 +14,7 @@ Security Features:
 
 import logging
 import re
-from typing import Any, Dict, List, Optional, Set
+from typing import Any
 
 from ..constants import (
     DANGEROUS_OPERATORS,
@@ -46,7 +46,7 @@ class QueryValidator:
         max_pipeline_stages: int = MAX_PIPELINE_STAGES,
         max_regex_length: int = MAX_REGEX_LENGTH,
         max_regex_complexity: int = MAX_REGEX_COMPLEXITY,
-        dangerous_operators: Optional[Set[str]] = None,
+        dangerous_operators: set[str] | None = None,
     ):
         """
         Initialize the query validator.
@@ -80,7 +80,7 @@ class QueryValidator:
                 else DANGEROUS_OPERATORS
             )
 
-    def validate_filter(self, filter: Optional[Dict[str, Any]], path: str = "") -> None:
+    def validate_filter(self, filter: dict[str, Any] | None, path: str = "") -> None:
         """
         Validate a MongoDB query filter.
 
@@ -105,7 +105,7 @@ class QueryValidator:
         self._check_dangerous_operators(filter, path)
         self._check_query_depth(filter, path, depth=0)
 
-    def validate_pipeline(self, pipeline: List[Dict[str, Any]]) -> None:
+    def validate_pipeline(self, pipeline: list[dict[str, Any]]) -> None:
         """
         Validate an aggregation pipeline.
 
@@ -201,7 +201,7 @@ class QueryValidator:
                 path=path,
             ) from e
 
-    def validate_sort(self, sort: Optional[Any]) -> None:
+    def validate_sort(self, sort: Any | None) -> None:
         """
         Validate a sort specification.
 
@@ -228,7 +228,7 @@ class QueryValidator:
             )
 
     def _check_dangerous_operators(
-        self, query: Dict[str, Any], path: str = "", depth: int = 0
+        self, query: dict[str, Any], path: str = "", depth: int = 0
     ) -> None:
         """
         Recursively check for dangerous operators in a query.
@@ -287,7 +287,7 @@ class QueryValidator:
                 # Direct $regex value (less common but possible)
                 self.validate_regex(value, current_path)
 
-    def _check_query_depth(self, query: Dict[str, Any], path: str = "", depth: int = 0) -> None:
+    def _check_query_depth(self, query: dict[str, Any], path: str = "", depth: int = 0) -> None:
         """
         Check query nesting depth.
 
@@ -348,7 +348,7 @@ class QueryValidator:
 
         return complexity
 
-    def _extract_sort_fields(self, sort: Any) -> List[str]:
+    def _extract_sort_fields(self, sort: Any) -> list[str]:
         """
         Extract field names from a sort specification.
 

mdb_engine/database/resource_limiter.py

@@ -12,7 +12,7 @@ Features:
 """
 
 import logging
-from typing import Any, Dict, Optional
+from typing import Any
 
 from bson import encode as bson_encode
 from bson.errors import InvalidDocument
@@ -65,8 +65,8 @@ class ResourceLimiter:
         self.max_document_size = max_document_size
 
     def enforce_query_timeout(
-        self, kwargs: Dict[str, Any], default_timeout: Optional[int] = None
-    ) -> Dict[str, Any]:
+        self, kwargs: dict[str, Any], default_timeout: int | None = None
+    ) -> dict[str, Any]:
         """
         Enforce query timeout by adding maxTimeMS if not present.
 
@@ -97,7 +97,7 @@ class ResourceLimiter:
 
         return kwargs
 
-    def enforce_result_limit(self, limit: Optional[int], max_limit: Optional[int] = None) -> int:
+    def enforce_result_limit(self, limit: int | None, max_limit: int | None = None) -> int:
         """
         Enforce maximum result limit.
 
@@ -122,7 +122,7 @@ class ResourceLimiter:
 
         return limit
 
-    def enforce_batch_size(self, batch_size: Optional[int], max_batch: Optional[int] = None) -> int:
+    def enforce_batch_size(self, batch_size: int | None, max_batch: int | None = None) -> int:
         """
         Enforce maximum batch size for cursor operations.
 
@@ -147,7 +147,7 @@ class ResourceLimiter:
 
         return batch_size
 
-    def validate_document_size(self, document: Dict[str, Any]) -> None:
+    def validate_document_size(self, document: dict[str, Any]) -> None:
         """
         Validate that a document doesn't exceed size limits.
 
@@ -180,7 +180,7 @@ class ResourceLimiter:
             # MongoDB will catch this anyway during actual insert
             logger.warning(f"Could not encode document as BSON for size validation: {e}")
 
-    def validate_documents_size(self, documents: list[Dict[str, Any]]) -> None:
+    def validate_documents_size(self, documents: list[dict[str, Any]]) -> None:
         """
         Validate that multiple documents don't exceed size limits.