Coverage for python/lsst/daf/relation/sql/_engine.py: 12%
296 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-13 10:31 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-13 10:31 +0000
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/>.
22from __future__ import annotations
24__all__ = ("Engine",)
26import dataclasses
27from collections.abc import Callable, Iterable, Mapping, Sequence, Set
28from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast
30import sqlalchemy
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 .._materialization import Materialization
55from .._operation_relations import BinaryOperationRelation, UnaryOperationRelation
56from .._operations import (
57 Calculation,
58 Chain,
59 Deduplication,
60 Join,
61 PartialJoin,
62 Projection,
63 Selection,
64 Slice,
65 Sort,
66 SortTerm,
67)
68from .._transfer import Transfer
69from .._unary_operation import Identity, UnaryOperation
70from ._payload import Payload
71from ._select import Select
73if TYPE_CHECKING: 73 ↛ 74line 73 didn't jump to line 74, because the condition on line 73 was never true
74 from .._relation import Relation
77_L = TypeVar("_L")
80@dataclasses.dataclass(repr=False, eq=False, kw_only=True)
81class Engine(
82 GenericConcreteEngine[Callable[..., sqlalchemy.sql.ColumnElement]],
83 Generic[_L],
84):
85 """A concrete engine class for relations backed by a SQL database.
87 See the `.sql` module documentation for details.
88 """
90 name: str = "sql"
92 EMPTY_COLUMNS_NAME: ClassVar[str] = "IGNORED"
93 """Name of the column added to a SQL ``SELECT`` query in order to represent
94 relations that have no real columns.
95 """
97 EMPTY_COLUMNS_TYPE: ClassVar[type] = sqlalchemy.Boolean
98 """Type of the column added to a SQL ``SELECT`` query in order to represent
99 relations that have no real columns.
100 """
102 def __str__(self) -> str:
103 return self.name
105 def __repr__(self) -> str:
106 return f"lsst.daf.relation.sql.Engine({self.name!r})@{id(self):0x}"
108 def make_leaf(
109 self,
110 columns: Set[ColumnTag],
111 payload: Payload[_L],
112 *,
113 min_rows: int = 0,
114 max_rows: int | None = None,
115 name: str = "",
116 messages: Sequence[str] = (),
117 name_prefix: str = "leaf",
118 parameters: Any = None,
119 ) -> Relation:
120 """Create a nontrivial leaf relation in this engine.
122 This is a convenience method that simply forwards all arguments to
123 the `.LeafRelation` constructor, and then wraps the result in a
124 `Select`; see `LeafRelation` for details.
125 """
126 return Select.apply_skip(
127 LeafRelation(
128 self,
129 frozenset(columns),
130 payload,
131 min_rows=min_rows,
132 max_rows=max_rows,
133 messages=messages,
134 name=name,
135 name_prefix=name_prefix,
136 parameters=parameters,
137 )
138 )
140 def conform(self, relation: Relation) -> Select:
141 # Docstring inherited.
142 match relation:
143 case Select():
144 return relation
145 case UnaryOperationRelation(operation=operation, target=target):
146 conformed_target = self.conform(target)
147 return self._append_unary_to_select(operation, conformed_target)
148 case BinaryOperationRelation(operation=operation, lhs=lhs, rhs=rhs):
149 conformed_lhs = self.conform(lhs)
150 conformed_rhs = self.conform(rhs)
151 return self._append_binary_to_select(operation, conformed_lhs, conformed_rhs)
152 case Transfer() | Materialization() | LeafRelation():
153 # We always conform upstream of transfers and materializations,
154 # so no need to recurse.
155 return Select.apply_skip(relation)
156 case _:
157 # Other marker relations.
158 conformed_target = self.conform(target)
159 return Select.apply_skip(relation)
160 raise AssertionError("Match should be exhaustive and all branches return.")
162 def materialize(
163 self, target: Relation, name: str | None = None, name_prefix: str = "materialization_"
164 ) -> Select:
165 # Docstring inherited.
166 conformed_target = self.conform(target)
167 if conformed_target.has_sort and not conformed_target.has_slice:
168 raise RelationalAlgebraError(
169 f"Materializing relation {conformed_target} will not preserve row order."
170 )
171 return Select.apply_skip(super().materialize(conformed_target, name, name_prefix))
173 def transfer(self, target: Relation, payload: Any | None = None) -> Select:
174 # Docstring inherited.
175 return Select.apply_skip(super().transfer(target, payload))
177 def make_doomed_relation(
178 self, columns: Set[ColumnTag], messages: Sequence[str], name: str = "0"
179 ) -> Relation:
180 # Docstring inherited.
181 return Select.apply_skip(super().make_doomed_relation(columns, messages, name))
183 def make_join_identity_relation(self, name: str = "I") -> Relation:
184 # Docstring inherited.
185 return Select.apply_skip(super().make_join_identity_relation(name))
187 def get_join_identity_payload(self) -> Payload[_L]:
188 # Docstring inherited.
189 select_columns: list[sqlalchemy.sql.ColumnElement] = []
190 self.handle_empty_columns(select_columns)
191 return Payload[_L](sqlalchemy.sql.select(*select_columns).subquery())
193 def get_doomed_payload(self, columns: Set[ColumnTag]) -> Payload[_L]:
194 # Docstring inherited.
195 select_columns = [sqlalchemy.sql.literal(None).label(tag.qualified_name) for tag in columns]
196 self.handle_empty_columns(select_columns)
197 subquery = sqlalchemy.sql.select(*select_columns).subquery()
198 return Payload(
199 subquery,
200 where=[sqlalchemy.sql.literal(False)],
201 columns_available=self.extract_mapping(columns, subquery.columns),
202 )
204 def append_unary(self, operation: UnaryOperation, target: Relation) -> Select:
205 # Docstring inherited.
206 conformed_target = self.conform(target)
207 return self._append_unary_to_select(operation, conformed_target)
209 def _append_unary_to_select(self, operation: UnaryOperation, select: Select) -> Select:
210 """Internal recursive implementation of `append_unary`.
212 This method should not be called by external code, but it may be
213 overridden and called recursively by derived engine classes.
215 Parameters
216 ----------
217 operation : `UnaryOperation`
218 Operation to add to the tree.
219 select : `Select`
220 Existing already-conformed relation tree.
222 Returns
223 -------
224 appended : `Select`
225 Conformed relation tree that includes the given operation.
226 """
227 match operation:
228 case Calculation(tag=tag):
229 if select.is_compound:
230 # This Select wraps a Chain operation in order to represent
231 # a SQL UNION or UNION ALL, and we trust the user's intent
232 # in putting those upstream of this operation, so we also
233 # add a nested subquery here.
234 return Select.apply_skip(operation._finish_apply(select))
235 elif select.has_projection:
236 return select.reapply_skip(
237 after=operation,
238 projection=Projection(frozenset(select.columns | {tag})),
239 )
240 else:
241 return select.reapply_skip(after=operation)
242 case Deduplication():
243 if not select.has_deduplication:
244 if select.has_slice:
245 # There was a Slice upstream, which needs to be applied
246 # before this Deduplication, so we nest the existing
247 # subquery within a new one that has deduplication.
248 return Select.apply_skip(select, deduplication=operation)
249 else:
250 # Move the Deduplication into the Select's
251 # operations.
252 return select.reapply_skip(deduplication=operation)
253 else:
254 # There was already another (redundant)
255 # Deduplication upstream.
256 return select
257 case Projection():
258 if select.has_deduplication:
259 # There was a Duplication upstream, so we need to ensure
260 # that is applied before this Projection via a nested
261 # subquery. Instead of applying the existing subquery
262 # operations, though, we apply one without any sort or
263 # slice that might exist, and save those for the new outer
264 # query, since putting those in a subquery would destroy
265 # the ordering.
266 subquery = select.reapply_skip(sort=None, slice=None)
267 return Select.apply_skip(
268 subquery,
269 projection=operation,
270 sort=select.sort,
271 slice=select.slice,
272 )
273 else:
274 # No Deduplication, so we can just add the Projection to
275 # the existing Select and reapply it.
276 match select.skip_to:
277 case BinaryOperationRelation(operation=Chain() as chain, lhs=lhs, rhs=rhs):
278 # ... unless the skip_to relation is a Chain; we
279 # want to move the Projection inside the Chain, to
280 # make it easier to avoid subqueries in UNION [ALL]
281 # constructs later.
282 return select.reapply_skip(
283 skip_to=chain._finish_apply(operation.apply(lhs), operation.apply(rhs)),
284 projection=None,
285 )
286 case _:
287 return select.reapply_skip(projection=operation)
288 case Selection():
289 if select.has_slice:
290 # There was a Slice upstream, which needs to be applied
291 # before this Selection via a nested subquery (which means
292 # we just apply the Selection to target, then add a new
293 # empty subquery marker).
294 return Select.apply_skip(operation._finish_apply(select))
295 elif select.is_compound:
296 # This Select wraps a Chain operation in order to represent
297 # a SQL UNION or UNION ALL, and we trust the user's intent
298 # in putting those upstream of this operation, so we also
299 # add a nested subquery here.
300 return Select.apply_skip(operation._finish_apply(select))
301 else:
302 return select.reapply_skip(after=operation)
303 case Slice():
304 return select.reapply_skip(slice=select.slice.then(operation))
305 case Sort():
306 if select.has_slice:
307 # There was a Slice upstream, which needs to be applied
308 # before this Sort via a nested subquery (which means we
309 # apply a new Sort-only Select to 'select' itself).
310 return Select.apply_skip(select, sort=operation)
311 else:
312 return select.reapply_skip(sort=select.sort.then(operation))
313 case PartialJoin(binary=binary, fixed=fixed, fixed_is_lhs=fixed_is_lhs):
314 if fixed_is_lhs:
315 return self.append_binary(binary, fixed, select)
316 else:
317 return self.append_binary(binary, select, fixed)
318 case Identity():
319 return select
320 raise NotImplementedError(f"Unsupported operation type {operation} for engine {self}.")
322 def append_binary(self, operation: BinaryOperation, lhs: Relation, rhs: Relation) -> Select:
323 # Docstring inherited.
324 conformed_lhs = self.conform(lhs)
325 conformed_rhs = self.conform(rhs)
326 return self._append_binary_to_select(operation, conformed_lhs, conformed_rhs)
328 def _append_binary_to_select(self, operation: BinaryOperation, lhs: Select, rhs: Select) -> Select:
329 """Internal recursive implementation of `append_binary`.
331 This method should not be called by external code, but it may be
332 overridden and called recursively by derived engine classes.
334 Parameters
335 ----------
336 operation : `UnaryOperation`
337 Operation to add to the tree.
338 lhs : `Select`
339 Existing already-conformed relation tree.
340 rhs : `Select`
341 The other existing already-conformed relation tree.
343 Returns
344 -------
345 appended : `Select`
346 Conformed relation tree that includes the given operation.
347 """
348 if lhs.has_sort and not lhs.has_slice:
349 raise RelationalAlgebraError(
350 f"Applying binary operation {operation} to relation {lhs} will not preserve row order."
351 )
352 if rhs.has_sort and not rhs.has_slice:
353 raise RelationalAlgebraError(
354 f"Applying binary operation {operation} to relation {rhs} will not preserve row order."
355 )
356 match operation:
357 case Chain():
358 # Chain operands should always be Selects, but they can't have
359 # Sorts or Slices, at least not directly; if those exist we
360 # need to move them into subqueries at which point we might as
361 # well move any Projection or Dedulication there as well). But
362 # because of the check on Sorts-without-Slices above, there has
363 # to be a Slice here for there to be a Sort here.
364 if lhs.has_slice:
365 lhs = Select.apply_skip(lhs)
366 if rhs.has_slice:
367 rhs = Select.apply_skip(rhs)
368 return Select.apply_skip(operation._finish_apply(lhs, rhs))
369 case Join():
370 # For Joins, we move Projections after the Join, unless we have
371 # another reason to make a subquery before the join, and the
372 # operands are only Selects if they need to be subqueries.
373 new_lhs, new_lhs_needs_projection = lhs.strip()
374 new_rhs, new_rhs_needs_projection = rhs.strip()
375 if new_lhs_needs_projection or new_rhs_needs_projection:
376 projection = Projection(frozenset(lhs.columns | rhs.columns))
377 else:
378 projection = None
379 return Select.apply_skip(operation._finish_apply(new_lhs, new_rhs), projection=projection)
380 case IgnoreOne(ignore_lhs=ignore_lhs):
381 if ignore_lhs:
382 return rhs
383 else:
384 return lhs
385 raise AssertionError(f"Match on {operation} should be exhaustive and all branches return..")
387 def extract_mapping(
388 self, tags: Iterable[ColumnTag], sql_columns: sqlalchemy.sql.ColumnCollection
389 ) -> dict[ColumnTag, _L]:
390 """Extract a mapping with `.ColumnTag` keys and logical column values
391 from a SQLAlchemy column collection.
393 Parameters
394 ----------
395 tags : `Iterable`
396 Set of `.ColumnTag` objects whose logical columns should be
397 extracted.
398 sql_columns : `sqlalchemy.sql.ColumnCollection`
399 SQLAlchemy collection of columns, such as
400 `sqlalchemy.sql.FromClause.columns`.
402 Returns
403 -------
404 logical_columns : `dict`
405 Dictionary mapping `.ColumnTag` to logical column type.
407 Notes
408 -----
409 This method must be overridden to support a custom logical columns.
410 """
411 return {tag: cast(_L, sql_columns[tag.qualified_name]) for tag in tags}
413 def select_items(
414 self,
415 items: Iterable[tuple[ColumnTag, _L]],
416 sql_from: sqlalchemy.sql.FromClause,
417 *extra: sqlalchemy.sql.ColumnElement,
418 ) -> sqlalchemy.sql.Select:
419 """Construct a SQLAlchemy representation of a SELECT query.
421 Parameters
422 ----------
423 items : `Iterable` [ `tuple` ]
424 Iterable of (`.ColumnTag`, logical column) pairs. This is
425 typically the ``items()`` of a mapping returned by
426 `extract_mapping` or obtained from `Payload.columns_available`.
427 sql_from : `sqlalchemy.sql.FromClause`
428 SQLAlchemy representation of a FROM clause, such as a single table,
429 aliased subquery, or join expression. Must provide all columns
430 referenced by ``items``.
431 *extra : `sqlalchemy.sql.ColumnElement`
432 Additional SQL column expressions to include.
434 Returns
435 -------
436 select : `sqlalchemy.sql.Select`
437 SELECT query.
439 Notes
440 -----
441 This method is responsible for handling the case where ``items`` is
442 empty, typically by delegating to `handle_empty_columns`.
444 This method must be overridden to support a custom logical columns.
445 """
446 select_columns = [
447 cast(sqlalchemy.sql.ColumnElement, logical_column).label(tag.qualified_name)
448 for tag, logical_column in items
449 ]
450 select_columns.extend(extra)
451 self.handle_empty_columns(select_columns)
452 return sqlalchemy.sql.select(*select_columns).select_from(sql_from)
454 def handle_empty_columns(self, columns: list[sqlalchemy.sql.ColumnElement]) -> None:
455 """Handle the edge case where a SELECT statement has no columns, by
456 adding a literal column that should be ignored.
458 Parameters
459 ----------
460 columns : `list` [ `sqlalchemy.sql.ColumnElement` ]
461 List of SQLAlchemy column objects. This may have no elements when
462 this method is called, and must always have at least one element
463 when it returns.
464 """
465 if not columns:
466 columns.append(sqlalchemy.sql.literal(True).label(self.EMPTY_COLUMNS_NAME))
468 def to_executable(
469 self, relation: Relation, extra_columns: Iterable[sqlalchemy.sql.ColumnElement] = ()
470 ) -> sqlalchemy.sql.expression.SelectBase:
471 """Convert a relation tree to an executable SQLAlchemy expression.
473 Parameters
474 ----------
475 relation : `Relation`
476 The relation tree to convert.
477 extra_columns : `~collections.abc.Iterable`
478 Iterable of additional SQLAlchemy column objects to include
479 directly in the ``SELECT`` clause.
481 Returns
482 -------
483 select : `sqlalchemy.sql.expression.SelectBase`
484 A SQLAlchemy ``SELECT`` or compound ``SELECT`` query.
486 Notes
487 -----
488 This method requires all relations in the tree to have the same engine
489 (``self``). It also cannot handle `.Materialization` operations
490 unless they have already been processed once already (and hence have
491 a payload attached). Use the `.Processor` function to handle both of
492 these cases.
493 """
494 if relation.engine != self:
495 raise EngineError(
496 f"Engine {self!r} cannot operate on relation {relation} with engine {relation.engine!r}. "
497 "Use lsst.daf.relation.Processor to evaluate transfers first."
498 )
499 relation = self.conform(relation)
500 return self._select_to_executable(relation, extra_columns)
502 def _select_to_executable(
503 self,
504 select: Select,
505 extra_columns: Iterable[sqlalchemy.sql.ColumnElement],
506 ) -> sqlalchemy.sql.expression.SelectBase:
507 """Internal recursive implementation of `to_executable`.
509 This method should not be called by external code, but it may be
510 overridden and called recursively by derived engine classes.
512 Parameters
513 ----------
514 columns : `~collections.abc.Set` [ `ColumnTag` ]
515 Columns to include in the ``SELECT`` clause.
516 select : `Select`
517 Already-conformed relation tree to convert.
518 extra_columns : `~collections.abc.Iterable`
519 Iterable of additional SQLAlchemy column objects to include
520 directly in the ``SELECT`` clause.
522 Returns
523 -------
524 executable : `sqlalchemy.sql.expression.SelectBase`
525 SQLAlchemy ``SELECT`` or compound ``SELECT`` object.
527 Notes
528 -----
529 This method handles trees terminated with `Select`, the operation
530 relation types managed by `Select`, and `Chain` operations nested
531 directly as the `skip_to` target of a `Select`. It delegates to
532 `to_payload` for all other relation types.
533 """
534 columns_available: Mapping[ColumnTag, _L] | None = None
535 match select.skip_to:
536 case BinaryOperationRelation(operation=Chain(), lhs=lhs, rhs=rhs):
537 lhs_executable = self._select_to_executable(cast(Select, lhs), extra_columns)
538 rhs_executable = self._select_to_executable(cast(Select, rhs), extra_columns)
539 if select.has_deduplication:
540 executable = sqlalchemy.sql.union(lhs_executable, rhs_executable)
541 else:
542 executable = sqlalchemy.sql.union_all(lhs_executable, rhs_executable)
543 case _:
544 if select.skip_to.payload is not None:
545 payload: Payload[_L] = select.skip_to.payload
546 else:
547 payload = self.to_payload(select.skip_to)
548 if not select.columns and not extra_columns:
549 extra_columns = list(extra_columns)
550 self.handle_empty_columns(extra_columns)
551 columns_available = payload.columns_available
552 columns_projected = {tag: columns_available[tag] for tag in select.columns}
553 executable = self.select_items(columns_projected.items(), payload.from_clause, *extra_columns)
554 if len(payload.where) == 1:
555 executable = executable.where(payload.where[0])
556 elif payload.where:
557 executable = executable.where(sqlalchemy.sql.and_(*payload.where))
558 if select.has_deduplication:
559 executable = executable.distinct()
560 if select.has_sort:
561 if columns_available is None:
562 columns_available = self.extract_mapping(select.skip_to.columns, executable.selected_columns)
563 executable = executable.order_by(
564 *[self.convert_sort_term(term, columns_available) for term in select.sort.terms]
565 )
566 if select.slice.start:
567 executable = executable.offset(select.slice.start)
568 if select.slice.limit is not None:
569 executable = executable.limit(select.slice.limit)
570 return executable
572 def to_payload(self, relation: Relation) -> Payload[_L]:
573 """Internal recursive implementation of `to_executable`.
575 This method should not be called by external code, but it may be
576 overridden and called recursively by derived engine classes.
578 Parameters
579 ----------
580 relation : `Relation`
581 Relation to convert. This method handles all operation relation
582 types other than `Chain` and the operations managed by `Select`.
584 Returns
585 -------
586 payload : `Payload`
587 Struct containing a SQLAlchemy represenation of a simple ``SELECT``
588 query.
589 """
590 assert relation.engine == self, "Should be guaranteed by callers."
591 if relation.payload is not None: # Should cover all LeafRelations
592 return relation.payload
593 match relation:
594 case UnaryOperationRelation(operation=operation, target=target):
595 match operation:
596 case Calculation(tag=tag, expression=expression):
597 result = self.to_payload(target).copy()
598 result.columns_available[tag] = self.convert_column_expression(
599 expression, result.columns_available
600 )
601 return result
602 case Selection(predicate=predicate):
603 result = self.to_payload(target).copy()
604 result.where.extend(
605 self.convert_flattened_predicate(predicate, result.columns_available)
606 )
607 return result
608 case BinaryOperationRelation(
609 operation=Join(predicate=predicate, common_columns=common_columns), lhs=lhs, rhs=rhs
610 ):
611 lhs_payload = self.to_payload(lhs)
612 rhs_payload = self.to_payload(rhs)
613 assert common_columns is not None, "Guaranteed by Join.apply and PartialJoin.apply."
614 on_terms: list[sqlalchemy.sql.ColumnElement] = []
615 if common_columns:
616 on_terms.extend(
617 lhs_payload.columns_available[tag] == rhs_payload.columns_available[tag]
618 for tag in common_columns
619 )
620 columns_available = {**lhs_payload.columns_available, **rhs_payload.columns_available}
621 if predicate.as_trivial() is not True:
622 on_terms.extend(self.convert_flattened_predicate(predicate, columns_available))
623 on_clause: sqlalchemy.sql.ColumnElement
624 if not on_terms:
625 on_clause = sqlalchemy.sql.literal(True)
626 elif len(on_terms) == 1:
627 on_clause = on_terms[0]
628 else:
629 on_clause = sqlalchemy.sql.and_(*on_terms)
630 return Payload(
631 from_clause=lhs_payload.from_clause.join(rhs_payload.from_clause, onclause=on_clause),
632 where=lhs_payload.where + rhs_payload.where,
633 columns_available=columns_available,
634 )
635 case Select():
636 from_clause = self._select_to_executable(relation, ()).subquery()
637 columns_available = self.extract_mapping(relation.columns, from_clause.columns)
638 return Payload(from_clause, columns_available=columns_available)
639 case Materialization(name=name):
640 raise EngineError(
641 f"Cannot persist materialization {name!r} during SQL conversion; "
642 "use `lsst.daf.relation.Processor` first to handle this operation."
643 )
644 case Transfer(target=target):
645 raise EngineError(
646 f"Cannot handle transfer from {target.engine} during SQL conversion; "
647 "use `lsst.daf.relation.Processor` first to handle this operation."
648 )
649 raise NotImplementedError(f"Unsupported relation type {relation} for engine {self}.")
651 def convert_column_expression(
652 self, expression: ColumnExpression, columns_available: Mapping[ColumnTag, _L]
653 ) -> _L:
654 """Convert a `.ColumnExpression` to a logical column.
656 Parameters
657 ----------
658 expression : `.ColumnExpression`
659 Expression to convert.
660 columns_available : `~collections.abc.Mapping`
661 Mapping from `.ColumnTag` to logical column, typically produced by
662 `extract_mapping` or obtained from `Payload.columns_available`.
664 Returns
665 -------
666 logical_column
667 SQLAlchemy expression object or other logical column value.
669 See Also
670 --------
671 :ref:`lsst.daf.relation-sql-logical-columns`
672 """
673 match expression:
674 case ColumnLiteral(value=value):
675 return self.convert_column_literal(value)
676 case ColumnReference(tag=tag):
677 return columns_available[tag]
678 case ColumnFunction(name=name, args=args):
679 sql_args = [self.convert_column_expression(arg, columns_available) for arg in args]
680 if (function := self.get_function(name)) is not None:
681 return function(*sql_args)
682 return getattr(sql_args[0], name)(*sql_args[1:])
683 raise AssertionError(
684 f"matches should be exhaustive and all branches should return; got {expression!r}."
685 )
687 def convert_column_literal(self, value: Any) -> _L:
688 """Convert a Python literal value to a logical column.
690 Parameters
691 ----------
692 value
693 Python value to convert.
695 Returns
696 -------
697 logical_column
698 SQLAlchemy expression object or other logical column value.
700 Notes
701 -----
702 This method must be overridden to support a custom logical columns.
704 See Also
705 --------
706 :ref:`lsst.daf.relation-sql-logical-columns`
707 """
708 return sqlalchemy.sql.literal(value)
710 def expect_column_scalar(self, logical_column: _L) -> sqlalchemy.sql.ColumnElement:
711 """Convert a logical column value to a SQLAlchemy expression.
713 Parameters
714 ----------
715 logical_column
716 SQLAlchemy expression object or other logical column value.
718 Returns
719 -------
720 sql : `sqlalchemy.sql.ColumnElement`
721 SQLAlchemy expression object.
723 Notes
724 -----
725 The default implementation assumes the logical column type is just a
726 SQLAlchemy type and returns the given object unchanged. Subclasses
727 with a custom logical column type should override to at least assert
728 that the value is in fact a SQLAlchemy expression. This is only called
729 in contexts where true SQLAlchemy expressions are required, such as in
730 ``ORDER BY`` or ``WHERE`` clauses.
732 See Also
733 --------
734 :ref:`lsst.daf.relation-sql-logical-columns`
735 """
736 return cast(sqlalchemy.sql.ColumnElement, logical_column)
738 def convert_flattened_predicate(
739 self, predicate: Predicate, columns_available: Mapping[ColumnTag, _L]
740 ) -> list[sqlalchemy.sql.ColumnElement]:
741 """Flatten all logical AND operators in a `.Predicate` and convert each
742 to a boolean SQLAlchemy expression.
744 Parameters
745 ----------
746 predicate : `.Predicate`
747 Predicate to convert.
748 columns_available : `~collections.abc.Mapping`
749 Mapping from `.ColumnTag` to logical column, typically produced by
750 `extract_mapping` or obtained from `Payload.columns_available`.
752 Returns
753 -------
754 sql : `list` [ `sqlalchemy.sql.ColumnElement` ]
755 List of boolean SQLAlchemy expressions to be combined with the
756 ``AND`` operator.
757 """
758 if (flattened := flatten_logical_and(predicate)) is False:
759 return [sqlalchemy.sql.literal(False)]
760 else:
761 return [self.convert_predicate(p, columns_available) for p in flattened]
763 def convert_predicate(
764 self, predicate: Predicate, columns_available: Mapping[ColumnTag, _L]
765 ) -> sqlalchemy.sql.ColumnElement:
766 """Convert a `.Predicate` to a SQLAlchemy expression.
768 Parameters
769 ----------
770 predicate : `.Predicate`
771 Predicate to convert.
772 columns_available : `~collections.abc.Mapping`
773 Mapping from `.ColumnTag` to logical column, typically produced by
774 `extract_mapping` or obtained from `Payload.columns_available`.
776 Returns
777 -------
778 sql : `sqlalchemy.sql.ColumnElement`
779 Boolean SQLAlchemy expression.
780 """
781 match predicate:
782 case PredicateFunction(name=name, args=args):
783 sql_args = [self.convert_column_expression(arg, columns_available) for arg in args]
784 if (function := self.get_function(name)) is not None:
785 return function(*sql_args)
786 return getattr(sql_args[0], name)(*sql_args[1:])
787 case LogicalAnd(operands=operands):
788 if not operands:
789 return sqlalchemy.sql.literal(True)
790 if len(operands) == 1:
791 return self.convert_predicate(operands[0], columns_available)
792 else:
793 return sqlalchemy.sql.and_(
794 *[self.convert_predicate(operand, columns_available) for operand in operands]
795 )
796 case LogicalOr(operands=operands):
797 if not operands:
798 return sqlalchemy.sql.literal(False)
799 if len(operands) == 1:
800 return self.convert_predicate(operands[0], columns_available)
801 else:
802 return sqlalchemy.sql.or_(
803 *[self.convert_predicate(operand, columns_available) for operand in operands]
804 )
805 case LogicalNot(operand=operand):
806 return sqlalchemy.sql.not_(self.convert_predicate(operand, columns_available))
807 case PredicateReference(tag=tag):
808 return self.expect_column_scalar(columns_available[tag])
809 case PredicateLiteral(value=value):
810 return sqlalchemy.sql.literal(value)
811 case ColumnInContainer(item=item, container=container):
812 sql_item = self.expect_column_scalar(self.convert_column_expression(item, columns_available))
813 match container:
814 case ColumnRangeLiteral(value=range(start=start, stop=stop_exclusive, step=step)):
815 # The convert_column_literal calls below should just
816 # call sqlalchemy.sql.literal(int), which would also
817 # happen automatically internal to any of the other
818 # sqlalchemy function calls, but they get the typing
819 # right, reflecting the fact that the derived engine is
820 # supposed to have final say over how we convert
821 # literals.
822 stop_inclusive = stop_exclusive - 1
823 if start == stop_inclusive:
824 return sql_item == self.convert_column_literal(start)
825 else:
826 target = sqlalchemy.sql.between(
827 sql_item,
828 self.convert_column_literal(start),
829 self.convert_column_literal(stop_inclusive),
830 )
831 if step != 1:
832 return sqlalchemy.sql.and_(
833 *[
834 target,
835 sql_item % self.convert_column_literal(step)
836 == self.convert_column_literal(start % step),
837 ]
838 )
839 else:
840 return target
841 case ColumnExpressionSequence(items=items):
842 return sql_item.in_(
843 [self.convert_column_expression(item, columns_available) for item in items]
844 )
845 raise AssertionError(
846 f"matches should be exhaustive and all branches should return; got {predicate!r}."
847 )
849 def convert_sort_term(
850 self, term: SortTerm, columns_available: Mapping[ColumnTag, _L]
851 ) -> sqlalchemy.sql.ColumnElement:
852 """Convert a `.SortTerm` to a SQLAlchemy expression.
854 Parameters
855 ----------
856 term : `.SortTerm`
857 Sort term to convert.
858 columns_available : `~collections.abc.Mapping`
859 Mapping from `.ColumnTag` to logical column, typically produced by
860 `extract_mapping` or obtained from `Payload.columns_available`.
862 Returns
863 -------
864 sql : `sqlalchemy.sql.ColumnElement`
865 Scalar SQLAlchemy expression.
866 """
867 result = self.expect_column_scalar(self.convert_column_expression(term.expression, columns_available))
868 if term.ascending:
869 return result
870 else:
871 return result.desc()