Coverage for python/lsst/daf/butler/registry/collections/nameKey.py: 96%

46 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-01-04 02:04 -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__ = ["NameKeyCollectionManager"] 

24 

25from typing import TYPE_CHECKING, Any 

26 

27import sqlalchemy 

28 

29from ...core import TimespanDatabaseRepresentation, ddl 

30from ..interfaces import 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 CollectionRecord, Database, DimensionRecordStorageManager, StaticTablesContext 

40 

41 

42_KEY_FIELD_SPEC = ddl.FieldSpec("name", dtype=sqlalchemy.String, length=64, primaryKey=True) 

43 

44 

45# This has to be updated on every schema change 

46_VERSION = VersionTuple(2, 0, 0) 

47 

48 

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

50 return CollectionTablesTuple( 

51 collection=ddl.TableSpec( 

52 fields=[ 

53 _KEY_FIELD_SPEC, 

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

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

56 ], 

57 ), 

58 run=makeRunTableSpec("name", sqlalchemy.String, TimespanReprClass), 

59 collection_chain=makeCollectionChainTableSpec("name", sqlalchemy.String), 

60 ) 

61 

62 

63class NameKeyCollectionManager(DefaultCollectionManager): 

64 """A `CollectionManager` implementation that uses collection names for 

65 primary/foreign keys and aggressively loads all collection/run records in 

66 the database into memory. 

67 

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

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

70 table schema. 

71 """ 

72 

73 @classmethod 

74 def initialize( 

75 cls, 

76 db: Database, 

77 context: StaticTablesContext, 

78 *, 

79 dimensions: DimensionRecordStorageManager, 

80 ) -> NameKeyCollectionManager: 

81 # Docstring inherited from CollectionManager. 

82 return cls( 

83 db, 

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

85 collectionIdName="name", 

86 dimensions=dimensions, 

87 ) 

88 

89 @classmethod 

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

91 # Docstring inherited from CollectionManager. 

92 return f"{prefix}_name" 

93 

94 @classmethod 

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

96 # Docstring inherited from CollectionManager. 

97 return f"{prefix}_name" 

98 

99 @classmethod 

100 def addCollectionForeignKey( 

101 cls, 

102 tableSpec: ddl.TableSpec, 

103 *, 

104 prefix: str = "collection", 

105 onDelete: str | None = None, 

106 constraint: bool = True, 

107 **kwargs: Any, 

108 ) -> ddl.FieldSpec: 

109 # Docstring inherited from CollectionManager. 

110 original = _KEY_FIELD_SPEC 

111 copy = ddl.FieldSpec( 

112 cls.getCollectionForeignKeyName(prefix), dtype=original.dtype, length=original.length, **kwargs 

113 ) 

114 tableSpec.fields.add(copy) 

115 if constraint: 

116 tableSpec.foreignKeys.append( 

117 ddl.ForeignKeySpec( 

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

119 ) 

120 ) 

121 return copy 

122 

123 @classmethod 

124 def addRunForeignKey( 

125 cls, 

126 tableSpec: ddl.TableSpec, 

127 *, 

128 prefix: str = "run", 

129 onDelete: str | None = 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.getRunForeignKeyName(prefix), dtype=original.dtype, length=original.length, **kwargs 

137 ) 

138 tableSpec.fields.add(copy) 

139 if constraint: 

140 tableSpec.foreignKeys.append( 

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

142 ) 

143 return copy 

144 

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

146 # Docstring inherited from DefaultCollectionManager. 

147 return self._records.get(name) 

148 

149 @classmethod 

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

151 # Docstring inherited from VersionedExtension. 

152 return _VERSION 

153 

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

155 # Docstring inherited from VersionedExtension. 

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