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

67 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-12-01 19:54 +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"""The default concrete implementations of the classes that manage 

23opaque tables for `Registry`. 

24""" 

25 

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

27 

28import itertools 

29from typing import ( 

30 Any, 

31 ClassVar, 

32 Dict, 

33 Iterable, 

34 Iterator, 

35 List, 

36 Optional, 

37) 

38 

39import sqlalchemy 

40 

41from ..core.ddl import TableSpec, FieldSpec 

42from .interfaces import ( 

43 Database, 

44 OpaqueTableStorageManager, 

45 OpaqueTableStorage, 

46 StaticTablesContext, 

47 VersionTuple 

48) 

49 

50 

51# This has to be updated on every schema change 

52_VERSION = VersionTuple(0, 2, 0) 

53 

54 

55class ByNameOpaqueTableStorage(OpaqueTableStorage): 

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

57 table for each different named opaque logical table. 

58 

59 A `ByNameOpaqueTableStorageManager` instance should always be used to 

60 construct and manage instances of this class. 

61 

62 Parameters 

63 ---------- 

64 db : `Database` 

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

66 name : `str` 

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

68 table : `sqlalchemy.schema.Table` 

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

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

71 of `ByNameOpaqueTableStorageManager`). 

72 """ 

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

74 super().__init__(name=name) 

75 self._db = db 

76 self._table = table 

77 

78 def insert(self, *data: dict) -> None: 

79 # Docstring inherited from OpaqueTableStorage. 

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

81 

82 def fetch(self, **where: Any) -> Iterator[dict]: 

83 # Docstring inherited from OpaqueTableStorage. 

84 

85 def _batch_in_clause(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 """ 

89 in_limit = 1000 

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

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

92 values = sorted(set(values)) 

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

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

95 yield in_clause 

96 

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

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

99 items in IN clauses. 

100 """ 

101 batches: List[Iterable[Any]] = [] 

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

103 column = self._table.columns[k] 

104 if isinstance(v, (list, tuple, set)): 

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

106 else: 

107 # single "batch" for a regular eq operator 

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

109 

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

111 yield sqlalchemy.sql.and_(*clauses) 

112 

113 sql = self._table.select() 

114 if where: 

115 # Split long IN clauses into shorter batches 

116 for clause in _batch_in_clauses(**where): 

117 sql_where = sql.where(clause) 

118 for row in self._db.query(sql_where): 

119 yield row._asdict() 

120 else: 

121 for row in self._db.query(sql): 

122 yield row._asdict() 

123 

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

125 # Docstring inherited from OpaqueTableStorage. 

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

127 

128 

129class ByNameOpaqueTableStorageManager(OpaqueTableStorageManager): 

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

131 true table for each different named opaque logical table. 

132 

133 Instances of this class should generally be constructed via the 

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

135 

136 Parameters 

137 ---------- 

138 db : `Database` 

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

140 metaTable : `sqlalchemy.schema.Table` 

141 SQLAlchemy representation of the table that records which opaque 

142 logical tables exist. 

143 """ 

144 def __init__(self, db: Database, metaTable: sqlalchemy.schema.Table): 

145 self._db = db 

146 self._metaTable = metaTable 

147 self._storage: Dict[str, OpaqueTableStorage] = {} 

148 

149 _META_TABLE_NAME: ClassVar[str] = "opaque_meta" 

150 

151 _META_TABLE_SPEC: ClassVar[TableSpec] = TableSpec( 

152 fields=[ 

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

154 ], 

155 ) 

156 

157 @classmethod 

158 def initialize(cls, db: Database, context: StaticTablesContext) -> OpaqueTableStorageManager: 

159 # Docstring inherited from OpaqueTableStorageManager. 

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

161 return cls(db=db, metaTable=metaTable) 

162 

163 def get(self, name: str) -> Optional[OpaqueTableStorage]: 

164 # Docstring inherited from OpaqueTableStorageManager. 

165 return self._storage.get(name) 

166 

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

168 # Docstring inherited from OpaqueTableStorageManager. 

169 result = self._storage.get(name) 

170 if result is None: 170 ↛ 180line 170 didn't jump to line 180, because the condition on line 170 was never false

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

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

173 # was initialized, that's fine. 

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

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

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

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

178 result = ByNameOpaqueTableStorage(name=name, table=table, db=self._db) 

179 self._storage[name] = result 

180 return result 

181 

182 @classmethod 

183 def currentVersion(cls) -> Optional[VersionTuple]: 

184 # Docstring inherited from VersionedExtension. 

185 return _VERSION 

186 

187 def schemaDigest(self) -> Optional[str]: 

188 # Docstring inherited from VersionedExtension. 

189 return self._defaultSchemaDigest([self._metaTable], self._db.dialect)