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

58 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-06-02 02:15 -0700

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__all__ = ["SynthIntKeyCollectionManager"] 

24 

25from collections.abc import Iterable 

26from typing import TYPE_CHECKING, Any 

27 

28import sqlalchemy 

29 

30from ...core import TimespanDatabaseRepresentation, ddl 

31from ..interfaces import CollectionRecord, VersionTuple 

32from ._base import ( 

33 CollectionTablesTuple, 

34 DefaultCollectionManager, 

35 makeCollectionChainTableSpec, 

36 makeRunTableSpec, 

37) 

38 

39if TYPE_CHECKING: 

40 from ..interfaces import Database, DimensionRecordStorageManager, StaticTablesContext 

41 

42 

43_KEY_FIELD_SPEC = ddl.FieldSpec( 

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

45) 

46 

47 

48# This has to be updated on every schema change 

49_VERSION = VersionTuple(2, 0, 0) 

50 

51 

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

53 return CollectionTablesTuple( 

54 collection=ddl.TableSpec( 

55 fields=[ 

56 _KEY_FIELD_SPEC, 

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

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

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

60 ], 

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

62 ), 

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

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

65 ) 

66 

67 

68class SynthIntKeyCollectionManager(DefaultCollectionManager): 

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

70 (auto-incremented integer) for collections table. 

71 

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

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

74 table schema. 

75 

76 Parameters 

77 ---------- 

78 db : `Database` 

79 Interface to the underlying database engine and namespace. 

80 tables : `NamedTuple` 

81 Named tuple of SQLAlchemy table objects. 

82 collectionIdName : `str` 

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

84 dimensions : `DimensionRecordStorageManager` 

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

86 """ 

87 

88 def __init__( 

89 self, 

90 db: Database, 

91 tables: CollectionTablesTuple, 

92 collectionIdName: str, 

93 dimensions: DimensionRecordStorageManager, 

94 registry_schema_version: VersionTuple | None = None, 

95 ): 

96 super().__init__( 

97 db=db, 

98 tables=tables, 

99 collectionIdName=collectionIdName, 

100 dimensions=dimensions, 

101 registry_schema_version=registry_schema_version, 

102 ) 

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

104 

105 @classmethod 

106 def initialize( 

107 cls, 

108 db: Database, 

109 context: StaticTablesContext, 

110 *, 

111 dimensions: DimensionRecordStorageManager, 

112 registry_schema_version: VersionTuple | None = None, 

113 ) -> SynthIntKeyCollectionManager: 

114 # Docstring inherited from CollectionManager. 

115 return cls( 

116 db, 

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

118 collectionIdName="collection_id", 

119 dimensions=dimensions, 

120 registry_schema_version=registry_schema_version, 

121 ) 

122 

123 @classmethod 

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

125 # Docstring inherited from CollectionManager. 

126 return f"{prefix}_id" 

127 

128 @classmethod 

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

130 # Docstring inherited from CollectionManager. 

131 return f"{prefix}_id" 

132 

133 @classmethod 

134 def addCollectionForeignKey( 

135 cls, 

136 tableSpec: ddl.TableSpec, 

137 *, 

138 prefix: str = "collection", 

139 onDelete: str | None = None, 

140 constraint: bool = True, 

141 **kwargs: Any, 

142 ) -> ddl.FieldSpec: 

143 # Docstring inherited from CollectionManager. 

144 original = _KEY_FIELD_SPEC 

145 copy = ddl.FieldSpec( 

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

147 ) 

148 tableSpec.fields.add(copy) 

149 if constraint: 

150 tableSpec.foreignKeys.append( 

151 ddl.ForeignKeySpec( 

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

153 ) 

154 ) 

155 return copy 

156 

157 @classmethod 

158 def addRunForeignKey( 

159 cls, 

160 tableSpec: ddl.TableSpec, 

161 *, 

162 prefix: str = "run", 

163 onDelete: str | None = None, 

164 constraint: bool = True, 

165 **kwargs: Any, 

166 ) -> ddl.FieldSpec: 

167 # Docstring inherited from CollectionManager. 

168 original = _KEY_FIELD_SPEC 

169 copy = ddl.FieldSpec( 

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

171 ) 

172 tableSpec.fields.add(copy) 

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

174 tableSpec.foreignKeys.append( 

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

176 ) 

177 return copy 

178 

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

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

181 old cached records will be removed. 

182 """ 

183 self._records = {} 

184 self._nameCache = {} 

185 for record in records: 

186 self._records[record.key] = record 

187 self._nameCache[record.name] = record 

188 

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

190 """Add single record to cache.""" 

191 self._records[record.key] = record 

192 self._nameCache[record.name] = record 

193 

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

195 """Remove single record from cache.""" 

196 del self._records[record.key] 

197 del self._nameCache[record.name] 

198 

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

200 # Docstring inherited from DefaultCollectionManager. 

201 return self._nameCache.get(name) 

202 

203 @classmethod 

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

205 # Docstring inherited from VersionedExtension. 

206 return [_VERSION]