diff --git a/alembic/ddl/impl.py b/alembic/ddl/impl.py index 46180246..cda0a4f6 100644 --- a/alembic/ddl/impl.py +++ b/alembic/ddl/impl.py @@ -48,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 @@ -141,7 +139,9 @@ 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, version_table_schema, version_table_pk): + def version_table_impl( + self, version_table, version_table_schema, version_table_pk + ) -> Table: """create the Table object for the version_table. Provided as part of impl so that third party dialects can override diff --git a/alembic/runtime/migration.py b/alembic/runtime/migration.py index f652808d..757bd75f 100644 --- a/alembic/runtime/migration.py +++ b/alembic/runtime/migration.py @@ -24,7 +24,6 @@ from sqlalchemy import Column from sqlalchemy import literal_column -from sqlalchemy import select from sqlalchemy.engine import Engine from sqlalchemy.engine import url as sqla_url from sqlalchemy.engine.strategies import MockEngineStrategy @@ -531,7 +530,9 @@ def get_current_heads(self) -> Tuple[str, ...]: return () assert self.connection is not None return tuple( - row[0] for row in self.connection.execute(select(self._version.c.version_num)) + row[0] for row in self.connection.execute( + self._version.select().with_only_columns( + self._version.c.version_num)) ) def _ensure_version_table(self, purge: bool = False) -> None: diff --git a/tests/test_version_table.py b/tests/test_version_table.py index e2fad65d..7e4f562e 100644 --- a/tests/test_version_table.py +++ b/tests/test_version_table.py @@ -23,8 +23,6 @@ ) - - def _up(from_, to_, branch_presence_changed=False): return migration.StampStep(from_, to_, True, branch_presence_changed) @@ -385,7 +383,9 @@ class CustomVersionTableTest(TestMigrationContext): class MyDialectImpl(impl.DefaultImpl): - def version_table_impl(self, version_table, version_table_schema, version_table_pk): + def version_table_impl( + self, version_table, version_table_schema, version_table_pk + ): vt = Table( version_table, MetaData(), @@ -401,9 +401,9 @@ def version_table_impl(self, version_table, version_table_schema, version_table_ ) return vt - def setUp(self): - # nasty hack to get the sqlite dialect to use our custom dialect implementation + # nasty hack to get the sqlite dialect + # to use our custom dialect implementation impl._impls["sqlite_bak"] = impl._impls["sqlite"] impl._impls["sqlite"] = self.MyDialectImpl super().setUp()