NameError Importing local packages in Alembic Operations 1.7.2

[daniel-butler]

Describe the bug
Importing local packages creates NameError in version 1.7.2 was not the case in 1.6.5

Expected behavior
Ability to use locally installed packages in alembic revisions

To Reproduce
Local Package "fun_stuff" installed via -e

The details on this file come from replaceable objects cookbook

fun_stuff/alembic_utils.py

from typing import Any, Optional

from alembic.operations import Operations, MigrateOperation


class ReplaceableObject:
    def __init__(self, name: str, sqltext: str) -> None:
        self.name = name
        self.sqltext = sqltext


class ReversibleOp(MigrateOperation):
    """This is the base of our “replaceable” operation, which includes not just a base operation for emitting
    CREATE and DROP instructions on a ReplaceableObject, it also assumes a certain model of “reversibility” which
    makes use of references to other migration files in order to refer to the “previous” version of an object.

    https://alembic.sqlalchemy.org/en/latest/cookbook.html#replaceable-objects
    """
    def __init__(self, target: ReplaceableObject) -> None:
        self.target = target

    @classmethod
    def invoke_for_target(cls, operations: Operations, target: ReplaceableObject):
        op = cls(target)
        return operations.invoke(op)

    def reverse(self):
        raise NotImplementedError

    @classmethod
    def _get_object_from_version(cls, operations: Operations, ident: str) -> Any:
        version, objectname = ident.split(".")
        module = operations.get_context().script.get_revision(version).module
        return getattr(module, objectname)

    @classmethod
    def replace(
            cls,
            operations: Operations,
            target: ReplaceableObject,
            replaces: Optional[str] = None,
            replace_with: Optional[str] = None,
    ) -> None:
        if replaces is None and replace_with is None:
            raise TypeError("replaces or replace_with is required")
        old_obj = cls._get_object_from_version(
            operations,
            replaces if replaces is not None else replace_with,
        )
        drop_old = cls(old_obj).reverse()
        create_new = cls(target)
        operations.invoke(drop_old)
        operations.invoke(create_new)


"""To create usable operations from this base, we will build a series of stub classes and use 
[Operations.register_operation()](https://alembic.sqlalchemy.org/en/latest/ops.html#alembic.operations.Operations.register_operation)
 to make them part of the op.* namespace"""


@Operations.register_operation("create_view", "invoke_for_target")
@Operations.register_operation("replace_view", "replace")
class CreateViewOp(ReversibleOp):
    def reverse(self):
        return DropViewOp(self.target)


@Operations.register_operation("drop_view", "invoke_for_target")
class DropViewOp(ReversibleOp):
    def reverse(self):
        return CreateViewOp(self.target)


@Operations.implementation_for(CreateViewOp)
def create_view(operations: Operations, operation: ReversibleOp) -> None:
    operations.execute(f"CREATE VIEW {operation.target.name} AS {operation.target.sqltext}")


@Operations.implementation_for(DropViewOp)
def drop_view(operations: Operations, operation: ReversibleOp) -> None:
    operations.execute(f"DROP VIEW {operation.target.name}")

24627210e3e6_swap_old_with_new.py

"""adjust view 

Revision ID: 24627210e3e6
Revises: bcf089c643e3
Create Date: 2021-07-16 17:03:05.673405

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy import orm

from fun_stuff import alembic_utils


# revision identifiers, used by Alembic.
revision = '24627210e3e6'
down_revision = 'bcf089c643e3'
branch_labels = None
depends_on = None


new_view = alembic_utils.ReplaceableObject(
    name="E_CrazyFunView",
    sqltext="""SELECT * FROM WowTable""",
)

old_view = alembic_utils.ReplaceableObject(
    name="E_CrazyFunView",
    sqltext="""SELECT * FROM NotWowTable""",
)


def upgrade():
    op.drop_view(old_view)
    op.create_view(new_view)


def downgrade():
    op.drop_view(new_view)
    op.create_view(old_view)

The error below occurs when running
alembic history

Error

Traceback (most recent call last):
  File "/usr/local/bin/alembic", line 8, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.8/site-packages/alembic/config.py", line 588, in main
    CommandLine(prog=prog).main(argv=argv)
  File "/usr/local/lib/python3.8/site-packages/alembic/config.py", line 582, in main
    self.run_cmd(cfg, options)
  File "/usr/local/lib/python3.8/site-packages/alembic/config.py", line 559, in run_cmd
    fn(
  File "/usr/local/lib/python3.8/site-packages/alembic/command.py", line 461, in history
    _display_history(config, script, base, head)
  File "/usr/local/lib/python3.8/site-packages/alembic/command.py", line 429, in _display_history
    for sc in script.walk_revisions(
  File "/usr/local/lib/python3.8/site-packages/alembic/script/base.py", line 277, in walk_revisions
    for rev in self.revision_map.iterate_revisions(
  File "/usr/local/lib/python3.8/site-packages/alembic/script/revision.py", line 793, in iterate_revisions
    revisions, heads = fn(
  File "/usr/local/lib/python3.8/site-packages/alembic/script/revision.py", line 1393, in _collect_upgrade_revisions
    targets: Collection["Revision"] = self._parse_upgrade_target(
  File "/usr/local/lib/python3.8/site-packages/alembic/script/revision.py", line 1193, in _parse_upgrade_target
    return self.get_revisions(target)
  File "/usr/local/lib/python3.8/site-packages/alembic/script/revision.py", line 527, in get_revisions
    resolved_id, branch_label = self._resolve_revision_number(
  File "/usr/local/lib/python3.8/site-packages/alembic/script/revision.py", line 747, in _resolve_revision_number
    self._revision_map
  File "/usr/local/lib/python3.8/site-packages/sqlalchemy/util/langhelpers.py", line 1113, in __get__
    obj.__dict__[self.__name__] = result = self.fget(obj)
  File "/usr/local/lib/python3.8/site-packages/alembic/script/revision.py", line 189, in _revision_map
    for revision in self._generator():
  File "/usr/local/lib/python3.8/site-packages/alembic/script/base.py", line 136, in _load_revisions
    script = Script._from_filename(self, vers, file_)
  File "/usr/local/lib/python3.8/site-packages/alembic/script/base.py", line 999, in _from_filename
    module = util.load_python_file(dir_, filename)
  File "/usr/local/lib/python3.8/site-packages/alembic/util/pyfiles.py", line 92, in load_python_file
    module = load_module_py(module_id, path)
  File "/usr/local/lib/python3.8/site-packages/alembic/util/pyfiles.py", line 108, in load_module_py
    spec.loader.exec_module(module)  # type: ignore
  File "<frozen importlib._bootstrap_external>", line 843, in exec_module
  File "<frozen importlib._bootstrap>", line 219, in _call_with_frames_removed
  File "/src/employees/alembic/versions/24627210e3e6_swap_old_with_new.py", line 13, in <module>
    from fun_stuff import alembic_utils
  File "/src/fun_stuff/alembic_utils.py", line 67, in <module>
    class CreateViewOp(ReversibleOp):
  File "/usr/local/lib/python3.8/site-packages/alembic/operations/base.py", line 163, in register
    exec(func_text, globals_, lcl)
  File "<string>", line 1, in <module>
NameError: name 'fun_stuff' is not defined

Versions.

• OS: Python3.8-buster image
• Python: 3.8.11
• Alembic: 1.7.2
• SQLAlchemy: 1.4.23
• Database: MSSQL

Additional context
This is not a problem in alembic 1.6.5. I love the project!

[zzzeek]

well here's the problem. we use exec() around your registered function. in 1.6.5, this is the text we exec'ed:

def replace_view(self, target, replaces=None, replace_with=None):
    None
    return op_cls.replace(self, target, replaces=replaces, replace_with=replace_with)

in 1.7.x, it's this:

def replace_view(self, target: fun_stuff.alembic_utils.ReplaceableObject, replaces: Optional[str]=None, replace_with: Optional[str]=None) -> None:
    None
    return op_cls.replace(self, target, replaces=replaces, replace_with=replace_with)

hence the problem.

I'd have to trace out how we made it add the annotations like that, which I believe we need for internal ops, for now it seems best that external ops would not render annotations or we'd have some option for imports. @CaselIT any ideas, but ill keep looking in a bit

[CaselIT]

A simple one could be to just render all annotations as strings

[zzzeek]

ooh we can do that too. i was just going to have it not render the annotations but yes

[zzzeek]

we need to render the annotations for our op.pyi thing right?

[sqla-tester]

Mike Bayer has proposed a fix for this issue in the master branch:

render 3rd party module annotations as forward references https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/3082

[CaselIT]

we need to render the annotations for our op.pyi thing right?

I think they could work the same, but I think having a flag there to enable/disable it would not be too hard

[zzzeek]

this should work well

[daniel-butler]

Wow you all are fast. Thank you!

[stephenknoth]

@zzzeek - this still happens when registering an operation that has a parameter with a type annotation of something like Dict[str, str]