snapshot of https://github.com/zulip/zulip.git @ 48c69911c7dd4e03b20552eb3bd77ad8151c1357 (2026-09-15)

zerver/lib/migrate.py

@@ -0,0 +1,139 @@
+import time
+from collections.abc import Callable
+from typing import Any
+
+from django.db import connection
+from django.db.backends.base.schema import BaseDatabaseSchemaEditor
+from django.db.backends.utils import CursorWrapper
+from django.db.migrations.state import StateApps
+from psycopg2.sql import SQL, Composable, Identifier
+
+
+def do_batch_update(
+    cursor: CursorWrapper,
+    table: str,
+    assignments: list[Composable],
+    batch_size: int = 10000,
+    sleep: float = 0.1,
+) -> None:
+    # The string substitution below is complicated by our need to
+    # support multiple PostgreSQL versions.
+    stmt = SQL(
+        """
+        UPDATE {}
+        SET {}
+        WHERE id >= %s AND id < %s
+    """
+    ).format(
+        Identifier(table),
+        SQL(", ").join(assignments),
+    )
+
+    cursor.execute(SQL("SELECT MIN(id), MAX(id) FROM {}").format(Identifier(table)))
+    (min_id, max_id) = cursor.fetchone()
+    if min_id is None:
+        return
+
+    print(f"\n    Range of rows to update: [{min_id}, {max_id}]")
+    while min_id <= max_id:
+        lower = min_id
+        upper = min_id + batch_size
+        print(f"    Updating range [{lower},{upper})")
+        cursor.execute(stmt, [lower, upper])
+
+        min_id = upper
+        time.sleep(sleep)
+
+        # Once we've finished, check if any new rows were inserted to the table
+        if min_id > max_id:
+            cursor.execute(SQL("SELECT MAX(id) FROM {}").format(Identifier(table)))
+            (max_id,) = cursor.fetchone()
+
+    print("    Finishing...", end="")
+
+
+def rename_indexes_constraints(
+    old_table: str, new_table: str
+) -> Callable[[StateApps, BaseDatabaseSchemaEditor], None]:
+    # NOTE: `get_constraints`, which this uses, does not include any
+    # information about if the index name was manually set from
+    # Django, nor if it is a partial index.  This has the theoretical
+    # possibility to cause false positives in the duplicate-index
+    # check below, which would incorrectly drop one of the wanted
+    # indexes.
+    #
+    # Before using this on a table, ensure that it does not use
+    # partial indexes, nor manually-named indexes.
+    def inner_migration(apps: StateApps, schema_editor: Any) -> None:
+        seen_indexes = set()
+        with connection.cursor() as cursor:
+            constraints = connection.introspection.get_constraints(cursor, old_table)
+
+            for old_name, constraint in constraints.items():
+                if constraint["check"]:
+                    suffix = "_check"
+                    is_index = False
+                elif constraint["foreign_key"] is not None:
+                    is_index = False
+                    to_table, to_column = constraint["foreign_key"]
+                    suffix = f"_fk_{to_table}_{to_column}"
+                elif constraint["primary_key"]:
+                    suffix = "_pk"
+                    is_index = True
+                elif constraint["unique"]:
+                    suffix = "_uniq"
+                    is_index = True
+                else:
+                    suffix = "_idx" if len(constraint["columns"]) > 1 else ""
+                    is_index = True
+                new_name = schema_editor._create_index_name(
+                    new_table, constraint["columns"], suffix
+                )
+                if new_name in seen_indexes:
+                    # This index duplicates one we already renamed,
+                    # and attempting to rename it would cause a
+                    # conflict.  Drop the duplicated index.
+                    if is_index:
+                        raw_query = SQL("DROP INDEX {old_name}").format(
+                            old_name=Identifier(old_name)
+                        )
+                    else:
+                        raw_query = SQL(
+                            "ALTER TABLE {table_name} DROP CONSTRAINT {old_name}"
+                        ).format(table_name=Identifier(old_table), old_name=Identifier(old_name))
+                    cursor.execute(raw_query)
+                    continue
+
+                seen_indexes.add(new_name)
+                if is_index:
+                    raw_query = SQL("ALTER INDEX {old_name} RENAME TO {new_name}").format(
+                        old_name=Identifier(old_name), new_name=Identifier(new_name)
+                    )
+                else:
+                    raw_query = SQL(
+                        "ALTER TABLE {old_table} RENAME CONSTRAINT {old_name} TO {new_name}"
+                    ).format(
+                        old_table=Identifier(old_table),
+                        old_name=Identifier(old_name),
+                        new_name=Identifier(new_name),
+                    )
+                cursor.execute(raw_query)
+
+            for sequence in connection.introspection.get_sequences(cursor, old_table):
+                old_name = sequence["name"]
+                column = sequence["column"]
+                new_name = f"{new_table}_{column}_seq"
+
+                raw_query = SQL("ALTER SEQUENCE {old_name} RENAME TO {new_name}").format(
+                    old_name=Identifier(old_name),
+                    new_name=Identifier(new_name),
+                )
+                cursor.execute(raw_query)
+
+            cursor.execute(
+                SQL("ALTER TABLE {old_table} RENAME TO {new_table}").format(
+                    old_table=Identifier(old_table), new_table=Identifier(new_table)
+                )
+            )
+
+    return inner_migration