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):