WIP: Fix managing tasks when in multi-database mode

Author: emyller

docker/docker-compose.local.yml

@@ -4,11 +4,11 @@ name: flagsmith
 
 volumes:
   pg_data:
+  task_processor_pg_data:
 
 services:
   db:
     image: postgres:15.5-alpine
-    pull_policy: always
     restart: unless-stopped
     volumes:
       - pg_data:/var/lib/postgresql/data
@@ -17,3 +17,21 @@ services:
     environment:
       POSTGRES_DB: flagsmith
       POSTGRES_PASSWORD: password
+    healthcheck:
+      test: pg_isready -Upostgres
+      interval: 1s
+      timeout: 30s
+
+  task-processor-db:
+    image: postgres:15.5-alpine
+    restart: unless-stopped
+    volumes:
+      - task_processor_pg_data:/var/lib/postgresql/data
+    ports:
+      - 5433:5432
+    environment:
+      POSTGRES_HOST_AUTH_METHOD: trust
+    healthcheck:
+      test: pg_isready -Upostgres
+      interval: 1s
+      timeout: 30s

settings/dev.py

@@ -29,9 +29,17 @@
         env(
             "DATABASE_URL",
             default="postgresql://postgres:password@localhost:5432/flagsmith",
-        )
-    )
+        ),
+    ),
+    "task_processor": dj_database_url.parse(
+        env(
+            "TASK_PROCESSOR_DATABASE_URL",
+            default="postgresql://postgres@localhost:5433/postgres",
+        ),
+    ),
 }
+DATABASE_ROUTERS = ["task_processor.routers.TaskProcessorRouter"]
+TASK_PROCESSOR_DATABASES = ["default"]
 INSTALLED_APPS = [
     "django.contrib.auth",
     "django.contrib.contenttypes",
@@ -62,5 +70,3 @@
 
 # Avoid models.W042 warnings
 DEFAULT_AUTO_FIELD = "django.db.models.AutoField"
-
-TASK_PROCESSOR_DATABASES = ["default"]

src/task_processor/managers.py

@@ -7,21 +7,17 @@
 
 
 class TaskManager(Manager["Task"]):
-    def get_tasks_to_process(
-        self,
-        database: str,
-        num_tasks: int,
-    ) -> typing.List["Task"]:
+    def get_tasks_to_process(self, num_tasks: int) -> typing.List["Task"]:
         return list(
-            self.using(database).raw(
+            self.raw(
                 "SELECT * FROM get_tasks_to_process(%s)",
                 [num_tasks],
             ),
         )
 
 
 class RecurringTaskManager(Manager["RecurringTask"]):
-    def get_tasks_to_process(self, database: str) -> typing.List["RecurringTask"]:
+    def get_tasks_to_process(self) -> typing.List["RecurringTask"]:
         return list(
-            self.using(database).raw("SELECT * FROM get_recurringtasks_to_process()"),
+            self.raw("SELECT * FROM get_recurringtasks_to_process()"),
         )

src/task_processor/processor.py

@@ -9,6 +9,7 @@
 from django.utils import timezone
 
 from task_processor import metrics
+from task_processor.managers import TaskManager
 from task_processor.models import (
     AbstractBaseTask,
     RecurringTask,
@@ -31,7 +32,8 @@ def run_tasks(database: str, num_tasks: int = 1) -> list[TaskRun]:
     if num_tasks < 1:
         raise ValueError("Number of tasks to process must be at least one")
 
-    tasks = Task.objects.get_tasks_to_process(database, num_tasks)
+    task_manager: TaskManager = Task.objects.db_manager(database)
+    tasks = task_manager.get_tasks_to_process(num_tasks)
     if tasks:
         logger.debug(f"Running {len(tasks)} task(s) from database '{database}'")
 
@@ -46,13 +48,13 @@ def run_tasks(database: str, num_tasks: int = 1) -> list[TaskRun]:
             task_runs.append(task_run)
 
         if executed_tasks:
-            Task.objects.bulk_update(
+            Task.objects.using(database).bulk_update(
                 executed_tasks,
                 fields=["completed", "num_failures", "is_locked"],
             )
 
         if task_runs:
-            TaskRun.objects.bulk_create(task_runs)
+            TaskRun.objects.using(database).bulk_create(task_runs)
             logger.debug(
                 f"Finished running {len(task_runs)} task(s) from database '{database}'"
             )
@@ -66,7 +68,7 @@ def run_recurring_tasks(database: str) -> list[RecurringTaskRun]:
     # NOTE: We will probably see a lot of delay in the execution of recurring tasks
     # if the tasks take longer then `run_every` to execute. This is not
     # a problem for now, but we should be mindful of this limitation
-    tasks = RecurringTask.objects.get_tasks_to_process(database)
+    tasks = RecurringTask.objects.db_manager(database).get_tasks_to_process()
     if tasks:
         logger.debug(f"Running {len(tasks)} recurring task(s)")
 

src/task_processor/routers.py

@@ -0,0 +1,77 @@
+from django.conf import settings
+from django.db.models import Model
+
+
+class TaskProcessorRouter:
+    """
+    Routing of database operations for task processor models
+    """
+
+    route_app_labels = ["task_processor"]
+
+    @property
+    def is_enabled(self) -> bool:
+        return "task_processor" in settings.TASK_PROCESSOR_DATABASES
+
+    def db_for_read(self, model: type[Model], **hints: None) -> str | None:
+        """
+        If enabled, route "task_processor" models to the a se database
+        """
+        if not self.is_enabled:
+            return None
+
+        if model._meta.app_label in self.route_app_labels:
+            return "task_processor"
+
+        return None
+
+    def db_for_write(self, model: type[Model], **hints: None) -> str | None:
+        """
+        Attempts to write task processor models go to 'task_processor' database.
+        """
+        if not self.is_enabled:
+            return None
+
+        if model._meta.app_label in self.route_app_labels:
+            return "task_processor"
+
+        return None
+
+    def allow_relation(self, obj1: Model, obj2: Model, **hints: None) -> bool | None:
+        """
+        Relations between objects are allowed if both objects are
+        in the task processor database.
+        """
+        if not self.is_enabled:
+            return None
+
+        both_objects_from_task_processor = (
+            obj1._meta.app_label in self.route_app_labels
+            and obj2._meta.app_label in self.route_app_labels
+        )
+
+        if both_objects_from_task_processor:
+            return True
+
+        return None
+
+    def allow_migrate(
+        self,
+        db: str,
+        app_label: str,
+        **hints: None,
+    ) -> bool | None:
+        """
+        Allow migrations for task processor models to run in both databases
+
+        NOTE: Even if, from a fresh install, the task processor tables are not
+        required in both databases, this is required to allow for easier
+        transition between a single database and a multi-database setup.
+        """
+        if not self.is_enabled:
+            return None
+
+        if app_label in self.route_app_labels:
+            return db in ["default", "task_processor"]
+
+        return None

src/task_processor/threads.py

@@ -104,6 +104,18 @@ def run_iteration(self) -> None:
         separate database setup.
         """
         database_is_separate = "task_processor" in settings.TASK_PROCESSOR_DATABASES
+        if database_is_separate:
+            assert (
+                "task_processor.routers.TaskProcessorRouter"
+                in settings.DATABASE_ROUTERS
+            ), (
+                "DATABASE_ROUTERS must include 'task_processor.routers.TaskProcessorRouter' "
+                "when using a separate task processor database."
+            )  # This is for our own sanity
+            assert "task_processor" in settings.DATABASES, (
+                "DATABASES must include 'task_processor' when using a separate task processor database."
+            )  # ¯\_(ツ)_/¯ One has to read the documentation and fix it: https://docs.flagsmith.com/deployment/configuration/task-processor
+
         for database in settings.TASK_PROCESSOR_DATABASES:
             try:
                 run_tasks(database, self.queue_pop_size)

tests/unit/task_processor/test_unit_task_processor_processor.py

@@ -1,6 +1,5 @@
 import logging
 import time
-import typing
 import uuid
 from datetime import timedelta
 from threading import Thread
@@ -9,6 +8,7 @@
 from django.core.cache import cache
 from django.utils import timezone
 from freezegun import freeze_time
+from pytest_django.fixtures import SettingsWrapper
 from pytest_mock import MockerFixture
 
 from common.test_tools import AssertMetricFixture
@@ -37,8 +37,7 @@
 
 
 @pytest.fixture(autouse=True)
-def reset_cache() -> typing.Generator[None, None, None]:
-    yield
+def reset_cache() -> None:
     cache.clear()
 
 
@@ -73,38 +72,49 @@ def _sleep_task(seconds: int) -> None:
     return _sleep_task
 
 
+@pytest.mark.django_db(databases=["default", "task_processor"])
 @pytest.mark.task_processor_mode
+@pytest.mark.parametrize(
+    "database",
+    ["default", "task_processor"],
+)
 def test_run_task_runs_task_and_creates_task_run_object_when_success(
+    database: str,
     dummy_task: TaskHandler[[str, str]],
 ) -> None:
     # Given
-    task = Task.create(
-        dummy_task.task_identifier,
-        scheduled_for=timezone.now(),
-    )
-    task.save()
+    task = Task.create(dummy_task.task_identifier, scheduled_for=timezone.now())
+    task.save(using=database)
 
     # When
-    task_runs = run_tasks("default")
+    task_runs = run_tasks(database)
 
     # Then
     assert cache.get(DEFAULT_CACHE_KEY)
 
-    assert len(task_runs) == TaskRun.objects.filter(task=task).count() == 1
+    assert (
+        len(task_runs) == TaskRun.objects.using(database).filter(task=task).count() == 1
+    )
     task_run = task_runs[0]
     assert task_run.result == TaskResult.SUCCESS.value
     assert task_run.started_at
     assert task_run.finished_at
     assert task_run.error_details is None
 
-    task.refresh_from_db()
+    task.refresh_from_db(using=database)
     assert task.completed
 
 
+@pytest.mark.django_db(databases=["default", "task_processor"])
 @pytest.mark.task_processor_mode
+@pytest.mark.parametrize(
+    "database",
+    ["default", "task_processor"],
+)
 def test_run_task_kills_task_after_timeout(
-    sleep_task: TaskHandler[[int]],
     caplog: pytest.LogCaptureFixture,
+    database: str,
+    sleep_task: TaskHandler[[int]],
 ) -> None:
     # Given
     task = Task.create(
@@ -113,21 +123,23 @@ def test_run_task_kills_task_after_timeout(
         args=(1,),
         timeout=timedelta(microseconds=1),
     )
-    task.save()
+    task.save(using=database)
 
     # When
-    task_runs = run_tasks("default")
+    task_runs = run_tasks(database)
 
     # Then
-    assert len(task_runs) == TaskRun.objects.filter(task=task).count() == 1
+    assert (
+        len(task_runs) == TaskRun.objects.using(database).filter(task=task).count() == 1
+    )
     task_run = task_runs[0]
     assert task_run.result == TaskResult.FAILURE.value
     assert task_run.started_at
     assert task_run.finished_at is None
     assert task_run.error_details
     assert "TimeoutError" in task_run.error_details
 
-    task.refresh_from_db()
+    task.refresh_from_db(using=database)
 
     assert task.completed is False
     assert task.num_failures == 1
@@ -139,12 +151,20 @@ def test_run_task_kills_task_after_timeout(
     )
 
 
-@pytest.mark.django_db
+@pytest.mark.django_db(databases=["default", "task_processor"])
 @pytest.mark.task_processor_mode
+@pytest.mark.parametrize(
+    "database",
+    ["default", "task_processor"],
+)
 def test_run_recurring_task_kills_task_after_timeout(
     caplog: pytest.LogCaptureFixture,
+    database: str,
+    settings: SettingsWrapper,
 ) -> None:
     # Given
+    settings.TASK_PROCESSOR_DATABASES = [database]
+
     @register_recurring_task(
         run_every=timedelta(seconds=1), timeout=timedelta(microseconds=1)
     )
@@ -157,18 +177,22 @@ def _dummy_recurring_task() -> None:
         task_identifier="test_unit_task_processor_processor._dummy_recurring_task",
     )
     # When
-    task_runs = run_recurring_tasks("default")
+    task_runs = run_recurring_tasks(database)
 
     # Then
-    assert len(task_runs) == RecurringTaskRun.objects.filter(task=task).count() == 1
+    assert (
+        len(task_runs)
+        == RecurringTaskRun.objects.using(database).filter(task=task).count()
+        == 1
+    )
     task_run = task_runs[0]
     assert task_run.result == TaskResult.FAILURE.value
     assert task_run.started_at
     assert task_run.finished_at is None
     assert task_run.error_details
     assert "TimeoutError" in task_run.error_details
 
-    task.refresh_from_db()
+    task.refresh_from_db(using=database)
 
     assert task.locked_at is None
     assert task.is_locked is False
@@ -179,6 +203,9 @@ def _dummy_recurring_task() -> None:
     )
 
 
+# TODO: Need to parametrize all/most tests below to run on both databases
+
+
 @pytest.mark.django_db
 @pytest.mark.task_processor_mode
 def test_run_recurring_tasks_runs_task_and_creates_recurring_task_run_object_when_success() -> (

tests/unit/task_processor/test_unit_task_processor_routers.py

@@ -0,0 +1 @@
+# TODO