Skip to content

Commit

Permalink
Add Alembic helpers for Add/Drop spatial columns with SQLite (#362)
Browse files Browse the repository at this point in the history
  • Loading branch information
adrien-berchet authored Mar 4, 2022
1 parent c151f95 commit 2bf9e1e
Show file tree
Hide file tree
Showing 3 changed files with 211 additions and 0 deletions.
74 changes: 74 additions & 0 deletions doc/alembic.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
# ...
129 changes: 129 additions & 0 deletions geoalchemy2/alembic_helpers.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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(")
8 changes: 8 additions & 0 deletions geoalchemy2/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down

0 comments on commit 2bf9e1e

Please sign in to comment.