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_CHECKING, 

31) 

32 

33import sqlalchemy 

34 

35from ._base import ( 

36 CollectionTablesTuple, 

37 DefaultCollectionManager, 

38 makeRunTableSpec, 

39 makeCollectionChainTableSpec, 

40) 

41from ...core import ddl 

42from ..interfaces import CollectionRecord, VersionTuple 

43 

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

45 from ..interfaces import Database, StaticTablesContext 

46 

47 

48_TABLES_SPEC = CollectionTablesTuple( 

49 collection=ddl.TableSpec( 

50 fields=[ 

51 ddl.FieldSpec("collection_id", dtype=sqlalchemy.BigInteger, primaryKey=True, autoincrement=True), 

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

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

54 ], 

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

56 ), 

57 run=makeRunTableSpec("collection_id", sqlalchemy.BigInteger), 

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

59) 

60 

61# This has to be updated on every schema change 

62_VERSION = VersionTuple(0, 1, 0) 

63 

64 

65class SynthIntKeyCollectionManager(DefaultCollectionManager): 

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

67 (auto-incremented integer) for collections table. 

68 

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

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

71 table schema. 

72 

73 Parameters 

74 ---------- 

75 db : `Database` 

76 Interface to the underlying database engine and namespace. 

77 tables : `NamedTuple` 

78 Named tuple of SQLAlchemy table objects. 

79 collectionIdName : `str` 

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

81 """ 

82 def __init__(self, db: Database, tables: CollectionTablesTuple, collectionIdName: str): 

83 super().__init__(db=db, tables=tables, collectionIdName=collectionIdName) 

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

85 

86 @classmethod 

87 def initialize(cls, db: Database, context: StaticTablesContext) -> SynthIntKeyCollectionManager: 

88 # Docstring inherited from CollectionManager. 

89 return cls(db, tables=context.addTableTuple(_TABLES_SPEC), # type: ignore 

90 collectionIdName="collection_id") 

91 

92 @classmethod 

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

94 # Docstring inherited from CollectionManager. 

95 return f"{prefix}_id" 

96 

97 @classmethod 

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

99 # Docstring inherited from CollectionManager. 

100 return f"{prefix}_id" 

101 

102 @classmethod 

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

104 onDelete: Optional[str] = None, 

105 constraint: bool = True, 

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

107 # Docstring inherited from CollectionManager. 

108 original = _TABLES_SPEC.collection.fields["collection_id"] 

109 copy = ddl.FieldSpec(cls.getCollectionForeignKeyName(prefix), dtype=original.dtype, **kwargs) 

110 tableSpec.fields.add(copy) 

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

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

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

114 return copy 

115 

116 @classmethod 

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

118 onDelete: Optional[str] = None, 

119 constraint: bool = True, 

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

121 # Docstring inherited from CollectionManager. 

122 original = _TABLES_SPEC.run.fields["collection_id"] 

123 copy = ddl.FieldSpec(cls.getRunForeignKeyName(prefix), dtype=original.dtype, **kwargs) 

124 tableSpec.fields.add(copy) 

125 if constraint: 

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

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

128 return copy 

129 

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

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

132 old cached records will be removed. 

133 """ 

134 self._records = {} 

135 self._nameCache = {} 

136 for record in records: 

137 self._records[record.key] = record 

138 self._nameCache[record.name] = record 

139 

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

141 """Add single record to cache. 

142 """ 

143 self._records[record.key] = record 

144 self._nameCache[record.name] = record 

145 

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

147 """Remove single record from cache. 

148 """ 

149 del self._records[record.key] 

150 del self._nameCache[record.name] 

151 

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

153 # Docstring inherited from DefaultCollectionManager. 

154 return self._nameCache.get(name) 

155 

156 @classmethod 

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

158 # Docstring inherited from VersionedExtension. 

159 return _VERSION 

160 

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

162 # Docstring inherited from VersionedExtension. 

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