Note
Go to the end to download the full example code.
Automatically use a function at insert or selectΒΆ
Sometimes the application wants to apply a function in an insert or in a select.
For example, the application might need the geometry with lat/lon coordinates while they
are projected in the DB. To avoid having to always tweak the query with a
ST_Transform(), it is possible to define a TypeDecorator
11 import re
12 from typing import Any
13
14 import shapely
15 from sqlalchemy import Column
16 from sqlalchemy import Integer
17 from sqlalchemy import MetaData
18 from sqlalchemy import func
19 from sqlalchemy import text
20 from sqlalchemy.orm import declarative_base
21 from sqlalchemy.types import TypeDecorator
22
23 from geoalchemy2 import Geometry
24 from geoalchemy2 import shape
25
26 # Tests imports
27 from geoalchemy2.elements import WKTElement
28 from tests import test_only_with_dialects
29
30 metadata = MetaData()
31
32 Base = declarative_base(metadata=metadata)
33
34
35 class TransformedGeometry(TypeDecorator):
36 """This class is used to insert a ST_Transform() in each insert or select."""
37
38 impl = Geometry
39
40 cache_ok = True
41
42 def __init__(self, db_srid, app_srid, **kwargs):
43 kwargs["srid"] = db_srid
44 super().__init__(**kwargs)
45 self.app_srid = app_srid
46 self.db_srid = db_srid
47
48 def column_expression(self, col):
49 """Return the column expression with the correct type.
50
51 This is needed so that the returned element will also be decorated. In this case we don't
52 want to transform it again afterwards so we set the same SRID to both the ``db_srid`` and
53 ``app_srid`` arguments.
54 Without this the SRID of the WKBElement would be wrong.
55 """
56 return getattr(func, self.impl.as_binary)(
57 func.ST_Transform(col, self.app_srid),
58 type_=self.__class__(db_srid=self.app_srid, app_srid=self.app_srid),
59 )
60
61 def bind_expression(self, bindvalue):
62 return func.ST_Transform(
63 self.impl.bind_expression(bindvalue),
64 self.db_srid,
65 type_=self,
66 )
67
68
69 class ThreeDGeometry(TypeDecorator):
70 """This class is used to insert a ST_Force3D() in each insert."""
71
72 impl = Geometry
73
74 cache_ok = True
75
76 def bind_expression(self, bindvalue):
77 return func.ST_Force3D(
78 self.impl.bind_expression(bindvalue),
79 type=self,
80 )
81
82
83 class ShapelyGeometry(TypeDecorator):
84 """Custom type that should return Shapely objects from ORM reads."""
85
86 impl = Geometry
87
88 cache_ok = True
89
90 def process_bind_param(self, value, dialect):
91 if value is None:
92 return None
93 return shape.from_shape(value, srid=self.impl.srid)
94
95 def process_result_value(self, value, dialect):
96 if value is None:
97 return None
98 return shape.to_shape(value)
99
100
101 class Point(Base): # type: ignore
102 __tablename__ = "point"
103 id = Column(Integer, primary_key=True)
104 raw_geom = Column(Geometry(srid=4326, geometry_type="POINT"))
105 geom: Column[Any] = Column(
106 TransformedGeometry(db_srid=2154, app_srid=4326, geometry_type="POINT")
107 )
108 three_d_geom: Column = Column(ThreeDGeometry(srid=4326, geometry_type="POINTZ", dimension=3))
109 shapely_geom: Column = Column(ShapelyGeometry(srid=4326, geometry_type="POINT"))
110
111
112 def check_wkb(wkb, x, y):
113 pt = shape.to_shape(wkb)
114 assert round(pt.x, 5) == x
115 assert round(pt.y, 5) == y
116
117
118 @test_only_with_dialects("postgresql")
119 class TestTypeDecorator:
120 def _create_one_point(self, session, conn):
121 metadata.drop_all(conn, checkfirst=True)
122 metadata.create_all(conn)
123
124 # Create new point instance
125 p = Point()
126 wkt = "SRID=4326;POINT(5 45)"
127 p.raw_geom = wkt
128 p.geom = wkt
129 p.three_d_geom = wkt # Insert 2D geometry into 3D column
130 p.shapely_geom = shape.to_shape(
131 WKTElement(wkt)
132 ) # Insert Shapely geometry into ShapelyGeometry column
133
134 # Insert point
135 session.add(p)
136 session.flush()
137 session.expire(p)
138
139 return p.id
140
141 def test_transform(self, session, conn):
142 self._create_one_point(session, conn)
143
144 # Query the point and check the result
145 pt = session.query(Point).one()
146 assert pt.id == 1
147 assert pt.raw_geom.srid == 4326
148 check_wkb(pt.raw_geom, 5, 45)
149
150 assert pt.geom.srid == 4326
151 check_wkb(pt.geom, 5, 45)
152
153 # Check that the data is correct in DB using raw query
154 q = text("SELECT id, ST_AsEWKT(geom) AS geom FROM point;")
155 res_q = session.execute(q).fetchone()
156 assert res_q.id == 1
157 assert re.match(
158 r"SRID=2154;POINT\(857581\.8993196681? 6435414\.7478354[0-9]*\)", res_q.geom
159 )
160
161 # Compare geom, raw_geom with auto transform and explicit transform
162 pt_trans = session.query(
163 Point,
164 Point.raw_geom,
165 func.ST_Transform(Point.raw_geom, 2154).label("trans"),
166 ).one()
167
168 assert pt_trans[0].id == 1
169
170 assert pt_trans[0].geom.srid == 4326
171 check_wkb(pt_trans[0].geom, 5, 45)
172
173 assert pt_trans[0].raw_geom.srid == 4326
174 check_wkb(pt_trans[0].raw_geom, 5, 45)
175
176 assert pt_trans[1].srid == 4326
177 check_wkb(pt_trans[1], 5, 45)
178
179 assert pt_trans[2].srid == 2154
180 check_wkb(pt_trans[2], 857581.89932, 6435414.74784)
181
182 def test_force_3d(self, session, conn):
183 self._create_one_point(session, conn)
184
185 # Query the point and check the result
186 pt = session.query(Point).one()
187
188 assert pt.id == 1
189 assert pt.three_d_geom.srid == 4326
190 assert pt.three_d_geom.desc.lower() == (
191 "01010000a0e6100000000000000000144000000000008046400000000000000000"
192 )
193
194 def test_shapely(self, session, conn):
195 self._create_one_point(session, conn)
196
197 # Query the point and check the result
198 pt = session.query(Point).one()
199
200 assert pt.id == 1
201 assert isinstance(pt.shapely_geom, shapely.Point)
202 assert pt.shapely_geom.geom_type == "Point"
203 assert round(pt.shapely_geom.x, 5) == 5
204 assert round(pt.shapely_geom.y, 5) == 45
