Coverage for python/lsst/daf/butler/registry/opaque.py: 97%
88 statements
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-06 09:46 +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/>.
28"""The default concrete implementations of the classes that manage
29opaque tables for `Registry`.
30"""
32from __future__ import annotations
34__all__ = ["ByNameOpaqueTableStorage", "ByNameOpaqueTableStorageManager"]
36import itertools
37from collections.abc import Iterable, Iterator, Sequence
38from typing import TYPE_CHECKING, Any, ClassVar
40import sqlalchemy
41import sqlalchemy.exc
43from .._utilities.thread_safe_cache import ThreadSafeCache
44from ..ddl import FieldSpec, TableSpec
45from ._exceptions import ConflictingDefinitionError
46from .interfaces import (
47 Database,
48 OpaqueTableStorage,
49 OpaqueTableStorageManager,
50 StaticTablesContext,
51 VersionTuple,
52)
54if TYPE_CHECKING:
55 from ..datastore import DatastoreTransaction
57# This has to be updated on every schema change
58_VERSION = VersionTuple(0, 2, 0)
61class ByNameOpaqueTableStorage(OpaqueTableStorage):
62 """An implementation of `OpaqueTableStorage` that simply creates a true
63 table for each different named opaque logical table.
65 A `ByNameOpaqueTableStorageManager` instance should always be used to
66 construct and manage instances of this class.
68 Parameters
69 ----------
70 db : `Database`
71 Database engine interface for the namespace in which this table lives.
72 name : `str`
73 Name of the logical table (also used as the name of the actual table).
74 table : `sqlalchemy.schema.Table`
75 SQLAlchemy representation of the table, which must have already been
76 created in the namespace managed by ``db`` (this is the responsibility
77 of `ByNameOpaqueTableStorageManager`).
78 """
80 def __init__(self, *, db: Database, name: str, table: sqlalchemy.schema.Table):
81 super().__init__(name=name)
82 self._db = db
83 self._table = table
85 def insert(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
86 # Docstring inherited from OpaqueTableStorage.
87 # The provided transaction object can be ignored since we rely on
88 # the database itself providing any rollback functionality.
89 try:
90 self._db.insert(self._table, *data)
91 except sqlalchemy.exc.IntegrityError as exc:
92 raise ConflictingDefinitionError(
93 f"Table {self._table.name} already contains matching record(s)."
94 ) from exc
96 def ensure(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
97 # Docstring inherited from OpaqueTableStorage.
98 # The provided transaction object can be ignored since we rely on
99 # the database itself providing any rollback functionality.
100 self._db.ensure(self._table, *data)
102 def replace(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
103 # Docstring inherited from OpaqueTableStorage.
104 # The provided transaction object can be ignored since we rely on
105 # the database itself providing any rollback functionality.
106 self._db.replace(self._table, *data)
108 def fetch(
109 self,
110 **where: Any,
111 ) -> Iterator[sqlalchemy.RowMapping]:
112 # Docstring inherited from OpaqueTableStorage.
113 for batch in self.fetch_batches(**where):
114 yield from batch
116 def fetch_batches(
117 self,
118 **where: Any,
119 ) -> Iterator[Sequence[sqlalchemy.RowMapping]]:
120 def _batch_in_clause(
121 column: sqlalchemy.schema.Column, values: Iterable[Any]
122 ) -> Iterator[sqlalchemy.sql.expression.ClauseElement]:
123 """Split one long IN clause into a series of shorter ones."""
124 in_limit = 1000
125 # We have to remove possible duplicates from values; and in many
126 # cases it should be helpful to order the items in the clause.
127 values = sorted(set(values))
128 for iposn in range(0, len(values), in_limit):
129 in_clause = self._db.make_in_array_constraint(column, values[iposn : iposn + in_limit])
130 yield in_clause
132 def _batch_in_clauses(**where: Any) -> Iterator[sqlalchemy.sql.expression.ColumnElement]:
133 """Generate a sequence of WHERE clauses with a limited number of
134 items in IN clauses.
135 """
136 batches: list[Iterable[Any]] = []
137 for k, v in where.items():
138 column = self._table.columns[k]
139 if isinstance(v, list | tuple | set):
140 batches.append(_batch_in_clause(column, v))
141 else:
142 if isinstance(v, str) and v.endswith("%"):
143 # Special case prefix queries.
144 batches.append([column.startswith(v[:-1])])
145 else:
146 # single "batch" for a regular eq operator
147 batches.append([column == v])
149 for clauses in itertools.product(*batches):
150 yield sqlalchemy.sql.and_(*clauses)
152 sql = self._table.select()
153 if where:
154 # Split long IN clauses into shorter batches
155 batched_sql = [sql.where(clause) for clause in _batch_in_clauses(**where)]
156 else:
157 batched_sql = [sql]
158 for sql_batch in batched_sql:
159 with self._db.query(sql_batch) as sql_result:
160 sql_mappings = sql_result.mappings().fetchall()
161 yield sql_mappings
163 def delete(self, columns: Iterable[str], *rows: dict) -> None:
164 # Docstring inherited from OpaqueTableStorage.
165 self._db.delete(self._table, columns, *rows)
168class ByNameOpaqueTableStorageManager(OpaqueTableStorageManager):
169 """An implementation of `OpaqueTableStorageManager` that simply creates a
170 true table for each different named opaque logical table.
172 Instances of this class should generally be constructed via the
173 `initialize` class method instead of invoking ``__init__`` directly.
175 Parameters
176 ----------
177 db : `Database`
178 Database engine interface for the namespace in which this table lives.
179 metaTable : `sqlalchemy.schema.Table`
180 SQLAlchemy representation of the table that records which opaque
181 logical tables exist.
182 tables : `ThreadSafeCache` [`str`, `~sqlalchemy.schema.Table`]
183 Mapping from string to table, to track which tables have already been
184 created. This mapping is shared between cloned instances of this
185 manager.
186 registry_schema_version : `VersionTuple` or `None`, optional
187 Version of registry schema.
188 """
190 def __init__(
191 self,
192 db: Database,
193 metaTable: sqlalchemy.schema.Table,
194 tables: ThreadSafeCache[str, sqlalchemy.schema.Table],
195 registry_schema_version: VersionTuple | None = None,
196 ):
197 super().__init__(registry_schema_version=registry_schema_version)
198 self._db = db
199 self._metaTable = metaTable
200 self._tables = tables
202 def clone(self, db: Database) -> ByNameOpaqueTableStorageManager:
203 return ByNameOpaqueTableStorageManager(
204 db, self._metaTable, self._tables, self._registry_schema_version
205 )
207 _META_TABLE_NAME: ClassVar[str] = "opaque_meta"
209 _META_TABLE_SPEC: ClassVar[TableSpec] = TableSpec(
210 fields=[
211 FieldSpec("table_name", dtype=sqlalchemy.String, length=128, primaryKey=True),
212 ],
213 )
215 @classmethod
216 def initialize(
217 cls, db: Database, context: StaticTablesContext, registry_schema_version: VersionTuple | None = None
218 ) -> OpaqueTableStorageManager:
219 # Docstring inherited from OpaqueTableStorageManager.
220 metaTable = context.addTable(cls._META_TABLE_NAME, cls._META_TABLE_SPEC)
221 return cls(
222 db=db,
223 metaTable=metaTable,
224 tables=ThreadSafeCache(),
225 registry_schema_version=registry_schema_version,
226 )
228 def get(self, name: str) -> OpaqueTableStorage | None:
229 # Docstring inherited from OpaqueTableStorageManager.
230 table = self._tables.get(name)
231 if table is None:
232 return None
233 return ByNameOpaqueTableStorage(name=name, table=table, db=self._db)
235 def register(self, name: str, spec: TableSpec) -> OpaqueTableStorage:
236 # Docstring inherited from OpaqueTableStorageManager.
237 result = self.get(name)
238 if result is None:
239 # Create the table itself. If it already exists but wasn't in
240 # the dict because it was added by another client since this one
241 # was initialized, that's fine.
242 table = self._db.ensureTableExists(name, spec)
243 # Add a row to the meta table so we can find this table in the
244 # future. Also okay if that already exists, so we use sync.
245 self._db.sync(self._metaTable, keys={"table_name": name})
246 self._tables.set_or_get(name, table)
247 result = self.get(name)
248 assert result is not None
249 return result
251 @classmethod
252 def currentVersions(cls) -> list[VersionTuple]:
253 # Docstring inherited from VersionedExtension.
254 return [_VERSION]