Hide keyboard shortcuts

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 

22 

23__all__ = ["PostgresqlDatabase"] 

24 

25from contextlib import contextmanager, closing 

26from typing import Optional 

27 

28import sqlalchemy 

29 

30from ..interfaces import Database, ReadOnlyDatabaseError 

31from ..nameShrinker import NameShrinker 

32 

33 

34class PostgresqlDatabase(Database): 

35 """An implementation of the `Database` interface for PostgreSQL. 

36 

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``. 

51 

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 """ 

59 

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) 

74 

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() 

78 

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) 

83 

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 

91 

92 def isWriteable(self) -> bool: 

93 return self._writeable 

94 

95 def __str__(self) -> str: 

96 return f"PostgreSQL@{self.dbname}:{self.namespace}" 

97 

98 def shrinkDatabaseEntityName(self, original: str) -> str: 

99 return self._shrinker.shrink(original) 

100 

101 def expandDatabaseEntityName(self, shrunk: str) -> str: 

102 return self._shrinker.expand(shrunk) 

103 

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)