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

42 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-06-14 09:10 +0000

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: 

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 registry_schema_version: VersionTuple | None = None, 

81 ) -> NameKeyCollectionManager: 

82 # Docstring inherited from CollectionManager. 

83 return cls( 

84 db, 

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

86 collectionIdName="name", 

87 dimensions=dimensions, 

88 registry_schema_version=registry_schema_version, 

89 ) 

90 

91 @classmethod 

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

93 # Docstring inherited from CollectionManager. 

94 return f"{prefix}_name" 

95 

96 @classmethod 

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

98 # Docstring inherited from CollectionManager. 

99 return f"{prefix}_name" 

100 

101 @classmethod 

102 def addCollectionForeignKey( 

103 cls, 

104 tableSpec: ddl.TableSpec, 

105 *, 

106 prefix: str = "collection", 

107 onDelete: str | None = None, 

108 constraint: bool = True, 

109 **kwargs: Any, 

110 ) -> ddl.FieldSpec: 

111 # Docstring inherited from CollectionManager. 

112 original = _KEY_FIELD_SPEC 

113 copy = ddl.FieldSpec( 

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

115 ) 

116 tableSpec.fields.add(copy) 

117 if constraint: 

118 tableSpec.foreignKeys.append( 

119 ddl.ForeignKeySpec( 

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

121 ) 

122 ) 

123 return copy 

124 

125 @classmethod 

126 def addRunForeignKey( 

127 cls, 

128 tableSpec: ddl.TableSpec, 

129 *, 

130 prefix: str = "run", 

131 onDelete: str | None = None, 

132 constraint: bool = True, 

133 **kwargs: Any, 

134 ) -> ddl.FieldSpec: 

135 # Docstring inherited from CollectionManager. 

136 original = _KEY_FIELD_SPEC 

137 copy = ddl.FieldSpec( 

138 cls.getRunForeignKeyName(prefix), dtype=original.dtype, length=original.length, **kwargs 

139 ) 

140 tableSpec.fields.add(copy) 

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

142 tableSpec.foreignKeys.append( 

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

144 ) 

145 return copy 

146 

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

148 # Docstring inherited from DefaultCollectionManager. 

149 return self._records.get(name) 

150 

151 @classmethod 

152 def currentVersions(cls) -> list[VersionTuple]: 

153 # Docstring inherited from VersionedExtension. 

154 return [_VERSION]