realm: Add uuid and associated secret columns - backfill migrations.

This commit is contained in:
Mateusz Mandera 2023-10-12 20:14:57 +02:00 committed by Tim Abbott
parent 3afe585922
commit 73cc680289
3 changed files with 96 additions and 0 deletions

View File

@ -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

View File

@ -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
),
]

View File

@ -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),
),
]