Coverage for gws-app/gws/base/database/provider.py: 89%
184 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-23 20:06 +0200
1"""Base database provider."""
3import contextlib
4import threading
5from typing import Optional, cast
7import gws
8import gws.lib.sa as sa
10from . import connection
13class Config(gws.Config):
14 """Database provider"""
16 schemaCacheLifeTime: gws.Duration = '3600'
17 """How long table structures read from the database are cached."""
18 withPool: Optional[bool] = False
19 """Keep and reuse database connections in a pool."""
20 pool: Optional[dict]
21 """Connection pool options."""
24_thread_local = threading.local()
27def _connections() -> dict[str, sa.Connection]:
28 """Return the open connections of the current thread, keyed by provider uid."""
29 d = getattr(_thread_local, 'connections', None)
30 if d is None:
31 d = {}
32 _thread_local.connections = d
33 return d
36class Object(gws.DatabaseProvider):
37 """Base database provider.
39 Manages the SQLAlchemy engine and the per-thread connection, reflects table
40 structures and describes tables and columns, and runs plain SQL text.
42 Subclasses provide ``url``, ``split_table_name``, ``join_table_name`` and
43 ``table_bounds``, and extend ``describe_column`` for database-specific types.
44 """
46 saEngine: sa.Engine
47 """SQLAlchemy engine."""
48 saMetaMap: dict[str, sa.MetaData]
49 """Reflected metadata, keyed by schema name."""
51 def __getstate__(self):
52 """Return the state for pickling, without the engine and the metadata."""
53 return gws.u.omit(vars(self), 'saMetaMap', 'saEngine')
55 def configure(self):
56 # init a dummy engine just to check things
57 self.saEngine = self.create_engine(poolclass=sa.NullPool)
58 self.saMetaMap = {}
60 def activate(self):
61 self.saEngine = self.create_engine()
62 self.saMetaMap = {}
64 def engine(self):
65 eng = getattr(self, 'saEngine', None)
66 if eng is not None:
67 return eng
68 self.saEngine = self.create_engine()
69 return self.saEngine
71 def create_engine(self, **kwargs):
72 eng = sa.create_engine(self.url(), **self.engine_options(**kwargs))
73 # setattr(eng, '_connection_cls', connection.Object)
74 return eng
76 def engine_options(self, **kwargs):
77 if self.root.app.developer_option('db.engine_echo'):
78 kwargs.setdefault('echo', True)
79 kwargs.setdefault('echo_pool', True)
81 if self.cfg('withPool') is False:
82 kwargs.setdefault('poolclass', sa.NullPool)
83 return kwargs
85 pool = self.cfg('pool') or {}
86 p = pool.get('disabled')
87 if p is True:
88 kwargs.setdefault('poolclass', sa.NullPool)
89 return kwargs
91 p = pool.get('pre_ping')
92 if p is True:
93 kwargs.setdefault('pool_pre_ping', True)
94 p = pool.get('size')
95 if isinstance(p, int):
96 kwargs.setdefault('pool_size', p)
97 p = pool.get('recycle')
98 if isinstance(p, int):
99 kwargs.setdefault('pool_recycle', p)
100 p = pool.get('timeout')
101 if isinstance(p, int):
102 kwargs.setdefault('pool_timeout', p)
104 return kwargs
106 def inspect_schema(self, schema, options=None):
107 if options and options.refresh:
108 self.saMetaMap.pop(schema, None)
110 if schema in self.saMetaMap:
111 return
113 def _load():
114 md = sa.MetaData(schema=schema)
116 # introspecting the whole schema is generally faster
117 # but what if we only need a single table from a big schema?
118 # @TODO add options for reflection
120 gws.debug.time_start(f'AUTOLOAD {self.uid=} {schema=}')
121 with self.begin() as conn:
122 md.reflect(conn.saConn, schema, resolve_fks=False, views=True)
123 gws.debug.time_end()
124 return md
126 life_time = self.cfg('schemaCacheLifeTime', 0)
127 if options and options.cacheLifeTime is not None:
128 life_time = options.cacheLifeTime
129 if not life_time:
130 self.saMetaMap[schema] = _load()
131 else:
132 self.saMetaMap[schema] = gws.u.get_cached_object(f'database_metadata_schema_{schema}', life_time, _load)
134 @contextlib.contextmanager
135 def begin(self, nested=False):
136 conns = _connections()
137 sa_conn = conns.get(self.uid)
139 if sa_conn is not None:
140 if not nested:
141 yield connection.Object(self, sa_conn)
142 return
143 with sa_conn.begin_nested():
144 yield connection.Object(self, sa_conn)
145 return
147 sa_conn = self.engine().connect()
148 conns[self.uid] = sa_conn
149 try:
150 with sa_conn.begin():
151 yield connection.Object(self, sa_conn)
152 finally:
153 conns.pop(self.uid, None)
154 sa_conn.close()
156 def connect(self):
157 return self.begin()
159 @contextlib.contextmanager
160 def autocommit_connection(self):
161 if self.uid in _connections():
162 raise gws.Error(f'db {self.uid!r}: autocommit_connection inside a transaction')
163 sa_conn = self.engine().connect().execution_options(isolation_level='AUTOCOMMIT')
164 try:
165 yield connection.Object(self, sa_conn)
166 finally:
167 sa_conn.close()
169 def _sa_connection(self) -> sa.Connection | None:
170 """Return the open connection of the current thread, if any."""
171 return _connections().get(self.uid)
173 def table(self, table, **kwargs):
174 tab = self._sa_table(table)
175 if tab is None:
176 raise sa.Error(f'table not found: {table!r}')
177 return tab
179 def count(self, table):
180 tab = self._sa_table(table)
181 if tab is None:
182 return 0
183 sql = sa.select(sa.func.count()).select_from(tab)
184 with self.begin() as conn:
185 return conn.fetch_int(sql)
187 def has_schema(self, schema):
188 return schema in self.schema_names()
190 def schema_names(self):
191 inspector = sa.inspect(self.engine())
192 return inspector.get_schema_names()
194 def has_table(self, table_name: str):
195 tab = self._sa_table(table_name)
196 return tab is not None
198 def _sa_table(self, tab_or_name) -> sa.Table | None:
199 """Return a reflected table by name, or ``None`` if the table does not exist."""
200 if isinstance(tab_or_name, sa.Table):
201 return tab_or_name
202 schema, name = self.split_table_name(tab_or_name)
203 self.inspect_schema(schema)
204 # see _get_table_key in sqlalchemy/sql/schema.py
205 table_key = schema + '.' + name
206 sm = self.saMetaMap.get(schema)
207 if sm is None:
208 raise sa.Error(f'schema {schema!r} not found')
209 return sm.tables.get(table_key)
211 def column(self, table, column_name):
212 tab = self.table(table)
213 try:
214 return tab.columns[column_name]
215 except KeyError:
216 raise sa.Error(f'column {str(table)}.{column_name!r} not found')
218 def has_column(self, table, column_name):
219 tab = self._sa_table(table)
220 return tab is not None and column_name in tab.columns
222 SA_TO_ATTR = {
223 # common: sqlalchemy.sql.sqltypes
224 'BIGINT': gws.AttributeType.int,
225 'BOOLEAN': gws.AttributeType.bool,
226 'CHAR': gws.AttributeType.str,
227 'DATE': gws.AttributeType.date,
228 'DOUBLE_PRECISION': gws.AttributeType.float,
229 'INTEGER': gws.AttributeType.int,
230 'NUMERIC': gws.AttributeType.float,
231 'REAL': gws.AttributeType.float,
232 'SMALLINT': gws.AttributeType.int,
233 'TEXT': gws.AttributeType.str,
234 # 'UUID': ...,
235 'VARCHAR': gws.AttributeType.str,
236 # postgres specific: sqlalchemy.dialects.postgresql.types
237 # 'JSON': ...,
238 # 'JSONB': ...,
239 # 'BIT': ...,
240 'BYTEA': gws.AttributeType.bytes,
241 # 'CIDR': ...,
242 # 'INET': ...,
243 # 'MACADDR': ...,
244 # 'MACADDR8': ...,
245 # 'MONEY': ...,
246 'TIME': gws.AttributeType.time,
247 'TIMESTAMP': gws.AttributeType.datetime,
248 }
249 """Attribute types for SQLAlchemy type names."""
251 # @TODO proper support for Z/M geoms
253 SA_TO_GEOM = {
254 'POINT': gws.GeometryType.point,
255 'POINTM': gws.GeometryType.point,
256 'POINTZ': gws.GeometryType.point,
257 'POINTZM': gws.GeometryType.point,
258 'LINESTRING': gws.GeometryType.linestring,
259 'LINESTRINGM': gws.GeometryType.linestring,
260 'LINESTRINGZ': gws.GeometryType.linestring,
261 'LINESTRINGZM': gws.GeometryType.linestring,
262 'POLYGON': gws.GeometryType.polygon,
263 'POLYGONM': gws.GeometryType.polygon,
264 'POLYGONZ': gws.GeometryType.polygon,
265 'POLYGONZM': gws.GeometryType.polygon,
266 'MULTIPOINT': gws.GeometryType.multipoint,
267 'MULTIPOINTM': gws.GeometryType.multipoint,
268 'MULTIPOINTZ': gws.GeometryType.multipoint,
269 'MULTIPOINTZM': gws.GeometryType.multipoint,
270 'MULTILINESTRING': gws.GeometryType.multilinestring,
271 'MULTILINESTRINGM': gws.GeometryType.multilinestring,
272 'MULTILINESTRINGZ': gws.GeometryType.multilinestring,
273 'MULTILINESTRINGZM': gws.GeometryType.multilinestring,
274 'MULTIPOLYGON': gws.GeometryType.multipolygon,
275 # 'GEOMETRYCOLLECTION': gws.GeometryType.geometrycollection,
276 # 'CURVE': gws.GeometryType.curve,
277 }
278 """Geometry types for database geometry type names."""
280 UNKNOWN_TYPE = gws.AttributeType.str
281 """Attribute type for columns of unknown types."""
282 UNKNOWN_ARRAY_TYPE = gws.AttributeType.strlist
283 """Attribute type for array columns of unknown item types."""
285 def describe(self, table):
286 tab = self._sa_table(table)
287 if tab is None:
288 raise sa.Error(f'table not found: {table!r}')
290 schema = tab.schema
291 name = tab.name
293 desc = gws.DataSetDescription(
294 columns=[],
295 columnMap={},
296 fullName=self.join_table_name(schema or '', name),
297 geometryName='',
298 geometrySrid=0,
299 geometryType='',
300 name=name,
301 schema=schema,
302 )
304 for n, sa_col in enumerate(cast(list[sa.Column], tab.columns)):
305 col = self.describe_column(table, sa_col.name)
306 col.columnIndex = n
307 desc.columns.append(col)
308 desc.columnMap[col.name] = col
310 for col in desc.columns:
311 if col.geometryType:
312 desc.geometryName = col.name
313 desc.geometryType = col.geometryType
314 desc.geometrySrid = col.geometrySrid
315 break
317 return desc
319 def describe_column(self, table, column_name):
320 sa_col = self.column(table, column_name)
322 col = gws.ColumnDescription(
323 columnIndex=0,
324 comment=str(sa_col.comment or ''),
325 geometrySrid=0,
326 geometryType='',
327 isAutoincrement=bool(sa_col.autoincrement),
328 isNullable=bool(sa_col.nullable),
329 isPrimaryKey=bool(sa_col.primary_key),
330 isUnique=bool(sa_col.unique),
331 hasDefault=sa_col.server_default is not None,
332 name=str(sa_col.name),
333 nativeType='',
334 type='',
335 )
337 col.nativeType = type(sa_col.type).__name__.upper()
338 col.type = self.SA_TO_ATTR.get(col.nativeType, self.UNKNOWN_TYPE)
340 return col
343##