Merge pull request #92 from Pogchamp-company/develop

Version 1.5.0

README.md

@@ -1,5 +1,3 @@
-from alembic_postgresql_enum.configuration import Config
-
 # alembic-postgresql-enum
 [<img src="https://img.shields.io/pypi/pyversions/alembic-postgresql-enum">](https://pypi.org/project/alembic-postgresql-enum/)
 [<img src="https://img.shields.io/pypi/v/alembic-postgresql-enum">](https://pypi.org/project/alembic-postgresql-enum/)
@@ -32,32 +30,19 @@ import alembic_postgresql_enum
 
 To the top of your migrations/env.py file.
 
-## Configuration
-
-You can configure this extension to disable parts of it, or to enable some feature flags
-
-To do so you need to call set_configuration function after the import:
-
-```python
-import alembic_postgresql_enum
-
-alembic_postgresql_enum.set_configuration(
-    alembic_postgresql_enum.Config(
-        add_type_ignore=True,
-    )
-)
-```
 ## Features
 
 * [Creation of enums](#creation-of-enum)
 * [Deletion of unreferenced enums](#deletion-of-unreferenced-enum)
-* [Creation of new enum values](#creation-of-new-enum-values)
-* [Deletion of enums values](#deletion-of-enums-values)
-* [Renaming of enum values](#rename-enum-value)
+* [Detection of enum values changes](#detection-of-enum-values-changes)
+  * [Creation of new enum values](#creation-of-new-enum-values)
+  * [Deletion of enums values](#deletion-of-enums-values)
+  * [Renaming of enum values](#rename-enum-value)
+* [Omitting managing enums](#omitting-managing-enums)
 
-### Creation of enum<a id="creation-of-enum"></a>
+## Creation of enum<a id="creation-of-enum"></a>
 
-#### When table is created
+### When table is created
 
 ```python
 class MyEnum(enum.Enum):
@@ -94,7 +79,7 @@ def downgrade():
     # ### end Alembic commands ###
 ```
 
-#### When column is added
+### When column is added
 ```python
 class MyEnum(enum.Enum):
     one = 1
@@ -126,7 +111,7 @@ def downgrade():
     # ### end Alembic commands ###
 ```
 
-### Deletion of unreferenced enum<a id="deletion-of-unreferenced-enum"></a>
+## Deletion of unreferenced enum<a id="deletion-of-unreferenced-enum"></a>
 If enum is defined in postgres schema, but its mentions removed from code - It will be automatically removed
 ```python
 class ExampleTable(BaseModel):
@@ -149,6 +134,10 @@ def downgrade():
     # ### end Alembic commands ###
 ```
 
+## Detection of enum values changes<a id="detection-of-enum-values-changes"></a>
+
+***Can be disabled with `detect_enum_values_changes` configuration flag turned off***
+
 ### Creation of new enum values<a id="creation-of-new-enum-values"></a>
 
 If new enum value is defined sync_enum_values function call will be added to migration to account for it
@@ -222,6 +211,7 @@ def downgrade():
     # ### end Alembic commands ###
 ```
 
+
 ### Rename enum value<a id="rename-enum-value"></a>
 In this case you must manually edit migration
 
@@ -286,3 +276,51 @@ def downgrade():
 Do not forget to switch places old and new values for downgrade
 
 All defaults in postgres will be renamed automatically as well
+
+## Omitting managing enums<a id="omitting-managing-enums"></a>
+
+If configured `include_name` function returns `False` given enum will be not managed.
+```python
+import alembic_postgresql_enum
+
+def include_name(name: str) -> bool:
+    return name not in ['enum-to-ignore', 'some-internal-enum']
+
+alembic_postgresql_enum.set_configuration(
+    alembic_postgresql_enum.Config(
+        include_name=include_name,
+    )
+)
+```
+
+Feature is similar to [sqlalchemy feature for tables](https://alembic.sqlalchemy.org/en/latest/autogenerate.html#omitting-table-names-from-the-autogenerate-process)
+
+## Configuration
+
+You can configure this extension to disable parts of it, or to enable some feature flags
+
+To do so you need to call set_configuration function after the import:
+
+```python
+import alembic_postgresql_enum
+
+alembic_postgresql_enum.set_configuration(
+    alembic_postgresql_enum.Config(
+        add_type_ignore=True,
+    )
+)
+```
+
+Available options:
+
+- `add_type_ignore` (`False` by default) - flag that can be turned on 
+to add `# type: ignore[attr-defined]` at the end of generated `op.sync_enum_values` calls.
+This is helpful if you are using type checker such as `mypy`.
+`type: ignore` is needed because there is no way to add new function to an existing alembic's `op`.
+
+- `include_name` (`lambda _: True` bby default) - it adds ability to ignore process enum by name in similar way alembic allows to define `include_name` function. 
+This property accepts function that takes enum name and returns whether it should be processed.  
+
+- `drop_unused_enums` (`True` by default) - feature flag that can be turned off to disable clean up of undeclared enums
+
+- `detect_enum_values_changes` (`True` by default) - feature flag that can be turned off to disable generation of `op.sync_enum_values`.

alembic_postgresql_enum/compare_dispatch.py

@@ -15,7 +15,7 @@
     drop_unused_enums,
 )
 from alembic_postgresql_enum.get_enum_data import get_defined_enums, get_declared_enums
-
+from alembic_postgresql_enum.configuration import get_configuration
 
 log = logging.getLogger(f"alembic.{__name__}")
 
@@ -48,6 +48,8 @@ def compare_enums(
     add_create_type_false(upgrade_ops)
     add_postgres_using_to_text(upgrade_ops)
 
+    configuration = get_configuration()
+
     schema_names = list(schema_names)
 
     # Issue #40
@@ -61,19 +63,27 @@ def compare_enums(
         if schema is None:
             schema = default_schema
 
-        definitions = get_defined_enums(autogen_context.connection, schema)
+        definitions = get_defined_enums(autogen_context.connection, schema, configuration.include_name)
+
         declarations = get_declared_enums(
-            autogen_context.metadata, schema, default_schema, autogen_context.connection, upgrade_ops
+            autogen_context.metadata,
+            schema,
+            default_schema,
+            autogen_context.connection,
+            upgrade_ops,
+            configuration.include_name,
         )
 
         create_new_enums(definitions, declarations.enum_values, schema, upgrade_ops)
 
-        drop_unused_enums(definitions, declarations.enum_values, schema, upgrade_ops)
+        if configuration.drop_unused_enums:
+            drop_unused_enums(definitions, declarations.enum_values, schema, upgrade_ops)
 
-        sync_changed_enums(
-            definitions,
-            declarations.enum_values,
-            declarations.enum_table_references,
-            schema,
-            upgrade_ops,
-        )
+        if configuration.detect_enum_values_changes:
+            sync_changed_enums(
+                definitions,
+                declarations.enum_values,
+                declarations.enum_table_references,
+                schema,
+                upgrade_ops,
+            )

alembic_postgresql_enum/configuration.py

@@ -1,9 +1,13 @@
 from dataclasses import dataclass
+from typing import Callable
 
 
 @dataclass
 class Config:
     add_type_ignore: bool = False
+    include_name: Callable[[str], bool] = lambda _: True
+    drop_unused_enums: bool = True
+    detect_enum_values_changes: bool = True
 
 
 _config = Config()

alembic_postgresql_enum/get_enum_data/declared_enums.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import Tuple, Any, Set, Union, List, TYPE_CHECKING, cast, Optional
+from typing import Tuple, Any, Set, Union, List, TYPE_CHECKING, cast, Optional, Callable
 
 import sqlalchemy
 from alembic.operations.ops import UpgradeOps
@@ -52,6 +52,7 @@ def get_declared_enums(
     default_schema: str,
     connection: "Connection",
     upgrade_ops: Optional[UpgradeOps] = None,
+    include_name: Callable[[str], bool] = lambda _: True,
 ) -> DeclaredEnumValues:
     """
     Return a dict mapping SQLAlchemy declared enumeration types to the set of their values
@@ -100,6 +101,9 @@ def get_declared_enums(
                 if not column_type_is_enum(column_type):
                     continue
 
+                if not include_name(column_type.name):
+                    continue
+
                 column_type_schema = column_type.schema or default_schema  # type: ignore[attr-defined]
                 if column_type_schema != schema:
                     continue

alembic_postgresql_enum/get_enum_data/defined_enums.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 
 from alembic_postgresql_enum.get_enum_data.types import EnumNamesToValues
 from alembic_postgresql_enum.sql_commands.enum_type import get_all_enums
@@ -16,17 +16,28 @@ def _remove_schema_prefix(enum_name: str, schema: str) -> str:
     return enum_name
 
 
-def get_defined_enums(connection: "Connection", schema: str) -> EnumNamesToValues:
+def get_defined_enums(
+    connection: "Connection", schema: str, include_name: Callable[[str], bool] = lambda _: True
+) -> EnumNamesToValues:
     """
     Return a dict mapping PostgreSQL defined enumeration types to the set of their
     defined values.
     :param conn:
         SQLAlchemy connection instance.
     :param str schema:
         Schema name (e.g. "public").
+    :param Callable[[str], bool] include_name:
+        Callable that returns True if the enum name should be included
     :returns DeclaredEnumValues:
         enum_definitions={
             "my_enum": tuple(["a", "b", "c"]),
         }
     """
-    return {_remove_schema_prefix(name, schema): tuple(values) for name, values in get_all_enums(connection, schema)}
+
+    return {
+        enum_name: tuple(values)
+        for enum_name, values in (
+            (_remove_schema_prefix(name, schema), values) for name, values in get_all_enums(connection, schema)
+        )
+        if include_name(enum_name)
+    }

alembic_postgresql_enum/sql_commands/column_default.py

@@ -66,7 +66,6 @@ def rename_default_if_required(
         return column_default_value
 
     if default_value.endswith("[]"):
-
         # remove old type postfix
         column_default_value = default_value[: default_value.find("::")]
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "alembic-postgresql-enum"
-version = "1.4.0"
+version = "1.5.0"
 description = "Alembic autogenerate support for creation, alteration and deletion of enums"
 authors = ["RustyGuard"]
 license = "MIT"

tests/get_enum_data/test_get_declared_enums.py

@@ -1,3 +1,4 @@
+import random
 from typing import TYPE_CHECKING
 
 from alembic_postgresql_enum.get_enum_data import TableReference, get_declared_enums
@@ -40,3 +41,27 @@ def test_with_multiple_enums(connection: "Connection"):
     function_result = get_declared_enums(declared_schema, DEFAULT_SCHEMA, DEFAULT_SCHEMA, connection)
 
     assert function_result == declared_enum_values
+
+
+def test_can_filter_by_name(connection: "Connection"):
+    declared_enum_values = get_declared_enum_values_with_orders_and_users()
+    declared_schema = get_schema_by_declared_enum_values(declared_enum_values)
+
+    all_declared_enums = declared_enum_values.enum_values
+    enum_to_exclude = random.choice(list(all_declared_enums))
+    enum_to_include = [e for e in all_declared_enums if e != enum_to_exclude]
+
+    function_result = get_declared_enums(
+        declared_schema,
+        DEFAULT_SCHEMA,
+        DEFAULT_SCHEMA,
+        connection,
+        None,
+        lambda enum_name: enum_name not in [enum_to_exclude],
+    )
+
+    assert enum_to_exclude not in function_result.enum_values
+    assert enum_to_exclude not in function_result.enum_table_references
+
+    assert set(enum_to_include) == set(function_result.enum_values)
+    assert set(enum_to_include) == set(function_result.enum_table_references)

tests/get_enum_data/test_get_defined_enums.py

@@ -1,3 +1,4 @@
+import random
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -31,3 +32,17 @@ def test_with_multiple_enums(connection: "Connection"):
     function_result = get_defined_enums(connection, DEFAULT_SCHEMA)
 
     assert function_result == declared_enum_values.enum_values
+
+
+def test_get_defined_enums_will_filter_by_name(connection: "Connection"):
+    declared_enum_values = get_declared_enum_values_with_orders_and_users()
+    defined_schema = get_schema_by_declared_enum_values(declared_enum_values)
+    all_declared_enums = declared_enum_values.enum_values
+    enum_to_exclude = random.choice(list(all_declared_enums))
+
+    defined_schema.create_all(connection)
+
+    function_result = get_defined_enums(
+        connection, DEFAULT_SCHEMA, lambda enum_name: enum_name not in [enum_to_exclude]
+    )
+    assert enum_to_exclude not in function_result

tests/sync_enum_values/test_render.py

@@ -26,6 +26,37 @@
 from tests.utils.migration_context import create_migration_context
 
 
+class TestDisabledSyncEnumValues(CompareAndRunTestCase):
+    """Check that enum variants are updated when new variant is added"""
+
+    config = Config(detect_enum_values_changes=False)
+
+    old_enum_variants = ["active", "passive"]
+    new_enum_variants = old_enum_variants + ["banned"]
+
+    def get_database_schema(self) -> MetaData:
+        schema = get_schema_with_enum_variants(self.old_enum_variants)
+        return schema
+
+    def get_target_schema(self) -> MetaData:
+        schema = get_schema_with_enum_variants(self.new_enum_variants)
+        return schema
+
+    def get_expected_upgrade(self) -> str:
+        return f"""
+        # ### commands auto generated by Alembic - please adjust! ###
+        pass
+        # ### end Alembic commands ###
+        """
+
+    def get_expected_downgrade(self) -> str:
+        return f"""
+        # ### commands auto generated by Alembic - please adjust! ###
+        pass
+        # ### end Alembic commands ###
+        """
+
+
 class TestAddNewEnumValueRender(CompareAndRunTestCase):
     """Check that enum variants are updated when new variant is added"""