Update exposure table for group changes

migrations/dimensions-config/1fae088c80b6.py

@@ -23,15 +23,15 @@
 
 def upgrade() -> None:
     """Perform schema upgrade."""
-    _create_group_table()
+    _migrate_groups()
 
 
 def downgrade() -> None:
     """Perform schema downgrade."""
     raise NotImplementedError()
 
 
-def _create_group_table() -> None:
+def _migrate_groups() -> None:
     mig_context = context.get_context()
     schema = mig_context.version_table_schema
     bind = mig_context.bind
@@ -56,7 +56,6 @@ def _create_group_table() -> None:
         exposure_table.columns["instrument"],
         exposure_table.columns["group_name"],
     ).distinct()
-
     op.execute(
         table.insert().from_select(
             [
@@ -66,3 +65,23 @@ def _create_group_table() -> None:
             select,
         )
     )
+
+    # Update the exposure table to reference the group table.
+    with op.batch_alter_table("exposure", schema=schema) as batch_op:
+        batch_op.alter_column("group_name", new_column_name="group")
+        batch_op.drop_column("group_id")
+
+    # In theory we should do this create_foreign_key as part of the batch
+    # above.  However, there is some undocumented weirdness with the column
+    # rename from "group_name" to "group".  When done in the batch above, this
+    # 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:
+        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,
+        )