Structure control migration fix ups (#110)

* Remove measure_variable from memory_control and table_control
* Rename control_measure_map to measure_map and control_measure_location to measure_location

CHANGES.rst

@@ -5,7 +5,8 @@ Changelog of threedi-schema
 0.226.7 (unreleased)
 --------------------
 
-- Nothing changed yet.
+- Remove measure_variable column from tables memory_control and table_control
+- Rename control_measure_map to measure_map and control_measure_location to measure_location
 
 
 0.226.6 (2024-10-03)

threedi_schema/domain/models.py

@@ -36,7 +36,7 @@ class BoundaryConditions2D(Base):
 
 
 class ControlMeasureLocation(Base):
-    __tablename__ = "control_measure_location"
+    __tablename__ = "measure_location"
     id = Column(Integer, primary_key=True)
     connection_node_id = Column(Integer)
     measure_variable = Column(VarcharEnum(constants.MeasureVariables))
@@ -47,9 +47,9 @@ class ControlMeasureLocation(Base):
 
 
 class ControlMeasureMap(Base):
-    __tablename__ = "control_measure_map"
+    __tablename__ = "measure_map"
     id = Column(Integer, primary_key=True)
-    control_measure_location_id = Column(Integer)
+    measure_location_id = Column(Integer)
     control_type = Column(VarcharEnum(constants.ControlType), nullable=False)
     control_id = Column(Integer)
     weight = Column(Float)
@@ -62,7 +62,6 @@ class ControlMeasureMap(Base):
 class ControlMemory(Base):
     __tablename__ = "memory_control"
     id = Column(Integer, primary_key=True)
-    measure_variable = Column(VarcharEnum(constants.MeasureVariables))
     upper_threshold = Column(Float)
     lower_threshold = Column(Float)
     action_type = Column(VarcharEnum(constants.ControlTableActionTypes))
@@ -83,7 +82,6 @@ class ControlTable(Base):
     id = Column(Integer, primary_key=True)
     action_table = Column(Text)
     action_type = Column(VarcharEnum(constants.ControlTableActionTypes))
-    measure_variable = Column(VarcharEnum(constants.MeasureVariables))
     measure_operator = Column(VarcharEnum(constants.MeasureOperators))
     target_type = Column(VarcharEnum(constants.StructureControlTypes))
     target_id = Column(Integer, nullable=False)

threedi_schema/migrations/versions/0227_fixups_structure_control.py

@@ -0,0 +1,59 @@
+"""Upgrade settings in schema
+
+Revision ID: 0227
+Revises:
+Create Date: 2024-09-24 15:10
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "0227"
+down_revision = "0226"
+branch_labels = None
+depends_on = None
+
+TABLES = ['memory_control', 'table_control']
+RENAME_TABLES = [('control_measure_location', 'measure_location'),
+                 ('control_measure_map', 'measure_map'), ]
+
+
+def fix_geometries(downgrade: bool=False):
+    op.execute(sa.text("SELECT RecoverGeometryColumn('memory_control', 'geom', 4326, 'POINT', 'XY')"))
+    op.execute(sa.text("SELECT RecoverGeometryColumn('table_control', 'geom', 4326, 'POINT', 'XY')"))
+    if downgrade:
+        op.execute(sa.text("SELECT RecoverGeometryColumn('control_measure_location', 'geom', 4326, 'POINT', 'XY')"))
+        op.execute(sa.text("SELECT RecoverGeometryColumn('control_measure_map', 'geom', 4326, 'LINESTRING', 'XY')"))
+    else:
+        op.execute(sa.text("SELECT RecoverGeometryColumn('measure_location', 'geom', 4326, 'POINT', 'XY')"))
+        op.execute(sa.text("SELECT RecoverGeometryColumn('measure_map', 'geom', 4326, 'LINESTRING', 'XY')"))
+
+
+def upgrade():
+    # remove measure variable from memory_control and table_control
+    for table_name in TABLES:
+        with op.batch_alter_table(table_name) as batch_op:
+            batch_op.drop_column('measure_variable')
+    # rename column
+    with op.batch_alter_table('control_measure_map') as batch_op:
+        batch_op.alter_column('control_measure_location_id', new_column_name='measure_location_id')
+    # rename tables
+    for old_table_name, new_table_name in RENAME_TABLES:
+        op.rename_table(old_table_name, new_table_name)
+    fix_geometries()
+
+
+def downgrade():
+    # undo remove measure variable from memory_control and table_control
+    for table_name in TABLES:
+        with op.batch_alter_table(table_name) as batch_op:
+            batch_op.add_column(sa.Column("measure_variable", sa.Text, server_default="water_level"))
+    # undo rename columns
+    with op.batch_alter_table('measure_map') as batch_op:
+        batch_op.alter_column('measure_location_id', new_column_name='control_measure_location_id')
+    # rename tables
+    for old_table_name, new_table_name in RENAME_TABLES:
+        op.rename_table(new_table_name, old_table_name)
+    fix_geometries(downgrade=True)

threedi_schema/tests/test_migration.py

@@ -150,7 +150,7 @@ class TestMigration224:
                           'v2_control_measure_group', 'v2_control_measure_map',
                           'v2_control_pid', 'v2_control_timed',
                           'v2_control_memory', 'v2_control_table'])
-    added_tables = set(['memory_control', 'table_control', 'control_measure_location', 'control_measure_map'])
+    added_tables = set(['memory_control', 'table_control', 'measure_map', 'measure_location'])
 
     def test_tables(self, schema_ref, schema_upgraded):
         # Test whether the added tables are present