233 lines
8.7 KiB
Python
233 lines
8.7 KiB
Python
|
"""This module defines specific functions for Postgresql dialect."""
|
||
|
|
|
||
|
|
import sqlalchemy
|
||
|
|
from packaging import version
|
||
|
|
from sqlalchemy import Index
|
||
|
|
from sqlalchemy import text
|
||
|
|
from sqlalchemy.dialects.postgresql.base import ischema_names as _postgresql_ischema_names
|
||
|
|
from sqlalchemy.ext.compiler import compiles
|
||
|
|
from sqlalchemy.sql import func
|
||
|
|
from sqlalchemy.sql import select
|
||
|
|
|
||
|
|
from geoalchemy2 import functions
|
||
|
|
from geoalchemy2.admin.dialects.common import _check_spatial_type
|
||
|
|
from geoalchemy2.admin.dialects.common import _format_select_args
|
||
|
|
from geoalchemy2.admin.dialects.common import _spatial_idx_name
|
||
|
|
from geoalchemy2.admin.dialects.common import compile_bin_literal
|
||
|
|
from geoalchemy2.admin.dialects.common import setup_create_drop
|
||
|
|
from geoalchemy2.types import Geography
|
||
|
|
from geoalchemy2.types import Geometry
|
||
|
|
from geoalchemy2.types import Raster
|
||
|
|
|
||
|
|
_SQLALCHEMY_VERSION_BEFORE_2 = version.parse(sqlalchemy.__version__) < version.parse("2")
|
||
|
|
|
||
|
|
# Register Geometry, Geography and Raster to SQLAlchemy's reflection subsystems.
|
||
|
|
_postgresql_ischema_names["geometry"] = Geometry
|
||
|
|
_postgresql_ischema_names["geography"] = Geography
|
||
|
|
_postgresql_ischema_names["raster"] = Raster
|
||
|
|
|
||
|
|
|
||
|
|
def check_management(column):
|
||
|
|
"""Check if the column should be managed."""
|
||
|
|
if _check_spatial_type(column.type, Raster):
|
||
|
|
# Raster columns are not managed
|
||
|
|
return _SQLALCHEMY_VERSION_BEFORE_2
|
||
|
|
return getattr(column.type, "use_typmod", None) is False
|
||
|
|
|
||
|
|
|
||
|
|
def create_spatial_index(bind, table, col):
|
||
|
|
"""Create spatial index on the given column."""
|
||
|
|
if col.type.use_N_D_index:
|
||
|
|
postgresql_ops = {col.name: "gist_geometry_ops_nd"}
|
||
|
|
else:
|
||
|
|
postgresql_ops = {}
|
||
|
|
if _check_spatial_type(col.type, Raster):
|
||
|
|
col_func = func.ST_ConvexHull(col)
|
||
|
|
else:
|
||
|
|
col_func = col
|
||
|
|
idx = Index(
|
||
|
|
_spatial_idx_name(table.name, col.name),
|
||
|
|
col_func,
|
||
|
|
postgresql_using="gist",
|
||
|
|
postgresql_ops=postgresql_ops,
|
||
|
|
_column_flag=True,
|
||
|
|
)
|
||
|
|
if bind is not None:
|
||
|
|
idx.create(bind=bind)
|
||
|
|
return idx
|
||
|
|
|
||
|
|
|
||
|
|
def reflect_geometry_column(inspector, table, column_info):
|
||
|
|
"""Reflect a column of type Geometry with Postgresql dialect."""
|
||
|
|
if not _check_spatial_type(column_info.get("type"), (Geometry, Geography, Raster)):
|
||
|
|
return
|
||
|
|
geo_type = column_info["type"]
|
||
|
|
geometry_type = geo_type.geometry_type
|
||
|
|
coord_dimension = geo_type.dimension
|
||
|
|
if geometry_type is not None:
|
||
|
|
if geometry_type.endswith("ZM"):
|
||
|
|
coord_dimension = 4
|
||
|
|
elif geometry_type[-1] in ["Z", "M"]:
|
||
|
|
coord_dimension = 3
|
||
|
|
|
||
|
|
# Query to check a given column has spatial index
|
||
|
|
if table.schema is not None:
|
||
|
|
schema_part = " AND nspname = '{}'".format(table.schema)
|
||
|
|
else:
|
||
|
|
schema_part = ""
|
||
|
|
|
||
|
|
# Check if the column has a spatial index (the regular expression checks for the column name
|
||
|
|
# in the index definition, which is required for functional indexes)
|
||
|
|
has_index_query = """SELECT EXISTS (
|
||
|
|
SELECT 1
|
||
|
|
FROM pg_class t
|
||
|
|
JOIN pg_namespace n ON n.oid = t.relnamespace
|
||
|
|
JOIN pg_index ix ON t.oid = ix.indrelid
|
||
|
|
JOIN pg_class i ON i.oid = ix.indexrelid
|
||
|
|
JOIN pg_am am ON i.relam = am.oid
|
||
|
|
WHERE
|
||
|
|
t.relname = '{table_name}'{schema_part}
|
||
|
|
AND am.amname = 'gist'
|
||
|
|
AND (
|
||
|
|
EXISTS (
|
||
|
|
SELECT 1
|
||
|
|
FROM pg_attribute a
|
||
|
|
WHERE a.attrelid = t.oid
|
||
|
|
AND a.attnum = ANY(ix.indkey)
|
||
|
|
AND a.attname = '{col_name}'
|
||
|
|
)
|
||
|
|
OR pg_get_indexdef(
|
||
|
|
ix.indexrelid
|
||
|
|
) ~ '(^|[^a-zA-Z0-9_])("?{col_name}"?)($|[^a-zA-Z0-9_])'
|
||
|
|
)
|
||
|
|
);""".format(
|
||
|
|
table_name=table.name, col_name=column_info["name"], schema_part=schema_part
|
||
|
|
)
|
||
|
|
spatial_index = inspector.bind.execute(text(has_index_query)).scalar()
|
||
|
|
|
||
|
|
# Set attributes
|
||
|
|
if not _check_spatial_type(column_info["type"], Raster):
|
||
|
|
column_info["type"].geometry_type = geometry_type
|
||
|
|
column_info["type"].dimension = coord_dimension
|
||
|
|
column_info["type"].spatial_index = bool(spatial_index)
|
||
|
|
|
||
|
|
# Spatial indexes are automatically reflected with PostgreSQL dialect
|
||
|
|
column_info["type"]._spatial_index_reflected = True
|
||
|
|
|
||
|
|
|
||
|
|
def before_create(table, bind, **kw):
|
||
|
|
"""Handle spatial indexes during the before_create event."""
|
||
|
|
dialect, gis_cols, regular_cols = setup_create_drop(table, bind, check_management)
|
||
|
|
|
||
|
|
# Remove the spatial indexes from the table metadata because they should not be
|
||
|
|
# created during the table.create() step since the associated columns do not exist
|
||
|
|
# at this time.
|
||
|
|
table.info["_after_create_indexes"] = []
|
||
|
|
current_indexes = set(table.indexes)
|
||
|
|
for idx in current_indexes:
|
||
|
|
for col in table.info["_saved_columns"]:
|
||
|
|
if (
|
||
|
|
_check_spatial_type(col.type, (Geometry, Raster), dialect) and check_management(col)
|
||
|
|
) and col in idx.columns.values():
|
||
|
|
table.indexes.remove(idx)
|
||
|
|
if idx.name != _spatial_idx_name(table.name, col.name) or not getattr(
|
||
|
|
col.type, "spatial_index", False
|
||
|
|
):
|
||
|
|
table.info["_after_create_indexes"].append(idx)
|
||
|
|
|
||
|
|
|
||
|
|
def after_create(table, bind, **kw):
|
||
|
|
"""Handle spatial indexes during the after_create event."""
|
||
|
|
# Restore original column list including managed Geometry columns
|
||
|
|
dialect = bind.dialect
|
||
|
|
|
||
|
|
table.columns = table.info.pop("_saved_columns")
|
||
|
|
|
||
|
|
for col in table.columns:
|
||
|
|
# Add the managed Geometry columns with AddGeometryColumn()
|
||
|
|
if _check_spatial_type(col.type, Geometry, dialect) and check_management(col):
|
||
|
|
dimension = col.type.dimension
|
||
|
|
args = [table.schema] if table.schema else []
|
||
|
|
args.extend([table.name, col.name, col.type.srid, col.type.geometry_type, dimension])
|
||
|
|
if col.type.use_typmod is not None:
|
||
|
|
args.append(col.type.use_typmod)
|
||
|
|
|
||
|
|
stmt = select(*_format_select_args(func.AddGeometryColumn(*args)))
|
||
|
|
stmt = stmt.execution_options(autocommit=True)
|
||
|
|
bind.execute(stmt)
|
||
|
|
|
||
|
|
# Add spatial indices for the Geometry, Geography and Raster columns
|
||
|
|
if (
|
||
|
|
_check_spatial_type(col.type, (Geometry, Geography, Raster), dialect)
|
||
|
|
and col.type.spatial_index is True
|
||
|
|
):
|
||
|
|
# If the index does not exist, define it and create it
|
||
|
|
if not [i for i in table.indexes if col in i.columns.values()] and check_management(
|
||
|
|
col
|
||
|
|
):
|
||
|
|
create_spatial_index(bind, table, col)
|
||
|
|
|
||
|
|
for idx in table.info.pop("_after_create_indexes"):
|
||
|
|
table.indexes.add(idx)
|
||
|
|
idx.create(bind=bind)
|
||
|
|
|
||
|
|
|
||
|
|
def before_drop(table, bind, **kw):
|
||
|
|
"""Handle spatial indexes during the before_drop event."""
|
||
|
|
dialect, gis_cols, regular_cols = setup_create_drop(table, bind, check_management)
|
||
|
|
|
||
|
|
# Drop the managed Geometry columns
|
||
|
|
for col in gis_cols:
|
||
|
|
if _check_spatial_type(col.type, Raster):
|
||
|
|
# Raster columns are dropped with the table, no need to drop them separately
|
||
|
|
continue
|
||
|
|
args = [table.schema] if table.schema else []
|
||
|
|
args.extend([table.name, col.name])
|
||
|
|
|
||
|
|
stmt = select(*_format_select_args(func.DropGeometryColumn(*args)))
|
||
|
|
stmt = stmt.execution_options(autocommit=True)
|
||
|
|
bind.execute(stmt)
|
||
|
|
|
||
|
|
|
||
|
|
def after_drop(table, bind, **kw):
|
||
|
|
"""Handle spatial indexes during the after_drop event."""
|
||
|
|
# Restore original column list including managed Geometry columns
|
||
|
|
saved_cols = table.info.pop("_saved_columns", None)
|
||
|
|
if saved_cols is not None:
|
||
|
|
table.columns = saved_cols
|
||
|
|
|
||
|
|
|
||
|
|
def _compile_GeomFromWKB_Postgresql(element, compiler, **kw):
|
||
|
|
# Store the SRID
|
||
|
|
clauses = list(element.clauses)
|
||
|
|
try:
|
||
|
|
srid = clauses[1].value
|
||
|
|
except (IndexError, TypeError, ValueError):
|
||
|
|
srid = element.type.srid
|
||
|
|
|
||
|
|
if kw.get("literal_binds", False):
|
||
|
|
wkb_clause = compile_bin_literal(clauses[0])
|
||
|
|
prefix = "decode("
|
||
|
|
suffix = ", 'hex')"
|
||
|
|
else:
|
||
|
|
wkb_clause = clauses[0]
|
||
|
|
prefix = ""
|
||
|
|
suffix = ""
|
||
|
|
|
||
|
|
compiled = compiler.process(wkb_clause, **kw)
|
||
|
|
|
||
|
|
if srid > 0:
|
||
|
|
return "{}({}{}{}, {})".format(element.identifier, prefix, compiled, suffix, srid)
|
||
|
|
else:
|
||
|
|
return "{}({}{}{})".format(element.identifier, prefix, compiled, suffix)
|
||
|
|
|
||
|
|
|
||
|
|
@compiles(functions.ST_GeomFromWKB, "postgresql") # type: ignore
|
||
|
|
def _PostgreSQL_ST_GeomFromWKB(element, compiler, **kw):
|
||
|
|
return _compile_GeomFromWKB_Postgresql(element, compiler, **kw)
|
||
|
|
|
||
|
|
|
||
|
|
@compiles(functions.ST_GeomFromEWKB, "postgresql") # type: ignore
|
||
|
|
def _PostgreSQL_ST_GeomFromEWKB(element, compiler, **kw):
|
||
|
|
return _compile_GeomFromWKB_Postgresql(element, compiler, **kw)
|