asynciomotor/db fixes (preciser types)

Author: HardMax71

backend/app/core/database_context.py

@@ -7,15 +7,21 @@
 from motor.motor_asyncio import (
     AsyncIOMotorClient,
     AsyncIOMotorClientSession,
+    AsyncIOMotorCollection,
+    AsyncIOMotorCursor,
     AsyncIOMotorDatabase,
 )
 from pymongo.errors import ServerSelectionTimeoutError
 
 from app.core.logging import logger
 
 # Python 3.12 type aliases using the new 'type' statement
-type DBClient = AsyncIOMotorClient[Any]
-type Database = AsyncIOMotorDatabase[Any]
+# MongoDocument represents the raw document type returned by Motor operations
+type MongoDocument = dict[str, Any]
+type DBClient = AsyncIOMotorClient[MongoDocument]
+type Database = AsyncIOMotorDatabase[MongoDocument]
+type Collection = AsyncIOMotorCollection[MongoDocument]
+type Cursor = AsyncIOMotorCursor[MongoDocument]
 type DBSession = AsyncIOMotorClientSession
 
 # Type variable for generic database provider
@@ -102,7 +108,7 @@ async def connect(self) -> None:
         # Always explicitly bind to current event loop for consistency
         import asyncio
 
-        client: AsyncIOMotorClient = AsyncIOMotorClient(
+        client: DBClient = AsyncIOMotorClient(
             self._config.mongodb_url,
             serverSelectionTimeoutMS=self._config.server_selection_timeout_ms,
             connectTimeoutMS=self._config.connect_timeout_ms,

backend/app/core/dishka_lifespan.py

@@ -4,8 +4,8 @@
 import redis.asyncio as redis
 from dishka import AsyncContainer
 from fastapi import FastAPI
-from motor.motor_asyncio import AsyncIOMotorDatabase
 
+from app.core.database_context import Database
 from app.core.logging import logger
 from app.core.startup import initialize_metrics_context, initialize_rate_limits
 from app.core.tracing import init_tracing
@@ -64,7 +64,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
     await initialize_event_schemas(schema_registry)
 
     # Initialize database schema at application scope using app-scoped DB
-    database = await container.get(AsyncIOMotorDatabase)
+    database = await container.get(Database)
     schema_manager = SchemaManager(database)
     await schema_manager.apply_all()
     logger.info("Database schema ensured by SchemaManager")

backend/app/core/providers.py

@@ -2,10 +2,10 @@
 
 import redis.asyncio as redis
 from dishka import Provider, Scope, provide
-from motor.motor_asyncio import AsyncIOMotorDatabase
 
 from app.core.database_context import (
     AsyncDatabaseConnection,
+    Database,
     DatabaseConfig,
     create_database_connection,
 )
@@ -105,7 +105,7 @@ async def get_database_connection(self, settings: Settings) -> AsyncIterator[Asy
             await db_connection.disconnect()
 
     @provide
-    def get_database(self, db_connection: AsyncDatabaseConnection) -> AsyncIOMotorDatabase:
+    def get_database(self, db_connection: AsyncDatabaseConnection) -> Database:
         return db_connection.database
 
 
@@ -174,7 +174,7 @@ async def get_kafka_producer(
             await producer.stop()
 
     @provide
-    async def get_dlq_manager(self, database: AsyncIOMotorDatabase) -> AsyncIterator[DLQManager]:
+    async def get_dlq_manager(self, database: Database) -> AsyncIterator[DLQManager]:
         manager = create_dlq_manager(database)
         await manager.start()
         try:
@@ -210,7 +210,7 @@ def get_schema_registry(self) -> SchemaRegistryManager:
     @provide
     async def get_event_store(
             self,
-            database: AsyncIOMotorDatabase,
+            database: Database,
             schema_registry: SchemaRegistryManager
     ) -> EventStore:
         store = create_event_store(
@@ -332,7 +332,7 @@ async def get_sse_kafka_redis_bridge(
     @provide
     def get_sse_repository(
             self,
-            database: AsyncIOMotorDatabase
+            database: Database
     ) -> SSERepository:
         return SSERepository(database)
 
@@ -365,7 +365,7 @@ class AuthProvider(Provider):
     scope = Scope.APP
 
     @provide
-    def get_user_repository(self, database: AsyncIOMotorDatabase) -> UserRepository:
+    def get_user_repository(self, database: Database) -> UserRepository:
         return UserRepository(database)
 
     @provide
@@ -377,11 +377,11 @@ class UserServicesProvider(Provider):
     scope = Scope.APP
 
     @provide
-    def get_user_settings_repository(self, database: AsyncIOMotorDatabase) -> UserSettingsRepository:
+    def get_user_settings_repository(self, database: Database) -> UserSettingsRepository:
         return UserSettingsRepository(database)
 
     @provide
-    def get_event_repository(self, database: AsyncIOMotorDatabase) -> EventRepository:
+    def get_event_repository(self, database: Database) -> EventRepository:
         return EventRepository(database)
 
     @provide
@@ -415,7 +415,7 @@ class AdminServicesProvider(Provider):
     scope = Scope.APP
 
     @provide
-    def get_admin_events_repository(self, database: AsyncIOMotorDatabase) -> AdminEventsRepository:
+    def get_admin_events_repository(self, database: Database) -> AdminEventsRepository:
         return AdminEventsRepository(database)
 
     @provide(scope=Scope.REQUEST)
@@ -427,7 +427,7 @@ def get_admin_events_service(
         return AdminEventsService(admin_events_repository, replay_service)
 
     @provide
-    def get_admin_settings_repository(self, database: AsyncIOMotorDatabase) -> AdminSettingsRepository:
+    def get_admin_settings_repository(self, database: Database) -> AdminSettingsRepository:
         return AdminSettingsRepository(database)
 
     @provide
@@ -438,15 +438,15 @@ def get_admin_settings_service(
         return AdminSettingsService(admin_settings_repository)
 
     @provide
-    def get_admin_user_repository(self, database: AsyncIOMotorDatabase) -> AdminUserRepository:
+    def get_admin_user_repository(self, database: Database) -> AdminUserRepository:
         return AdminUserRepository(database)
 
     @provide
-    def get_saga_repository(self, database: AsyncIOMotorDatabase) -> SagaRepository:
+    def get_saga_repository(self, database: Database) -> SagaRepository:
         return SagaRepository(database)
 
     @provide
-    def get_notification_repository(self, database: AsyncIOMotorDatabase) -> NotificationRepository:
+    def get_notification_repository(self, database: Database) -> NotificationRepository:
         return NotificationRepository(database)
 
     @provide
@@ -482,23 +482,23 @@ class BusinessServicesProvider(Provider):
     scope = Scope.REQUEST
 
     @provide
-    def get_execution_repository(self, database: AsyncIOMotorDatabase) -> ExecutionRepository:
+    def get_execution_repository(self, database: Database) -> ExecutionRepository:
         return ExecutionRepository(database)
 
     @provide
-    def get_resource_allocation_repository(self, database: AsyncIOMotorDatabase) -> ResourceAllocationRepository:
+    def get_resource_allocation_repository(self, database: Database) -> ResourceAllocationRepository:
         return ResourceAllocationRepository(database)
 
     @provide
-    def get_saved_script_repository(self, database: AsyncIOMotorDatabase) -> SavedScriptRepository:
+    def get_saved_script_repository(self, database: Database) -> SavedScriptRepository:
         return SavedScriptRepository(database)
 
     @provide
-    def get_dlq_repository(self, database: AsyncIOMotorDatabase) -> DLQRepository:
+    def get_dlq_repository(self, database: Database) -> DLQRepository:
         return DLQRepository(database)
 
     @provide
-    def get_replay_repository(self, database: AsyncIOMotorDatabase) -> ReplayRepository:
+    def get_replay_repository(self, database: Database) -> ReplayRepository:
         return ReplayRepository(database)
 
     @provide
@@ -623,5 +623,5 @@ class ResultProcessorProvider(Provider):
     scope = Scope.APP
 
     @provide
-    def get_execution_repository(self, database: AsyncIOMotorDatabase) -> ExecutionRepository:
+    def get_execution_repository(self, database: Database) -> ExecutionRepository:
         return ExecutionRepository(database)

backend/app/db/repositories/admin/admin_events_repository.py

@@ -1,9 +1,9 @@
 from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, List
 
-from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase
 from pymongo import ReturnDocument
 
+from app.core.database_context import Collection, Database
 from app.domain.admin import (
     ReplayQuery,
     ReplaySession,
@@ -43,18 +43,18 @@
 class AdminEventsRepository:
     """Repository for admin event operations using domain models."""
 
-    def __init__(self, db: AsyncIOMotorDatabase):
+    def __init__(self, db: Database):
         self.db = db
-        self.events_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.EVENTS)
-        self.event_store_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.EVENT_STORE)
+        self.events_collection: Collection = self.db.get_collection(CollectionNames.EVENTS)
+        self.event_store_collection: Collection = self.db.get_collection(CollectionNames.EVENT_STORE)
         # Bind related collections used by this repository
-        self.executions_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.EXECUTIONS)
-        self.events_archive_collection: AsyncIOMotorCollection = self.db.get_collection(
+        self.executions_collection: Collection = self.db.get_collection(CollectionNames.EXECUTIONS)
+        self.events_archive_collection: Collection = self.db.get_collection(
             CollectionNames.EVENTS_ARCHIVE
         )
         self.replay_mapper = ReplaySessionMapper()
         self.replay_query_mapper = ReplayQueryMapper()
-        self.replay_sessions_collection: AsyncIOMotorCollection = self.db.get_collection(
+        self.replay_sessions_collection: Collection = self.db.get_collection(
             CollectionNames.REPLAY_SESSIONS)
         self.mapper = EventMapper()
         self.summary_mapper = EventSummaryMapper()

backend/app/db/repositories/admin/admin_settings_repository.py

@@ -1,7 +1,6 @@
 from datetime import datetime, timezone
 
-from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase
-
+from app.core.database_context import Collection, Database
 from app.core.logging import logger
 from app.domain.admin import (
     AuditAction,
@@ -12,10 +11,10 @@
 
 
 class AdminSettingsRepository:
-    def __init__(self, db: AsyncIOMotorDatabase):
+    def __init__(self, db: Database):
         self.db = db
-        self.settings_collection: AsyncIOMotorCollection = self.db.get_collection("system_settings")
-        self.audit_log_collection: AsyncIOMotorCollection = self.db.get_collection("audit_log")
+        self.settings_collection: Collection = self.db.get_collection("system_settings")
+        self.audit_log_collection: Collection = self.db.get_collection("audit_log")
         self.settings_mapper = SettingsMapper()
         self.audit_mapper = AuditLogMapper()
 

backend/app/db/repositories/admin/admin_user_repository.py

@@ -1,7 +1,6 @@
 from datetime import datetime, timezone
 
-from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase
-
+from app.core.database_context import Collection, Database
 from app.core.security import SecurityService
 from app.domain.enums import UserRole
 from app.domain.events.event_models import CollectionNames
@@ -17,17 +16,17 @@
 
 
 class AdminUserRepository:
-    def __init__(self, db: AsyncIOMotorDatabase):
+    def __init__(self, db: Database):
         self.db = db
-        self.users_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.USERS)
+        self.users_collection: Collection = self.db.get_collection(CollectionNames.USERS)
 
         # Related collections used by this repository (e.g., cascade deletes)
-        self.executions_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.EXECUTIONS)
-        self.saved_scripts_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.SAVED_SCRIPTS)
-        self.notifications_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.NOTIFICATIONS)
-        self.user_settings_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.USER_SETTINGS)
-        self.events_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.EVENTS)
-        self.sagas_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.SAGAS)
+        self.executions_collection: Collection = self.db.get_collection(CollectionNames.EXECUTIONS)
+        self.saved_scripts_collection: Collection = self.db.get_collection(CollectionNames.SAVED_SCRIPTS)
+        self.notifications_collection: Collection = self.db.get_collection(CollectionNames.NOTIFICATIONS)
+        self.user_settings_collection: Collection = self.db.get_collection(CollectionNames.USER_SETTINGS)
+        self.events_collection: Collection = self.db.get_collection(CollectionNames.EVENTS)
+        self.sagas_collection: Collection = self.db.get_collection(CollectionNames.SAGAS)
         self.security_service = SecurityService()
         self.mapper = UserMapper()
 

backend/app/db/repositories/dlq_repository.py

@@ -1,8 +1,7 @@
 from datetime import datetime, timezone
 from typing import Dict, List, Mapping
 
-from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase
-
+from app.core.database_context import Collection, Database
 from app.core.logging import logger
 from app.dlq import (
     AgeStatistics,
@@ -24,9 +23,9 @@
 
 
 class DLQRepository:
-    def __init__(self, db: AsyncIOMotorDatabase):
+    def __init__(self, db: Database):
         self.db = db
-        self.dlq_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.DLQ_MESSAGES)
+        self.dlq_collection: Collection = self.db.get_collection(CollectionNames.DLQ_MESSAGES)
 
     async def get_dlq_stats(self) -> DLQStatistics:
         # Get counts by status

backend/app/db/repositories/event_repository.py

@@ -3,9 +3,9 @@
 from types import MappingProxyType
 from typing import Any, AsyncIterator, Mapping
 
-from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase
 from pymongo import ASCENDING, DESCENDING
 
+from app.core.database_context import Collection, Database
 from app.core.logging import logger
 from app.core.tracing import EventAttributes
 from app.core.tracing.utils import add_span_attributes
@@ -25,10 +25,10 @@
 
 
 class EventRepository:
-    def __init__(self, database: AsyncIOMotorDatabase) -> None:
+    def __init__(self, database: Database) -> None:
         self.database = database
         self.mapper = EventMapper()
-        self._collection: AsyncIOMotorCollection = self.database.get_collection(CollectionNames.EVENTS)
+        self._collection: Collection = self.database.get_collection(CollectionNames.EVENTS)
 
     def _build_time_filter(
             self,

backend/app/db/repositories/execution_repository.py

@@ -1,18 +1,17 @@
 from datetime import datetime, timezone
 
-from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase
-
+from app.core.database_context import Collection, Database
 from app.core.logging import logger
 from app.domain.enums.execution import ExecutionStatus
 from app.domain.events.event_models import CollectionNames
 from app.domain.execution import DomainExecution, ExecutionResultDomain, ResourceUsageDomain
 
 
 class ExecutionRepository:
-    def __init__(self, db: AsyncIOMotorDatabase):
+    def __init__(self, db: Database):
         self.db = db
-        self.collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.EXECUTIONS)
-        self.results_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.EXECUTION_RESULTS)
+        self.collection: Collection = self.db.get_collection(CollectionNames.EXECUTIONS)
+        self.results_collection: Collection = self.db.get_collection(CollectionNames.EXECUTION_RESULTS)
 
     async def create_execution(self, execution: DomainExecution) -> DomainExecution:
         execution_dict = {

backend/app/db/repositories/notification_repository.py

@@ -1,8 +1,8 @@
 from datetime import UTC, datetime, timedelta
 
-from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase
 from pymongo import ASCENDING, DESCENDING, IndexModel
 
+from app.core.database_context import Collection, Database
 from app.core.logging import logger
 from app.domain.enums.notification import (
     NotificationChannel,
@@ -16,11 +16,11 @@
 
 
 class NotificationRepository:
-    def __init__(self, database: AsyncIOMotorDatabase):
-        self.db: AsyncIOMotorDatabase = database
+    def __init__(self, database: Database):
+        self.db: Database = database
 
-        self.notifications_collection: AsyncIOMotorCollection = self.db.get_collection(CollectionNames.NOTIFICATIONS)
-        self.subscriptions_collection: AsyncIOMotorCollection = self.db.get_collection(
+        self.notifications_collection: Collection = self.db.get_collection(CollectionNames.NOTIFICATIONS)
+        self.subscriptions_collection: Collection = self.db.get_collection(
             CollectionNames.NOTIFICATION_SUBSCRIPTIONS)
         self.mapper = NotificationMapper()