Coverage for python/lsst/daf/relation/sql/_engine.py: 12%
292 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-07 02:01 -0800
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-07 02:01 -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/>.
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.has_projection:
230 return select.reapply_skip(
231 after=operation,
232 projection=Projection(frozenset(select.columns | {tag})),
233 )
234 else:
235 return select.reapply_skip(after=operation)
236 case Deduplication():
237 if not select.has_deduplication:
238 if select.has_slice:
239 # There was a Slice upstream, which needs to be applied
240 # before this Deduplication, so we nest the existing
241 # subquery within a new one that has deduplication.
242 return Select.apply_skip(select, deduplication=operation)
243 else:
244 # Move the Deduplication into the Select's
245 # operations.
246 return select.reapply_skip(deduplication=operation)
247 else:
248 # There was already another (redundant)
249 # Deduplication upstream.
250 return select
251 case Projection():
252 if select.has_deduplication:
253 # There was a Duplication upstream, so we need to ensure
254 # that is applied before this Projection via a nested
255 # subquery. Instead of applying the existing subquery
256 # operations, though, we apply one without any sort or
257 # slice that might exist, and save those for the new outer
258 # query, since putting those in a subquery would destroy
259 # the ordering.
260 subquery = select.reapply_skip(sort=None, slice=None)
261 return Select.apply_skip(
262 subquery,
263 projection=operation,
264 sort=select.sort,
265 slice=select.slice,
266 )
267 else:
268 # No Deduplication, so we can just add the Projection to
269 # the existing Select and reapply it.
270 match select.skip_to:
271 case BinaryOperationRelation(operation=Chain() as chain, lhs=lhs, rhs=rhs):
272 # ... unless the skip_to relation is a Chain; we
273 # want to move the Projection inside the Chain, to
274 # make it easier to avoid subqueries in UNION [ALL]
275 # constructs later.
276 return select.reapply_skip(
277 skip_to=chain._finish_apply(operation.apply(lhs), operation.apply(rhs)),
278 projection=None,
279 )
280 case _:
281 return select.reapply_skip(projection=operation)
282 case Selection():
283 if select.has_slice:
284 # There was a Slice upstream, which needs to be applied
285 # before this Selection via a nested subquery (which means
286 # we just apply the Selection to target, then add a new
287 # empty subquery marker).
288 return Select.apply_skip(operation._finish_apply(select))
289 else:
290 return select.reapply_skip(after=operation)
291 case Slice():
292 return select.reapply_skip(slice=select.slice.then(operation))
293 case Sort():
294 if select.has_slice:
295 # There was a Slice upstream, which needs to be applied
296 # before this Sort via a nested subquery (which means we
297 # apply a new Sort-only Select to 'select' itself).
298 return Select.apply_skip(select, sort=operation)
299 else:
300 return select.reapply_skip(sort=select.sort.then(operation))
301 case PartialJoin(binary=binary, fixed=fixed, fixed_is_lhs=fixed_is_lhs):
302 if fixed_is_lhs:
303 return self.append_binary(binary, fixed, select)
304 else:
305 return self.append_binary(binary, select, fixed)
306 case Identity():
307 return select
308 raise NotImplementedError(f"Unsupported operation type {operation} for engine {self}.")
310 def append_binary(self, operation: BinaryOperation, lhs: Relation, rhs: Relation) -> Select:
311 # Docstring inherited.
312 conformed_lhs = self.conform(lhs)
313 conformed_rhs = self.conform(rhs)
314 return self._append_binary_to_select(operation, conformed_lhs, conformed_rhs)
316 def _append_binary_to_select(self, operation: BinaryOperation, lhs: Select, rhs: Select) -> Select:
317 """Internal recursive implementation of `append_binary`.
319 This method should not be called by external code, but it may be
320 overridden and called recursively by derived engine classes.
322 Parameters
323 ----------
324 operation : `UnaryOperation`
325 Operation to add to the tree.
326 lhs : `Select`
327 Existing already-conformed relation tree.
328 rhs : `Select`
329 The other existing already-conformed relation tree.
331 Returns
332 -------
333 appended : `Select`
334 Conformed relation tree that includes the given operation.
335 """
336 if lhs.has_sort and not lhs.has_slice:
337 raise RelationalAlgebraError(
338 f"Applying binary operation {operation} to relation {lhs} will not preserve row order."
339 )
340 if rhs.has_sort and not rhs.has_slice:
341 raise RelationalAlgebraError(
342 f"Applying binary operation {operation} to relation {rhs} will not preserve row order."
343 )
344 match operation:
345 case Chain():
346 # Chain operands should always be Selects, but they can't have
347 # Sorts or Slices, at least not directly; if those exist we
348 # need to move them into subqueries at which point we might as
349 # well move any Projection or Dedulication there as well). But
350 # because of the check on Sorts-without-Slices above, there has
351 # to be a Slice here for there to be a Sort here.
352 if lhs.has_slice:
353 lhs = Select.apply_skip(lhs)
354 if rhs.has_slice:
355 rhs = Select.apply_skip(rhs)
356 return Select.apply_skip(operation._finish_apply(lhs, rhs))
357 case Join():
358 # For Joins, we move Projections after the Join, unless we have
359 # another reason to make a subquery before the join, and the
360 # operands are only Selects if they need to be subqueries.
361 new_lhs, new_lhs_needs_projection = lhs.strip()
362 new_rhs, new_rhs_needs_projection = rhs.strip()
363 if new_lhs_needs_projection or new_rhs_needs_projection:
364 projection = Projection(frozenset(lhs.columns | rhs.columns))
365 else:
366 projection = None
367 return Select.apply_skip(operation._finish_apply(new_lhs, new_rhs), projection=projection)
368 case IgnoreOne(ignore_lhs=ignore_lhs):
369 if ignore_lhs:
370 return rhs
371 else:
372 return lhs
373 raise AssertionError(f"Match on {operation} should be exhaustive and all branches return..")
375 def extract_mapping(
376 self, tags: Iterable[ColumnTag], sql_columns: sqlalchemy.sql.ColumnCollection
377 ) -> dict[ColumnTag, _L]:
378 """Extract a mapping with `.ColumnTag` keys and logical column values
379 from a SQLAlchemy column collection.
381 Parameters
382 ----------
383 tags : `Iterable`
384 Set of `.ColumnTag` objects whose logical columns should be
385 extracted.
386 sql_columns : `sqlalchemy.sql.ColumnCollection`
387 SQLAlchemy collection of columns, such as
388 `sqlalchemy.sql.FromClause.columns`.
390 Returns
391 -------
392 logical_columns : `dict`
393 Dictionary mapping `.ColumnTag` to logical column type.
395 Notes
396 -----
397 This method must be overridden to support a custom logical columns.
398 """
399 return {tag: cast(_L, sql_columns[tag.qualified_name]) for tag in tags}
401 def select_items(
402 self,
403 items: Iterable[tuple[ColumnTag, _L]],
404 sql_from: sqlalchemy.sql.FromClause,
405 *extra: sqlalchemy.sql.ColumnElement,
406 ) -> sqlalchemy.sql.Select:
407 """Construct a SQLAlchemy representation of a SELECT query.
409 Parameters
410 ----------
411 items : `Iterable` [ `tuple` ]
412 Iterable of (`.ColumnTag`, logical column) pairs. This is
413 typically the ``items()`` of a mapping returned by
414 `extract_mapping` or obtained from `Payload.columns_available`.
415 sql_from : `sqlalchemy.sql.FromClause`
416 SQLAlchemy representation of a FROM clause, such as a single table,
417 aliased subquery, or join expression. Must provide all columns
418 referenced by ``items``.
419 *extra : `sqlalchemy.sql.ColumnElement`
420 Additional SQL column expressions to include.
422 Returns
423 -------
424 select : `sqlalchemy.sql.Select`
425 SELECT query.
427 Notes
428 -----
429 This method is responsible for handling the case where ``items`` is
430 empty, typically by delegating to `handle_empty_columns`.
432 This method must be overridden to support a custom logical columns.
433 """
434 select_columns = [
435 cast(sqlalchemy.sql.ColumnElement, logical_column).label(tag.qualified_name)
436 for tag, logical_column in items
437 ]
438 select_columns.extend(extra)
439 self.handle_empty_columns(select_columns)
440 return sqlalchemy.sql.select(*select_columns).select_from(sql_from)
442 def handle_empty_columns(self, columns: list[sqlalchemy.sql.ColumnElement]) -> None:
443 """Handle the edge case where a SELECT statement has no columns, by
444 adding a literal column that should be ignored.
446 Parameters
447 ----------
448 columns : `list` [ `sqlalchemy.sql.ColumnElement` ]
449 List of SQLAlchemy column objects. This may have no elements when
450 this method is called, and must always have at least one element
451 when it returns.
452 """
453 if not columns:
454 columns.append(sqlalchemy.sql.literal(True).label(self.EMPTY_COLUMNS_NAME))
456 def to_executable(
457 self, relation: Relation, extra_columns: Iterable[sqlalchemy.sql.ColumnElement] = ()
458 ) -> sqlalchemy.sql.expression.SelectBase:
459 """Convert a relation tree to an executable SQLAlchemy expression.
461 Parameters
462 ----------
463 relation : `Relation`
464 The relation tree to convert.
465 extra_columns : `~collections.abc.Iterable`
466 Iterable of additional SQLAlchemy column objects to include
467 directly in the ``SELECT`` clause.
469 Returns
470 -------
471 select : `sqlalchemy.sql.expression.SelectBase`
472 A SQLAlchemy ``SELECT`` or compound ``SELECT`` query.
474 Notes
475 -----
476 This method requires all relations in the tree to have the same engine
477 (``self``). It also cannot handle `.Materialization` operations
478 unless they have already been processed once already (and hence have
479 a payload attached). Use the `.Processor` function to handle both of
480 these cases.
481 """
482 if relation.engine != self:
483 raise EngineError(
484 f"Engine {self!r} cannot operate on relation {relation} with engine {relation.engine!r}. "
485 "Use lsst.daf.relation.Processor to evaluate transfers first."
486 )
487 relation = self.conform(relation)
488 return self._select_to_executable(relation, extra_columns)
490 def _select_to_executable(
491 self,
492 select: Select,
493 extra_columns: Iterable[sqlalchemy.sql.ColumnElement],
494 ) -> sqlalchemy.sql.expression.SelectBase:
495 """Internal recursive implementation of `to_executable`.
497 This method should not be called by external code, but it may be
498 overridden and called recursively by derived engine classes.
500 Parameters
501 ----------
502 columns : `~collections.abc.Set` [ `ColumnTag` ]
503 Columns to include in the ``SELECT`` clause.
504 select : `Select`
505 Already-conformed relation tree to convert.
506 extra_columns : `~collections.abc.Iterable`
507 Iterable of additional SQLAlchemy column objects to include
508 directly in the ``SELECT`` clause.
510 Returns
511 -------
512 executable : `sqlalchemy.sql.expression.SelectBase`
513 SQLAlchemy ``SELECT`` or compound ``SELECT`` object.
515 Notes
516 -----
517 This method handles trees terminated with `Select`, the operation
518 relation types managed by `Select`, and `Chain` operations nested
519 directly as the `skip_to` target of a `Select`. It delegates to
520 `to_payload` for all other relation types.
521 """
522 columns_available: Mapping[ColumnTag, _L] | None = None
523 match select.skip_to:
524 case BinaryOperationRelation(operation=Chain(), lhs=lhs, rhs=rhs):
525 lhs_executable = self._select_to_executable(cast(Select, lhs), extra_columns)
526 rhs_executable = self._select_to_executable(cast(Select, rhs), extra_columns)
527 if select.has_deduplication:
528 executable = sqlalchemy.sql.union(lhs_executable, rhs_executable)
529 else:
530 executable = sqlalchemy.sql.union_all(lhs_executable, rhs_executable)
531 case _:
532 if select.skip_to.payload is not None:
533 payload: Payload[_L] = select.skip_to.payload
534 else:
535 payload = self.to_payload(select.skip_to)
536 if not select.columns and not extra_columns:
537 extra_columns = list(extra_columns)
538 self.handle_empty_columns(extra_columns)
539 columns_available = payload.columns_available
540 columns_projected = {tag: columns_available[tag] for tag in select.columns}
541 executable = self.select_items(columns_projected.items(), payload.from_clause, *extra_columns)
542 if len(payload.where) == 1:
543 executable = executable.where(payload.where[0])
544 elif payload.where:
545 executable = executable.where(sqlalchemy.sql.and_(*payload.where))
546 if select.has_deduplication:
547 executable = executable.distinct()
548 if select.has_sort:
549 if columns_available is None:
550 columns_available = self.extract_mapping(select.skip_to.columns, executable.selected_columns)
551 executable = executable.order_by(
552 *[self.convert_sort_term(term, columns_available) for term in select.sort.terms]
553 )
554 if select.slice.start:
555 executable = executable.offset(select.slice.start)
556 if select.slice.limit is not None:
557 executable = executable.limit(select.slice.limit)
558 return executable
560 def to_payload(self, relation: Relation) -> Payload[_L]:
561 """Internal recursive implementation of `to_executable`.
563 This method should not be called by external code, but it may be
564 overridden and called recursively by derived engine classes.
566 Parameters
567 ----------
568 relation : `Relation`
569 Relation to convert. This method handles all operation relation
570 types other than `Chain` and the operations managed by `Select`.
572 Returns
573 -------
574 payload : `Payload`
575 Struct containing a SQLAlchemy represenation of a simple ``SELECT``
576 query.
577 """
578 assert relation.engine == self, "Should be guaranteed by callers."
579 if relation.payload is not None: # Should cover all LeafRelations
580 return relation.payload
581 match relation:
582 case UnaryOperationRelation(operation=operation, target=target):
583 match operation:
584 case Calculation(tag=tag, expression=expression):
585 result = self.to_payload(target).copy()
586 result.columns_available[tag] = self.convert_column_expression(
587 expression, result.columns_available
588 )
589 return result
590 case Selection(predicate=predicate):
591 result = self.to_payload(target).copy()
592 result.where.extend(
593 self.convert_flattened_predicate(predicate, result.columns_available)
594 )
595 return result
596 case BinaryOperationRelation(
597 operation=Join(predicate=predicate, common_columns=common_columns), lhs=lhs, rhs=rhs
598 ):
599 lhs_payload = self.to_payload(lhs)
600 rhs_payload = self.to_payload(rhs)
601 assert common_columns is not None, "Guaranteed by Join.apply and PartialJoin.apply."
602 on_terms: list[sqlalchemy.sql.ColumnElement] = []
603 if common_columns:
604 on_terms.extend(
605 lhs_payload.columns_available[tag] == rhs_payload.columns_available[tag]
606 for tag in common_columns
607 )
608 columns_available = {**lhs_payload.columns_available, **rhs_payload.columns_available}
609 if predicate.as_trivial() is not True:
610 on_terms.extend(self.convert_flattened_predicate(predicate, columns_available))
611 on_clause: sqlalchemy.sql.ColumnElement
612 if not on_terms:
613 on_clause = sqlalchemy.sql.literal(True)
614 elif len(on_terms) == 1:
615 on_clause = on_terms[0]
616 else:
617 on_clause = sqlalchemy.sql.and_(*on_terms)
618 return Payload(
619 from_clause=lhs_payload.from_clause.join(rhs_payload.from_clause, onclause=on_clause),
620 where=lhs_payload.where + rhs_payload.where,
621 columns_available=columns_available,
622 )
623 case Select():
624 from_clause = self._select_to_executable(relation, ()).subquery()
625 columns_available = self.extract_mapping(relation.columns, from_clause.columns)
626 return Payload(from_clause, columns_available=columns_available)
627 case Materialization(name=name):
628 raise EngineError(
629 f"Cannot persist materialization {name!r} during SQL conversion; "
630 "use `lsst.daf.relation.Processor` first to handle this operation."
631 )
632 case Transfer(target=target):
633 raise EngineError(
634 f"Cannot handle transfer from {target.engine} during SQL conversion; "
635 "use `lsst.daf.relation.Processor` first to handle this operation."
636 )
637 raise NotImplementedError(f"Unsupported relation type {relation} for engine {self}.")
639 def convert_column_expression(
640 self, expression: ColumnExpression, columns_available: Mapping[ColumnTag, _L]
641 ) -> _L:
642 """Convert a `.ColumnExpression` to a logical column.
644 Parameters
645 ----------
646 expression : `.ColumnExpression`
647 Expression to convert.
648 columns_available : `~collections.abc.Mapping`
649 Mapping from `.ColumnTag` to logical column, typically produced by
650 `extract_mapping` or obtained from `Payload.columns_available`.
652 Returns
653 -------
654 logical_column
655 SQLAlchemy expression object or other logical column value.
657 See Also
658 --------
659 :ref:`lsst.daf.relation-sql-logical-columns`
660 """
661 match expression:
662 case ColumnLiteral(value=value):
663 return self.convert_column_literal(value)
664 case ColumnReference(tag=tag):
665 return columns_available[tag]
666 case ColumnFunction(name=name, args=args):
667 sql_args = [self.convert_column_expression(arg, columns_available) for arg in args]
668 if (function := self.get_function(name)) is not None:
669 return function(*sql_args)
670 return getattr(sql_args[0], name)(*sql_args[1:])
671 raise AssertionError(
672 f"matches should be exhaustive and all branches should return; got {expression!r}."
673 )
675 def convert_column_literal(self, value: Any) -> _L:
676 """Convert a Python literal value to a logical column.
678 Parameters
679 ----------
680 value
681 Python value to convert.
683 Returns
684 -------
685 logical_column
686 SQLAlchemy expression object or other logical column value.
688 Notes
689 -----
690 This method must be overridden to support a custom logical columns.
692 See Also
693 --------
694 :ref:`lsst.daf.relation-sql-logical-columns`
695 """
696 return sqlalchemy.sql.literal(value)
698 def expect_column_scalar(self, logical_column: _L) -> sqlalchemy.sql.ColumnElement:
699 """Convert a logical column value to a SQLAlchemy expression.
701 Parameters
702 ----------
703 logical_column
704 SQLAlchemy expression object or other logical column value.
706 Returns
707 -------
708 sql : `sqlalchemy.sql.ColumnElement`
709 SQLAlchemy expression object.
711 Notes
712 -----
713 The default implementation assumes the logical column type is just a
714 SQLAlchemy type and returns the given object unchanged. Subclasses
715 with a custom logical column type should override to at least assert
716 that the value is in fact a SQLAlchemy expression. This is only called
717 in contexts where true SQLAlchemy expressions are required, such as in
718 ``ORDER BY`` or ``WHERE`` clauses.
720 See Also
721 --------
722 :ref:`lsst.daf.relation-sql-logical-columns`
723 """
724 return cast(sqlalchemy.sql.ColumnElement, logical_column)
726 def convert_flattened_predicate(
727 self, predicate: Predicate, columns_available: Mapping[ColumnTag, _L]
728 ) -> list[sqlalchemy.sql.ColumnElement]:
729 """Flatten all logical AND operators in a `.Predicate` and convert each
730 to a boolean SQLAlchemy expression.
732 Parameters
733 ----------
734 predicate : `.Predicate`
735 Predicate to convert.
736 columns_available : `~collections.abc.Mapping`
737 Mapping from `.ColumnTag` to logical column, typically produced by
738 `extract_mapping` or obtained from `Payload.columns_available`.
740 Returns
741 -------
742 sql : `list` [ `sqlalchemy.sql.ColumnElement` ]
743 List of boolean SQLAlchemy expressions to be combined with the
744 ``AND`` operator.
745 """
746 if (flattened := flatten_logical_and(predicate)) is False:
747 return [sqlalchemy.sql.literal(False)]
748 else:
749 return [self.convert_predicate(p, columns_available) for p in flattened]
751 def convert_predicate(
752 self, predicate: Predicate, columns_available: Mapping[ColumnTag, _L]
753 ) -> sqlalchemy.sql.ColumnElement:
754 """Convert a `.Predicate` to a SQLAlchemy expression.
756 Parameters
757 ----------
758 predicate : `.Predicate`
759 Predicate to convert.
760 columns_available : `~collections.abc.Mapping`
761 Mapping from `.ColumnTag` to logical column, typically produced by
762 `extract_mapping` or obtained from `Payload.columns_available`.
764 Returns
765 -------
766 sql : `sqlalchemy.sql.ColumnElement`
767 Boolean SQLAlchemy expression.
768 """
769 match predicate:
770 case PredicateFunction(name=name, args=args):
771 sql_args = [self.convert_column_expression(arg, columns_available) for arg in args]
772 if (function := self.get_function(name)) is not None:
773 return function(*sql_args)
774 return getattr(sql_args[0], name)(*sql_args[1:])
775 case LogicalAnd(operands=operands):
776 if not operands:
777 return sqlalchemy.sql.literal(True)
778 if len(operands) == 1:
779 return self.convert_predicate(operands[0], columns_available)
780 else:
781 return sqlalchemy.sql.and_(
782 *[self.convert_predicate(operand, columns_available) for operand in operands]
783 )
784 case LogicalOr(operands=operands):
785 if not operands:
786 return sqlalchemy.sql.literal(False)
787 if len(operands) == 1:
788 return self.convert_predicate(operands[0], columns_available)
789 else:
790 return sqlalchemy.sql.or_(
791 *[self.convert_predicate(operand, columns_available) for operand in operands]
792 )
793 case LogicalNot(operand=operand):
794 return sqlalchemy.sql.not_(self.convert_predicate(operand, columns_available))
795 case PredicateReference(tag=tag):
796 return self.expect_column_scalar(columns_available[tag])
797 case PredicateLiteral(value=value):
798 return sqlalchemy.sql.literal(value)
799 case ColumnInContainer(item=item, container=container):
800 sql_item = self.expect_column_scalar(self.convert_column_expression(item, columns_available))
801 match container:
802 case ColumnRangeLiteral(value=range(start=start, stop=stop_exclusive, step=step)):
803 # The convert_column_literal calls below should just
804 # call sqlalchemy.sql.literal(int), which would also
805 # happen automatically internal to any of the other
806 # sqlalchemy function calls, but they get the typing
807 # right, reflecting the fact that the derived engine is
808 # supposed to have final say over how we convert
809 # literals.
810 stop_inclusive = stop_exclusive - 1
811 if start == stop_inclusive:
812 return sql_item == self.convert_column_literal(start)
813 else:
814 target = sqlalchemy.sql.between(
815 sql_item,
816 self.convert_column_literal(start),
817 self.convert_column_literal(stop_inclusive),
818 )
819 if step != 1:
820 return sqlalchemy.sql.and_(
821 *[
822 target,
823 sql_item % self.convert_column_literal(step)
824 == self.convert_column_literal(start % step),
825 ]
826 )
827 else:
828 return target
829 case ColumnExpressionSequence(items=items):
830 return sql_item.in_(
831 [self.convert_column_expression(item, columns_available) for item in items]
832 )
833 raise AssertionError(
834 f"matches should be exhaustive and all branches should return; got {predicate!r}."
835 )
837 def convert_sort_term(
838 self, term: SortTerm, columns_available: Mapping[ColumnTag, _L]
839 ) -> sqlalchemy.sql.ColumnElement:
840 """Convert a `.SortTerm` to a SQLAlchemy expression.
842 Parameters
843 ----------
844 term : `.SortTerm`
845 Sort term to convert.
846 columns_available : `~collections.abc.Mapping`
847 Mapping from `.ColumnTag` to logical column, typically produced by
848 `extract_mapping` or obtained from `Payload.columns_available`.
850 Returns
851 -------
852 sql : `sqlalchemy.sql.ColumnElement`
853 Scalar SQLAlchemy expression.
854 """
855 result = self.expect_column_scalar(self.convert_column_expression(term.expression, columns_available))
856 if term.ascending:
857 return result
858 else:
859 return result.desc()