Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/core_api/routes/ui/partitioned_dag_runs.py: 19%
136 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.
17from __future__ import annotations
19from typing import TYPE_CHECKING, NamedTuple, TypeAlias, cast
21import structlog
22from fastapi import Depends, HTTPException, status
23from sqlalchemy import and_, select
25from airflow.api_fastapi.common.db.common import SessionDep, apply_filters_to_select
26from airflow.api_fastapi.common.parameters import (
27 QueryPartitionedDagRunDagIdFilter,
28 QueryPartitionedDagRunHasCreatedDagRunIdFilter,
29)
30from airflow.api_fastapi.common.partition_helpers import (
31 load_partitioned_timetable,
32 load_partitioned_timetables,
33)
34from airflow.api_fastapi.common.router import AirflowRouter
35from airflow.api_fastapi.core_api.datamodels.ui.partitioned_dag_runs import (
36 PartitionedDagRunAssetResponse,
37 PartitionedDagRunCollectionResponse,
38 PartitionedDagRunDetailResponse,
39 PartitionedDagRunResponse,
40)
41from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
42from airflow.api_fastapi.core_api.security import (
43 ReadableDagsFilterDep,
44 requires_access_asset,
45 requires_access_dag,
46)
47from airflow.models import DagModel
48from airflow.models.asset import (
49 AssetActive,
50 AssetModel,
51 AssetPartitionDagRun,
52 DagScheduleAssetReference,
53 PartitionedAssetKeyLog,
54)
55from airflow.models.dagrun import DagRun
57if TYPE_CHECKING: 57 ↛ 58line 57 didn't jump to line 58 because the condition on line 57 was never true
58 from sqlalchemy.orm import Session
60 from airflow.partition_mappers.base import RollupMapper
61 from airflow.timetables.simple import PartitionedAssetTimetable
64log = structlog.get_logger(logger_name=__name__)
67AssetNameUri: TypeAlias = tuple[str, str]
68"""A ``(name, uri)`` pair identifying an asset."""
71def _fetch_active_assets_per_dag(
72 dag_ids: list[str], session: Session
73) -> dict[str, tuple[list[AssetNameUri], dict[int, AssetNameUri]]]:
74 """
75 Batch-fetch required assets for multiple Dags in a single query.
77 Returns ``{dag_id: ([(name, uri), ...], {asset_id: (name, uri)})}``.
78 Dags with no references are still included with empty containers
79 so callers can index by ``dag_id`` without ``KeyError``.
81 Inactive (deactivated) assets are still included so list-route totals stay
82 symmetric with the detail-route response; the per-asset ``asset_inactive``
83 flag (detail route only) surfaces the freeze state.
84 """
85 rows = session.execute(
86 select(
87 DagScheduleAssetReference.dag_id,
88 AssetModel.id,
89 AssetModel.name,
90 AssetModel.uri,
91 )
92 .join(DagScheduleAssetReference, DagScheduleAssetReference.asset_id == AssetModel.id)
93 .where(DagScheduleAssetReference.dag_id.in_(dag_ids))
94 ).all()
95 result: dict[str, tuple[list[AssetNameUri], dict[int, AssetNameUri]]] = {
96 dag_id: ([], {}) for dag_id in dag_ids
97 }
98 for row in rows:
99 info, id_to_info = result[row.dag_id]
100 info.append((row.name, row.uri))
101 id_to_info[row.id] = (row.name, row.uri)
102 return result
105class _RollupResolution(NamedTuple):
106 """
107 Outcome of resolving an asset's upstream-key requirement for one partition key.
109 Three states, distinguished so callers can match the scheduler's
110 ``_resolve_asset_partition_status`` semantics:
112 - ``keys`` is a ``frozenset`` and ``mapper_failed`` is ``False``: rollup
113 asset, mapper succeeded — use ``keys`` as the required set.
114 - ``keys`` is ``None`` and ``mapper_failed`` is ``False``: not a rollup
115 asset — a single received event satisfies it.
116 - ``keys`` is ``None`` and ``mapper_failed`` is ``True``: rollup asset
117 whose mapper raised — the scheduler treats it as not-yet-satisfied; the
118 UI must not credit any received event either, otherwise progress would
119 silently show "ready" for a run the scheduler will never fire.
120 """
122 keys: frozenset[str] | None = None
123 mapper_failed: bool = False
126def _resolve_rollup_status(
127 dag_model: DagModel | None,
128 rollup_timetable: PartitionedAssetTimetable | None,
129 name: str,
130 uri: str,
131 partition_key: str,
132) -> _RollupResolution:
133 """
134 Resolve the rollup state for *(name, uri)* under the given partition key.
136 The ``dag_model is None`` / ``rollup_timetable is None`` cases short-circuit
137 to "not rollup" because there is nothing to evaluate against, not because
138 the asset is mis-configured.
139 """
140 if dag_model is None or rollup_timetable is None or not dag_model.is_rollup_asset(name=name, uri=uri):
141 return _RollupResolution()
142 try:
143 mapper = rollup_timetable.get_partition_mapper(name=name, uri=uri)
144 return _RollupResolution(keys=frozenset(cast("RollupMapper", mapper).to_upstream(partition_key)))
145 except Exception:
146 # Mismatch with the scheduler's rollup contract. The scheduler writes a
147 # Log row for the same condition (once per misconfiguration); this path
148 # is per-request and lighter.
149 log.warning(
150 "Failed to evaluate rollup mapper; treating asset as not-yet-satisfied",
151 dag_id=dag_model.dag_id,
152 asset_name=name,
153 asset_uri=uri,
154 partition_key=partition_key,
155 exc_info=True,
156 )
157 return _RollupResolution(mapper_failed=True)
160def _build_asset_resolutions(
161 dag_model: DagModel | None,
162 rollup_timetable: PartitionedAssetTimetable | None,
163 asset_info: list[AssetNameUri],
164 partition_key: str,
165) -> dict[AssetNameUri, _RollupResolution]:
166 """
167 Resolve each ``(name, uri)`` asset once per APDR row.
169 Shared between ``_compute_total_required`` and ``_compute_received_count``.
170 ``_resolve_rollup_status`` runs ``to_upstream`` for rollup assets, which
171 yields up to 60 keys (HourWindow) per mapper call. Without this cache each
172 row evaluated the same ``(name, uri)`` twice.
173 """
174 return {
175 (name, uri): _resolve_rollup_status(dag_model, rollup_timetable, name, uri, partition_key)
176 for name, uri in asset_info
177 }
180def _compute_total_required(resolutions: dict[AssetNameUri, _RollupResolution]) -> int:
181 """
182 Sum required upstream events across all assets, using to_upstream for rollup mappers.
184 Non-rollup assets and broken-mapper assets both count as 1: non-rollup needs
185 one event to satisfy, broken-mapper counts as 1 unit of "blocked" so the
186 asset still contributes to the totals (received side credits 0, keeping the
187 progress short of "ready" as the scheduler intends).
188 """
189 return sum(len(res.keys) if res.keys is not None else 1 for res in resolutions.values())
192def _compute_received_count(
193 received_by_asset: dict[int, set[str]],
194 asset_id_to_info: dict[int, AssetNameUri],
195 resolutions: dict[AssetNameUri, _RollupResolution],
196) -> int:
197 """
198 Count received events using rollup-aware deduplication.
200 For rollup assets whose mapper succeeded: count distinct upstream keys that
201 intersect the required set. For non-rollup assets: any logged event
202 satisfies the asset. For rollup assets whose mapper raised: do not credit
203 anything received, matching the scheduler's not-yet-satisfied verdict.
204 """
205 total = 0
206 for asset_id, received_keys in received_by_asset.items():
207 # Logs may reference assets that have been removed from the Dag's
208 # reference list entirely (asset removed from declaration). The skip
209 # filter avoids counting those. Inactive-but-still-declared assets stay
210 # in ``asset_id_to_info`` so their historical received events keep
211 # contributing to the totals.
212 if (info := asset_id_to_info.get(asset_id)) is None:
213 continue
214 res = resolutions[info]
215 if res.keys is not None:
216 total += len(received_keys & res.keys)
217 elif res.mapper_failed:
218 # Match scheduler: a broken mapper means the run is held; don't
219 # let UI counts march toward "ready" while the scheduler holds.
220 continue
221 else:
222 total += 1 if received_keys else 0
223 return total
226partitioned_dag_runs_router = AirflowRouter(tags=["PartitionedDagRun"])
229def _build_response(row, required_count: int, received_count: int) -> PartitionedDagRunResponse:
230 return PartitionedDagRunResponse(
231 id=row.id,
232 dag_id=row.target_dag_id,
233 partition_key=row.partition_key,
234 created_at=row.created_at.isoformat() if row.created_at else None,
235 total_received=received_count,
236 total_required=required_count,
237 state=row.dag_run_state if row.created_dag_run_id else "pending",
238 created_dag_run_id=row.dag_run_id,
239 )
242@partitioned_dag_runs_router.get(
243 "/partitioned_dag_runs",
244 responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
245 dependencies=[Depends(requires_access_asset(method="GET"))],
246)
247def get_partitioned_dag_runs(
248 session: SessionDep,
249 readable_dags_filter: ReadableDagsFilterDep,
250 dag_id: QueryPartitionedDagRunDagIdFilter,
251 has_created_dag_run_id: QueryPartitionedDagRunHasCreatedDagRunIdFilter,
252) -> PartitionedDagRunCollectionResponse:
253 """Return PartitionedDagRuns. Filter by dag_id and/or has_created_dag_run_id."""
254 # The Dag-existence check is intentionally deferred to the empty-results branch
255 # below. In the happy path (rows exist), filtering by dag_id already restricts
256 # to that Dag, so an extra DagModel lookup just to validate existence wastes a
257 # query. We only consult DagModel when we have no rows to report — that's the
258 # only case where the distinction (404 vs empty) matters.
260 query = select(
261 AssetPartitionDagRun.id,
262 AssetPartitionDagRun.target_dag_id,
263 AssetPartitionDagRun.partition_key,
264 AssetPartitionDagRun.created_at,
265 AssetPartitionDagRun.created_dag_run_id,
266 DagRun.run_id.label("dag_run_id"),
267 DagRun.state.label("dag_run_state"),
268 ).outerjoin(DagRun, AssetPartitionDagRun.created_dag_run_id == DagRun.id)
269 query = apply_filters_to_select(statement=query, filters=[dag_id, has_created_dag_run_id])
270 readable_dag_ids = readable_dags_filter.value
271 if readable_dag_ids is not None:
272 query = query.where(AssetPartitionDagRun.target_dag_id.in_(readable_dag_ids))
273 query = query.order_by(AssetPartitionDagRun.created_at.desc())
275 if not (rows := session.execute(query).all()):
276 if dag_id.value is not None:
277 # An unreadable-but-existing Dag must return 404 too — otherwise the caller
278 # can probe by dag_id and learn which Dags exist outside their permitted set.
279 if readable_dag_ids is not None:
280 if dag_id.value not in readable_dag_ids:
281 raise HTTPException(
282 status.HTTP_404_NOT_FOUND, f"Dag with id {dag_id.value} was not found"
283 )
284 elif session.scalar(select(DagModel.dag_id).where(DagModel.dag_id == dag_id.value)) is None:
285 raise HTTPException(status.HTTP_404_NOT_FOUND, f"Dag with id {dag_id.value} was not found")
286 return PartitionedDagRunCollectionResponse(partitioned_dag_runs=[], total=0)
288 # Batch-fetch DagModels (for cached partition_mapper_info), required assets,
289 # and APDR log entries in three single queries instead of N per-Dag queries.
290 # Timetables are only loaded for Dags that actually have rollup mappers,
291 # since that's the only case where ``to_upstream`` evaluation is needed.
292 # A SQL count subquery for total_received cannot honour rollup windows
293 # without running the mapper, so the rollup-aware Python computation runs
294 # uniformly across single-Dag and global views.
295 unique_dag_ids = list({row.target_dag_id for row in rows})
296 dag_models: dict[str, DagModel] = {
297 dm.dag_id: dm
298 for dm in session.scalars(select(DagModel).where(DagModel.dag_id.in_(unique_dag_ids))).all()
299 }
300 assets_by_dag = _fetch_active_assets_per_dag(unique_dag_ids, session)
301 # Batch-fetch timetables for only the Dags that need rollup evaluation so the
302 # query stays bounded by ``has_rollup_mappers`` rather than the full Dag set.
303 rollup_dag_ids = [
304 d_id for d_id in unique_dag_ids if (dm := dag_models.get(d_id)) is not None and dm.has_rollup_mappers
305 ]
306 rollup_timetables_by_dag: dict[str, PartitionedAssetTimetable | None] = {
307 d_id: None for d_id in unique_dag_ids
308 }
309 rollup_timetables_by_dag.update(load_partitioned_timetables(rollup_dag_ids, session))
311 apdr_ids = [row.id for row in rows]
312 log_by_apdr: dict[int, dict[int, set[str]]] = {}
313 for pakl_row in session.execute(
314 select(
315 PartitionedAssetKeyLog.asset_partition_dag_run_id,
316 PartitionedAssetKeyLog.asset_id,
317 PartitionedAssetKeyLog.source_partition_key,
318 ).where(PartitionedAssetKeyLog.asset_partition_dag_run_id.in_(apdr_ids))
319 ).all():
320 log_by_apdr.setdefault(pakl_row.asset_partition_dag_run_id, {}).setdefault(
321 pakl_row.asset_id, set()
322 ).add(pakl_row.source_partition_key)
324 results = []
325 for row in rows:
326 asset_info, asset_id_to_info = assets_by_dag[row.target_dag_id]
327 resolutions = _build_asset_resolutions(
328 dag_models.get(row.target_dag_id),
329 rollup_timetables_by_dag[row.target_dag_id],
330 asset_info,
331 row.partition_key,
332 )
333 results.append(
334 _build_response(
335 row,
336 _compute_total_required(resolutions),
337 _compute_received_count(log_by_apdr.get(row.id, {}), asset_id_to_info, resolutions),
338 )
339 )
341 asset_expressions: dict[str, dict | None] | None = None
342 if dag_id.value is None:
343 asset_expressions = {dm.dag_id: dm.asset_expression for dm in dag_models.values()}
345 return PartitionedDagRunCollectionResponse(
346 partitioned_dag_runs=results,
347 total=len(results),
348 asset_expressions=asset_expressions,
349 )
352@partitioned_dag_runs_router.get(
353 "/pending_partitioned_dag_run/{dag_id}",
354 dependencies=[Depends(requires_access_asset(method="GET")), Depends(requires_access_dag(method="GET"))],
355)
356def get_pending_partitioned_dag_run(
357 dag_id: str,
358 partition_key: str,
359 session: SessionDep,
360) -> PartitionedDagRunDetailResponse:
361 """Return full details for pending PartitionedDagRun."""
362 # partition_key is a query param, not a path segment: it is a free-form key
363 # (up to 250 chars) that may itself contain "/", which would otherwise be
364 # ambiguous (or mis-routed) as a path segment.
365 partitioned_dag_run = session.execute(
366 select(
367 AssetPartitionDagRun.id,
368 AssetPartitionDagRun.target_dag_id,
369 AssetPartitionDagRun.partition_key,
370 AssetPartitionDagRun.created_at,
371 AssetPartitionDagRun.updated_at,
372 DagRun.run_id.label("created_dag_run_id"),
373 )
374 .outerjoin(DagRun, AssetPartitionDagRun.created_dag_run_id == DagRun.id)
375 .where(
376 AssetPartitionDagRun.target_dag_id == dag_id,
377 AssetPartitionDagRun.partition_key == partition_key,
378 AssetPartitionDagRun.created_dag_run_id.is_(None),
379 )
380 # Duplicate pending rows for the same (dag_id, partition_key) can exist
381 # after a crash; mirror _get_or_create_apdr and work on the latest one.
382 .order_by(AssetPartitionDagRun.id.desc())
383 .limit(1)
384 ).first()
386 if partitioned_dag_run is None:
387 raise HTTPException(
388 status.HTTP_404_NOT_FOUND,
389 f"No PartitionedDagRun for dag={dag_id} partition={partition_key}",
390 )
392 # Collect received upstream partition keys per asset for this partition run.
393 # Use a set to deduplicate: multiple events for the same key count as one.
394 received_keys_by_asset: dict[int, set[str]] = {}
395 for row in session.execute(
396 select(
397 PartitionedAssetKeyLog.asset_id,
398 PartitionedAssetKeyLog.source_partition_key,
399 ).where(PartitionedAssetKeyLog.asset_partition_dag_run_id == partitioned_dag_run.id)
400 ):
401 received_keys_by_asset.setdefault(row.asset_id, set()).add(row.source_partition_key)
403 dag_model = session.get(DagModel, dag_id)
404 asset_rows = session.execute(
405 select(
406 AssetModel.id,
407 AssetModel.uri,
408 AssetModel.name,
409 AssetActive.name.label("active_name"),
410 )
411 .join(DagScheduleAssetReference, DagScheduleAssetReference.asset_id == AssetModel.id)
412 .outerjoin(
413 AssetActive,
414 and_(AssetActive.name == AssetModel.name, AssetActive.uri == AssetModel.uri),
415 )
416 .where(DagScheduleAssetReference.dag_id == dag_id)
417 .order_by(AssetModel.uri)
418 ).all()
420 # Skip the timetable load when no rollup mapper is configured — the cached
421 # ``partition_mapper_info`` already tells us whether we will need
422 # ``to_upstream`` evaluation, which is the only thing the timetable adds here.
423 has_rollup_mappers = dag_model is not None and dag_model.has_rollup_mappers
424 rollup_timetable = load_partitioned_timetable(dag_id, session) if has_rollup_mappers else None
426 assets = []
427 for asset_row in asset_rows:
428 received_keys = sorted(received_keys_by_asset.get(asset_row.id, set()))
429 required_keys: list[str] = [partition_key]
430 is_rollup = (
431 has_rollup_mappers
432 and dag_model is not None
433 and dag_model.is_rollup_asset(name=asset_row.name, uri=asset_row.uri)
434 )
435 mapper_failed = False
436 if is_rollup and rollup_timetable is not None:
437 try:
438 mapper = rollup_timetable.get_partition_mapper(name=asset_row.name, uri=asset_row.uri)
439 required_keys = sorted(cast("RollupMapper", mapper).to_upstream(partition_key))
440 except Exception:
441 # Mirror the scheduler: a misconfigured rollup mapper holds the
442 # run, so the detail view must not show "received" for an asset
443 # the scheduler treats as not-yet-satisfied.
444 log.warning(
445 "Failed to evaluate rollup mapper; treating asset as not-yet-satisfied",
446 dag_id=dag_id,
447 asset_name=asset_row.name,
448 asset_uri=asset_row.uri,
449 partition_key=partition_key,
450 exc_info=True,
451 )
452 mapper_failed = True
453 if mapper_failed:
454 received_keys = []
455 required_keys = []
456 received_count = 0
457 required_count = 1
458 elif is_rollup:
459 received_count = len(received_keys)
460 required_count = len(required_keys)
461 else:
462 # Match the list route's _compute_received_count: a non-rollup asset is
463 # satisfied by any single received event, so credit caps at 1 even if
464 # several distinct upstream keys mapped onto this one target key.
465 required_count = len(required_keys)
466 received_count = 1 if received_keys else 0
467 assets.append(
468 PartitionedDagRunAssetResponse(
469 asset_id=asset_row.id,
470 asset_name=asset_row.name,
471 asset_uri=asset_row.uri,
472 received=received_count >= required_count and required_count > 0,
473 received_count=received_count,
474 required_count=required_count,
475 received_keys=received_keys,
476 required_keys=required_keys,
477 is_rollup=is_rollup,
478 mapper_error=mapper_failed,
479 asset_inactive=(asset_row.active_name is None),
480 )
481 )
483 total_received = sum(a.received_count for a in assets)
484 total_required = sum(a.required_count for a in assets)
485 asset_expression = dag_model.asset_expression if dag_model is not None else None
487 return PartitionedDagRunDetailResponse(
488 id=partitioned_dag_run.id,
489 dag_id=dag_id,
490 partition_key=partition_key,
491 created_at=partitioned_dag_run.created_at.isoformat() if partitioned_dag_run.created_at else None,
492 updated_at=partitioned_dag_run.updated_at.isoformat() if partitioned_dag_run.updated_at else None,
493 created_dag_run_id=partitioned_dag_run.created_dag_run_id,
494 assets=assets,
495 total_required=total_required,
496 total_received=total_received,
497 asset_expression=asset_expression,
498 )