Coverage for python / lsst / daf / butler / registry / opaque.py: 29%

80 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-04-24 08:17 +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/>. 

27 

28"""The default concrete implementations of the classes that manage 

29opaque tables for `Registry`. 

30""" 

31 

32from __future__ import annotations 

33 

34__all__ = ["ByNameOpaqueTableStorage", "ByNameOpaqueTableStorageManager"] 

35 

36import itertools 

37from collections.abc import Iterable, Iterator 

38from typing import TYPE_CHECKING, Any, ClassVar 

39 

40import sqlalchemy 

41 

42from .._utilities.thread_safe_cache import ThreadSafeCache 

43from ..ddl import FieldSpec, TableSpec 

44from .interfaces import ( 

45 Database, 

46 OpaqueTableStorage, 

47 OpaqueTableStorageManager, 

48 StaticTablesContext, 

49 VersionTuple, 

50) 

51 

52if TYPE_CHECKING: 

53 from ..datastore import DatastoreTransaction 

54 

55# This has to be updated on every schema change 

56_VERSION = VersionTuple(0, 2, 0) 

57 

58 

59class ByNameOpaqueTableStorage(OpaqueTableStorage): 

60 """An implementation of `OpaqueTableStorage` that simply creates a true 

61 table for each different named opaque logical table. 

62 

63 A `ByNameOpaqueTableStorageManager` instance should always be used to 

64 construct and manage instances of this class. 

65 

66 Parameters 

67 ---------- 

68 db : `Database` 

69 Database engine interface for the namespace in which this table lives. 

70 name : `str` 

71 Name of the logical table (also used as the name of the actual table). 

72 table : `sqlalchemy.schema.Table` 

73 SQLAlchemy representation of the table, which must have already been 

74 created in the namespace managed by ``db`` (this is the responsibility 

75 of `ByNameOpaqueTableStorageManager`). 

76 """ 

77 

78 def __init__(self, *, db: Database, name: str, table: sqlalchemy.schema.Table): 

79 super().__init__(name=name) 

80 self._db = db 

81 self._table = table 

82 

83 def insert(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None: 

84 # Docstring inherited from OpaqueTableStorage. 

85 # The provided transaction object can be ignored since we rely on 

86 # the database itself providing any rollback functionality. 

87 self._db.insert(self._table, *data) 

88 

89 def ensure(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None: 

90 # Docstring inherited from OpaqueTableStorage. 

91 # The provided transaction object can be ignored since we rely on 

92 # the database itself providing any rollback functionality. 

93 self._db.ensure(self._table, *data) 

94 

95 def replace(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None: 

96 # Docstring inherited from OpaqueTableStorage. 

97 # The provided transaction object can be ignored since we rely on 

98 # the database itself providing any rollback functionality. 

99 self._db.replace(self._table, *data) 

100 

101 def fetch( 

102 self, 

103 **where: Any, 

104 ) -> Iterator[sqlalchemy.RowMapping]: 

105 # Docstring inherited from OpaqueTableStorage. 

106 

107 def _batch_in_clause( 

108 column: sqlalchemy.schema.Column, values: Iterable[Any] 

109 ) -> Iterator[sqlalchemy.sql.expression.ClauseElement]: 

110 """Split one long IN clause into a series of shorter ones.""" 

111 in_limit = 1000 

112 # We have to remove possible duplicates from values; and in many 

113 # cases it should be helpful to order the items in the clause. 

114 values = sorted(set(values)) 

115 for iposn in range(0, len(values), in_limit): 

116 in_clause = column.in_(values[iposn : iposn + in_limit]) 

117 yield in_clause 

118 

119 def _batch_in_clauses(**where: Any) -> Iterator[sqlalchemy.sql.expression.ColumnElement]: 

120 """Generate a sequence of WHERE clauses with a limited number of 

121 items in IN clauses. 

122 """ 

123 batches: list[Iterable[Any]] = [] 

124 for k, v in where.items(): 

125 column = self._table.columns[k] 

126 if isinstance(v, list | tuple | set): 

127 batches.append(_batch_in_clause(column, v)) 

128 else: 

129 if isinstance(v, str) and v.endswith("%"): 

130 # Special case prefix queries. 

131 batches.append([column.startswith(v[:-1])]) 

132 else: 

133 # single "batch" for a regular eq operator 

134 batches.append([column == v]) 

135 

136 for clauses in itertools.product(*batches): 

137 yield sqlalchemy.sql.and_(*clauses) 

138 

139 sql = self._table.select() 

140 if where: 

141 # Split long IN clauses into shorter batches 

142 batched_sql = [sql.where(clause) for clause in _batch_in_clauses(**where)] 

143 else: 

144 batched_sql = [sql] 

145 for sql_batch in batched_sql: 

146 with self._db.query(sql_batch) as sql_result: 

147 sql_mappings = sql_result.mappings().fetchall() 

148 yield from sql_mappings 

149 

150 def delete(self, columns: Iterable[str], *rows: dict) -> None: 

151 # Docstring inherited from OpaqueTableStorage. 

152 self._db.delete(self._table, columns, *rows) 

153 

154 

155class ByNameOpaqueTableStorageManager(OpaqueTableStorageManager): 

156 """An implementation of `OpaqueTableStorageManager` that simply creates a 

157 true table for each different named opaque logical table. 

158 

159 Instances of this class should generally be constructed via the 

160 `initialize` class method instead of invoking ``__init__`` directly. 

161 

162 Parameters 

163 ---------- 

164 db : `Database` 

165 Database engine interface for the namespace in which this table lives. 

166 metaTable : `sqlalchemy.schema.Table` 

167 SQLAlchemy representation of the table that records which opaque 

168 logical tables exist. 

169 tables : `ThreadSafeCache` [`str`, `~sqlalchemy.schema.Table`] 

170 Mapping from string to table, to track which tables have already been 

171 created. This mapping is shared between cloned instances of this 

172 manager. 

173 registry_schema_version : `VersionTuple` or `None`, optional 

174 Version of registry schema. 

175 """ 

176 

177 def __init__( 

178 self, 

179 db: Database, 

180 metaTable: sqlalchemy.schema.Table, 

181 tables: ThreadSafeCache[str, sqlalchemy.schema.Table], 

182 registry_schema_version: VersionTuple | None = None, 

183 ): 

184 super().__init__(registry_schema_version=registry_schema_version) 

185 self._db = db 

186 self._metaTable = metaTable 

187 self._tables = tables 

188 

189 def clone(self, db: Database) -> ByNameOpaqueTableStorageManager: 

190 return ByNameOpaqueTableStorageManager( 

191 db, self._metaTable, self._tables, self._registry_schema_version 

192 ) 

193 

194 _META_TABLE_NAME: ClassVar[str] = "opaque_meta" 

195 

196 _META_TABLE_SPEC: ClassVar[TableSpec] = TableSpec( 

197 fields=[ 

198 FieldSpec("table_name", dtype=sqlalchemy.String, length=128, primaryKey=True), 

199 ], 

200 ) 

201 

202 @classmethod 

203 def initialize( 

204 cls, db: Database, context: StaticTablesContext, registry_schema_version: VersionTuple | None = None 

205 ) -> OpaqueTableStorageManager: 

206 # Docstring inherited from OpaqueTableStorageManager. 

207 metaTable = context.addTable(cls._META_TABLE_NAME, cls._META_TABLE_SPEC) 

208 return cls( 

209 db=db, 

210 metaTable=metaTable, 

211 tables=ThreadSafeCache(), 

212 registry_schema_version=registry_schema_version, 

213 ) 

214 

215 def get(self, name: str) -> OpaqueTableStorage | None: 

216 # Docstring inherited from OpaqueTableStorageManager. 

217 table = self._tables.get(name) 

218 if table is None: 

219 return None 

220 return ByNameOpaqueTableStorage(name=name, table=table, db=self._db) 

221 

222 def register(self, name: str, spec: TableSpec) -> OpaqueTableStorage: 

223 # Docstring inherited from OpaqueTableStorageManager. 

224 result = self.get(name) 

225 if result is None: 

226 # Create the table itself. If it already exists but wasn't in 

227 # the dict because it was added by another client since this one 

228 # was initialized, that's fine. 

229 table = self._db.ensureTableExists(name, spec) 

230 # Add a row to the meta table so we can find this table in the 

231 # future. Also okay if that already exists, so we use sync. 

232 self._db.sync(self._metaTable, keys={"table_name": name}) 

233 self._tables.set_or_get(name, table) 

234 result = self.get(name) 

235 assert result is not None 

236 return result 

237 

238 @classmethod 

239 def currentVersions(cls) -> list[VersionTuple]: 

240 # Docstring inherited from VersionedExtension. 

241 return [_VERSION]