Coverage for python/lsst/daf/butler/registry/queries/_sql_query

1# This file is part of daf_butler.

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.

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

29__all__ = ("SqlQueryBackend",)

31import warnings

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

33from contextlib import AbstractContextManager

34from typing import TYPE_CHECKING, Any, cast

36from lsst.daf.relation import Relation

37from lsst.utils.introspection import find_outside_stacklevel

39from ..._collection_type import CollectionType

40from ..._dataset_type import DatasetType

41from ..._exceptions import DataIdValueError, MissingDatasetTypeError

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

43from ...dimensions.record_cache import DimensionRecordCache

44from ..interfaces import CollectionRecord, Database

45from ._query_backend import QueryBackend

46from ._sql_query_context import SqlQueryContext

48if TYPE_CHECKING:

49 from ..managers import RegistryManagerInstances

52class SqlQueryBackend(QueryBackend[SqlQueryContext]):

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

55 Parameters

56 ----------

57 db : `Database`

58 Object that abstracts the database engine.

59 managers : `RegistryManagerInstances`

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

61 dimension_record_cache : `DimensionRecordCache`

62 Cache of all records for dimension elements with

63 `~DimensionElement.is_cached` `True`.

64 """

66 def __init__(

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

68 ):

69 self._db = db

70 self._managers = managers

71 self._dimension_record_cache = dimension_record_cache

73 @property

74 def universe(self) -> DimensionUniverse:

75 # Docstring inherited.

76 return self._managers.dimensions.universe

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

79 # Docstring inherited.

80 return self._managers.caching_context_manager()

82 def context(self) -> SqlQueryContext:

83 # Docstring inherited.

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

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

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

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

89 )

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

92 def resolve_collection_wildcard(

93 self,

94 expression: Any,

95 *,

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

97 flatten_chains: bool = True,

98 include_chains: bool | None = None,

99 ) -> list[CollectionRecord]:

100 # Docstring inherited.

101 return self._managers.collections.resolve_wildcard(

102 expression,

103 collection_types=collection_types,

104 flatten_chains=flatten_chains,

105 include_chains=include_chains,

106 )

107

108 def resolve_dataset_type_wildcard(

109 self,

110 expression: Any,

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

112 explicit_only: bool = False,

113 ) -> list[DatasetType]:

114 # Docstring inherited.

115 return self._managers.datasets.resolve_wildcard(

116 expression,

117 missing,

118 explicit_only,

119 )

120

121 def filter_dataset_collections(

122 self,

123 dataset_types: Iterable[DatasetType],

124 collections: Sequence[CollectionRecord],

125 *,

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

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

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

129 # Docstring inherited.

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

131 dataset_type: [] for dataset_type in dataset_types

132 }

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

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

135 for collection_record in collections:

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

137 if rejections is not None:

138 rejections.append(

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

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

141 )

142 else:

143 collection_summary = summaries[collection_record.key]

144 if collection_summary.is_compatible_with(

145 dataset_type,

146 governor_constraints,

147 rejections=rejections,

148 name=collection_record.name,

149 ):

150 filtered_collections.append(collection_record)

151 return result

152

153 def _make_dataset_query_relation_impl(

154 self,

155 dataset_type: DatasetType,

156 collections: Sequence[CollectionRecord],

157 columns: Set[str],

158 context: SqlQueryContext,

159 ) -> Relation:

160 # Docstring inherited.

161 assert len(collections) > 0, (

162 "Caller is responsible for handling the case of all collections being rejected (we can't "

163 "write a good error message without knowing why collections were rejected)."

164 )

165 try:

166 return self._managers.datasets.make_relation(

167 dataset_type,

168 *collections,

169 columns=columns,

170 context=context,

171 )

172 except MissingDatasetTypeError:

173 pass

174 # Unrecognized dataset type means no results.

175 return self.make_doomed_dataset_relation(

176 dataset_type,

177 columns,

178 messages=[

179 f"Dataset type {dataset_type.name!r} is not registered, "

180 "so no instances of it can exist in any collection."

181 ],

182 context=context,

183 )

184

185 def resolve_governor_constraints(

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

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

188 # Docstring inherited.

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

190 for dimension_name in dimensions.governors:

191 all_values = {

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

193 for record in self._dimension_record_cache[dimension_name]

194 }

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

196 if not (constraint_values <= all_values):

197 warnings.warn(

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

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

200 FutureWarning,

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

202 )

203 raise DataIdValueError(

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

205 f"{constraint_values - all_values}."

206 )

207 result[dimension_name] = constraint_values

208 else:

209 result[dimension_name] = all_values

210 return result

211

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

213 return (

214 self._dimension_record_cache[element_name]

215 if element_name in self._dimension_record_cache

216 else None

217 )

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

67 statements