Feat: Add support of spatial indexes for Raster type for PostgreSQL

geoalchemy2/admin/dialects/postgresql.py

@@ -1,5 +1,7 @@
 """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
@@ -17,6 +19,8 @@
 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
@@ -25,6 +29,9 @@
 
 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
 
 
@@ -34,59 +41,74 @@ def create_spatial_index(bind, table, col):
         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,
+        col_func,
         postgresql_using="gist",
         postgresql_ops=postgresql_ops,
         _column_flag=True,
     )
-    idx.create(bind=bind)
+    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 isinstance(column_info.get("type"), Geometry):
+    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.endswith("ZM"):
-        coord_dimension = 4
-    elif geometry_type[-1] in ["Z", "M"]:
-        coord_dimension = 3
+    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 = ""
 
-    has_index_query = """SELECT (indexrelid IS NOT NULL) AS has_index
-        FROM (
-            SELECT
-                    n.nspname,
-                    c.relname,
-                    c.oid AS relid,
-                    a.attname,
-                    a.attnum
-            FROM pg_attribute a
-            INNER JOIN pg_class c ON (a.attrelid=c.oid)
-            INNER JOIN pg_type t ON (a.atttypid=t.oid)
-            INNER JOIN pg_namespace n ON (c.relnamespace=n.oid)
-            WHERE t.typname='geometry'
-                    AND c.relkind='r'
-        ) g
-        LEFT JOIN pg_index i ON (g.relid = i.indrelid AND g.attnum = ANY(i.indkey))
-        WHERE relname = '{}' AND attname = '{}'{};
-    """.format(
-        table.name, column_info["name"], 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
-    column_info["type"].geometry_type = geometry_type
-    column_info["type"].dimension = coord_dimension
+    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
@@ -105,7 +127,7 @@ def before_create(table, bind, **kw):
     for idx in current_indexes:
         for col in table.info["_saved_columns"]:
             if (
-                _check_spatial_type(col.type, Geometry, dialect) and check_management(col)
+                _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(
@@ -134,9 +156,9 @@ def after_create(table, bind, **kw):
             stmt = stmt.execution_options(autocommit=True)
             bind.execute(stmt)
 
-        # Add spatial indices for the Geometry and Geography columns
+        # Add spatial indices for the Geometry, Geography and Raster columns
         if (
-            _check_spatial_type(col.type, (Geometry, Geography), dialect)
+            _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
@@ -156,6 +178,9 @@ def before_drop(table, bind, **kw):
 
     # 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])
 

geoalchemy2/types/__init__.py

@@ -345,7 +345,14 @@ class Raster(_GISType):
     cache_ok = True
     """ Enable cache for this type. """
 
-    def __init__(self, spatial_index=True, from_text=None, name=None, nullable=True) -> None:
+    def __init__(
+        self,
+        spatial_index=True,
+        from_text=None,
+        name=None,
+        nullable=True,
+        _spatial_index_reflected=None,
+    ) -> None:
         # Enforce default values
         super(Raster, self).__init__(
             geometry_type=None,
@@ -357,6 +364,7 @@ def __init__(self, spatial_index=True, from_text=None, name=None, nullable=True)
             from_text=from_text,
             name=name,
             nullable=nullable,
+            _spatial_index_reflected=_spatial_index_reflected,
         )
         self.extended = None
 

tests/schema_fixtures.py

@@ -193,5 +193,12 @@ class Lake(base):
                 geom_z = Column(Geometry(geometry_type="LINESTRINGZ", srid=4326, dimension=3))
                 geom_m = Column(Geometry(geometry_type="LINESTRINGM", srid=4326, dimension=3))
                 geom_zm = Column(Geometry(geometry_type="LINESTRINGZM", srid=4326, dimension=4))
+            if dialect_name in ["postgresql"]:
+                geom_geog = Column(Geography(geometry_type="LINESTRING"))
+                geom_geog_no_idx = Column(
+                    Geography(geometry_type="LINESTRING", spatial_index=False)
+                )
+                rast = Column(Raster(spatial_index=True))
+                rast_no_idx = Column(Raster(spatial_index=False))
 
     return metadata

tests/test_functional.py

@@ -1230,6 +1230,31 @@ def test_reflection(self, conn, setup_reflection_tables, dialect_name):
                 assert type_.srid == 4326
                 assert type_.dimension == 4
 
+            if dialect_name == "postgresql":
+                type_ = t.c.geom_geog.type
+                assert isinstance(type_, Geography)
+                assert type_.geometry_type == "LINESTRING"
+                assert type_.srid == 4326
+                assert type_.dimension == 2
+
+                type_ = t.c.geom_geog_no_idx.type
+                assert isinstance(type_, Geography)
+                assert type_.geometry_type == "LINESTRING"
+                assert type_.srid == 4326
+                assert type_.dimension == 2
+
+                type_ = t.c.rast.type
+                assert isinstance(type_, Raster)
+                assert type_.geometry_type is None
+                assert type_.srid == -1
+                assert type_.dimension is None
+
+                type_ = t.c.rast_no_idx.type
+                assert isinstance(type_, Raster)
+                assert type_.geometry_type is None
+                assert type_.srid == -1
+                assert type_.dimension is None
+
         # Drop the table
         t.drop(bind=conn)
 
@@ -1281,6 +1306,10 @@ def test_reflection(self, conn, setup_reflection_tables, dialect_name):
                         "idx_lake_geom",
                         "CREATE INDEX idx_lake_geom ON gis.lake USING gist (geom)",
                     ),
+                    (
+                        "idx_lake_geom_geog",
+                        "CREATE INDEX idx_lake_geom_geog ON gis.lake USING gist (geom_geog)",
+                    ),
                     (
                         "idx_lake_geom_m",
                         "CREATE INDEX idx_lake_geom_m ON gis.lake USING gist (geom_m)",
@@ -1293,6 +1322,10 @@ def test_reflection(self, conn, setup_reflection_tables, dialect_name):
                         "idx_lake_geom_zm",
                         "CREATE INDEX idx_lake_geom_zm ON gis.lake USING gist (geom_zm)",
                     ),
+                    (
+                        "idx_lake_rast",
+                        "CREATE INDEX idx_lake_rast ON gis.lake USING gist (st_convexhull(rast))",
+                    ),
                     (
                         "lake_pkey",
                         "CREATE UNIQUE INDEX lake_pkey ON gis.lake USING btree (id)",

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = py{310,311,312}-sqla{14,latest}, pypy3-sqla{14,latest}, lint, coverage, docs
+envlist = py{310,311,312,313}-sqla{14,latest}, pypy3-sqla{14,latest}, lint, coverage, docs
 requires=
     setuptools>42