Use CompositeTypeΒΆ

Some functions return composite types. This example shows how to deal with this kind of functions.

  8 import pytest
  9 from pkg_resources import parse_version
 10 from sqlalchemy import Column
 11 from sqlalchemy import Float
 12 from sqlalchemy import Integer
 13 from sqlalchemy import MetaData
 14 from sqlalchemy import __version__ as SA_VERSION
 15 from sqlalchemy.ext.declarative import declarative_base
 16
 17 from geoalchemy2 import Raster
 18 from geoalchemy2 import WKTElement
 19 from geoalchemy2.functions import GenericFunction
 20 from geoalchemy2.types import CompositeType
 21
 22 # Tests imports
 23 from tests import select
 24 from tests import test_only_with_dialects
 25
 26
 27 class SummaryStatsCustomType(CompositeType):
 28     """Define the composite type returned by the function ST_SummaryStatsAgg."""
 29
 30     typemap = {
 31         "count": Integer,
 32         "sum": Float,
 33         "mean": Float,
 34         "stddev": Float,
 35         "min": Float,
 36         "max": Float,
 37     }
 38
 39     cache_ok = True
 40
 41
 42 class ST_SummaryStatsAgg(GenericFunction):
 43     type = SummaryStatsCustomType
 44     # Set a specific identifier to not override the actual ST_SummaryStatsAgg function
 45     identifier = "ST_SummaryStatsAgg_custom"
 46
 47     inherit_cache = True
 48
 49
 50 metadata = MetaData()
 51 Base = declarative_base(metadata=metadata)
 52
 53
 54 class Ocean(Base):
 55     __tablename__ = "ocean"
 56     id = Column(Integer, primary_key=True)
 57     rast = Column(Raster)
 58
 59     def __init__(self, rast):
 60         self.rast = rast
 61
 62
 63 @test_only_with_dialects("postgresql")
 64 class TestSTSummaryStatsAgg:
 65     @pytest.mark.skipif(
 66         parse_version(SA_VERSION) < parse_version("1.4"),
 67         reason="requires SQLAlchely>1.4",
 68     )
 69     def test_st_summary_stats_agg(self, session, conn):
 70         metadata.drop_all(conn, checkfirst=True)
 71         metadata.create_all(conn)
 72
 73         # Create a new raster
 74         polygon = WKTElement("POLYGON((0 0,1 1,0 1,0 0))", srid=4326)
 75         o = Ocean(polygon.ST_AsRaster(5, 6))
 76         session.add(o)
 77         session.flush()
 78
 79         # Define the query to compute stats
 80         stats_agg = select([Ocean.rast.ST_SummaryStatsAgg_custom(1, True, 1).label("stats")])
 81         stats_agg_alias = stats_agg.alias("stats_agg")
 82
 83         # Use these stats
 84         query = select(
 85             [
 86                 stats_agg_alias.c.stats.count.label("count"),
 87                 stats_agg_alias.c.stats.sum.label("sum"),
 88                 stats_agg_alias.c.stats.mean.label("mean"),
 89                 stats_agg_alias.c.stats.stddev.label("stddev"),
 90                 stats_agg_alias.c.stats.min.label("min"),
 91                 stats_agg_alias.c.stats.max.label("max"),
 92             ]
 93         )
 94
 95         # Check the query
 96         assert str(query.compile(dialect=session.bind.dialect)) == (
 97             "SELECT "
 98             "(stats_agg.stats).count AS count, "
 99             "(stats_agg.stats).sum AS sum, "
100             "(stats_agg.stats).mean AS mean, "
101             "(stats_agg.stats).stddev AS stddev, "
102             "(stats_agg.stats).min AS min, "
103             "(stats_agg.stats).max AS max \n"
104             "FROM ("
105             "SELECT "
106             "ST_SummaryStatsAgg("
107             "ocean.rast, "
108             "%(ST_SummaryStatsAgg_1)s, %(ST_SummaryStatsAgg_2)s, %(ST_SummaryStatsAgg_3)s"
109             ") AS stats \n"
110             "FROM ocean) AS stats_agg"
111         )
112
113         # Execute the query
114         res = session.execute(query).fetchall()
115
116         # Check the result
117         assert res == [(15, 15.0, 1.0, 0.0, 1.0, 1.0)]

Gallery generated by Sphinx-Gallery