feat: Separate task processor database

.github/workflows/python-test.yml

@@ -13,14 +13,20 @@ jobs:
         python-version: ["3.11", "3.12"]
 
     services:
-      postgres:
+      default-database:
         image: postgres:15.5-alpine
         env:
-          POSTGRES_PASSWORD: password
-          POSTGRES_DB: flagsmith
+          POSTGRES_HOST_AUTH_METHOD: trust
         ports: ['5432:5432']
         options: --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5
 
+      task-processor-database:
+        image: postgres:15.5-alpine
+        env:
+          POSTGRES_HOST_AUTH_METHOD: trust
+        ports: ['5433:5432']
+        options: --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5
+
     steps:
       - uses: actions/checkout@v4
 

docker/docker-compose.local.yml

@@ -1,19 +1,39 @@
-# A Compose file with minimal dependencies to be able to run Flagsmith, including its test suite, locally (not in Docker).
+# Minimal dependencies to run Flagsmith non-Dockerized, including its test
+# suite. Currently, settings are analogous to CI for compatibility.
+
+# TODO: Use Docker Compose in CI to run tests.
 
 name: flagsmith
 
 volumes:
-  pg_data:
+  default-database:
+  task-processor-database:
 
 services:
-  db:
+  default-database:
     image: postgres:15.5-alpine
-    pull_policy: always
     restart: unless-stopped
     volumes:
-      - pg_data:/var/lib/postgresql/data
+      - default-database:/var/lib/postgresql/data
     ports:
       - 5432:5432
     environment:
-      POSTGRES_DB: flagsmith
-      POSTGRES_PASSWORD: password
+      POSTGRES_HOST_AUTH_METHOD: trust
+    healthcheck:
+      test: pg_isready -Upostgres
+      interval: 1s
+      timeout: 30s
+
+  task-processor-database:
+    image: postgres:15.5-alpine
+    restart: unless-stopped
+    volumes:
+      - task-processor-database:/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

@@ -28,10 +28,18 @@
     "default": dj_database_url.parse(
         env(
             "DATABASE_URL",
-            default="postgresql://postgres:password@localhost:5432/flagsmith",
-        )
-    )
+            default="postgresql://postgres@localhost:5432/postgres",
+        ),
+    ),
+    "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",

src/task_processor/processor.py

@@ -9,6 +9,7 @@
 from django.utils import timezone
 
 from task_processor import metrics
+from task_processor.managers import RecurringTaskManager, TaskManager
 from task_processor.models import (
     AbstractBaseTask,
     RecurringTask,
@@ -27,14 +28,14 @@
 UNREGISTERED_RECURRING_TASK_GRACE_PERIOD = timedelta(minutes=30)
 
 
-def run_tasks(num_tasks: int = 1) -> list[TaskRun]:
+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(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)")
+        logger.debug(f"Running {len(tasks)} task(s) from database '{database}'")
 
         executed_tasks = []
         task_runs = []
@@ -47,25 +48,28 @@ def run_tasks(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)
-            logger.debug(f"Finished running {len(task_runs)} task(s)")
+            TaskRun.objects.using(database).bulk_create(task_runs)
+            logger.debug(
+                f"Finished running {len(task_runs)} task(s) from database '{database}'"
+            )
 
         return task_runs
 
     return []
 
 
-def run_recurring_tasks() -> list[RecurringTaskRun]:
+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()
+    task_manager: RecurringTaskManager = RecurringTask.objects.db_manager(database)
+    tasks = task_manager.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

@@ -4,6 +4,7 @@
 from datetime import datetime, timedelta
 from threading import Thread
 
+from django.conf import settings
 from django.db import close_old_connections
 from django.utils import timezone
 
@@ -95,18 +96,46 @@ def run(self) -> None:
             time.sleep(self.sleep_interval_millis / 1000)
 
     def run_iteration(self) -> None:
-        try:
-            run_tasks(self.queue_pop_size)
-            run_recurring_tasks()
-        except Exception as e:
-            # To prevent task threads from dying if they get an error retrieving the tasks from the
-            # database this will allow the thread to continue trying to retrieve tasks if it can
-            # successfully re-establish a connection to the database.
-            # TODO: is this also what is causing tasks to get stuck as locked? Can we unlock
-            #  tasks here?
-
-            logger.error("Received error retrieving tasks: %s.", e, exc_info=e)
-            close_old_connections()
+        """
+        Consume and execute tasks from the queue, and run recurring tasks
+
+        This method tries to consume tasks from multiple databases as to ensure
+        that any remaining tasks are processed after opting in or out of a
+        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)
+
+                # Recurring tasks are only run on one database
+                if (database == "default") ^ database_is_separate:
+                    run_recurring_tasks(database)
+            except Exception as exception:
+                # To prevent task threads from dying if they get an error retrieving the tasks from the
+                # database this will allow the thread to continue trying to retrieve tasks if it can
+                # successfully re-establish a connection to the database.
+                # TODO: is this also what is causing tasks to get stuck as locked? Can we unlock
+                #  tasks here?
+                exception_repr = f"{exception.__class__.__module__}.{repr(exception)}"
+                logger.error(
+                    f"Error handling tasks from database '{database}': {exception_repr}",
+                    exc_info=exception,
+                )
+
+                close_old_connections()
 
     def stop(self) -> None:
         self._stopped = True