Fix mypy type errors in migration scripts

Author: nuwang

lib/galaxy/model/migrations/alembic/versions_gxy/724237cc4cf0_migrate_custos_to_psa_tokens.py

@@ -43,13 +43,12 @@ def get_custos_table(connection):
 
     This function can be reused by both migration scripts and tests.
     """
+    from sqlalchemy import MetaData
+
+    metadata = MetaData()
     return Table(
         CUSTOS_TABLE,
-        (
-            op.get_bind().dialect.get_columns(op.get_bind(), CUSTOS_TABLE)
-            if hasattr(op.get_bind().dialect, "get_columns")
-            else []
-        ),
+        metadata,
         Column("id", Integer, primary_key=True),
         Column("user_id", Integer),
         Column("external_user_id", String(255)),
@@ -70,13 +69,12 @@ def get_psa_table(connection):
 
     This function can be reused by both migration scripts and tests.
     """
+    from sqlalchemy import MetaData
+
+    metadata = MetaData()
     return Table(
         PSA_TABLE,
-        (
-            op.get_bind().dialect.get_columns(op.get_bind(), PSA_TABLE)
-            if hasattr(op.get_bind().dialect, "get_columns")
-            else []
-        ),
+        metadata,
         Column("id", Integer, primary_key=True),
         Column("user_id", Integer, ForeignKey("galaxy_user.id")),
         Column("uid", VARCHAR(255)),

test/unit/data/model/test_migrate_custos_to_psa.py

@@ -13,6 +13,7 @@
     timedelta,
 )
 from pathlib import Path
+from typing import Any
 
 import pytest
 from sqlalchemy import (
@@ -33,17 +34,24 @@
 
 import galaxy.model.migrations
 
-migrations_path = Path(galaxy.model.migrations.__file__).parent / "alembic" / "versions_gxy" / "724237cc4cf0_migrate_custos_to_psa_tokens.py"
+migrations_path = (
+    Path(galaxy.model.migrations.__file__).parent
+    / "alembic"
+    / "versions_gxy"
+    / "724237cc4cf0_migrate_custos_to_psa_tokens.py"
+)
 spec = importlib.util.spec_from_file_location("migration_module", migrations_path)
+assert spec is not None, f"Could not load spec from {migrations_path}"
+assert spec.loader is not None, f"Spec has no loader for {migrations_path}"
 migration_module = importlib.util.module_from_spec(spec)
 spec.loader.exec_module(migration_module)
 migrate_custos_tokens_to_psa = migration_module.migrate_custos_tokens_to_psa
 
 # Create test tables
-Base = declarative_base()
+_Base: Any = declarative_base()
 
 
-class CustosAuthnzTokenTest(Base):
+class CustosAuthnzTokenTest(_Base):
     """Test model for custos_authnz_token table."""
 
     __tablename__ = "custos_authnz_token"
@@ -59,7 +67,7 @@ class CustosAuthnzTokenTest(Base):
     refresh_expiration_time = Column(DateTime)
 
 
-class UserAuthnzTokenTest(Base):
+class UserAuthnzTokenTest(_Base):
     """Test model for oidc_user_authnz_tokens table."""
 
     __tablename__ = "oidc_user_authnz_tokens"
@@ -77,7 +85,7 @@ class UserAuthnzTokenTest(Base):
 def db_session():
     """Create an in-memory SQLite database for testing."""
     engine = create_engine("sqlite:///:memory:")
-    Base.metadata.create_all(engine)
+    _Base.metadata.create_all(engine)
     Session = sessionmaker(bind=engine)
     session = Session()
     yield session