uof with di for handlers

Author: vladimirgubarik

app/application/repository/interfaces.py

@@ -6,6 +6,26 @@
 from app.domain.entities.sample_record import SampleRecord
 
 
+class IUnitOfWork(ABC):
+    """Interface for unit of work operations."""
+
+    @abstractmethod
+    async def __aenter__(self):
+        return self
+
+    @abstractmethod
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+
+class ISampleRecordUnitOfWork(IUnitOfWork, ABC):
+    """Interface for record unit of work operations."""
+
+    @abstractmethod
+    def get_sample_record_repository(self) -> "ISampleRecordRepository":
+        pass
+
+
 class ISampleRecordRepository(ABC):
     """Interface for record repository operations."""
 
@@ -60,3 +80,8 @@ async def get_by_id(self, record_id: int) -> SampleRecord:
     async def get_all(self) -> List[SampleRecord]:
         """Get all records from the database"""
         pass
+
+    @abstractmethod
+    async def commit(self) -> None:
+        """Commit changes to the database"""
+        pass

app/application/use_cases/record_use_cases.py

@@ -23,8 +23,11 @@ async def create_record(
         domain_object = SampleRecord(
             record_data=request_object.record_data, name=request_object.name
         )
-        created_record = await self._repo.create(domain_object)
-        return SampleRecordResponseSchema.from_orm(created_record)
+        created_record = SampleRecordResponseSchema.from_orm(
+            await self._repo.create(domain_object)
+        )
+
+        return created_record
 
     async def update_record(
         self, update_request: SampleRecordUpdateRequestSchema
@@ -33,10 +36,12 @@ async def update_record(
         domain_object = SampleRecord(
             record_data=update_request.record_data,
             id=update_request.id,
-            name=update_request.name
+            name=update_request.name,
+        )
+        updated_record = SampleRecordResponseSchema.from_orm(
+            await self._repo.update(domain_object)
         )
-        updated_record = await self._repo.update(domain_object)
-        return SampleRecordResponseSchema.from_orm(updated_record)
+        return updated_record
 
     async def delete_record(self, record_id: int) -> None:
         """Delete a record."""

app/infrastructure/containers.py

@@ -1,4 +1,5 @@
 import asyncio
+from importlib import import_module
 
 from dependency_injector import containers, providers
 from dependency_injector.providers import Callable, Factory, Singleton
@@ -14,7 +15,10 @@
     PubsubExceptionHandler,
 )
 from app.infrastructure.repositories.caching.redis_repo import RedisRepo
-from app.infrastructure.repositories.sample_record import SampleRecordRepository
+from app.infrastructure.repositories.unit_of_work import (
+    ReadOnlySampleRecordUnitOfWork,
+    WriteSampleRecordUnitOfWork,
+)
 from app.logger import logger
 
 from app.presentation.bot.error_handlers.internal_error_handler import (
@@ -27,32 +31,18 @@
 
 
 class BotSampleRecordCommandContainer(containers.DeclarativeContainer):
-    record_use_cases_factory = Callable(
-        lambda session: SampleRecordUseCases(
-            record_repo=SampleRecordRepository(session=session)
-        )
-    )
+    session_factory = providers.Dependency()
 
+    ro_unit_of_work: Factory[ReadOnlySampleRecordUnitOfWork] = Factory(
+        ReadOnlySampleRecordUnitOfWork, session_factory
+    )
+    rw_unit_of_work: Factory[WriteSampleRecordUnitOfWork] = Factory(
+        WriteSampleRecordUnitOfWork, session_factory
+    )
 
-# class StorageContainer(containers.DeclarativeContainer):
-#     wiring_config = containers.WiringConfiguration(
-#         modules=["app.presentation.bot.commands.sample_records"]
-#     )
-#
-#     # Provider that returns a factory to create sessions
-#     session_factory = Factory(build_db_session_factory)
-#
-#     # Provider that creates a session (e.g., AsyncSession instance)
-#     session = Resource(session_factory)
-#
-#     # Provider that creates the SampleRecordUseCases, injecting the session
-#     record_use_cases = Factory(
-#         SampleRecordUseCases,
-#         record_repo=Factory(
-#             SampleRecordRepository,
-#             session=session
-#         )
-#     )
+    record_use_cases_factory = Callable(
+        lambda repository: SampleRecordUseCases(record_repo=repository)
+    )
 
 
 class CallbackTaskManager:
@@ -117,11 +107,16 @@ class ApplicationStartupContainer(containers.DeclarativeContainer):
         {} if not settings.RAISE_BOT_EXCEPTIONS else {Exception: internal_error_handler}
     )
 
-    from app.presentation.bot.commands import common, sample_record
+    # Ленивая загрузка коллекторов
+    @staticmethod
+    def get_collectors():
+        common = import_module("app.presentation.bot.commands.common")
+        sample_record = import_module("app.presentation.bot.commands.sample_record")
+        return [common.collector, sample_record.collector]
 
     bot = providers.Singleton(
         Bot,
-        collectors=[common.collector, sample_record.collector],
+        collectors=Callable(get_collectors),
         bot_accounts=settings.BOT_CREDENTIALS,
         exception_handlers=exception_handlers,  # type: ignore
         default_callback_timeout=settings.BOTX_CALLBACK_TIMEOUT_IN_SECONDS,
@@ -146,7 +141,6 @@ class ApplicationStartupContainer(containers.DeclarativeContainer):
     )
 
 
-
 class WorkerStartupContainer(containers.DeclarativeContainer):
     redis_client = Singleton(lambda: aioredis.from_url(settings.REDIS_DSN))
 

app/infrastructure/db/sqlalchemy.py

@@ -41,7 +41,6 @@ def make_url_sync(url: str) -> str:
 
 Base = declarative_base(metadata=MetaData(naming_convention=convention))
 
-
 @lru_cache(maxsize=1)
 def get_engine() -> AsyncEngine:
     """Lazily initialize and cache a single SQLAlchemy async engine."""
@@ -54,33 +53,31 @@ def get_engine() -> AsyncEngine:
     )
 
 
-session_factory = async_sessionmaker(bind=get_engine(), expire_on_commit=False)
-
-
-async def close_db_connections() -> None:
-    await get_engine().dispose()
-
-
-def provide_session(func: Callable) -> Callable:
-    """
-    Provides a database session to an async function if one is not already passed.
-
-    :param func: The asynchronous function to wrap. It must accept a `session`
-        keyword argument.
-    :return: The wrapped function with automatic session provisioning."""
-
-    @wraps(func)
-    async def wrapper(*args: Any, **kwargs: Any) -> Any:
-        if kwargs.get("session"):
-            return await func(*args, **kwargs)
-
-        async with session_factory() as session:
-            try:
-                return await func(*args, **kwargs, session=session)
-            except Exception:
-                await session.rollback()
-                raise
-            finally:
-                await session.close()
-
-    return wrapper
+def get_session_factory() -> async_sessionmaker:
+    engine = get_engine()
+    return async_sessionmaker(bind=engine, expire_on_commit=False)
+
+#
+# def provide_session(func: Callable) -> Callable:
+#     """
+#     Provides a database session to an async function if one is not already passed.
+#
+#     :param func: The asynchronous function to wrap. It must accept a `session`
+#         keyword argument.
+#     :return: The wrapped function with automatic session provisioning."""
+#
+#     @wraps(func)
+#     async def wrapper(*args: Any, **kwargs: Any) -> Any:
+#         if kwargs.get("session"):
+#             return await func(*args, **kwargs)
+#
+#         async with session_factory() as session:
+#             try:
+#                 return await func(*args, **kwargs, session=session)
+#             except Exception:
+#                 await session.rollback()
+#                 raise
+#             finally:
+#                 await session.close()
+#
+#     return wrapper

app/infrastructure/repositories/sample_record.py

@@ -13,7 +13,7 @@
     RecordUpdateError,
     RecordAlreadyExistsError,
     ForeignKeyError,
-    ValidationError,
+    ValidationError, BaseRepositoryError,
 )
 from app.application.repository.interfaces import ISampleRecordRepository
 from app.decorators.mapper.exception_mapper import (
@@ -58,6 +58,15 @@ def __init__(self, session: AsyncSession):
         """
         self._session = session
 
+    @ExceptionMapper(
+        {
+            Exception: EnrichedExceptionFactory(BaseRepositoryError),
+        },
+        is_bound_method=True,
+    )
+    async def commit(self) -> None:
+        await self._session.commit()
+
     @ExceptionMapper(
         {
             IntegrityError: IntegrityErrorFactory(RecordCreateError),

app/infrastructure/repositories/unit_of_work.py

@@ -0,0 +1,51 @@
+import asyncio
+
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from app.application.repository.interfaces import (
+    ISampleRecordUnitOfWork,
+    ISampleRecordRepository,
+)
+from app.infrastructure.repositories.sample_record import SampleRecordRepository
+
+
+class ReadOnlySampleRecordUnitOfWork(ISampleRecordUnitOfWork):
+    def get_sample_record_repository(self) -> ISampleRecordRepository:
+        if not self._session:
+            raise RuntimeError("Session is not initialized")
+
+        return SampleRecordRepository(self._session)
+
+    def __init__(self, session_factory: async_sessionmaker):
+        super().__init__()
+        self.session_factory = session_factory
+        self._session: AsyncSession | None = None
+
+    async def __aenter__(self):
+        self._session = self.session_factory()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        try:
+            # Recommended for implicit resources cleanup
+            await self._session.rollback()
+        finally:
+            await self._session.close()
+
+
+class WriteSampleRecordUnitOfWork(ReadOnlySampleRecordUnitOfWork):
+    """Unit of Work for write operations with full transaction management."""
+
+    async def __aenter__(self):
+        self._session = self.session_factory()
+        await asyncio.wait_for(self._session.begin(), timeout=5)
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        try:
+            if exc_type:
+                await self._session.rollback()
+            else:
+                await self._session.commit()
+        finally:
+            await self._session.close()

app/main.py

@@ -12,41 +12,48 @@
     BotSampleRecordCommandContainer,
     CallbackTaskManager,
 )
-from app.infrastructure.db.sqlalchemy import close_db_connections
+from app.infrastructure.db.sqlalchemy import (
+    get_session_factory, get_engine,
+)
 from app.presentation.api.routers import router
 from app.presentation.bot.resources import strings
 
 
 async def startup(
-    bot: Bot = Provide[ApplicationStartupContainer.bot],
+    bot: Bot,
 ) -> None:
     await bot.startup()
 
 
 async def shutdown(
-    callback_task_manager: CallbackTaskManager = Provide[
-        ApplicationStartupContainer.callback_task_manager
-    ],
-    bot: Bot = Provide[ApplicationStartupContainer.bot],
-    redis_client: Redis = Provide[ApplicationStartupContainer.redis_client],
+    container: ApplicationStartupContainer = Provide[ApplicationStartupContainer],
 ) -> None:
-    await bot.shutdown()
+    await container.bot().shutdown()
 
-    await callback_task_manager.shutdown()
+    await container.callback_task_manager().shutdown()
 
-    await redis_client.aclose()
-    await close_db_connections()
+    await container.redis_client().aclose()
+    await container.shutdown_resources()
+    await get_engine().dispose()
 
 
 def get_application() -> FastAPI:
     """Create configured server application instance."""
 
     # Initialize the main application container
     main_container = ApplicationStartupContainer()
-    main_container.wire(modules=["app.main", "app.presentation.api.botx"])
+    main_container.wire(
+        modules=[
+            "app.main",
+            "app.presentation.api.botx",
+            "app.presentation.bot.commands.sample_record",
+        ]
+    )
 
     # Initialize the SampleRecord commands container
-    sample_record_commands_container = BotSampleRecordCommandContainer()
+    sample_record_commands_container = BotSampleRecordCommandContainer(
+        session_factory=get_session_factory()
+    )
     sample_record_commands_container.wire(
         modules=["app.presentation.bot.commands.sample_record"]
     )
@@ -64,9 +71,10 @@ def get_application() -> FastAPI:
         "shutdown",
         partial(
             shutdown,
-            callback_task_manager=main_container.callback_task_manager(),
-            bot=main_container.bot(),
-            redis_client=main_container.redis_client(),
+            # callback_task_manager=main_container.callback_task_manager(),
+            # bot=main_container.bot(),
+            # redis_client=main_container.redis_client(),
+            container=main_container,
         ),
     )