feat(scheduled flag changes): implement scheduling of changes

latest_migrations.manifest

@@ -5,7 +5,7 @@ contenttypes: 0002_remove_content_type_name
 ee: 0015_add_verified_properties
 otp_static: 0002_throttling
 otp_totp: 0002_auto_20190420_0723
-posthog: 0378_alter_user_theme_mode
+posthog: 0379_alter_scheduledchange
 sessions: 0001_initial
 social_django: 0010_uid_db_index
 two_factor: 0007_auto_20201201_1019

posthog/celery.py

@@ -271,6 +271,13 @@ def setup_periodic_tasks(sender: Celery, **kwargs):
         name="recalculate cohorts",
     )
 
+    add_periodic_task_with_expiry(
+        sender,
+        120,
+        process_scheduled_changes.s(),
+        name="process scheduled changes",
+    )
+
     if clear_clickhouse_crontab := get_crontab(settings.CLEAR_CLICKHOUSE_REMOVED_DATA_SCHEDULE_CRON):
         sender.add_periodic_task(
             clear_clickhouse_crontab,
@@ -871,6 +878,13 @@ def calculate_cohort():
     calculate_cohorts()
 
 
+@app.task(ignore_result=True)
+def process_scheduled_changes():
+    from posthog.tasks.process_scheduled_changes import process_scheduled_changes
+
+    process_scheduled_changes()
+
+
 @app.task(ignore_result=True)
 def sync_insight_cache_states_task():
     from posthog.caching.insight_caching_state import sync_insight_cache_states

posthog/migrations/0379_alter_scheduledchange.py

@@ -0,0 +1,27 @@
+# Generated by Django 3.2.19 on 2023-12-21 11:41
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("posthog", "0378_alter_user_theme_mode"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="scheduledchange",
+            name="executed_at",
+            field=models.DateTimeField(blank=True, db_index=True, null=True),
+        ),
+        migrations.AlterField(
+            model_name="scheduledchange",
+            name="record_id",
+            field=models.CharField(max_length=200),
+        ),
+        migrations.AlterField(
+            model_name="scheduledchange",
+            name="scheduled_at",
+            field=models.DateTimeField(db_index=True),
+        ),
+    ]

posthog/models/feature_flag/feature_flag.py

@@ -297,6 +297,31 @@ def get_cohort_ids(
 
         return list(cohort_ids)
 
+    def scheduled_changes_dispatcher(self, payload):
+        from posthog.api.feature_flag import FeatureFlagSerializer
+
+        if "operation" not in payload or "value" not in payload:
+            raise Exception("Invalid payload")
+
+        context = {
+            "request": {"user": self.created_by},
+            "team_id": self.team_id,
+        }
+        serializer_data = {}
+
+        if payload["operation"] == "add_release_condition":
+            existing_groups = self.filters.get("groups", [])
+            new_groups = payload["value"].get("groups", [])
+            serializer_data["filters"] = {"groups": existing_groups + new_groups}
+        elif payload["operation"] == "update_status":
+            serializer_data["active"] = payload["value"]
+        else:
+            raise Exception(f"Unrecognized operation: {payload['operation']}")
+
+        serializer = FeatureFlagSerializer(self, data=serializer_data, context=context, partial=True)
+        if serializer.is_valid(raise_exception=True):
+            serializer.save()
+
     @property
     def uses_cohorts(self) -> bool:
         for condition in self.conditions:

posthog/models/scheduled_change.py

@@ -1,17 +1,16 @@
 from django.db import models
-from django.utils import timezone
 
 
 class ScheduledChange(models.Model):
     class AllowedModels(models.TextChoices):
         FEATURE_FLAG = "FeatureFlag", "feature flag"
 
     id = models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")
-    record_id = models.IntegerField()
+    record_id: models.CharField = models.CharField(max_length=200)
     model_name: models.CharField = models.CharField(max_length=100, choices=AllowedModels.choices)
     payload: models.JSONField = models.JSONField(default=dict)
-    scheduled_at: models.DateTimeField = models.DateTimeField(default=timezone.now)
-    executed_at: models.DateTimeField = models.DateTimeField(null=True, blank=True)
+    scheduled_at: models.DateTimeField = models.DateTimeField(db_index=True)
+    executed_at: models.DateTimeField = models.DateTimeField(null=True, blank=True, db_index=True)
     failure_reason = models.CharField(max_length=400, null=True, blank=True)
 
     team: models.ForeignKey = models.ForeignKey("Team", on_delete=models.CASCADE)

posthog/tasks/__init__.py

@@ -7,6 +7,7 @@
     demo_create_data,
     email,
     exporter,
+    process_scheduled_changes,
     split_person,
     sync_all_organization_available_features,
     usage_report,
@@ -20,6 +21,7 @@
     "demo_create_data",
     "email",
     "exporter",
+    "process_scheduled_changes",
     "split_person",
     "sync_all_organization_available_features",
     "user_identify",

posthog/tasks/process_scheduled_changes.py

@@ -0,0 +1,39 @@
+from posthog.models import ScheduledChange
+from django.utils import timezone
+from posthog.models import FeatureFlag
+from django.db import transaction, OperationalError
+
+models = {"FeatureFlag": FeatureFlag}
+
+
+def process_scheduled_changes() -> None:
+    try:
+        with transaction.atomic():
+            scheduled_changes = (
+                ScheduledChange.objects.select_for_update(nowait=True)
+                .filter(
+                    executed_at__isnull=True,
+                    scheduled_at__lt=timezone.now(),
+                )
+                .order_by("scheduled_at")[:10000]
+            )
+
+            for scheduled_change in scheduled_changes:
+                try:
+                    # Execute the change on the model instance
+                    model = models[scheduled_change.model_name]
+                    instance = model.objects.get(id=scheduled_change.record_id)
+                    instance.scheduled_changes_dispatcher(scheduled_change.payload)
+
+                    # Mark scheduled change completed
+                    scheduled_change.executed_at = timezone.now()
+                    scheduled_change.save()
+
+                except Exception as e:
+                    # Store the failure reason
+                    scheduled_change.failure_reason = str(e)
+                    scheduled_change.executed_at = timezone.now()
+                    scheduled_change.save()
+    except OperationalError:
+        # Failed to obtain the lock
+        pass