Merge pull request #59 from Pogchamp-company/develop

Version 1.1.0

.github/workflows/black.yml

@@ -0,0 +1,10 @@
+name: Lint
+
+on: [push, pull_request]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: psf/black@stable

alembic_postgresql_enum/add_create_type_false.py

@@ -1,7 +1,14 @@
 import logging
 
 import sqlalchemy
-from alembic.operations.ops import UpgradeOps, ModifyTableOps, AddColumnOp, CreateTableOp, DropColumnOp, DropTableOp
+from alembic.operations.ops import (
+    UpgradeOps,
+    ModifyTableOps,
+    AddColumnOp,
+    CreateTableOp,
+    DropColumnOp,
+    DropTableOp,
+)
 from sqlalchemy import Column
 from sqlalchemy.dialects import postgresql
 
@@ -10,13 +17,12 @@ class ReprWorkaround(postgresql.ENUM):
     """
     As postgresql.ENUM does not include create_type inside __repr__, we have to swap it with custom type
     """
-    __module__ = 'sqlalchemy.dialects.postgresql'
+
+    __module__ = "sqlalchemy.dialects.postgresql"
 
     def __repr__(self):
-        return (
-            f'{super().__repr__()[:-1]}, create_type=False)'
-            .replace('ReprWorkaround', 'ENUM')
-            .replace(', metadata=MetaData()', '')
+        return f"{super().__repr__()[:-1]}, create_type=False)".replace("ReprWorkaround", "ENUM").replace(
+            ", metadata=MetaData()", ""
         )
 
 
@@ -26,7 +32,7 @@ def inject_repr_into_enums(column: Column):
         if not column.type.native_enum:
             return
         log.info("%r converted into postgresql.ENUM", column.type)
-        column.type = eval(repr(column.type).replace('Enum', 'postgresql.ENUM'))
+        column.type = eval(repr(column.type).replace("Enum", "postgresql.ENUM"))
     if isinstance(column.type, postgresql.ENUM):
         if column.type.create_type:
             log.info("create_type=False injected into %r", column.type.name)
@@ -36,7 +42,7 @@ def inject_repr_into_enums(column: Column):
         column.type = replacement_enum_type
 
 
-log = logging.getLogger(f'alembic.{__name__}')
+log = logging.getLogger(f"alembic.{__name__}")
 
 
 def add_create_type_false(upgrade_ops: UpgradeOps):

alembic_postgresql_enum/add_postgres_using_to_text.py

@@ -14,17 +14,15 @@ class PostgresUsingAlterColumnOp(AlterColumnOp):
 
     def reverse(self):
         reversed_object = super().reverse()
-        reversed_object.kw.pop('postgresql_using', None)
+        reversed_object.kw.pop("postgresql_using", None)
         return reversed_object
 
 
 @renderers.dispatch_for(PostgresUsingAlterColumnOp)
-def _postgres_using_alter_column(
-        autogen_context: AutogenContext, op: ops.AlterColumnOp
-) -> str:
+def _postgres_using_alter_column(autogen_context: AutogenContext, op: ops.AlterColumnOp) -> str:
     alter_column_expression = render._alter_column(autogen_context, op)
 
-    postgresql_using = op.kw.get('postgresql_using', None)
+    postgresql_using = op.kw.get("postgresql_using", None)
     indent = " " * 11
 
     # To remove closing bracket
@@ -37,11 +35,11 @@ def _postgres_using_alter_column(
     return alter_column_expression
 
 
-log = logging.getLogger(f'alembic.{__name__}')
+log = logging.getLogger(f"alembic.{__name__}")
 
 
 def add_postgres_using_to_alter_operation(op: AlterColumnOp):
-    op.kw['postgresql_using'] = f'{op.column_name}::{op.modify_type.name}'
+    op.kw["postgresql_using"] = f"{op.column_name}::{op.modify_type.name}"
     log.info("postgresql_using added to %r.%r alteration", op.table_name, op.column_name)
     op.__class__ = PostgresUsingAlterColumnOp
 

alembic_postgresql_enum/compare_dispatch.py

@@ -5,13 +5,23 @@
 from alembic.operations.ops import UpgradeOps, CreateTableOp
 
 from alembic_postgresql_enum.add_create_type_false import add_create_type_false
-from alembic_postgresql_enum.add_postgres_using_to_text import add_postgres_using_to_text
-from alembic_postgresql_enum.detection_of_changes import sync_changed_enums, create_new_enums, drop_unused_enums
+from alembic_postgresql_enum.add_postgres_using_to_text import (
+    add_postgres_using_to_text,
+)
+from alembic_postgresql_enum.detection_of_changes import (
+    sync_changed_enums,
+    create_new_enums,
+    drop_unused_enums,
+)
 from alembic_postgresql_enum.get_enum_data import get_defined_enums, get_declared_enums
 
 
 @alembic.autogenerate.comparators.dispatch_for("schema")
-def compare_enums(autogen_context: AutogenContext, upgrade_ops: UpgradeOps, schema_names: Iterable[Union[str, None]]):
+def compare_enums(
+    autogen_context: AutogenContext,
+    upgrade_ops: UpgradeOps,
+    schema_names: Iterable[Union[str, None]],
+):
     """
     Walk the declared SQLAlchemy schema for every referenced Enum, walk the PG
     schema for every defined Enum, then generate SyncEnumValuesOp migrations
@@ -41,5 +51,10 @@ def compare_enums(autogen_context: AutogenContext, upgrade_ops: UpgradeOps, sche
 
         drop_unused_enums(definitions, declarations.enum_values, schema, upgrade_ops)
 
-        sync_changed_enums(definitions, declarations.enum_values,
-                           declarations.enum_table_references, schema, upgrade_ops)
+        sync_changed_enums(
+            definitions,
+            declarations.enum_values,
+            declarations.enum_table_references,
+            schema,
+            upgrade_ops,
+        )

alembic_postgresql_enum/detection_of_changes/enum_alteration.py

@@ -3,22 +3,27 @@
 SQLAlchemy enums.
 
 """
+
 import logging
 
 from alembic.operations.ops import UpgradeOps
 
-from alembic_postgresql_enum.get_enum_data import EnumNamesToValues, EnumNamesToTableReferences
+from alembic_postgresql_enum.get_enum_data import (
+    EnumNamesToValues,
+    EnumNamesToTableReferences,
+)
 from alembic_postgresql_enum.operations.sync_enum_values import SyncEnumValuesOp
 
-log = logging.getLogger(f'alembic.{__name__}')
+log = logging.getLogger(f"alembic.{__name__}")
 
 
-def sync_changed_enums(defined_enums: EnumNamesToValues,
-                       declared_enums: EnumNamesToValues,
-                       table_references: EnumNamesToTableReferences,
-                       schema: str,
-                       upgrade_ops: UpgradeOps,
-                       ):
+def sync_changed_enums(
+    defined_enums: EnumNamesToValues,
+    declared_enums: EnumNamesToValues,
+    table_references: EnumNamesToTableReferences,
+    schema: str,
+    upgrade_ops: UpgradeOps,
+):
     for enum_name, new_values in declared_enums.items():
         if enum_name not in defined_enums:
             # That is work for create_new_enums function
@@ -30,8 +35,18 @@ def sync_changed_enums(defined_enums: EnumNamesToValues,
             # Enum definition and declaration are in sync
             continue
 
-        log.info("Detected changed enum values in %r\nWas: %r\nBecome: %r", enum_name,
-                 list(old_values), list(new_values))
+        log.info(
+            "Detected changed enum values in %r\nWas: %r\nBecome: %r",
+            enum_name,
+            list(old_values),
+            list(new_values),
+        )
         affected_columns = table_references[enum_name]
-        op = SyncEnumValuesOp(schema, enum_name, list(old_values), list(new_values), list(affected_columns))
+        op = SyncEnumValuesOp(
+            schema,
+            enum_name,
+            list(old_values),
+            list(new_values),
+            list(affected_columns),
+        )
         upgrade_ops.ops.append(op)

alembic_postgresql_enum/detection_of_changes/enum_creation.py

@@ -5,11 +5,15 @@
 from alembic_postgresql_enum.get_enum_data import EnumNamesToValues
 from alembic_postgresql_enum.operations.create_enum import CreateEnumOp
 
-log = logging.getLogger(f'alembic.{__name__}')
+log = logging.getLogger(f"alembic.{__name__}")
 
 
-def create_new_enums(defined_enums: EnumNamesToValues, declared_enums: EnumNamesToValues,
-                     schema: str, upgrade_ops: UpgradeOps):
+def create_new_enums(
+    defined_enums: EnumNamesToValues,
+    declared_enums: EnumNamesToValues,
+    schema: str,
+    upgrade_ops: UpgradeOps,
+):
     """
     Create enums that are not in Postgres schema
     """

alembic_postgresql_enum/detection_of_changes/enum_deletion.py

@@ -5,11 +5,15 @@
 from alembic_postgresql_enum.get_enum_data import EnumNamesToValues
 from alembic_postgresql_enum.operations.drop_enum import DropEnumOp
 
-log = logging.getLogger(f'alembic.{__name__}')
+log = logging.getLogger(f"alembic.{__name__}")
 
 
-def drop_unused_enums(defined_enums: EnumNamesToValues, declared_enums: EnumNamesToValues,
-                      schema: str, upgrade_ops: UpgradeOps):
+def drop_unused_enums(
+    defined_enums: EnumNamesToValues,
+    declared_enums: EnumNamesToValues,
+    schema: str,
+    upgrade_ops: UpgradeOps,
+):
     """
     Drop enums that are in Postgres schema but not declared in SqlAlchemy schema
     """

alembic_postgresql_enum/get_enum_data/__init__.py

@@ -1,3 +1,9 @@
-from .types import ColumnType, DeclaredEnumValues, EnumNamesToValues, EnumNamesToTableReferences, TableReference
+from .types import (
+    ColumnType,
+    DeclaredEnumValues,
+    EnumNamesToValues,
+    EnumNamesToTableReferences,
+    TableReference,
+)
 from .defined_enums import get_defined_enums
 from .declared_enums import get_declared_enums

alembic_postgresql_enum/get_enum_data/declared_enums.py

@@ -10,22 +10,26 @@
 if TYPE_CHECKING:
     from sqlalchemy.engine import Connection
 
-from alembic_postgresql_enum.get_enum_data import DeclaredEnumValues, TableReference, ColumnType
+from alembic_postgresql_enum.get_enum_data import (
+    DeclaredEnumValues,
+    TableReference,
+    ColumnType,
+)
 
 
-def get_enum_values(enum_type: sqlalchemy.Enum) -> 'Tuple[str, ...]':
+def get_enum_values(enum_type: sqlalchemy.Enum) -> "Tuple[str, ...]":
     # For specific case when types.TypeDecorator is used
     if isinstance(enum_type, sqlalchemy.types.TypeDecorator):
         dialect = postgresql.dialect
 
         def value_processor(value):
-            return enum_type.process_bind_param(
-                enum_type.impl.result_processor(dialect, enum_type)(value),
-                dialect
-            )
+            return enum_type.process_bind_param(enum_type.impl.result_processor(dialect, enum_type)(value), dialect)
+
     else:
+
         def value_processor(enum_value):
             return enum_value
+
     return tuple(value_processor(value) for value in enum_type.enums)
 
 
@@ -34,17 +38,18 @@ def column_type_is_enum(column_type: Any) -> bool:
         return column_type.native_enum
 
     # For specific case when types.TypeDecorator is used
-    if isinstance(getattr(column_type, 'impl', None), sqlalchemy.Enum):
+    if isinstance(getattr(column_type, "impl", None), sqlalchemy.Enum):
         return True
 
     return False
 
 
-def get_declared_enums(metadata: Union[MetaData, List[MetaData]],
-                       schema: str,
-                       default_schema: str,
-                       connection: 'Connection',
-                       ) -> DeclaredEnumValues:
+def get_declared_enums(
+    metadata: Union[MetaData, List[MetaData]],
+    schema: str,
+    default_schema: str,
+    connection: "Connection",
+) -> DeclaredEnumValues:
     """
     Return a dict mapping SQLAlchemy declared enumeration types to the set of their values
     with columns where enums are used.
@@ -95,14 +100,22 @@ def get_declared_enums(metadata: Union[MetaData, List[MetaData]],
                 if column_type.name not in enum_name_to_values:
                     enum_name_to_values[column_type.name] = get_enum_values(column_type)
 
-                column_default = get_column_default(connection, schema, table.name, column.name)
+                table_schema = table.schema or default_schema
+                column_default = get_column_default(connection, table.schema, table.name, column.name)
                 enum_name_to_table_references[column_type.name].add(
-                    TableReference(table.name, column.name, column_type_wrapper, column_default)
+                    TableReference(
+                        table_schema=table_schema,
+                        table_name=table.name,
+                        column_name=column.name,
+                        column_type=column_type_wrapper,
+                        existing_server_default=column_default,
+                    )
                 )
 
     return DeclaredEnumValues(
         enum_values=enum_name_to_values,
-        enum_table_references={enum_name: frozenset(table_references)
-                               for enum_name, table_references
-                               in enum_name_to_table_references.items()},
+        enum_table_references={
+            enum_name: frozenset(table_references)
+            for enum_name, table_references in enum_name_to_table_references.items()
+        },
     )

alembic_postgresql_enum/get_enum_data/defined_enums.py

@@ -8,15 +8,15 @@
 
 
 def _remove_schema_prefix(enum_name: str, schema: str) -> str:
-    schema_prefix = f'{schema}.'
+    schema_prefix = f"{schema}."
 
     if enum_name.startswith(schema_prefix):
-        enum_name = enum_name[len(schema_prefix):]
+        enum_name = enum_name[len(schema_prefix) :]
 
     return enum_name
 
 
-def get_defined_enums(connection: 'Connection', schema: str) -> EnumNamesToValues:
+def get_defined_enums(connection: "Connection", schema: str) -> EnumNamesToValues:
     """
     Return a dict mapping PostgreSQL defined enumeration types to the set of their
     defined values.
@@ -29,7 +29,4 @@ def get_defined_enums(connection: 'Connection', schema: str) -> EnumNamesToValue
             "my_enum": tuple(["a", "b", "c"]),
         }
     """
-    return {
-        _remove_schema_prefix(name, schema): tuple(values)
-        for name, values in get_all_enums(connection, schema)
-    }
+    return {_remove_schema_prefix(name, schema): tuple(values) for name, values in get_all_enums(connection, schema)}