atciss: use taskiq 💡 instead of celery 🥦

.dockerignore

@@ -13,7 +13,7 @@ htmlcov
 .cache
 nosetests.xml
 coverage.xml
-*,cover
+*.cover
 .hypothesis
 xunit-*.xml
 .idea
@@ -25,5 +25,7 @@ site
 .git
 .github
 result
-
 atciss-frontend/
+screenshots/
+nixos/
+docker-compose.*.yml

Dockerfile

@@ -1,17 +1,27 @@
-FROM python:3.11-slim@sha256:23f52205321f806c2cc742cefbf837e0d25101388c043e860c7817985230565c as base
+FROM python:3.11-slim@sha256:53d6284a40eae6b625f22870f5faba6c54f2a28db9027408f4dee111f1e885a2 as base
 
 FROM base as requirements-stage
 WORKDIR /tmp
 RUN pip install poetry
 COPY ./pyproject.toml ./poetry.lock* /tmp/
+
+FROM requirements-stage as requirements-dev
+RUN poetry export --with=dev -f requirements.txt --output requirements.txt --without-hashes
+
+FROM requirements-stage as requirements-prod
 RUN poetry export -f requirements.txt --output requirements.txt --without-hashes
 
 FROM base as dev
 WORKDIR /code
-COPY --from=requirements-stage /tmp/requirements.txt /code/requirements.txt
+COPY --from=requirements-dev /tmp/requirements.txt /code/requirements.txt
 RUN apt-get update && apt-get install git -y
 RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
-ENTRYPOINT ["python3", "-c", "from atciss.cli import cli; cli()"]
+ENTRYPOINT ["python3", "-m", "atciss.cli"]
 
 FROM dev as production
+WORKDIR /code
+COPY --from=requirements-prod /tmp/requirements.txt /code/requirements.txt
+RUN apt-get update && apt-get install git -y
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 COPY . /code
+ENTRYPOINT ["python3", "-m", "atciss.cli"]

README.md

@@ -31,15 +31,14 @@ using [Nix Flakes](https://nixos.wiki/wiki/Flakes) and
 atciss serve
 ```
 
-### Run the worker
+### Run the worker & task scheduler
 
 ```
-celery -A atciss worker
+atciss worker
 ```
 
-Automatic scheduling of tasks:
 ```
-celery -A atciss beat
+atciss scheduler
 ```
 
 ### Run the frontend

alembic/env.py

@@ -3,12 +3,10 @@
 from sqlalchemy import pool
 from sqlalchemy.engine import Connection
 from sqlalchemy.ext.asyncio import async_engine_from_config
-
-from alembic import context
-
 from sqlmodel import SQLModel
-import atciss.app.models  # noqa: F401
 
+import atciss.app.models  # noqa: F401
+from alembic import context
 from atciss.config import settings
 from atciss.log import setup_logging
 

atciss/__init__.py

@@ -1,7 +0,0 @@
-"""ATCISS."""
-
-from atciss import celery
-
-from .version import __version__
-
-__all__ = ("__version__", "celery")

atciss/app/asgi.py

@@ -15,15 +15,18 @@
 from fastapi import FastAPI
 from loguru import logger
 from prometheus_fastapi_instrumentator import Instrumentator as PrometheusInstrumentator
-from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
 
 import alembic.command
 import alembic.config
 
-from ..config import settings
-from ..log import setup_logging
-from .router import root_api_router
-from .utils import RedisClient
+# for tasks discovery by broker
+import atciss.tasks  # pyright: ignore # noqa: F401
+from atciss.app.router import root_api_router
+from atciss.app.utils import RedisClient
+from atciss.config import settings
+from atciss.log import setup_logging
+from atciss.tkq import broker
 
 
 @dataclass
@@ -40,7 +43,7 @@ async def __call__(
             return await self.app(scope, receive, send)
 
 
-def run_migrations(connection, cfg):
+def run_migrations(connection: AsyncSession, cfg: alembic.config.Config):
     cfg.attributes["connection"] = connection
     alembic.command.upgrade(cfg, "head")
 
@@ -57,15 +60,19 @@ async def run_async_migrations():
 
 @asynccontextmanager
 async def lifespan(_app: FastAPI):
-    await run_async_migrations()
+    if not broker.is_worker_process:
+        await run_async_migrations()
+        await broker.startup()
     yield
+    if not broker.is_worker_process:
+        await broker.shutdown()
     await RedisClient.close()
 
 
 def get_application() -> FastAPI:
     """Initialize FastAPI application."""
     setup_logging()
-    logger.debug("Initialize FastAPI application node.")
+    logger.debug("Initialize FastAPI application.")
 
     app = FastAPI(
         title=settings.PROJECT_NAME,
@@ -77,7 +84,7 @@ def get_application() -> FastAPI:
 
     _ = PrometheusInstrumentator().instrument(app).expose(app, tags=["monitoring"])
 
-    app.add_middleware(CorrelationIdLogMiddleware)
+    app.add_middleware(CorrelationIdLogMiddleware)  # pyright: ignore
     app.add_middleware(
         CorrelationIdMiddleware,
         generator=(lambda: uuid4().hex) if not settings.DEBUG else (lambda: uuid4().hex[:8]),

atciss/app/controllers/admin.py

@@ -1,11 +1,10 @@
 from typing import Annotated
 
-import celery
 from fastapi import APIRouter, Depends
 
-import atciss.celery as ac
 from atciss.app.controllers.auth import get_admin
 from atciss.app.models import User
+from atciss.tkq import broker
 
 router = APIRouter()
 
@@ -14,12 +13,12 @@
 async def get_tasks(
     user: Annotated[User, Depends(get_admin)],
 ) -> list[str]:
-    return [task for task in dir(ac) if isinstance(getattr(ac, task), celery.Task)]
+    return list(broker.get_all_tasks().keys())
 
 
 @router.post("/admin/task/{task}")
 async def trigger_task(
     task: str,
     user: Annotated[User, Depends(get_admin)],
 ):
-    getattr(ac, task).apply_async()
+    _ = await broker.get_all_tasks()[task].kiq()

atciss/app/controllers/airspace.py

@@ -6,7 +6,7 @@
 from loguru import logger
 from pydantic import TypeAdapter
 
-from atciss.app.tasks.sectors import Airport, Airspace, Position, SectorData
+from atciss.app.views.sector import Airport, Airspace, Position, SectorData
 
 from ...config import settings
 from ..controllers.auth import get_user

atciss/app/controllers/auth.py

@@ -137,7 +137,7 @@ async def get_user(
         user = await session.scalar(stmt)
 
         if user is None:
-            raise HTTPException(500, "User not not found in database")
+            raise HTTPException(500, "User not found in database")
 
     except JWTError as exc:
         raise credentials_exception from exc

atciss/app/models.py

@@ -2,8 +2,8 @@
 
 from sqlmodel import Field, SQLModel
 
-from .views.booking import Booking  # noqa: F401 pylint: disable=unused-import
-from .views.dfs_aixm import (  # noqa: F401 pylint: disable=unused-import
+from .views.booking import Booking  # noqa: F401
+from .views.dfs_aixm import (  # noqa: F401
     Aerodrome,
     Navaid,
     Runway,

atciss/cli/__init__.py

@@ -1,16 +1 @@
-"""
-Command-line interface.
-
-The ``cli`` submodule defines Click command-line interface root and its
-commands.
-
-Resources:
-    1. `Click documentation`_
-
-.. _Click documentation:
-    https://click.palletsprojects.com/en/8.1.x/
-"""
-
-from atciss.cli.cli import cli
-
-__all__ = ("cli",)
+from .cli import cli  # pyright: ignore # noqa: F401

atciss/cli/__main__.py

@@ -0,0 +1,3 @@
+from atciss.cli.cli import cli
+
+cli()

atciss/cli/cli.py

@@ -4,7 +4,9 @@
 
 import click
 
+from .scheduler import scheduler
 from .serve import serve
+from .worker import worker
 
 CMD_HELP = "ATCISS CLI root."
 
@@ -15,6 +17,5 @@ def cli(**_: dict[str, Any]) -> None:
 
 
 cli.add_command(serve)
-
-if __name__ == "__main__":
-    cli()
+cli.add_command(worker)
+cli.add_command(scheduler)