Ability to avoid duplicate alembic comparator dispatch with `registered_entities` about alembic_utils HOT 10 CLOSED

Thanks for this improvement, we literally went from 10 minutes to 4 minutes in our CI just upgrading to alembic_utils 0.7.7 🎉
FYI we are using both libraries, pytest-alembic and alembic-utils !

from alembic_utils.

Love to see it, just upgraded and all seems well for our migrations. thanks!

from alembic_utils.

I see you're a maintainer of pytest-alembic. Could this be solved there using a pytest fixture to reset the contents of comparator between tests/calls into env.py?

from alembic_utils.

Probably not, the internal comparators.dispatch_for calls are at the module-level, and so resetting to pre-env.py state would end up being pretty wonky, if at all possible.

All pytest-alembic is doing is basically just calling parts of alembic's public api for programmatic use of alembic. The fact that using alembic's "command" API happens to reevaluate the env.py is, i guess, a kind of implementation detail of how alembic works, so other users of that api in combination with alembic_utils will run into the same thing.

I'd consider this more of a bug in my env.py for not being idempotent (Thus my feature request to enable one to write an env.py that is idempotent with this lib). Moving this call to, i.e. the models package would be another option to achieve the same thing but would unfortunately tie my models to alembic, which they are otherwise not.

As I kind of suggested above, alembic's internals seem to really want you to register comparators statically, and I haven't been able to find any hooks at the env.py level which would only get invoked once

from alembic_utils.

ok, makes sense

If you have any ideas for a backwards compatible way to refactor alembic_utils to avoid this issue I'd be happy to review. Getting the extensions to play nice together would be great

from alembic_utils.

class ReplaceableEntityRegistry:
    def __init__(self):
        self.entities: List[Type[ReplaceableEntity]] = []
        self.schemas: List[str] = []
        self.exclude_schemas: List[str] = []
        self.entity_typesList[Type[ReplaceableEntity]] = []

    def clear(self):
        self.entities.clear()
        self.schemas.clear()
        self.exclude_schemas.clear()
        self.entity_types.clear()

    def register(
        self,
        entities: List[T],
        schemas: Optional[List[str]] = None,
        exclude_schemas: Optional[List[str]] = None,
        entity_types: Optional[List[Type[ReplaceableEntity]]] = None,
    ) -> None:
         entities = 
         self.entities.extend(entites)

         if schemas:
             self.schemas.extend(schemas)

         if exclude_schemas:
             self.exclude_schemas.extend(exclude_schemas)

         if entity_types:
             if not self.entity_types:
                 self.entity_types.extend(entity_types)
             else:
                 self.entity_types.extend(collect_subclasses(alembic_utils, ReplaceableEntity))

registry = ReplaceableEntityRegistry()

def register_entities(
    entities: List[T],
    schemas: Optional[List[str]] = None,
    exclude_schemas: Optional[List[str]] = None,
    entity_types: Optional[List[Type[ReplaceableEntity]]] = None,
) -> None:
    registry.register(entities, schemas, exclude_schemas, entity_types)

@comparators.dispatch_for("schema")
def compare_registered_entities(
    autogen_context: AutogenContext,
    upgrade_ops,
    sqla_schemas: Optional[List[Optional[str]]],
):
    # identical to the inline function now, but uses of entities, schemas, exclude_schemas, and entity_types
    # all become prefixed with `registry.`

This is the method that seems most alike to how alembic itself works. The main semantic difference is that comparators.dispatch_for is registered at the module level (like how alembic itself does it), so there's only
ever one call to

As a fringe benefit (at least imo) this would make it easier to do other things which are neater by calling register_entites more than once, like

import foo.models
import bar models

register_entities(collect_instances(foo.models, PGView))
register_entities(collect_instances(bar.models, PGView))

I'd be happy to make a PR to this effect if you were into it

Alternatively, you could avoid using dispatch_for and do it manually, enabling a check

from alembic import comparators

def register_entities(...) -> None:
    def compare_registered_entities(...):

    registered_comparators = comparators._registry[("schema", "default")]
    names = {c.__name__ for c in registered_comparators}
    if compare_registered_entities.__name__ not in names:
        registered_comparators.append(compare_registered_entities)  # or vice versa where you overwrite it

this is kind of bad for a few reasons, but is more similar to how it works today

Alternatively, i could imagine something like

def register_entities(
    entities: List[T],
    schemas: Optional[List[str]] = None,
    exclude_schemas: Optional[List[str]] = None,
    entity_types: Optional[List[Type[ReplaceableEntity]]] = None,
) -> None:
    dispatcher = comparators.dispatch_for("schema")

    compare_registered_entities = compare_registered_entities_factory(entities, schemas, exclude_schemas, entity_types)
    dispatcher(compare_registered_entities)


def compare_registered_entities_factory(entities, schemas, exclude_schemas, entity_types):
    def compare_registered_entities(
        autogen_context: AutogenContext,
        upgrade_ops,
        sqla_schemas: Optional[List[Optional[str]]],
    ):
        # identical to the inline function now, but not decorated

    return compare_registered_entities

but this has a bunch of other drawbacks (which i can get into, if you really dont like the first solution), which means it's not ideal.

from alembic_utils.

thanks for taking the time to write that up 🙌

The ReplaceableEntityRegistry solution looks great

I'd be happy to make a PR to this effect if you were into it

yes please

Any chance you could also tweak the logic in ReplaceableEntityRegistry.register to deduplicate entities, schemas, exclude_schemas, and entity_types in case it is called multiple times?

from alembic_utils.

great! I'll see if I can get to that sometime later today!

from alembic_utils.

thanks! your patch is available in alembic_utils==0.7.7

from alembic_utils.

Thanks for doing this! It has saved me quite a few manual hacks!

from alembic_utils.