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

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 

22 

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

24opaque tables for `Registry`. 

25""" 

26 

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

28 

29import itertools 

30from collections.abc import Iterable, Iterator 

31from typing import TYPE_CHECKING, Any, ClassVar 

32 

33import sqlalchemy 

34 

35from ..core.ddl import FieldSpec, TableSpec 

36from .interfaces import ( 

37 Database, 

38 OpaqueTableStorage, 

39 OpaqueTableStorageManager, 

40 StaticTablesContext, 

41 VersionTuple, 

42) 

43 

44if TYPE_CHECKING: 

45 from ..core.datastore import DatastoreTransaction 

46 

47# This has to be updated on every schema change 

48_VERSION = VersionTuple(0, 2, 0) 

49 

50 

51class ByNameOpaqueTableStorage(OpaqueTableStorage): 

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

53 table for each different named opaque logical table. 

54 

55 A `ByNameOpaqueTableStorageManager` instance should always be used to 

56 construct and manage instances of this class. 

57 

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 """ 

69 

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 

74 

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) 

80 

81 def fetch(self, **where: Any) -> Iterator[sqlalchemy.RowMapping]: 

82 # Docstring inherited from OpaqueTableStorage. 

83 

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 

95 

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]) 

108 

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

110 yield sqlalchemy.sql.and_(*clauses) 

111 

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 

122 

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

124 # Docstring inherited from OpaqueTableStorage. 

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

126 

127 

128class ByNameOpaqueTableStorageManager(OpaqueTableStorageManager): 

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

130 true table for each different named opaque logical table. 

131 

132 Instances of this class should generally be constructed via the 

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

134 

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 """ 

143 

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] = {} 

154 

155 _META_TABLE_NAME: ClassVar[str] = "opaque_meta" 

156 

157 _META_TABLE_SPEC: ClassVar[TableSpec] = TableSpec( 

158 fields=[ 

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

160 ], 

161 ) 

162 

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) 

170 

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

172 # Docstring inherited from OpaqueTableStorageManager. 

173 return self._storage.get(name) 

174 

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 

189 

190 @classmethod 

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

192 # Docstring inherited from VersionedExtension. 

193 return [_VERSION]