Coverage for python/lsst/daf/butler/registry/databases/postgresql.py : 32%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
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 program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
21from __future__ import annotations
23__all__ = ["PostgresqlDatabase"]
25from contextlib import contextmanager, closing
26from typing import Optional
28import sqlalchemy
30from ..interfaces import Database, ReadOnlyDatabaseError
31from ..nameShrinker import NameShrinker
34class PostgresqlDatabase(Database):
35 """An implementation of the `Database` interface for PostgreSQL.
37 Parameters
38 ----------
39 connection : `sqlalchemy.engine.Connection`
40 An existing connection created by a previous call to `connect`.
41 origin : `int`
42 An integer ID that should be used as the default for any datasets,
43 quanta, or other entities that use a (autoincrement, origin) compound
44 primary key.
45 namespace : `str`, optional
46 The namespace (schema) this database is associated with. If `None`,
47 the default schema for the connection is used (which may be `None`).
48 writeable : `bool`, optional
49 If `True`, allow write operations on the database, including
50 ``CREATE TABLE``.
52 Notes
53 -----
54 This currently requires the psycopg2 driver to be used as the backend for
55 SQLAlchemy. Running the tests for this class requires the
56 ``testing.postgresql`` be installed, which we assume indicates that a
57 PostgreSQL server is installed and can be run locally in userspace.
58 """
60 def __init__(self, *, connection: sqlalchemy.engine.Connection, origin: int,
61 namespace: Optional[str] = None, writeable: bool = True):
62 super().__init__(origin=origin, connection=connection, namespace=namespace)
63 dbapi = connection.connection
64 try:
65 dsn = dbapi.get_dsn_parameters()
66 except (AttributeError, KeyError) as err:
67 raise RuntimeError("Only the psycopg2 driver for PostgreSQL is supported.") from err
68 if namespace is None:
69 namespace = connection.execute("SELECT current_schema();").scalar()
70 self.namespace = namespace
71 self.dbname = dsn.get("dbname")
72 self._writeable = writeable
73 self._shrinker = NameShrinker(connection.engine.dialect.max_identifier_length)
75 @classmethod
76 def connect(cls, uri: str, *, writeable: bool = True) -> sqlalchemy.engine.Connection:
77 return sqlalchemy.engine.create_engine(uri, pool_size=1).connect()
79 @classmethod
80 def fromConnection(cls, connection: sqlalchemy.engine.Connection, *, origin: int,
81 namespace: Optional[str] = None, writeable: bool = True) -> Database:
82 return cls(connection=connection, origin=origin, namespace=namespace, writeable=writeable)
84 @contextmanager
85 def transaction(self, *, interrupting: bool = False) -> None:
86 with super().transaction(interrupting=interrupting):
87 if not self.isWriteable():
88 with closing(self._connection.connection.cursor()) as cursor:
89 cursor.execute("SET TRANSACTION READ ONLY")
90 yield
92 def isWriteable(self) -> bool:
93 return self._writeable
95 def __str__(self) -> str:
96 return f"PostgreSQL@{self.dbname}:{self.namespace}"
98 def shrinkDatabaseEntityName(self, original: str) -> str:
99 return self._shrinker.shrink(original)
101 def expandDatabaseEntityName(self, shrunk: str) -> str:
102 return self._shrinker.expand(shrunk)
104 def replace(self, table: sqlalchemy.schema.Table, *rows: dict):
105 if not self.isWriteable():
106 raise ReadOnlyDatabaseError(f"Attempt to replace into read-only database '{self}'.")
107 if not rows:
108 return
109 # This uses special support for UPSERT in PostgreSQL backend:
110 # https://docs.sqlalchemy.org/en/13/dialects/postgresql.html#insert-on-conflict-upsert
111 query = sqlalchemy.dialects.postgresql.dml.insert(table)
112 # In the SET clause assign all columns using special `excluded`
113 # pseudo-table. If some column in the table does not appear in the
114 # INSERT list this will set it to NULL.
115 excluded = query.excluded
116 data = {column.name: getattr(excluded, column.name)
117 for column in table.columns
118 if column.name not in table.primary_key}
119 query = query.on_conflict_do_update(constraint=table.primary_key, set_=data)
120 self._connection.execute(query, *rows)