AlertRuleProject backfill migration (#69343)

Backfill migration follow up for
https://github.com/getsentry/sentry/pull/69323

migrations_lockfile.txt

@@ -9,5 +9,5 @@ feedback: 0004_index_together
 hybridcloud: 0016_add_control_cacheversion
 nodestore: 0002_nodestore_no_dictfield
 replays: 0004_index_together
-sentry: 0701_backfill_alertrule_user_team
+sentry: 0702_alert_rule_project_backfill_migration_2
 social_auth: 0002_default_auto_field

src/sentry/migrations/0702_alert_rule_project_backfill_migration_2.py

@@ -0,0 +1,108 @@
+# Generated by Django 5.0.3 on 2024-04-19 17:53
+
+
+import logging
+
+from django.db import migrations
+from django.db.backends.base.schema import BaseDatabaseSchemaEditor
+from django.db.migrations.state import StateApps
+
+from sentry.new_migrations.migrations import CheckedMigration
+from sentry.utils.query import RangeQuerySetWrapper
+
+logger = logging.getLogger(__name__)
+
+
+def _backfill_alert_rule_projects(apps: StateApps, schema_editor: BaseDatabaseSchemaEditor) -> None:
+    QuerySubscriptions = apps.get_model("sentry", "QuerySubscription")
+    AlertRuleProjects = apps.get_model("sentry", "AlertRuleProjects")
+
+    # use RangeQuerySetWrapper to avoid loading all subscriptions into memory
+    for subscription in RangeQuerySetWrapper(
+        QuerySubscriptions.objects.all().prefetch_related("snuba_query__alertrule_set")
+    ):
+
+        snuba_query = subscription.snuba_query
+        if not snuba_query:
+            logger.warning(
+                "QuerySubscription found with no snuba_query",
+                extra={"query_subscription_id": subscription.id},
+            )
+            continue
+
+        alert_rule_set = list(snuba_query.alertrule_set.all())
+        if not len(alert_rule_set):
+            logger.warning(
+                "QuerySubscription + SnubaQuery found with no alert_rule",
+                extra={
+                    "query_subscription_id": subscription.id,
+                    "snuba_query_id": snuba_query.id,
+                },
+            )
+            continue
+        elif len(alert_rule_set) > 1:
+            logger.warning(
+                "QuerySubscription + SnubaQuery found with multiple alert_rules",
+                extra={
+                    "query_subscription_id": subscription.id,
+                    "snuba_query_id": snuba_query.id,
+                    "alert_rule_ids": [alert_rule.id for alert_rule in alert_rule_set],
+                },
+            )
+
+        # Default to the first alert rule
+        alert_rule = alert_rule_set[0]
+
+        existing_alert_rule_projects = list(AlertRuleProjects.objects.filter(alert_rule=alert_rule))
+        should_create_new = True
+
+        if len(existing_alert_rule_projects) > 0:
+            for arp in existing_alert_rule_projects:
+                if arp.project_id != subscription.project_id:
+                    logger.warning(
+                        "AlertRuleProject found with different project than subscription",
+                        extra={
+                            "alert_rule_id": alert_rule.id,
+                            "subscription_id": subscription.id,
+                            "subscription_project": subscription.project_id,
+                            "alert_rule_project": arp.project_id,
+                        },
+                    )
+                    arp.delete()
+                else:
+                    should_create_new = False
+
+        if should_create_new:
+            AlertRuleProjects.objects.create(
+                alert_rule=alert_rule,
+                project=subscription.project,
+            )
+
+
+class Migration(CheckedMigration):
+    # This flag is used to mark that a migration shouldn't be automatically run in production.
+    # This should only be used for operations where it's safe to run the migration after your
+    # code has deployed. So this should not be used for most operations that alter the schema
+    # of a table.
+    # Here are some things that make sense to mark as post deployment:
+    # - Large data migrations. Typically we want these to be run manually so that they can be
+    #   monitored and not block the deploy for a long period of time while they run.
+    # - Adding indexes to large tables. Since this can take a long time, we'd generally prefer to
+    #   run this outside deployments so that we don't block them. Note that while adding an index
+    #   is a schema change, it's completely safe to run the operation after the code has deployed.
+    # Once deployed, run these manually via: https://develop.sentry.dev/database-migrations/#migration-deployment
+
+    is_post_deployment = True
+
+    dependencies = [
+        ("sentry", "0701_backfill_alertrule_user_team"),
+    ]
+
+    operations = [
+        # Run the data migration
+        migrations.RunPython(
+            _backfill_alert_rule_projects,
+            migrations.RunPython.noop,
+            hints={"tables": ["sentry_alertruleprojects", "sentry_querysubscription"]},
+        ),
+    ]

tests/sentry/migrations/test_0687_alert_rule_project_backfill_migration.py → tests/sentry/migrations/test_0702_alert_rule_project_backfill_migration_2.py

@@ -8,8 +8,8 @@ from sentry.testutils.cases import TestMigrations
 
 
 class AlertRuleProjectBackfillTest(TestMigrations):
-    migrate_from = "0686_remove_config_from_checkin_state_operation"
-    migrate_to = "0687_alert_rule_project_backfill_migration"
+    migrate_from = "0701_backfill_alertrule_user_team"
+    migrate_to = "0702_alert_rule_project_backfill_migration_2"
 
     def setup_before_migration(self, app):
         self.snuba_query = create_snuba_query(