Coverage for python/lsst/daf/butler/registry/databases/oracle.py : 23%

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__ = ["OracleDatabase"]
25from contextlib import closing, contextmanager
26from typing import Any, Iterator, Optional
28import sqlalchemy
29import sqlalchemy.ext.compiler
31from ..interfaces import Database, ReadOnlyDatabaseError
32from ..nameShrinker import NameShrinker
35class _Merge(sqlalchemy.sql.expression.Executable, sqlalchemy.sql.ClauseElement):
36 """A SQLAlchemy query that compiles to a MERGE invocation that is the
37 equivalent of PostgreSQL and SQLite's INSERT ... ON CONFLICT REPLACE on the
38 primary key constraint for the table.
39 """
41 def __init__(self, table: sqlalchemy.schema.Table):
42 super().__init__()
43 self.table = table
46# Hard to infer what types these should be from SQLAlchemy docs; just disable
47# static typing by calling everything "Any".
48@sqlalchemy.ext.compiler.compiles(_Merge, "oracle")
49def _merge(merge: Any, compiler: Any, **kwargs: Any) -> Any:
50 """Generate MERGE query for inserting or updating records.
51 """
52 table = merge.table
53 preparer = compiler.preparer
55 allColumns = [col.name for col in table.columns]
56 pkColumns = [col.name for col in table.primary_key]
57 nonPkColumns = [col for col in allColumns if col not in pkColumns]
59 # To properly support type decorators defined in core/ddl.py we need
60 # to pass column type to `bindparam`.
61 selectColumns = [sqlalchemy.sql.bindparam(col.name, type_=col.type).label(col.name)
62 for col in table.columns]
63 selectClause = sqlalchemy.sql.select(selectColumns)
65 tableAlias = table.alias("t")
66 tableAliasText = compiler.process(tableAlias, asfrom=True, **kwargs)
67 selectAlias = selectClause.alias("d")
68 selectAliasText = compiler.process(selectAlias, asfrom=True, **kwargs)
70 condition = sqlalchemy.sql.and_(
71 *[tableAlias.columns[col] == selectAlias.columns[col] for col in pkColumns]
72 )
73 conditionText = compiler.process(condition, **kwargs)
75 query = f"MERGE INTO {tableAliasText}" \
76 f"\nUSING {selectAliasText}" \
77 f"\nON ({conditionText})"
78 updates = []
79 for col in nonPkColumns:
80 src = compiler.process(selectAlias.columns[col], **kwargs)
81 dst = compiler.process(tableAlias.columns[col], **kwargs)
82 updates.append(f"{dst} = {src}")
83 text = ", ".join(updates)
84 query += f"\nWHEN MATCHED THEN UPDATE SET {text}"
86 insertColumns = ", ".join([preparer.format_column(col) for col in table.columns])
87 insertValues = ", ".join([compiler.process(selectAlias.columns[col], **kwargs) for col in allColumns])
89 query += f"\nWHEN NOT MATCHED THEN INSERT ({insertColumns}) VALUES ({insertValues})"
90 return query
93class OracleDatabase(Database):
94 """An implementation of the `Database` interface for Oracle.
96 Parameters
97 ----------
98 connection : `sqlalchemy.engine.Connection`
99 An existing connection created by a previous call to `connect`.
100 origin : `int`
101 An integer ID that should be used as the default for any datasets,
102 quanta, or other entities that use a (autoincrement, origin) compound
103 primary key.
104 namespace : `str`, optional
105 The namespace (schema) this database is associated with. If `None`,
106 the default schema for the connection is used (which may be `None`).
107 writeable : `bool`, optional
108 If `True`, allow write operations on the database, including
109 ``CREATE TABLE``.
110 prefix : `str`, optional
111 Prefix to add to all table names, effectively defining a virtual
112 schema that can coexist with others within the same actual database
113 schema. This prefix must not be used in the un-prefixed names of
114 tables.
116 Notes
117 -----
118 To use a prefix from standardized factory functions like `Database.fromUri`
119 and `Database.fromConnectionStruct`, a '+' character in the namespace will
120 be interpreted as a combination of ``namespace`` (first) and ``prefix``
121 (second). Either may be empty. This does *not* work when constructing
122 an `OracleDatabase` instance directly.
123 """
125 def __init__(self, *, connection: sqlalchemy.engine.Connection, origin: int,
126 namespace: Optional[str] = None, writeable: bool = True, prefix: Optional[str] = None):
127 # Get the schema that was included/implicit in the URI we used to
128 # connect.
129 dbapi = connection.engine.raw_connection()
130 namespace = dbapi.current_schema
131 super().__init__(connection=connection, origin=origin, namespace=namespace)
132 self._writeable = writeable
133 self.dsn = dbapi.dsn
134 self.prefix = prefix
135 self._shrinker = NameShrinker(connection.engine.dialect.max_identifier_length)
137 @classmethod
138 def connect(cls, uri: str, *, writeable: bool = True) -> sqlalchemy.engine.Connection:
139 connection = sqlalchemy.engine.create_engine(uri, pool_size=1, max_identifier_length=128).connect()
140 return connection
142 @classmethod
143 def fromConnection(cls, connection: sqlalchemy.engine.Connection, *, origin: int,
144 namespace: Optional[str] = None, writeable: bool = True) -> Database:
145 prefix: Optional[str]
146 if namespace and "+" in namespace:
147 namespace, prefix = namespace.split("+")
148 if not namespace:
149 namespace = None
150 if not prefix:
151 prefix = None
152 else:
153 prefix = None
154 return cls(connection=connection, origin=origin, writeable=writeable, namespace=namespace,
155 prefix=prefix)
157 @contextmanager
158 def transaction(self, *, interrupting: bool = False) -> Iterator[None]:
159 with super().transaction(interrupting=interrupting):
160 if not self.isWriteable():
161 with closing(self._connection.connection.cursor()) as cursor:
162 cursor.execute("SET TRANSACTION READ ONLY")
163 yield
165 def isWriteable(self) -> bool:
166 return self._writeable
168 def __str__(self) -> str:
169 if self.namespace is None:
170 name = self.dsn
171 else:
172 name = f"{self.dsn:self.namespace}"
173 return f"Oracle@{name}"
175 def shrinkDatabaseEntityName(self, original: str) -> str:
176 return self._shrinker.shrink(original)
178 def expandDatabaseEntityName(self, shrunk: str) -> str:
179 return self._shrinker.expand(shrunk)
181 def _mangleTableName(self, name: str) -> str:
182 if self.prefix is not None and not name.startswith(self.prefix):
183 name = self.prefix + name
184 return name
186 def replace(self, table: sqlalchemy.schema.Table, *rows: dict) -> None:
187 if not self.isWriteable():
188 raise ReadOnlyDatabaseError(f"Attempt to replace into read-only database '{self}'.")
189 if not rows:
190 return
191 self._connection.execute(_Merge(table), *rows)
193 prefix: Optional[str]
194 """A prefix included in all table names to simulate a database namespace
195 (`str` or `None`).
196 """
198 dsn: str
199 """The TNS entry of the database this instance is connected to (`str`).
200 """