Add type checking to alembic/

Author: Iain-S

.pre-commit-config.yaml

@@ -38,7 +38,6 @@ repos:
         entry: poetry run mypy --install-types --non-interactive
         language: system
         types: ['python']
-        exclude: ^alembic/
       - id: pylint
         name: Pylint
         entry: poetry run pylint --rcfile=.python-lint

alembic/env.py

@@ -1,6 +1,8 @@
 from logging.config import fileConfig
+from typing import Any, Literal, Optional
 
 from sqlalchemy import engine_from_config, pool
+from sqlalchemy.sql.schema import SchemaItem
 
 from alembic import context
 from rctab.crud.models import metadata
@@ -13,6 +15,7 @@
 config.set_main_option("sqlalchemy.url", str(settings.postgres_dsn))
 # Interpret the config file for Python logging.
 # This line sets up loggers basically.
+assert config.config_file_name is not None
 fileConfig(config.config_file_name)
 
 # add your model's MetaData object here
@@ -26,16 +29,29 @@
 # my_important_option = config.get_main_option("my_important_option")
 # ... etc.
 
+SchemaType = Literal[
+    "schema", "table", "column", "index", "unique_constraint", "foreign_key_constraint"
+]
 
-def include_object(object, name, type_, reflected, compare_to):
+
+def include_object(
+    object: SchemaItem,
+    name: Optional[str],
+    type_: SchemaType,
+    reflected: bool,
+    compare_to: Optional[SchemaItem],
+) -> bool:
     """
     Exclude views from Alembic's consideration.
     """
 
-    return not object.info.get("is_view", False)
+    if object.info:
+        if object.info.get("is_view", False):
+            return False
+    return True
 
 
-def run_migrations_offline():
+def run_migrations_offline() -> None:
     """Run migrations in 'offline' mode.
 
     This configures the context with just a URL
@@ -61,7 +77,7 @@ def run_migrations_offline():
         context.run_migrations()
 
 
-def run_migrations_online():
+def run_migrations_online() -> None:
     """Run migrations in 'online' mode.
 
     In this scenario we need to create an Engine

alembic/versions/b65796c99771_squash.py

@@ -67,7 +67,7 @@ class FinanceHistorySql:
     DROP_FUNCTION = "DROP FUNCTION {schema}.finance_changed();"
 
 
-def upgrade():
+def upgrade() -> None:
     op.execute(text("create schema accounting"))
     op.create_table(
         "user_cache",
@@ -362,9 +362,7 @@ def upgrade():
         sa.Column("subscription_id", postgresql.UUID(), nullable=False),
         sa.Column("display_name", sa.String(), nullable=False),
         sa.Column("state", sa.String(), nullable=False),
-        sa.Column(
-            "role_assignments", postgresql.JSONB(astext_type=sa.Text()), nullable=True
-        ),
+        sa.Column("role_assignments", postgresql.JSONB(), nullable=True),
         sa.Column(
             "time_created",
             sa.DateTime(timezone=True),
@@ -385,7 +383,7 @@ def upgrade():
         sa.Column("id", sa.String(), nullable=False),
         sa.Column("name", sa.String(), nullable=True),
         sa.Column("type", sa.String(), nullable=True),
-        sa.Column("tags", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
+        sa.Column("tags", postgresql.JSONB(), nullable=True),
         sa.Column("billing_account_id", sa.String(), nullable=True),
         sa.Column("billing_account_name", sa.String(), nullable=True),
         sa.Column("billing_period_start_date", sa.Date(), nullable=True),
@@ -488,7 +486,7 @@ def upgrade():
     )
 
 
-def downgrade():
+def downgrade() -> None:
     op.drop_table("cost_recovery", schema="accounting")
     op.execute(
         "DROP MATERIALIZED VIEW {schema}.usage_view;".format(schema="accounting")