flagsmith-common 2.2.4__py3-none-any.whl

task_processor/health.py

@@ -0,0 +1,44 @@
+import uuid
+
+import backoff
+from health_check.backends import BaseHealthCheckBackend  # type: ignore[import-untyped]
+from health_check.exceptions import HealthCheckException  # type: ignore[import-untyped]
+
+from task_processor.models import HealthCheckModel
+from task_processor.tasks import create_health_check_model
+
+
+def is_processor_healthy(max_tries: int = 5, factor: float = 0.1) -> bool:
+    health_check_model_uuid = str(uuid.uuid4())
+
+    create_health_check_model.delay(args=(health_check_model_uuid,))
+
+    @backoff.on_predicate(
+        backoff.expo,
+        lambda m: m is None,
+        max_tries=max_tries,
+        factor=factor,
+        jitter=None,
+    )
+    def get_health_check_model() -> HealthCheckModel | None:
+        return HealthCheckModel.objects.filter(uuid=health_check_model_uuid).first()
+
+    health_check_model = get_health_check_model()
+    if health_check_model:
+        health_check_model.delete()
+        return True
+
+    return False
+
+
+class TaskProcessorHealthCheckBackend(BaseHealthCheckBackend):  # type: ignore[misc]
+    #: The status endpoints will respond with a 200 status code
+    #: even if the check errors.
+    critical_service = False
+
+    def check_status(self) -> None:
+        if not is_processor_healthy():
+            raise HealthCheckException("Task processor is unable to process tasks.")
+
+    def identifier(self) -> str:
+        return self.__class__.__name__  # Display name on the endpoint.

task_processor/managers.py

@@ -0,0 +1,18 @@
+import typing
+
+from django.db.models import Manager
+
+if typing.TYPE_CHECKING:
+    from django.db.models.query import RawQuerySet
+
+    from task_processor.models import RecurringTask, Task
+
+
+class TaskManager(Manager["Task"]):
+    def get_tasks_to_process(self, num_tasks: int) -> "RawQuerySet[Task]":
+        return self.raw("SELECT * FROM get_tasks_to_process(%s)", [num_tasks])
+
+
+class RecurringTaskManager(Manager["RecurringTask"]):
+    def get_tasks_to_process(self) -> "RawQuerySet[RecurringTask]":
+        return self.raw("SELECT * FROM get_recurringtasks_to_process()")

task_processor/metrics.py

@@ -0,0 +1,22 @@
+import prometheus_client
+from django.conf import settings
+
+from common.prometheus import Histogram
+
+flagsmith_task_processor_enqueued_tasks_total = prometheus_client.Counter(
+    "flagsmith_task_processor_enqueued_tasks_total",
+    "Total number of enqueued tasks.",
+    ["task_identifier"],
+)
+
+if settings.DOCGEN_MODE or settings.TASK_PROCESSOR_MODE:
+    flagsmith_task_processor_finished_tasks_total = prometheus_client.Counter(
+        "flagsmith_task_processor_finished_tasks_total",
+        "Total number of finished tasks. Only collected by Task Processor. `task_type` label is either `recurring` or `standard`.",
+        ["task_identifier", "task_type", "result"],
+    )
+    flagsmith_task_processor_task_duration_seconds = Histogram(
+        "flagsmith_task_processor_task_duration_seconds",
+        "Task processor task duration in seconds. Only collected by Task Processor. `task_type` label is either `recurring` or `standard`.",
+        ["task_identifier", "task_type", "result"],
+    )

task_processor/migrations/0001_initial.py

@@ -0,0 +1,44 @@
+# Generated by Django 3.2.14 on 2022-08-02 11:25
+
+from django.db import migrations, models
+import django.db.models.deletion
+import django.utils.timezone
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = [
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='Task',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('uuid', models.UUIDField(default=uuid.uuid4, unique=True)),
+                ('created_at', models.DateTimeField(auto_now_add=True)),
+                ('scheduled_for', models.DateTimeField(blank=True, default=django.utils.timezone.now, null=True)),
+                ('task_identifier', models.CharField(max_length=200)),
+                ('serialized_args', models.TextField(blank=True, null=True)),
+                ('serialized_kwargs', models.TextField(blank=True, null=True)),
+                ('num_failures', models.IntegerField(default=0)),
+            ],
+            options={
+                'index_together': {('scheduled_for', 'num_failures')},
+            },
+        ),
+        migrations.CreateModel(
+            name='TaskRun',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('started_at', models.DateTimeField()),
+                ('finished_at', models.DateTimeField(blank=True, null=True)),
+                ('result', models.CharField(blank=True, choices=[('SUCCESS', 'Success'), ('FAILURE', 'Failure')], db_index=True, max_length=50, null=True)),
+                ('error_details', models.TextField(blank=True, null=True)),
+                ('task', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='task_runs', to='task_processor.task')),
+            ],
+        ),
+    ]

task_processor/migrations/0002_healthcheckmodel.py

@@ -0,0 +1,21 @@
+# Generated by Django 3.2.14 on 2022-08-12 11:39
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('task_processor', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='HealthCheckModel',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('created_at', models.DateTimeField(auto_now_add=True)),
+                ('uuid', models.UUIDField(unique=True)),
+            ],
+        ),
+    ]

task_processor/migrations/0003_add_completed_to_task.py

@@ -0,0 +1,22 @@
+# Generated by Django 3.2.15 on 2022-08-24 13:53
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('task_processor', '0002_healthcheckmodel'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='task',
+            name='completed',
+            field=models.BooleanField(default=False),
+        ),
+        migrations.AlterIndexTogether(
+            name='task',
+            index_together={('scheduled_for', 'num_failures', 'completed')},
+        ),
+    ]

task_processor/migrations/0004_recreate_task_indexes.py

@@ -0,0 +1,43 @@
+# Generated by Django 3.2.15 on 2022-10-07 09:53
+
+from django.db import migrations, models
+
+from common.migrations.helpers import PostgresOnlyRunSQL
+
+
+class Migration(migrations.Migration):
+
+    atomic = False
+
+    dependencies = [
+        ("task_processor", "0003_add_completed_to_task"),
+    ]
+
+    operations = [
+        migrations.SeparateDatabaseAndState(
+            state_operations=[
+                migrations.AlterIndexTogether(
+                    name="task",
+                    index_together=set(),
+                ),
+                migrations.AddIndex(
+                    model_name="task",
+                    index=models.Index(
+                        condition=models.Q(("completed", False)),
+                        fields=["num_failures", "scheduled_for"],
+                        name="incomplete_tasks_idx",
+                    ),
+                ),
+            ],
+            database_operations=[
+                PostgresOnlyRunSQL(
+                    "DROP INDEX CONCURRENTLY task_processor_task_scheduled_for_num_failur_17d6dc77_idx;",
+                    reverse_sql='CREATE INDEX "task_processor_task_scheduled_for_num_failur_17d6dc77_idx" ON "task_processor_task" ("scheduled_for", "num_failures", "completed");',
+                ),
+                PostgresOnlyRunSQL(
+                    'CREATE INDEX CONCURRENTLY "incomplete_tasks_idx" ON "task_processor_task" ("num_failures", "scheduled_for") WHERE NOT "completed";',
+                    reverse_sql='DROP INDEX CONCURRENTLY "incomplete_tasks_idx";',
+                ),
+            ],
+        ),
+    ]

task_processor/migrations/0005_update_conditional_index_conditions.py

@@ -0,0 +1,45 @@
+# Generated by Django 3.2.15 on 2022-10-07 11:16
+
+from django.db import migrations, models
+
+from common.migrations.helpers import PostgresOnlyRunSQL
+
+
+class Migration(migrations.Migration):
+
+    atomic = False
+
+    dependencies = [
+        ("task_processor", "0004_recreate_task_indexes"),
+    ]
+
+    operations = [
+        migrations.SeparateDatabaseAndState(
+            state_operations=[
+                migrations.RemoveIndex(
+                    model_name="task",
+                    name="incomplete_tasks_idx",
+                ),
+                migrations.AddIndex(
+                    model_name="task",
+                    index=models.Index(
+                        condition=models.Q(
+                            ("completed", False), ("num_failures__lt", 3)
+                        ),
+                        fields=["scheduled_for"],
+                        name="incomplete_tasks_idx",
+                    ),
+                ),
+            ],
+            database_operations=[
+                PostgresOnlyRunSQL(
+                    'DROP INDEX CONCURRENTLY "incomplete_tasks_idx";',
+                    reverse_sql='CREATE INDEX CONCURRENTLY "incomplete_tasks_idx" ON "task_processor_task" ("num_failures", "scheduled_for") WHERE NOT "completed";',
+                ),
+                PostgresOnlyRunSQL(
+                    'CREATE INDEX CONCURRENTLY "incomplete_tasks_idx" ON "task_processor_task" ("scheduled_for") WHERE (NOT "completed" and "num_failures" < 3);',
+                    reverse_sql='DROP INDEX CONCURRENTLY "incomplete_tasks_idx";',
+                ),
+            ],
+        )
+    ]

task_processor/migrations/0006_auto_20230221_0802.py

@@ -0,0 +1,45 @@
+# Generated by Django 3.2.16 on 2023-02-21 08:02
+
+from django.db import migrations, models
+import django.db.models.deletion
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('task_processor', '0005_update_conditional_index_conditions'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='RecurringTask',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('uuid', models.UUIDField(default=uuid.uuid4, unique=True)),
+                ('created_at', models.DateTimeField(auto_now_add=True)),
+                ('task_identifier', models.CharField(max_length=200)),
+                ('serialized_args', models.TextField(blank=True, null=True)),
+                ('serialized_kwargs', models.TextField(blank=True, null=True)),
+                ('run_every', models.DurationField()),
+            ],
+        ),
+        migrations.CreateModel(
+            name='RecurringTaskRun',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('started_at', models.DateTimeField()),
+                ('finished_at', models.DateTimeField(blank=True, null=True)),
+                ('result', models.CharField(blank=True, choices=[('SUCCESS', 'Success'), ('FAILURE', 'Failure')], db_index=True, max_length=50, null=True)),
+                ('error_details', models.TextField(blank=True, null=True)),
+                ('task', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='task_runs', to='task_processor.recurringtask')),
+            ],
+            options={
+                'abstract': False,
+            },
+        ),
+        migrations.AddConstraint(
+            model_name='recurringtask',
+            constraint=models.UniqueConstraint(fields=('task_identifier', 'run_every'), name='unique_run_every_tasks'),
+        ),
+    ]

task_processor/migrations/0007_add_is_locked.py

@@ -0,0 +1,23 @@
+# Generated by Django 3.2.18 on 2023-04-20 02:43
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('task_processor', '0006_auto_20230221_0802'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='recurringtask',
+            name='is_locked',
+            field=models.BooleanField(default=False),
+        ),
+        migrations.AddField(
+            model_name='task',
+            name='is_locked',
+            field=models.BooleanField(default=False),
+        ),
+    ]

task_processor/migrations/0008_add_get_task_to_process_function.py

@@ -0,0 +1,31 @@
+# Generated by Django 3.2.18 on 2023-04-20 02:45
+
+from django.db import migrations
+
+from common.migrations.helpers import PostgresOnlyRunSQL
+import os
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("task_processor", "0007_add_is_locked"),
+    ]
+
+    operations = [
+        PostgresOnlyRunSQL.from_sql_file(
+            os.path.join(
+                os.path.dirname(__file__),
+                "sql",
+                "0008_get_tasks_to_process.sql",
+            ),
+            reverse_sql="DROP FUNCTION IF EXISTS get_tasks_to_process",
+        ),
+        PostgresOnlyRunSQL.from_sql_file(
+            os.path.join(
+                os.path.dirname(__file__),
+                "sql",
+                "0008_get_recurring_tasks_to_process.sql",
+            ),
+            reverse_sql="DROP FUNCTION IF EXISTS get_recurringtasks_to_process",
+        ),
+    ]

task_processor/migrations/0009_add_recurring_task_run_first_run_at.py

@@ -0,0 +1,18 @@
+# Generated by Django 3.2.18 on 2023-04-05 13:47
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('task_processor', '0008_add_get_task_to_process_function'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='recurringtask',
+            name='first_run_time',
+            field=models.TimeField(blank=True, null=True),
+        ),
+    ]

task_processor/migrations/0010_task_priority.py

@@ -0,0 +1,27 @@
+# Generated by Django 3.2.20 on 2023-10-13 06:04
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("task_processor", "0009_add_recurring_task_run_first_run_at"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="task",
+            name="priority",
+            field=models.SmallIntegerField(
+                choices=[
+                    (100, "Lower"),
+                    (75, "Low"),
+                    (50, "Normal"),
+                    (25, "High"),
+                    (0, "Highest"),
+                ],
+                default=None,
+                null=True,
+            ),
+        ),
+    ]

task_processor/migrations/0011_add_priority_to_get_tasks_to_process.py

@@ -0,0 +1,27 @@
+# Generated by Django 3.2.20 on 2023-10-13 04:44
+
+from django.db import migrations
+
+from common.migrations.helpers import PostgresOnlyRunSQL
+import os
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("task_processor", "0010_task_priority"),
+    ]
+
+    operations = [
+        PostgresOnlyRunSQL.from_sql_file(
+            os.path.join(
+                os.path.dirname(__file__),
+                "sql",
+                "0011_get_tasks_to_process.sql",
+            ),
+            reverse_sql=os.path.join(
+                os.path.dirname(__file__),
+                "sql",
+                "0008_get_tasks_to_process.sql",
+            ),
+        ),
+    ]

task_processor/migrations/0012_add_locked_at_and_timeout.py

@@ -0,0 +1,40 @@
+# Generated by Django 3.2.23 on 2025-01-06 04:51
+
+import datetime
+from django.db import migrations, models
+import os
+
+from common.migrations.helpers import PostgresOnlyRunSQL
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("task_processor", "0011_add_priority_to_get_tasks_to_process"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="recurringtask",
+            name="locked_at",
+            field=models.DateTimeField(blank=True, null=True),
+        ),
+        migrations.AddField(
+            model_name="recurringtask",
+            name="timeout",
+            field=models.DurationField(default=datetime.timedelta(minutes=30)),
+        ),
+        migrations.AddField(
+            model_name="task",
+            name="timeout",
+            field=models.DurationField(blank=True, null=True),
+        ),
+        PostgresOnlyRunSQL.from_sql_file(
+            os.path.join(
+                os.path.dirname(__file__),
+                "sql",
+                "0012_get_recurringtasks_to_process.sql",
+            ),
+            reverse_sql="DROP FUNCTION IF EXISTS get_recurringtasks_to_process()",
+        ),
+    ]

task_processor/migrations/0013_add_last_picked_at.py

@@ -0,0 +1,34 @@
+# Generated by Django 4.2.18 on 2025-04-03 07:34
+
+from django.db import migrations, models
+import django.utils.timezone
+import os
+from common.migrations.helpers import PostgresOnlyRunSQL
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("task_processor", "0012_add_locked_at_and_timeout"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="recurringtask",
+            name="last_picked_at",
+            field=models.DateTimeField(blank=True, null=True),
+            preserve_default=False,
+        ),
+        PostgresOnlyRunSQL.from_sql_file(
+            os.path.join(
+                os.path.dirname(__file__),
+                "sql",
+                "0013_get_recurringtasks_to_process.sql",
+            ),
+            reverse_sql=os.path.join(
+                os.path.dirname(__file__),
+                "sql",
+                "0012_get_recurringtasks_to_process.sql",
+            ),
+        ),
+    ]

task_processor/migrations/sql/0008_get_recurring_tasks_to_process.sql

@@ -0,0 +1,30 @@
+CREATE OR REPLACE FUNCTION get_recurringtasks_to_process(num_tasks integer)
+RETURNS SETOF task_processor_recurringtask AS $$
+DECLARE
+    row_to_return task_processor_recurringtask;
+BEGIN
+    -- Select the tasks that needs to be processed
+    FOR row_to_return IN
+        SELECT *
+        FROM task_processor_recurringtask
+        WHERE is_locked = FALSE
+        ORDER BY id
+        LIMIT num_tasks
+        -- Select for update to ensure that no other workers can select these tasks while in this transaction block
+        FOR UPDATE SKIP LOCKED
+    LOOP
+        -- Lock every selected task(by updating `is_locked` to true)
+        UPDATE task_processor_recurringtask
+        -- Lock this row by setting is_locked True, so that no other workers can select these tasks after this
+        -- transaction is complete (but the tasks are still being executed by the current worker)
+        SET is_locked = TRUE
+        WHERE id = row_to_return.id;
+        -- If we don't explicitly update the `is_locked` column here, the client will receive the row that is actually locked but has the `is_locked` value set to `False`.
+        row_to_return.is_locked := TRUE;
+        RETURN NEXT row_to_return;
+    END LOOP;
+
+    RETURN;
+END;
+$$ LANGUAGE plpgsql
+

task_processor/migrations/sql/0008_get_tasks_to_process.sql

@@ -0,0 +1,30 @@
+CREATE OR REPLACE FUNCTION get_tasks_to_process(num_tasks integer)
+RETURNS SETOF task_processor_task AS $$
+DECLARE
+    row_to_return task_processor_task;
+BEGIN
+    -- Select the tasks that needs to be processed
+    FOR row_to_return IN
+        SELECT *
+        FROM task_processor_task
+        WHERE num_failures < 3 AND scheduled_for < NOW() AND completed = FALSE AND is_locked = FALSE
+        ORDER BY scheduled_for ASC, created_at ASC
+        LIMIT num_tasks
+        -- Select for update to ensure that no other workers can select these tasks while in this transaction block
+        FOR UPDATE SKIP LOCKED
+    LOOP
+        -- Lock every selected task(by updating `is_locked` to true)
+        UPDATE task_processor_task
+        -- Lock this row by setting is_locked True, so that no other workers can select these tasks after this
+        -- transaction is complete (but the tasks are still being executed by the current worker)
+        SET is_locked = TRUE
+        WHERE id = row_to_return.id;
+        -- If we don't explicitly update the `is_locked` column here, the client will receive the row that is actually locked but has the `is_locked` value set to `False`.
+        row_to_return.is_locked := TRUE;
+        RETURN NEXT row_to_return;
+    END LOOP;
+
+    RETURN;
+END;
+$$ LANGUAGE plpgsql
+

task_processor/migrations/sql/0011_get_tasks_to_process.sql

@@ -0,0 +1,30 @@
+CREATE OR REPLACE FUNCTION get_tasks_to_process(num_tasks integer)
+RETURNS SETOF task_processor_task AS $$
+DECLARE
+    row_to_return task_processor_task;
+BEGIN
+    -- Select the tasks that needs to be processed
+    FOR row_to_return IN
+        SELECT *
+        FROM task_processor_task
+        WHERE num_failures < 3 AND scheduled_for < NOW() AND completed = FALSE AND is_locked = FALSE
+        ORDER BY priority ASC, scheduled_for ASC, created_at ASC
+        LIMIT num_tasks
+        -- Select for update to ensure that no other workers can select these tasks while in this transaction block
+        FOR UPDATE SKIP LOCKED
+    LOOP
+        -- Lock every selected task(by updating `is_locked` to true)
+        UPDATE task_processor_task
+        -- Lock this row by setting is_locked True, so that no other workers can select these tasks after this
+        -- transaction is complete (but the tasks are still being executed by the current worker)
+        SET is_locked = TRUE
+        WHERE id = row_to_return.id;
+        -- If we don't explicitly update the `is_locked` column here, the client will receive the row that is actually locked but has the `is_locked` value set to `False`.
+        row_to_return.is_locked := TRUE;
+        RETURN NEXT row_to_return;
+    END LOOP;
+
+    RETURN;
+END;
+$$ LANGUAGE plpgsql
+

task_processor/migrations/sql/0012_get_recurringtasks_to_process.sql

@@ -0,0 +1,33 @@
+CREATE OR REPLACE FUNCTION get_recurringtasks_to_process()
+RETURNS SETOF task_processor_recurringtask AS $$
+DECLARE
+    row_to_return task_processor_recurringtask;
+BEGIN
+    -- Select the tasks that needs to be processed
+    FOR row_to_return IN
+        SELECT *
+        FROM task_processor_recurringtask
+        -- Add one minute to the timeout as a grace period for overhead
+        WHERE is_locked = FALSE OR (locked_at IS NOT NULL AND locked_at < NOW() - timeout + INTERVAL '1 minute')
+        ORDER BY id
+        LIMIT 1
+        -- Select for update to ensure that no other workers can select these tasks while in this transaction block
+        FOR UPDATE SKIP LOCKED
+    LOOP
+        -- Lock every selected task(by updating `is_locked` to true)
+        UPDATE task_processor_recurringtask
+        -- Lock this row by setting is_locked True, so that no other workers can select these tasks after this
+        -- transaction is complete (but the tasks are still being executed by the current worker)
+        SET is_locked = TRUE, locked_at = NOW()
+        WHERE id = row_to_return.id;
+        -- If we don't explicitly update the columns here, the client will receive a row
+        -- that is locked but still shows `is_locked` as `False` and `locked_at` as `None`.
+        row_to_return.is_locked := TRUE;
+        row_to_return.locked_at := NOW();
+        RETURN NEXT row_to_return;
+    END LOOP;
+
+    RETURN;
+END;
+$$ LANGUAGE plpgsql
+