781 lines
27 KiB
Python
781 lines
27 KiB
Python
"""Some helpers to use with Alembic migration tool."""
|
|
|
|
from alembic.autogenerate import renderers
|
|
from alembic.autogenerate import rewriter
|
|
from alembic.autogenerate.render import _add_column
|
|
from alembic.autogenerate.render import _add_index
|
|
from alembic.autogenerate.render import _add_table
|
|
from alembic.autogenerate.render import _drop_column
|
|
from alembic.autogenerate.render import _drop_index
|
|
from alembic.autogenerate.render import _drop_table
|
|
from alembic.ddl.base import RenameTable
|
|
from alembic.ddl.base import format_table_name
|
|
from alembic.ddl.base import visit_rename_table
|
|
from alembic.ddl.sqlite import SQLiteImpl
|
|
from alembic.operations import BatchOperations
|
|
from alembic.operations import Operations
|
|
from alembic.operations import ops
|
|
from sqlalchemy import Column
|
|
from sqlalchemy import text
|
|
from sqlalchemy.dialects.mysql.base import MySQLDialect
|
|
from sqlalchemy.dialects.sqlite.base import SQLiteDialect
|
|
from sqlalchemy.ext.compiler import compiles
|
|
from sqlalchemy.schema import DropTable
|
|
from sqlalchemy.sql import func
|
|
from sqlalchemy.types import TypeDecorator
|
|
|
|
from geoalchemy2 import Geography
|
|
from geoalchemy2 import Geometry
|
|
from geoalchemy2 import Raster
|
|
from geoalchemy2.admin.dialects.common import _check_spatial_type
|
|
from geoalchemy2.admin.dialects.common import _get_gis_cols
|
|
from geoalchemy2.admin.dialects.common import _spatial_idx_name
|
|
|
|
writer = rewriter.Rewriter()
|
|
"""Rewriter object for Alembic."""
|
|
|
|
_SPATIAL_TABLES = set()
|
|
|
|
|
|
class GeoPackageImpl(SQLiteImpl):
|
|
"""Class to copy the Alembic implementation from SQLite to GeoPackage."""
|
|
|
|
__dialect__ = "geopackage"
|
|
|
|
|
|
def _monkey_patch_get_indexes_for_sqlite():
|
|
"""Monkey patch SQLAlchemy to fix spatial index reflection."""
|
|
normal_behavior = SQLiteDialect.get_indexes
|
|
|
|
def spatial_behavior(self, connection, table_name, schema=None, **kw):
|
|
indexes = self._get_indexes_normal_behavior(connection, table_name, schema=None, **kw)
|
|
|
|
is_gpkg = connection.dialect.name == "geopackage"
|
|
|
|
try:
|
|
# Check that SpatiaLite was loaded into the DB
|
|
is_spatial_db = connection.exec_driver_sql(
|
|
"""PRAGMA main.table_info({})""".format(
|
|
"gpkg_geometry_columns" if is_gpkg else "geometry_columns"
|
|
)
|
|
).fetchall()
|
|
if not is_spatial_db:
|
|
return indexes
|
|
except AttributeError:
|
|
return indexes
|
|
|
|
# Get spatial indexes
|
|
if is_gpkg:
|
|
spatial_index_query = text(
|
|
"""SELECT A.table_name, A.column_name, IFNULL(B.has_index, 0) AS has_index
|
|
FROM "gpkg_geometry_columns"
|
|
AS A
|
|
LEFT JOIN (
|
|
SELECT table_name, column_name, COUNT(*) AS has_index
|
|
FROM gpkg_extensions
|
|
WHERE LOWER(table_name) = LOWER('{table_name}')
|
|
AND extension_name = 'gpkg_rtree_index'
|
|
) AS B
|
|
ON LOWER(A.table_name) = LOWER(B.table_name)
|
|
WHERE LOWER(A.table_name) = LOWER('{table_name}');
|
|
""".format(
|
|
table_name=table_name
|
|
)
|
|
)
|
|
else:
|
|
spatial_index_query = text(
|
|
"""SELECT *
|
|
FROM geometry_columns
|
|
WHERE f_table_name = '{}'
|
|
ORDER BY f_table_name, f_geometry_column;""".format(
|
|
table_name
|
|
)
|
|
)
|
|
|
|
spatial_indexes = connection.execute(spatial_index_query).fetchall()
|
|
|
|
if spatial_indexes:
|
|
reflected_names = set([i["name"] for i in indexes])
|
|
for idx in spatial_indexes:
|
|
idx_col = idx[1]
|
|
idx_name = _spatial_idx_name(table_name, idx_col)
|
|
if not bool(idx[-1]) or idx_name in reflected_names:
|
|
continue
|
|
indexes.append(
|
|
{
|
|
"name": idx_name,
|
|
"column_names": [idx_col],
|
|
"unique": 0,
|
|
"dialect_options": {"_column_flag": True},
|
|
}
|
|
)
|
|
reflected_names.add(idx_name)
|
|
|
|
return indexes
|
|
|
|
spatial_behavior.__doc__ = normal_behavior.__doc__
|
|
SQLiteDialect.get_indexes = spatial_behavior
|
|
SQLiteDialect._get_indexes_normal_behavior = normal_behavior
|
|
|
|
|
|
_monkey_patch_get_indexes_for_sqlite()
|
|
|
|
|
|
def _monkey_patch_get_indexes_for_mysql():
|
|
"""Monkey patch SQLAlchemy to fix spatial index reflection."""
|
|
normal_behavior = MySQLDialect.get_indexes
|
|
|
|
def spatial_behavior(self, connection, table_name, schema=None, **kw):
|
|
indexes = self._get_indexes_normal_behavior(connection, table_name, schema=None, **kw)
|
|
|
|
# Get spatial indexes
|
|
has_index_query = """SELECT DISTINCT
|
|
COLUMN_NAME
|
|
FROM INFORMATION_SCHEMA.STATISTICS
|
|
WHERE TABLE_NAME = '{}' AND INDEX_TYPE = 'SPATIAL'""".format(
|
|
table_name
|
|
)
|
|
if schema is not None:
|
|
has_index_query += """ AND TABLE_SCHEMA = '{}'""".format(schema)
|
|
spatial_indexes = connection.execute(text(has_index_query)).fetchall()
|
|
|
|
if spatial_indexes:
|
|
reflected_names = set([i["name"] for i in indexes])
|
|
for idx in spatial_indexes:
|
|
idx_col = idx[0]
|
|
idx_name = _spatial_idx_name(table_name, idx_col)
|
|
if idx_name in reflected_names:
|
|
continue
|
|
indexes.append(
|
|
{
|
|
"name": idx_name,
|
|
"column_names": [idx_col],
|
|
"unique": 0,
|
|
"dialect_options": {"_column_flag": True},
|
|
}
|
|
)
|
|
reflected_names.add(idx_name)
|
|
|
|
return indexes
|
|
|
|
spatial_behavior.__doc__ = normal_behavior.__doc__
|
|
MySQLDialect.get_indexes = spatial_behavior
|
|
MySQLDialect._get_indexes_normal_behavior = normal_behavior
|
|
|
|
|
|
_monkey_patch_get_indexes_for_mysql()
|
|
|
|
|
|
def render_item(obj_type, obj, autogen_context):
|
|
"""Add proper imports for spatial types."""
|
|
if obj_type == "type" and isinstance(obj, (Geometry, Geography, Raster)):
|
|
import_name = obj.__class__.__name__
|
|
autogen_context.imports.add(f"from geoalchemy2 import {import_name}")
|
|
return "%r" % obj
|
|
|
|
# Default rendering for other objects
|
|
return False
|
|
|
|
|
|
def include_object(obj, name, obj_type, reflected, compare_to):
|
|
"""Do not include internal tables of spatial extensions.
|
|
|
|
.. warning::
|
|
This function only checks the table names, so it might exclude tables that should not be.
|
|
In such case, you should create your own function to handle your specific table names.
|
|
|
|
"""
|
|
if obj_type == "table" and (
|
|
name.startswith("geometry_columns")
|
|
or name.startswith("spatial_ref_sys")
|
|
or name.startswith("spatialite_history")
|
|
or name.startswith("sqlite_sequence")
|
|
or name.startswith("views_geometry_columns")
|
|
or name.startswith("virts_geometry_columns")
|
|
or name.startswith("idx_")
|
|
or name.startswith("gpkg_")
|
|
or name.startswith("vgpkg_")
|
|
):
|
|
return False
|
|
return True
|
|
|
|
|
|
@Operations.register_operation("add_geospatial_column")
|
|
@BatchOperations.register_operation("add_geospatial_column", "batch_add_geospatial_column")
|
|
class AddGeospatialColumnOp(ops.AddColumnOp):
|
|
"""Add a Geospatial Column in an Alembic migration context.
|
|
|
|
This method 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 DropGeospatialColumnOp.from_column_and_tablename(
|
|
self.schema, self.table_name, self.column.name
|
|
)
|
|
|
|
@classmethod
|
|
def batch_add_geospatial_column(
|
|
cls,
|
|
operations,
|
|
column,
|
|
insert_before=None,
|
|
insert_after=None,
|
|
):
|
|
"""Issue an "add column" instruction using the current batch migration context."""
|
|
kw = {}
|
|
if insert_before:
|
|
kw["insert_before"] = insert_before
|
|
if insert_after:
|
|
kw["insert_after"] = insert_after
|
|
|
|
op = cls(
|
|
operations.impl.table_name,
|
|
column,
|
|
schema=operations.impl.schema,
|
|
**kw,
|
|
)
|
|
return operations.invoke(op)
|
|
|
|
|
|
@Operations.register_operation("drop_geospatial_column")
|
|
@BatchOperations.register_operation("drop_geospatial_column", "batch_drop_geospatial_column")
|
|
class DropGeospatialColumnOp(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 AddGeospatialColumnOp.from_column_and_tablename(
|
|
self.schema, self.table_name, self.column
|
|
)
|
|
|
|
@classmethod
|
|
def batch_drop_geospatial_column(cls, operations, column_name, **kw):
|
|
"""Issue a "drop column" instruction using the current batch migration context."""
|
|
op = cls(
|
|
operations.impl.table_name,
|
|
column_name,
|
|
schema=operations.impl.schema,
|
|
**kw,
|
|
)
|
|
return operations.invoke(op)
|
|
|
|
|
|
@Operations.implementation_for(AddGeospatialColumnOp)
|
|
def add_geospatial_column(operations, operation):
|
|
"""Handle the actual column addition according to the dialect backend.
|
|
|
|
Args:
|
|
operations: Operations object from alembic base, defining high level migration operations.
|
|
operation: AddGeospatialColumnOp 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 dialect.name == "sqlite":
|
|
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
|
|
operations.execute(
|
|
func.AddGeometryColumn(
|
|
table_name,
|
|
column_name,
|
|
geospatial_core_type.srid,
|
|
geospatial_core_type.geometry_type,
|
|
geospatial_core_type.dimension,
|
|
not geospatial_core_type.nullable,
|
|
)
|
|
)
|
|
elif "postgresql" in dialect.name:
|
|
operations.impl.add_column(
|
|
table_name,
|
|
operation.column,
|
|
schema=operation.schema,
|
|
)
|
|
|
|
|
|
@Operations.implementation_for(DropGeospatialColumnOp)
|
|
def drop_geospatial_column(operations, operation):
|
|
"""Handle the actual column removal according to the dialect backend.
|
|
|
|
Args:
|
|
operations: Operations object from alembic base, defining high level migration operations.
|
|
operation: AddGeospatialColumnOp call, with attributes for table_name, column_name,
|
|
column_type, and optional keywords.
|
|
"""
|
|
table_name = operation.table_name
|
|
column = operation.to_column(operations.migration_context)
|
|
|
|
dialect = operations.get_bind().dialect
|
|
|
|
if dialect.name == "sqlite":
|
|
_SPATIAL_TABLES.add(table_name)
|
|
operations.impl.drop_column(table_name, column, schema=operation.schema, **operation.kw)
|
|
|
|
|
|
@compiles(RenameTable, "sqlite")
|
|
def visit_rename_geospatial_table(element, compiler, **kw):
|
|
"""Specific compilation rule to rename spatial tables with SQLite dialect."""
|
|
table_is_spatial = element.table_name in _SPATIAL_TABLES
|
|
new_table_is_spatial = element.new_table_name in _SPATIAL_TABLES
|
|
|
|
if table_is_spatial or new_table_is_spatial:
|
|
# Here we suppose there is only one DB attached to the current engine, so the prefix
|
|
# is set to NULL
|
|
return "SELECT RenameTable(NULL, '%s', '%s')" % (
|
|
format_table_name(compiler, element.table_name, element.schema),
|
|
format_table_name(compiler, element.new_table_name, element.schema),
|
|
)
|
|
else:
|
|
return visit_rename_table(element, compiler, **kw)
|
|
|
|
|
|
@compiles(DropTable, "sqlite")
|
|
def visit_drop_geospatial_table(element, compiler, **kw):
|
|
"""Specific compilation rule to drop spatial tables with SQLite dialect."""
|
|
table_is_spatial = element.element.name in _SPATIAL_TABLES
|
|
|
|
if table_is_spatial:
|
|
# Here we suppose there is only one DB attached to the current engine
|
|
return "SELECT DropTable(NULL, '%s')" % (
|
|
format_table_name(compiler, element.element.name, None),
|
|
)
|
|
else:
|
|
return compiler.visit_drop_table(element, **kw)
|
|
|
|
|
|
@renderers.dispatch_for(AddGeospatialColumnOp)
|
|
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(DropGeospatialColumnOp)
|
|
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(")
|
|
|
|
|
|
@writer.rewrites(ops.AddColumnOp)
|
|
def add_geo_column(context, revision, op):
|
|
"""Replace 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)):
|
|
op.column.type.spatial_index = False
|
|
op.column.type._spatial_index_reflected = None
|
|
new_op = AddGeospatialColumnOp(op.table_name, op.column, schema=op.schema)
|
|
else:
|
|
new_op = op
|
|
return new_op
|
|
|
|
|
|
@writer.rewrites(ops.DropColumnOp)
|
|
def drop_geo_column(context, revision, op):
|
|
"""Replace 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 = DropGeospatialColumnOp(op.table_name, op.column_name, schema=op.schema)
|
|
else:
|
|
new_op = op
|
|
return new_op
|
|
|
|
|
|
@Operations.register_operation("create_geospatial_table")
|
|
class CreateGeospatialTableOp(ops.CreateTableOp):
|
|
"""Create a Geospatial Table in an Alembic migration context.
|
|
|
|
This method originates from:
|
|
https://alembic.sqlalchemy.org/en/latest/api/operations.html#operation-plugins
|
|
"""
|
|
|
|
@classmethod
|
|
def create_geospatial_table(cls, operations, table_name, *columns, **kw):
|
|
"""Handle the different situations arising from creating geospatial table to a DB."""
|
|
op = cls(table_name, columns, **kw)
|
|
return operations.invoke(op)
|
|
|
|
def reverse(self):
|
|
"""Used to autogenerate the downgrade function."""
|
|
return DropGeospatialColumnOp.from_table(
|
|
self.to_table(),
|
|
_namespace_metadata=self._namespace_metadata,
|
|
)
|
|
|
|
@classmethod
|
|
def from_table(cls, table, _namespace_metadata=None):
|
|
obj = super().from_table(table, _namespace_metadata)
|
|
return obj
|
|
|
|
def to_table(self, migration_context=None):
|
|
table = super().to_table(migration_context)
|
|
|
|
# Set spatial_index attribute to False so the indexes are created explicitly
|
|
for col in table.columns:
|
|
try:
|
|
if col.type.spatial_index:
|
|
col.type.spatial_index = False
|
|
except AttributeError:
|
|
pass
|
|
return table
|
|
|
|
|
|
@Operations.register_operation("drop_geospatial_table")
|
|
class DropGeospatialTableOp(ops.DropTableOp):
|
|
@classmethod
|
|
def drop_geospatial_table(cls, operations, table_name, schema=None, **kw):
|
|
"""Handle the different situations arising from dropping geospatial table from a DB."""
|
|
op = cls(table_name, schema=schema, table_kw=kw)
|
|
return operations.invoke(op)
|
|
|
|
def reverse(self):
|
|
"""Used to autogenerate the downgrade function."""
|
|
return CreateGeospatialTableOp.from_table(
|
|
self.to_table(),
|
|
_namespace_metadata=self._namespace_metadata,
|
|
)
|
|
|
|
@classmethod
|
|
def from_table(cls, table, _namespace_metadata=None):
|
|
obj = super().from_table(table, _namespace_metadata)
|
|
return obj
|
|
|
|
def to_table(self, migration_context=None):
|
|
table = super().to_table(migration_context)
|
|
return table
|
|
|
|
|
|
@Operations.implementation_for(CreateGeospatialTableOp)
|
|
def create_geospatial_table(operations, operation):
|
|
"""Handle the actual table creation according to the dialect backend.
|
|
|
|
Args:
|
|
operations: Operations object from alembic base, defining high level migration operations.
|
|
operation: CreateGeospatialTableOp call, with attributes for table_name, column_name,
|
|
column_type, and optional keywords.
|
|
"""
|
|
table_name = operation.table_name
|
|
bind = operations.get_bind()
|
|
|
|
# For now the default events defined in geoalchemy2 are enough to handle table creation
|
|
operations.create_table(table_name, *operation.columns, schema=operation.schema, **operation.kw)
|
|
|
|
if bind.dialect.name == "sqlite":
|
|
_SPATIAL_TABLES.add(table_name)
|
|
|
|
|
|
@Operations.implementation_for(DropGeospatialTableOp)
|
|
def drop_geospatial_table(operations, operation):
|
|
"""Handle the actual table removal according to the dialect backend.
|
|
|
|
Args:
|
|
operations: Operations object from alembic base, defining high level migration operations.
|
|
operation: DropGeospatialTableOp call, with attributes for table_name, column_name,
|
|
column_type, and optional keywords.
|
|
"""
|
|
table_name = operation.table_name
|
|
bind = operations.get_bind()
|
|
dialect = bind.dialect
|
|
|
|
if dialect.name == "sqlite":
|
|
_SPATIAL_TABLES.add(table_name)
|
|
operations.drop_table(table_name, schema=operation.schema, **operation.table_kw)
|
|
|
|
|
|
@renderers.dispatch_for(CreateGeospatialTableOp)
|
|
def render_create_geo_table(autogen_context, op):
|
|
"""Render the create_geospatial_table operation in migration script."""
|
|
table_render = _add_table(autogen_context, op)
|
|
return table_render.replace(".create_table(", ".create_geospatial_table(")
|
|
|
|
|
|
@renderers.dispatch_for(DropGeospatialTableOp)
|
|
def render_drop_geo_table(autogen_context, op):
|
|
"""Render the drop_geospatial_table operation in migration script."""
|
|
table_render = _drop_table(autogen_context, op)
|
|
return table_render.replace(".drop_table(", ".drop_geospatial_table(")
|
|
|
|
|
|
@writer.rewrites(ops.CreateTableOp)
|
|
def create_geo_table(context, revision, op):
|
|
"""Replace the default CreateTableOp by a geospatial-specific one."""
|
|
dialect = context.bind.dialect
|
|
gis_cols = _get_gis_cols(op, (Geometry, Geography, Raster), dialect)
|
|
|
|
if gis_cols:
|
|
new_op = CreateGeospatialTableOp(
|
|
op.table_name,
|
|
op.columns,
|
|
schema=op.schema,
|
|
_namespace_metadata=op._namespace_metadata,
|
|
_constraints_included=op._constraints_included,
|
|
**op.kw,
|
|
)
|
|
else:
|
|
new_op = op
|
|
|
|
return new_op
|
|
|
|
|
|
@writer.rewrites(ops.DropTableOp)
|
|
def drop_geo_table(context, revision, op):
|
|
"""Replace the default DropTableOp by a geospatial-specific one."""
|
|
dialect = context.bind.dialect
|
|
table = op.to_table()
|
|
gis_cols = _get_gis_cols(table, (Geometry, Geography, Raster), dialect)
|
|
|
|
if gis_cols:
|
|
new_op = DropGeospatialTableOp(op.table_name, schema=op.schema)
|
|
else:
|
|
new_op = op
|
|
|
|
return new_op
|
|
|
|
|
|
@Operations.register_operation("create_geospatial_index")
|
|
@BatchOperations.register_operation("create_geospatial_index", "batch_create_geospatial_index")
|
|
class CreateGeospatialIndexOp(ops.CreateIndexOp):
|
|
@classmethod
|
|
def create_geospatial_index(
|
|
cls,
|
|
operations,
|
|
index_name,
|
|
table_name,
|
|
columns,
|
|
schema=None,
|
|
unique=False,
|
|
**kw,
|
|
):
|
|
"""Handle the different situations arising from creating geospatial index into a DB."""
|
|
op = cls(index_name, table_name, columns, schema=schema, unique=unique, **kw)
|
|
return operations.invoke(op)
|
|
|
|
def reverse(self):
|
|
"""Used to autogenerate the downgrade function."""
|
|
return DropGeospatialIndexOp(
|
|
self.index_name,
|
|
self.table_name,
|
|
column_name=self.columns[0].name,
|
|
schema=self.schema,
|
|
)
|
|
|
|
@classmethod
|
|
def batch_create_geospatial_index(
|
|
cls,
|
|
operations,
|
|
index_name,
|
|
columns,
|
|
**kw,
|
|
):
|
|
"""Issue a "create index" instruction using the current batch migration context."""
|
|
op = cls(
|
|
index_name,
|
|
operations.impl.table_name,
|
|
columns,
|
|
schema=operations.impl.schema,
|
|
**kw,
|
|
)
|
|
return operations.invoke(op)
|
|
|
|
|
|
@Operations.register_operation("drop_geospatial_index")
|
|
@BatchOperations.register_operation("drop_geospatial_index", "batch_drop_geospatial_index")
|
|
class DropGeospatialIndexOp(ops.DropIndexOp):
|
|
def __init__(self, *args, column_name, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.column_name = column_name
|
|
|
|
@classmethod
|
|
def drop_geospatial_index(
|
|
cls,
|
|
operations,
|
|
index_name,
|
|
table_name,
|
|
column_name,
|
|
schema=None,
|
|
unique=False,
|
|
**kw,
|
|
):
|
|
"""Handle the different situations arising from dropping geospatial index from a DB."""
|
|
op = cls(
|
|
index_name,
|
|
table_name=table_name,
|
|
column_name=column_name,
|
|
schema=schema,
|
|
unique=unique,
|
|
**kw,
|
|
)
|
|
return operations.invoke(op)
|
|
|
|
def reverse(self):
|
|
"""Used to autogenerate the downgrade function."""
|
|
return CreateGeospatialIndexOp(
|
|
self.index_name,
|
|
self.table_name,
|
|
column_name=self.column_name,
|
|
schema=self.schema,
|
|
_reverse=self,
|
|
**self.kw,
|
|
)
|
|
|
|
@classmethod
|
|
def from_index(cls, index):
|
|
assert index.table is not None
|
|
assert len(index.columns) == 1, "A spatial index must be set on one column only"
|
|
return cls(
|
|
index.name,
|
|
index.table.name,
|
|
column_name=index.columns[0].name,
|
|
schema=index.table.schema,
|
|
_reverse=CreateGeospatialIndexOp.from_index(index),
|
|
**index.kwargs,
|
|
)
|
|
|
|
@classmethod
|
|
def batch_drop_geospatial_index(cls, operations, index_name, **kw):
|
|
"""Issue a "drop index" instruction using the current batch migration context."""
|
|
op = cls(
|
|
index_name,
|
|
table_name=operations.impl.table_name,
|
|
schema=operations.impl.schema,
|
|
**kw,
|
|
)
|
|
return operations.invoke(op)
|
|
|
|
|
|
@Operations.implementation_for(CreateGeospatialIndexOp)
|
|
def create_geospatial_index(operations, operation):
|
|
"""Handle the actual index creation according to the dialect backend.
|
|
|
|
Args:
|
|
operations: Operations object from alembic base, defining high level migration operations.
|
|
operation: CreateGeospatialIndexOp call, with attributes for table_name, column_name,
|
|
column_type, and optional keywords.
|
|
"""
|
|
bind = operations.get_bind()
|
|
dialect = bind.dialect
|
|
|
|
if dialect.name == "sqlite":
|
|
assert len(operation.columns) == 1, "A spatial index must be set on one column only"
|
|
operations.execute(func.CreateSpatialIndex(operation.table_name, operation.columns[0]))
|
|
else:
|
|
idx = operation.to_index(operations.migration_context)
|
|
operations.impl.create_index(idx)
|
|
|
|
|
|
@Operations.implementation_for(DropGeospatialIndexOp)
|
|
def drop_geospatial_index(operations, operation):
|
|
"""Handle the actual index drop according to the dialect backend.
|
|
|
|
Args:
|
|
operations: Operations object from alembic base, defining high level migration operations.
|
|
operation: DropGeospatialIndexOp call, with attributes for table_name, column_name,
|
|
column_type, and optional keywords.
|
|
"""
|
|
bind = operations.get_bind()
|
|
dialect = bind.dialect
|
|
|
|
if dialect.name == "sqlite":
|
|
operations.execute(func.DisableSpatialIndex(operation.table_name, operation.column_name))
|
|
else:
|
|
operations.impl.drop_index(operation.to_index(operations.migration_context))
|
|
|
|
|
|
@renderers.dispatch_for(CreateGeospatialIndexOp)
|
|
def render_create_geo_index(autogen_context, op):
|
|
"""Render the create_geospatial_index operation in migration script."""
|
|
idx_render = _add_index(autogen_context, op)
|
|
return idx_render.replace(".create_index(", ".create_geospatial_index(")
|
|
|
|
|
|
@renderers.dispatch_for(DropGeospatialIndexOp)
|
|
def render_drop_geo_index(autogen_context, op):
|
|
"""Render the drop_geospatial_index operation in migration script."""
|
|
idx_render = _drop_index(autogen_context, op)
|
|
|
|
# Replace function name
|
|
text = idx_render.replace(".drop_index(", ".drop_geospatial_index(")
|
|
|
|
# Add column name as keyword argument
|
|
text = text[:-1] + ", column_name='%s')" % (op.column_name,)
|
|
|
|
return text
|
|
|
|
|
|
@writer.rewrites(ops.CreateIndexOp)
|
|
def create_geo_index(context, revision, op):
|
|
"""Replace the default CreateIndexOp by a geospatial-specific one."""
|
|
dialect = context.bind.dialect
|
|
|
|
if len(op.columns) == 1:
|
|
col = op.columns[0]
|
|
if isinstance(col, Column) and _check_spatial_type(
|
|
col.type, (Geometry, Geography, Raster), dialect
|
|
):
|
|
# Fix index properties
|
|
op.kw["postgresql_using"] = op.kw.get("postgresql_using", "gist")
|
|
if col.type.use_N_D_index:
|
|
postgresql_ops = {col.name: "gist_geometry_ops_nd"}
|
|
else:
|
|
postgresql_ops = {}
|
|
op.kw["postgresql_ops"] = op.kw.get("postgresql_ops", postgresql_ops)
|
|
|
|
return CreateGeospatialIndexOp(
|
|
op.index_name,
|
|
op.table_name,
|
|
op.columns,
|
|
schema=op.schema,
|
|
unique=op.unique,
|
|
**op.kw,
|
|
)
|
|
|
|
return op
|
|
|
|
|
|
@writer.rewrites(ops.DropIndexOp)
|
|
def drop_geo_index(context, revision, op):
|
|
"""Replace the default DropIndexOp by a geospatial-specific one."""
|
|
dialect = context.bind.dialect
|
|
idx = op.to_index()
|
|
|
|
if len(idx.columns) == 1:
|
|
col = idx.columns[0]
|
|
if isinstance(col, Column) and _check_spatial_type(
|
|
col.type, (Geometry, Geography, Raster), dialect
|
|
):
|
|
return DropGeospatialIndexOp(
|
|
op.index_name,
|
|
table_name=op.table_name,
|
|
column_name=col.name,
|
|
schema=op.schema,
|
|
**op.kw,
|
|
)
|
|
|
|
return op
|