Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/common/parameters.py: 82%
687 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 14:22 +0000
1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements. See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership. The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License. You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied. See the License for the
15# specific language governing permissions and limitations
16# under the License.
18from __future__ import annotations
20from abc import ABC, abstractmethod
21from collections.abc import Callable, Iterable, Sequence
22from datetime import datetime
23from enum import Enum
24from typing import (
25 TYPE_CHECKING,
26 Annotated,
27 Any,
28 Generic,
29 Literal,
30 TypeVar,
31 cast,
32 overload,
33)
35from fastapi import Depends, HTTPException, Query, status
36from pendulum.parsing.exceptions import ParserError
37from pydantic import AfterValidator, BaseModel, NonNegativeInt
38from sqlalchemy import Column, String, and_, func, not_, or_, select as sql_select, true as sql_true
39from sqlalchemy.ext.compiler import compiles
40from sqlalchemy.inspection import inspect
41from sqlalchemy.orm import aliased
42from sqlalchemy.sql.functions import FunctionElement
44from airflow._shared.timezones import timezone
45from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT
46from airflow.api_fastapi.core_api.base import OrmClause
47from airflow.api_fastapi.core_api.security import GetUserDep
48from airflow.configuration import conf
49from airflow.models import Base
50from airflow.models.asset import (
51 AssetAliasModel,
52 AssetEvent,
53 AssetModel,
54 AssetPartitionDagRun,
55 DagScheduleAssetReference,
56 TaskInletAssetReference,
57 TaskOutletAssetReference,
58 association_table,
59)
60from airflow.models.connection import Connection
61from airflow.models.dag import DagModel, DagTag
62from airflow.models.dag_favorite import DagFavorite
63from airflow.models.dag_version import DagVersion
64from airflow.models.dagrun import DagRun
65from airflow.models.errors import ParseImportError
66from airflow.models.hitl import HITLDetail
67from airflow.models.pool import Pool
68from airflow.models.taskinstance import TaskInstance
69from airflow.models.variable import Variable
70from airflow.models.xcom import XComModel
71from airflow.typing_compat import Self
72from airflow.utils.state import DagRunState, TaskInstanceState
73from airflow.utils.types import DagRunType
75if TYPE_CHECKING: 75 ↛ 76line 75 didn't jump to line 76 because the condition on line 75 was never true
76 from sqlalchemy.orm.attributes import InstrumentedAttribute
77 from sqlalchemy.sql import ColumnElement, Select
78 from sqlalchemy.sql.compiler import SQLCompiler
80 from airflow.serialization.definitions.dag import SerializedDAG
82T = TypeVar("T")
84_FALLBACK_PAGE_LIMIT: int = conf.getint("api", "fallback_page_limit")
87class _MySQLCollate(FunctionElement):
88 """
89 Wraps a SQL expression so that on MySQL it is emitted with an explicit ``COLLATE`` clause.
91 On every other dialect the expression is passed through unchanged.
93 This is needed when a computed expression (e.g. a ``CASE … END`` that mixes
94 a stored ``VARCHAR`` column with a ``CAST(integer AS CHAR)``) ends up with
95 MySQL coercibility ``NONE`` because the two branches carry different implicit
96 collations. Comparing such an expression with a bound parameter fails with
97 "Illegal mix of collations". Wrapping the expression in an explicit
98 ``COLLATE`` gives it ``EXPLICIT`` coercibility, which MySQL accepts in all
99 comparison operators.
100 """
102 type = String()
103 inherit_cache = True
105 def __init__(self, expr: ColumnElement[Any], collation: str) -> None:
106 super().__init__(expr)
107 self.collation = collation
110@compiles(_MySQLCollate)
111def _compile_mysql_collate_default(element: _MySQLCollate, compiler: SQLCompiler, **kw: Any) -> str:
112 """Non-MySQL: render the inner expression without any COLLATE clause."""
113 (expr,) = element.clauses
114 return compiler.process(expr, **kw)
117@compiles(_MySQLCollate, "mysql")
118def _compile_mysql_collate_mysql(element: _MySQLCollate, compiler: SQLCompiler, **kw: Any) -> str:
119 """MySQL: wrap the inner expression with the requested COLLATE clause."""
120 (expr,) = element.clauses
121 inner = compiler.process(expr, **kw)
122 return f"({inner}) COLLATE {element.collation}"
125class BaseParam(OrmClause[T], ABC):
126 """Base class for path or query parameters with ORM transformation."""
128 def __init__(self, value: T | None = None, skip_none: bool = True) -> None:
129 super().__init__(value)
130 self.attribute: ColumnElement | InstrumentedAttribute | None = None
131 self.skip_none = skip_none
133 def set_value(self, value: T | None) -> Self:
134 self.value = value
135 return self
137 @classmethod
138 @abstractmethod
139 def depends(cls, *args: Any, **kwargs: Any) -> Self:
140 pass
143class LimitFilter(BaseParam[NonNegativeInt]):
144 """Filter on the limit."""
146 def to_orm(self, select: Select) -> Select:
147 if self.value is None and self.skip_none: 147 ↛ 148line 147 didn't jump to line 148 because the condition on line 147 was never true
148 return select
150 return select.limit(self.value)
152 @classmethod
153 def depends(cls, limit: NonNegativeInt = _FALLBACK_PAGE_LIMIT) -> LimitFilter:
154 return cls().set_value(min(limit, conf.getint("api", "maximum_page_limit")))
157class OffsetFilter(BaseParam[NonNegativeInt]):
158 """Filter on offset."""
160 def to_orm(self, select: Select) -> Select:
161 if self.value is None and self.skip_none: 161 ↛ 162line 161 didn't jump to line 162 because the condition on line 161 was never true
162 return select
163 return select.offset(self.value)
165 @classmethod
166 def depends(
167 cls,
168 offset: NonNegativeInt = 0,
169 ) -> OffsetFilter:
170 return cls().set_value(offset)
173class _FavoriteFilter(BaseParam[bool]):
174 """Filter Dags by favorite status."""
176 def __init__(self, user_id: str, value: T | None = None, skip_none: bool = True) -> None:
177 super().__init__(skip_none=skip_none)
178 self.user_id = user_id
180 def to_orm(self, select_stmt: Select) -> Select:
181 if self.value is None and self.skip_none:
182 return select_stmt
184 if self.value:
185 select_stmt = select_stmt.join(DagFavorite, DagFavorite.dag_id == DagModel.dag_id).where(
186 DagFavorite.user_id == self.user_id
187 )
188 else:
189 select_stmt = select_stmt.where(
190 not_(
191 sql_select(DagFavorite)
192 .where(and_(DagFavorite.dag_id == DagModel.dag_id, DagFavorite.user_id == self.user_id))
193 .exists()
194 )
195 )
197 return select_stmt
199 @classmethod
200 def depends(cls, user: GetUserDep, is_favorite: bool | None = Query(None)) -> _FavoriteFilter:
201 instance = cls(user_id=str(user.get_id())).set_value(is_favorite)
202 return instance
205class _ExcludeStaleFilter(BaseParam[bool]):
206 """Filter on is_stale."""
208 def to_orm(self, select: Select) -> Select:
209 if self.value and self.skip_none:
210 return select.where(DagModel.is_stale != self.value)
211 return select
213 @classmethod
214 def depends(cls, exclude_stale: bool = True) -> _ExcludeStaleFilter:
215 return cls().set_value(exclude_stale)
218class _PrefixPatternParam(BaseParam[str], ABC):
219 """
220 Shared prefix pattern: pipe ``|`` for OR, ``~`` → empty (match all), Unicode prefix range.
222 .. note::
223 Trailing non-alphanumeric characters in a search term are stripped before the range
224 is computed. A range scan with a punctuation-terminated upper bound is unsafe under
225 PostgreSQL's default locale-aware collation (``en_US.utf8`` sorts punctuation in a
226 way that breaks the range), and additionally stopping the range at an alphanumeric
227 character keeps the upper bound alphanumeric too, so the predicate stays usable by
228 default btree indexes. A user who asks for prefix ``"test_"`` gets matches starting
229 with ``"test"`` — a small over-match trade-off made explicit in the public
230 ``*_prefix_pattern`` query-param description.
231 """
233 pipe_as_or: bool = True
235 @staticmethod
236 def _prefix_range_upper(term: str) -> str | None:
237 """
238 Compute the exclusive upper bound for a prefix range scan.
240 Returns ``None`` if the term has no alphanumeric characters. Trailing non-alphanumeric
241 characters are dropped before bumping the last character so the resulting upper bound
242 is itself alphanumeric and behaves predictably under locale-aware collations. If
243 incrementing would land outside the alphanumeric range (e.g. ``'z' → '{'``), we drop
244 that character and retry.
245 """
246 while term and not term[-1].isalnum():
247 term = term[:-1]
248 if not term:
249 return None
250 last = ord(term[-1])
251 if last >= 0x10FFFF: 251 ↛ 252line 251 didn't jump to line 252 because the condition on line 251 was never true
252 return _PrefixPatternParam._prefix_range_upper(term[:-1])
253 bumped = chr(last + 1)
254 if not bumped.isalnum():
255 return _PrefixPatternParam._prefix_range_upper(term[:-1])
256 return term[:-1] + bumped
258 @staticmethod
259 def _prefix_lower_bound(term: str) -> str:
260 """Return the matching lower bound: strip trailing non-alphanumeric chars to pair with the upper."""
261 while term and not term[-1].isalnum():
262 term = term[:-1]
263 return term
265 @abstractmethod
266 def _prefix_clause(self, term: str):
267 """Return the SQL boolean for one prefix term (including empty string after ``~`` alias)."""
269 def to_orm(self, select: Select) -> Select:
270 # ``skip_none`` only gates the "no value" behavior for the callers that must keep
271 # the filter slot present (e.g. ``QueryDagIdPrefixPatternSearchWithNone``); applying
272 # a ``None`` value as a filter produces nonsense predicates, so always skip it here.
273 if self.value is None:
274 return select
276 val_str = str(self.value)
277 if self.pipe_as_or and "|" in val_str:
278 search_terms = [term.strip() for term in val_str.split("|") if term.strip()]
279 if search_terms: 279 ↛ 282line 279 didn't jump to line 282 because the condition on line 279 was always true
280 return select.where(or_(*(self._prefix_clause(term) for term in search_terms)))
282 return select.where(self._prefix_clause(val_str))
284 def transform_aliases(self, value: str | None) -> str | None:
285 if value == "~": 285 ↛ 286line 285 didn't jump to line 286 because the condition on line 285 was never true
286 value = ""
287 return value
290def _build_pipe_clause(pipe_as_or: bool) -> str:
291 """Build the per-parameter pipe note. OR is the documented default (see the API description), so only the literal exception is spelled out."""
292 return "" if pipe_as_or else "Here `|` is matched literally, not as OR. "
295_LIKE_ESCAPE_CHAR = "\\"
298def _escape_like_pattern(value: str) -> str:
299 r"""
300 Escape SQL ``LIKE`` / ``ILIKE`` metacharacters in a user-supplied value.
302 Use together with ``column.ilike(f"%{_escape_like_pattern(value)}%", escape="\\")`` on filter
303 parameters that intend literal substring matching (so a user-supplied ``%`` or ``_`` does not
304 widen the match beyond what the filter semantics promise). Search parameters that explicitly
305 expose wildcard semantics (see :class:`_SearchParam`) must not call this — they want the
306 metacharacters to pass through.
307 """
308 return (
309 value.replace(_LIKE_ESCAPE_CHAR, _LIKE_ESCAPE_CHAR * 2)
310 .replace("%", _LIKE_ESCAPE_CHAR + "%")
311 .replace("_", _LIKE_ESCAPE_CHAR + "_")
312 )
315class _SearchParam(BaseParam[str]):
316 """
317 Substring search on a column using ``ILIKE '%term%'`` (case-insensitive).
319 .. note::
320 This full-match substring search most of the time prevents the database
321 from using B-tree indexes on ``attribute``, which can be very slow on
322 large tables. Prefer :class:`_PrefixSearchParam` (the ``*_prefix_pattern``
323 query-param counterpart) when matching from the beginning of the value
324 is acceptable.
325 """
327 def __init__(self, attribute: ColumnElement, skip_none: bool = True, pipe_as_or: bool = True) -> None:
328 super().__init__(skip_none=skip_none)
329 self.attribute: ColumnElement = attribute
330 self.pipe_as_or = pipe_as_or
332 def to_orm(self, select: Select) -> Select:
333 if self.value is None and self.skip_none:
334 return select
336 val_str = str(self.value)
337 if self.pipe_as_or and "|" in val_str:
338 search_terms = [term.strip() for term in val_str.split("|") if term.strip()]
339 if search_terms: 339 ↛ 342line 339 didn't jump to line 342 because the condition on line 339 was always true
340 return select.where(or_(*(self.attribute.ilike(f"%{term}%") for term in search_terms)))
342 return select.where(self.attribute.ilike(f"%{val_str}%"))
344 def transform_aliases(self, value: str | None) -> str | None:
345 if value == "~": 345 ↛ 346line 345 didn't jump to line 346 because the condition on line 345 was never true
346 value = "%"
347 return value
349 @classmethod
350 def depends(cls, *args: Any, **kwargs: Any) -> Self:
351 raise NotImplementedError("Use search_param_factory instead , depends is not implemented.")
354class _PrefixSearchParam(_PrefixPatternParam):
355 """
356 Prefix search on a column using range comparison (case-sensitive, index-friendly).
358 Unlike :class:`_SearchParam`, wildcard characters are treated as literals and the query
359 plan can use the column's default B-tree index for the range scan. Trailing
360 non-alphanumeric characters in ``term`` are stripped first (see
361 :class:`_PrefixPatternParam` for why).
362 """
364 def __init__(self, attribute: ColumnElement, skip_none: bool = True, pipe_as_or: bool = True) -> None:
365 super().__init__(skip_none=skip_none)
366 self.attribute: ColumnElement = attribute
367 self.pipe_as_or = pipe_as_or
369 def _prefix_clause(self, term: str):
370 lower = self._prefix_lower_bound(term)
371 if not lower:
372 return self.attribute.is_not(None)
373 upper = self._prefix_range_upper(term)
374 if upper is None:
375 return self.attribute >= lower
376 return and_(self.attribute >= lower, self.attribute < upper)
378 @classmethod
379 def depends(cls, *args: Any, **kwargs: Any) -> Self:
380 raise NotImplementedError("Use prefix_search_param_factory instead, depends is not implemented.")
383class _TaskDisplayNamePrefixPatternParam(_PrefixPatternParam):
384 """
385 Prefix filter equivalent to :attr:`TaskInstance.task_display_name`, rewritten for composite-index use.
387 The hybrid expression ``coalesce(_task_display_property_value, task_id)`` cannot use those indexes;
388 this implementation applies an equivalent ``OR`` of simpler range predicates instead. Trailing
389 non-alphanumeric characters in ``term`` are stripped first (see :class:`_PrefixPatternParam`).
390 """
392 def _prefix_clause(self, term: str):
393 lower = self._prefix_lower_bound(term)
394 if not lower:
395 return sql_true()
396 upper = self._prefix_range_upper(term)
397 if upper is None:
398 return or_(
399 and_(
400 TaskInstance._task_display_property_value.is_(None),
401 TaskInstance.task_id >= lower,
402 ),
403 and_(
404 TaskInstance._task_display_property_value.is_not(None),
405 TaskInstance._task_display_property_value >= lower,
406 ),
407 )
408 return or_(
409 and_(
410 TaskInstance._task_display_property_value.is_(None),
411 TaskInstance.task_id >= lower,
412 TaskInstance.task_id < upper,
413 ),
414 and_(
415 TaskInstance._task_display_property_value.is_not(None),
416 TaskInstance._task_display_property_value >= lower,
417 TaskInstance._task_display_property_value < upper,
418 ),
419 )
421 @classmethod
422 def depends(
423 cls,
424 task_display_name_prefix_pattern: str | None = Query(
425 default=None,
426 description=(
427 "Case-sensitive prefix match on task display name (`_task_display_property_value` else "
428 "`task_id`). Index-friendly alternative to `task_display_name_pattern`; on large databases "
429 "combine with `dag_id_prefix_pattern` (or a specific Dag in the path) so composite indexes "
430 'apply. See "Filtering with pattern parameters".'
431 ),
432 ),
433 ) -> Self:
434 param = cls()
435 return param.set_value(param.transform_aliases(task_display_name_prefix_pattern))
438class QueryTaskInstanceTaskGroupFilter(BaseParam[str]):
439 """Task group filter - returns all tasks in the specified group."""
441 def __init__(self, dag=None, skip_none: bool = True):
442 super().__init__(skip_none=skip_none)
443 self._dag: None | SerializedDAG = dag
445 @property
446 def dag(self) -> None | SerializedDAG:
447 return self._dag
449 @dag.setter
450 def dag(self, value: None | SerializedDAG) -> None:
451 self._dag = value
453 def to_orm(self, select: Select) -> Select:
454 if self.value is None and self.skip_none:
455 return select
457 if not self.dag: 457 ↛ 458line 457 didn't jump to line 458 because the condition on line 457 was never true
458 raise ValueError("Dag must be set before calling to_orm")
460 if not hasattr(self.dag, "task_group"): 460 ↛ 461line 460 didn't jump to line 461 because the condition on line 460 was never true
461 return select
463 # Exact matching on group_id
464 task_groups = self.dag.task_group.get_task_group_dict()
465 task_group = task_groups.get(self.value)
466 if not task_group: 466 ↛ 475line 466 didn't jump to line 475 because the condition on line 466 was always true
467 raise HTTPException(
468 status.HTTP_404_NOT_FOUND,
469 detail={
470 "reason": "not_found",
471 "message": f"Task group {self.value} not found",
472 },
473 )
475 return select.where(TaskInstance.task_id.in_(task.task_id for task in task_group.iter_tasks()))
477 @classmethod
478 def depends(
479 cls,
480 value: str | None = Query(
481 alias="task_group_id",
482 default=None,
483 description="Filter by exact task group ID. Returns all tasks within the specified task group.",
484 ),
485 ) -> QueryTaskInstanceTaskGroupFilter:
486 return cls(dag=None).set_value(value)
489def search_param_factory(
490 attribute: ColumnElement,
491 pattern_name: str,
492 skip_none: bool = True,
493 pipe_as_or: bool = True,
494) -> Callable[[str | None], _SearchParam]:
495 prefix_pattern_name = pattern_name.replace("_pattern", "_prefix_pattern")
496 DESCRIPTION = (
497 "Case-insensitive substring match (SQL `ILIKE`). "
498 f"{_build_pipe_clause(pipe_as_or)}"
499 f'Slower than `{prefix_pattern_name}` on large tables — see "Filtering with pattern parameters".'
500 )
502 def depends_search(
503 value: str | None = Query(alias=pattern_name, default=None, description=DESCRIPTION),
504 ) -> _SearchParam:
505 search_parm = _SearchParam(attribute, skip_none, pipe_as_or=pipe_as_or)
506 value = search_parm.transform_aliases(value)
507 return search_parm.set_value(value)
509 return depends_search
512def prefix_search_param_factory(
513 attribute: ColumnElement,
514 prefix_pattern_name: str,
515 skip_none: bool = True,
516 pipe_as_or: bool = True,
517) -> Callable[[str | None], _PrefixSearchParam]:
518 """
519 Build a FastAPI ``Depends`` returning a :class:`_PrefixSearchParam` for prefix matching.
521 Prefer this over :func:`search_param_factory` for performance: prefix matching uses a
522 B-tree index range scan, while substring matching requires a full table scan.
523 """
524 DESCRIPTION = (
525 "Case-sensitive, index-friendly prefix match. "
526 f"{_build_pipe_clause(pipe_as_or)}"
527 'See "Filtering with pattern parameters".'
528 )
530 def depends_prefix_search(
531 value: str | None = Query(alias=prefix_pattern_name, default=None, description=DESCRIPTION),
532 ) -> _PrefixSearchParam:
533 search_parm = _PrefixSearchParam(attribute, skip_none, pipe_as_or=pipe_as_or)
534 value = search_parm.transform_aliases(value)
535 return search_parm.set_value(value)
537 return depends_prefix_search
540class SortParam(BaseParam[list[str]]):
541 """Order result by the attribute."""
543 MAX_SORT_PARAMS = 10
545 def __init__(
546 self,
547 allowed_attrs: list[str],
548 model: Base,
549 to_replace: dict[str, str | Column | list[Column]] | None = None,
550 ) -> None:
551 super().__init__()
552 self.allowed_attrs = allowed_attrs
553 self.model = model
554 self.to_replace = to_replace
555 self._cached_resolution: list[tuple[str, ColumnElement, bool]] | None = None
557 def set_value(self, value: list[str] | None) -> Self:
558 self._cached_resolution = None
559 return super().set_value(value)
561 def _resolve(self) -> list[tuple[str, ColumnElement, bool]]:
562 """Resolve sort columns as (attr_name, column, is_descending) tuples. Cached after first call."""
563 if self._cached_resolution is not None:
564 return self._cached_resolution
566 if self.value is None:
567 self.value = [self.get_primary_key_string()]
569 order_by_values = self.value
570 if len(order_by_values) > self.MAX_SORT_PARAMS:
571 raise HTTPException(
572 400,
573 f"Ordering with more than {self.MAX_SORT_PARAMS} parameters is not allowed. Provided: {order_by_values}",
574 )
576 resolved: list[tuple[str, ColumnElement, bool]] = []
577 for order_by_value in order_by_values:
578 lstriped_orderby = order_by_value.lstrip("-")
579 # Store the user-facing name in the resolved tuple. ``row_value`` resolves
580 # it back to the actual row accessor via ``to_replace`` when reading values
581 # for cursor encoding.
582 attr_name = lstriped_orderby
583 column: Column | None = None
584 if self.to_replace:
585 replacement = self.to_replace.get(lstriped_orderby, lstriped_orderby)
586 if isinstance(replacement, str): 586 ↛ 588line 586 didn't jump to line 588 because the condition on line 586 was always true
587 lstriped_orderby = replacement
588 elif isinstance(replacement, list):
589 # Compound sort: expand the list into multiple sort entries.
590 # Each column's ORM key becomes its attr_name so that
591 # row_value() can read the corresponding attribute via
592 # getattr(row, attr_name) without further to_replace lookups.
593 is_desc = order_by_value.startswith("-")
594 for col in replacement:
595 col_attr_name = col.key
596 resolved.append((col_attr_name, col, is_desc))
597 continue
598 else:
599 column = replacement
601 if (self.allowed_attrs and lstriped_orderby not in self.allowed_attrs) and column is None:
602 raise HTTPException(
603 400,
604 f"Ordering with '{lstriped_orderby}' is disallowed or "
605 f"the attribute does not exist on the model",
606 )
607 if column is None: 607 ↛ 610line 607 didn't jump to line 610 because the condition on line 607 was always true
608 column = getattr(self.model, lstriped_orderby)
610 resolved.append((attr_name, column, order_by_value.startswith("-")))
612 primary_key_column = self.get_primary_key_column()
613 pk_name = self.get_primary_key_string()
614 resolved_column_keys = {getattr(col, "key", None) for _, col, _ in resolved}
615 if pk_name not in resolved_column_keys:
616 pk_desc = bool(order_by_values and order_by_values[0].startswith("-"))
617 resolved.append((pk_name, primary_key_column, pk_desc))
619 self._cached_resolution = resolved
620 return self._cached_resolution
622 def to_orm(self, select: Select, *, reversed: bool = False) -> Select:
623 if self.skip_none is False: 623 ↛ 624line 623 didn't jump to line 624 because the condition on line 623 was never true
624 raise ValueError(f"Cannot set 'skip_none' to False on a {type(self)}")
626 resolved = self._resolve()
627 if reversed: 627 ↛ 628line 627 didn't jump to line 628 because the condition on line 627 was never true
628 columns = [col.asc() if is_desc else col.desc() for _, col, is_desc in resolved]
629 else:
630 columns = [col.desc() if is_desc else col.asc() for _, col, is_desc in resolved]
631 return select.order_by(None).order_by(*columns)
633 def get_resolved_columns(self) -> list[tuple[str, ColumnElement, bool]]:
634 """Return resolved sort columns as (attr_name, column_element, is_descending) tuples."""
635 return self._resolve()
637 def row_value(self, row: Any, name: str) -> Any:
638 """
639 Extract the sort-key value for ``name`` from a result row.
641 Resolves the accessor through ``to_replace`` for string aliases
642 (e.g. ``{"dag_run_id": "run_id"}``). For column-form mappings
643 (e.g. ``{"run_after": DagRun.run_after}``), resolves through the
644 primary model's attribute so association proxies can still be used
645 for cursor values. Raises ``NotImplementedError`` when the model
646 exposes no such attribute rather than emitting a ``None`` cursor token.
647 """
648 if self.to_replace:
649 replacement = self.to_replace.get(name)
650 if isinstance(replacement, str):
651 return getattr(row, replacement, None)
652 if replacement is not None and not isinstance(replacement, list):
653 # Column-form mapping resolves through the primary model's attribute,
654 # often an association proxy onto the joined entity
655 # (``TaskInstance.run_after`` -> ``dag_run.run_after``). Fail loudly if the
656 # model exposes no such attribute, rather than emitting a ``None`` cursor token.
657 try:
658 return getattr(row, name)
659 except AttributeError:
660 raise NotImplementedError(
661 f"Cursor pagination cannot resolve column-form ``to_replace`` for "
662 f"``{name}``: the primary model exposes no such attribute. Add an "
663 f"association proxy, use a string alias, or sort by a primary-model column."
664 )
665 # List-form replacements are expanded in _resolve() into individual entries
666 # each using the column's own ORM key as attr_name, so ``name`` at this point
667 # is already a concrete model attribute (e.g. ``_rendered_map_index`` or
668 # ``map_index``) — fall through to the getattr below.
669 return getattr(row, name, None)
671 def get_primary_key_column(self) -> Column:
672 """Get the primary key column of the model of SortParam object."""
673 return inspect(self.model).primary_key[0]
675 def get_primary_key_string(self) -> str:
676 """Get the primary key string of the model of SortParam object."""
677 return self.get_primary_key_column().name
679 @classmethod
680 def depends(cls, *args: Any, **kwargs: Any) -> Self:
681 raise NotImplementedError("Use dynamic_depends, depends not implemented.")
683 def dynamic_depends(self, default: str | Sequence[str] | None = None) -> Callable:
684 # Include to_replace keys that are not already in allowed_attrs to avoid
685 # duplicate entries in the spec description.
686 allowed_set = set(self.allowed_attrs)
687 to_replace_attrs = [k for k in self.to_replace if k not in allowed_set] if self.to_replace else []
689 all_attrs = self.allowed_attrs + to_replace_attrs
691 if default is None:
692 default_list = [self.get_primary_key_string()]
693 elif isinstance(default, str):
694 default_list = [default]
695 else:
696 default_list = list(default)
698 _order_by_query = Query(
699 default=default_list,
700 description=f"Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. "
701 f"Supported attributes: `{', '.join(all_attrs) if all_attrs else self.get_primary_key_string()}`",
702 )
704 def inner(order_by: list[str] = _order_by_query) -> SortParam:
705 return SortParam(self.allowed_attrs, self.model, self.to_replace).set_value(order_by)
707 return inner
710class FilterOptionEnum(Enum):
711 """Filter options for FilterParam."""
713 EQUAL = "eq"
714 NOT_EQUAL = "ne"
715 LESS_THAN = "lt"
716 LESS_THAN_EQUAL = "le"
717 GREATER_THAN = "gt"
718 GREATER_THAN_EQUAL = "ge"
719 IN = "in"
720 NOT_IN = "not_in"
721 ANY_EQUAL = "any_eq"
722 ALL_EQUAL = "all_eq"
723 IS_NONE = "is_none"
724 CONTAINS = "contains"
727class FilterParam(BaseParam[T]):
728 """Filter on attribute."""
730 def __init__(
731 self,
732 attribute: InstrumentedAttribute,
733 value: T | None = None,
734 filter_option: FilterOptionEnum = FilterOptionEnum.EQUAL,
735 skip_none: bool = True,
736 ) -> None:
737 super().__init__(value, skip_none)
738 self.attribute: InstrumentedAttribute = attribute
739 self.value: T | None = value
740 self.filter_option: FilterOptionEnum = filter_option
742 def to_orm(self, select: Select) -> Select:
743 if isinstance(self.value, (list, str)) and not self.value and self.skip_none:
744 return select
745 if self.value is None and self.skip_none:
746 return select
748 if isinstance(self.value, list):
749 if self.filter_option == FilterOptionEnum.IN:
750 return select.where(self.attribute.in_(self.value))
751 if self.filter_option == FilterOptionEnum.NOT_IN:
752 return select.where(self.attribute.notin_(self.value))
753 if self.filter_option == FilterOptionEnum.ANY_EQUAL: 753 ↛ 756line 753 didn't jump to line 756 because the condition on line 753 was always true
754 conditions = [self.attribute == val for val in self.value]
755 return select.where(or_(*conditions))
756 if self.filter_option == FilterOptionEnum.ALL_EQUAL:
757 conditions = [self.attribute == val for val in self.value]
758 return select.where(and_(*conditions))
759 raise HTTPException(
760 400, f"Invalid filter option {self.filter_option} for list value {self.value}"
761 )
763 if self.filter_option == FilterOptionEnum.EQUAL:
764 return select.where(self.attribute == self.value)
765 if self.filter_option == FilterOptionEnum.NOT_EQUAL: 765 ↛ 766line 765 didn't jump to line 766 because the condition on line 765 was never true
766 return select.where(self.attribute != self.value)
767 if self.filter_option == FilterOptionEnum.LESS_THAN:
768 return select.where(self.attribute < self.value)
769 if self.filter_option == FilterOptionEnum.LESS_THAN_EQUAL: 769 ↛ 770line 769 didn't jump to line 770 because the condition on line 769 was never true
770 return select.where(self.attribute <= self.value)
771 if self.filter_option == FilterOptionEnum.GREATER_THAN:
772 return select.where(self.attribute > self.value)
773 if self.filter_option == FilterOptionEnum.GREATER_THAN_EQUAL: 773 ↛ 774line 773 didn't jump to line 774 because the condition on line 773 was never true
774 return select.where(self.attribute >= self.value)
775 if self.filter_option == FilterOptionEnum.IS_NONE: 775 ↛ 776line 775 didn't jump to line 776 because the condition on line 775 was never true
776 if self.value is None:
777 return select
778 if self.value is False:
779 return select.where(self.attribute.is_not(None))
780 if self.value is True:
781 return select.where(self.attribute.is_(None))
782 if self.filter_option == FilterOptionEnum.CONTAINS: 782 ↛ 789line 782 didn't jump to line 789 because the condition on line 782 was always true
783 # For JSON/JSONB columns, convert to text before applying LIKE
784 from sqlalchemy import Text, cast
786 if str(self.attribute.type).upper() in ("JSON", "JSONB"): 786 ↛ 788line 786 didn't jump to line 788 because the condition on line 786 was always true
787 return select.where(cast(self.attribute, Text).contains(self.value))
788 return select.where(self.attribute.contains(self.value))
789 raise ValueError(f"Invalid filter option {self.filter_option} for value {self.value}")
791 @classmethod
792 def depends(cls, *args: Any, **kwargs: Any) -> Self:
793 raise NotImplementedError("Use filter_param_factory instead , depends is not implemented.")
796def filter_param_factory(
797 attribute: ColumnElement | InstrumentedAttribute,
798 _type: type,
799 filter_option: FilterOptionEnum = FilterOptionEnum.EQUAL,
800 filter_name: str | None = None,
801 default_value: T | None = None,
802 default_factory: Callable[[], T | None] | None = None,
803 skip_none: bool = True,
804 transform_callable: Callable[[T | None], Any] | None = None,
805 *,
806 description: str | None = None,
807) -> Callable[[T | None], FilterParam[T | None]]:
808 # if filter_name is not provided, use the attribute name as the default
809 filter_name = filter_name or getattr(attribute, "name", str(attribute))
810 # can only set either default_value or default_factory
811 query = (
812 Query(alias=filter_name, default_factory=default_factory, description=description)
813 if default_factory is not None
814 else Query(alias=filter_name, default=default_value, description=description)
815 )
817 def depends_filter(value: T | None = query) -> FilterParam[T | None]:
818 if transform_callable:
819 value = transform_callable(value)
820 # Cast to InstrumentedAttribute for type compatibility
821 attr = cast("InstrumentedAttribute", attribute)
822 return FilterParam(attr, value, filter_option, skip_none)
824 # add type hint to value at runtime
825 depends_filter.__annotations__["value"] = _type
827 return depends_filter
830class _TagFilterModel(BaseModel):
831 """Tag Filter Model with a match mode parameter."""
833 tags: list[str]
834 tags_match_mode: Literal["any", "all"] | None
837class _TagsFilter(BaseParam[_TagFilterModel]):
838 """Filter on tags."""
840 def to_orm(self, select: Select) -> Select:
841 if self.skip_none is False: 841 ↛ 842line 841 didn't jump to line 842 because the condition on line 841 was never true
842 raise ValueError(f"Cannot set 'skip_none' to False on a {type(self)}")
844 if not self.value or not self.value.tags:
845 return select
847 conditions = [DagModel.tags.any(DagTag.name == tag) for tag in self.value.tags]
848 operator = or_ if not self.value.tags_match_mode or self.value.tags_match_mode == "any" else and_
849 return select.where(operator(*conditions))
851 @classmethod
852 def depends(
853 cls,
854 tags: list[str] = Query(default_factory=list),
855 tags_match_mode: Literal["any", "all"] | None = None,
856 ) -> _TagsFilter:
857 return cls().set_value(_TagFilterModel(tags=tags, tags_match_mode=tags_match_mode))
860class _OwnersFilter(BaseParam[list[str]]):
861 """Filter on owners."""
863 def to_orm(self, select: Select) -> Select:
864 if self.skip_none is False: 864 ↛ 865line 864 didn't jump to line 865 because the condition on line 864 was never true
865 raise ValueError(f"Cannot set 'skip_none' to False on a {type(self)}")
867 if not self.value:
868 return select
870 conditions = [
871 DagModel.owners.ilike(f"%{_escape_like_pattern(owner)}%", escape=_LIKE_ESCAPE_CHAR)
872 for owner in self.value
873 ]
874 return select.where(or_(*conditions))
876 @classmethod
877 def depends(cls, owners: list[str] = Query(default_factory=list)) -> _OwnersFilter:
878 return cls().set_value(owners)
881def _safe_parse_datetime(date_to_check: str) -> datetime:
882 """
883 Parse datetime and raise error for invalid dates.
885 :param date_to_check: the string value to be parsed
886 """
887 if not date_to_check:
888 raise ValueError(f"{date_to_check} cannot be None.")
889 return _safe_parse_datetime_optional(date_to_check)
892@overload
893def _safe_parse_datetime_optional(date_to_check: str) -> datetime: ... 893 ↛ exitline 893 didn't return from function '_safe_parse_datetime_optional' because
896@overload
897def _safe_parse_datetime_optional(date_to_check: None) -> None: ... 897 ↛ exitline 897 didn't return from function '_safe_parse_datetime_optional' because
900def _safe_parse_datetime_optional(date_to_check: str | None) -> datetime | None:
901 """
902 Parse datetime and raise error for invalid dates.
904 Allow None values.
906 :param date_to_check: the string value to be parsed
907 """
908 if date_to_check is None: 908 ↛ 909line 908 didn't jump to line 909 because the condition on line 908 was never true
909 return None
910 try:
911 return timezone.parse(date_to_check, strict=True)
912 except (TypeError, ParserError):
913 raise HTTPException(
914 400, f"Invalid datetime: {date_to_check!r}. Please check the date parameter have this value."
915 )
918class _DagIdAssetReferenceFilter(BaseParam[list[str]]):
919 """Search on dag_id."""
921 def __init__(self, skip_none: bool = True) -> None:
922 super().__init__(skip_none=skip_none)
924 @classmethod
925 def depends(cls, dag_ids: list[str] = Query(None)) -> _DagIdAssetReferenceFilter:
926 # needed to handle cases where dag_ids=a1,b1
927 if dag_ids and len(dag_ids) == 1 and "," in dag_ids[0]: 927 ↛ 928line 927 didn't jump to line 928 because the condition on line 927 was never true
928 dag_ids = dag_ids[0].split(",")
929 return cls().set_value(dag_ids)
931 def to_orm(self, select: Select) -> Select:
932 if self.value is None and self.skip_none:
933 return select
935 # At this point, self.value is either a list[str] or None -> coerce falsy None to an empty list
936 dag_ids = self.value or []
937 return select.where(
938 (AssetModel.scheduled_dags.any(DagScheduleAssetReference.dag_id.in_(dag_ids)))
939 | (AssetModel.producing_tasks.any(TaskOutletAssetReference.dag_id.in_(dag_ids)))
940 | (AssetModel.consuming_tasks.any(TaskInletAssetReference.dag_id.in_(dag_ids)))
941 )
944class Range(BaseModel, Generic[T]):
945 """Range with a lower and upper bound."""
947 lower_bound_gte: T | None
948 lower_bound_gt: T | None
949 upper_bound_lte: T | None
950 upper_bound_lt: T | None
953class RangeFilter(BaseParam[Range]):
954 """Filter on range in between the lower and upper bound."""
956 def __init__(self, value: Range | None, attribute: InstrumentedAttribute) -> None:
957 super().__init__(value)
958 self.attribute: InstrumentedAttribute = attribute
960 def to_orm(self, select: Select) -> Select:
961 if self.skip_none is False: 961 ↛ 962line 961 didn't jump to line 962 because the condition on line 961 was never true
962 raise ValueError(f"Cannot set 'skip_none' to False on a {type(self)}")
964 if self.value is None: 964 ↛ 965line 964 didn't jump to line 965 because the condition on line 964 was never true
965 return select
967 if self.value.lower_bound_gte:
968 select = select.where(self.attribute >= self.value.lower_bound_gte)
969 if self.value.lower_bound_gt:
970 select = select.where(self.attribute > self.value.lower_bound_gt)
971 if self.value.upper_bound_lte:
972 select = select.where(self.attribute <= self.value.upper_bound_lte)
973 if self.value.upper_bound_lt:
974 select = select.where(self.attribute < self.value.upper_bound_lt)
976 return select
978 @classmethod
979 def depends(cls, *args: Any, **kwargs: Any) -> Self:
980 raise NotImplementedError("Use the `range_filter_factory` function to create the dependency")
982 def is_active(self) -> bool:
983 """Check if the range filter has any active bounds."""
984 return self.value is not None and (
985 self.value.lower_bound_gte is not None
986 or self.value.lower_bound_gt is not None
987 or self.value.upper_bound_lte is not None
988 or self.value.upper_bound_lt is not None
989 )
992class NullableDatetimeRangeFilter(RangeFilter):
993 """
994 RangeFilter for nullable datetime columns (``start_date``, ``end_date``), rewritten for index use.
996 ``COALESCE(column, now())`` wraps the column in a function call that prevents PostgreSQL from
997 using btree indexes, forcing sequential scans on large tables. This class emits equivalent
998 ``OR`` predicates so each branch can be satisfied by an independent index scan.
1000 NULL semantics: ``start_date=NULL`` means the task has not started yet; ``end_date=NULL`` means
1001 the task is still running. For lower bounds the NULL branch passes unconditionally — a not-yet-
1002 started/ended task will eventually satisfy any past lower bound. For upper bounds the NULL branch
1003 is ``col IS NULL AND now() <= x``, preserving the COALESCE(col, now()) semantics without the
1004 function-wrap index penalty.
1005 """
1007 def to_orm(self, select: Select) -> Select:
1008 if self.skip_none is False: 1008 ↛ 1009line 1008 didn't jump to line 1009 because the condition on line 1008 was never true
1009 raise ValueError(f"Cannot set 'skip_none' to False on a {type(self)}")
1011 if self.value is None: 1011 ↛ 1012line 1011 didn't jump to line 1012 because the condition on line 1011 was never true
1012 return select
1014 if self.value.lower_bound_gte:
1015 x = self.value.lower_bound_gte
1016 select = select.where(or_(self.attribute >= x, self.attribute.is_(None)))
1017 if self.value.lower_bound_gt:
1018 x = self.value.lower_bound_gt
1019 select = select.where(or_(self.attribute > x, self.attribute.is_(None)))
1020 if self.value.upper_bound_lte:
1021 x = self.value.upper_bound_lte
1022 select = select.where(or_(self.attribute <= x, and_(self.attribute.is_(None), func.now() <= x)))
1023 if self.value.upper_bound_lt:
1024 x = self.value.upper_bound_lt
1025 select = select.where(or_(self.attribute < x, and_(self.attribute.is_(None), func.now() < x)))
1027 return select
1030def datetime_range_filter_factory(
1031 filter_name: str, model: Base, attribute_name: str | None = None
1032) -> Callable[[datetime | None, datetime | None, datetime | None, datetime | None], RangeFilter]:
1033 def depends_datetime(
1034 lower_bound_gte: datetime | None = Query(alias=f"{filter_name}_gte", default=None),
1035 lower_bound_gt: datetime | None = Query(alias=f"{filter_name}_gt", default=None),
1036 upper_bound_lte: datetime | None = Query(alias=f"{filter_name}_lte", default=None),
1037 upper_bound_lt: datetime | None = Query(alias=f"{filter_name}_lt", default=None),
1038 ) -> RangeFilter:
1039 attr = getattr(model, attribute_name or filter_name)
1040 range_val = Range(
1041 lower_bound_gte=lower_bound_gte,
1042 lower_bound_gt=lower_bound_gt,
1043 upper_bound_lte=upper_bound_lte,
1044 upper_bound_lt=upper_bound_lt,
1045 )
1046 if filter_name in ("start_date", "end_date"):
1047 return NullableDatetimeRangeFilter(range_val, attr)
1048 return RangeFilter(range_val, attr)
1050 return depends_datetime
1053def float_range_filter_factory(
1054 filter_name: str, model: Base
1055) -> Callable[[float | None, float | None, float | None, float | None], RangeFilter]:
1056 def depends_float(
1057 lower_bound_gte: float | None = Query(alias=f"{filter_name}_gte", default=None),
1058 lower_bound_gt: float | None = Query(alias=f"{filter_name}_gt", default=None),
1059 upper_bound_lte: float | None = Query(alias=f"{filter_name}_lte", default=None),
1060 upper_bound_lt: float | None = Query(alias=f"{filter_name}_lt", default=None),
1061 ) -> RangeFilter:
1062 return RangeFilter(
1063 Range(
1064 lower_bound_gte=lower_bound_gte,
1065 lower_bound_gt=lower_bound_gt,
1066 upper_bound_lte=upper_bound_lte,
1067 upper_bound_lt=upper_bound_lt,
1068 ),
1069 getattr(model, filter_name),
1070 )
1072 return depends_float
1075# Common Safe DateTime
1076DateTimeQuery = Annotated[str, AfterValidator(_safe_parse_datetime)]
1077OptionalDateTimeQuery = Annotated[str | None, AfterValidator(_safe_parse_datetime_optional)]
1079# Dag
1080QueryLimit = Annotated[LimitFilter, Depends(LimitFilter.depends)]
1081QueryOffset = Annotated[OffsetFilter, Depends(OffsetFilter.depends)]
1082QueryPausedFilter = Annotated[
1083 FilterParam[bool | None],
1084 Depends(filter_param_factory(DagModel.is_paused, bool | None, filter_name="paused")),
1085]
1086QueryHasImportErrorsFilter = Annotated[
1087 FilterParam[bool | None],
1088 Depends(
1089 filter_param_factory(
1090 DagModel.has_import_errors,
1091 bool | None,
1092 filter_name="has_import_errors",
1093 description="Filter Dags by having import errors. Only Dags that have been successfully loaded before will be returned.",
1094 )
1095 ),
1096]
1097QueryFavoriteFilter = Annotated[_FavoriteFilter, Depends(_FavoriteFilter.depends)]
1098QueryExcludeStaleFilter = Annotated[_ExcludeStaleFilter, Depends(_ExcludeStaleFilter.depends)]
1099QueryDagIdPatternSearch = Annotated[
1100 _SearchParam, Depends(search_param_factory(DagModel.dag_id, "dag_id_pattern"))
1101]
1102QueryDagIdPrefixPatternSearch = Annotated[
1103 _PrefixSearchParam, Depends(prefix_search_param_factory(DagModel.dag_id, "dag_id_prefix_pattern"))
1104]
1105QueryDagDisplayNamePatternSearch = Annotated[
1106 _SearchParam, Depends(search_param_factory(DagModel.dag_display_name, "dag_display_name_pattern"))
1107]
1108QueryDagDisplayNamePrefixPatternSearch = Annotated[
1109 _PrefixSearchParam,
1110 Depends(prefix_search_param_factory(DagModel.dag_display_name, "dag_display_name_prefix_pattern")),
1111]
1112QueryBundleNameFilter = Annotated[
1113 FilterParam[str | None],
1114 Depends(filter_param_factory(DagModel.bundle_name, str | None, filter_name="bundle_name")),
1115]
1116QueryBundleVersionFilter = Annotated[
1117 FilterParam[str | None],
1118 Depends(filter_param_factory(DagModel.bundle_version, str | None, filter_name="bundle_version")),
1119]
1120QueryDagIdPatternSearchWithNone = Annotated[
1121 _SearchParam, Depends(search_param_factory(DagModel.dag_id, "dag_id_pattern", False))
1122]
1123QueryDagIdPrefixPatternSearchWithNone = Annotated[
1124 _PrefixSearchParam,
1125 Depends(prefix_search_param_factory(DagModel.dag_id, "dag_id_prefix_pattern", False)),
1126]
1127QueryTagsFilter = Annotated[_TagsFilter, Depends(_TagsFilter.depends)]
1128QueryOwnersFilter = Annotated[_OwnersFilter, Depends(_OwnersFilter.depends)]
1131class _HasAssetScheduleFilter(BaseParam[bool]):
1132 """Filter Dags that have asset-based scheduling."""
1134 def to_orm(self, select: Select) -> Select:
1135 if self.value is None and self.skip_none:
1136 return select
1138 asset_ref_subquery = sql_select(DagScheduleAssetReference.dag_id).distinct()
1140 if self.value:
1141 # Filter Dags that have asset-based scheduling
1142 return select.where(DagModel.dag_id.in_(asset_ref_subquery))
1144 # Filter Dags that do NOT have asset-based scheduling
1145 return select.where(DagModel.dag_id.notin_(asset_ref_subquery))
1147 @classmethod
1148 def depends(
1149 cls,
1150 has_asset_schedule: bool | None = Query(None, description="Filter Dags with asset-based scheduling"),
1151 ) -> _HasAssetScheduleFilter:
1152 return cls().set_value(has_asset_schedule)
1155class _AssetDependencyFilter(BaseParam[str]):
1156 """Filter Dags by specific asset dependencies."""
1158 def to_orm(self, select: Select) -> Select:
1159 if self.value is None:
1160 return select
1162 escaped = _escape_like_pattern(self.value)
1163 asset_dag_subquery = (
1164 sql_select(DagScheduleAssetReference.dag_id)
1165 .join(AssetModel, DagScheduleAssetReference.asset_id == AssetModel.id)
1166 .where(
1167 or_(
1168 AssetModel.name.ilike(f"%{escaped}%", escape=_LIKE_ESCAPE_CHAR),
1169 AssetModel.uri.ilike(f"%{escaped}%", escape=_LIKE_ESCAPE_CHAR),
1170 )
1171 )
1172 .distinct()
1173 )
1175 return select.where(DagModel.dag_id.in_(asset_dag_subquery))
1177 @classmethod
1178 def depends(
1179 cls,
1180 asset_dependency: str | None = Query(
1181 None, description="Filter Dags by asset dependency (name or URI)"
1182 ),
1183 ) -> _AssetDependencyFilter:
1184 return cls().set_value(asset_dependency)
1187QueryHasAssetScheduleFilter = Annotated[_HasAssetScheduleFilter, Depends(_HasAssetScheduleFilter.depends)]
1188QueryAssetDependencyFilter = Annotated[_AssetDependencyFilter, Depends(_AssetDependencyFilter.depends)]
1191class _ConsumingAssetFilter(BaseParam[str | None]):
1192 """Filter Dag runs by consuming asset (name or URI)."""
1194 def to_orm(self, select: Select) -> Select:
1195 if not self.value:
1196 return select
1198 escaped = _escape_like_pattern(self.value)
1199 event_subquery = (
1200 sql_select(AssetEvent.id)
1201 .join(AssetModel, AssetEvent.asset_id == AssetModel.id)
1202 .where(
1203 or_(
1204 AssetModel.name.ilike(f"%{escaped}%", escape=_LIKE_ESCAPE_CHAR),
1205 AssetModel.uri.ilike(f"%{escaped}%", escape=_LIKE_ESCAPE_CHAR),
1206 )
1207 )
1208 .distinct()
1209 )
1211 dagrun_subquery = (
1212 sql_select(association_table.c.dag_run_id)
1213 .where(association_table.c.event_id.in_(event_subquery))
1214 .distinct()
1215 )
1217 return select.where(DagRun.id.in_(dagrun_subquery))
1219 @classmethod
1220 def depends(
1221 cls,
1222 consuming_asset_pattern: str | None = Query(
1223 None,
1224 description=(
1225 "Case-insensitive substring match against the consuming asset name or URI. "
1226 "Unlike the wildcard `*_pattern` parameters, `%` and `_` are matched literally, "
1227 "`|` is not an OR separator, and `~` does not match everything."
1228 ),
1229 ),
1230 ) -> _ConsumingAssetFilter:
1231 return cls().set_value(consuming_asset_pattern)
1234QueryConsumingAssetPatternSearch = Annotated[_ConsumingAssetFilter, Depends(_ConsumingAssetFilter.depends)]
1237class _PendingActionsFilter(BaseParam[bool]):
1238 """Filter Dags by having pending HITL actions (more than 1)."""
1240 def to_orm(self, select: Select) -> Select:
1241 if self.value is None and self.skip_none:
1242 return select
1244 from airflow.models.hitl import HITLDetail
1245 from airflow.models.taskinstance import TaskInstance
1247 # Join with HITLDetail and TaskInstance to find Dags
1248 pending_actions_count_subquery = (
1249 sql_select(func.count(HITLDetail.ti_id))
1250 .join(TaskInstance, HITLDetail.ti_id == TaskInstance.id)
1251 .where(
1252 HITLDetail.responded_at.is_(None),
1253 TaskInstance.state.in_((TaskInstanceState.DEFERRED, TaskInstanceState.AWAITING_INPUT)),
1254 )
1255 .where(TaskInstance.dag_id == DagModel.dag_id)
1256 .scalar_subquery()
1257 )
1259 if self.value is True:
1260 # Filter to show only Dags with pending actions
1261 where_clause = pending_actions_count_subquery >= 1
1262 else:
1263 # Filter to show only Dags without pending actions
1264 where_clause = pending_actions_count_subquery == 0
1266 return select.where(where_clause)
1268 @classmethod
1269 def depends(cls, has_pending_actions: bool | None = Query(None)) -> _PendingActionsFilter:
1270 return cls().set_value(has_pending_actions)
1273QueryPendingActionsFilter = Annotated[_PendingActionsFilter, Depends(_PendingActionsFilter.depends)]
1276class _AnyDagRunStateFilter(BaseParam[DagRunState | None]):
1277 """Filter Dags that have any DagRun in the given state, not only the latest one."""
1279 def to_orm(self, select: Select) -> Select:
1280 if self.value is None and self.skip_none:
1281 return select
1283 # Alias DagRun so this EXISTS subquery cannot auto-correlate to a DagRun the outer query
1284 # may already reference (e.g. the last_dag_run_state filter), which would strip the
1285 # subquery's FROM and raise. EXISTS resolves each Dag via the (dag_id, state) index.
1286 any_run = aliased(DagRun)
1287 has_run_in_state = (
1288 sql_select(any_run.dag_id)
1289 .where(any_run.dag_id == DagModel.dag_id, any_run.state == self.value)
1290 .exists()
1291 )
1292 return select.where(has_run_in_state)
1294 @classmethod
1295 def depends(
1296 cls,
1297 dag_run_state: DagRunState | None = Query(
1298 None,
1299 description="Filter Dags that have any DagRun in the given state.",
1300 ),
1301 ) -> _AnyDagRunStateFilter:
1302 return cls().set_value(dag_run_state)
1305# DagRun
1306QueryLastDagRunStateFilter = Annotated[
1307 FilterParam[DagRunState | None],
1308 Depends(filter_param_factory(DagRun.state, DagRunState | None, filter_name="last_dag_run_state")),
1309]
1310QueryAnyDagRunStateFilter = Annotated[_AnyDagRunStateFilter, Depends(_AnyDagRunStateFilter.depends)]
1313def _transform_dag_run_states(states: Iterable[str] | None) -> list[DagRunState | None] | None:
1314 try:
1315 if not states:
1316 return None
1317 return [None if s in ("none", None) else DagRunState(s) for s in states]
1318 except ValueError:
1319 raise HTTPException(
1320 status_code=HTTP_422_UNPROCESSABLE_CONTENT,
1321 detail=f"Invalid value for state. Valid values are {', '.join(DagRunState)}",
1322 )
1325QueryDagRunStateFilter = Annotated[
1326 FilterParam[list[str]],
1327 Depends(
1328 filter_param_factory(
1329 DagRun.state,
1330 list[str],
1331 FilterOptionEnum.ANY_EQUAL,
1332 default_factory=list,
1333 transform_callable=_transform_dag_run_states,
1334 )
1335 ),
1336]
1339def _transform_dag_run_types(types: list[str] | None) -> list[DagRunType | None] | None:
1340 try:
1341 if not types:
1342 return None
1343 return [None if run_type in ("none", None) else DagRunType(run_type) for run_type in types]
1344 except ValueError:
1345 raise HTTPException(
1346 status_code=HTTP_422_UNPROCESSABLE_CONTENT,
1347 detail=f"Invalid value for run type. Valid values are {', '.join(DagRunType)}",
1348 )
1351QueryDagRunRunTypesFilter = Annotated[
1352 FilterParam[list[str]],
1353 Depends(
1354 filter_param_factory(
1355 attribute=DagRun.run_type,
1356 _type=list[str],
1357 filter_option=FilterOptionEnum.ANY_EQUAL,
1358 default_factory=list,
1359 transform_callable=_transform_dag_run_types,
1360 )
1361 ),
1362]
1364QueryDagRunTriggeringUserSearch = Annotated[
1365 _SearchParam, Depends(search_param_factory(DagRun.triggering_user_name, "triggering_user"))
1366]
1367QueryDagRunTriggeringUserPrefixSearch = Annotated[
1368 _PrefixSearchParam,
1369 Depends(prefix_search_param_factory(DagRun.triggering_user_name, "triggering_user_prefix")),
1370]
1371QueryDagRunPartitionKeySearch = Annotated[
1372 _SearchParam,
1373 Depends(search_param_factory(DagRun.partition_key, "partition_key_pattern", pipe_as_or=False)),
1374]
1375QueryDagRunPartitionKeyPrefixSearch = Annotated[
1376 _PrefixSearchParam,
1377 Depends(
1378 prefix_search_param_factory(DagRun.partition_key, "partition_key_prefix_pattern", pipe_as_or=False)
1379 ),
1380]
1382# DagTags
1383QueryDagTagPatternSearch = Annotated[
1384 _SearchParam, Depends(search_param_factory(DagTag.name, "tag_name_pattern"))
1385]
1386QueryDagTagPrefixPatternSearch = Annotated[
1387 _PrefixSearchParam, Depends(prefix_search_param_factory(DagTag.name, "tag_name_prefix_pattern"))
1388]
1391# TI
1392def _transform_ti_states(states: list[str] | None) -> list[TaskInstanceState | None] | None:
1393 """Transform a list of state strings into a list of TaskInstanceState enums handling special 'None' cases."""
1394 if not states:
1395 return None
1397 try:
1398 return [None if s in ("no_status", "none", None) else TaskInstanceState(s) for s in states]
1399 except ValueError:
1400 raise HTTPException(
1401 status_code=HTTP_422_UNPROCESSABLE_CONTENT,
1402 detail=f"Invalid value for state. Valid values are {', '.join(TaskInstanceState)}",
1403 )
1406QueryTIStateFilter = Annotated[
1407 FilterParam[list[str]],
1408 Depends(
1409 filter_param_factory(
1410 TaskInstance.state,
1411 list[str],
1412 FilterOptionEnum.ANY_EQUAL,
1413 default_factory=list,
1414 transform_callable=_transform_ti_states,
1415 )
1416 ),
1417]
1418QueryTIPoolFilter = Annotated[
1419 FilterParam[list[str]],
1420 Depends(
1421 filter_param_factory(TaskInstance.pool, list[str], FilterOptionEnum.ANY_EQUAL, default_factory=list)
1422 ),
1423]
1424QueryTIQueueFilter = Annotated[
1425 FilterParam[list[str]],
1426 Depends(
1427 filter_param_factory(TaskInstance.queue, list[str], FilterOptionEnum.ANY_EQUAL, default_factory=list)
1428 ),
1429]
1430QueryTIPoolNamePatternSearch = Annotated[
1431 _SearchParam,
1432 Depends(search_param_factory(TaskInstance.pool, "pool_name_pattern")),
1433]
1434QueryTIPoolNamePrefixPatternSearch = Annotated[
1435 _PrefixSearchParam,
1436 Depends(prefix_search_param_factory(TaskInstance.pool, "pool_name_prefix_pattern")),
1437]
1439QueryTIQueueNamePatternSearch = Annotated[
1440 _SearchParam,
1441 Depends(search_param_factory(TaskInstance.queue, "queue_name_pattern")),
1442]
1443QueryTIQueueNamePrefixPatternSearch = Annotated[
1444 _PrefixSearchParam,
1445 Depends(prefix_search_param_factory(TaskInstance.queue, "queue_name_prefix_pattern")),
1446]
1447QueryTIExecutorFilter = Annotated[
1448 FilterParam[list[str]],
1449 Depends(
1450 filter_param_factory(
1451 TaskInstance.executor, list[str], FilterOptionEnum.ANY_EQUAL, default_factory=list
1452 )
1453 ),
1454]
1455QueryTITaskDisplayNamePatternSearch = Annotated[
1456 _SearchParam,
1457 Depends(search_param_factory(TaskInstance.task_display_name, "task_display_name_pattern")),
1458]
1459QueryTITaskDisplayNamePrefixPatternSearch = Annotated[
1460 _TaskDisplayNamePrefixPatternParam, Depends(_TaskDisplayNamePrefixPatternParam.depends)
1461]
1462QueryTITaskGroupFilter = Annotated[
1463 QueryTaskInstanceTaskGroupFilter, Depends(QueryTaskInstanceTaskGroupFilter.depends)
1464]
1465QueryTIDagVersionFilter = Annotated[
1466 FilterParam[list[int]],
1467 Depends(
1468 filter_param_factory(
1469 DagVersion.version_number,
1470 list[int],
1471 FilterOptionEnum.ANY_EQUAL,
1472 default_factory=list,
1473 )
1474 ),
1475]
1476QueryDagRunVersionFilter = Annotated[
1477 FilterParam[list[int]],
1478 Depends(
1479 filter_param_factory(
1480 DagVersion.version_number,
1481 list[int],
1482 FilterOptionEnum.ANY_EQUAL,
1483 default_factory=list,
1484 filter_name="dag_version",
1485 )
1486 ),
1487]
1488QueryTITryNumberFilter = Annotated[
1489 FilterParam[list[int]],
1490 Depends(
1491 filter_param_factory(
1492 TaskInstance.try_number, list[int], FilterOptionEnum.ANY_EQUAL, default_factory=list
1493 )
1494 ),
1495]
1497QueryTIOperatorFilter = Annotated[
1498 FilterParam[list[str]],
1499 Depends(
1500 filter_param_factory(
1501 TaskInstance.operator, list[str], FilterOptionEnum.ANY_EQUAL, default_factory=list
1502 )
1503 ),
1504]
1505QueryTIOperatorNamePatternSearch = Annotated[
1506 _SearchParam,
1507 Depends(
1508 search_param_factory(
1509 TaskInstance.custom_operator_name,
1510 "operator_name_pattern",
1511 )
1512 ),
1513]
1514QueryTIOperatorNamePrefixPatternSearch = Annotated[
1515 _PrefixSearchParam,
1516 Depends(
1517 prefix_search_param_factory(
1518 TaskInstance.custom_operator_name,
1519 "operator_name_prefix_pattern",
1520 )
1521 ),
1522]
1524QueryTIMapIndexFilter = Annotated[
1525 FilterParam[list[int]],
1526 Depends(
1527 filter_param_factory(
1528 TaskInstance.map_index, list[int], FilterOptionEnum.ANY_EQUAL, default_factory=list
1529 )
1530 ),
1531]
1532# On MySQL the CASE expression that backs rendered_map_index mixes a stored
1533# VARCHAR column (utf8mb4_bin, IMPLICIT) with CAST(map_index AS CHAR)
1534# (utf8mb4_0900_ai_ci, IMPLICIT), which gives the whole expression NONE
1535# coercibility. Comparing it against a bound parameter then fails with
1536# "Illegal mix of collations". _MySQLCollate wraps the expression so that
1537# on MySQL an explicit COLLATE clause is emitted (giving EXPLICIT coercibility);
1538# on PostgreSQL and SQLite the wrapper is transparent.
1539_rendered_map_index_collated = _MySQLCollate(
1540 cast("ColumnElement[Any]", TaskInstance.rendered_map_index), "utf8mb4_0900_ai_ci"
1541)
1543QueryTIRenderedMapIndexPatternSearch = Annotated[
1544 _SearchParam,
1545 Depends(
1546 search_param_factory(
1547 _rendered_map_index_collated,
1548 "rendered_map_index_pattern",
1549 )
1550 ),
1551]
1552QueryTIRenderedMapIndexPrefixPatternSearch = Annotated[
1553 _PrefixSearchParam,
1554 Depends(
1555 prefix_search_param_factory(
1556 _rendered_map_index_collated,
1557 "rendered_map_index_prefix_pattern",
1558 )
1559 ),
1560]
1562# XCom
1563QueryXComKeyPatternSearch = Annotated[
1564 _SearchParam, Depends(search_param_factory(XComModel.key, "xcom_key_pattern"))
1565]
1566QueryXComKeyPrefixPatternSearch = Annotated[
1567 _PrefixSearchParam, Depends(prefix_search_param_factory(XComModel.key, "xcom_key_prefix_pattern"))
1568]
1570QueryXComDagDisplayNamePatternSearch = Annotated[
1571 _SearchParam, Depends(search_param_factory(DagModel.dag_display_name, "dag_display_name_pattern"))
1572]
1573QueryXComDagDisplayNamePrefixPatternSearch = Annotated[
1574 _PrefixSearchParam,
1575 Depends(prefix_search_param_factory(DagModel.dag_display_name, "dag_display_name_prefix_pattern")),
1576]
1577QueryXComRunIdPatternSearch = Annotated[
1578 _SearchParam, Depends(search_param_factory(XComModel.run_id, "run_id_pattern"))
1579]
1580QueryXComRunIdPrefixPatternSearch = Annotated[
1581 _PrefixSearchParam, Depends(prefix_search_param_factory(XComModel.run_id, "run_id_prefix_pattern"))
1582]
1583QueryXComTaskIdPatternSearch = Annotated[
1584 _SearchParam, Depends(search_param_factory(XComModel.task_id, "task_id_pattern"))
1585]
1586QueryXComTaskIdPrefixPatternSearch = Annotated[
1587 _PrefixSearchParam, Depends(prefix_search_param_factory(XComModel.task_id, "task_id_prefix_pattern"))
1588]
1590# Assets
1591QueryAssetNamePatternSearch = Annotated[
1592 _SearchParam, Depends(search_param_factory(AssetModel.name, "name_pattern"))
1593]
1594QueryAssetNamePrefixPatternSearch = Annotated[
1595 _PrefixSearchParam, Depends(prefix_search_param_factory(AssetModel.name, "name_prefix_pattern"))
1596]
1597QueryUriPatternSearch = Annotated[_SearchParam, Depends(search_param_factory(AssetModel.uri, "uri_pattern"))]
1598QueryUriPrefixPatternSearch = Annotated[
1599 _PrefixSearchParam, Depends(prefix_search_param_factory(AssetModel.uri, "uri_prefix_pattern"))
1600]
1601QueryAssetAliasNamePatternSearch = Annotated[
1602 _SearchParam, Depends(search_param_factory(AssetAliasModel.name, "name_pattern"))
1603]
1604QueryAssetAliasNamePrefixPatternSearch = Annotated[
1605 _PrefixSearchParam, Depends(prefix_search_param_factory(AssetAliasModel.name, "name_prefix_pattern"))
1606]
1607QueryAssetDagIdPatternSearch = Annotated[
1608 _DagIdAssetReferenceFilter, Depends(_DagIdAssetReferenceFilter.depends)
1609]
1610QueryPartitionedDagRunHasCreatedDagRunIdFilter = Annotated[
1611 FilterParam[bool | None],
1612 Depends(
1613 filter_param_factory(
1614 AssetPartitionDagRun.created_dag_run_id,
1615 bool | None,
1616 FilterOptionEnum.IS_NONE,
1617 filter_name="has_created_dag_run_id",
1618 transform_callable=lambda v: not v if v is not None else None,
1619 )
1620 ),
1621]
1622QueryPartitionedDagRunDagIdFilter = Annotated[
1623 FilterParam[str | None],
1624 Depends(
1625 filter_param_factory(
1626 AssetPartitionDagRun.target_dag_id,
1627 str | None,
1628 filter_name="dag_id",
1629 )
1630 ),
1631]
1633# Variables
1634QueryVariableKeyPatternSearch = Annotated[
1635 _SearchParam, Depends(search_param_factory(Variable.key, "variable_key_pattern"))
1636]
1637QueryVariableKeyPrefixPatternSearch = Annotated[
1638 _PrefixSearchParam,
1639 Depends(prefix_search_param_factory(Variable.key, "variable_key_prefix_pattern")),
1640]
1642# Pools
1643QueryPoolNamePatternSearch = Annotated[
1644 _SearchParam, Depends(search_param_factory(Pool.pool, "pool_name_pattern"))
1645]
1646QueryPoolNamePrefixPatternSearch = Annotated[
1647 _PrefixSearchParam, Depends(prefix_search_param_factory(Pool.pool, "pool_name_prefix_pattern"))
1648]
1651# UI Shared
1652def _optional_boolean(value: bool | None) -> bool | None:
1653 return value if value is not None else False
1656QueryIncludeUpstream = Annotated[bool, AfterValidator(_optional_boolean)]
1657QueryIncludeDownstream = Annotated[bool, AfterValidator(_optional_boolean)]
1659state_priority: list[None | TaskInstanceState] = [
1660 TaskInstanceState.FAILED,
1661 TaskInstanceState.UPSTREAM_FAILED,
1662 TaskInstanceState.UP_FOR_RETRY,
1663 TaskInstanceState.UP_FOR_RESCHEDULE,
1664 TaskInstanceState.RUNNING,
1665 TaskInstanceState.RESTARTING,
1666 TaskInstanceState.DEFERRED,
1667 TaskInstanceState.AWAITING_INPUT,
1668 TaskInstanceState.QUEUED,
1669 TaskInstanceState.SCHEDULED,
1670 None,
1671 TaskInstanceState.SUCCESS,
1672 TaskInstanceState.SKIPPED,
1673 TaskInstanceState.REMOVED,
1674]
1676# Connections
1677QueryConnectionIdPatternSearch = Annotated[
1678 _SearchParam, Depends(search_param_factory(Connection.conn_id, "connection_id_pattern"))
1679]
1680QueryConnectionIdPrefixPatternSearch = Annotated[
1681 _PrefixSearchParam,
1682 Depends(prefix_search_param_factory(Connection.conn_id, "connection_id_prefix_pattern")),
1683]
1685# Human in the loop
1686QueryHITLDetailDagIdPatternSearch = Annotated[
1687 _SearchParam,
1688 Depends(
1689 search_param_factory(
1690 TaskInstance.dag_id,
1691 "dag_id_pattern",
1692 )
1693 ),
1694]
1695QueryHITLDetailDagIdPrefixPatternSearch = Annotated[
1696 _PrefixSearchParam,
1697 Depends(
1698 prefix_search_param_factory(
1699 TaskInstance.dag_id,
1700 "dag_id_prefix_pattern",
1701 )
1702 ),
1703]
1704QueryHITLDetailTaskIdPatternSearch = Annotated[
1705 _SearchParam,
1706 Depends(
1707 search_param_factory(
1708 TaskInstance.task_id,
1709 "task_id_pattern",
1710 )
1711 ),
1712]
1713QueryHITLDetailTaskIdPrefixPatternSearch = Annotated[
1714 _PrefixSearchParam,
1715 Depends(
1716 prefix_search_param_factory(
1717 TaskInstance.task_id,
1718 "task_id_prefix_pattern",
1719 )
1720 ),
1721]
1722QueryHITLDetailTaskIdFilter = Annotated[
1723 FilterParam[str | None],
1724 Depends(
1725 filter_param_factory(
1726 TaskInstance.task_id,
1727 str | None,
1728 filter_name="task_id",
1729 )
1730 ),
1731]
1732QueryHITLDetailMapIndexFilter = Annotated[
1733 FilterParam[int | None],
1734 Depends(
1735 filter_param_factory(
1736 TaskInstance.map_index,
1737 int | None,
1738 filter_name="map_index",
1739 )
1740 ),
1741]
1742QueryHITLDetailSubjectSearch = Annotated[
1743 _SearchParam,
1744 Depends(
1745 search_param_factory(
1746 HITLDetail.subject,
1747 "subject_search",
1748 )
1749 ),
1750]
1751QueryHITLDetailBodySearch = Annotated[
1752 _SearchParam,
1753 Depends(
1754 search_param_factory(
1755 HITLDetail.body,
1756 "body_search",
1757 )
1758 ),
1759]
1760QueryHITLDetailResponseReceivedFilter = Annotated[
1761 FilterParam[bool | None],
1762 Depends(
1763 filter_param_factory(
1764 HITLDetail.response_received,
1765 bool | None,
1766 filter_name="response_received",
1767 )
1768 ),
1769]
1770QueryHITLDetailRespondedUserIdFilter = Annotated[
1771 FilterParam[list[str]],
1772 Depends(
1773 filter_param_factory(
1774 HITLDetail.responded_by_user_id,
1775 list[str],
1776 FilterOptionEnum.ANY_EQUAL,
1777 default_factory=list,
1778 filter_name="responded_by_user_id",
1779 )
1780 ),
1781]
1782QueryHITLDetailRespondedUserNameFilter = Annotated[
1783 FilterParam[list[str]],
1784 Depends(
1785 filter_param_factory(
1786 HITLDetail.responded_by_user_name,
1787 list[str],
1788 FilterOptionEnum.ANY_EQUAL,
1789 default_factory=list,
1790 filter_name="responded_by_user_name",
1791 )
1792 ),
1793]
1795# Parse Import Errors
1796QueryParseImportErrorFilenamePatternSearch = Annotated[
1797 _SearchParam, Depends(search_param_factory(ParseImportError.filename, "filename_pattern"))
1798]
1799QueryParseImportErrorFilenamePrefixPatternSearch = Annotated[
1800 _PrefixSearchParam,
1801 Depends(prefix_search_param_factory(ParseImportError.filename, "filename_prefix_pattern")),
1802]
1803QueryParseImportErrorFilenameFilter = Annotated[
1804 FilterParam,
1805 Depends(
1806 filter_param_factory(
1807 ParseImportError.filename,
1808 str | None,
1809 filter_name="filename",
1810 description="Exact filename match. Returns only the import error for this specific file path.",
1811 )
1812 ),
1813]
1814QueryParseImportErrorBundleNameFilter = Annotated[
1815 FilterParam,
1816 Depends(
1817 filter_param_factory(
1818 ParseImportError.bundle_name,
1819 str | None,
1820 filter_name="bundle_name",
1821 description="Exact bundle name match. Returns only import errors from this specific bundle.",
1822 )
1823 ),
1824]