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

59 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-10-27 09:43 +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/>. 

27from __future__ import annotations 

28 

29from ... import ddl 

30 

31__all__ = ["SynthIntKeyCollectionManager"] 

32 

33from collections.abc import Iterable 

34from typing import TYPE_CHECKING, Any 

35 

36import sqlalchemy 

37 

38from ..._timespan import TimespanDatabaseRepresentation 

39from ..interfaces import CollectionRecord, VersionTuple 

40from ._base import ( 

41 CollectionTablesTuple, 

42 DefaultCollectionManager, 

43 makeCollectionChainTableSpec, 

44 makeRunTableSpec, 

45) 

46 

47if TYPE_CHECKING: 

48 from ..interfaces import Database, DimensionRecordStorageManager, StaticTablesContext 

49 

50 

51_KEY_FIELD_SPEC = ddl.FieldSpec( 

52 "collection_id", dtype=sqlalchemy.BigInteger, primaryKey=True, autoincrement=True 

53) 

54 

55 

56# This has to be updated on every schema change 

57_VERSION = VersionTuple(2, 0, 0) 

58 

59 

60def _makeTableSpecs(TimespanReprClass: type[TimespanDatabaseRepresentation]) -> CollectionTablesTuple: 

61 return CollectionTablesTuple( 

62 collection=ddl.TableSpec( 

63 fields=[ 

64 _KEY_FIELD_SPEC, 

65 ddl.FieldSpec("name", dtype=sqlalchemy.String, length=64, nullable=False), 

66 ddl.FieldSpec("type", dtype=sqlalchemy.SmallInteger, nullable=False), 

67 ddl.FieldSpec("doc", dtype=sqlalchemy.Text, nullable=True), 

68 ], 

69 unique=[("name",)], 

70 ), 

71 run=makeRunTableSpec("collection_id", sqlalchemy.BigInteger, TimespanReprClass), 

72 collection_chain=makeCollectionChainTableSpec("collection_id", sqlalchemy.BigInteger), 

73 ) 

74 

75 

76class SynthIntKeyCollectionManager(DefaultCollectionManager): 

77 """A `CollectionManager` implementation that uses synthetic primary key 

78 (auto-incremented integer) for collections table. 

79 

80 Most of the logic, including caching policy, is implemented in the base 

81 class, this class only adds customizations specific to this particular 

82 table schema. 

83 

84 Parameters 

85 ---------- 

86 db : `Database` 

87 Interface to the underlying database engine and namespace. 

88 tables : `NamedTuple` 

89 Named tuple of SQLAlchemy table objects. 

90 collectionIdName : `str` 

91 Name of the column in collections table that identifies it (PK). 

92 dimensions : `DimensionRecordStorageManager` 

93 Manager object for the dimensions in this `Registry`. 

94 """ 

95 

96 def __init__( 

97 self, 

98 db: Database, 

99 tables: CollectionTablesTuple, 

100 collectionIdName: str, 

101 dimensions: DimensionRecordStorageManager, 

102 registry_schema_version: VersionTuple | None = None, 

103 ): 

104 super().__init__( 

105 db=db, 

106 tables=tables, 

107 collectionIdName=collectionIdName, 

108 dimensions=dimensions, 

109 registry_schema_version=registry_schema_version, 

110 ) 

111 self._nameCache: dict[str, CollectionRecord] = {} # indexed by collection name 

112 

113 @classmethod 

114 def initialize( 

115 cls, 

116 db: Database, 

117 context: StaticTablesContext, 

118 *, 

119 dimensions: DimensionRecordStorageManager, 

120 registry_schema_version: VersionTuple | None = None, 

121 ) -> SynthIntKeyCollectionManager: 

122 # Docstring inherited from CollectionManager. 

123 return cls( 

124 db, 

125 tables=context.addTableTuple(_makeTableSpecs(db.getTimespanRepresentation())), # type: ignore 

126 collectionIdName="collection_id", 

127 dimensions=dimensions, 

128 registry_schema_version=registry_schema_version, 

129 ) 

130 

131 @classmethod 

132 def getCollectionForeignKeyName(cls, prefix: str = "collection") -> str: 

133 # Docstring inherited from CollectionManager. 

134 return f"{prefix}_id" 

135 

136 @classmethod 

137 def getRunForeignKeyName(cls, prefix: str = "run") -> str: 

138 # Docstring inherited from CollectionManager. 

139 return f"{prefix}_id" 

140 

141 @classmethod 

142 def addCollectionForeignKey( 

143 cls, 

144 tableSpec: ddl.TableSpec, 

145 *, 

146 prefix: str = "collection", 

147 onDelete: str | None = None, 

148 constraint: bool = True, 

149 **kwargs: Any, 

150 ) -> ddl.FieldSpec: 

151 # Docstring inherited from CollectionManager. 

152 original = _KEY_FIELD_SPEC 

153 copy = ddl.FieldSpec( 

154 cls.getCollectionForeignKeyName(prefix), dtype=original.dtype, autoincrement=False, **kwargs 

155 ) 

156 tableSpec.fields.add(copy) 

157 if constraint: 

158 tableSpec.foreignKeys.append( 

159 ddl.ForeignKeySpec( 

160 "collection", source=(copy.name,), target=(original.name,), onDelete=onDelete 

161 ) 

162 ) 

163 return copy 

164 

165 @classmethod 

166 def addRunForeignKey( 

167 cls, 

168 tableSpec: ddl.TableSpec, 

169 *, 

170 prefix: str = "run", 

171 onDelete: str | None = None, 

172 constraint: bool = True, 

173 **kwargs: Any, 

174 ) -> ddl.FieldSpec: 

175 # Docstring inherited from CollectionManager. 

176 original = _KEY_FIELD_SPEC 

177 copy = ddl.FieldSpec( 

178 cls.getRunForeignKeyName(prefix), dtype=original.dtype, autoincrement=False, **kwargs 

179 ) 

180 tableSpec.fields.add(copy) 

181 if constraint: 181 ↛ 185line 181 didn't jump to line 185, because the condition on line 181 was never false

182 tableSpec.foreignKeys.append( 

183 ddl.ForeignKeySpec("run", source=(copy.name,), target=(original.name,), onDelete=onDelete) 

184 ) 

185 return copy 

186 

187 def _setRecordCache(self, records: Iterable[CollectionRecord]) -> None: 

188 """Set internal record cache to contain given records, 

189 old cached records will be removed. 

190 """ 

191 self._records = {} 

192 self._nameCache = {} 

193 for record in records: 

194 self._records[record.key] = record 

195 self._nameCache[record.name] = record 

196 

197 def _addCachedRecord(self, record: CollectionRecord) -> None: 

198 """Add single record to cache.""" 

199 self._records[record.key] = record 

200 self._nameCache[record.name] = record 

201 

202 def _removeCachedRecord(self, record: CollectionRecord) -> None: 

203 """Remove single record from cache.""" 

204 del self._records[record.key] 

205 del self._nameCache[record.name] 

206 

207 def _getByName(self, name: str) -> CollectionRecord | None: 

208 # Docstring inherited from DefaultCollectionManager. 

209 return self._nameCache.get(name) 

210 

211 @classmethod 

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

213 # Docstring inherited from VersionedExtension. 

214 return [_VERSION]