Add Alembic helpers for Add/Drop spatial columns with SQLite (#362)

doc/alembic.rst

@@ -204,3 +204,77 @@ in ``my_package.custom_types``, you just have to edit the ``env.py`` file like t
         # ...
 
 Then the proper imports will be automatically added in the migration scripts.
+
+
+Add / Drop columns
+------------------
+
+Some dialects (like SQLite) require some specific management to alter columns of a table. In this
+case, other dedicated helpers are provided to handle this. For example, if one wants to add and drop
+columns in a SQLite database, the ``env.py`` file should look like the following:
+
+.. code-block:: python
+
+from alembic.autogenerate import rewriter
+
+writer = rewriter.Rewriter()
+
+
+@writer.rewrites(ops.AddColumnOp)
+def add_geo_column(context, revision, op):
+"""This function replaces the default AddColumnOp by a geospatial-specific one."""
+    col_type = op.column.type
+    if isinstance(col_type, TypeDecorator):
+        dialect = context.bind().dialect
+        col_type = col_type.load_dialect_impl(dialect)
+    if isinstance(col_type, (Geometry, Geography, Raster)):
+        new_op = AddGeospatialColumn(op.table_name, op.column, op.schema)
+    else:
+        new_op = op
+    return new_op
+
+
+@writer.rewrites(ops.DropColumnOp)
+def drop_geo_column(context, revision, op):
+"""This function replaces the default DropColumnOp by a geospatial-specific one."""
+    col_type = op.to_column().type
+    if isinstance(col_type, TypeDecorator):
+        dialect = context.bind.dialect
+        col_type = col_type.load_dialect_impl(dialect)
+    if isinstance(col_type, (Geometry, Geography, Raster)):
+        new_op = DropGeospatialColumn(op.table_name, op.column_name, op.schema)
+    else:
+        new_op = op
+    return new_op
+
+
+def load_spatialite(dbapi_conn, connection_record):
+    """Load SpatiaLite extension in SQLite DB."""
+    dbapi_conn.enable_load_extension(True)
+    dbapi_conn.load_extension(os.environ['SPATIALITE_LIBRARY_PATH'])
+    dbapi_conn.enable_load_extension(False)
+    dbapi_conn.execute('SELECT InitSpatialMetaData()')
+
+
+def run_migrations_offline():
+    # ...
+    context.configure(
+        # ...
+        process_revision_directives=writer,
+    )
+    # ...
+
+
+def run_migrations_online():
+    # ...
+    if connectable.dialect.name == "sqlite":
+        # Load the SpatiaLite extension when the engine connects to the DB
+        listen(connectable, 'connect', load_spatialite)
+
+    with connectable.connect() as connection:
+        # ...
+        context.configure(
+            # ...
+            process_revision_directives=writer,
+        )
+        # ...

geoalchemy2/alembic_helpers.py

@@ -1,8 +1,19 @@
 """Some helpers to use with Alembic migration tool."""
+from alembic.autogenerate import renderers
+from alembic.autogenerate.render import _add_column
+from alembic.autogenerate.render import _drop_column
+from alembic.operations import Operations
+from alembic.operations import ops
+from packaging.version import parse as parse_version
+from sqlalchemy import text
+from sqlalchemy.types import TypeDecorator
+
+from geoalchemy2 import Column
 from geoalchemy2 import Geography
 from geoalchemy2 import Geometry
 from geoalchemy2 import Raster
 from geoalchemy2 import _check_spatial_type
+from geoalchemy2 import func
 
 
 def render_item(obj_type, obj, autogen_context):
@@ -33,3 +44,121 @@ def include_object(obj, name, obj_type, reflected, compare_to):
     if (obj_type == "table" and name == "spatial_ref_sys"):
         return False
     return True
+
+
+@Operations.register_operation("add_geospatial_column")
+class AddGeospatialColumn(ops.AddColumnOp):
+    """
+    Add a Geospatial Column in an Alembic migration context. This methodology originates from:
+    https://alembic.sqlalchemy.org/en/latest/api/operations.html#operation-plugins
+    """
+
+    @classmethod
+    def add_geospatial_column(cls, operations, table_name, column, schema=None):
+        """Handle the different situations arising from adding geospatial column to a DB."""
+        op = cls(table_name, column, schema=schema)
+        return operations.invoke(op)
+
+    def reverse(self):
+        """Used to autogenerate the downgrade function."""
+        return DropGeospatialColumn.from_column_and_tablename(
+            self.schema, self.table_name, self.column.name
+        )
+
+
+@Operations.register_operation("drop_geospatial_column")
+class DropGeospatialColumn(ops.DropColumnOp):
+    """Drop a Geospatial Column in an Alembic migration context."""
+
+    @classmethod
+    def drop_geospatial_column(cls, operations, table_name, column_name, schema=None, **kw):
+        """Handle the different situations arising from dropping geospatial column from a DB."""
+
+        op = cls(table_name, column_name, schema=schema, **kw)
+        return operations.invoke(op)
+
+    def reverse(self):
+        """Used to autogenerate the downgrade function."""
+        return AddGeospatialColumn.from_column_and_tablename(
+            self.schema, self.table_name, self.column
+        )
+
+
+@Operations.implementation_for(AddGeospatialColumn)
+def add_geospatial_column(operations, operation):
+    """Handle the actual column addition according to the dialect backend.
+
+    Parameters:
+        operations: Operations object from alembic base, defining high level migration operations
+        operation: AddGeospatialColumn call, with attributes for table_name, column_name,
+            column_type, and optional keywords.
+    """
+
+    table_name = operation.table_name
+    column_name = operation.column.name
+
+    dialect = operations.get_bind().dialect
+
+    if isinstance(operation.column, TypeDecorator):
+        # Will be either geoalchemy2.types.Geography or geoalchemy2.types.Geometry, if using a
+        # custom type
+        geospatial_core_type = operation.column.type.load_dialect_impl(dialect)
+    else:
+        geospatial_core_type = operation.column.type
+
+    if "sqlite" in dialect.name:
+        operations.execute(func.AddGeometryColumn(
+            table_name,
+            column_name,
+            geospatial_core_type.srid,
+            geospatial_core_type.geometry_type
+        ))
+    elif "postgresql" in dialect.name:
+        operations.add_column(
+            table_name,
+            Column(
+                column_name,
+                operation.column
+            )
+        )
+
+
+@Operations.implementation_for(DropGeospatialColumn)
+def drop_geospatial_column(operations, operation):
+    """
+    Handles the actual column removal by checking for the dialect backend and issuing proper
+    commands.
+    """
+
+    table_name = operation.table_name
+    column_name = operation.column_name
+
+    dialect = operations.get_bind().dialect
+
+    if "sqlite" in dialect.name:
+        operations.execute(func.DiscardGeometryColumn(table_name, column_name))
+        # This second drop column call is necessary; SpatiaLite was designed for a SQLite that did
+        # not support dropping columns from tables at all. DiscardGeometryColumn removes associated
+        # metadata and triggers from the DB associated with a geospatial column, without removing
+        # the column itself. The next call actually removes the geospatial column, IF the underlying
+        # SQLite package version >= 3.35
+        conn = operations.get_bind()
+        sqlite_version = conn.execute(text("SELECT sqlite_version();")).scalar()
+        if parse_version(sqlite_version) >= parse_version("3.35"):
+            operations.drop_column(table_name, column_name)
+    elif "postgresql" in dialect.name:
+        operations.drop_column(table_name, column_name)
+
+
+@renderers.dispatch_for(AddGeospatialColumn)
+def render_add_geo_column(autogen_context, op):
+    """Render the add_geospatial_column operation in migration script."""
+    col_render = _add_column(autogen_context, op)
+    return col_render.replace(".add_column(", ".add_geospatial_column(")
+
+
+@renderers.dispatch_for(DropGeospatialColumn)
+def render_drop_geo_column(autogen_context, op):
+    """Render the drop_geospatial_column operation in migration script."""
+    col_render = _drop_column(autogen_context, op)
+    return col_render.replace(".drop_column(", ".drop_geospatial_column(")

geoalchemy2/types.py

@@ -432,6 +432,14 @@ class GeometryDump(CompositeType):
 postgresql_ischema_names['raster'] = Raster
 
 sqlite_ischema_names['GEOMETRY'] = Geometry
+sqlite_ischema_names['POINT'] = Geometry
+sqlite_ischema_names['LINESTRING'] = Geometry
+sqlite_ischema_names['POLYGON'] = Geometry
+sqlite_ischema_names['MULTIPOINT'] = Geometry
+sqlite_ischema_names['MULTILINESTRING'] = Geometry
+sqlite_ischema_names['MULTIPOLYGON'] = Geometry
+sqlite_ischema_names['CURVE'] = Geometry
+sqlite_ischema_names['GEOMETRYCOLLECTION'] = Geometry
 sqlite_ischema_names['RASTER'] = Raster