From 18ad18c0253f7585e68f4bcca9c84ad3d1ea4f26 Mon Sep 17 00:00:00 2001
From: Mateusz Mandera <mateusz.mandera@zulip.com>
Date: Mon, 2 Aug 2021 20:53:16 +0200
Subject: [PATCH] confirmation: Migration to add expiry_date step 2 - backfill.

---
 .../0009_confirmation_expiry_date_backfill.py | 68 +++++++++++++++++++
 1 file changed, 68 insertions(+)
 create mode 100644 confirmation/migrations/0009_confirmation_expiry_date_backfill.py

diff --git a/confirmation/migrations/0009_confirmation_expiry_date_backfill.py b/confirmation/migrations/0009_confirmation_expiry_date_backfill.py
new file mode 100644
index 0000000000..bc7a74b9df
--- /dev/null
+++ b/confirmation/migrations/0009_confirmation_expiry_date_backfill.py
@@ -0,0 +1,68 @@
+# Generated by Django 3.1.7 on 2021-03-31 20:47
+
+import time
+from datetime import timedelta
+
+from django.conf import settings
+from django.db import migrations, transaction
+from django.db.backends.postgresql.schema import DatabaseSchemaEditor
+from django.db.migrations.state import StateApps
+
+
+def set_expiry_date_for_existing_confirmations(
+    apps: StateApps, schema_editor: DatabaseSchemaEditor
+) -> None:
+    Confirmation = apps.get_model("confirmation", "Confirmation")
+    if not Confirmation.objects.exists():
+        return
+
+    # The values at the time of this migration
+    INVITATION = 2
+    UNSUBSCRIBE = 4
+    MULTIUSE_INVITE = 6
+
+    @transaction.atomic()
+    def backfill_confirmations_between(lower_bound: int, upper_bound: int) -> None:
+        confirmations = Confirmation.objects.filter(id__gte=lower_bound, id__lte=upper_bound)
+        for confirmation in confirmations:
+            if confirmation.type in (INVITATION, MULTIUSE_INVITE):
+                confirmation.expiry_date = confirmation.date_sent + timedelta(
+                    days=settings.INVITATION_LINK_VALIDITY_DAYS
+                )
+            elif confirmation.type == UNSUBSCRIBE:
+                # Unsubscribe links never expire, which we apparently implement as in 1M days.
+                confirmation.expiry_date = confirmation.date_sent + timedelta(days=1000000)
+            else:
+                confirmation.expiry_date = confirmation.date_sent + timedelta(
+                    days=settings.CONFIRMATION_LINK_DEFAULT_VALIDITY_DAYS
+                )
+        Confirmation.objects.bulk_update(confirmations, ["expiry_date"])
+
+    BATCH_SIZE = 1000
+
+    first_id = Confirmation.objects.earliest("id").id
+    last_id = Confirmation.objects.latest("id").id
+
+    id_range_lower_bound = first_id
+    id_range_upper_bound = first_id + BATCH_SIZE
+    while id_range_lower_bound <= last_id:
+        backfill_confirmations_between(id_range_lower_bound, id_range_upper_bound)
+        id_range_lower_bound = id_range_upper_bound + 1
+        id_range_upper_bound = id_range_lower_bound + BATCH_SIZE
+        time.sleep(0.1)
+
+
+class Migration(migrations.Migration):
+    atomic = False
+
+    dependencies = [
+        ("confirmation", "0008_confirmation_expiry_date"),
+    ]
+
+    operations = [
+        migrations.RunPython(
+            set_expiry_date_for_existing_confirmations,
+            reverse_code=migrations.RunPython.noop,
+            elidable=True,
+        ),
+    ]