Factor out context object

migrations/dimensions-config/1fae088c80b6.py

@@ -7,9 +7,11 @@
 """
 
 import logging
+from dataclasses import dataclass
 
+import alembic
 import sqlalchemy as sa
-from alembic import context, op
+from alembic import op
 
 # revision identifiers, used by Alembic.
 revision = "1fae088c80b6"
@@ -32,10 +34,7 @@ def downgrade() -> None:
 
 
 def _migrate_groups() -> None:
-    mig_context = context.get_context()
-    schema = mig_context.version_table_schema
-    bind = mig_context.bind
-    assert bind is not None
+    ctx = _Context()
 
     table = op.create_table(
         "group",
@@ -46,12 +45,11 @@ def _migrate_groups() -> None:
             refcolumns=["instrument.name"],
             name="fkey_group_instrument_name_instrument",
         ),
-        schema=schema,
+        schema=ctx.schema,
     )
 
     # Populate the new table based on the data in the exposure table.
-    metadata = sa.schema.MetaData(schema=schema)
-    exposure_table = sa.schema.Table("exposure", metadata, autoload_with=bind, schema=schema)
+    exposure_table = ctx.get_table("exposure")
     select = sa.select(
         exposure_table.columns["instrument"],
         exposure_table.columns["group_name"],
@@ -67,7 +65,7 @@ def _migrate_groups() -> None:
     )
 
     # Update the exposure table to reference the group table.
-    with op.batch_alter_table("exposure", schema=schema) as batch_op:
+    with op.batch_alter_table("exposure", schema=ctx.schema) as batch_op:
         batch_op.alter_column("group_name", new_column_name="group")
         batch_op.drop_column("group_id")
 
@@ -77,11 +75,23 @@ def _migrate_groups() -> None:
     # foreign key only works if you specify the original column name instead of
     # the final one.  This seems fragile (and is likely incompatible with
     # Postgres, which ignores the batching). So do it in a separate batch.
-    with op.batch_alter_table("exposure", schema=schema) as batch_op:
+    with op.batch_alter_table("exposure", schema=ctx.schema) as batch_op:
         batch_op.create_foreign_key(
             constraint_name="fkey_exposure_group_instrument_name_instrument_group",
             referent_table="group",
             local_cols=["instrument", "group"],
             remote_cols=["instrument", "name"],
-            referent_schema=schema,
+            referent_schema=ctx.schema,
         )
+
+
+class _Context:
+    def __init__(self) -> None:
+        self.mig_context = alembic.context.get_context()
+        self.schema = self.mig_context.version_table_schema
+        self.bind = self.mig_context.bind
+        assert self.bind is not None
+        self.metadata = sa.schema.MetaData(schema=self.schema)
+
+    def get_table(self, table_name: str) -> sa.Table:
+        return sa.schema.Table(table_name, self.metadata, autoload_with=self.bind, schema=self.schema)