fix: Issue with runtime version_table_schema option.

`version_table_schema` appears to be (one of?) the only options that could
meaningfully affect the runtime behavior as-was, because it changes
where an un-configured (through env.py) MigrationContext instance would
point at an incorrect bit of data.

Now we run the process through env.py to ensure that setting it picked
up.

In order to attempt to avoid excess calls through the env.py, the
defaults tests now internally avoid making calls to `self.current`
wherever possible.

examples/test_version_table_schema/alembic.ini

@@ -0,0 +1,36 @@
+[alembic]
+script_location = migrations
+
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S

examples/test_version_table_schema/conftest.py

@@ -0,0 +1,3 @@
+from pytest_mock_resources import create_postgres_fixture, Statements
+
+alembic_engine = create_postgres_fixture(Statements("CREATE SCHEMA version_table_schema"))

examples/test_version_table_schema/migrations/env.py

@@ -0,0 +1,29 @@
+from logging.config import fileConfig
+
+from alembic import context
+from sqlalchemy import engine_from_config, pool
+
+from models import Base
+
+fileConfig(context.config.config_file_name)
+target_metadata = Base.metadata
+
+
+connectable = context.config.attributes.get("connection", None)
+
+if connectable is None:
+    connectable = engine_from_config(
+        context.config.get_section(context.config.config_ini_section),
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+    )
+
+with connectable.connect() as connection:
+    context.configure(
+        connection=connection,
+        target_metadata=target_metadata,
+        version_table_schema="version_table_schema",
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()

examples/test_version_table_schema/migrations/script.py.mako

@@ -0,0 +1,24 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision = ${repr(up_revision)}
+down_revision = ${repr(down_revision)}
+branch_labels = ${repr(branch_labels)}
+depends_on = ${repr(depends_on)}
+
+
+def upgrade():
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade():
+    ${downgrades if downgrades else "pass"}

examples/test_version_table_schema/migrations/versions/aaaaaaaaaaaa_create_foo.py

@@ -0,0 +1,25 @@
+import sqlalchemy as sa
+from alembic import op
+
+revision = "aaaaaaaaaaaa"
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_table(
+        "foo",
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("(CURRENT_TIMESTAMP)"),
+            nullable=False,
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+
+
+def downgrade():
+    op.drop_table("foo")

examples/test_version_table_schema/models.py

@@ -0,0 +1,17 @@
+import sqlalchemy
+from sqlalchemy import Column, types
+from sqlalchemy.ext.declarative import declarative_base
+
+Base = declarative_base()
+
+
+class CreatedAt(Base):
+    __tablename__ = "foo"
+
+    id = Column(types.Integer(), autoincrement=True, primary_key=True)
+
+    created_at = sqlalchemy.Column(
+        sqlalchemy.types.DateTime(timezone=True),
+        server_default=sqlalchemy.text("CURRENT_TIMESTAMP"),
+        nullable=False,
+    )

examples/test_version_table_schema/setup.cfg

@@ -0,0 +1,3 @@
+[tool:pytest]
+pytest_alembic_exclude = single_head_revision
+pytest_alembic_include_experimental = downgrade_leaves_no_trace,all_models_register_on_metadata

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytest-alembic"
-version = "0.10.5"
+version = "0.10.6"
 description = "A pytest plugin for verifying alembic migrations."
 authors = [
     "Dan Cardin <[email protected]>",

src/pytest_alembic/executor.py

@@ -36,6 +36,10 @@ def configure(self, **kwargs):
         for key, value in kwargs.items():
             self.alembic_config.attributes[key] = value
 
+    def execute_fn(self, fn):
+        with EnvironmentContext(self.alembic_config, self.script, fn=fn):
+            self.script.run_env()
+
     def run_command(self, command, *args, **kwargs):
         self.stream_position = self.stdout.tell()
 

src/pytest_alembic/runner.py

@@ -70,15 +70,17 @@ def heads(self) -> List[str]:
     @property
     def current(self) -> str:
         """Get the list of revision heads."""
+        current = "base"
 
-        def get_current(connection):
-            context = alembic.migration.MigrationContext.configure(connection)
-            heads = context.get_current_heads()
-            if heads:
-                return heads[0]
-            return None
+        def get_current(rev, _):
+            nonlocal current
+            if rev:
+                current = rev[0]
+
+            return []
+
+        self.command_executor.execute_fn(get_current)
 
-        current = self.connection_executor.run_task(get_current)
         if current:
             return current
         return "base"
@@ -138,9 +140,11 @@ def raw_command(self, *args, **kwargs):
         """Execute a raw alembic command."""
         return self.command_executor.run_command(*args, **kwargs)
 
-    def managed_upgrade(self, dest_revision):
+    def managed_upgrade(self, dest_revision, *, current=None, return_current=True):
         """Perform an upgrade one migration at a time, inserting static data at the given points."""
-        current = self.current
+        if current is None:
+            current = self.current
+
         for current_revision, next_revision in self.history.revision_window(current, dest_revision):
             before_upgrade_data = self.revision_data.get_before(next_revision)
             self.insert_into(data=before_upgrade_data, revision=current_revision, table=None)
@@ -153,12 +157,16 @@ def managed_upgrade(self, dest_revision):
             at_upgrade_data = self.revision_data.get_at(next_revision)
             self.insert_into(data=at_upgrade_data, revision=next_revision, table=None)
 
-        current = self.current
-        return current
+        if return_current:
+            current = self.current
+            return current
+        return None
 
-    def managed_downgrade(self, dest_revision):
+    def managed_downgrade(self, dest_revision, *, current=None, return_current=True):
         """Perform an downgrade, one migration at a time."""
-        current = self.current
+        if current is None:
+            current = self.current
+
         for next_revision, current_revision in reversed(
             self.history.revision_window(dest_revision, current)
         ):
@@ -173,23 +181,25 @@ def managed_downgrade(self, dest_revision):
                     else:
                         raise
 
-        current = self.current
-        return current
+        if return_current:
+            current = self.current
+            return current
+        return None
 
     def migrate_up_before(self, revision):
         """Migrate up to, but not including the given `revision`."""
         preceeding_revision = self.history.previous_revision(revision)
         return self.managed_upgrade(preceeding_revision)
 
-    def migrate_up_to(self, revision):
+    def migrate_up_to(self, revision, *, return_current: bool = True):
         """Migrate up to, and including the given `revision`."""
-        return self.managed_upgrade(revision)
+        return self.managed_upgrade(revision, return_current=return_current)
 
     def migrate_up_one(self):
         """Migrate up by exactly one revision."""
         current = self.current
         next_revision = self.history.next_revision(current)
-        new_revision = self.managed_upgrade(next_revision)
+        new_revision = self.managed_upgrade(next_revision, current=current)
         if current == new_revision:
             return None
         return new_revision
@@ -199,16 +209,17 @@ def migrate_down_before(self, revision):
         next_revision = self.history.next_revision(revision)
         return self.migrate_down_to(next_revision)
 
-    def migrate_down_to(self, revision):
+    def migrate_down_to(self, revision, *, return_current: bool = True):
         """Migrate down to, and including the given `revision`."""
         self.history.validate_revision(revision)
-        self.managed_downgrade(revision)
+        self.managed_downgrade(revision, return_current=return_current)
         return revision
 
     def migrate_down_one(self):
         """Migrate down by exactly one revision."""
-        previous_revision = self.history.previous_revision(self.current)
-        self.managed_downgrade(previous_revision)
+        current = self.current
+        previous_revision = self.history.previous_revision(current)
+        self.managed_downgrade(previous_revision, current=current)
         return previous_revision
 
     def roundtrip_next_revision(self):

src/pytest_alembic/tests/default.py

@@ -42,7 +42,7 @@ def test_single_head_revision(alembic_runner):
 def test_upgrade(alembic_runner):
     """Assert that the revision history can be run through from base to head."""
     try:
-        alembic_runner.migrate_up_to("heads")
+        alembic_runner.migrate_up_to("heads", return_current=False)
     except RuntimeError as e:
         message = (
             "Failed to upgrade to the head revision. This means the historical chain from an "
@@ -110,7 +110,7 @@ def test_up_down_consistency(alembic_runner):
     """
     for revision in alembic_runner.history.revisions:
         try:
-            alembic_runner.migrate_up_to(revision)
+            alembic_runner.migrate_up_to(revision, return_current=False)
         except Exception as e:
             message = "Failed to upgrade through each revision individually."
             raise AlembicTestFailure(
@@ -128,7 +128,7 @@ def test_up_down_consistency(alembic_runner):
             break
 
         try:
-            alembic_runner.migrate_down_to(revision)
+            alembic_runner.migrate_down_to(revision, return_current=False)
         except NotImplementedError:
             # In the event of a `NotImplementedError`, we should have the same semantics,
             # as-if `minimum_downgrade_revision` was specified, but we'll emit a warning
@@ -148,7 +148,7 @@ def test_up_down_consistency(alembic_runner):
 
     for revision in reversed(down_revisions):
         try:
-            alembic_runner.migrate_up_to(revision)
+            alembic_runner.migrate_up_to(revision, return_current=False)
         except Exception as e:
             message = (
                 "Failed to upgrade through each revision individually after performing a "

tests/test_runner.py

@@ -270,3 +270,8 @@ def test_skip_revision(pytester):
 def test_pytest_alembic_tests_path(pytester):
     """Assert the pytest_alembic_tests_path can be overridden."""
     run_pytest(pytester, passed=4, args=["-vv", "--test-alembic", "tests_"])
+
+
+def test_version_table_schema(pytester):
+    """Assert the setting the version_table_schema option functions correctly."""
+    run_pytest(pytester, passed=5)