From 32a5c422e933b570e01ffbf18cb2ac1786656f95 Mon Sep 17 00:00:00 2001 From: Prakhar Pratyush Date: Wed, 6 Dec 2023 10:39:58 +0530 Subject: [PATCH] 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 | 1 + zerver/lib/migrate.py | 90 +++++++++++++++++- .../0443_userpresence_new_table_schema.py | 93 +------------------ 3 files changed, 92 insertions(+), 92 deletions(-) diff --git a/tools/semgrep-py.yml b/tools/semgrep-py.yml index 4dfcbed31c..851e8a2f9c 100644 --- a/tools/semgrep-py.yml +++ b/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 diff --git a/zerver/lib/migrate.py b/zerver/lib/migrate.py index 5c341e6e01..f800f9ea9f 100644 --- a/zerver/lib/migrate.py +++ b/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 diff --git a/zerver/migrations/0443_userpresence_new_table_schema.py b/zerver/migrations/0443_userpresence_new_table_schema.py index 534e4153c9..18cecc2a62 100644 --- a/zerver/migrations/0443_userpresence_new_table_schema.py +++ b/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.backends.base.schema import BaseDatabaseSchemaEditor -from django.db.migrations.state import StateApps -from psycopg2.sql import SQL, Identifier +from django.db import migrations, models - -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 +from zerver.lib.migrate import rename_indexes_constraints class Migration(migrations.Migration):