Coverage for python/lsst/daf/butler/registry/opaque.py: 33%
72 statements
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-16 10:44 +0000
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-16 10:44 +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
38from typing import TYPE_CHECKING, Any, ClassVar
40import sqlalchemy
42from ..ddl import FieldSpec, TableSpec
43from .interfaces import (
44 Database,
45 OpaqueTableStorage,
46 OpaqueTableStorageManager,
47 StaticTablesContext,
48 VersionTuple,
49)
51if TYPE_CHECKING:
52 from ..datastore import DatastoreTransaction
54# This has to be updated on every schema change
55_VERSION = VersionTuple(0, 2, 0)
58class ByNameOpaqueTableStorage(OpaqueTableStorage):
59 """An implementation of `OpaqueTableStorage` that simply creates a true
60 table for each different named opaque logical table.
62 A `ByNameOpaqueTableStorageManager` instance should always be used to
63 construct and manage instances of this class.
65 Parameters
66 ----------
67 db : `Database`
68 Database engine interface for the namespace in which this table lives.
69 name : `str`
70 Name of the logical table (also used as the name of the actual table).
71 table : `sqlalchemy.schema.Table`
72 SQLAlchemy representation of the table, which must have already been
73 created in the namespace managed by ``db`` (this is the responsibility
74 of `ByNameOpaqueTableStorageManager`).
75 """
77 def __init__(self, *, db: Database, name: str, table: sqlalchemy.schema.Table):
78 super().__init__(name=name)
79 self._db = db
80 self._table = table
82 def insert(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
83 # Docstring inherited from OpaqueTableStorage.
84 # The provided transaction object can be ignored since we rely on
85 # the database itself providing any rollback functionality.
86 self._db.insert(self._table, *data)
88 def ensure(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
89 # Docstring inherited from OpaqueTableStorage.
90 # The provided transaction object can be ignored since we rely on
91 # the database itself providing any rollback functionality.
92 self._db.ensure(self._table, *data)
94 def replace(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
95 # Docstring inherited from OpaqueTableStorage.
96 # The provided transaction object can be ignored since we rely on
97 # the database itself providing any rollback functionality.
98 self._db.replace(self._table, *data)
100 def fetch(self, **where: Any) -> Iterator[sqlalchemy.RowMapping]:
101 # Docstring inherited from OpaqueTableStorage.
103 def _batch_in_clause(
104 column: sqlalchemy.schema.Column, values: Iterable[Any]
105 ) -> Iterator[sqlalchemy.sql.expression.ClauseElement]:
106 """Split one long IN clause into a series of shorter ones."""
107 in_limit = 1000
108 # We have to remove possible duplicates from values; and in many
109 # cases it should be helpful to order the items in the clause.
110 values = sorted(set(values))
111 for iposn in range(0, len(values), in_limit):
112 in_clause = column.in_(values[iposn : iposn + in_limit])
113 yield in_clause
115 def _batch_in_clauses(**where: Any) -> Iterator[sqlalchemy.sql.expression.ColumnElement]:
116 """Generate a sequence of WHERE clauses with a limited number of
117 items in IN clauses.
118 """
119 batches: list[Iterable[Any]] = []
120 for k, v in where.items():
121 column = self._table.columns[k]
122 if isinstance(v, list | tuple | set):
123 batches.append(_batch_in_clause(column, v))
124 else:
125 # single "batch" for a regular eq operator
126 batches.append([column == v])
128 for clauses in itertools.product(*batches):
129 yield sqlalchemy.sql.and_(*clauses)
131 sql = self._table.select()
132 if where:
133 # Split long IN clauses into shorter batches
134 batched_sql = [sql.where(clause) for clause in _batch_in_clauses(**where)]
135 else:
136 batched_sql = [sql]
137 for sql_batch in batched_sql:
138 with self._db.query(sql_batch) as sql_result:
139 sql_mappings = sql_result.mappings().fetchall()
140 yield from sql_mappings
142 def delete(self, columns: Iterable[str], *rows: dict) -> None:
143 # Docstring inherited from OpaqueTableStorage.
144 self._db.delete(self._table, columns, *rows)
147class ByNameOpaqueTableStorageManager(OpaqueTableStorageManager):
148 """An implementation of `OpaqueTableStorageManager` that simply creates a
149 true table for each different named opaque logical table.
151 Instances of this class should generally be constructed via the
152 `initialize` class method instead of invoking ``__init__`` directly.
154 Parameters
155 ----------
156 db : `Database`
157 Database engine interface for the namespace in which this table lives.
158 metaTable : `sqlalchemy.schema.Table`
159 SQLAlchemy representation of the table that records which opaque
160 logical tables exist.
161 registry_schema_version : `VersionTuple` or `None`, optional
162 Version of registry schema.
163 """
165 def __init__(
166 self,
167 db: Database,
168 metaTable: sqlalchemy.schema.Table,
169 registry_schema_version: VersionTuple | None = None,
170 ):
171 super().__init__(registry_schema_version=registry_schema_version)
172 self._db = db
173 self._metaTable = metaTable
174 self._storage: dict[str, OpaqueTableStorage] = {}
176 _META_TABLE_NAME: ClassVar[str] = "opaque_meta"
178 _META_TABLE_SPEC: ClassVar[TableSpec] = TableSpec(
179 fields=[
180 FieldSpec("table_name", dtype=sqlalchemy.String, length=128, primaryKey=True),
181 ],
182 )
184 @classmethod
185 def initialize(
186 cls, db: Database, context: StaticTablesContext, registry_schema_version: VersionTuple | None = None
187 ) -> OpaqueTableStorageManager:
188 # Docstring inherited from OpaqueTableStorageManager.
189 metaTable = context.addTable(cls._META_TABLE_NAME, cls._META_TABLE_SPEC)
190 return cls(db=db, metaTable=metaTable, registry_schema_version=registry_schema_version)
192 def get(self, name: str) -> OpaqueTableStorage | None:
193 # Docstring inherited from OpaqueTableStorageManager.
194 return self._storage.get(name)
196 def register(self, name: str, spec: TableSpec) -> OpaqueTableStorage:
197 # Docstring inherited from OpaqueTableStorageManager.
198 result = self._storage.get(name)
199 if result is None:
200 # Create the table itself. If it already exists but wasn't in
201 # the dict because it was added by another client since this one
202 # was initialized, that's fine.
203 table = self._db.ensureTableExists(name, spec)
204 # Add a row to the meta table so we can find this table in the
205 # future. Also okay if that already exists, so we use sync.
206 self._db.sync(self._metaTable, keys={"table_name": name})
207 result = ByNameOpaqueTableStorage(name=name, table=table, db=self._db)
208 self._storage[name] = result
209 return result
211 @classmethod
212 def currentVersions(cls) -> list[VersionTuple]:
213 # Docstring inherited from VersionedExtension.
214 return [_VERSION]