Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

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 ( 

26 Any, 

27 Dict, 

28 Iterable, 

29 Optional, 

30 Type, 

31 TYPE_CHECKING, 

32) 

33 

34import sqlalchemy 

35 

36from ._base import ( 

37 CollectionTablesTuple, 

38 DefaultCollectionManager, 

39 makeRunTableSpec, 

40 makeCollectionChainTableSpec, 

41) 

42from ...core import TimespanDatabaseRepresentation, ddl 

43from ..interfaces import CollectionRecord, VersionTuple 

44 

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

46 from ..interfaces import Database, DimensionRecordStorageManager, StaticTablesContext 

47 

48 

49_KEY_FIELD_SPEC = ddl.FieldSpec("collection_id", dtype=sqlalchemy.BigInteger, primaryKey=True, 

50 autoincrement=True) 

51 

52 

53# This has to be updated on every schema change 

54_VERSION = VersionTuple(2, 0, 0) 

55 

56 

57def _makeTableSpecs(tsRepr: Type[TimespanDatabaseRepresentation]) -> CollectionTablesTuple: 

58 return CollectionTablesTuple( 

59 collection=ddl.TableSpec( 

60 fields=[ 

61 _KEY_FIELD_SPEC, 

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

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

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

65 ], 

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

67 ), 

68 run=makeRunTableSpec("collection_id", sqlalchemy.BigInteger, tsRepr), 

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

70 ) 

71 

72 

73class SynthIntKeyCollectionManager(DefaultCollectionManager): 

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

75 (auto-incremented integer) for collections table. 

76 

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

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

79 table schema. 

80 

81 Parameters 

82 ---------- 

83 db : `Database` 

84 Interface to the underlying database engine and namespace. 

85 tables : `NamedTuple` 

86 Named tuple of SQLAlchemy table objects. 

87 collectionIdName : `str` 

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

89 dimensions : `DimensionRecordStorageManager` 

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

91 """ 

92 def __init__( 

93 self, 

94 db: Database, 

95 tables: CollectionTablesTuple, 

96 collectionIdName: str, 

97 dimensions: DimensionRecordStorageManager, 

98 ): 

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

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

101 

102 @classmethod 

103 def initialize( 

104 cls, 

105 db: Database, 

106 context: StaticTablesContext, *, 

107 dimensions: DimensionRecordStorageManager, 

108 ) -> SynthIntKeyCollectionManager: 

109 # Docstring inherited from CollectionManager. 

110 return cls( 

111 db, 

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

113 collectionIdName="collection_id", 

114 dimensions=dimensions, 

115 ) 

116 

117 @classmethod 

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

119 # Docstring inherited from CollectionManager. 

120 return f"{prefix}_id" 

121 

122 @classmethod 

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

124 # Docstring inherited from CollectionManager. 

125 return f"{prefix}_id" 

126 

127 @classmethod 

128 def addCollectionForeignKey(cls, tableSpec: ddl.TableSpec, *, prefix: str = "collection", 

129 onDelete: Optional[str] = None, 

130 constraint: bool = True, 

131 **kwargs: Any) -> ddl.FieldSpec: 

132 # Docstring inherited from CollectionManager. 

133 original = _KEY_FIELD_SPEC 

134 copy = ddl.FieldSpec(cls.getCollectionForeignKeyName(prefix), dtype=original.dtype, 

135 autoincrement=False, **kwargs) 

136 tableSpec.fields.add(copy) 

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

138 tableSpec.foreignKeys.append(ddl.ForeignKeySpec("collection", source=(copy.name,), 

139 target=(original.name,), onDelete=onDelete)) 

140 return copy 

141 

142 @classmethod 

143 def addRunForeignKey(cls, tableSpec: ddl.TableSpec, *, prefix: str = "run", 

144 onDelete: Optional[str] = None, 

145 constraint: bool = True, 

146 **kwargs: Any) -> ddl.FieldSpec: 

147 # Docstring inherited from CollectionManager. 

148 original = _KEY_FIELD_SPEC 

149 copy = ddl.FieldSpec(cls.getRunForeignKeyName(prefix), dtype=original.dtype, 

150 autoincrement=False, **kwargs) 

151 tableSpec.fields.add(copy) 

152 if constraint: 

153 tableSpec.foreignKeys.append(ddl.ForeignKeySpec("run", source=(copy.name,), 

154 target=(original.name,), onDelete=onDelete)) 

155 return copy 

156 

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

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

159 old cached records will be removed. 

160 """ 

161 self._records = {} 

162 self._nameCache = {} 

163 for record in records: 

164 self._records[record.key] = record 

165 self._nameCache[record.name] = record 

166 

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

168 """Add single record to cache. 

169 """ 

170 self._records[record.key] = record 

171 self._nameCache[record.name] = record 

172 

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

174 """Remove single record from cache. 

175 """ 

176 del self._records[record.key] 

177 del self._nameCache[record.name] 

178 

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

180 # Docstring inherited from DefaultCollectionManager. 

181 return self._nameCache.get(name) 

182 

183 @classmethod 

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

185 # Docstring inherited from VersionedExtension. 

186 return _VERSION 

187 

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

189 # Docstring inherited from VersionedExtension. 

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