Add ability to configure alembic_version table in DialectImpl

Added a new hook to the :class:`.DefaultImpl`
:meth:`.DefaultImpl.version_table_impl`.  This allows third party dialects
to define the exact structure of the alembic_version table, to include use
cases where the table requires special directives and/or additional columns
so that it may function correctly on a particular backend.  This is not
intended as a user-expansion hook, only a dialect implementation hook to
produce a working alembic_version table. Pull request courtesy Maciek
Bryński.

This will be 1.14 so this also version bumps

Fixes: #1560
Closes: #1563
Pull-request: #1563
Pull-request-sha: e70fdc8

Change-Id: I5e565dff60a979526608d2a1c0c620fbca269a3f

Author: maver1ck

alembic/ddl/impl.py

@@ -21,7 +21,12 @@
 from typing import Union
 
 from sqlalchemy import cast
+from sqlalchemy import Column
+from sqlalchemy import MetaData
+from sqlalchemy import PrimaryKeyConstraint
 from sqlalchemy import schema
+from sqlalchemy import String
+from sqlalchemy import Table
 from sqlalchemy import text
 
 from . import _autogen
@@ -43,11 +48,9 @@
     from sqlalchemy.sql import Executable
     from sqlalchemy.sql.elements import ColumnElement
     from sqlalchemy.sql.elements import quoted_name
-    from sqlalchemy.sql.schema import Column
     from sqlalchemy.sql.schema import Constraint
     from sqlalchemy.sql.schema import ForeignKeyConstraint
     from sqlalchemy.sql.schema import Index
-    from sqlalchemy.sql.schema import Table
     from sqlalchemy.sql.schema import UniqueConstraint
     from sqlalchemy.sql.selectable import TableClause
     from sqlalchemy.sql.type_api import TypeEngine
@@ -136,6 +139,40 @@ def static_output(self, text: str) -> None:
         self.output_buffer.write(text + "\n\n")
         self.output_buffer.flush()
 
+    def version_table_impl(
+        self,
+        *,
+        version_table: str,
+        version_table_schema: Optional[str],
+        version_table_pk: bool,
+        **kw: Any,
+    ) -> Table:
+        """Generate a :class:`.Table` object which will be used as the
+        structure for the Alembic version table.
+
+        Third party dialects may override this hook to provide an alternate
+        structure for this :class:`.Table`; requirements are only that it
+        be named based on the ``version_table`` parameter and contains
+        at least a single string-holding column named ``version_num``.
+
+        .. versionadded:: 1.14
+
+        """
+        vt = Table(
+            version_table,
+            MetaData(),
+            Column("version_num", String(32), nullable=False),
+            schema=version_table_schema,
+        )
+        if version_table_pk:
+            vt.append_constraint(
+                PrimaryKeyConstraint(
+                    "version_num", name=f"{version_table}_pkc"
+                )
+            )
+
+        return vt
+
     def requires_recreate_in_batch(
         self, batch_op: BatchOperationsImpl
     ) -> bool:

alembic/runtime/migration.py

@@ -24,10 +24,6 @@
 
 from sqlalchemy import Column
 from sqlalchemy import literal_column
-from sqlalchemy import MetaData
-from sqlalchemy import PrimaryKeyConstraint
-from sqlalchemy import String
-from sqlalchemy import Table
 from sqlalchemy.engine import Engine
 from sqlalchemy.engine import url as sqla_url
 from sqlalchemy.engine.strategies import MockEngineStrategy
@@ -36,6 +32,7 @@
 from .. import util
 from ..util import sqla_compat
 from ..util.compat import EncodedIO
+from ..util.sqla_compat import _select
 
 if TYPE_CHECKING:
     from sqlalchemy.engine import Dialect
@@ -190,18 +187,6 @@ def __init__(
         self.version_table_schema = version_table_schema = opts.get(
             "version_table_schema", None
         )
-        self._version = Table(
-            version_table,
-            MetaData(),
-            Column("version_num", String(32), nullable=False),
-            schema=version_table_schema,
-        )
-        if opts.get("version_table_pk", True):
-            self._version.append_constraint(
-                PrimaryKeyConstraint(
-                    "version_num", name="%s_pkc" % version_table
-                )
-            )
 
         self._start_from_rev: Optional[str] = opts.get("starting_rev")
         self.impl = ddl.DefaultImpl.get_by_dialect(dialect)(
@@ -212,6 +197,13 @@ def __init__(
             self.output_buffer,
             opts,
         )
+
+        self._version = self.impl.version_table_impl(
+            version_table=version_table,
+            version_table_schema=version_table_schema,
+            version_table_pk=opts.get("version_table_pk", True),
+        )
+
         log.info("Context impl %s.", self.impl.__class__.__name__)
         if self.as_sql:
             log.info("Generating static SQL")
@@ -540,7 +532,10 @@ def get_current_heads(self) -> Tuple[str, ...]:
                 return ()
         assert self.connection is not None
         return tuple(
-            row[0] for row in self.connection.execute(self._version.select())
+            row[0]
+            for row in self.connection.execute(
+                _select(self._version.c.version_num)
+            )
         )
 
     def _ensure_version_table(self, purge: bool = False) -> None:

docs/build/changelog.rst

@@ -4,7 +4,7 @@ Changelog
 ==========
 
 .. changelog::
-    :version: 1.13.4
+    :version: 1.14.0
     :include_notes_from: unreleased
 
 .. changelog::

docs/build/unreleased/1560.rst

@@ -0,0 +1,12 @@
+.. change::
+    :tags: usecase, runtime
+    :tickets: 1560
+
+    Added a new hook to the :class:`.DefaultImpl`
+    :meth:`.DefaultImpl.version_table_impl`.  This allows third party dialects
+    to define the exact structure of the alembic_version table, to include use
+    cases where the table requires special directives and/or additional columns
+    so that it may function correctly on a particular backend.  This is not
+    intended as a user-expansion hook, only a dialect implementation hook to
+    produce a working alembic_version table. Pull request courtesy Maciek
+    Bryński.

tests/test_version_table.py

@@ -1,10 +1,15 @@
 from sqlalchemy import Column
 from sqlalchemy import inspect
+from sqlalchemy import Integer
 from sqlalchemy import MetaData
+from sqlalchemy import PrimaryKeyConstraint
 from sqlalchemy import String
 from sqlalchemy import Table
+from sqlalchemy.dialects import registry
+from sqlalchemy.engine import default
 
 from alembic import migration
+from alembic.ddl import impl
 from alembic.testing import assert_raises
 from alembic.testing import assert_raises_message
 from alembic.testing import config
@@ -373,3 +378,44 @@ def test_delete_multi_match_no_sane_rowcount(self):
                 self.connection.dialect, "supports_sane_rowcount", False
             ):
                 self.updater.update_to_step(_down("a", None, True))
+
+
+registry.register("custom_version", __name__, "CustomVersionDialect")
+
+
+class CustomVersionDialect(default.DefaultDialect):
+    name = "custom_version"
+
+
+class CustomVersionTableImpl(impl.DefaultImpl):
+    __dialect__ = "custom_version"
+
+    def version_table_impl(
+        self,
+        *,
+        version_table,
+        version_table_schema,
+        version_table_pk,
+        **kw,
+    ):
+        vt = Table(
+            version_table,
+            MetaData(),
+            Column("id", Integer, autoincrement=True),
+            Column("version_num", String(32), nullable=False),
+            schema=version_table_schema,
+        )
+        if version_table_pk:
+            vt.append_constraint(
+                PrimaryKeyConstraint("id", name=f"{version_table}_pkc")
+            )
+        return vt
+
+
+class CustomVersionTableTest(TestMigrationContext):
+
+    def test_custom_version_table(self):
+        context = migration.MigrationContext.configure(
+            dialect_name="custom_version",
+        )
+        eq_(len(context._version.columns), 2)