From 9ea9847ee485a6bd16c41cdfe8efae965aa7388d Mon Sep 17 00:00:00 2001 From: rusty Date: Wed, 29 Jan 2025 21:01:10 +0400 Subject: [PATCH 1/5] Add `__all__` to `__init__.py` --- alembic_postgresql_enum/__init__.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/alembic_postgresql_enum/__init__.py b/alembic_postgresql_enum/__init__.py index 8b62712..23bbe4a 100644 --- a/alembic_postgresql_enum/__init__.py +++ b/alembic_postgresql_enum/__init__.py @@ -1,3 +1,10 @@ -from .compare_dispatch import compare_enums +from .compare_dispatch import compare_enums as _ from .get_enum_data import ColumnType, TableReference from .configuration import set_configuration, Config + +__all__ = ( + "ColumnType", + "TableReference", + "set_configuration", + "Config", +) From ff59f9b58cca6f79ac441be05a0903e006ba67c7 Mon Sep 17 00:00:00 2001 From: rusty Date: Wed, 29 Jan 2025 21:07:57 +0400 Subject: [PATCH 2/5] Downgrade test runner os version to ubuntu-22.04 --- .github/workflows/test_on_push.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test_on_push.yaml b/.github/workflows/test_on_push.yaml index f0299e9..cb88d16 100644 --- a/.github/workflows/test_on_push.yaml +++ b/.github/workflows/test_on_push.yaml @@ -13,7 +13,7 @@ on: jobs: run_tests: - runs-on: [ubuntu-latest] + runs-on: [ubuntu-22.04] strategy: matrix: sqlalchemy: [ "1.4", "2.0" ] From 191c4f1bc07dd0c33964281b72faa60f126fafd1 Mon Sep 17 00:00:00 2001 From: rusty Date: Wed, 29 Jan 2025 22:33:18 +0400 Subject: [PATCH 3/5] Escape schema, column and enum names with "quotes" --- .../get_enum_data/types.py | 6 +- .../operations/sync_enum_values.py | 15 ++-- .../sql_commands/column_default.py | 20 +++--- .../sql_commands/enum_type.py | 28 ++++---- tests/fixtures/db.py | 4 +- tests/schemas.py | 1 + .../test_exotic_column_name.py | 70 ++++++++++++++++++ .../sync_enum_values/test_exotic_enum_name.py | 71 +++++++++++++++++++ .../test_exotic_schema_name.py | 71 +++++++++++++++++++ 9 files changed, 253 insertions(+), 33 deletions(-) create mode 100644 tests/sync_enum_values/test_exotic_column_name.py create mode 100644 tests/sync_enum_values/test_exotic_enum_name.py create mode 100644 tests/sync_enum_values/test_exotic_schema_name.py diff --git a/alembic_postgresql_enum/get_enum_data/types.py b/alembic_postgresql_enum/get_enum_data/types.py index 904c09b..8e32908 100644 --- a/alembic_postgresql_enum/get_enum_data/types.py +++ b/alembic_postgresql_enum/get_enum_data/types.py @@ -45,11 +45,15 @@ def is_column_type_import_needed(self): @property def table_name_with_schema(self): if self.table_schema: - prefix = f"{self.table_schema}." + prefix = f'"{self.table_schema}".' else: prefix = "" return f'{prefix}"{self.table_name}"' + @property + def escaped_column_name(self): + return f'"{self.column_name}"' + EnumNamesToValues = Dict[str, Tuple[str, ...]] EnumNamesToTableReferences = Dict[str, FrozenSet[TableReference]] diff --git a/alembic_postgresql_enum/operations/sync_enum_values.py b/alembic_postgresql_enum/operations/sync_enum_values.py index a642107..427a162 100644 --- a/alembic_postgresql_enum/operations/sync_enum_values.py +++ b/alembic_postgresql_enum/operations/sync_enum_values.py @@ -76,11 +76,11 @@ def _set_enum_values( affected_columns: List[TableReference], enum_values_to_rename: List[Tuple[str, str]], ): - enum_type_name = f"{enum_schema}.{enum_name}" + enum_type_name = f'"{enum_schema}"."{enum_name}"' temporary_enum_name = f"{enum_name}_old" - rename_type(connection, enum_schema, enum_name, temporary_enum_name) - create_type(connection, enum_schema, enum_name, new_values) + rename_type(connection, enum_type_name, temporary_enum_name) + create_type(connection, enum_type_name, new_values) create_comparison_operators(connection, enum_schema, enum_name, temporary_enum_name, enum_values_to_rename) @@ -88,7 +88,7 @@ def _set_enum_values( column_default = table_reference.existing_server_default if column_default is not None: - drop_default(connection, table_reference.table_name_with_schema, table_reference.column_name) + drop_default(connection, table_reference) try: cast_old_enum_type_to_new(connection, table_reference, enum_type_name, enum_values_to_rename) @@ -104,12 +104,11 @@ def _set_enum_values( enum_schema, column_default, enum_name, enum_values_to_rename ) - set_default( - connection, table_reference.table_name_with_schema, table_reference.column_name, column_default - ) + set_default(connection, table_reference, column_default) drop_comparison_operators(connection, enum_schema, enum_name, temporary_enum_name) - drop_type(connection, enum_schema, temporary_enum_name) + temporary_enum_type_name = f'"{enum_schema}"."{temporary_enum_name}"' + drop_type(connection, temporary_enum_type_name) @classmethod def sync_enum_values( diff --git a/alembic_postgresql_enum/sql_commands/column_default.py b/alembic_postgresql_enum/sql_commands/column_default.py index f191cc9..278d5b0 100644 --- a/alembic_postgresql_enum/sql_commands/column_default.py +++ b/alembic_postgresql_enum/sql_commands/column_default.py @@ -3,6 +3,8 @@ import sqlalchemy +from alembic_postgresql_enum.get_enum_data.types import TableReference + if TYPE_CHECKING: from sqlalchemy.engine import Connection @@ -26,30 +28,32 @@ def get_column_default( return default_value -def drop_default(connection: "Connection", table_name_with_schema: str, column_name: str): +def drop_default( + connection: "Connection", + table_reference: TableReference, +): connection.execute( sqlalchemy.text( - f"""ALTER TABLE {table_name_with_schema} - ALTER COLUMN {column_name} DROP DEFAULT""" + f"""ALTER TABLE {table_reference.table_name_with_schema} + ALTER COLUMN {table_reference.escaped_column_name} DROP DEFAULT""" ) ) def set_default( connection: "Connection", - table_name_with_schema: str, - column_name: str, + table_reference: TableReference, default_value: str, ): connection.execute( sqlalchemy.text( - f"""ALTER TABLE {table_name_with_schema} - ALTER COLUMN {column_name} SET DEFAULT {default_value}""" + f"""ALTER TABLE {table_reference.table_name_with_schema} + ALTER COLUMN {table_reference.escaped_column_name} SET DEFAULT {default_value}""" ) ) -def rename_default_if_required( +def rename_default_if_required( # todo smells like teen shit schema: str, default_value: str, enum_name: str, diff --git a/alembic_postgresql_enum/sql_commands/enum_type.py b/alembic_postgresql_enum/sql_commands/enum_type.py index 3e63aeb..2a58ca1 100644 --- a/alembic_postgresql_enum/sql_commands/enum_type.py +++ b/alembic_postgresql_enum/sql_commands/enum_type.py @@ -14,7 +14,7 @@ def cast_old_array_enum_type_to_new( enum_type_name: str, enum_values_to_rename: List[Tuple[str, str]], ): - cast_clause = f"{table_reference.column_name}::text[]" + cast_clause = f"{table_reference.escaped_column_name}::text[]" for old_value, new_value in enum_values_to_rename: cast_clause = f"""array_replace({cast_clause}, '{old_value}', '{new_value}')""" @@ -22,7 +22,7 @@ def cast_old_array_enum_type_to_new( connection.execute( sqlalchemy.text( f"""ALTER TABLE {table_reference.table_name_with_schema} - ALTER COLUMN {table_reference.column_name} TYPE {enum_type_name}[] + ALTER COLUMN {table_reference.escaped_column_name} TYPE {enum_type_name}[] USING {cast_clause}::{enum_type_name}[] """ ) @@ -43,13 +43,13 @@ def cast_old_enum_type_to_new( connection.execute( sqlalchemy.text( f"""ALTER TABLE {table_reference.table_name_with_schema} - ALTER COLUMN {table_reference.column_name} TYPE {enum_type_name} + ALTER COLUMN {table_reference.escaped_column_name} TYPE {enum_type_name} USING CASE {' '.join( - f"WHEN {table_reference.column_name}::text = '{old_value}' THEN '{new_value}'::{enum_type_name}" + f"WHEN {table_reference.escaped_column_name}::text = '{old_value}' THEN '{new_value}'::{enum_type_name}" for old_value, new_value in enum_values_to_rename)} - ELSE {table_reference.column_name}::text::{enum_type_name} + ELSE {table_reference.escaped_column_name}::text::{enum_type_name} END """ ) @@ -58,26 +58,24 @@ def cast_old_enum_type_to_new( connection.execute( sqlalchemy.text( f"""ALTER TABLE {table_reference.table_name_with_schema} - ALTER COLUMN {table_reference.column_name} TYPE {enum_type_name} - USING {table_reference.column_name}::text::{enum_type_name} + ALTER COLUMN {table_reference.escaped_column_name} TYPE {enum_type_name} + USING {table_reference.escaped_column_name}::text::{enum_type_name} """ ) ) -def drop_type(connection: "Connection", schema: str, type_name: str): - connection.execute(sqlalchemy.text(f"""DROP TYPE {schema}.{type_name}""")) +def drop_type(connection: "Connection", enum_type_name: str): + connection.execute(sqlalchemy.text(f"""DROP TYPE {enum_type_name}""")) -def rename_type(connection: "Connection", schema: str, type_name: str, new_type_name: str): - connection.execute(sqlalchemy.text(f"""ALTER TYPE {schema}.{type_name} RENAME TO {new_type_name}""")) +def rename_type(connection: "Connection", enum_type_name: str, new_type_name: str): + connection.execute(sqlalchemy.text(f"""ALTER TYPE {enum_type_name} RENAME TO {new_type_name}""")) -def create_type(connection: "Connection", schema: str, type_name: str, enum_values: List[str]): +def create_type(connection: "Connection", enum_type_name: str, enum_values: List[str]): connection.execute( - sqlalchemy.text( - f"""CREATE TYPE {schema}.{type_name} AS ENUM({', '.join(f"'{value}'" for value in enum_values)})""" - ) + sqlalchemy.text(f"""CREATE TYPE {enum_type_name} AS ENUM({', '.join(f"'{value}'" for value in enum_values)})""") ) diff --git a/tests/fixtures/db.py b/tests/fixtures/db.py index 7b79703..2f38bd2 100644 --- a/tests/fixtures/db.py +++ b/tests/fixtures/db.py @@ -5,7 +5,7 @@ import sqlalchemy from sqlalchemy import create_engine -from tests.schemas import ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA +from tests.schemas import ANOTHER_SCHEMA_NAME, DEFAULT_SCHEMA, KEYWORD_SCHEMA_NAME try: import dotenv @@ -27,6 +27,8 @@ def connection() -> Generator: CREATE SCHEMA {DEFAULT_SCHEMA}; DROP SCHEMA IF EXISTS {ANOTHER_SCHEMA_NAME} CASCADE; CREATE SCHEMA {ANOTHER_SCHEMA_NAME}; + DROP SCHEMA IF EXISTS "{KEYWORD_SCHEMA_NAME}" CASCADE; + CREATE SCHEMA "{KEYWORD_SCHEMA_NAME}"; """ ) ) diff --git a/tests/schemas.py b/tests/schemas.py index 374ecb1..f972922 100644 --- a/tests/schemas.py +++ b/tests/schemas.py @@ -20,6 +20,7 @@ CAR_COLORS_ENUM_NAME = "car_color" ANOTHER_SCHEMA_NAME = "another" +KEYWORD_SCHEMA_NAME = "default" def get_schema_with_enum_variants(variants: List[str]) -> MetaData: diff --git a/tests/sync_enum_values/test_exotic_column_name.py b/tests/sync_enum_values/test_exotic_column_name.py new file mode 100644 index 0000000..60aa52b --- /dev/null +++ b/tests/sync_enum_values/test_exotic_column_name.py @@ -0,0 +1,70 @@ +from typing import Optional + +from sqlalchemy import MetaData, Table, Column, Integer +from sqlalchemy.dialects import postgresql + +from tests.schemas import USER_TABLE_NAME, USER_STATUS_ENUM_NAME +from tests.base.run_migration_test_abc import CompareAndRunTestCase + + +class TestExoticColumnNameRender(CompareAndRunTestCase): + """https://github.com/Pogchamp-company/alembic-postgresql-enum/issues/95""" + + old_enum_variants = ["active", "passive"] + new_enum_variants = old_enum_variants + ["banned"] + + def get_database_schema(self) -> MetaData: + schema = MetaData() + + Table( + USER_TABLE_NAME, + schema, + Column("id", Integer, primary_key=True), + Column( + "case", + postgresql.ENUM(*self.old_enum_variants, name=USER_STATUS_ENUM_NAME), + ), + ) + + return schema + + def get_target_schema(self) -> MetaData: + schema = MetaData() + + Table( + USER_TABLE_NAME, + schema, + Column("id", Integer, primary_key=True), + Column( + "case", + postgresql.ENUM(*self.new_enum_variants, name=USER_STATUS_ENUM_NAME), + ), + ) + + return schema + + def get_expected_upgrade(self) -> str: + return f""" + # ### commands auto generated by Alembic - please adjust! ### + op.sync_enum_values( + enum_schema='public', + enum_name='user_status', + new_values=['active', 'passive', 'banned'], + affected_columns=[TableReference(table_schema='public', table_name='users', column_name='case')], + enum_values_to_rename=[], + ) + # ### end Alembic commands ### + """ + + def get_expected_downgrade(self) -> Optional[str]: + return f""" + # ### commands auto generated by Alembic - please adjust! ### + op.sync_enum_values( + enum_schema='public', + enum_name='user_status', + new_values=['active', 'passive'], + affected_columns=[TableReference(table_schema='public', table_name='users', column_name='case')], + enum_values_to_rename=[], + ) + # ### end Alembic commands ### + """ diff --git a/tests/sync_enum_values/test_exotic_enum_name.py b/tests/sync_enum_values/test_exotic_enum_name.py new file mode 100644 index 0000000..9364ac4 --- /dev/null +++ b/tests/sync_enum_values/test_exotic_enum_name.py @@ -0,0 +1,71 @@ +from typing import Optional + +from sqlalchemy import MetaData, Table, Column, Integer +from sqlalchemy.dialects import postgresql + +from tests.schemas import USER_STATUS_COLUMN_NAME +from tests.schemas import USER_TABLE_NAME +from tests.base.run_migration_test_abc import CompareAndRunTestCase + + +class TestExoticColumnNameRender(CompareAndRunTestCase): + """Test for enum names that are keyword in postgres""" + + old_enum_variants = ["active", "passive"] + new_enum_variants = old_enum_variants + ["banned"] + + def get_database_schema(self) -> MetaData: + schema = MetaData() + + Table( + USER_TABLE_NAME, + schema, + Column("id", Integer, primary_key=True), + Column( + USER_STATUS_COLUMN_NAME, + postgresql.ENUM(*self.old_enum_variants, name="type"), + ), + ) + + return schema + + def get_target_schema(self) -> MetaData: + schema = MetaData() + + Table( + USER_TABLE_NAME, + schema, + Column("id", Integer, primary_key=True), + Column( + USER_STATUS_COLUMN_NAME, + postgresql.ENUM(*self.new_enum_variants, name="type"), + ), + ) + + return schema + + def get_expected_upgrade(self) -> str: + return f""" + # ### commands auto generated by Alembic - please adjust! ### + op.sync_enum_values( + enum_schema='public', + enum_name='type', + new_values=['active', 'passive', 'banned'], + affected_columns=[TableReference(table_schema='public', table_name='users', column_name='status')], + enum_values_to_rename=[], + ) + # ### end Alembic commands ### + """ + + def get_expected_downgrade(self) -> Optional[str]: + return f""" + # ### commands auto generated by Alembic - please adjust! ### + op.sync_enum_values( + enum_schema='public', + enum_name='type', + new_values=['active', 'passive'], + affected_columns=[TableReference(table_schema='public', table_name='users', column_name='status')], + enum_values_to_rename=[], + ) + # ### end Alembic commands ### + """ diff --git a/tests/sync_enum_values/test_exotic_schema_name.py b/tests/sync_enum_values/test_exotic_schema_name.py new file mode 100644 index 0000000..c6ae918 --- /dev/null +++ b/tests/sync_enum_values/test_exotic_schema_name.py @@ -0,0 +1,71 @@ +from typing import Optional + +from sqlalchemy import MetaData, Table, Column, Integer +from sqlalchemy.dialects import postgresql + +from tests.schemas import KEYWORD_SCHEMA_NAME +from tests.base.run_migration_test_abc import CompareAndRunTestCase + + +class TestExoticSchemaNameRender(CompareAndRunTestCase): + """https://github.com/Pogchamp-company/alembic-postgresql-enum/issues/95#issuecomment-2619738045""" + + old_variants = ("Walter White", "Jesse Pinkman", "Tuco") + new_variants = ("Walter White", "Jesse Pinkman") + enum_name = "dealers" + + def get_database_schema(self) -> MetaData: + metadata = MetaData() + Table( + "pricing", + metadata, + Column("id", Integer, primary_key=True), + Column( + "dealer", + postgresql.ENUM(*self.old_variants, name=self.enum_name), + ), + schema=KEYWORD_SCHEMA_NAME, + ) + + return metadata + + def get_target_schema(self) -> MetaData: + metadata = MetaData() + Table( + "pricing", + metadata, + Column("id", Integer, primary_key=True), + Column( + "dealer", + postgresql.ENUM(*self.new_variants, name=self.enum_name), + ), + schema=KEYWORD_SCHEMA_NAME, + ) + + return metadata + + def get_expected_upgrade(self) -> str: + return f""" + # ### commands auto generated by Alembic - please adjust! ### + op.sync_enum_values( + enum_schema='public', + enum_name='dealers', + new_values=['Walter White', 'Jesse Pinkman'], + affected_columns=[TableReference(table_schema='default', table_name='pricing', column_name='dealer')], + enum_values_to_rename=[], + ) + # ### end Alembic commands ### + """ + + def get_expected_downgrade(self) -> Optional[str]: + return f""" + # ### commands auto generated by Alembic - please adjust! ### + op.sync_enum_values( + enum_schema='public', + enum_name='dealers', + new_values=['Walter White', 'Jesse Pinkman', 'Tuco'], + affected_columns=[TableReference(table_schema='default', table_name='pricing', column_name='dealer')], + enum_values_to_rename=[], + ) + # ### end Alembic commands ### + """ From a3693fa8ec2d3c214160526cedf17df49748d0c7 Mon Sep 17 00:00:00 2001 From: rusty Date: Wed, 29 Jan 2025 22:34:28 +0400 Subject: [PATCH 4/5] Fix black version in job --- .github/workflows/black.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/black.yml b/.github/workflows/black.yml index 9065b5e..1e661e0 100644 --- a/.github/workflows/black.yml +++ b/.github/workflows/black.yml @@ -7,4 +7,4 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 - - uses: psf/black@stable + - uses: psf/black@25.1.0 From de2ab82d87f5ba9349891c3287cc2b33fcb6346a Mon Sep 17 00:00:00 2001 From: rusty Date: Wed, 29 Jan 2025 22:37:35 +0400 Subject: [PATCH 5/5] Bump version to 1.6.0 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 06b5926..4847b13 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "alembic-postgresql-enum" -version = "1.5.0" +version = "1.6.0" description = "Alembic autogenerate support for creation, alteration and deletion of enums" authors = ["RustyGuard"] license = "MIT"