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