feat: add worker (#34)

Author: alexhook

.gitlab-ci.yml.jinja

@@ -1,6 +1,7 @@
 # YAML objects named with dot is anchors, which are not recognized as jobs
 .run_bot: &run_bot
   - docker rm -f ${CONTAINER_NAME} || true
+  {% if add_worker -%}- docker rm -f ${CONTAINER_NAME}-worker || true{%- endif %}
   - docker pull ${CONTAINER_RELEASE_IMAGE}
   # Add new envs here. Don't forget to add them in example.env and docker-compose files.
   - docker run
@@ -17,6 +18,21 @@
     -e BOT_CREDENTIALS="${BOT_CREDENTIALS}"
     -e DEBUG="${DEBUG:-false}"
     $CONTAINER_RELEASE_IMAGE
+  {% if add_worker -%}
+  # Add envs for worker here
+  - docker run
+    -d
+    --name ${CONTAINER_NAME}-worker
+    --restart always
+    --log-opt max-size=10m
+    --log-opt max-file=5
+    -e POSTGRES_DSN="${POSTGRES_DSN}"
+    -e REDIS_DSN="${REDIS_DSN}"
+    -e BOT_CREDENTIALS="${BOT_CREDENTIALS}"
+    -e DEBUG="${DEBUG:-false}"
+    ${CONTAINER_RELEASE_IMAGE}
+    bash -c 'PYTHONPATH="$PYTHONPATH:$PWD" saq app.worker.worker.settings'
+  {%- endif %}
 
 .create_db: &create_db
   - psql -c "create user \"${POSTGRES_USER}\"" postgres || true
@@ -147,6 +163,9 @@ deploy.botstest.stop:
   extends: deploy.botstest
   script:
     - docker rm -f ${CONTAINER_NAME} || true
+    {% if add_worker -%}
+    - docker rm -f ${CONTAINER_NAME} ${CONTAINER_NAME}-worker || true
+    {%- endif %}
     - psql -c "select pg_terminate_backend(pid) from pg_stat_activity \
       where datname = '${POSTGRES_DB}';" postgres || true
     - psql -c "drop database \"${POSTGRES_DB}\"" postgres || true

app/api/dependencies/healthcheck.py renamed to app/api/dependencies/healthcheck.py.jinja

@@ -1,11 +1,19 @@
 """Bot dependency for healthcheck."""
 
+{% if add_worker -%}
+from asyncio.exceptions import TimeoutError
+{%- endif %}
 from typing import Optional
 
 from fastapi import Depends, Request
 from pybotx import Bot
 from sqlalchemy.sql import text
 
+{% if add_worker -%}
+from app.settings import settings
+from app.worker.worker import queue
+{%- endif %}
+
 
 async def check_db_connection(request: Request) -> Optional[str]:
     assert isinstance(request.app.state.bot, Bot)
@@ -33,3 +41,24 @@ async def check_redis_connection(request: Request) -> Optional[str]:
 
 
 check_redis_connection_dependency = Depends(check_redis_connection)
+{%- if add_worker %}
+
+
+async def check_worker_status() -> Optional[str]:
+    job = await queue.enqueue("healthcheck")
+
+    if not job:
+        return None
+
+    try:
+        await job.refresh(settings.WORKER_TIMEOUT_SEC)
+    except TimeoutError:
+        return "Worker is overloaded or not launched"
+    except Exception as exc:
+        return str(exc)
+
+    return None
+
+
+check_worker_status_dependency = Depends(check_worker_status)
+{%- endif %}

app/api/endpoints/healthcheck.py renamed to app/api/endpoints/healthcheck.py.jinja

@@ -7,6 +7,9 @@
 from app.api.dependencies.healthcheck import (
     check_db_connection_dependency,
     check_redis_connection_dependency,
+    {% if add_worker -%}
+    check_worker_status_dependency,
+    {%- endif %}
 )
 from app.services.healthcheck import (
     HealthCheckResponse,
@@ -21,6 +24,9 @@
 async def healthcheck(
     redis_connection_error: Optional[str] = check_redis_connection_dependency,
     db_connection_error: Optional[str] = check_db_connection_dependency,
+    {% if add_worker -%}
+    worker_status_error: Optional[str] = check_worker_status_dependency,
+    {%- endif %}
 ) -> HealthCheckResponse:
     """Check the health of the bot and services."""
     healthcheck_builder = HealthCheckResponseBuilder()
@@ -30,5 +36,10 @@ async def healthcheck(
     healthcheck_builder.add_healthcheck_result(
         HealthCheckServiceResult(name="redis", error=redis_connection_error)
     )
+    {% if add_worker -%}
+    healthcheck_builder.add_healthcheck_result(
+        HealthCheckServiceResult(name="worker", error=worker_status_error)
+    )
+    {%- endif %}
 
     return healthcheck_builder.build()

app/settings.py renamed to app/settings.py.jinja

@@ -30,6 +30,11 @@ class Config:  # noqa: WPS431
     # redis
     REDIS_DSN: str
 
+    {% if add_worker -%}
+    # healthcheck
+    WORKER_TIMEOUT_SEC: float = 4
+    {%- endif %}
+
     @validator("BOT_CREDENTIALS", pre=True)
     @classmethod
     def parse_bot_credentials(cls, raw_credentials: Any) -> List[BotAccountWithSecret]:

app/{% if add_worker %}worker{% endif %}/__init__.py

@@ -0,0 +1,51 @@
+"""Tasks worker configuration."""
+
+from typing import Any, Dict, Literal
+
+from pybotx import Bot
+from redis import asyncio as aioredis
+from saq import Queue
+
+from app.caching.callback_redis_repo import CallbackRedisRepo
+from app.logger import logger
+
+# `saq` import its own settings and hides our module
+from app.settings import settings as app_settings
+
+SaqCtx = Dict[str, Any]
+
+
+async def startup(ctx: SaqCtx) -> None:
+    from app.bot.bot import get_bot  # noqa: WPS433
+
+    callback_repo = CallbackRedisRepo(aioredis.from_url(app_settings.REDIS_DSN))
+    bot = get_bot(callback_repo)
+
+    await bot.startup(fetch_tokens=False)
+
+    ctx["bot"] = bot
+
+    logger.info("Worker started")
+
+
+async def shutdown(ctx: SaqCtx) -> None:
+    bot: Bot = ctx["bot"]
+    await bot.shutdown()
+
+    logger.info("Worker stopped")
+
+
+async def healthcheck(_: SaqCtx) -> Literal[True]:
+    return True
+
+
+queue = Queue(aioredis.from_url(app_settings.REDIS_DSN), name="{{bot_project_name}}")
+
+settings = {
+    "queue": queue,
+    "functions": [healthcheck],
+    "cron_jobs": [],
+    "concurrency": 8,
+    "startup": startup,
+    "shutdown": shutdown,
+}

app/{% if add_worker %}worker{% endif %}/worker.py.jinja

@@ -11,6 +11,11 @@ bot_description:
     help: Description for README.md. First line will be added pyproject.toml
     default: TODO
 
+add_worker:
+    help: Include tasks worker in `docker-compose.yml`
+    type: bool
+    default: yes
+
 
 bot_name_underscored:
     default: "{{bot_project_name|replace('-', '_')}}"

copier.yaml

@@ -4,28 +4,41 @@ services:
   {{bot_project_name}}:
     build: .
     container_name: {{bot_project_name}}
-    environment:
+    environment: &environment
       - BOT_CREDENTIALS=cts_host@secret_key@bot_id
       - POSTGRES_DSN=postgres://postgres:postgres@{{bot_project_name}}-postgres/{{bot_name_underscored}}_db
       - REDIS_DSN=redis://{{bot_project_name}}-redis/0
       - DEBUG=true
     ports:
       - "8000:8000"  # Отредактируйте порт хоста (первый), если он уже занят
     restart: always
-    depends_on:
+    depends_on: &depends_on
       - postgres
       - redis
-    logging:
+    logging: &logging
       driver: "json-file"
       options:
         max-size: "10m"
         max-file: "10"
-    ulimits:
+    ulimits: &ulimits
       nproc: 65535
       nofile:
         soft: 20000
         hard: 40000
 
+  {% if add_worker -%}
+  {{bot_project_name}}-worker:
+    build: .
+    container_name: {{bot_project_name}}-worker
+    # '$$' prevents docker-compose from interpolating a value
+    command: /bin/sh -c 'PYTHONPATH="$$PYTHONPATH:$$PWD" saq app.worker.worker.settings'
+    environment: *environment
+    restart: always
+    depends_on: *depends_on
+    logging: *logging
+    ulimits: *ulimits
+
+  {% endif -%}
   postgres:
     image: postgres:15.3-alpine
     container_name: {{bot_project_name}}-postgres