Coverage for python/lsst/daf/butler/registry/queries/_sql_query_backend.py: 36%

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 software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <http://www.gnu.org/licenses/>. 

27from __future__ import annotations 

28 

29__all__ = ("SqlQueryBackend",) 

30 

31import warnings 

32from collections.abc import Iterable, Mapping, Sequence, Set 

33from contextlib import AbstractContextManager 

34from typing import TYPE_CHECKING, Any, cast 

35 

36from lsst.utils.introspection import find_outside_stacklevel 

37 

38from ..._collection_type import CollectionType 

39from ..._dataset_type import DatasetType 

40from ..._exceptions import DataIdValueError 

41from ...dimensions import DimensionGroup, DimensionRecordSet, DimensionUniverse 

42from ...dimensions.record_cache import DimensionRecordCache 

43from ..interfaces import CollectionRecord, Database 

44from ._query_backend import QueryBackend 

45from ._sql_query_context import SqlQueryContext 

46 

47if TYPE_CHECKING: 

48 from ..managers import RegistryManagerInstances 

49 

50 

51class SqlQueryBackend(QueryBackend[SqlQueryContext]): 

52 """An implementation of `QueryBackend` for `SqlRegistry`. 

53 

54 Parameters 

55 ---------- 

56 db : `Database` 

57 Object that abstracts the database engine. 

58 managers : `RegistryManagerInstances` 

59 Struct containing the manager objects that back a `SqlRegistry`. 

60 dimension_record_cache : `DimensionRecordCache` 

61 Cache of all records for dimension elements with 

62 `~DimensionElement.is_cached` `True`. 

63 """ 

64 

65 def __init__( 

66 self, db: Database, managers: RegistryManagerInstances, dimension_record_cache: DimensionRecordCache 

67 ): 

68 self._db = db 

69 self._managers = managers 

70 self._dimension_record_cache = dimension_record_cache 

71 

72 @property 

73 def universe(self) -> DimensionUniverse: 

74 # Docstring inherited. 

75 return self._managers.dimensions.universe 

76 

77 def caching_context(self) -> AbstractContextManager[None]: 

78 # Docstring inherited. 

79 return self._managers.caching_context_manager() 

80 

81 def context(self) -> SqlQueryContext: 

82 # Docstring inherited. 

83 return SqlQueryContext(self._db, self._managers.column_types) 

84 

85 def get_collection_name(self, key: Any) -> str: 

86 assert self._managers.caching_context.collection_records is not None, ( 

87 "Collection-record caching should already been enabled any time this is called." 

88 ) 

89 return self._managers.collections[key].name 

90 

91 def resolve_collection_wildcard( 

92 self, 

93 expression: Any, 

94 *, 

95 collection_types: Set[CollectionType] = CollectionType.all(), 

96 flatten_chains: bool = True, 

97 include_chains: bool | None = None, 

98 ) -> list[CollectionRecord]: 

99 # Docstring inherited. 

100 return self._managers.collections.resolve_wildcard( 

101 expression, 

102 collection_types=collection_types, 

103 flatten_chains=flatten_chains, 

104 include_chains=include_chains, 

105 ) 

106 

107 def resolve_dataset_type_wildcard( 

108 self, 

109 expression: Any, 

110 missing: list[str] | None = None, 

111 explicit_only: bool = False, 

112 ) -> list[DatasetType]: 

113 # Docstring inherited. 

114 return self._managers.datasets.resolve_wildcard( 

115 expression, 

116 missing, 

117 explicit_only, 

118 ) 

119 

120 def filter_dataset_collections( 

121 self, 

122 dataset_types: Iterable[DatasetType], 

123 collections: Sequence[CollectionRecord], 

124 *, 

125 governor_constraints: Mapping[str, Set[str]], 

126 rejections: list[str] | None = None, 

127 ) -> dict[DatasetType, list[CollectionRecord]]: 

128 # Docstring inherited. 

129 result: dict[DatasetType, list[CollectionRecord]] = { 

130 dataset_type: [] for dataset_type in dataset_types 

131 } 

132 summaries = self._managers.datasets.fetch_summaries(collections, result.keys()) 

133 for dataset_type, filtered_collections in result.items(): 

134 for collection_record in collections: 

135 if not dataset_type.isCalibration() and collection_record.type is CollectionType.CALIBRATION: 

136 if rejections is not None: 

137 rejections.append( 

138 f"Not searching for non-calibration dataset of type {dataset_type.name!r} " 

139 f"in CALIBRATION collection {collection_record.name!r}." 

140 ) 

141 else: 

142 collection_summary = summaries[collection_record.key] 

143 if collection_summary.is_compatible_with( 

144 dataset_type, 

145 governor_constraints, 

146 rejections=rejections, 

147 name=collection_record.name, 

148 ): 

149 filtered_collections.append(collection_record) 

150 return result 

151 

152 def resolve_governor_constraints( 

153 self, dimensions: DimensionGroup, constraints: Mapping[str, Set[str]] 

154 ) -> Mapping[str, Set[str]]: 

155 # Docstring inherited. 

156 result: dict[str, Set[str]] = {} 

157 for dimension_name in dimensions.governors: 

158 all_values = { 

159 cast(str, record.dataId[dimension_name]) 

160 for record in self._dimension_record_cache[dimension_name] 

161 } 

162 if (constraint_values := constraints.get(dimension_name)) is not None: 

163 if not (constraint_values <= all_values): 

164 warnings.warn( 

165 "DataIdValueError will no longer be raised for invalid governor dimension values." 

166 " Instead, an empty list will be returned. Will be changed after v28.", 

167 FutureWarning, 

168 stacklevel=find_outside_stacklevel("lsst.daf.butler"), 

169 ) 

170 raise DataIdValueError( 

171 f"Unknown values specified for governor dimension {dimension_name}: " 

172 f"{constraint_values - all_values}." 

173 ) 

174 result[dimension_name] = constraint_values 

175 else: 

176 result[dimension_name] = all_values 

177 return result 

178 

179 def get_dimension_record_cache(self, element_name: str) -> DimensionRecordSet | None: 

180 return ( 

181 self._dimension_record_cache[element_name] 

182 if element_name in self._dimension_record_cache 

183 else None 

184 )