Coverage for python/lsst/daf/butler/direct_query_driver/_driver.py: 87%

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/>. 

27 

28from __future__ import annotations 

29 

30__all__ = ("DirectQueryDriver",) 

31 

32import dataclasses 

33import itertools 

34import logging 

35import sys 

36import uuid 

37from collections import defaultdict 

38from collections.abc import Iterable, Iterator, Mapping, Set 

39from contextlib import ExitStack 

40from typing import TYPE_CHECKING, Any, TypeVar, cast, overload 

41 

42import sqlalchemy 

43 

44from .. import ddl 

45from .._collection_type import CollectionType 

46from .._dataset_type import DatasetType 

47from .._exceptions import InvalidQueryError 

48from ..dimensions import DataCoordinate, DataIdValue, DimensionElement, DimensionGroup, DimensionUniverse 

49from ..dimensions.record_cache import DimensionRecordCache 

50from ..queries import tree as qt 

51from ..queries.driver import ( 

52 DataCoordinateResultPage, 

53 DatasetRefResultPage, 

54 DimensionRecordResultPage, 

55 GeneralResultPage, 

56 QueryDriver, 

57 ResultPage, 

58) 

59from ..queries.predicate_constraints_summary import PredicateConstraintsSummary 

60from ..queries.result_specs import ( 

61 DataCoordinateResultSpec, 

62 DatasetRefResultSpec, 

63 DimensionRecordResultSpec, 

64 GeneralResultSpec, 

65 ResultSpec, 

66) 

67from ..registry import CollectionSummary, NoDefaultCollectionError 

68from ..registry.interfaces import ChainedCollectionRecord, CollectionRecord 

69from ..registry.managers import RegistryManagerInstances 

70from ..registry.wildcards import CollectionWildcard 

71from ._postprocessing import Postprocessing 

72from ._query_analysis import ( 

73 QueryCollectionAnalysis, 

74 QueryFindFirstAnalysis, 

75 QueryJoinsAnalysis, 

76 QueryTreeAnalysis, 

77 ResolvedDatasetSearch, 

78) 

79from ._query_builder import QueryBuilder, SingleSelectQueryBuilder, UnionQueryBuilder 

80from ._result_page_converter import ( 

81 DataCoordinateResultPageConverter, 

82 DatasetRefResultPageConverter, 

83 DimensionRecordResultPageConverter, 

84 GeneralResultPageConverter, 

85 ResultPageConverter, 

86 ResultPageConverterContext, 

87) 

88from ._sql_builders import SqlJoinsBuilder, SqlSelectBuilder, make_table_spec 

89from ._sql_column_visitor import SqlColumnVisitor 

90 

91if TYPE_CHECKING: 

92 from ..registry.interfaces import Database 

93 

94 

95_LOG = logging.getLogger(__name__) 

96 

97_T = TypeVar("_T", bound=str | qt.AnyDatasetType) 

98 

99 

100class DirectQueryDriver(QueryDriver): 

101 """The `QueryDriver` implementation for `DirectButler`. 

102 

103 Parameters 

104 ---------- 

105 db : `Database` 

106 Abstraction for the SQL database. 

107 universe : `DimensionUniverse` 

108 Definitions of all dimensions. 

109 managers : `RegistryManagerInstances` 

110 Struct of registry manager objects. 

111 dimension_record_cache : `DimensionRecordCache` 

112 Cache of dimension records for infrequently-changing, commonly-used 

113 dimensions. 

114 default_collections : `~collections.abc.Sequence` [ `str` ] 

115 Default collection search path. 

116 default_data_id : `DataCoordinate` 

117 Default governor dimension values. 

118 raw_page_size : `int`, optional 

119 Number of database rows to fetch for each result page. The actual 

120 number of rows in a page may be smaller due to postprocessing. 

121 constant_rows_limit : `int`, optional 

122 Maximum number of uploaded rows to include in queries via 

123 `Database.constant_rows`; above this limit a temporary table is used 

124 instead. 

125 postprocessing_filter_factor : `int`, optional 

126 The number of database rows we expect to have to fetch to yield a 

127 single output row for queries that involve postprocessing. This is 

128 purely a performance tuning parameter that attempts to balance between 

129 fetching too much and requiring multiple fetches; the true value is 

130 highly dependent on the actual query. 

131 """ 

132 

133 def __init__( 

134 self, 

135 db: Database, 

136 universe: DimensionUniverse, 

137 managers: RegistryManagerInstances, 

138 dimension_record_cache: DimensionRecordCache, 

139 default_collections: Iterable[str], 

140 default_data_id: DataCoordinate, 

141 # Increasing raw_page_size increases memory usage for queries on 

142 # Butler server, so if you increase this you may need to increase the 

143 # memory allocation for the server in Phalanx as well. 

144 raw_page_size: int = 2000, 

145 constant_rows_limit: int = 1000, 

146 postprocessing_filter_factor: int = 10, 

147 ): 

148 self.db = db 

149 self.managers = managers 

150 self._dimension_record_cache = dimension_record_cache 

151 self._universe = universe 

152 self._default_collections = tuple(default_collections) 

153 self._default_data_id = default_data_id 

154 self._materializations: dict[qt.MaterializationKey, _MaterializationState] = {} 

155 self._upload_tables: dict[qt.DataCoordinateUploadKey, sqlalchemy.FromClause] = {} 

156 self._exit_stack: ExitStack | None = None 

157 self._raw_page_size = raw_page_size 

158 self._postprocessing_filter_factor = postprocessing_filter_factor 

159 self._constant_rows_limit = min(constant_rows_limit, db.get_constant_rows_max()) 

160 self._cursors: set[_Cursor] = set() 

161 

162 def __enter__(self) -> None: 

163 self._exit_stack = ExitStack() 

164 # It might be nice to defer opening a transaction here until first use 

165 # to reduce the time spent in transactions. But it's worth noting that 

166 # this is the default low-level behavior of the Python SQLite driver, 

167 # and it makes it incredibly prone to deadlocks. We might be okay 

168 # here, because Query doesn't do true write operations - just temp 

169 # table writes - but I'm not confident that's enough to make delayed 

170 # transaction starts safe against deadlocks, and it'd be more 

171 # complicated to implement anyway. 

172 # 

173 # We start a transaction rather than just opening a connection to make 

174 # temp table and cursors work with pg_bouncer transaction affinity. 

175 self._exit_stack.enter_context(self.db.transaction(for_temp_tables=True)) 

176 

177 def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: 

178 assert self._exit_stack is not None 

179 self._materializations.clear() 

180 self._upload_tables.clear() 

181 # Transfer open cursors' close methods to exit stack, this will help 

182 # with the cleanup in case a cursor raises an exceptions on close. 

183 for cursor in self._cursors: 

184 self._exit_stack.push(cursor.close) 

185 self._exit_stack.__exit__(exc_type, exc_value, traceback) 

186 self._cursors = set() 

187 self._exit_stack = None 

188 

189 @property 

190 def universe(self) -> DimensionUniverse: 

191 return self._universe 

192 

193 @overload 

194 def execute( 194 ↛ exitline 194 didn't return from function 'execute' because

195 self, result_spec: DataCoordinateResultSpec, tree: qt.QueryTree 

196 ) -> Iterator[DataCoordinateResultPage]: ... 

197 

198 @overload 

199 def execute( 199 ↛ exitline 199 didn't return from function 'execute' because

200 self, result_spec: DimensionRecordResultSpec, tree: qt.QueryTree 

201 ) -> Iterator[DimensionRecordResultPage]: ... 

202 

203 @overload 

204 def execute( 204 ↛ exitline 204 didn't return from function 'execute' because

205 self, result_spec: DatasetRefResultSpec, tree: qt.QueryTree 

206 ) -> Iterator[DatasetRefResultPage]: ... 

207 

208 @overload 

209 def execute(self, result_spec: GeneralResultSpec, tree: qt.QueryTree) -> Iterator[GeneralResultPage]: ... 209 ↛ exitline 209 didn't return from function 'execute' because

210 

211 def execute(self, result_spec: ResultSpec, tree: qt.QueryTree) -> Iterator[ResultPage]: 

212 # Docstring inherited. 

213 if self._exit_stack is None: 213 ↛ 214line 213 didn't jump to line 214 because the condition on line 213 was never true

214 raise RuntimeError("QueryDriver context must be entered before queries can be executed.") 

215 builder = self.build_query( 

216 tree, 

217 final_columns=result_spec.get_result_columns(), 

218 order_by=result_spec.order_by, 

219 find_first_dataset=result_spec.find_first_dataset, 

220 allow_duplicate_overlaps=result_spec.allow_duplicate_overlaps, 

221 ) 

222 sql_select, sql_columns = builder.finish_select() 

223 if result_spec.order_by: 

224 visitor = SqlColumnVisitor(sql_columns, self) 

225 sql_select = sql_select.order_by(*[visitor.expect_scalar(term) for term in result_spec.order_by]) 

226 if result_spec.limit is not None: 

227 if builder.postprocessing: 

228 builder.postprocessing.limit = result_spec.limit 

229 else: 

230 sql_select = sql_select.limit(result_spec.limit) 

231 if builder.postprocessing.limit is not None: 

232 # We might want to fetch many fewer rows than the default page 

233 # size if we have to implement limit in postprocessing. 

234 raw_page_size = min( 

235 self._postprocessing_filter_factor * builder.postprocessing.limit, 

236 self._raw_page_size, 

237 ) 

238 else: 

239 raw_page_size = self._raw_page_size 

240 # Execute the query by initializing a _Cursor object that manages the 

241 # lifetime of the result. 

242 cursor = _Cursor( 

243 self.db, 

244 sql_select, 

245 postprocessing=builder.postprocessing, 

246 raw_page_size=raw_page_size, 

247 page_converter=self._create_result_page_converter(result_spec, builder.final_columns), 

248 ) 

249 # Since this function isn't a context manager and the caller could stop 

250 # iterating before we retrieve all the results, we have to track open 

251 # cursors to ensure we can close them as part of higher-level cleanup. 

252 self._cursors.add(cursor) 

253 

254 # Return the iterator as a separate function, so that all the code 

255 # above runs immediately instead of later when we first read from the 

256 # iterator. This ensures that any exceptions that are triggered during 

257 # set-up for this query occur immediately. 

258 return self._read_results(cursor) 

259 

260 def _read_results(self, cursor: _Cursor) -> Iterator[ResultPage]: 

261 """Read out all of the result pages from the database.""" 

262 try: 

263 while (result_page := cursor.next()) is not None: 

264 yield result_page 

265 finally: 

266 self._cursors.discard(cursor) 

267 cursor.close() 

268 

269 def _create_result_page_converter(self, spec: ResultSpec, columns: qt.ColumnSet) -> ResultPageConverter: 

270 context = ResultPageConverterContext( 

271 db=self.db, 

272 column_order=columns.get_column_order(), 

273 dimension_record_cache=self._dimension_record_cache, 

274 ) 

275 match spec: 

276 case DimensionRecordResultSpec(): 

277 return DimensionRecordResultPageConverter(spec, context) 

278 case DataCoordinateResultSpec(): 

279 return DataCoordinateResultPageConverter(spec, context) 

280 case DatasetRefResultSpec(): 

281 return DatasetRefResultPageConverter( 

282 spec, self.get_dataset_type(spec.dataset_type_name), context 

283 ) 

284 case GeneralResultSpec(): 

285 return GeneralResultPageConverter(spec, context) 

286 case _: 

287 raise NotImplementedError(f"Result type '{spec.result_type}' not yet implemented") 

288 

289 def materialize( 

290 self, 

291 tree: qt.QueryTree, 

292 dimensions: DimensionGroup, 

293 datasets: frozenset[str], 

294 allow_duplicate_overlaps: bool = False, 

295 key: qt.MaterializationKey | None = None, 

296 ) -> qt.MaterializationKey: 

297 # Docstring inherited. 

298 if self._exit_stack is None: 298 ↛ 299line 298 didn't jump to line 299 because the condition on line 298 was never true

299 raise RuntimeError("QueryDriver context must be entered before 'materialize' is called.") 

300 plan = self.build_query( 

301 tree, qt.ColumnSet(dimensions), allow_duplicate_overlaps=allow_duplicate_overlaps 

302 ) 

303 # Current implementation ignores 'datasets' aside from remembering 

304 # them, because figuring out what to put in the temporary table for 

305 # them is tricky, especially if calibration collections are involved. 

306 # That's okay because: 

307 # 

308 # - the query whose results we materialize includes the dataset 

309 # searches as constraints; 

310 # 

311 # - we still (in Query.materialize) join the dataset searches back in 

312 # anyway, and given materialized data IDs the join to the dataset 

313 # search is straightforward and definitely well-indexed, and not much 

314 # (if at all) worse than joining back in on a materialized UUID. 

315 # 

316 sql_select, _ = plan.finish_select(return_columns=False) 

317 table = self._exit_stack.enter_context( 

318 self.db.temporary_table( 

319 make_table_spec(plan.final_columns, self.db, plan.postprocessing, make_indices=True) 

320 ) 

321 ) 

322 self.db.insert(table, select=sql_select) 

323 if key is None: 323 ↛ 325line 323 didn't jump to line 325 because the condition on line 323 was always true

324 key = uuid.uuid4() 

325 self._materializations[key] = _MaterializationState(table, datasets, plan.postprocessing) 

326 return key 

327 

328 def upload_data_coordinates( 

329 self, 

330 dimensions: DimensionGroup, 

331 rows: Iterable[tuple[DataIdValue, ...]], 

332 key: qt.DataCoordinateUploadKey | None = None, 

333 ) -> qt.DataCoordinateUploadKey: 

334 # Docstring inherited. 

335 if self._exit_stack is None: 335 ↛ 336line 335 didn't jump to line 336 because the condition on line 335 was never true

336 raise RuntimeError( 

337 "QueryDriver context must be entered before 'upload_data_coordinates' is called." 

338 ) 

339 columns = qt.ColumnSet(dimensions).drop_implied_dimension_keys() 

340 table_spec = ddl.TableSpec( 

341 [columns.get_column_spec(logical_table, field).to_sql_spec() for logical_table, field in columns] 

342 ) 

343 dict_rows: list[dict[str, Any]] 

344 if not columns: 

345 table_spec.fields.add( 

346 ddl.FieldSpec( 

347 SqlSelectBuilder.EMPTY_COLUMNS_NAME, 

348 dtype=SqlSelectBuilder.EMPTY_COLUMNS_TYPE, 

349 nullable=True, 

350 ) 

351 ) 

352 dict_rows = [{SqlSelectBuilder.EMPTY_COLUMNS_NAME: None}] 

353 else: 

354 dict_rows = [dict(zip(dimensions.required, values)) for values in rows] 

355 from_clause: sqlalchemy.FromClause 

356 if self.db.supports_temporary_tables and len(dict_rows) > self._constant_rows_limit: 

357 table = self._exit_stack.enter_context(self.db.temporary_table(table_spec)) 

358 self.db.insert(table, *dict_rows) 

359 from_clause = table 

360 else: 

361 from_clause = self.db.constant_rows(table_spec.fields, *dict_rows) 

362 if key is None: 362 ↛ 364line 362 didn't jump to line 364 because the condition on line 362 was always true

363 key = uuid.uuid4() 

364 self._upload_tables[key] = from_clause 

365 return key 

366 

367 def count( 

368 self, 

369 tree: qt.QueryTree, 

370 result_spec: ResultSpec, 

371 *, 

372 exact: bool, 

373 discard: bool, 

374 ) -> int: 

375 # Docstring inherited. 

376 columns = result_spec.get_result_columns() 

377 builder = self.build_query(tree, columns, find_first_dataset=result_spec.find_first_dataset) 

378 if not all(d.collection_records for d in builder.joins_analysis.datasets.values()): 

379 return 0 

380 # No need to do similar check on 

381 if not exact: 

382 builder.postprocessing = Postprocessing() 

383 if builder.postprocessing: 

384 if not discard: 

385 raise InvalidQueryError("Cannot count query rows exactly without discarding them.") 

386 sql_select, _ = builder.finish_select(return_columns=False) 

387 builder.postprocessing.limit = result_spec.limit 

388 n = 0 

389 with self.db.query(sql_select.execution_options(yield_per=self._raw_page_size)) as results: 

390 for _ in builder.postprocessing.apply(results): 

391 n += 1 

392 return n 

393 # If the query has DISTINCT, GROUP BY, or UNION [ALL], nest it in a 

394 # subquery so we count deduplicated rows. 

395 select_builder = builder.finish_nested() 

396 # Replace the columns of the query with just COUNT(*). 

397 select_builder.columns = qt.ColumnSet(self._universe.empty) 

398 select_builder.joins.special.clear() 

399 count_func: sqlalchemy.ColumnElement[int] = sqlalchemy.func.count() 

400 select_builder.joins.special["_ROWCOUNT"] = count_func 

401 # Render and run the query. 

402 sql_select = select_builder.select(builder.postprocessing) 

403 with self.db.query(sql_select) as result: 

404 count = cast(int, result.scalar()) 

405 if result_spec.limit is not None: 

406 count = min(count, result_spec.limit) 

407 return count 

408 

409 def any(self, tree: qt.QueryTree, *, execute: bool, exact: bool) -> bool: 

410 # Docstring inherited. 

411 builder = self.build_query(tree, qt.ColumnSet(tree.dimensions), allow_duplicate_overlaps=True) 

412 if not all(d.collection_records for d in builder.joins_analysis.datasets.values()): 

413 return False 

414 if not execute: 

415 if exact: 

416 raise InvalidQueryError("Cannot obtain exact result for 'any' without executing.") 

417 return True 

418 if builder.postprocessing and exact: 

419 sql_select, _ = builder.finish_select(return_columns=False) 

420 with self.db.query( 

421 sql_select.execution_options(yield_per=self._postprocessing_filter_factor) 

422 ) as result: 

423 for _ in builder.postprocessing.apply(result): 

424 return True 

425 return False 

426 sql_select, _ = builder.finish_select() 

427 with self.db.query(sql_select.limit(1)) as result: 

428 return result.first() is not None 

429 

430 def explain_no_results(self, tree: qt.QueryTree, execute: bool) -> Iterable[str]: 

431 # Docstring inherited. 

432 plan = self.build_query(tree, qt.ColumnSet(tree.dimensions), analyze_only=True) 

433 if plan.joins_analysis.messages or not execute: 

434 return plan.joins_analysis.messages 

435 # TODO: guess at ways to split up query that might fail or succeed if 

436 # run separately, execute them with LIMIT 1 and report the results. 

437 return [] 

438 

439 def get_dataset_type(self, name: str) -> DatasetType: 

440 # Docstring inherited 

441 return self.managers.datasets.get_dataset_type(name) 

442 

443 def get_default_collections(self) -> tuple[str, ...]: 

444 # Docstring inherited. 

445 if not self._default_collections: 

446 raise NoDefaultCollectionError("No collections provided and no default collections.") 

447 return self._default_collections 

448 

449 def build_query( 

450 self, 

451 tree: qt.QueryTree, 

452 final_columns: qt.ColumnSet, 

453 *, 

454 order_by: Iterable[qt.OrderExpression] = (), 

455 find_first_dataset: str | qt.AnyDatasetType | None = None, 

456 analyze_only: bool = False, 

457 allow_duplicate_overlaps: bool = False, 

458 ) -> QueryBuilder: 

459 """Convert a query description into a nearly-complete builder object 

460 for the SQL version of that query. 

461 

462 Parameters 

463 ---------- 

464 tree : `.queries.tree.QueryTree` 

465 Description of the joins and row filters in the query. 

466 final_columns : `.queries.tree.ColumnSet` 

467 Final output columns that should be emitted by the SQL query. 

468 order_by : `~collections.abc.Iterable` [ \ 

469 `.queries.tree.OrderExpression` ], optional 

470 Column expressions to sort by. 

471 find_first_dataset : `str`, ``ANY_DATASET``, or `None`, optional 

472 Name of a dataset type for which only one result row for each data 

473 ID should be returned, with the collections searched in order. 

474 ``ANY_DATASET`` is used to represent the search for all dataset 

475 types with a particular set of dimensions in ``tree.any_dataset``. 

476 analyze_only : `bool`, optional 

477 If `True`, perform the initial analysis needed to construct the 

478 builder, but do not call methods that build its SQL form. This can 

479 be useful for obtaining diagnostic information about the query that 

480 would be generated. 

481 allow_duplicate_overlaps : `bool`, optional 

482 If set to `True` then query will be allowed to generate 

483 non-distinct rows for spatial overlaps. 

484 

485 Returns 

486 ------- 

487 builder : `QueryBuilder` 

488 An object that contains an analysis of the queries columns and 

489 general structure, SQLAlchemy representations of most of the 

490 constructed query, and a description of Python-side postprocessing 

491 to be performed after executing it. Callers should generally only 

492 have to call `finish_select` or `finish_nested`; all ``apply`` 

493 methods will have already been called. 

494 

495 Notes 

496 ----- 

497 The SQL queries produced by this driver are built in three steps: 

498 

499 - First we "analyze" the entire query, fetching extra information as 

500 needed to identify the tables we'll join in, the columns we'll have 

501 at each level of what may be a nested SELECT, and the degree to which 

502 we'll have to use GROUP BY / DISTINCT or window functions to obtain 

503 just the rows we want. This process initializes a `QueryBuilder` 

504 instance, but does not call any of its ``apply*`` methods to actually 

505 build the SQLAlchemy version of the query, and it is all that is done 

506 when ``analyze_only=True``. 

507 - In the next step we call `QueryBuilder` ``apply*`` methods to mutate 

508 and/or replace the `QueryBuilder` object's nested `SqlSelectBuilder` 

509 instances to reflect that analysis, building the SQL from the inside 

510 out (subqueries before their parents). This is also done by the 

511 `build_query` method by default 

512 - The returned builder can be used by calling either `finish_select` 

513 (to get the full executable query) or `finish_nested` (to wrap 

514 that query as a subquery - if needed - in order to do related 

515 queries, like ``COUNT(*)`` or ``LIMIT 1`` checks). 

516 

517 Within both the analysis and application steps, we further split the 

518 query *structure* and building process into three stages, mapping 

519 roughly to levels of (potential) subquery nesting: 

520 

521 - In the "joins" stage, we join all the tables we need columns or 

522 constraints from and apply the predicate as a WHERE clause. All 

523 non-calculated columns are included in the query at this stage. 

524 - In the optional "projection" stage, we apply a GROUP BY or DISTINCT 

525 to reduce the set of columns and/or rows, and in some cases add new 

526 calculated columns. 

527 - In the optional "find first" stage, we use a common table expression 

528 with PARTITION ON to search for datasets in a sequence of collections 

529 in the order those collections were provided. 

530 

531 In addition, there are two different overall structures for butler 

532 queries, modeled as two different `QueryBuilder` subclasses. 

533 

534 - `SingleSelectQueryBuilder` produces a single possibly-nested SELECT 

535 query structured directly according to the stages above. This is 

536 used for almost all queries - specifically whenever 

537 `QueryTree.any_dataset` is `None`. 

538 - `UnionQueryBuilder` implements queries for datasets of multiple types 

539 with the same dimensions as a UNION ALL combination of SELECTs, in 

540 which each component SELECT has the joins/projection/find-first 

541 structure. This cannot be implemented as a sequence of 

542 `SingleSelectQueryBuilder` objects, however, since the need for the 

543 terms to share a single `Postprocessing` object and column list means 

544 that they are not independent. 

545 

546 The `build_query` method delegates to methods of the `QueryBuilder` 

547 object when "applying" to let them handle these differences. These 

548 often call back to other methods on the driver (so code shared by both 

549 builders mostly lives in the driver class). 

550 """ 

551 # Analyze the dimensions, dataset searches, and other join operands 

552 # that will go into the query. This also initializes a 

553 # SqlSelectBuilder and Postprocessing with spatial/temporal constraints 

554 # potentially transformed by the dimensions manager (but none of the 

555 # rest of the analysis reflected in that SqlSelectBuilder). 

556 query_tree_analysis = self._analyze_query_tree(tree, allow_duplicate_overlaps) 

557 # The "projection" columns differ from the final columns by not 

558 # omitting any dimension keys (this keeps queries for different result 

559 # types more similar during construction), including any columns needed 

560 # only by order_by terms, and including the collection key if we need 

561 # it for GROUP BY or DISTINCT. 

562 projection_columns = final_columns.copy() 

563 projection_columns.restore_dimension_keys() 

564 for term in order_by: 

565 term.gather_required_columns(projection_columns) 

566 # There are two kinds of query builders: simple SELECTS and UNIONs 

567 # over dataset types. 

568 builder: QueryBuilder 

569 if tree.any_dataset is not None: 569 ↛ 570line 569 didn't jump to line 570 because the condition on line 569 was never true

570 builder = UnionQueryBuilder( 

571 query_tree_analysis, 

572 union_dataset_dimensions=tree.any_dataset.dimensions, 

573 projection_columns=projection_columns, 

574 final_columns=final_columns, 

575 find_first_dataset=find_first_dataset, 

576 ) 

577 else: 

578 assert find_first_dataset != qt.AnyDatasetType.ANY_DATASET 

579 builder = SingleSelectQueryBuilder( 

580 tree_analysis=query_tree_analysis, 

581 projection_columns=projection_columns, 

582 final_columns=final_columns, 

583 find_first_dataset=find_first_dataset, 

584 ) 

585 # Finish setting up the projection part of the builder. 

586 builder.analyze_projection() 

587 # The joins-stage query also needs to include all columns needed by the 

588 # downstream projection query. Note that this: 

589 # - never adds new dimensions to the joins stage (since those are 

590 # always a superset of the projection-stage dimensions); 

591 # - does not affect our previous determination of 

592 # needs_dataset_distinct, because any dataset fields being added to 

593 # the joins stage here are already in the projection. 

594 builder.joins_analysis.columns.update(builder.projection_columns) 

595 # Set up the find-first part of the builder. 

596 if find_first_dataset is not None: 

597 builder.analyze_find_first() 

598 # At this point, analysis is complete, and we can proceed to making 

599 # the select_builder(s) reflect that analysis. 

600 if not analyze_only: 

601 builder.apply_joins(self) 

602 builder.apply_projection(self, order_by) 

603 builder.apply_find_first(self) 

604 return builder 

605 

606 def _analyze_query_tree(self, tree: qt.QueryTree, allow_duplicate_overlaps: bool) -> QueryTreeAnalysis: 

607 """Analyze a `.queries.tree.QueryTree` as the first step in building 

608 a SQL query. 

609 

610 Parameters 

611 ---------- 

612 tree : `.queries.tree.QueryTree` 

613 Description of the joins and row filters in the query. 

614 allow_duplicate_overlaps : `bool`, optional 

615 If set to `True` then query will be allowed to generate 

616 non-distinct rows for spatial overlaps. 

617 

618 Returns 

619 ------- 

620 tree_analysis : `QueryTreeAnalysis` 

621 Struct containing additional information need to build the joins 

622 stage of a query. 

623 

624 Notes 

625 ----- 

626 See `build_query` for the complete picture of how SQL queries are 

627 constructed. This method is the very first step for all queries. 

628 

629 The fact that this method returns both a QueryPlan and an initial 

630 SqlSelectBuilder (rather than just a QueryPlan) is a tradeoff that lets 

631 DimensionRecordStorageManager.process_query_overlaps (which is called 

632 by the `_analyze_query_tree` call below) pull out overlap expressions 

633 from the predicate at the same time it turns them into SQL table joins 

634 (in the builder). 

635 """ 

636 # Fetch the records and summaries for any collections we might be 

637 # searching for datasets and organize them for the kind of lookups 

638 # we'll do later. 

639 collection_analysis = self._analyze_collections(tree) 

640 # Extract the data ID implied by the predicate; we can use the governor 

641 # dimensions in that to constrain the collections we search for 

642 # datasets later. 

643 predicate_constraints = PredicateConstraintsSummary(tree.predicate) 

644 # Use the default data ID to apply additional constraints where needed. 

645 predicate_constraints.apply_default_data_id( 

646 self._default_data_id, 

647 tree.dimensions, 

648 validate_governor_constraints=tree.validateGovernorConstraints, 

649 ) 

650 predicate = predicate_constraints.predicate 

651 # Delegate to the dimensions manager to rewrite the predicate and start 

652 # a SqlSelectBuilder to cover any spatial overlap joins or constraints. 

653 # We'll return that SqlSelectBuilder (or copies of it) at the end. 

654 ( 

655 predicate, 

656 select_builder, 

657 postprocessing, 

658 ) = self.managers.dimensions.process_query_overlaps( 

659 tree.dimensions, 

660 predicate, 

661 tree.get_joined_dimension_groups(), 

662 collection_analysis.calibration_dataset_types, 

663 allow_duplicate_overlaps, 

664 predicate_constraints.constraint_data_id, 

665 ) 

666 # Initialize the plan we're return at the end of the method. 

667 joins = QueryJoinsAnalysis(predicate=predicate, columns=select_builder.columns) 

668 joins.messages.extend(predicate_constraints.messages) 

669 # Add columns required by postprocessing. 

670 postprocessing.gather_columns_required(joins.columns) 

671 # Add materializations, which can also bring in more postprocessing. 

672 for m_key, m_dimensions in tree.materializations.items(): 

673 m_state = self._materializations[m_key] 

674 joins.materializations[m_key] = m_dimensions 

675 # When a query is materialized, the new tree has an empty 

676 # (trivially true) predicate because the original was used to make 

677 # the materialized rows. But the original postprocessing isn't 

678 # executed when the materialization happens, so we have to include 

679 # it here. 

680 postprocessing.spatial_join_filtering.extend(m_state.postprocessing.spatial_join_filtering) 

681 postprocessing.spatial_where_filtering.extend(m_state.postprocessing.spatial_where_filtering) 

682 postprocessing.spatial_expression_filtering.extend( 

683 m_state.postprocessing.spatial_expression_filtering 

684 ) 

685 # Add data coordinate uploads. 

686 joins.data_coordinate_uploads.update(tree.data_coordinate_uploads) 

687 # Add dataset_searches and filter out collections that don't have the 

688 # right dataset type or governor dimensions. We re-resolve dataset 

689 # searches now that we have a constraint data ID. 

690 for dataset_type_name, dataset_search in tree.datasets.items(): 

691 resolved_dataset_search = self._resolve_dataset_search( 

692 dataset_type_name, 

693 dataset_search, 

694 predicate_constraints.constraint_data_id, 

695 collection_analysis.summaries_by_dataset_type[dataset_type_name], 

696 ) 

697 if resolved_dataset_search.dimensions != self.get_dataset_type(dataset_type_name).dimensions: 697 ↛ 701line 697 didn't jump to line 701 because the condition on line 697 was never true

698 # This is really for server-side defensiveness; it's hard to 

699 # imagine the query getting different dimensions for a dataset 

700 # type in two calls to the same query driver. 

701 raise InvalidQueryError( 

702 f"Incorrect dimensions {resolved_dataset_search.dimensions} for dataset " 

703 f"{dataset_type_name!r} in query " 

704 f"(vs. {self.get_dataset_type(dataset_type_name).dimensions})." 

705 ) 

706 joins.datasets[dataset_type_name] = resolved_dataset_search 

707 if not resolved_dataset_search.collection_records: 

708 joins.messages.append( 

709 f"Search for dataset type {resolved_dataset_search.name!r} in " 

710 f"{list(dataset_search.collections)} is doomed to fail." 

711 ) 

712 joins.messages.extend(resolved_dataset_search.messages) 

713 # Process the special any_dataset search, if there is one. This entails 

714 # making a modified copy of the plan for each distinct post-filtering 

715 # collection search path. 

716 if tree.any_dataset is None: 716 ↛ 720line 716 didn't jump to line 720 because the condition on line 716 was always true

717 return QueryTreeAnalysis( 

718 joins, union_datasets=[], initial_select_builder=select_builder, postprocessing=postprocessing 

719 ) 

720 union_datasets = self._resolve_union_datasets(tree.any_dataset.dimensions, collection_analysis) 

721 return QueryTreeAnalysis( 

722 joins, 

723 union_datasets=union_datasets, 

724 initial_select_builder=select_builder, 

725 postprocessing=postprocessing, 

726 ) 

727 

728 def _resolve_union_datasets( 

729 self, dimensions: DimensionGroup, collection_analysis: QueryCollectionAnalysis 

730 ) -> list[ResolvedDatasetSearch[list[str]]]: 

731 """Resolve searches for union datasets. 

732 

733 Parameters 

734 ---------- 

735 dimensions : `DimensionGroup` 

736 Dimensions of the union dataset types in this query. 

737 collection_analysis : `CollectionAnalysis` 

738 Information about the collections appearing in this collection. 

739 

740 Returns 

741 ------- 

742 searches : `list` [ `ResolvedDatasetSearch` ] 

743 Resolved dataset searches for all union dataset types with these 

744 dimensions. Each item in the list groups dataset types with the 

745 same collection search path. 

746 """ 

747 # Gather the filtered collection search path for each union dataset 

748 # type. 

749 collections_by_dataset_type = defaultdict[str, list[str]](list) 

750 for collection_record, collection_summary in collection_analysis.summaries_by_dataset_type[ 

751 qt.ANY_DATASET 

752 ]: 

753 for dataset_type in collection_summary.dataset_types: 

754 if dataset_type.dimensions == dimensions: 

755 collections_by_dataset_type[dataset_type.name].append(collection_record.name) 

756 # Reverse the lookup order on the mapping we just made to group 

757 # dataset types by their collection search path. Each such group 

758 # yields a term in a union query builder 

759 dataset_searches_by_collections: dict[tuple[str, ...], ResolvedDatasetSearch[list[str]]] = {} 

760 for dataset_type_name, collection_path in collections_by_dataset_type.items(): 

761 key = tuple(collection_path) 

762 if (resolved_search := dataset_searches_by_collections.get(key)) is None: 

763 resolved_search = ResolvedDatasetSearch[list[str]]( 

764 [], 

765 dimensions=dimensions, 

766 collection_records=[ 

767 collection_analysis.collection_records[collection_name] 

768 for collection_name in collection_path 

769 ], 

770 messages=[], 

771 ) 

772 resolved_search.is_calibration_search = any( 

773 r.type is CollectionType.CALIBRATION for r in resolved_search.collection_records 

774 ) 

775 dataset_searches_by_collections[key] = resolved_search 

776 resolved_search.name.append(dataset_type_name) 

777 return list(dataset_searches_by_collections.values()) 

778 

779 def apply_initial_query_joins( 

780 self, 

781 select_builder: SqlSelectBuilder, 

782 joins_analysis: QueryJoinsAnalysis, 

783 union_dataset_dimensions: DimensionGroup | None, 

784 ) -> None: 

785 """Apply most of the "joins" stage of query construction to a single 

786 SELECT. 

787 

788 This method is expected to be invoked by `QueryBuilder.apply_joins` 

789 implementations. It handles all tables and subqueries in the FROM 

790 clause, except: 

791 

792 - the `QueryTree.any_dataset` search (handled by `UnionQueryBuilder` 

793 directly); 

794 - joins of dimension tables needed only for their keys (handled by 

795 `apply_missing_dimension_joins`). 

796 

797 Parameters 

798 ---------- 

799 select_builder : `SqlSelectBuilder` 

800 Low-level SQL builder for a single SELECT term, modified in place. 

801 joins_analysis : `QueryJoinsAnalysis` 

802 Information about the joins stage of query construction. 

803 union_dataset_dimensions : `DimensionGroup` or `None` 

804 Dimensions of the union dataset types, or `None` if this is not 

805 a union dataset query. 

806 """ 

807 # Process data coordinate upload joins. 

808 for upload_key, upload_dimensions in joins_analysis.data_coordinate_uploads.items(): 

809 select_builder.joins.join( 

810 SqlJoinsBuilder(db=self.db, from_clause=self._upload_tables[upload_key]).extract_dimensions( 

811 upload_dimensions.required 

812 ) 

813 ) 

814 # Process materialization joins. We maintain a set of dataset types 

815 # that were included in a materialization; searches for these datasets 

816 # can be dropped if they are only present to provide a constraint on 

817 # data IDs, since that's already embedded in a materialization. 

818 materialized_datasets: set[str] = set() 

819 for materialization_key, materialization_dimensions in joins_analysis.materializations.items(): 

820 materialized_datasets.update( 

821 self._join_materialization( 

822 select_builder.joins, materialization_key, materialization_dimensions 

823 ) 

824 ) 

825 # Process dataset joins (not including any union dataset). Datasets 

826 # searches included in materialization can be skipped unless we need 

827 # something from their tables. 

828 materialized_datasets = set() 

829 for m_state in self._materializations.values(): 

830 materialized_datasets.update(m_state.datasets) 

831 for dataset_type_name, dataset_search in joins_analysis.datasets.items(): 

832 if ( 

833 dataset_type_name not in materialized_datasets 

834 or dataset_type_name in select_builder.columns.dataset_fields 

835 ): 

836 self.join_dataset_search( 

837 select_builder.joins, 

838 dataset_search, 

839 joins_analysis.columns.dataset_fields[dataset_search.name], 

840 ) 

841 # Join in dimension element tables that we know we need relationships 

842 # or columns from. 

843 for element in joins_analysis.iter_mandatory(union_dataset_dimensions): 

844 select_builder.joins.join( 

845 self.managers.dimensions.make_joins_builder( 

846 element, joins_analysis.columns.dimension_fields[element.name] 

847 ) 

848 ) 

849 

850 def apply_missing_dimension_joins( 

851 self, select_builder: SqlSelectBuilder, joins_analysis: QueryJoinsAnalysis 

852 ) -> None: 

853 """Apply dimension-table joins to a single SQL SELECT builder to ensure 

854 the full set of desired dimension keys is preset. 

855 

856 This method is expected to be invoked by `QueryBuilder.apply_joins` 

857 implementations. 

858 

859 Parameters 

860 ---------- 

861 select_builder : `SqlSelectBuilder` 

862 Low-level SQL builder for a single SELECT term, modified in place. 

863 joins_analysis : `QueryJoinsAnalysis` 

864 Information about the joins stage of query construction. 

865 """ 

866 # See if any dimension keys are still missing, and if so join in 

867 # their tables. Note that we know there are no fields needed from 

868 # these. 

869 while not (select_builder.joins.dimension_keys.keys() >= joins_analysis.columns.dimensions.names): 

870 # Look for opportunities to join in multiple dimensions via 

871 # single table, to reduce the total number of tables joined in. 

872 missing_dimension_names = ( 

873 joins_analysis.columns.dimensions.names - select_builder.joins.dimension_keys.keys() 

874 ) 

875 best = self._universe[ 

876 max( 

877 missing_dimension_names, 

878 key=lambda name: len(self._universe[name].dimensions.names & missing_dimension_names), 

879 ) 

880 ] 

881 to_join = self.managers.dimensions.make_joins_builder(best, frozenset()) 

882 select_builder.joins.join(to_join) 

883 # Add the WHERE clause to the builder. 

884 select_builder.joins.where( 

885 joins_analysis.predicate.visit(SqlColumnVisitor(select_builder.joins, self))