Coverage for python/lsst/daf/butler/registry/dimensions/caching.py: 96%

57 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-12-02 02:00 +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__ = ["CachingDimensionRecordStorage"] 

24 

25from typing import Any, Dict, Iterable, Mapping, Optional, Set, Union 

26 

27import sqlalchemy 

28from lsst.utils import doImportType 

29 

30from ...core import ( 

31 DatabaseDimensionElement, 

32 DataCoordinate, 

33 DataCoordinateIterable, 

34 DataCoordinateSet, 

35 DimensionElement, 

36 DimensionRecord, 

37 GovernorDimension, 

38 NamedKeyDict, 

39 NamedKeyMapping, 

40 SpatialRegionDatabaseRepresentation, 

41 TimespanDatabaseRepresentation, 

42) 

43from ..interfaces import ( 

44 Database, 

45 DatabaseDimensionRecordStorage, 

46 GovernorDimensionRecordStorage, 

47 StaticTablesContext, 

48) 

49from ..queries import QueryBuilder 

50 

51 

52class CachingDimensionRecordStorage(DatabaseDimensionRecordStorage): 

53 """A record storage implementation that adds caching to some other nested 

54 storage implementation. 

55 

56 Parameters 

57 ---------- 

58 nested : `DatabaseDimensionRecordStorage` 

59 The other storage to cache fetches from and to delegate all other 

60 operations to. 

61 """ 

62 

63 def __init__(self, nested: DatabaseDimensionRecordStorage): 

64 self._nested = nested 

65 self._cache: Dict[DataCoordinate, Optional[DimensionRecord]] = {} 

66 

67 @classmethod 

68 def initialize( 

69 cls, 

70 db: Database, 

71 element: DatabaseDimensionElement, 

72 *, 

73 context: Optional[StaticTablesContext] = None, 

74 config: Mapping[str, Any], 

75 governors: NamedKeyMapping[GovernorDimension, GovernorDimensionRecordStorage], 

76 ) -> DatabaseDimensionRecordStorage: 

77 # Docstring inherited from DatabaseDimensionRecordStorage. 

78 config = config["nested"] 

79 NestedClass = doImportType(config["cls"]) 

80 if not hasattr(NestedClass, "initialize"): 80 ↛ 81line 80 didn't jump to line 81, because the condition on line 80 was never true

81 raise TypeError(f"Nested class {config['cls']} does not have an initialize() method.") 

82 nested = NestedClass.initialize(db, element, context=context, config=config, governors=governors) 

83 return cls(nested) 

84 

85 @property 

86 def element(self) -> DatabaseDimensionElement: 

87 # Docstring inherited from DimensionRecordStorage.element. 

88 return self._nested.element 

89 

90 def clearCaches(self) -> None: 

91 # Docstring inherited from DimensionRecordStorage.clearCaches. 

92 self._cache.clear() 

93 self._nested.clearCaches() 

94 

95 def join( 

96 self, 

97 builder: QueryBuilder, 

98 *, 

99 regions: Optional[NamedKeyDict[DimensionElement, SpatialRegionDatabaseRepresentation]] = None, 

100 timespans: Optional[NamedKeyDict[DimensionElement, TimespanDatabaseRepresentation]] = None, 

101 ) -> None: 

102 # Docstring inherited from DimensionRecordStorage. 

103 return self._nested.join(builder, regions=regions, timespans=timespans) 

104 

105 def insert(self, *records: DimensionRecord, replace: bool = False, skip_existing: bool = False) -> None: 

106 # Docstring inherited from DimensionRecordStorage.insert. 

107 self._nested.insert(*records, replace=replace, skip_existing=skip_existing) 

108 for record in records: 

109 # We really shouldn't ever get into a situation where the record 

110 # here differs from the one in the DB, but the last thing we want 

111 # is to make it harder to debug by making the cache different from 

112 # the DB. 

113 if skip_existing: 

114 self._cache.setdefault(record.dataId, record) 

115 else: 

116 self._cache[record.dataId] = record 

117 

118 def sync(self, record: DimensionRecord, update: bool = False) -> Union[bool, Dict[str, Any]]: 

119 # Docstring inherited from DimensionRecordStorage.sync. 

120 inserted_or_updated = self._nested.sync(record, update=update) 

121 if inserted_or_updated: 

122 self._cache[record.dataId] = record 

123 return inserted_or_updated 

124 

125 def fetch(self, dataIds: DataCoordinateIterable) -> Iterable[DimensionRecord]: 

126 # Docstring inherited from DimensionRecordStorage.fetch. 

127 missing: Set[DataCoordinate] = set() 

128 for dataId in dataIds: 

129 # Use ... as sentinal value so we can also cache None == "no such 

130 # record exists". 

131 record = self._cache.get(dataId, ...) 

132 if record is ...: 

133 missing.add(dataId) 

134 elif record is not None: 134 ↛ 128line 134 didn't jump to line 128, because the condition on line 134 was never false

135 # Unclear why MyPy can't tell that this isn't ..., but it 

136 # thinks it's still a possibility. 

137 yield record 

138 if missing: 

139 toFetch = DataCoordinateSet(missing, graph=self.element.graph) 

140 for record in self._nested.fetch(toFetch): 

141 self._cache[record.dataId] = record 

142 yield record 

143 missing -= self._cache.keys() 

144 for dataId in missing: 

145 self._cache[dataId] = None 

146 

147 def digestTables(self) -> Iterable[sqlalchemy.schema.Table]: 

148 # Docstring inherited from DimensionRecordStorage.digestTables. 

149 return self._nested.digestTables()