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

61 statements  

« prev     ^ index     » next       coverage.py v6.4.2, created at 2022-07-15 02:34 -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 typing import TYPE_CHECKING, Any, Dict, Iterable, Optional, Type 

26 

27import sqlalchemy 

28 

29from ...core import TimespanDatabaseRepresentation, ddl 

30from ..interfaces import CollectionRecord, VersionTuple 

31from ._base import ( 

32 CollectionTablesTuple, 

33 DefaultCollectionManager, 

34 makeCollectionChainTableSpec, 

35 makeRunTableSpec, 

36) 

37 

38if TYPE_CHECKING: 38 ↛ 39line 38 didn't jump to line 39, because the condition on line 38 was never true

39 from ..interfaces import Database, DimensionRecordStorageManager, StaticTablesContext 

40 

41 

42_KEY_FIELD_SPEC = ddl.FieldSpec( 

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

44) 

45 

46 

47# This has to be updated on every schema change 

48_VERSION = VersionTuple(2, 0, 0) 

49 

50 

51def _makeTableSpecs(TimespanReprClass: Type[TimespanDatabaseRepresentation]) -> CollectionTablesTuple: 

52 return CollectionTablesTuple( 

53 collection=ddl.TableSpec( 

54 fields=[ 

55 _KEY_FIELD_SPEC, 

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

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

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

59 ], 

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

61 ), 

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

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

64 ) 

65 

66 

67class SynthIntKeyCollectionManager(DefaultCollectionManager): 

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

69 (auto-incremented integer) for collections table. 

70 

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

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

73 table schema. 

74 

75 Parameters 

76 ---------- 

77 db : `Database` 

78 Interface to the underlying database engine and namespace. 

79 tables : `NamedTuple` 

80 Named tuple of SQLAlchemy table objects. 

81 collectionIdName : `str` 

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

83 dimensions : `DimensionRecordStorageManager` 

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

85 """ 

86 

87 def __init__( 

88 self, 

89 db: Database, 

90 tables: CollectionTablesTuple, 

91 collectionIdName: str, 

92 dimensions: DimensionRecordStorageManager, 

93 ): 

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

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

96 

97 @classmethod 

98 def initialize( 

99 cls, 

100 db: Database, 

101 context: StaticTablesContext, 

102 *, 

103 dimensions: DimensionRecordStorageManager, 

104 ) -> SynthIntKeyCollectionManager: 

105 # Docstring inherited from CollectionManager. 

106 return cls( 

107 db, 

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

109 collectionIdName="collection_id", 

110 dimensions=dimensions, 

111 ) 

112 

113 @classmethod 

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

115 # Docstring inherited from CollectionManager. 

116 return f"{prefix}_id" 

117 

118 @classmethod 

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

120 # Docstring inherited from CollectionManager. 

121 return f"{prefix}_id" 

122 

123 @classmethod 

124 def addCollectionForeignKey( 

125 cls, 

126 tableSpec: ddl.TableSpec, 

127 *, 

128 prefix: str = "collection", 

129 onDelete: Optional[str] = None, 

130 constraint: bool = True, 

131 **kwargs: Any, 

132 ) -> ddl.FieldSpec: 

133 # Docstring inherited from CollectionManager. 

134 original = _KEY_FIELD_SPEC 

135 copy = ddl.FieldSpec( 

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

137 ) 

138 tableSpec.fields.add(copy) 

139 if constraint: 

140 tableSpec.foreignKeys.append( 

141 ddl.ForeignKeySpec( 

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

143 ) 

144 ) 

145 return copy 

146 

147 @classmethod 

148 def addRunForeignKey( 

149 cls, 

150 tableSpec: ddl.TableSpec, 

151 *, 

152 prefix: str = "run", 

153 onDelete: Optional[str] = None, 

154 constraint: bool = True, 

155 **kwargs: Any, 

156 ) -> ddl.FieldSpec: 

157 # Docstring inherited from CollectionManager. 

158 original = _KEY_FIELD_SPEC 

159 copy = ddl.FieldSpec( 

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

161 ) 

162 tableSpec.fields.add(copy) 

163 if constraint: 

164 tableSpec.foreignKeys.append( 

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

166 ) 

167 return copy 

168 

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

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

171 old cached records will be removed. 

172 """ 

173 self._records = {} 

174 self._nameCache = {} 

175 for record in records: 

176 self._records[record.key] = record 

177 self._nameCache[record.name] = record 

178 

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

180 """Add single record to cache.""" 

181 self._records[record.key] = record 

182 self._nameCache[record.name] = record 

183 

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

185 """Remove single record from cache.""" 

186 del self._records[record.key] 

187 del self._nameCache[record.name] 

188 

189 def _getByName(self, name: str) -> Optional[CollectionRecord]: 

190 # Docstring inherited from DefaultCollectionManager. 

191 return self._nameCache.get(name) 

192 

193 @classmethod 

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

195 # Docstring inherited from VersionedExtension. 

196 return _VERSION 

197 

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

199 # Docstring inherited from VersionedExtension. 

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