Skip to content

Commit

Permalink
features: fixes Pogchamp-company#33 allow to ignore enums by name
Browse files Browse the repository at this point in the history
  • Loading branch information
dswistowski committed Dec 2, 2024
1 parent bbb6729 commit 2d7925e
Show file tree
Hide file tree
Showing 6 changed files with 60 additions and 9 deletions.
19 changes: 19 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ alembic_postgresql_enum.set_configuration(
* [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>

Expand Down Expand Up @@ -284,3 +285,21 @@ 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)
7 changes: 5 additions & 2 deletions alembic_postgresql_enum/compare_dispatch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__}")

Expand Down Expand Up @@ -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
Expand All @@ -61,7 +63,8 @@ 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
)
Expand Down
2 changes: 2 additions & 0 deletions alembic_postgresql_enum/configuration.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
from dataclasses import dataclass
from typing import Callable


@dataclass
class Config:
add_type_ignore: bool = False
include_name: Callable[[str], bool] = lambda _: True


_config = Config()
Expand Down
17 changes: 14 additions & 3 deletions alembic_postgresql_enum/get_enum_data/defined_enums.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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)
}
12 changes: 8 additions & 4 deletions tests/base/run_migration_test_abc.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,19 +19,23 @@ class CompareAndRunTestCase(ABC):
config = Config()

@abstractmethod
def get_database_schema(self) -> MetaData: ...
def get_database_schema(self) -> MetaData:
...

@abstractmethod
def get_target_schema(self) -> MetaData: ...
def get_target_schema(self) -> MetaData:
...

def insert_migration_data(self, connection: "Connection", database_schema: MetaData) -> None:
pass

@abstractmethod
def get_expected_upgrade(self) -> str: ...
def get_expected_upgrade(self) -> str:
...

@abstractmethod
def get_expected_downgrade(self) -> str: ...
def get_expected_downgrade(self) -> str:
...

def get_expected_imports(self) -> Optional[str]:
return None
Expand Down
12 changes: 12 additions & 0 deletions tests/get_enum_data/test_get_defined_enums.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import random
from typing import TYPE_CHECKING

if TYPE_CHECKING:
Expand Down Expand Up @@ -31,3 +32,14 @@ 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

0 comments on commit 2d7925e

Please sign in to comment.