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

62 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-12-15 02:03 -0800

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: 39 ↛ 40line 39 didn't jump to line 40, because the condition on line 39 was never true

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

95 super().__init__(db=db, tables=tables, collectionIdName=collectionIdName, dimensions=dimensions) 

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

97 

98 @classmethod 

99 def initialize( 

100 cls, 

101 db: Database, 

102 context: StaticTablesContext, 

103 *, 

104 dimensions: DimensionRecordStorageManager, 

105 ) -> SynthIntKeyCollectionManager: 

106 # Docstring inherited from CollectionManager. 

107 return cls( 

108 db, 

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

110 collectionIdName="collection_id", 

111 dimensions=dimensions, 

112 ) 

113 

114 @classmethod 

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

116 # Docstring inherited from CollectionManager. 

117 return f"{prefix}_id" 

118 

119 @classmethod 

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

121 # Docstring inherited from CollectionManager. 

122 return f"{prefix}_id" 

123 

124 @classmethod 

125 def addCollectionForeignKey( 

126 cls, 

127 tableSpec: ddl.TableSpec, 

128 *, 

129 prefix: str = "collection", 

130 onDelete: str | None = None, 

131 constraint: bool = True, 

132 **kwargs: Any, 

133 ) -> ddl.FieldSpec: 

134 # Docstring inherited from CollectionManager. 

135 original = _KEY_FIELD_SPEC 

136 copy = ddl.FieldSpec( 

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

138 ) 

139 tableSpec.fields.add(copy) 

140 if constraint: 

141 tableSpec.foreignKeys.append( 

142 ddl.ForeignKeySpec( 

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

144 ) 

145 ) 

146 return copy 

147 

148 @classmethod 

149 def addRunForeignKey( 

150 cls, 

151 tableSpec: ddl.TableSpec, 

152 *, 

153 prefix: str = "run", 

154 onDelete: str | None = None, 

155 constraint: bool = True, 

156 **kwargs: Any, 

157 ) -> ddl.FieldSpec: 

158 # Docstring inherited from CollectionManager. 

159 original = _KEY_FIELD_SPEC 

160 copy = ddl.FieldSpec( 

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

162 ) 

163 tableSpec.fields.add(copy) 

164 if constraint: 

165 tableSpec.foreignKeys.append( 

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

167 ) 

168 return copy 

169 

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

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

172 old cached records will be removed. 

173 """ 

174 self._records = {} 

175 self._nameCache = {} 

176 for record in records: 

177 self._records[record.key] = record 

178 self._nameCache[record.name] = record 

179 

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

181 """Add single record to cache.""" 

182 self._records[record.key] = record 

183 self._nameCache[record.name] = record 

184 

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

186 """Remove single record from cache.""" 

187 del self._records[record.key] 

188 del self._nameCache[record.name] 

189 

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

191 # Docstring inherited from DefaultCollectionManager. 

192 return self._nameCache.get(name) 

193 

194 @classmethod 

195 def currentVersion(cls) -> VersionTuple | None: 

196 # Docstring inherited from VersionedExtension. 

197 return _VERSION 

198 

199 def schemaDigest(self) -> str | None: 

200 # Docstring inherited from VersionedExtension. 

201 return self._defaultSchemaDigest(self._tables, self._db.dialect)