From 73cc68028954ef019de2d23970873777681ba288 Mon Sep 17 00:00:00 2001
From: Mateusz Mandera <mateusz.mandera@zulip.com>
Date: Thu, 12 Oct 2023 20:14:57 +0200
Subject: [PATCH] realm: Add uuid and associated secret columns - backfill
 migrations.

---
 tools/semgrep-py.yml                          |  1 +
 .../0480_realm_backfill_uuid_and_secret.py    | 69 +++++++++++++++++++
 ...ealm_uuid_alter_realm_uuid_owner_secret.py | 26 +++++++
 3 files changed, 96 insertions(+)
 create mode 100644 zerver/migrations/0480_realm_backfill_uuid_and_secret.py
 create mode 100644 zerver/migrations/0481_alter_realm_uuid_alter_realm_uuid_owner_secret.py

diff --git a/tools/semgrep-py.yml b/tools/semgrep-py.yml
index c8a39fec2a..ccd8c849bd 100644
--- a/tools/semgrep-py.yml
+++ b/tools/semgrep-py.yml
@@ -50,6 +50,7 @@ rules:
       - pattern-not: from zerver.models import filter_pattern_validator
       - pattern-not: from zerver.models import url_template_validator
       - pattern-not: from zerver.models import generate_email_token_for_stream
+      - pattern-not: from zerver.models import generate_realm_uuid_owner_secret
       - pattern-either:
           - pattern: from zerver import $X
           - pattern: from analytics import $X
diff --git a/zerver/migrations/0480_realm_backfill_uuid_and_secret.py b/zerver/migrations/0480_realm_backfill_uuid_and_secret.py
new file mode 100644
index 0000000000..75d47c90d1
--- /dev/null
+++ b/zerver/migrations/0480_realm_backfill_uuid_and_secret.py
@@ -0,0 +1,69 @@
+import secrets
+import uuid
+
+from django.db import migrations, models
+from django.db.backends.base.schema import BaseDatabaseSchemaEditor
+from django.db.migrations.state import StateApps
+
+
+def generate_api_key() -> str:
+    """
+    This is a copy of zerver.lib.utils.generate_api_key. Importing code that's prone
+    to change in a migration is something we generally avoid, to ensure predictable,
+    consistent behavior of the migration across time.
+    """
+
+    api_key = ""
+    while len(api_key) < 32:
+        api_key += secrets.token_urlsafe(3 * 9).replace("_", "").replace("-", "")
+    return api_key[:32]
+
+
+def generate_realm_uuid_owner_secret() -> str:
+    token = generate_api_key()
+
+    return f"zuliprealm_{token}"
+
+
+def backfill_realm_uuid_and_secret(
+    apps: StateApps, schema_editor: BaseDatabaseSchemaEditor
+) -> None:
+    Realm = apps.get_model("zerver", "Realm")
+
+    max_id = Realm.objects.aggregate(models.Max("id"))["id__max"]
+    if max_id is None:
+        # Nothing to do if there are no realms yet.
+        return
+
+    BATCH_SIZE = 100
+    lower_bound = 0
+
+    while lower_bound < max_id:
+        realms_to_update = []
+        for realm in Realm.objects.filter(
+            id__gt=lower_bound,
+            id__lte=lower_bound + BATCH_SIZE,
+            # We're setting uuid and uuid_owner_secret together, so it's enough
+            # to query for one of them being None.
+            uuid=None,
+        ).only("id", "uuid", "uuid_owner_secret"):
+            realm.uuid = uuid.uuid4()
+            realm.uuid_owner_secret = generate_realm_uuid_owner_secret()
+            realms_to_update.append(realm)
+        lower_bound += BATCH_SIZE
+
+        Realm.objects.bulk_update(realms_to_update, ["uuid", "uuid_owner_secret"])
+
+
+class Migration(migrations.Migration):
+    atomic = False
+
+    dependencies = [
+        ("zerver", "0479_realm_uuid_realm_uuid_owner_secret"),
+    ]
+
+    operations = [
+        migrations.RunPython(
+            backfill_realm_uuid_and_secret, reverse_code=migrations.RunPython.noop
+        ),
+    ]
diff --git a/zerver/migrations/0481_alter_realm_uuid_alter_realm_uuid_owner_secret.py b/zerver/migrations/0481_alter_realm_uuid_alter_realm_uuid_owner_secret.py
new file mode 100644
index 0000000000..c30da3f7db
--- /dev/null
+++ b/zerver/migrations/0481_alter_realm_uuid_alter_realm_uuid_owner_secret.py
@@ -0,0 +1,26 @@
+# Generated by Django 4.2.5 on 2023-10-06 10:08
+
+import uuid
+
+from django.db import migrations, models
+
+from zerver.models import generate_realm_uuid_owner_secret
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("zerver", "0480_realm_backfill_uuid_and_secret"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="realm",
+            name="uuid",
+            field=models.UUIDField(default=uuid.uuid4, unique=True),
+        ),
+        migrations.AlterField(
+            model_name="realm",
+            name="uuid_owner_secret",
+            field=models.TextField(default=generate_realm_uuid_owner_secret),
+        ),
+    ]