Coverage for python/lsst/daf/butler/registry/databases/postgresql.py: 25%
264 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-10 08:54 +0000
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
27from __future__ import annotations
29from sqlalchemy.sql.expression import ColumnElement as ColumnElement
31from ... import ddl, time_utils
33__all__ = ["PostgresqlDatabase"]
35import re
36import uuid
37from collections.abc import Callable, Iterable, Iterator, Mapping
38from contextlib import closing, contextmanager
39from typing import Any, TypeVar, cast
41import sqlalchemy
42import sqlalchemy.dialects.postgresql
43from sqlalchemy import sql
44from sqlalchemy.dialects.postgresql import Range
46from ..._named import NamedValueAbstractSet
47from ..._timespan import Timespan
48from ...timespan_database_representation import TimespanDatabaseRepresentation
49from ..interfaces import Database, DatabaseMetadata
51_T = TypeVar("_T")
54class PostgresqlDatabase(Database):
55 """An implementation of the `Database` interface for PostgreSQL.
57 Parameters
58 ----------
59 engine : `sqlalchemy.engine.Engine`
60 Engine to use for this connection.
61 origin : `int`
62 An integer ID that should be used as the default for any datasets,
63 quanta, or other entities that use a (autoincrement, origin) compound
64 primary key.
65 namespace : `str`, optional
66 The namespace (schema) this database is associated with. If `None`,
67 the default schema for the connection is used (which may be `None`).
68 writeable : `bool`, optional
69 If `True`, allow write operations on the database, including
70 ``CREATE TABLE``.
71 allow_temporary_tables : `bool`, optional
72 If `True`, database operations will be allowed to use temporary tables.
73 If `False`, other SQL constructs will be used instead of temporary
74 tables when possible.
76 Notes
77 -----
78 Running the tests for this class requires the ``testing.postgresql`` be
79 installed, which we assume indicates that a PostgreSQL server is installed
80 and can be run locally in userspace.
82 Some functionality provided by this class (and used by `Registry`) requires
83 the ``btree_gist`` PostgreSQL server extension to be installed an enabled
84 on the database being connected to; this is checked at connection time.
85 """
87 def __init__(
88 self,
89 *,
90 engine: sqlalchemy.engine.Engine,
91 origin: int,
92 namespace: str | None = None,
93 writeable: bool = True,
94 allow_temporary_tables: bool = True,
95 ):
96 with engine.connect() as connection:
97 query = sql.select(sql.func.current_database())
98 dbname = cast(str, connection.execute(query).scalar())
99 if namespace is None:
100 query = sql.select(sql.func.current_schema())
101 namespace = connection.execute(query).scalar()
102 query_text = "SELECT COUNT(*) FROM pg_extension WHERE extname='btree_gist';"
103 if not connection.execute(sqlalchemy.text(query_text)).scalar():
104 raise RuntimeError(
105 "The Butler PostgreSQL backend requires the btree_gist extension. "
106 "As extensions are enabled per-database, this may require an administrator to run "
107 "`CREATE EXTENSION btree_gist;` in a database before a butler client for it is "
108 " initialized."
109 )
111 pg_version = get_postgres_server_version(connection)
112 self._init(
113 engine=engine,
114 origin=origin,
115 namespace=namespace,
116 writeable=writeable,
117 dbname=dbname,
118 metadata=None,
119 pg_version=pg_version,
120 allow_temporary_tables=allow_temporary_tables,
121 )
123 def _init(
124 self,
125 *,
126 engine: sqlalchemy.engine.Engine,
127 origin: int,
128 namespace: str | None = None,
129 writeable: bool = True,
130 dbname: str,
131 metadata: DatabaseMetadata | None,
132 pg_version: tuple[int, int],
133 allow_temporary_tables: bool = True,
134 ) -> None:
135 # Initialization logic shared between ``__init__`` and ``clone``.
136 super().__init__(
137 origin=origin,
138 engine=engine,
139 namespace=namespace,
140 metadata=metadata,
141 allow_temporary_tables=allow_temporary_tables,
142 )
143 self._writeable = writeable
144 self.dbname = dbname
145 self._pg_version = pg_version
147 def clone(self) -> PostgresqlDatabase:
148 clone = self.__new__(type(self))
149 clone._init(
150 origin=self.origin,
151 engine=self._engine,
152 namespace=self.namespace,
153 writeable=self._writeable,
154 dbname=self.dbname,
155 metadata=self._metadata,
156 pg_version=self._pg_version,
157 allow_temporary_tables=self._allow_temporary_tables,
158 )
159 return clone
161 @classmethod
162 def makeEngine(
163 cls, uri: str | sqlalchemy.engine.URL, *, writeable: bool = True
164 ) -> sqlalchemy.engine.Engine:
165 return sqlalchemy.engine.create_engine(
166 uri,
167 # Prevent stale database connections from throwing exeptions, at
168 # the expense of a round trip to the database server each time we
169 # check out a session. Many services using the Butler operate in
170 # networks where connections are dropped when idle for some time.
171 pool_pre_ping=True,
172 # This engine and database connection pool can be shared between
173 # multiple Butler instances created via Butler.clone() or
174 # LabeledButlerFactory, and typically these will be used from
175 # multiple threads simultaneously. So we need to configure
176 # SQLAlchemy to pool connections for multi-threaded usage.
177 #
178 # This pool size was chosen to work well for services using
179 # FastAPI. FastAPI uses a thread pool of 40 by default, so this
180 # gives us a connection for each thread in the pool. Because Butler
181 # is currently sync-only, we won't ever be executing more queries
182 # than we have threads.
183 #
184 # Connections are only created as they are needed, so in typical
185 # single-threaded Butler use only one connection will ever be
186 # created. Services with low peak concurrency may never create this
187 # many connections.
188 #
189 # The main Butler databases at SLAC are run behind pgbouncer, so we
190 # can support a larger number of simultaneous connections than if
191 # we were connecting directly to Postgres.
192 #
193 # See
194 # https://docs.sqlalchemy.org/en/20/core/pooling.html#sqlalchemy.pool.QueuePool.__init__
195 # for more information on the behavior of this parameter.
196 pool_size=40,
197 # If we are experiencing heavy enough load that we overflow the
198 # connection pool, it will be harmful to start creating extra
199 # connections that we disconnect immediately after use.
200 # Connecting from scratch is fairly expensive, which is why we have
201 # a pool in the first place.
202 max_overflow=0,
203 # If the pool is full, this is the maximum number of seconds we
204 # will wait for a connection to become available before giving up.
205 pool_timeout=60,
206 # In combination with pool_pre_ping, prevent SQLAlchemy from
207 # unnecessarily reviving pooled connections that have gone stale.
208 # Setting this to true makes it always re-use the most recent
209 # known-good connection when possible, instead of cycling to other
210 # connections in the pool that we may no longer need.
211 pool_use_lifo=True,
212 )
214 @classmethod
215 def fromEngine(
216 cls,
217 engine: sqlalchemy.engine.Engine,
218 *,
219 origin: int,
220 namespace: str | None = None,
221 writeable: bool = True,
222 ) -> Database:
223 return cls(engine=engine, origin=origin, namespace=namespace, writeable=writeable)
225 @contextmanager
226 def _transaction(
227 self,
228 *,
229 interrupting: bool = False,
230 savepoint: bool = False,
231 lock: Iterable[sqlalchemy.schema.Table] = (),
232 for_temp_tables: bool = False,
233 ) -> Iterator[tuple[bool, sqlalchemy.engine.Connection]]:
234 with super()._transaction(interrupting=interrupting, savepoint=savepoint, lock=lock) as (
235 is_new,
236 connection,
237 ):
238 if is_new:
239 # pgbouncer with transaction-level pooling (which we aim to
240 # support) says that SET cannot be used, except for a list of
241 # "Startup parameters" that includes "timezone" (see
242 # https://www.pgbouncer.org/features.html#fnref:0). But I
243 # don't see "timezone" in PostgreSQL's list of parameters
244 # passed when creating a new connection
245 # (https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-PARAMKEYWORDS).
246 # Given that the pgbouncer docs say, "PgBouncer detects their
247 # changes and so it can guarantee they remain consistent for
248 # the client", I assume we can use "SET TIMESPAN" and pgbouncer
249 # will take care of clients that share connections being set
250 # consistently. And if that assumption is wrong, we should
251 # still probably be okay, since all clients should be Butler
252 # clients, and they'll all be setting the same thing.
253 #
254 # The "SET TRANSACTION READ ONLY" should also be safe, because
255 # it only ever acts on the current transaction; I think it's
256 # not included in pgbouncer's declaration that SET is
257 # incompatible with transaction-level pooling because
258 # PostgreSQL actually considers SET TRANSACTION to be a
259 # fundamentally different statement from SET (they have their
260 # own distinct doc pages, at least).
261 with closing(connection.connection.cursor()) as cursor:
262 # PostgreSQL permits writing to temporary tables inside
263 # read-only transactions, but it doesn't permit creating
264 # them.
265 if not (self.isWriteable() or for_temp_tables):
266 cursor.execute("SET TRANSACTION READ ONLY")
267 # Make timestamps UTC, because we didn't use TIMESTAMPZ
268 # for the column type. When we can tolerate a schema
269 # change, we should change that type and remove this
270 # line.
271 cursor.execute("SET TIME ZONE 0")
272 # Using server-side cursors with complex queries frequently
273 # generates suboptimal query plan, setting
274 # cursor_tuple_fraction=1 helps for those cases.
275 cursor.execute("SET cursor_tuple_fraction = 1")
276 yield is_new, connection
278 @contextmanager
279 def temporary_table(
280 self, spec: ddl.TableSpec, name: str | None = None
281 ) -> Iterator[sqlalchemy.schema.Table]:
282 # Docstring inherited.
283 with self.transaction(for_temp_tables=True), super().temporary_table(spec, name) as table:
284 yield table
286 def _lockTables(
287 self, connection: sqlalchemy.engine.Connection, tables: Iterable[sqlalchemy.schema.Table] = ()
288 ) -> None:
289 # Docstring inherited.
290 for table in tables:
291 connection.execute(sqlalchemy.text(f"LOCK TABLE {table.key} IN EXCLUSIVE MODE"))
293 def isWriteable(self) -> bool:
294 return self._writeable
296 def __str__(self) -> str:
297 return f"PostgreSQL@{self.dbname}:{self.namespace}"
299 def shrinkDatabaseEntityName(self, original: str) -> str:
300 return self.name_shrinker.shrink(original)
302 def expandDatabaseEntityName(self, shrunk: str) -> str:
303 return self.name_shrinker.expand(shrunk)
305 def _convertExclusionConstraintSpec(
306 self,
307 table: str,
308 spec: tuple[str | type[TimespanDatabaseRepresentation], ...],
309 metadata: sqlalchemy.MetaData,
310 ) -> sqlalchemy.schema.Constraint:
311 # Docstring inherited.
312 args: list[tuple[sqlalchemy.schema.Column, str]] = []
313 names = ["excl"]
314 for item in spec:
315 if isinstance(item, str):
316 args.append((sqlalchemy.schema.Column(item), "="))
317 names.append(item)
318 elif issubclass(item, TimespanDatabaseRepresentation):
319 assert item is self.getTimespanRepresentation()
320 args.append((sqlalchemy.schema.Column(TimespanDatabaseRepresentation.NAME), "&&"))
321 names.append(TimespanDatabaseRepresentation.NAME)
322 return sqlalchemy.dialects.postgresql.ExcludeConstraint(
323 *args,
324 name=self.shrinkDatabaseEntityName("_".join(names)),
325 )
327 def _make_temporary_table(
328 self,
329 connection: sqlalchemy.engine.Connection,
330 spec: ddl.TableSpec,
331 name: str | None = None,
332 **kwargs: Any,
333 ) -> sqlalchemy.schema.Table:
334 # Docstring inherited
335 # Adding ON COMMIT DROP here is really quite defensive: we already
336 # manually drop the table at the end of the temporary_table context
337 # manager, and that will usually happen first. But this will guarantee
338 # that we drop the table at the end of the transaction even if the
339 # connection lasts longer, and that's good citizenship when connections
340 # may be multiplexed by e.g. pgbouncer.
341 return super()._make_temporary_table(connection, spec, name, postgresql_on_commit="DROP", **kwargs)
343 @classmethod
344 def getTimespanRepresentation(cls) -> type[TimespanDatabaseRepresentation]:
345 # Docstring inherited.
346 return _RangeTimespanRepresentation
348 def replace(self, table: sqlalchemy.schema.Table, *rows: dict) -> None:
349 self.assertTableWriteable(table, f"Cannot replace into read-only table {table}.")
350 if not rows:
351 return
352 # This uses special support for UPSERT in PostgreSQL backend:
353 # https://docs.sqlalchemy.org/en/13/dialects/postgresql.html#insert-on-conflict-upsert
354 query = sqlalchemy.dialects.postgresql.dml.insert(table)
355 # In the SET clause assign all columns using special `excluded`
356 # pseudo-table. If some column in the table does not appear in the
357 # INSERT list this will set it to NULL.
358 excluded = query.excluded
359 data = {
360 column.name: getattr(excluded, column.name)
361 for column in table.columns
362 if column.name not in table.primary_key
363 }
364 if not data:
365 self.ensure(table, *rows)
366 return
367 query = query.on_conflict_do_update(constraint=table.primary_key, set_=data)
368 with self._transaction() as (_, connection):
369 connection.execute(query, rows)
371 def ensure(self, table: sqlalchemy.schema.Table, *rows: dict, primary_key_only: bool = False) -> int:
372 # Docstring inherited.
373 self.assertTableWriteable(table, f"Cannot ensure into read-only table {table}.")
374 if not rows:
375 return 0
376 # Like `replace`, this uses UPSERT.
377 base_insert = sqlalchemy.dialects.postgresql.dml.insert(table)
378 if primary_key_only:
379 query = base_insert.on_conflict_do_nothing(constraint=table.primary_key)
380 else:
381 query = base_insert.on_conflict_do_nothing()
382 with self._transaction() as (_, connection):
383 return connection.execution_options(preserve_rowcount=True).execute(query, rows).rowcount
385 def constant_rows(
386 self,
387 fields: NamedValueAbstractSet[ddl.FieldSpec],
388 *rows: dict,
389 name: str | None = None,
390 ) -> sqlalchemy.sql.FromClause:
391 # Docstring inherited.
392 if name is None:
393 name = f"tmp_{uuid.uuid4().hex}"
395 # Convert rows to columns.
396 column_values = {name: [row[name] for row in rows] for name in fields.names}
398 columns: list[sqlalchemy.schema.Column] = []
399 column_data = []
400 for field in fields:
401 array_type: Any = sqlalchemy.dialects.postgresql.ARRAY(field.dtype)
402 column_data.append(sqlalchemy.literal(column_values[field.name], type_=array_type))
403 columns.append(sqlalchemy.schema.Column(field.name, type_=field.dtype))
405 unnest_expr = sqlalchemy.func.unnest(*column_data).table_valued(*columns).render_derived(name=name)
406 return unnest_expr
408 @property
409 def has_distinct_on(self) -> bool:
410 # Docstring inherited.
411 return True
413 @property
414 def has_any_aggregate(self) -> bool:
415 # Docstring inherited.
416 return self._pg_version >= (16, 0)
418 def apply_any_aggregate(self, column: sqlalchemy.ColumnElement[Any]) -> sqlalchemy.ColumnElement[Any]:
419 # Docstring inherited
421 # The cast is required to prevent sqlalchemy from forgetting the type
422 # of the initial column. Without the cast, for example, Base64Region
423 # would become a String column in the output.
424 return sqlalchemy.cast(sqlalchemy.func.any_value(column), column.type)
426 def glob_expression(
427 self, expression: sqlalchemy.ColumnElement[Any], pattern: str
428 ) -> sqlalchemy.ColumnElement[bool]:
429 # Docstring inherited.
430 # With Postgres we can use LIKE as it is case-sensitive, we only need
431 # to translate our simple glob language to LIKE patterns, escaping few
432 # things. Also we can try to optimize special cases when we can use
433 # ``starts_with()`` function or "*" pattern. ``starts_with()`` in
434 # Postgres14 does not use btree index, and LIKE with a constant prefix
435 # is using index, so we may not need ``starts_with()`` after all.
436 if pattern == "*":
437 # Simple case.
438 return sqlalchemy.literal(True)
440 def _escape(pattern: str) -> str:
441 """Transform glob pattern string into LIKE pattern."""
442 tokens = []
443 escape = ""
444 for char in pattern:
445 if not escape:
446 if char == "\\":
447 escape = char
448 elif char in "_%":
449 # Escape special SQL characters.
450 tokens.append(f"\\{char}")
451 elif char == "*":
452 tokens.append("%")
453 elif char == "?":
454 tokens.append("_")
455 else:
456 tokens.append(char)
457 else:
458 if char == "\\":
459 # Two backslashes become one.
460 tokens.append(char)
461 elif char in "*?":
462 # Handle escaped wildcards.
463 tokens.append(f"{char}")
464 else:
465 # Copy both backslash and character.
466 tokens.append(escape)
467 tokens.append(char)
468 escape = ""
469 # Trailing backslash added to pattern.
470 tokens.append(escape)
471 return "".join(tokens)
473 pattern = _escape(pattern)
474 return expression.op("LIKE")(sqlalchemy.literal(pattern))
476 def make_in_array_constraint(
477 self, column: sqlalchemy.ColumnElement[_T], values: Iterable[_T]
478 ) -> sqlalchemy.ColumnElement[bool]:
479 array_type = sqlalchemy.dialects.postgresql.ARRAY(column.type)
480 return column == sqlalchemy.any_(sqlalchemy.literal(list(values), array_type))
483class _RangeTimespanType(sqlalchemy.TypeDecorator):
484 """A single-column `Timespan` representation usable only with
485 PostgreSQL.
487 This type should be able to take advantage of PostgreSQL's built-in
488 range operators, and the indexing and EXCLUSION table constraints built
489 off of them.
490 """
492 impl = sqlalchemy.dialects.postgresql.INT8RANGE
494 cache_ok = True
496 def process_bind_param(self, value: Timespan | None, dialect: sqlalchemy.engine.Dialect) -> Range | None:
497 if value is None:
498 return None
499 if not isinstance(value, Timespan):
500 raise TypeError(f"Unsupported type: {type(value)}, expected Timespan.")
501 if value.isEmpty():
502 return Range(empty=True)
503 else:
504 converter = time_utils.TimeConverter()
505 assert value.nsec[0] >= converter.min_nsec, "Guaranteed by Timespan.__init__."
506 assert value.nsec[1] <= converter.max_nsec, "Guaranteed by Timespan.__init__."
507 lower = None if value.nsec[0] == converter.min_nsec else value.nsec[0]
508 upper = None if value.nsec[1] == converter.max_nsec else value.nsec[1]
509 return Range(lower=lower, upper=upper)
511 def process_result_value(
512 self, value: Range | None, dialect: sqlalchemy.engine.Dialect
513 ) -> Timespan | None:
514 if value is None:
515 return None
516 if value.isempty:
517 return Timespan.makeEmpty()
518 converter = time_utils.TimeConverter()
519 begin_nsec = converter.min_nsec if value.lower is None else value.lower
520 end_nsec = converter.max_nsec if value.upper is None else value.upper
521 return Timespan(begin=None, end=None, _nsec=(begin_nsec, end_nsec))
524class _RangeTimespanRepresentation(TimespanDatabaseRepresentation):
525 """An implementation of `TimespanDatabaseRepresentation` that uses
526 `_RangeTimespanType` to store a timespan in a single
527 PostgreSQL-specific field.
529 Parameters
530 ----------
531 column : `sqlalchemy.sql.ColumnElement`
532 SQLAlchemy object representing the column.
533 """
535 def __init__(self, column: sqlalchemy.sql.ColumnElement, name: str):
536 self.column = column
537 self._name = name
539 __slots__ = ("column", "_name")
541 @classmethod
542 def makeFieldSpecs(
543 cls, nullable: bool, name: str | None = None, **kwargs: Any
544 ) -> tuple[ddl.FieldSpec, ...]:
545 # Docstring inherited.
546 if name is None:
547 name = cls.NAME
548 return (
549 ddl.FieldSpec(
550 name,
551 dtype=_RangeTimespanType,
552 nullable=nullable,
553 default=(None if nullable else sqlalchemy.sql.text("'(,)'::int8range")),
554 **kwargs,
555 ),
556 )
558 @classmethod
559 def getFieldNames(cls, name: str | None = None) -> tuple[str, ...]:
560 # Docstring inherited.
561 if name is None:
562 name = cls.NAME
563 return (name,)
565 @classmethod
566 def update(
567 cls, extent: Timespan | None, name: str | None = None, result: dict[str, Any] | None = None
568 ) -> dict[str, Any]:
569 # Docstring inherited.
570 if name is None:
571 name = cls.NAME
572 if result is None:
573 result = {}
574 result[name] = extent
575 return result
577 @classmethod
578 def extract(cls, mapping: Mapping[str, Any], name: str | None = None) -> Timespan | None:
579 # Docstring inherited.
580 if name is None:
581 name = cls.NAME
582 return mapping[name]
584 @classmethod
585 def fromLiteral(cls, timespan: Timespan | None) -> _RangeTimespanRepresentation:
586 # Docstring inherited.
587 if timespan is None:
588 # Cast NULL to an expected type, helps Postgres to figure out
589 # column type when doing UNION.
590 return cls(
591 column=sqlalchemy.func.cast(sqlalchemy.sql.null(), type_=_RangeTimespanType), name=cls.NAME
592 )
593 return cls(
594 column=sqlalchemy.sql.cast(
595 sqlalchemy.sql.literal(timespan, type_=_RangeTimespanType), type_=_RangeTimespanType
596 ),
597 name=cls.NAME,
598 )
600 @classmethod
601 def from_columns(
602 cls, columns: sqlalchemy.sql.ColumnCollection, name: str | None = None
603 ) -> _RangeTimespanRepresentation:
604 # Docstring inherited.
605 if name is None:
606 name = cls.NAME
607 return cls(columns[name], name)
609 @property
610 def name(self) -> str:
611 # Docstring inherited.
612 return self._name
614 def isNull(self) -> sqlalchemy.sql.ColumnElement:
615 # Docstring inherited.
616 return self.column.is_(None)
618 def isEmpty(self) -> sqlalchemy.sql.ColumnElement:
619 # Docstring inherited
620 return sqlalchemy.sql.func.isempty(self.column)
622 def __lt__(
623 self, other: _RangeTimespanRepresentation | sqlalchemy.sql.ColumnElement
624 ) -> sqlalchemy.sql.ColumnElement:
625 # Docstring inherited.
626 if isinstance(other, sqlalchemy.sql.ColumnElement):
627 return sqlalchemy.sql.and_(
628 sqlalchemy.sql.not_(sqlalchemy.sql.func.upper_inf(self.column)),
629 sqlalchemy.sql.not_(sqlalchemy.sql.func.isempty(self.column)),
630 sqlalchemy.sql.func.upper(self.column) <= other,
631 )
632 else:
633 return self.column << other.column
635 def __gt__(
636 self, other: _RangeTimespanRepresentation | sqlalchemy.sql.ColumnElement
637 ) -> sqlalchemy.sql.ColumnElement:
638 # Docstring inherited.
639 if isinstance(other, sqlalchemy.sql.ColumnElement):
640 return sqlalchemy.sql.and_(
641 sqlalchemy.sql.not_(sqlalchemy.sql.func.lower_inf(self.column)),
642 sqlalchemy.sql.not_(sqlalchemy.sql.func.isempty(self.column)),
643 sqlalchemy.sql.func.lower(self.column) > other,
644 )
645 else:
646 return self.column >> other.column
648 def overlaps(
649 self, other: _RangeTimespanRepresentation | sqlalchemy.sql.ColumnElement
650 ) -> sqlalchemy.sql.ColumnElement:
651 # Docstring inherited.
652 if not isinstance(other, _RangeTimespanRepresentation):
653 return self.contains(other)
654 return self.column.overlaps(other.column)
656 def contains(
657 self, other: _RangeTimespanRepresentation | sqlalchemy.sql.ColumnElement
658 ) -> sqlalchemy.sql.ColumnElement:
659 # Docstring inherited
660 if isinstance(other, _RangeTimespanRepresentation):
661 return self.column.contains(other.column)
662 else:
663 return self.column.contains(other)
665 def lower(self) -> sqlalchemy.sql.ColumnElement:
666 # Docstring inherited.
667 return sqlalchemy.sql.functions.coalesce(
668 sqlalchemy.sql.func.lower(self.column), sqlalchemy.sql.literal(0)
669 )
671 def upper(self) -> sqlalchemy.sql.ColumnElement:
672 # Docstring inherited.
673 return sqlalchemy.sql.functions.coalesce(
674 sqlalchemy.sql.func.upper(self.column), sqlalchemy.sql.literal(0)
675 )
677 def flatten(self, name: str | None = None) -> tuple[sqlalchemy.sql.ColumnElement]:
678 # Docstring inherited.
679 if name is None:
680 return (self.column,)
681 else:
682 return (self.column.label(name),)
684 def apply_any_aggregate(
685 self, func: Callable[[ColumnElement[Any]], ColumnElement[Any]]
686 ) -> TimespanDatabaseRepresentation:
687 # Docstring inherited.
688 return _RangeTimespanRepresentation(func(self.column), self.name)
691def get_postgres_server_version(connection: sqlalchemy.Connection) -> tuple[int, int]: # numpydoc ignore=PR01
692 """Return the postgres DB server version as a tuple
693 (major version, minor version).
694 """
695 raw_pg_version = connection.execute(sqlalchemy.text("SHOW server_version")).scalar()
696 _SERVER_VERSION_REGEX = re.compile(r"(?P<major>\d+)\.(?P<minor>\d+)")
697 if raw_pg_version is not None and (m := _SERVER_VERSION_REGEX.search(raw_pg_version)):
698 return (int(m.group("major")), int(m.group("minor")))
699 else:
700 raise RuntimeError("Failed to get PostgreSQL server version.")