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

57 statements  

« prev     ^ index     » next       coverage.py v6.4.4, created at 2022-09-22 02:04 -0700

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

41) 

42from ..interfaces import ( 

43 Database, 

44 DatabaseDimensionRecordStorage, 

45 GovernorDimensionRecordStorage, 

46 StaticTablesContext, 

47) 

48from ..queries import QueryBuilder 

49 

50 

51class CachingDimensionRecordStorage(DatabaseDimensionRecordStorage): 

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

53 storage implementation. 

54 

55 Parameters 

56 ---------- 

57 nested : `DatabaseDimensionRecordStorage` 

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

59 operations to. 

60 """ 

61 

62 def __init__(self, nested: DatabaseDimensionRecordStorage): 

63 self._nested = nested 

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

65 

66 @classmethod 

67 def initialize( 

68 cls, 

69 db: Database, 

70 element: DatabaseDimensionElement, 

71 *, 

72 context: Optional[StaticTablesContext] = None, 

73 config: Mapping[str, Any], 

74 governors: NamedKeyMapping[GovernorDimension, GovernorDimensionRecordStorage], 

75 ) -> DatabaseDimensionRecordStorage: 

76 # Docstring inherited from DatabaseDimensionRecordStorage. 

77 config = config["nested"] 

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

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

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

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

82 return cls(nested) 

83 

84 @property 

85 def element(self) -> DatabaseDimensionElement: 

86 # Docstring inherited from DimensionRecordStorage.element. 

87 return self._nested.element 

88 

89 def clearCaches(self) -> None: 

90 # Docstring inherited from DimensionRecordStorage.clearCaches. 

91 self._cache.clear() 

92 self._nested.clearCaches() 

93 

94 def join( 

95 self, 

96 builder: QueryBuilder, 

97 *, 

98 regions: Optional[NamedKeyDict[DimensionElement, sqlalchemy.sql.ColumnElement]] = None, 

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

100 ) -> None: 

101 # Docstring inherited from DimensionRecordStorage. 

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

103 

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

105 # Docstring inherited from DimensionRecordStorage.insert. 

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

107 for record in records: 

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

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

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

111 # the DB. 

112 if skip_existing: 

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

114 else: 

115 self._cache[record.dataId] = record 

116 

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

118 # Docstring inherited from DimensionRecordStorage.sync. 

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

120 if inserted_or_updated: 

121 self._cache[record.dataId] = record 

122 return inserted_or_updated 

123 

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

125 # Docstring inherited from DimensionRecordStorage.fetch. 

126 missing: Set[DataCoordinate] = set() 

127 for dataId in dataIds: 

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

129 # record exists". 

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

131 if record is ...: 

132 missing.add(dataId) 

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

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

135 # thinks it's still a possibility. 

136 yield record # type: ignore 

137 if missing: 

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

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

140 self._cache[record.dataId] = record 

141 yield record 

142 missing -= self._cache.keys() 

143 for dataId in missing: 

144 self._cache[dataId] = None 

145 

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

147 # Docstring inherited from DimensionRecordStorage.digestTables. 

148 return self._nested.digestTables()