From 7607969f278af8d50f44dbe30acaa118ffa02dd1 Mon Sep 17 00:00:00 2001
From: Alex Vandiver <alexmv@zulip.com>
Date: Mon, 6 May 2024 15:26:17 +0000
Subject: [PATCH] zilencer: Add a unique index on RemoteRealm counts with
 RemoteRealm objects.

---
 ...emote_realm_installation_count_and_more.py | 81 +++++++++++++++++++
 zilencer/models.py                            | 16 +++-
 2 files changed, 96 insertions(+), 1 deletion(-)
 create mode 100644 zilencer/migrations/0060_remove_remoterealmcount_unique_remote_realm_installation_count_and_more.py

diff --git a/zilencer/migrations/0060_remove_remoterealmcount_unique_remote_realm_installation_count_and_more.py b/zilencer/migrations/0060_remove_remoterealmcount_unique_remote_realm_installation_count_and_more.py
new file mode 100644
index 0000000000..6374a1d4dc
--- /dev/null
+++ b/zilencer/migrations/0060_remove_remoterealmcount_unique_remote_realm_installation_count_and_more.py
@@ -0,0 +1,81 @@
+# Generated by Django 5.0.5 on 2024-05-06 15:17
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("zilencer", "0059_remoterealmauditlog_add_synced_billing_event_type_index"),
+    ]
+
+    operations = [
+        migrations.SeparateDatabaseAndState(
+            # Django does not provide a RenameConstraint operation.
+            # "Constraints" are created as indexes in PostgreSQL, so
+            # we rename the underlying indexes.
+            database_operations=[
+                migrations.RunSQL(
+                    sql=(
+                        "ALTER INDEX unique_remote_realm_installation_count "
+                        "RENAME TO unique_server_realm_installation_count"
+                    ),
+                    reverse_sql=(
+                        "ALTER INDEX unique_server_realm_installation_count "
+                        "RENAME TO unique_remote_realm_installation_count"
+                    ),
+                ),
+                migrations.RunSQL(
+                    sql=(
+                        "ALTER INDEX unique_remote_realm_installation_count_null_subgroup "
+                        "RENAME TO unique_server_realm_installation_count_null_subgroup"
+                    ),
+                    reverse_sql=(
+                        "ALTER INDEX unique_server_realm_installation_count_null_subgroup "
+                        "RENAME TO unique_remote_realm_installation_count_null_subgroup"
+                    ),
+                ),
+            ],
+            state_operations=[
+                migrations.RemoveConstraint(
+                    model_name="remoterealmcount",
+                    name="unique_remote_realm_installation_count",
+                ),
+                migrations.RemoveConstraint(
+                    model_name="remoterealmcount",
+                    name="unique_remote_realm_installation_count_null_subgroup",
+                ),
+                migrations.AddConstraint(
+                    model_name="remoterealmcount",
+                    constraint=models.UniqueConstraint(
+                        condition=models.Q(("subgroup__isnull", False)),
+                        fields=("server", "realm_id", "property", "subgroup", "end_time"),
+                        name="unique_server_realm_installation_count",
+                    ),
+                ),
+                migrations.AddConstraint(
+                    model_name="remoterealmcount",
+                    constraint=models.UniqueConstraint(
+                        condition=models.Q(("subgroup__isnull", True)),
+                        fields=("server", "realm_id", "property", "end_time"),
+                        name="unique_server_realm_installation_count_null_subgroup",
+                    ),
+                ),
+            ],
+        ),
+        migrations.AddConstraint(
+            model_name="remoterealmcount",
+            constraint=models.UniqueConstraint(
+                condition=models.Q(("subgroup__isnull", False)),
+                fields=("remote_realm_id", "property", "subgroup", "end_time"),
+                name="unique_remote_realm_installation_count",
+            ),
+        ),
+        migrations.AddConstraint(
+            model_name="remoterealmcount",
+            constraint=models.UniqueConstraint(
+                condition=models.Q(("subgroup__isnull", True)),
+                fields=("remote_realm_id", "property", "end_time"),
+                name="unique_remote_realm_installation_count_null_subgroup",
+            ),
+        ),
+    ]
diff --git a/zilencer/models.py b/zilencer/models.py
index 19023018db..095ba1ad4a 100644
--- a/zilencer/models.py
+++ b/zilencer/models.py
@@ -416,13 +416,27 @@ class RemoteRealmCount(BaseRemoteCount):
 
     class Meta:
         constraints = [
+            # These two constraints come from the information as
+            # provided by the remote server, for rows they provide.
             UniqueConstraint(
                 fields=["server", "realm_id", "property", "subgroup", "end_time"],
                 condition=Q(subgroup__isnull=False),
+                name="unique_server_realm_installation_count",
+            ),
+            UniqueConstraint(
+                fields=["server", "realm_id", "property", "end_time"],
+                condition=Q(subgroup__isnull=True),
+                name="unique_server_realm_installation_count_null_subgroup",
+            ),
+            # These two constraints come from our internal
+            # record-keeping, which has a RemoteRealm object.
+            UniqueConstraint(
+                fields=["remote_realm_id", "property", "subgroup", "end_time"],
+                condition=Q(subgroup__isnull=False),
                 name="unique_remote_realm_installation_count",
             ),
             UniqueConstraint(
-                fields=["server", "realm_id", "property", "end_time"],
+                fields=["remote_realm_id", "property", "end_time"],
                 condition=Q(subgroup__isnull=True),
                 name="unique_remote_realm_installation_count_null_subgroup",
             ),