Fix geometry types in migration 223 and 228 (#169)

CHANGES.rst

@@ -5,6 +5,7 @@ Changelog of threedi-schema
 0.229.1 (unreleased)
 --------------------
 
+- Fix setting of geometry columns for revision 223 and 228
 - Fix incorrect creation of geometry for dry weather flow and surface during migration
 
 

threedi_schema/migrations/versions/0223_upgrade_db_inflow.py

@@ -6,6 +6,7 @@
 
 """
 import json
+import uuid
 import warnings
 from pathlib import Path
 from typing import Dict, List, Tuple
@@ -483,6 +484,29 @@ def populate_dry_weather_flow_distribution():
     op.execute(sa.text(sql_query))
 
 
+def make_geom_col_notnull(table_name):
+    # Make control_measure_map.geom not nullable by creating a new table with
+    # not nullable geometry column, copying the data from control_measure_map
+    # to the new table, dropping the original and renaming the new one to control_measure_map
+    # For some reason, changing this via batch_op.alter_column does not seem to work
+
+    # Retrieve column names and types from table
+    # Note that it is expected that the geometry column is the last column!
+    connection = op.get_bind()
+    columns = connection.execute(sa.text(f"PRAGMA table_info('{table_name}')")).fetchall()
+    col_names = [col[1] for col in columns]
+    col_types = [col[2] for col in columns]
+    cols = (['id INTEGER PRIMARY KEY'] +
+            [f'{cname} {ctype}' for cname, ctype in zip(col_names[:-1], col_types[:-1]) if cname != 'id'] +
+            [f'geom {columns[-1][2]} NOT NULL'])
+    # Create new table, insert data, drop original and rename to table_name
+    temp_name = f'_temp_223_{uuid.uuid4().hex}'
+    op.execute(sa.text(f"CREATE TABLE {temp_name} ({','.join(cols)});"))
+    op.execute(sa.text(f"INSERT INTO {temp_name} ({','.join(col_names)}) SELECT {','.join(col_names)} FROM {table_name}"))
+    drop_geo_table(op, table_name)
+    op.execute(sa.text(f"ALTER TABLE {temp_name} RENAME TO {table_name};"))
+
+
 def fix_geometry_columns():
     GEO_COL_INFO = [
         ('dry_weather_flow', 'geom', 'POLYGON'),
@@ -491,8 +515,7 @@ def fix_geometry_columns():
         ('surface_map', 'geom', 'LINESTRING'),
     ]
     for table, column, geotype in GEO_COL_INFO:
-        with op.batch_alter_table(table) as batch_op:
-            batch_op.alter_column(column_name=column, nullable=False)
+        make_geom_col_notnull(table)
         migration_query = f"SELECT RecoverGeometryColumn('{table}', '{column}', {4326}, '{geotype}', 'XY')"
         op.execute(sa.text(migration_query))
 

threedi_schema/migrations/versions/0228_upgrade_db_1D.py

@@ -100,6 +100,17 @@
 
 RETYPE_COLUMNS = {}
 
+GEOM_TYPES = {'channel': 'LINESTRING',
+              'orifice': 'LINESTRING',
+              'weir': 'LINESTRING',
+              'culvert': 'LINESTRING',
+              'pipe': 'LINESTRING',
+              'pump': 'POINT',
+              'pump_map': 'LINESTRING',
+              'connection_node': 'POINT',
+              'cross_section_location': 'POINT',
+              'windshielding_1d': 'POINT',
+              }
 
 def add_columns_to_tables(table_columns: List[Tuple[str, Column]]):
     # no checks for existence are done, this will fail if any column already exists
@@ -135,14 +146,6 @@ def remove_tables(tables: List[str]):
     for table in tables:
         drop_geo_table(op, table)
 
-
-def get_geom_type(table_name, geo_col_name):
-    connection = op.get_bind()
-    columns = connection.execute(sa.text(f"PRAGMA table_info('{table_name}')")).fetchall()
-    for col in columns:
-        if col[1] == geo_col_name:
-            return col[2]
-
 def modify_table(old_table_name, new_table_name):
     # Create a new table named `new_table_name` by copying the
     # data from `old_table_name`.
@@ -157,7 +160,7 @@ def modify_table(old_table_name, new_table_name):
     col_names = [col[1] for col in columns]
     col_types = [col[2] for col in columns]
     # get type of the geometry column
-    geom_type = get_geom_type(old_table_name, 'the_geom')
+    geom_type = GEOM_TYPES[new_table_name]
     # create list of new columns and types for creating the new table
     # create list of old columns to copy to new table
     skip_cols = ['id', 'the_geom']
@@ -195,7 +198,7 @@ def fix_geometry_columns():
     tables = ['channel', 'connection_node', 'cross_section_location', 'culvert',
               'orifice', 'pipe', 'pump', 'pump_map', 'weir', 'windshielding_1d']
     for table in tables:
-        geom_type = get_geom_type(table, geo_col_name='geom')
+        geom_type = GEOM_TYPES[table]
         op.execute(sa.text(f"SELECT RecoverGeometryColumn('{table}', "
                            f"'geom', {4326}, '{geom_type}', 'XY')"))
         op.execute(sa.text(f"SELECT CreateSpatialIndex('{table}', 'geom')"))