feat: Add Task processor (#26)

* feat: Add Task processor
- Port `flagsmith/flagsmith-task-processor`
- Port `waitfordb` management command
- Introduce PostgreSQL for tests
- Bump Poetry from 2.0.1 to 2.1.1

---------

Co-authored-by: Matthew Elwell <[email protected]>

Author: khvn26

.github/workflows/python-test.yml

@@ -12,6 +12,15 @@ jobs:
       matrix:
         python-version: ["3.11", "3.12"]
 
+    services:
+      postgres:
+        image: postgres:15.5-alpine
+        env:
+          POSTGRES_PASSWORD: password
+          POSTGRES_DB: flagsmith
+        ports: ['5432:5432']
+        options: --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5
+
     steps:
       - uses: actions/checkout@v4
 
@@ -23,13 +32,17 @@ jobs:
         run: pipx install poetry
 
       - name: Install Dependencies
-        run: poetry install --with dev
+        env:
+          opts: --with dev
+        run: make install-packages
 
       - name: Check for missing migrations
-        run: poetry run python manage.py makemigrations --no-input --dry-run --check
+        env:
+          opts: --no-input --dry-run --check
+        run: make django-make-migrations
 
       - name: Check for new typing errors
-        run: poetry run mypy .
+        run: make typecheck
 
       - name: Run Tests
-        run: poetry run pytest
+        run: make test

Makefile

@@ -1,4 +1,9 @@
-POETRY_VERSION ?= 2.0.1
+.EXPORT_ALL_VARIABLES:
+
+POETRY_VERSION ?= 2.1.1
+
+COMPOSE_FILE ?= docker/docker-compose.local.yml
+COMPOSE_PROJECT_NAME ?= flagsmith-common
 
 .PHONY: install-pip
 install-pip:
@@ -10,11 +15,44 @@ install-poetry:
 
 .PHONY: install-packages
 install-packages:
-	poetry install --no-root $(opts)
+	poetry install $(opts)
 
 .PHONY: install
 install: install-pip install-poetry install-packages
 
 .PHONY: lint
 lint:
 	poetry run pre-commit run -a
+
+.PHONY: docker-up
+docker-up:
+	docker compose up --force-recreate --remove-orphans -d
+	docker compose ps
+
+.PHONY: docker-down
+docker-down:
+	docker compose down
+
+.PHONY: test
+test:
+	poetry run pytest $(opts)
+
+.PHONY: typecheck
+typecheck:
+	poetry run mypy .
+
+.PHONY: django-make-migrations
+django-make-migrations:
+	poetry run python manage.py waitfordb
+	poetry run python manage.py makemigrations $(opts)
+
+.PHONY: django-squash-migrations
+django-squash-migrations:
+	poetry run python manage.py waitfordb
+	poetry run python manage.py squashmigrations $(opts)
+
+.PHONY: django-migrate
+django-migrate:
+	poetry run python manage.py waitfordb
+	poetry run python manage.py migrate
+	poetry run python manage.py createcachetable

docker/docker-compose.local.yml

@@ -0,0 +1,19 @@
+# A Compose file with minimal dependencies to be able to run Flagsmith, including its test suite, locally (not in Docker).
+
+name: flagsmith
+
+volumes:
+  pg_data:
+
+services:
+  db:
+    image: postgres:15.5-alpine
+    pull_policy: always
+    restart: unless-stopped
+    volumes:
+      - pg_data:/var/lib/postgresql/data
+    ports:
+      - 5432:5432
+    environment:
+      POSTGRES_DB: flagsmith
+      POSTGRES_PASSWORD: password

poetry.lock

@@ -4,15 +4,19 @@ version = "1.5.2"
 description = "Flagsmith's common library"
 requires-python = ">=3.11,<4.0"
 dependencies = [
-    "django (<5)",
+    "backoff (>=2.2.1,<3.0.0)",
+    "django (>4,<5)",
     "django-health-check",
     "djangorestframework-recursive",
     "djangorestframework",
     "drf-writable-nested",
+    "drf-yasg (>=1.21.10,<2.0.0)",
+    "environs (<15)",
     "flagsmith-flag-engine",
     "gunicorn (>=19.1)",
     "prometheus-client (>=0.0.16)",
-    "environs (<15)",
+    "psycopg2 (>=2,<3)",
+    "simplejson (>=3,<4)",
 ]
 authors = [
     { name = "Matthew Elwell" },
@@ -29,6 +33,7 @@ readme = "README.md"
 dynamic = ["classifiers"]
 
 [project.urls]
+Changelog = "https://github.com/flagsmith/flagsmith-common/blob/main/CHANGELOG.md"
 Download = "https://github.com/flagsmith/flagsmith-common/releases"
 Homepage = "https://flagsmith.com"
 Issues = "https://github.com/flagsmith/flagsmith-common/issues"
@@ -48,9 +53,13 @@ classifiers = [
     "Intended Audience :: Developers",
     "License :: OSI Approved :: BSD License",
 ]
-packages = [{ include = "common", from = "src" }]
+packages = [
+    { include = "common", from = "src" },
+    { include = "task_processor", from = "src" },
+]
 
 [tool.poetry.group.dev.dependencies]
+dj-database-url = "^2.3.0"
 django-stubs = "^5.1.3"
 djangorestframework-stubs = "^3.15.3"
 mypy = "^1.15.0"
@@ -65,6 +74,7 @@ pytest-mock = "^3.14.0"
 requests = "^2.32.3"
 ruff = "*"
 setuptools = "^77.0.3"
+types-simplejson = "^3.20.0.20250326"
 
 [build-system]
 requires = ["poetry-core"]

pyproject.toml

@@ -1,4 +1,12 @@
+from datetime import time, timedelta
+
+import dj_database_url
 import prometheus_client
+from environs import Env
+
+from task_processor.task_run_method import TaskRunMethod
+
+env = Env()
 
 # Settings expected by `mypy_django_plugin`
 AWS_SES_REGION_ENDPOINT: str
@@ -8,15 +16,18 @@
 
 # Settings required for tests
 DATABASES = {
-    "default": {
-        "ENGINE": "django.db.backends.sqlite3",
-        "NAME": "common.sqlite3",
-    }
+    "default": dj_database_url.parse(
+        env(
+            "DATABASE_URL",
+            default="postgresql://postgres:password@localhost:5432/flagsmith",
+        )
+    )
 }
 INSTALLED_APPS = [
     "django.contrib.auth",
     "django.contrib.contenttypes",
     "common.core",
+    "task_processor",
 ]
 MIDDLEWARE = [
     "common.gunicorn.middleware.RouteLoggerMiddleware",
@@ -26,3 +37,16 @@
 ROOT_URLCONF = "common.core.urls"
 TIME_ZONE = "UTC"
 USE_TZ = True
+
+ENABLE_CLEAN_UP_OLD_TASKS = True
+ENABLE_TASK_PROCESSOR_HEALTH_CHECK = True
+RECURRING_TASK_RUN_RETENTION_DAYS = 15
+TASK_DELETE_BATCH_SIZE = 2000
+TASK_DELETE_INCLUDE_FAILED_TASKS = False
+TASK_DELETE_RETENTION_DAYS = 15
+TASK_DELETE_RUN_EVERY = timedelta(days=1)
+TASK_DELETE_RUN_TIME = time(5, 0, 0)
+TASK_RUN_METHOD = TaskRunMethod.TASK_PROCESSOR
+
+# Avoid models.W042 warnings
+DEFAULT_AUTO_FIELD = "django.db.models.AutoField"

settings/dev.py

@@ -1,40 +1,80 @@
+import contextlib
 import logging
 import os
 import sys
 import tempfile
+import typing
 
 from django.core.management import execute_from_command_line
 
 logger = logging.getLogger(__name__)
 
 
-def main() -> None:
+@contextlib.contextmanager
+def ensure_cli_env() -> typing.Generator[None, None, None]:
     """
-    The main entry point to the Flagsmith application.
+    Set up the environment for the main entry point of the application
+    and clean up after it's done.
 
-    An equivalent to Django's `manage.py` script, this module is used to run management commands.
+    Add environment-related code that needs to happen before and after Django is involved
+    to here.
 
-    It's installed as the `flagsmith` command.
+    Use as a context manager, e.g.:
 
-    Everything that needs to be run before Django is started should be done here.
+    ```python
+    with ensure_cli_env():
+        main()
+    ```
+    """
+    ctx = contextlib.ExitStack()
 
-    The end goal is to eventually replace Core API's `run-docker.sh` with this.
+    # TODO @khvn26 Move logging setup to here
 
-    Usage:
-    `flagsmith <command> [options]`
-    """
+    # Currently we don't install Flagsmith modules as a package, so we need to add
+    # $CWD to the Python path to be able to import them
+    sys.path.append(os.getcwd())
+
+    # TODO @khvn26 We should find a better way to pre-set the Django settings module
+    # without resorting to it being set outside of the application
     os.environ.setdefault("DJANGO_SETTINGS_MODULE", "settings.dev")
 
     # Set up Prometheus' multiprocess mode
     if "PROMETHEUS_MULTIPROC_DIR" not in os.environ:
-        prometheus_multiproc_dir = tempfile.TemporaryDirectory(
-            prefix="prometheus_multiproc",
+        prometheus_multiproc_dir_name = ctx.enter_context(
+            tempfile.TemporaryDirectory(
+                prefix="prometheus_multiproc",
+            )
         )
+
         logger.info(
             "Created %s for Prometheus multi-process mode",
-            prometheus_multiproc_dir.name,
+            prometheus_multiproc_dir_name,
         )
-        os.environ["PROMETHEUS_MULTIPROC_DIR"] = prometheus_multiproc_dir.name
+        os.environ["PROMETHEUS_MULTIPROC_DIR"] = prometheus_multiproc_dir_name
+
+    if "task-processor" in sys.argv:
+        # A hacky way to signal we're not running the API
+        os.environ["RUN_BY_PROCESSOR"] = "true"
+
+    with ctx:
+        yield
+
 
-    # Run Django
-    execute_from_command_line(sys.argv)
+def main() -> None:
+    """
+    The main entry point to the Flagsmith application.
+
+    An equivalent to Django's `manage.py` script, this module is used to run management commands.
+
+    It's installed as the `flagsmith` command.
+
+    Everything that needs to be run before Django is started should be done here.
+
+    The end goal is to eventually replace Core API's `run-docker.sh` with this.
+
+    Usage:
+    `flagsmith <command> [options]`
+    """
+    with ensure_cli_env():
+        # Run Django
+        execute_from_command_line(sys.argv)

src/common/core/main.py

@@ -3,7 +3,10 @@
 from django.core.management import BaseCommand, CommandParser
 from django.utils.module_loading import autodiscover_modules
 
-from common.gunicorn.utils import add_arguments, run_server
+from common.gunicorn.utils import add_arguments as add_gunicorn_arguments
+from common.gunicorn.utils import run_server
+from task_processor.utils import add_arguments as add_task_processor_arguments
+from task_processor.utils import start_task_processor
 
 
 class Command(BaseCommand):
@@ -13,20 +16,31 @@ def create_parser(self, *args: Any, **kwargs: Any) -> CommandParser:
         return super().create_parser(*args, conflict_handler="resolve", **kwargs)
 
     def add_arguments(self, parser: CommandParser) -> None:
-        add_arguments(parser)
+        add_gunicorn_arguments(parser)
 
         subparsers = parser.add_subparsers(
             title="sub-commands",
             required=True,
         )
+
         api_parser = subparsers.add_parser(
             "api",
             help="Start the Core API.",
         )
         api_parser.set_defaults(handle_method=self.handle_api)
 
+        task_processor_parser = subparsers.add_parser(
+            "task-processor",
+            help="Start the Task Processor.",
+        )
+        task_processor_parser.set_defaults(handle_method=self.handle_task_processor)
+        add_task_processor_arguments(task_processor_parser)
+
     def initialise(self) -> None:
-        autodiscover_modules("metrics")
+        autodiscover_modules(
+            "metrics",
+            "tasks",
+        )
 
     def handle(
         self,
@@ -39,3 +53,9 @@ def handle(
 
     def handle_api(self, *args: Any, **options: Any) -> None:
         run_server(options)
+
+    def handle_task_processor(self, *args: Any, **options: Any) -> None:
+        with start_task_processor(options):
+            # Delegate signal handling to Gunicorn.
+            # The task processor will finalise once Gunicorn is shut down.
+            run_server(options)