Coverage for python/lsst/daf/relation/sql/_engine.py: 12%

300 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-02-22 02:59 -0800

1# This file is part of daf_relation. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://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 <https://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24__all__ = ("Engine",) 

25 

26import dataclasses 

27from collections.abc import Callable, Iterable, Mapping, Sequence, Set 

28from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast 

29 

30import sqlalchemy 

31 

32from .._binary_operation import BinaryOperation, IgnoreOne 

33from .._columns import ( 

34 ColumnExpression, 

35 ColumnExpressionSequence, 

36 ColumnFunction, 

37 ColumnInContainer, 

38 ColumnLiteral, 

39 ColumnRangeLiteral, 

40 ColumnReference, 

41 ColumnTag, 

42 LogicalAnd, 

43 LogicalNot, 

44 LogicalOr, 

45 Predicate, 

46 PredicateFunction, 

47 PredicateLiteral, 

48 PredicateReference, 

49 flatten_logical_and, 

50) 

51from .._engine import GenericConcreteEngine 

52from .._exceptions import EngineError, RelationalAlgebraError 

53from .._leaf_relation import LeafRelation 

54from .._marker_relation import MarkerRelation 

55from .._materialization import Materialization 

56from .._operation_relations import BinaryOperationRelation, UnaryOperationRelation 

57from .._operations import ( 

58 Calculation, 

59 Chain, 

60 Deduplication, 

61 Join, 

62 PartialJoin, 

63 Projection, 

64 Selection, 

65 Slice, 

66 Sort, 

67 SortTerm, 

68) 

69from .._transfer import Transfer 

70from .._unary_operation import Identity, UnaryOperation 

71from ._payload import Payload 

72from ._select import Select 

73 

74if TYPE_CHECKING: 74 ↛ 75line 74 didn't jump to line 75, because the condition on line 74 was never true

75 from .._relation import Relation 

76 

77 

78_L = TypeVar("_L") 

79 

80 

81@dataclasses.dataclass(repr=False, eq=False, kw_only=True) 

82class Engine( 

83 GenericConcreteEngine[Callable[..., sqlalchemy.sql.ColumnElement[Any]]], 

84 Generic[_L], 

85): 

86 """A concrete engine class for relations backed by a SQL database. 

87 

88 See the `.sql` module documentation for details. 

89 """ 

90 

91 name: str = "sql" 

92 

93 EMPTY_COLUMNS_NAME: ClassVar[str] = "IGNORED" 

94 """Name of the column added to a SQL ``SELECT`` query in order to represent 

95 relations that have no real columns. 

96 """ 

97 

98 EMPTY_COLUMNS_TYPE: ClassVar[type] = sqlalchemy.Boolean 

99 """Type of the column added to a SQL ``SELECT`` query in order to represent 

100 relations that have no real columns. 

101 """ 

102 

103 def __str__(self) -> str: 

104 return self.name 

105 

106 def __repr__(self) -> str: 

107 return f"lsst.daf.relation.sql.Engine({self.name!r})@{id(self):0x}" 

108 

109 def make_leaf( 

110 self, 

111 columns: Set[ColumnTag], 

112 payload: Payload[_L], 

113 *, 

114 min_rows: int = 0, 

115 max_rows: int | None = None, 

116 name: str = "", 

117 messages: Sequence[str] = (), 

118 name_prefix: str = "leaf", 

119 parameters: Any = None, 

120 ) -> Relation: 

121 """Create a nontrivial leaf relation in this engine. 

122 

123 This is a convenience method that simply forwards all arguments to 

124 the `.LeafRelation` constructor, and then wraps the result in a 

125 `Select`; see `LeafRelation` for details. 

126 """ 

127 return Select.apply_skip( 

128 LeafRelation( 

129 self, 

130 frozenset(columns), 

131 payload, 

132 min_rows=min_rows, 

133 max_rows=max_rows, 

134 messages=messages, 

135 name=name, 

136 name_prefix=name_prefix, 

137 parameters=parameters, 

138 ) 

139 ) 

140 

141 def conform(self, relation: Relation) -> Select: 

142 # Docstring inherited. 

143 match relation: 

144 case Select(): 

145 return relation 

146 case UnaryOperationRelation(operation=operation, target=target): 

147 conformed_target = self.conform(target) 

148 return self._append_unary_to_select(operation, conformed_target) 

149 case BinaryOperationRelation(operation=operation, lhs=lhs, rhs=rhs): 

150 conformed_lhs = self.conform(lhs) 

151 conformed_rhs = self.conform(rhs) 

152 return self._append_binary_to_select(operation, conformed_lhs, conformed_rhs) 

153 case Transfer() | Materialization() | LeafRelation(): 

154 # We always conform upstream of transfers and materializations, 

155 # so no need to recurse. 

156 return Select.apply_skip(relation) 

157 case MarkerRelation(target=target): 

158 # Other marker relations. 

159 conformed_target = self.conform(target) 

160 return Select.apply_skip(conformed_target) 

161 raise AssertionError("Match should be exhaustive and all branches return.") 

162 

163 def materialize( 

164 self, target: Relation, name: str | None = None, name_prefix: str = "materialization_" 

165 ) -> Select: 

166 # Docstring inherited. 

167 conformed_target = self.conform(target) 

168 if conformed_target.has_sort and not conformed_target.has_slice: 

169 raise RelationalAlgebraError( 

170 f"Materializing relation {conformed_target} will not preserve row order." 

171 ) 

172 return Select.apply_skip(super().materialize(conformed_target, name, name_prefix)) 

173 

174 def transfer(self, target: Relation, payload: Any | None = None) -> Select: 

175 # Docstring inherited. 

176 return Select.apply_skip(super().transfer(target, payload)) 

177 

178 def make_doomed_relation( 

179 self, columns: Set[ColumnTag], messages: Sequence[str], name: str = "0" 

180 ) -> Relation: 

181 # Docstring inherited. 

182 return Select.apply_skip(super().make_doomed_relation(columns, messages, name)) 

183 

184 def make_join_identity_relation(self, name: str = "I") -> Relation: 

185 # Docstring inherited. 

186 return Select.apply_skip(super().make_join_identity_relation(name)) 

187 

188 def get_join_identity_payload(self) -> Payload[_L]: 

189 # Docstring inherited. 

190 select_columns: list[sqlalchemy.sql.ColumnElement] = [] 

191 self.handle_empty_columns(select_columns) 

192 return Payload[_L](sqlalchemy.sql.select(*select_columns).subquery()) 

193 

194 def get_doomed_payload(self, columns: Set[ColumnTag]) -> Payload[_L]: 

195 # Docstring inherited. 

196 select_columns: list[sqlalchemy.sql.ColumnElement] = [ 

197 sqlalchemy.sql.literal(None).label(self.get_identifier(tag)) for tag in columns 

198 ] 

199 self.handle_empty_columns(select_columns) 

200 subquery = sqlalchemy.sql.select(*select_columns).subquery() 

201 return Payload( 

202 subquery, 

203 where=[sqlalchemy.sql.literal(False)], 

204 columns_available=self.extract_mapping(columns, subquery.columns), 

205 ) 

206 

207 def append_unary(self, operation: UnaryOperation, target: Relation) -> Select: 

208 # Docstring inherited. 

209 conformed_target = self.conform(target) 

210 return self._append_unary_to_select(operation, conformed_target) 

211 

212 def _append_unary_to_select(self, operation: UnaryOperation, select: Select) -> Select: 

213 """Internal recursive implementation of `append_unary`. 

214 

215 This method should not be called by external code, but it may be 

216 overridden and called recursively by derived engine classes. 

217 

218 Parameters 

219 ---------- 

220 operation : `UnaryOperation` 

221 Operation to add to the tree. 

222 select : `Select` 

223 Existing already-conformed relation tree. 

224 

225 Returns 

226 ------- 

227 appended : `Select` 

228 Conformed relation tree that includes the given operation. 

229 """ 

230 match operation: 

231 case Calculation(tag=tag): 

232 if select.is_compound: 

233 # This Select wraps a Chain operation in order to represent 

234 # a SQL UNION or UNION ALL, and we trust the user's intent 

235 # in putting those upstream of this operation, so we also 

236 # add a nested subquery here. 

237 return Select.apply_skip(operation._finish_apply(select)) 

238 elif select.has_projection: 

239 return select.reapply_skip( 

240 after=operation, 

241 projection=Projection(frozenset(select.columns | {tag})), 

242 ) 

243 else: 

244 return select.reapply_skip(after=operation) 

245 case Deduplication(): 

246 if not select.has_deduplication: 

247 if select.has_slice: 

248 # There was a Slice upstream, which needs to be applied 

249 # before this Deduplication, so we nest the existing 

250 # subquery within a new one that has deduplication. 

251 return Select.apply_skip(select, deduplication=operation) 

252 else: 

253 # Move the Deduplication into the Select's 

254 # operations. 

255 return select.reapply_skip(deduplication=operation) 

256 else: 

257 # There was already another (redundant) 

258 # Deduplication upstream. 

259 return select 

260 case Projection(): 

261 if select.has_deduplication: 

262 # There was a Duplication upstream, so we need to ensure 

263 # that is applied before this Projection via a nested 

264 # subquery. Instead of applying the existing subquery 

265 # operations, though, we apply one without any sort or 

266 # slice that might exist, and save those for the new outer 

267 # query, since putting those in a subquery would destroy 

268 # the ordering. 

269 subquery = select.reapply_skip(sort=None, slice=None) 

270 return Select.apply_skip( 

271 subquery, 

272 projection=operation, 

273 sort=select.sort, 

274 slice=select.slice, 

275 ) 

276 else: 

277 # No Deduplication, so we can just add the Projection to 

278 # the existing Select and reapply it. 

279 match select.skip_to: 

280 case BinaryOperationRelation(operation=Chain() as chain, lhs=lhs, rhs=rhs): 

281 # ... unless the skip_to relation is a Chain; we 

282 # want to move the Projection inside the Chain, to 

283 # make it easier to avoid subqueries in UNION [ALL] 

284 # constructs later. 

285 return select.reapply_skip( 

286 skip_to=chain._finish_apply(operation.apply(lhs), operation.apply(rhs)), 

287 projection=None, 

288 ) 

289 case _: 

290 return select.reapply_skip(projection=operation) 

291 case Selection(): 

292 if select.has_slice: 

293 # There was a Slice upstream, which needs to be applied 

294 # before this Selection via a nested subquery (which means 

295 # we just apply the Selection to target, then add a new 

296 # empty subquery marker). 

297 return Select.apply_skip(operation._finish_apply(select)) 

298 elif select.is_compound: 

299 # This Select wraps a Chain operation in order to represent 

300 # a SQL UNION or UNION ALL, and we trust the user's intent 

301 # in putting those upstream of this operation, so we also 

302 # add a nested subquery here. 

303 return Select.apply_skip(operation._finish_apply(select)) 

304 else: 

305 return select.reapply_skip(after=operation) 

306 case Slice(): 

307 return select.reapply_skip(slice=select.slice.then(operation)) 

308 case Sort(): 

309 if select.has_slice: 

310 # There was a Slice upstream, which needs to be applied 

311 # before this Sort via a nested subquery (which means we 

312 # apply a new Sort-only Select to 'select' itself). 

313 return Select.apply_skip(select, sort=operation) 

314 else: 

315 return select.reapply_skip(sort=select.sort.then(operation)) 

316 case PartialJoin(binary=binary, fixed=fixed, fixed_is_lhs=fixed_is_lhs): 

317 if fixed_is_lhs: 

318 return self.append_binary(binary, fixed, select) 

319 else: 

320 return self.append_binary(binary, select, fixed) 

321 case Identity(): 

322 return select 

323 raise NotImplementedError(f"Unsupported operation type {operation} for engine {self}.") 

324 

325 def append_binary(self, operation: BinaryOperation, lhs: Relation, rhs: Relation) -> Select: 

326 # Docstring inherited. 

327 conformed_lhs = self.conform(lhs) 

328 conformed_rhs = self.conform(rhs) 

329 return self._append_binary_to_select(operation, conformed_lhs, conformed_rhs) 

330 

331 def _append_binary_to_select(self, operation: BinaryOperation, lhs: Select, rhs: Select) -> Select: 

332 """Internal recursive implementation of `append_binary`. 

333 

334 This method should not be called by external code, but it may be 

335 overridden and called recursively by derived engine classes. 

336 

337 Parameters 

338 ---------- 

339 operation : `UnaryOperation` 

340 Operation to add to the tree. 

341 lhs : `Select` 

342 Existing already-conformed relation tree. 

343 rhs : `Select` 

344 The other existing already-conformed relation tree. 

345 

346 Returns 

347 ------- 

348 appended : `Select` 

349 Conformed relation tree that includes the given operation. 

350 """ 

351 if lhs.has_sort and not lhs.has_slice: 

352 raise RelationalAlgebraError( 

353 f"Applying binary operation {operation} to relation {lhs} will not preserve row order." 

354 ) 

355 if rhs.has_sort and not rhs.has_slice: 

356 raise RelationalAlgebraError( 

357 f"Applying binary operation {operation} to relation {rhs} will not preserve row order." 

358 ) 

359 match operation: 

360 case Chain(): 

361 # Chain operands should always be Selects, but they can't have 

362 # Sorts or Slices, at least not directly; if those exist we 

363 # need to move them into subqueries at which point we might as 

364 # well move any Projection or Dedulication there as well). But 

365 # because of the check on Sorts-without-Slices above, there has 

366 # to be a Slice here for there to be a Sort here. 

367 if lhs.has_slice: 

368 lhs = Select.apply_skip(lhs) 

369 if rhs.has_slice: 

370 rhs = Select.apply_skip(rhs) 

371 return Select.apply_skip(operation._finish_apply(lhs, rhs)) 

372 case Join(): 

373 # For Joins, we move Projections after the Join, unless we have 

374 # another reason to make a subquery before the join, and the 

375 # operands are only Selects if they need to be subqueries. 

376 new_lhs, new_lhs_needs_projection = lhs.strip() 

377 new_rhs, new_rhs_needs_projection = rhs.strip() 

378 if new_lhs_needs_projection or new_rhs_needs_projection: 

379 projection = Projection(frozenset(lhs.columns | rhs.columns)) 

380 else: 

381 projection = None 

382 return Select.apply_skip(operation._finish_apply(new_lhs, new_rhs), projection=projection) 

383 case IgnoreOne(ignore_lhs=ignore_lhs): 

384 if ignore_lhs: 

385 return rhs 

386 else: 

387 return lhs 

388 raise AssertionError(f"Match on {operation} should be exhaustive and all branches return..") 

389 

390 def get_identifier(self, tag: ColumnTag) -> str: 

391 """Return the SQL identifier that should be used to represent the given 

392 column. 

393 

394 Parameters 

395 ---------- 

396 tag : `.ColumnTag` 

397 Object representing a column. 

398 

399 Returns 

400 ------- 

401 identifier : `str` 

402 SQL identifier for this column. 

403 

404 Notes 

405 ----- 

406 This method may be overridden to replace special characters not 

407 supported by a particular DBMS (even after quoting, which SQLAlchemy 

408 handles transparently), deal with case transformation, or ensure 

409 identifiers are not truncated (e.g. by PostgreSQL's 64-char limit). 

410 The default implementation returns ``tag.qualified_name`` unchanged. 

411 """ 

412 return tag.qualified_name 

413 

414 def extract_mapping( 

415 self, tags: Iterable[ColumnTag], sql_columns: sqlalchemy.sql.ColumnCollection 

416 ) -> dict[ColumnTag, _L]: 

417 """Extract a mapping with `.ColumnTag` keys and logical column values 

418 from a SQLAlchemy column collection. 

419 

420 Parameters 

421 ---------- 

422 tags : `Iterable` 

423 Set of `.ColumnTag` objects whose logical columns should be 

424 extracted. 

425 sql_columns : `sqlalchemy.sql.ColumnCollection` 

426 SQLAlchemy collection of columns, such as 

427 `sqlalchemy.sql.FromClause.columns`. 

428 

429 Returns 

430 ------- 

431 logical_columns : `dict` 

432 Dictionary mapping `.ColumnTag` to logical column type. 

433 

434 Notes 

435 ----- 

436 This method must be overridden to support a custom logical columns. 

437 """ 

438 return {tag: cast(_L, sql_columns[self.get_identifier(tag)]) for tag in tags} 

439 

440 def select_items( 

441 self, 

442 items: Iterable[tuple[ColumnTag, _L]], 

443 sql_from: sqlalchemy.sql.FromClause, 

444 *extra: sqlalchemy.sql.ColumnElement, 

445 ) -> sqlalchemy.sql.Select: 

446 """Construct a SQLAlchemy representation of a SELECT query. 

447 

448 Parameters 

449 ---------- 

450 items : `Iterable` [ `tuple` ] 

451 Iterable of (`.ColumnTag`, logical column) pairs. This is 

452 typically the ``items()`` of a mapping returned by 

453 `extract_mapping` or obtained from `Payload.columns_available`. 

454 sql_from : `sqlalchemy.sql.FromClause` 

455 SQLAlchemy representation of a FROM clause, such as a single table, 

456 aliased subquery, or join expression. Must provide all columns 

457 referenced by ``items``. 

458 *extra : `sqlalchemy.sql.ColumnElement` 

459 Additional SQL column expressions to include. 

460 

461 Returns 

462 ------- 

463 select : `sqlalchemy.sql.Select` 

464 SELECT query. 

465 

466 Notes 

467 ----- 

468 This method is responsible for handling the case where ``items`` is 

469 empty, typically by delegating to `handle_empty_columns`. 

470 

471 This method must be overridden to support a custom logical columns. 

472 """ 

473 select_columns: list[sqlalchemy.sql.ColumnElement] = [ 

474 cast(sqlalchemy.sql.ColumnElement, logical_column).label(self.get_identifier(tag)) 

475 for tag, logical_column in items 

476 ] 

477 select_columns.extend(extra) 

478 self.handle_empty_columns(select_columns) 

479 return sqlalchemy.sql.select(*select_columns).select_from(sql_from) 

480 

481 def handle_empty_columns(self, columns: list[sqlalchemy.sql.ColumnElement]) -> None: 

482 """Handle the edge case where a SELECT statement has no columns, by 

483 adding a literal column that should be ignored. 

484 

485 Parameters 

486 ---------- 

487 columns : `list` [ `sqlalchemy.sql.ColumnElement` ] 

488 List of SQLAlchemy column objects. This may have no elements when 

489 this method is called, and must always have at least one element 

490 when it returns. 

491 """ 

492 if not columns: 

493 columns.append(sqlalchemy.sql.literal(True).label(self.EMPTY_COLUMNS_NAME)) 

494 

495 def to_executable( 

496 self, relation: Relation, extra_columns: Iterable[sqlalchemy.sql.ColumnElement] = () 

497 ) -> sqlalchemy.sql.expression.SelectBase: 

498 """Convert a relation tree to an executable SQLAlchemy expression. 

499 

500 Parameters 

501 ---------- 

502 relation : `Relation` 

503 The relation tree to convert. 

504 extra_columns : `~collections.abc.Iterable` 

505 Iterable of additional SQLAlchemy column objects to include 

506 directly in the ``SELECT`` clause. 

507 

508 Returns 

509 ------- 

510 select : `sqlalchemy.sql.expression.SelectBase` 

511 A SQLAlchemy ``SELECT`` or compound ``SELECT`` query. 

512 

513 Notes 

514 ----- 

515 This method requires all relations in the tree to have the same engine 

516 (``self``). It also cannot handle `.Materialization` operations 

517 unless they have already been processed once already (and hence have 

518 a payload attached). Use the `.Processor` function to handle both of 

519 these cases. 

520 """ 

521 if relation.engine != self: 

522 raise EngineError( 

523 f"Engine {self!r} cannot operate on relation {relation} with engine {relation.engine!r}. " 

524 "Use lsst.daf.relation.Processor to evaluate transfers first." 

525 ) 

526 relation = self.conform(relation) 

527 return self._select_to_executable(relation, extra_columns) 

528 

529 def _select_to_executable( 

530 self, 

531 select: Select, 

532 extra_columns: Iterable[sqlalchemy.sql.ColumnElement], 

533 ) -> sqlalchemy.sql.expression.SelectBase: 

534 """Internal recursive implementation of `to_executable`. 

535 

536 This method should not be called by external code, but it may be 

537 overridden and called recursively by derived engine classes. 

538 

539 Parameters 

540 ---------- 

541 columns : `~collections.abc.Set` [ `ColumnTag` ] 

542 Columns to include in the ``SELECT`` clause. 

543 select : `Select` 

544 Already-conformed relation tree to convert. 

545 extra_columns : `~collections.abc.Iterable` 

546 Iterable of additional SQLAlchemy column objects to include 

547 directly in the ``SELECT`` clause. 

548 

549 Returns 

550 ------- 

551 executable : `sqlalchemy.sql.expression.SelectBase` 

552 SQLAlchemy ``SELECT`` or compound ``SELECT`` object. 

553 

554 Notes 

555 ----- 

556 This method handles trees terminated with `Select`, the operation 

557 relation types managed by `Select`, and `Chain` operations nested 

558 directly as the `skip_to` target of a `Select`. It delegates to 

559 `to_payload` for all other relation types. 

560 """ 

561 columns_available: Mapping[ColumnTag, _L] | None = None 

562 executable: sqlalchemy.sql.Select | sqlalchemy.sql.CompoundSelect 

563 match select.skip_to: 

564 case BinaryOperationRelation(operation=Chain(), lhs=lhs, rhs=rhs): 

565 lhs_executable = self._select_to_executable(cast(Select, lhs), extra_columns) 

566 rhs_executable = self._select_to_executable(cast(Select, rhs), extra_columns) 

567 if select.has_deduplication: 

568 executable = sqlalchemy.sql.union(lhs_executable, rhs_executable) 

569 else: 

570 executable = sqlalchemy.sql.union_all(lhs_executable, rhs_executable) 

571 case _: 

572 if select.skip_to.payload is not None: 

573 payload: Payload[_L] = select.skip_to.payload 

574 else: 

575 payload = self.to_payload(select.skip_to) 

576 if not select.columns and not extra_columns: 

577 extra_columns = list(extra_columns) 

578 self.handle_empty_columns(extra_columns) 

579 columns_available = payload.columns_available 

580 columns_projected = {tag: columns_available[tag] for tag in select.columns} 

581 executable = self.select_items(columns_projected.items(), payload.from_clause, *extra_columns) 

582 if len(payload.where) == 1: 

583 executable = executable.where(payload.where[0]) 

584 elif payload.where: 

585 executable = executable.where(sqlalchemy.sql.and_(*payload.where)) 

586 if select.has_deduplication: 

587 executable = executable.distinct() 

588 if select.has_sort: 

589 if columns_available is None: 

590 columns_available = self.extract_mapping(select.skip_to.columns, executable.selected_columns) 

591 executable = executable.order_by( 

592 *[self.convert_sort_term(term, columns_available) for term in select.sort.terms] 

593 ) 

594 if select.slice.start: 

595 executable = executable.offset(select.slice.start) 

596 if select.slice.limit is not None: 

597 executable = executable.limit(select.slice.limit) 

598 return executable 

599 

600 def to_payload(self, relation: Relation) -> Payload[_L]: 

601 """Internal recursive implementation of `to_executable`. 

602 

603 This method should not be called by external code, but it may be 

604 overridden and called recursively by derived engine classes. 

605 

606 Parameters 

607 ---------- 

608 relation : `Relation` 

609 Relation to convert. This method handles all operation relation 

610 types other than `Chain` and the operations managed by `Select`. 

611 

612 Returns 

613 ------- 

614 payload : `Payload` 

615 Struct containing a SQLAlchemy represenation of a simple ``SELECT`` 

616 query. 

617 """ 

618 assert relation.engine == self, "Should be guaranteed by callers." 

619 if relation.payload is not None: # Should cover all LeafRelations 

620 return relation.payload 

621 match relation: 

622 case UnaryOperationRelation(operation=operation, target=target): 

623 match operation: 

624 case Calculation(tag=tag, expression=expression): 

625 result = self.to_payload(target).copy() 

626 result.columns_available[tag] = self.convert_column_expression( 

627 expression, result.columns_available 

628 ) 

629 return result 

630 case Selection(predicate=predicate): 

631 result = self.to_payload(target).copy() 

632 result.where.extend( 

633 self.convert_flattened_predicate(predicate, result.columns_available) 

634 ) 

635 return result 

636 case BinaryOperationRelation( 

637 operation=Join(predicate=predicate, common_columns=common_columns), lhs=lhs, rhs=rhs 

638 ): 

639 lhs_payload = self.to_payload(lhs) 

640 rhs_payload = self.to_payload(rhs) 

641 assert common_columns is not None, "Guaranteed by Join.apply and PartialJoin.apply." 

642 on_terms: list[sqlalchemy.sql.ColumnElement] = [] 

643 if common_columns: 

644 on_terms.extend( 

645 cast( 

646 sqlalchemy.sql.ColumnElement, 

647 lhs_payload.columns_available[tag] == rhs_payload.columns_available[tag], 

648 ) 

649 for tag in common_columns 

650 ) 

651 columns_available = {**lhs_payload.columns_available, **rhs_payload.columns_available} 

652 if predicate.as_trivial() is not True: 

653 on_terms.extend(self.convert_flattened_predicate(predicate, columns_available)) 

654 on_clause: sqlalchemy.sql.ColumnElement 

655 if not on_terms: 

656 on_clause = sqlalchemy.sql.literal(True) 

657 elif len(on_terms) == 1: 

658 on_clause = on_terms[0] 

659 else: 

660 on_clause = sqlalchemy.sql.and_(*on_terms) 

661 return Payload( 

662 from_clause=lhs_payload.from_clause.join(rhs_payload.from_clause, onclause=on_clause), 

663 where=lhs_payload.where + rhs_payload.where, 

664 columns_available=columns_available, 

665 ) 

666 case Select(): 

667 from_clause = self._select_to_executable(relation, ()).subquery() 

668 columns_available = self.extract_mapping(relation.columns, from_clause.columns) 

669 return Payload(from_clause, columns_available=columns_available) 

670 case Materialization(name=name): 

671 raise EngineError( 

672 f"Cannot persist materialization {name!r} during SQL conversion; " 

673 "use `lsst.daf.relation.Processor` first to handle this operation." 

674 ) 

675 case Transfer(target=target): 

676 raise EngineError( 

677 f"Cannot handle transfer from {target.engine} during SQL conversion; " 

678 "use `lsst.daf.relation.Processor` first to handle this operation." 

679 ) 

680 raise NotImplementedError(f"Unsupported relation type {relation} for engine {self}.") 

681 

682 def convert_column_expression( 

683 self, expression: ColumnExpression, columns_available: Mapping[ColumnTag, _L] 

684 ) -> _L: 

685 """Convert a `.ColumnExpression` to a logical column. 

686 

687 Parameters 

688 ---------- 

689 expression : `.ColumnExpression` 

690 Expression to convert. 

691 columns_available : `~collections.abc.Mapping` 

692 Mapping from `.ColumnTag` to logical column, typically produced by 

693 `extract_mapping` or obtained from `Payload.columns_available`. 

694 

695 Returns 

696 ------- 

697 logical_column 

698 SQLAlchemy expression object or other logical column value. 

699 

700 See Also 

701 -------- 

702 :ref:`lsst.daf.relation-sql-logical-columns` 

703 """ 

704 match expression: 

705 case ColumnLiteral(value=value): 

706 return self.convert_column_literal(value) 

707 case ColumnReference(tag=tag): 

708 return columns_available[tag] 

709 case ColumnFunction(name=name, args=args): 

710 sql_args = [self.convert_column_expression(arg, columns_available) for arg in args] 

711 if (function := self.get_function(name)) is not None: 

712 return cast(_L, function(*sql_args)) 

713 return getattr(sql_args[0], name)(*sql_args[1:]) 

714 raise AssertionError( 

715 f"matches should be exhaustive and all branches should return; got {expression!r}." 

716 ) 

717 

718 def convert_column_literal(self, value: Any) -> _L: 

719 """Convert a Python literal value to a logical column. 

720 

721 Parameters 

722 ---------- 

723 value 

724 Python value to convert. 

725 

726 Returns 

727 ------- 

728 logical_column 

729 SQLAlchemy expression object or other logical column value. 

730 

731 Notes 

732 ----- 

733 This method must be overridden to support a custom logical columns. 

734 

735 See Also 

736 -------- 

737 :ref:`lsst.daf.relation-sql-logical-columns` 

738 """ 

739 return cast(_L, sqlalchemy.sql.literal(value)) 

740 

741 def expect_column_scalar(self, logical_column: _L) -> sqlalchemy.sql.ColumnElement: 

742 """Convert a logical column value to a SQLAlchemy expression. 

743 

744 Parameters 

745 ---------- 

746 logical_column 

747 SQLAlchemy expression object or other logical column value. 

748 

749 Returns 

750 ------- 

751 sql : `sqlalchemy.sql.ColumnElement` 

752 SQLAlchemy expression object. 

753 

754 Notes 

755 ----- 

756 The default implementation assumes the logical column type is just a 

757 SQLAlchemy type and returns the given object unchanged. Subclasses 

758 with a custom logical column type should override to at least assert 

759 that the value is in fact a SQLAlchemy expression. This is only called 

760 in contexts where true SQLAlchemy expressions are required, such as in 

761 ``ORDER BY`` or ``WHERE`` clauses. 

762 

763 See Also 

764 -------- 

765 :ref:`lsst.daf.relation-sql-logical-columns` 

766 """ 

767 return cast(sqlalchemy.sql.ColumnElement, logical_column) 

768 

769 def convert_flattened_predicate( 

770 self, predicate: Predicate, columns_available: Mapping[ColumnTag, _L] 

771 ) -> list[sqlalchemy.sql.ColumnElement]: 

772 """Flatten all logical AND operators in a `.Predicate` and convert each 

773 to a boolean SQLAlchemy expression. 

774 

775 Parameters 

776 ---------- 

777 predicate : `.Predicate` 

778 Predicate to convert. 

779 columns_available : `~collections.abc.Mapping` 

780 Mapping from `.ColumnTag` to logical column, typically produced by 

781 `extract_mapping` or obtained from `Payload.columns_available`. 

782 

783 Returns 

784 ------- 

785 sql : `list` [ `sqlalchemy.sql.ColumnElement` ] 

786 List of boolean SQLAlchemy expressions to be combined with the 

787 ``AND`` operator. 

788 """ 

789 if (flattened := flatten_logical_and(predicate)) is False: 

790 return [sqlalchemy.sql.literal(False)] 

791 else: 

792 return [self.convert_predicate(p, columns_available) for p in flattened] 

793 

794 def convert_predicate( 

795 self, predicate: Predicate, columns_available: Mapping[ColumnTag, _L] 

796 ) -> sqlalchemy.sql.ColumnElement: 

797 """Convert a `.Predicate` to a SQLAlchemy expression. 

798 

799 Parameters 

800 ---------- 

801 predicate : `.Predicate` 

802 Predicate to convert. 

803 columns_available : `~collections.abc.Mapping` 

804 Mapping from `.ColumnTag` to logical column, typically produced by 

805 `extract_mapping` or obtained from `Payload.columns_available`. 

806 

807 Returns 

808 ------- 

809 sql : `sqlalchemy.sql.ColumnElement` 

810 Boolean SQLAlchemy expression. 

811 """ 

812 match predicate: 

813 case PredicateFunction(name=name, args=args): 

814 sql_args = [self.convert_column_expression(arg, columns_available) for arg in args] 

815 if (function := self.get_function(name)) is not None: 

816 return function(*sql_args) 

817 return getattr(sql_args[0], name)(*sql_args[1:]) 

818 case LogicalAnd(operands=operands): 

819 if not operands: 

820 return sqlalchemy.sql.literal(True) 

821 if len(operands) == 1: 

822 return self.convert_predicate(operands[0], columns_available) 

823 else: 

824 return sqlalchemy.sql.and_( 

825 *[self.convert_predicate(operand, columns_available) for operand in operands] 

826 ) 

827 case LogicalOr(operands=operands): 

828 if not operands: 

829 return sqlalchemy.sql.literal(False) 

830 if len(operands) == 1: 

831 return self.convert_predicate(operands[0], columns_available) 

832 else: 

833 return sqlalchemy.sql.or_( 

834 *[self.convert_predicate(operand, columns_available) for operand in operands] 

835 ) 

836 case LogicalNot(operand=operand): 

837 return sqlalchemy.sql.not_(self.convert_predicate(operand, columns_available)) 

838 case PredicateReference(tag=tag): 

839 return self.expect_column_scalar(columns_available[tag]) 

840 case PredicateLiteral(value=value): 

841 return sqlalchemy.sql.literal(value) 

842 case ColumnInContainer(item=item, container=container): 

843 sql_item = self.expect_column_scalar(self.convert_column_expression(item, columns_available)) 

844 match container: 

845 case ColumnRangeLiteral(value=range(start=start, stop=stop_exclusive, step=step)): 

846 # The convert_column_literal calls below should just 

847 # call sqlalchemy.sql.literal(int), which would also 

848 # happen automatically internal to any of the other 

849 # sqlalchemy function calls, but they get the typing 

850 # right, reflecting the fact that the derived engine is 

851 # supposed to have final say over how we convert 

852 # literals. 

853 stop_inclusive = stop_exclusive - 1 

854 if start == stop_inclusive: 

855 return sql_item == self.convert_column_literal(start) 

856 else: 

857 target = sqlalchemy.sql.between( 

858 sql_item, 

859 self.convert_column_literal(start), 

860 self.convert_column_literal(stop_inclusive), 

861 ) 

862 if step != 1: 

863 return sqlalchemy.sql.and_( 

864 *[ 

865 target, 

866 sql_item % self.convert_column_literal(step) 

867 == self.convert_column_literal(start % step), 

868 ] 

869 ) 

870 else: 

871 return target 

872 case ColumnExpressionSequence(items=items): 

873 return sql_item.in_( 

874 [self.convert_column_expression(item, columns_available) for item in items] 

875 ) 

876 raise AssertionError( 

877 f"matches should be exhaustive and all branches should return; got {predicate!r}." 

878 ) 

879 

880 def convert_sort_term( 

881 self, term: SortTerm, columns_available: Mapping[ColumnTag, _L] 

882 ) -> sqlalchemy.sql.ColumnElement: 

883 """Convert a `.SortTerm` to a SQLAlchemy expression. 

884 

885 Parameters 

886 ---------- 

887 term : `.SortTerm` 

888 Sort term to convert. 

889 columns_available : `~collections.abc.Mapping` 

890 Mapping from `.ColumnTag` to logical column, typically produced by 

891 `extract_mapping` or obtained from `Payload.columns_available`. 

892 

893 Returns 

894 ------- 

895 sql : `sqlalchemy.sql.ColumnElement` 

896 Scalar SQLAlchemy expression. 

897 """ 

898 result = self.expect_column_scalar(self.convert_column_expression(term.expression, columns_available)) 

899 if term.ascending: 

900 return result 

901 else: 

902 return result.desc()