migration: Make 'rename_indexes_constraints' a lib function.

This prep commit moves the 'rename_indexes_constraints'
function to 'lib/migrate' as we're going to re-use it for
the 'UserHotspot' to 'OnboardingStep' table rename operation.

In general, this function would be helpful in migrations
involving table rename operations, subject to the caution
mentioned in the function via comments.

tools/semgrep-py.yml

@@ -68,6 +68,7 @@ rules:
         - zerver/migrations/0209_user_profile_no_empty_password.py
         - zerver/migrations/0260_missed_message_addresses_from_redis_to_db.py
         - zerver/migrations/0387_reupload_realmemoji_again.py
+        - zerver/migrations/0443_userpresence_new_table_schema.py
         - pgroonga/migrations/0002_html_escape_subject.py
 
   - id: html-format

zerver/lib/migrate.py

@@ -1,7 +1,10 @@
 import time
-from typing import List
+from typing import Any, Callable, List
 
+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
 
 
@@ -46,3 +49,88 @@ def do_batch_update(
             (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, infodict in constraints.items():
+                if infodict["check"]:
+                    suffix = "_check"
+                    is_index = False
+                elif infodict["foreign_key"] is not None:
+                    is_index = False
+                    to_table, to_column = infodict["foreign_key"]
+                    suffix = f"_fk_{to_table}_{to_column}"
+                elif infodict["primary_key"]:
+                    suffix = "_pk"
+                    is_index = True
+                elif infodict["unique"]:
+                    suffix = "_uniq"
+                    is_index = True
+                else:
+                    suffix = "_idx" if len(infodict["columns"]) > 1 else ""
+                    is_index = True
+                new_name = schema_editor._create_index_name(new_table, infodict["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 infodict in connection.introspection.get_sequences(cursor, old_table):
+                old_name = infodict["name"]
+                column = infodict["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

zerver/migrations/0443_userpresence_new_table_schema.py

@@ -1,98 +1,9 @@
-from typing import Any, Callable
-
 import django.db.models.deletion
 import django.utils.timezone
 from django.conf import settings
-from django.db import connection, migrations, models
+from django.db import migrations, models
-from django.db.backends.base.schema import BaseDatabaseSchemaEditor
-from django.db.migrations.state import StateApps
-from psycopg2.sql import SQL, Identifier
 
-
+from zerver.lib.migrate import rename_indexes_constraints
-def rename_indexes_constraints(
-    old_table: str, new_table: str
-) -> Callable[[StateApps, BaseDatabaseSchemaEditor], None]:
-    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)
-
-            # NOTE: `get_constraints` 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.  Neither partial indexes
-            # nor explicit index names are used on the UserPresence
-            # table as of when this migration runs, so this use is
-            # safe, but use caution if reusing this code.
-
-            for old_name, infodict in constraints.items():
-                if infodict["check"]:
-                    suffix = "_check"
-                    is_index = False
-                elif infodict["foreign_key"] is not None:
-                    is_index = False
-                    to_table, to_column = infodict["foreign_key"]
-                    suffix = f"_fk_{to_table}_{to_column}"
-                elif infodict["primary_key"]:
-                    suffix = "_pk"
-                    is_index = True
-                elif infodict["unique"]:
-                    suffix = "_uniq"
-                    is_index = True
-                else:
-                    suffix = "_idx" if len(infodict["columns"]) > 1 else ""
-                    is_index = True
-                new_name = schema_editor._create_index_name(new_table, infodict["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 infodict in connection.introspection.get_sequences(cursor, old_table):
-                old_name = infodict["name"]
-                column = infodict["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
 
 
 class Migration(migrations.Migration):