fix worker for di

Author: vladimirgubarik

app/infrastructure/containers.py

@@ -146,6 +146,7 @@ class ApplicationStartupContainer(containers.DeclarativeContainer):
     )
 
 
+
 class WorkerStartupContainer(containers.DeclarativeContainer):
     redis_client = Singleton(lambda: aioredis.from_url(settings.REDIS_DSN))
 

app/infrastructure/db/sqlalchemy.py

@@ -57,16 +57,11 @@ def get_engine() -> AsyncEngine:
 session_factory = async_sessionmaker(bind=get_engine(), expire_on_commit=False)
 
 
-async def verify_db_connection(engine: AsyncEngine) -> None:
-    connection = await engine.connect()
-    await connection.close()
-
-
 async def close_db_connections() -> None:
     await get_engine().dispose()
 
 
-def provide_transaction_session(func: Callable) -> Callable:
+def provide_session(func: Callable) -> Callable:
     """
     Provides a database session to an async function if one is not already passed.
 
@@ -78,7 +73,14 @@ def provide_transaction_session(func: Callable) -> Callable:
     async def wrapper(*args: Any, **kwargs: Any) -> Any:
         if kwargs.get("session"):
             return await func(*args, **kwargs)
+
         async with session_factory() as session:
-            return await func(*args, **kwargs, session=session)
+            try:
+                return await func(*args, **kwargs, session=session)
+            except Exception:
+                await session.rollback()
+                raise
+            finally:
+                await session.close()
 
     return wrapper

app/infrastructure/worker/tasks/simple_task.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import Any
 
 from dependency_injector.wiring import inject, Provide
@@ -13,6 +14,7 @@ async def heartbeat_task(
     ctx: dict[str, Any],
     bot: Bot = Provide[WorkerStartupContainer.bot],
 ):
-    # logger.info("Heartbeat task started")
-
-    logger.info(f"Heartbeat task executed {[account.id for account in bot.bot_accounts]}")
+    task_name = asyncio.current_task().get_name()
+    logger.info(f"Task {task_name} Heartbeat task executed start bot id {id(bot)}")
+    await asyncio.sleep(10)
+    logger.info(f"Task {task_name} Heartbeat task executed end bot id {id(bot)}")

app/infrastructure/worker/worker.py

@@ -15,14 +15,14 @@
 from app.infrastructure.worker.tasks.simple_task import heartbeat_task
 from app.logger import logger
 
-# `saq` import its own settings and hides our module
-from app.settings import settings as app_settings
+from app.settings import settings
 
 SaqCtx = Dict[str, Any]
 
 
+# queue = Queue(aioredis.from_url(settings.REDIS_DSN), name="bot_template")
+queue = Queue.from_url(settings.REDIS_DSN, name="bot_template_worker")
 
-queue = Queue(aioredis.from_url(app_settings.REDIS_DSN), name="bot_refactor")
 
 @inject
 async def _startup_with_injection(
@@ -45,11 +45,7 @@ async def _shutdown_with_injection(
 async def startup(ctx: SaqCtx) -> None:
     worker_startup_container = WorkerStartupContainer()
 
-    queue.add_cron_job(
-        CronJob(function=heartbeat_task, cron="*/5 * * * * *", unique=True)
-    )
     worker_startup_container.wire(modules=[__name__, "app.infrastructure.worker.tasks"])
-
     await _startup_with_injection()
 
     logger.info("Worker started")
@@ -60,21 +56,18 @@ async def shutdown(ctx: SaqCtx) -> None:
     logger.info("Worker stopped")
 
 
-settings = {
+saq_settings = {
     "queue": queue,
     "functions": [],
-    # "cron_jobs": [
-    #     CronJob(
-    #         function=heartbeat_task,
-    #         cron="*/5 * * * * *",
-    #         unique=True,
-    #         # timeout=app_settings.PERIODIC_TASKS_DEFAULT_TIMEOUT,
-    #         # heartbeat=app_settings.PERIODIC_TASKS_DEFAULT_HEARTBEAT,
-    #         # retries=app_settings.PERIODIC_TASKS_DEFAULT_RETRIES,
-    #         # ttl=app_settings.PERIODIC_TASKS_DEFAULT_TTL,
-    #     ),
-    # ],
-    "concurrency": 8,
+    "cron_jobs": [
+        CronJob(
+            function=heartbeat_task,
+            cron="* * * * * */5",
+            unique=False,
+            timeout=15,
+        ),
+    ],
+    "concurrency": settings.WORKER_CONCURRENCY,
     "startup": startup,
     "shutdown": shutdown,
 }

app/main.py

@@ -53,7 +53,7 @@ def get_application() -> FastAPI:
 
     application = FastAPI(title=strings.BOT_PROJECT_NAME, openapi_url=None)
 
-    # put bot to state for tests
+    # put bot to state only for tests
     application.state.bot = main_container.bot()
 
     application.add_event_handler(

app/presentation/bot/commands/sample_record.py

@@ -5,15 +5,14 @@
 
 from app.application.use_cases.interfaces import ISampleRecordUseCases
 from app.infrastructure.containers import BotSampleRecordCommandContainer
-from app.infrastructure.db.sqlalchemy import provide_transaction_session
+from app.infrastructure.db.sqlalchemy import provide_session
 from app.presentation.bot.commands.command_listing import SampleRecordCommands
 from app.presentation.bot.command_handlers.sample_record import CreateSampleRecordHandler
 
 collector = HandlerCollector()
 
-
 @collector.command(**SampleRecordCommands.CREATE_RECORD.command_data())
-@provide_transaction_session
+@provide_session
 @inject
 async def create_sample_record(
     message: IncomingMessage,

app/presentation/bot/decorators/__init__.py

@@ -77,11 +77,11 @@ def _build_credentials_from_string(
     REDIS_DSN: str
     REDIS_CONNECTION_POOL_SIZE: int = 10
 
-    # healthcheck
-    WORKER_TIMEOUT_SEC: float = 4
+    # worker
+    WORKER_CONCURRENCY: int = 2
+    WORKERS_COUNT: int = 1
 
     BOTX_CALLBACK_TIMEOUT_IN_SECONDS = 30
-
     BOT_ASYNC_CLIENT_TIMEOUT_IN_SECONDS = 60
 
 

app/presentation/bot/decorators/bot_exception_answer.py

@@ -21,7 +21,9 @@ services:
       - REDIS_DSN=redis://redis/0
       - DEBUG=true
     # '$$' prevents docker-compose from interpolating a value
-    command: /bin/sh -c 'PYTHONPATH="$$PYTHONPATH:$$PWD" saq app.infrastructure.worker.worker.settings'
+    command: /bin/sh -c 'PYTHONPATH="$$PYTHONPATH:$$PWD" saq app.infrastructure.worker.worker.saq_settings --web --workers 2'
+    ports:
+      - "8081:8080"
     env_file:
       - .env
     restart: always