diff --git a/.agents/skills/prepare-providers-documentation/SKILL.md b/.agents/skills/prepare-providers-documentation/SKILL.md index a02798ac4997e..fb5602f894a62 100644 --- a/.agents/skills/prepare-providers-documentation/SKILL.md +++ b/.agents/skills/prepare-providers-documentation/SKILL.md @@ -4,7 +4,8 @@ description: > Replace the manual commit-by-commit classification step in `breeze release-management prepare-provider-documentation` with AI-driven classification. For each provider with pending changes, analyze every PR - (using sub-agents per PR), pay special attention to potentially breaking + (batched into one sub-agent per provider, not one per PR), pay special + attention to potentially breaking changes by inspecting the actual diff, scope multi-provider PRs to the current provider's slice, ask the release manager when uncertain, and apply version bumps + changelog entries. Use during the regular provider @@ -107,49 +108,73 @@ date — running breeze the first time below will recreate and fetch it. The skill runs in five phases. Mark tasks with `TaskCreate` for each phase and tick them off as you go — the release manager wants to see progress. -### Phase 1 — Discover providers with pending changes +### Phase 1 — Discover and pre-classify pending changes (deterministic) -For each provider, the source of truth for "what changed since last release" -is the same git query breeze uses internally: commits between the latest -release tag for that provider (`providers-/`) and -`apache-https-for-providers/`, restricted to the provider's own -folders. +The source of truth for "what changed since last release" is the same git +query breeze uses internally: commits between the latest release tag for that +provider (`providers-/`) and `apache-https-for-providers/`, +restricted to the provider's own folders. -Discover in batch by running: +Run the **deterministic classifier** — it discovers every provider with pending +changes **and** pre-classifies each commit with hard-coded, high-confidence +rules, flagging only the genuinely ambiguous ones as `needs_llm`. No random +answers, nothing to discard: ```bash -breeze release-management prepare-provider-documentation \ - --non-interactive \ - --skip-changelog \ - --skip-readme \ - --release-date "$RELEASE_DATE" +breeze release-management classify-provider-changes \ + --base-branch main \ + --output-file /tmp/provider-changes.json +# scope to a subset by appending provider ids, e.g. ... amazon cncf.kubernetes ``` -> [!WARNING] -> Do **not** commit the result of that command. `--non-interactive` answers -> the classification prompts with random values — Claude will overwrite the -> changelog and version bumps in Phase 4 with real classifications. The only -> reason to run breeze first is to refresh the apache remote, regenerate -> build files, and confirm which providers have pending changes (read the -> "Summary of prepared documentation" block at the end). +The JSON it writes: + +```json +{ + "base_branch": "main", + "providers": { + "amazon": { + "current_version": "9.29.0", + "commits": [ + {"hash": "c2dbd7a75a", "pr": "67987", "subject": "Fix IDC domain S3 path resolution", + "classification": "needs_llm", "reason": "no high-confidence deterministic rule matched"}, + {"hash": "abc123", "pr": "68087", "subject": "Bump the edge-ui-package-updates group ...", + "classification": "misc", "reason": "dependency bump (subject starts with 'Bump')"} + ] + } + } +} +``` -Record from the summary: +How to read it: -- **Success** — providers that had real changes (these need classification). -- **Docs only** — providers with only documentation changes (already handled - by breeze; skip in Phase 2). -- **Skipped on no changes** — nothing to do. +- Providers under `providers` have pending changes (these need attention). +- `classification ∈ {documentation, skip, misc}` are **decided by rules — take + them as-is**, no sub-agent needed (doc-only → `documentation`, test/example + only → `skip`, `Bump …` dependency bump → `misc`). +- `classification == needs_llm` → **Phase 3 decides** with a sub-agent. These are + the only commits that need LLM analysis. +- A provider with a `note`/`error` (e.g. a brand-new provider with no prior + release tag) → treat as an **initial release** and classify by hand. -Reset the per-provider files that breeze touched but you'll be rewriting -yourself before continuing: +> [!NOTE] +> The classifier is deliberately conservative: `Fix …`/`Add …` subjects are +> **not** auto-classified (an "Add …" can be a breaking change), so they come +> back as `needs_llm`. The rules live in `classify_change_deterministically` +> (`dev/breeze/src/airflow_breeze/prepare_providers/provider_documentation.py`). + +Then regenerate the auto-generated build files (this does **no** classification, +so nothing random is produced): ```bash +breeze release-management prepare-provider-documentation \ + --reapply-templates-only --release-date "$RELEASE_DATE" git checkout -- $(git diff --name-only -- '**/provider.yaml' '**/changelog.rst') ``` This leaves the regenerated build files (`__init__.py`, `README.rst`, -`pyproject.toml`, `conf.py`, `get_provider_info.py`, `index.rst`) in place -and discards only the stuff Claude is about to rewrite. +`pyproject.toml`, `conf.py`, `get_provider_info.py`, `index.rst`) in place and +discards only the changelog/version files Claude is about to rewrite itself. ### Phase 2 — Per-provider commit list @@ -190,7 +215,7 @@ each commit. Note that some old providers also have legacy paths under `provider_details.possible_old_provider_paths` semantics by checking the provider's `provider.yaml` history if needed). -### Phase 3 — Classify each PR with sub-agents +### Phase 3 — Classify the PRs (inline, or batched per-provider sub-agents) For each commit, classify it into one of: @@ -204,35 +229,56 @@ For each commit, classify it into one of: | `s` | Skip (test/CI/example only — no user impact) | none | | `v` | Min Airflow version bump | minor (treated as misc + bump) | -#### Auto-classify cheap cases first - -Before spawning a sub-agent, apply the same fast heuristics breeze uses -(see `classify_provider_pr_files` in -`dev/breeze/src/airflow_breeze/prepare_providers/provider_documentation.py`): - -- All changed files match `providers//docs/**/*.rst` → **`d`** (docs). -- All changed files match `providers//tests/**` or - `providers//src/airflow/providers//example_dags/**` → **`s`** (skip). -- Subject contains `Bump minimum Airflow version` and only `__init__.py` / - `provider.yaml` changed → **`v`**. - -Note these classifications and move on — no sub-agent needed. - -#### Sub-agent per PR for the rest - -For the remaining commits, spawn sub-agents in parallel (batches of 5–10 to -avoid context pressure). Use the `Explore` agent type — they need read-only -access. Brief each sub-agent with: +#### Take the deterministic classifications from Phase 1 + +`classify-provider-changes` (Phase 1) already classified every commit it could +with hard-coded rules. Read `/tmp/provider-changes.json` and: + +- Use any commit whose `classification` is `documentation`, `skip`, or `misc` + **as-is** — these map to `d`, `s`, `m` respectively; no sub-agent needed. +- Only commits with `classification: needs_llm` go to a sub-agent (below). + +The deterministic rules (doc-only → `d`, test/example-only → `s`, `Bump …` +dependency bump → `m`) are exactly the cheap cases — now computed once by +breeze (`classify_change_deterministically`) instead of re-derived here. If you +ever need the min-Airflow-bump case (`v`), that one is still a `needs_llm` +judgement: a sub-agent should flag it when a PR bumps the provider's minimum +Airflow version. + +#### Classify the `needs_llm` commits — batched per provider, not one agent per PR + +Only the commits the classifier returned as `needs_llm` still need a sub-agent. +Classification is the token-heavy part of this skill, so spend sub-agents +sparingly. Do **not** spawn one sub-agent per PR — that is one agent per +commit and balloons to hundreds of agents on a normal release wave. Pick the +smallest fan-out that fits the volume: + +- **Few `needs_llm` commits remain (≲ 15 across all providers) → classify inline.** + Read each PR and its provider-scoped diff yourself, in this context. Spawn no + sub-agents at all. +- **More than that → one sub-agent per provider.** Each agent classifies that + provider's *entire* remaining `needs_llm` list in a single pass. This is the + natural unit: multi-provider PRs are classified independently per provider + anyway (see Cross-Cutting Rules), and one provider-scoped agent amortizes + the breaking-change checklist across all of that provider's commits instead + of paying a fresh agent spin-up per commit. Only split a provider across + more than one agent when its remaining list is large (> ~25 commits) — chunk + it then. This keeps the sub-agent count at roughly the number of providers + with pending changes, not the number of commits. + +Use the `Explore` agent type — they need read-only access. Brief each +sub-agent with its provider and the whole batch of commits it owns: ``` -Classify a single Apache Airflow provider PR. +Classify a batch of Apache Airflow provider PRs for ONE provider. -PR: # -Commit: -Subject: Provider: (path: providers//) +Commits to classify () — one row per PR: + # + # + …(this provider's full remaining list) -Tasks: +For EACH commit above: 1. Read the PR's title, body, and labels: `gh pr view --json title,body,labels,files` 2. Read the diff for the slice of the PR that touched @@ -250,14 +296,14 @@ Tasks: changes, type-hint cleanups, no user-visible behavior - skip: only tests/examples/CI for this provider's slice - min_airflow_bump: explicitly bumps the minimum Airflow version pin -4. Output strictly: - CLASSIFICATION: - CONFIDENCE: - JUSTIFICATION: - BREAKING_RISK: (set "maybe" when the diff has any - signal from the breaking-change - checklist, even if you think the - author intended otherwise) +4. Set BREAKING_RISK to "maybe" whenever the diff has any signal from the + breaking-change checklist below, even if you think the author intended + otherwise. + +Output one row per commit and nothing else, in this exact pipe format +( rows for commits): + + # | | | | Breaking-change checklist (any of these → BREAKING_RISK >= maybe; usually breaking unless clearly behind a deprecation shim): @@ -290,7 +336,8 @@ that removes a public method is breaking. A PR titled "BREAKING: rename foo" that only renames a private symbol is not. ``` -Collect all sub-agent results into a table. +Collect every sub-agent's rows (and any you classified inline) into one +classification table for Phase 3.5. ### Phase 3.5 — Confirm with the release manager @@ -542,9 +589,12 @@ If there are zero new commits for a provider, skip it. ### Incremental Phase 3 — Classify the new commits Same logic as Phase 3 of the initial run — including the auto-classify -heuristic for docs/test-only changes and the sub-agent-per-PR pattern with -the breaking-change checklist. The output is a per-provider table mapping -each new commit hash to a classification. +heuristic for docs/test-only changes and the batched classification (inline +when few commits remain, otherwise one sub-agent per provider) with the +breaking-change checklist. Incremental runs usually have only a handful of new +commits, so prefer classifying them inline rather than spawning any sub-agent. +The output is a per-provider table mapping each new commit hash to a +classification. ### Incremental Phase 3.5 — Decide whether to escalate the version bump @@ -612,8 +662,8 @@ When a single PR touches several providers (e.g. `Add Python 3.14 Support (#63520)` touches dozens), classify it **independently per provider**. The same PR can be `feature` in one provider (a real new capability) and `misc` in another (just a constraint bump in -`pyproject.toml`). Always scope the sub-agent's diff inspection to the -current provider's path: +`pyproject.toml`). Always scope the diff inspection (whether inline or in a +per-provider sub-agent) to the current provider's path: ```bash gh pr diff -- 'providers//**' diff --git a/.github/instructions/code-review.instructions.md b/.github/instructions/code-review.instructions.md index 08eec81085d32..da2674c4c1a32 100644 --- a/.github/instructions/code-review.instructions.md +++ b/.github/instructions/code-review.instructions.md @@ -32,7 +32,7 @@ Use these rules when reviewing pull requests to the Apache Airflow repository. - **Flag any `@lru_cache(maxsize=None)`.** This creates an unbounded cache — every unique argument set is cached forever. Note: `@lru_cache()` without arguments defaults to `maxsize=128` and is fine. - **Flag any heavy import** (e.g., `kubernetes.client`) in multi-process code paths that is not behind a `TYPE_CHECKING` guard. - **Flag any file, connection, or session opened without a context manager or `try/finally`.** -- **Flag any new `raise AirflowException` usage.** The community has stopped adding new ones (enforced by the `check-no-new-airflow-exceptions` prek hook) — prefer Python's standard exceptions (`ValueError`, `TypeError`, `OSError`), or a dedicated class in the appropriate `exceptions.py`. **Do not suggest changing specific exceptions back to `AirflowException`.** +- **Flag any new `raise AirflowException` usage.** The community is reducing direct `AirflowException` usage, not increasing it; new ones are not allowed (enforced by the `check-no-new-airflow-exceptions` prek hook) — prefer Python's standard exceptions (`ValueError`, `TypeError`, `OSError`), or a dedicated class in the appropriate `exceptions.py`. **The one exception is a pure relocation: an already-existing `AirflowException` moved verbatim during a refactor (e.g. code moved between files) is not a new usage — do not flag it, but confirm the diff removes the identical raise elsewhere and leaves it otherwise unchanged.** **Do not suggest changing specific exceptions back to `AirflowException`.** ## Testing Requirements diff --git a/AGENTS.md b/AGENTS.md index 4f8b4a37d44f6..f0eb03813f4b7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -139,7 +139,7 @@ reported as such are described in "What is NOT considered a security vulnerabili - In `airflow-core`, functions with a `session` parameter must not call `session.commit()`. Use keyword-only `session` parameters. - Imports at top of file. Valid exceptions: circular imports, lazy loading for worker isolation, `TYPE_CHECKING` blocks. - Guard heavy type-only imports (e.g., `kubernetes.client`) with `TYPE_CHECKING` in multi-process code paths. -- Define dedicated exception classes or use existing exceptions such as `ValueError` instead of raising the broad `AirflowException` directly. Each error case should have a specific exception type that conveys what went wrong. +- Define dedicated exception classes or use existing exceptions such as `ValueError` instead of raising the broad `AirflowException` directly. Each error case should have a specific exception type that conveys what went wrong. **Never add new direct `raise AirflowException(...)` usages — the community is actively reducing them, not adding more, and the `check-no-new-airflow-exceptions` prek hook enforces this across `airflow-core`, `airflow-ctl`, `task-sdk`, `providers`, and `shared`.** Prefer a Python built-in (`ValueError`, `TypeError`, `OSError`, …) or a dedicated class in the appropriate `exceptions.py`. The only acceptable way an `AirflowException` line may move is relocating an already-existing one verbatim during a refactor (e.g. moving code between files) — that is not a new usage. When you touch code that already raises `AirflowException`, prefer narrowing it to a more specific exception rather than leaving or duplicating it. - Translate domain-layer exceptions to `HTTPException` at FastAPI route boundaries. In `airflow-core/src/airflow/core_api/` route handlers, catch errors raised by domain code (e.g., `ValueError` from `airflow.state.metastore.MetastoreStateBackend` for a missing row or invalid input) and re-raise as `HTTPException` with the right status (`404` for not-found, `400` for invalid input). Otherwise they propagate as `500 Internal Server Error`, leaking internals and misleading clients. - Bulk `DELETE`/`UPDATE` in the scheduler loop or any synchronous interval task (e.g. `call_regular_interval` callbacks) must be batched with `LIMIT` and committed between batches — never issue a single unbounded bulk write against a user-driven table. Unbounded bulk writes hold row locks for the entire transaction (blocking concurrent writers) and stall the scheduler main loop. Filter columns used by the cleanup must be indexed. Follow the batching pattern in `airflow-core/src/airflow/utils/db_cleanup.py`. - Name functions and methods with action verbs: `get_`, `extract_`, `find_`, `compute_`, `build_`, etc. Avoid noun-only names like `_serialize_keys` or `_base_names` — they read as attributes, not callables. Predicates (`is_`, `has_`) are the one exception. diff --git a/airflow-core/src/airflow/api/common/trigger_dag.py b/airflow-core/src/airflow/api/common/trigger_dag.py index 4f8ca4c25e718..8dfd06c1ac9e8 100644 --- a/airflow-core/src/airflow/api/common/trigger_dag.py +++ b/airflow-core/src/airflow/api/common/trigger_dag.py @@ -118,6 +118,8 @@ def _trigger_dag( if dag_run := DagRun.find_duplicate(dag_id=dag_id, run_id=run_id): raise DagRunAlreadyExists(dag_run) + partition_date = dag.timetable.resolve_partition_date(partition_key) + run_conf = None if is_arg_set(conf): run_conf = _normalize_conf(conf) @@ -133,6 +135,7 @@ def _trigger_dag( note=note, state=DagRunState.QUEUED, partition_key=partition_key, + partition_date=partition_date, session=session, ) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py index 88ffceca6d7c7..033437e2ad98b 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py @@ -191,6 +191,9 @@ def validate_context(self, dag: SerializedDAG) -> dict: run_after=timezone.coerce_datetime(run_after), data_interval=data_interval, ) + + partition_date = dag.timetable.resolve_partition_date(self.partition_key) + return { "run_id": run_id, "logical_date": coerced_logical_date, @@ -199,6 +202,7 @@ def validate_context(self, dag: SerializedDAG) -> dict: "conf": self.conf, "note": self.note, "partition_key": self.partition_key, + "partition_date": partition_date, } diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py index 0860b29da9424..30ae090289c75 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -657,6 +657,7 @@ def trigger_dag_run( triggering_user_name=user.get_name(), state=DagRunState.QUEUED, partition_key=params["partition_key"], + partition_date=params["partition_date"], session=session, ) except (ParamValidationError, ValueError) as e: diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/connections.py index bc6e5e49f75b7..74f5d05df776e 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/connections.py @@ -121,7 +121,12 @@ def handle_bulk_create( if connection.connection_id in create_connection_ids: if connection.connection_id in matched_connection_ids: existed_connection = existed_connections_dict[connection.connection_id] - for key, val in connection.model_dump(by_alias=True).items(): + # Only overwrite fields the request actually provided (see pools.py for the + # full rationale). Plain ``model_dump()`` resets omitted fields to their + # defaults on the existing connection — e.g. silently nulling ``team_name`` + # multi-team ownership. ``exclude_unset=True`` writes only the fields present + # in the request body. + for key, val in connection.model_dump(by_alias=True, exclude_unset=True).items(): setattr(existed_connection, key, val) else: self.session.add(Connection(**connection.model_dump(by_alias=True))) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/pools.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/pools.py index dce54099772de..476c45b8f5ccf 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/pools.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/pools.py @@ -149,7 +149,14 @@ def handle_bulk_create(self, action: BulkCreateAction[PoolBody], results: BulkAc if pool.pool in create_pool_names: if pool.pool in matched_pool_names: existed_pool = existing_pools_dict[pool.pool] - for key, val in pool.model_dump().items(): + # Only overwrite fields the request actually provided. Plain ``model_dump()`` + # emits every field at its default, so an overwrite that omits e.g. + # ``team_name``/``description``/``include_deferred`` silently resets them on the + # existing pool — most damagingly nulling its multi-team ``team_name`` ownership. + # ``exclude_unset=True`` writes only fields present in the request body, so + # omitted fields keep their current value while an explicitly-set field (even + # ``None``) is still applied. + for key, val in pool.model_dump(exclude_unset=True).items(): setattr(existed_pool, key, val) else: self.session.add(Pool(**pool.model_dump())) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py index 872c74957527e..1092369e913bc 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py @@ -34,7 +34,7 @@ from airflow.api_fastapi.execution_api.datamodels.taskinstance import DagRun from airflow.api_fastapi.execution_api.datamodels.token import TIToken from airflow.api_fastapi.execution_api.security import CurrentTIToken -from airflow.exceptions import DagNotPartitionedError, DagRunAlreadyExists +from airflow.exceptions import DagNotPartitionedError, DagRunAlreadyExists, InvalidPartitionKeyError from airflow.models.dag import DagModel from airflow.models.dagrun import DagRun as DagRunModel from airflow.models.taskinstance import TaskInstance @@ -158,6 +158,11 @@ def trigger_dag_run( status.HTTP_400_BAD_REQUEST, detail={"reason": "not_partitioned", "message": str(e)}, ) + except InvalidPartitionKeyError as e: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, + detail={"reason": "invalid_partition_key", "message": str(e)}, + ) from e @router.post( diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index ae71e41cf8e5f..0f862e1625832 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -544,6 +544,16 @@ def _emit_task_span(ti, state): return dr_ctx = TraceContextTextMapPropagator().extract(ti.dag_run.context_carrier) + # Skip if the run was head-sampled out, so every span in the run agrees with the + # carrier's decision. A parent-based sampler would already drop this child span, + # but the explicit check also covers non-parent-based samplers (which ignore the + # parent and would re-sample it in) and short-circuits before building the span. + # An invalid/empty carrier (legacy/NULL) recorded no decision, so it falls through + # and still emits — preserving prior behavior. + dr_span_context = trace.get_current_span(context=dr_ctx).get_span_context() + if dr_span_context.is_valid and not dr_span_context.trace_flags.sampled: + return + ti_ctx = TraceContextTextMapPropagator().extract(ti.context_carrier) ti_span = trace.get_current_span(context=ti_ctx) span_context = ti_span.get_span_context() diff --git a/airflow-core/src/airflow/assets/manager.py b/airflow-core/src/airflow/assets/manager.py index f72c533c5a0e7..c8d2edef4f104 100644 --- a/airflow-core/src/airflow/assets/manager.py +++ b/airflow-core/src/airflow/assets/manager.py @@ -44,7 +44,7 @@ ) from airflow.models.log import Log from airflow.timetables.base import compute_rollup_fingerprint -from airflow.utils.helpers import is_container +from airflow.utils.helpers import is_container, prune_dict from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.sqlalchemy import get_dialect_name, with_row_locks @@ -397,7 +397,12 @@ def register_asset_change( ) ) - stats.incr("asset.updates") + team_name = None + if task_instance and conf.getboolean("core", "multi_team"): + from airflow.models.dag import DagModel + + team_name = DagModel.get_team_name(task_instance.dag_id, session=session) + stats.incr("asset.updates", tags=prune_dict({"team_name": team_name})) dags_to_queue = ( dags_to_queue_from_asset | dags_to_queue_from_asset_alias | dags_to_queue_from_asset_ref @@ -405,7 +410,6 @@ def register_asset_change( if conf.getboolean("core", "multi_team"): if task_instance: - team_name = DagModel.get_team_name(task_instance.dag_id, session=session) resolved_source_teams = {team_name} if team_name else set() # Resolve consumer-team filtering from the outlet reference outlet_ref = session.scalar( diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index a288c4fe6ed6d..86e422451a08b 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2453,7 +2453,8 @@ smtp: default: "30" smtp_retry_limit: description: | - Defines the maximum number of times Airflow will attempt to connect to the SMTP server. + Defines the number of times Airflow will attempt to connect to the SMTP server after the first + attempt. version_added: 2.0.0 type: integer example: ~ diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 92ca10c67fa68..b2428985a6506 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -69,6 +69,7 @@ from airflow.sdk.log import init_log_file, logging_processors from airflow.typing_compat import assert_never from airflow.utils.file import list_py_file_paths, might_contain_dag +from airflow.utils.helpers import prune_dict from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.net import get_hostname from airflow.utils.process_utils import ( @@ -84,7 +85,7 @@ ) if TYPE_CHECKING: - from collections.abc import Callable, Iterable, Iterator, Sequence + from collections.abc import Callable, Collection, Iterable, Iterator, Sequence from socket import socket from sqlalchemy.orm import Session @@ -265,6 +266,8 @@ class DagFileProcessorManager(LoggingMixin): _dag_bundles: list[BaseDagBundle] = attrs.field(factory=list, init=False) _bundle_versions: dict[str, str | None] = attrs.field(factory=dict, init=False) _bundle_version_data: dict[str, dict | None] = attrs.field(factory=dict, init=False) + _multi_team: bool = attrs.field(factory=lambda: conf.getboolean("core", "multi_team"), init=False) + _bundle_name_to_team_name: dict[str, str | None] = attrs.field(factory=dict, init=False) _processors: dict[DagFileInfo, DagFileProcessorProcess] = attrs.field(factory=dict, init=False) @@ -306,6 +309,19 @@ def register_exit_signals(self): # So that we ignore the debug dump signal, making it easier to send signal.signal(signal.SIGUSR2, signal.SIG_IGN) + def _get_team_names(self, bundle_names: Collection[str]) -> dict[str, str | None]: + if not self._multi_team or not bundle_names: + return {} + missing = [name for name in bundle_names if name not in self._bundle_name_to_team_name] + if missing: + queried = DagBundleModel.get_team_names(missing) + for name in missing: + self._bundle_name_to_team_name[name] = queried.get(name) + return {name: self._bundle_name_to_team_name.get(name) for name in bundle_names} + + def _get_team_name(self, bundle_name: str) -> str | None: + return self._get_team_names({bundle_name}).get(bundle_name) + def _exit_gracefully(self, signum, frame): """Clean up DAG file processors to avoid leaving orphan processes.""" self.log.info("Exiting gracefully upon receiving signal %s", signum) @@ -719,7 +735,8 @@ def _add_callback_to_queue(self, request: CallbackRequest) -> None: ) self._callback_to_execute[file_info].append(request) self._add_files_to_queue([file_info], mode="front") - stats.incr("dag_processing.other_callback_count") + team_name = self._get_team_name(file_info.bundle_name) + stats.incr("dag_processing.other_callback_count", tags=prune_dict({"team_name": team_name})) @provide_session def get_bundle_state(self, bundle_name: str, *, session: Session = NEW_SESSION) -> BundleState | None: @@ -885,6 +902,8 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): ) if any_refreshed: + # Bundle-to-team assignments can only change on bundle refresh, so clear the cache. + self._bundle_name_to_team_name = {} self.handle_removed_files(known_files=known_files) self._resort_file_queue() self._add_new_files_to_queue(known_files=known_files) @@ -1020,6 +1039,8 @@ def _log_file_processing_stats(self, known_files: dict[str, set[DagFileInfo]]): utcnow = timezone.utcnow() now = time.monotonic() + bundle_to_team = self._get_team_names({bundle_name for bundle_name in known_files}) + for files in known_files.values(): for file in files: stat = self._file_stats[file] @@ -1040,11 +1061,14 @@ def _log_file_processing_stats(self, known_files: dict[str, set[DagFileInfo]]): stats.gauge( "dag_processing.last_run.seconds_ago", seconds_ago, - tags={ - "file_path": file.normalized_file_path_for_stats, - "bundle_name": normalize_name_for_stats(file.bundle_name), - "file_name": file_name, - }, + tags=prune_dict( + { + "file_path": file.normalized_file_path_for_stats, + "bundle_name": normalize_name_for_stats(file.bundle_name), + "file_name": file_name, + "team_name": bundle_to_team.get(file.bundle_name), + } + ), ) rows.append( @@ -1136,6 +1160,9 @@ def remove_orphaned_file_stats(self, present: set[DagFileInfo]): def terminate_orphan_processes(self, present: set[DagFileInfo]): """Stop processors that are working on deleted files.""" present_keys = {file.presence_key for file in present} + + bundle_to_team = self._get_team_names({file.bundle_name for file in self._processors}) + for file in list(self._processors.keys()): if file.presence_key not in present_keys: processor = self._processors.pop(file, None) @@ -1145,7 +1172,13 @@ def terminate_orphan_processes(self, present: set[DagFileInfo]): self.log.warning("Stopping processor for %s", file_name) stats.decr( "dag_processing.processes", - tags={"file_path": file.normalized_file_path_for_stats, "action": "stop"}, + tags=prune_dict( + { + "file_path": file.normalized_file_path_for_stats, + "action": "stop", + "team_name": bundle_to_team.get(file.bundle_name), + } + ), ) processor.kill(signal.SIGKILL) processor.logger_filehandle.close() @@ -1183,6 +1216,7 @@ def handle_parsing_result( run_duration = time.monotonic() - proc.start_time finish_time = timezone.utcnow() + team_name = self._get_team_name(file.bundle_name) next_stat = process_parse_results( run_duration=run_duration, finish_time=finish_time, @@ -1191,6 +1225,7 @@ def handle_parsing_result( parsing_result=proc.parsing_result, is_callback_only=is_callback_only, relative_fileloc=str(file.rel_path), + team_name=team_name, ) if proc.parsing_result is not None: @@ -1358,6 +1393,8 @@ def _create_process(self, dag_file: DagFileInfo) -> DagFileProcessorProcess: def _start_new_processes(self): """Start more processors if we have enough slots and files to process.""" + bundle_to_team = self._get_team_names({file.bundle_name for file in self._file_queue}) + while self._parallelism > len(self._processors) and self._file_queue: file, _ = self._file_queue.popitem(last=False) # Stop creating duplicate processor i.e. processor with the same filepath @@ -1367,7 +1404,13 @@ def _start_new_processes(self): processor = self._create_process(file) stats.incr( "dag_processing.processes", - tags={"file_path": file.normalized_file_path_for_stats, "action": "start"}, + tags=prune_dict( + { + "file_path": file.normalized_file_path_for_stats, + "action": "start", + "team_name": bundle_to_team.get(file.bundle_name), + } + ), ) self._processors[file] = processor @@ -1533,6 +1576,9 @@ def _kill_timed_out_processors(self): """Kill any file processors that timeout to defend against process hangs.""" now = time.monotonic() processors_to_remove = [] + + bundle_to_team = self._get_team_names({file.bundle_name for file in self._processors}) + for file, processor in self._processors.items(): duration = now - processor.start_time if duration > self.processor_timeout: @@ -1544,8 +1590,17 @@ def _kill_timed_out_processors(self): self.processor_timeout, ) file_path_tag = file.normalized_file_path_for_stats - stats.decr("dag_processing.processes", tags={"file_path": file_path_tag, "action": "timeout"}) - stats.incr("dag_processing.processor_timeouts", tags={"file_path": file_path_tag}) + team_name = bundle_to_team.get(file.bundle_name) + stats.decr( + "dag_processing.processes", + tags=prune_dict( + {"file_path": file_path_tag, "action": "timeout", "team_name": team_name} + ), + ) + stats.incr( + "dag_processing.processor_timeouts", + tags=prune_dict({"file_path": file_path_tag, "team_name": team_name}), + ) processor.kill(signal.SIGKILL) processors_to_remove.append(file) @@ -1606,10 +1661,18 @@ def max_runs_reached(self): def terminate(self): """Stop all running processors.""" + bundle_to_team = self._get_team_names({file.bundle_name for file in self._processors}) + for file, processor in self._processors.items(): stats.decr( "dag_processing.processes", - tags={"file_path": file.normalized_file_path_for_stats, "action": "terminate"}, + tags=prune_dict( + { + "file_path": file.normalized_file_path_for_stats, + "action": "terminate", + "team_name": bundle_to_team.get(file.bundle_name), + } + ), ) # SIGTERM, wait 5s, SIGKILL if still alive processor.kill(signal.SIGTERM, escalation_delay=5.0) @@ -1642,6 +1705,7 @@ def process_parse_results( *, is_callback_only: bool = False, relative_fileloc: str | None = None, + team_name: str | None = None, ) -> DagFileStat: """ Create a DagFileStat from parsing results and emit metrics. @@ -1655,7 +1719,7 @@ def process_parse_results( last_duration=run_duration, run_count=run_count, # Don't increment for callback-only processing ) - stats.incr("dag_processing.callback_only_count") + stats.incr("dag_processing.callback_only_count", tags=prune_dict({"team_name": team_name})) else: # Actual DAG parsing or import error stat = DagFileStat( @@ -1673,7 +1737,9 @@ def process_parse_results( stats.timing( "dag_processing.last_duration", stat.last_duration, - tags={"bundle_name": normalized_bundle, "file_name": file_name}, + tags=prune_dict( + {"bundle_name": normalized_bundle, "file_name": file_name, "team_name": team_name} + ), ) if parsing_result is None: diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index f7b3affd00830..251b7b629747d 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -38,6 +38,7 @@ from airflow.configuration import conf from airflow.dag_processing.bundles.base import BundleVersionLock from airflow.dag_processing.dagbag import BundleDagBag, DagBag +from airflow.models.dag import DagModel from airflow.sdk.exceptions import TaskNotFound from airflow.sdk.execution_time.comms import ( ConnectionResult, @@ -90,6 +91,7 @@ from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG from airflow.utils.dag_version_inflation_checker import check_dag_file_stability from airflow.utils.file import iter_airflow_imports +from airflow.utils.helpers import prune_dict from airflow.utils.state import TaskInstanceState if TYPE_CHECKING: @@ -386,7 +388,19 @@ def _execute_dag_callbacks(dagbag: DagBag, request: DagCallbackRequest, log: Fil callback(context) except Exception: log.exception("Callback failed", dag_id=request.dag_id) - stats.incr("dag.callback_exceptions", tags={"dag_id": request.dag_id}) + stats.incr( + "dag.callback_exceptions", + tags=prune_dict( + { + "dag_id": request.dag_id, + "team_name": ( + DagModel.get_team_name(request.dag_id) + if conf.getboolean("core", "multi_team") + else None + ), + } + ), + ) def _execute_task_callbacks(dagbag: DagBag, request: TaskCallbackRequest, log: FilteringBoundLogger) -> None: diff --git a/airflow-core/src/airflow/example_dags/example_asset_partition.py b/airflow-core/src/airflow/example_dags/example_asset_partition.py index 65a078d3c3f9e..1b3e2aa31b6d5 100644 --- a/airflow-core/src/airflow/example_dags/example_asset_partition.py +++ b/airflow-core/src/airflow/example_dags/example_asset_partition.py @@ -28,6 +28,7 @@ FanOutMapper, FixedKeyMapper, IdentityMapper, + MinimumCount, MonthWindow, PartitionAtRuntime, PartitionedAssetTimetable, @@ -39,7 +40,9 @@ StartOfMonthMapper, StartOfWeekMapper, StartOfYearMapper, + WaitForAll, WeekWindow, + Window, asset, task, ) @@ -306,6 +309,10 @@ def multi_region_player_stats(self, outlet_events): default_partition_mapper=RollupMapper( upstream_mapper=StartOfDayMapper(), window=DayWindow(), + # Explicit default wait policy: hold the run until all 24 hourly + # partitions arrive. Identical to omitting wait_policy entirely; shown + # here as the counterpart to the early-firing MinimumCount example below. + wait_policy=WaitForAll(), ), ), catchup=False, @@ -396,6 +403,13 @@ def train_model(): default_partition_mapper=FanOutMapper( upstream_mapper=StartOfWeekMapper(), window=WeekWindow(), + # Safety cap on how many downstream keys one upstream event may create. + # WeekWindow has exactly seven members, so max_downstream_keys=7 sits at + # the boundary and never blocks. A smaller value would skip queuing the + # runs for that event and record a "partition fan-out exceeded" audit log + # entry instead. Omitting it falls back to the global + # ``[scheduler] partition_mapper_max_downstream_keys`` (default 1000). + max_downstream_keys=7, ), ), catchup=False, @@ -413,6 +427,38 @@ def run_inference(dag_run=None): run_inference() +# --- Fan-out over a trailing window (Window.Direction.BACKWARD) -------------- +# ``daily_inference`` above fans the weekly artifact FORWARD: the seven days +# *starting* at the upstream key. The same artifact can drive a trailing window — +# the seven days *ending* at the key — e.g. to score the week leading up to a +# model release. Direction is the only difference between the two Dags. + +with DAG( + dag_id="trailing_week_inference", + schedule=PartitionedAssetTimetable( + assets=weekly_model_artifact, + default_partition_mapper=FanOutMapper( + upstream_mapper=StartOfWeekMapper(), + # BACKWARD yields the trailing period ending at the upstream key — the + # mirror of the default FORWARD that daily_inference uses. + window=WeekWindow(direction=Window.Direction.BACKWARD), + ), + ), + catchup=False, + tags=["example", "model", "inference"], +): + """Run inference over the trailing week: the seven days ending at the weekly key.""" + + @task + def run_trailing_inference(dag_run=None): + """Run inference for one daily partition in the trailing week.""" + if TYPE_CHECKING: + assert dag_run + print(dag_run.partition_key) + + run_trailing_inference() + + # --- Segment (categorical) rollup ------------------------------------------- # ``multi_region_player_stats`` (defined above) emits one partition per region # (``us``, ``eu``, ``apac``) from a single run. The Dag below holds a downstream @@ -449,3 +495,41 @@ def aggregate_all_regions(dag_run=None): print(f"All region partitions received. Partition: {dag_run.partition_key}") aggregate_all_regions() + + +# --- Segment rollup with an early-fire wait policy (MinimumCount) ------------ +# ``segment_region_stats_rollup`` above waits for all three regions (WaitForAll). +# This sibling fires as soon as any two of the three have arrived, tolerating one +# slow or missing region rather than holding the downstream run indefinitely. + +with DAG( + dag_id="segment_region_stats_early_rollup", + schedule=PartitionedAssetTimetable( + assets=Asset.ref(name="multi_region_player_stats"), + default_partition_mapper=RollupMapper( + upstream_mapper=FixedKeyMapper("all_regions"), + window=SegmentWindow(["us", "eu", "apac"]), + # Fire once at least two of the three declared regions have arrived. + # MinimumCount(-1) ("at most one missing") is equivalent for this window. + wait_policy=MinimumCount(2), + ), + ), + catchup=False, + tags=["example", "player-stats", "rollup", "segment"], +): + """ + Categorical rollup that fires early. + + Produces the cross-region summary once two of the three regions have arrived + instead of waiting for all of them — the early-firing counterpart to + ``segment_region_stats_rollup``. + """ + + @task + def aggregate_available_regions(dag_run=None): + """Produce the cross-region summary once the minimum region count is met.""" + if TYPE_CHECKING: + assert dag_run + print(f"Minimum region partitions received. Partition: {dag_run.partition_key}") + + aggregate_available_regions() diff --git a/airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py b/airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py new file mode 100644 index 0000000000000..3fd8778af87c5 --- /dev/null +++ b/airflow-core/src/airflow/example_dags/example_task_state_store_mapped.py @@ -0,0 +1,59 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Example DAG with mapped tasks to demonstrate task state store isolation per map_index.""" + +from __future__ import annotations + +import random +from datetime import datetime, timezone + +from airflow.sdk import DAG, task + +TABLES = ["orders", "customers", "products"] + +with DAG( + dag_id="example_task_state_store_mapped", + schedule=None, + start_date=datetime(2026, 1, 1), + catchup=False, + tags=["example", "task-state-store"], + doc_md=__doc__, +) as dag: + + @task + def get_tables() -> list[str]: + """Return the list of tables to process.""" + return TABLES + + @task + def process_table(table: str, task_state_store=None, ti=None) -> dict: + """Process one table — each mapped instance gets its own task state.""" + row_count = random.randint(100, 10000) + result = { + "table": table, + "map_index": ti.map_index, + "row_count": row_count, + "processed_at": datetime.now(tz=timezone.utc).isoformat(timespec="seconds"), + } + task_state_store.set("status", "complete") + task_state_store.set("result", result) + + print(f"[map_index={ti.map_index}] Processed {table}: {row_count} rows") + return result + + tables = get_tables() + process_table.expand(table=tables) diff --git a/airflow-core/src/airflow/exceptions.py b/airflow-core/src/airflow/exceptions.py index ff45daebc02e2..addd4f8f2a547 100644 --- a/airflow-core/src/airflow/exceptions.py +++ b/airflow-core/src/airflow/exceptions.py @@ -146,7 +146,12 @@ class DagNotPartitionedError(ValueError): class InvalidPartitionKeyError(ValueError): - """Raise when a partition_key value is empty or exceeds the maximum allowed length.""" + """ + Raise when a partition_key value is invalid. + + 1. empty or exceeds the maximum allowed length + 2. cannot be decoded to a partition_date by the timetable + """ class DagRunAlreadyExists(AirflowBadRequest): diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index ec5b3248178c2..93b155f2257ac 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -575,6 +575,10 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - starved_pools = {pool_name for pool_name, stats in pools.items() if stats["open"] <= 0} + pool_to_team_name: dict[str, str | None] = {} + if self._multi_team: + pool_to_team_name = Pool.get_name_to_team_name_mapping(list(pools.keys()), session=session) + # dag_id to # of running tasks and (dag_id, task_id) to # of running tasks. concurrency_map = ConcurrencyMap() concurrency_map.load(session=session) @@ -749,6 +753,20 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - starved_pools.add(pool_name) continue + if pool_team := pool_to_team_name.get(pool_name): + dag_team = dag_id_to_team_name.get(task_instance.dag_id) + if dag_team != pool_team: + self.log.debug( + "Not executing %s. Pool '%s' is assigned to team '%s' " + "but task's DAG belongs to team '%s'", + task_instance, + pool_name, + pool_team, + dag_team, + ) + starved_tasks.add((task_instance.dag_id, task_instance.task_id)) + continue + # Make sure to emit metrics if pool has no starving tasks pool_num_starving_tasks.setdefault(pool_name, 0) @@ -2593,7 +2611,12 @@ def _create_dag_runs_asset_triggered( creating_job_id=self.job.id, session=session, ) - stats.incr("asset.triggered_dagruns") + team_name = ( + self._get_team_names_for_dag_ids([dag.dag_id], session).get(dag.dag_id) + if self._multi_team + else None + ) + stats.incr("asset.triggered_dagruns", tags=prune_dict({"team_name": team_name})) dag_run.consumed_asset_events.extend(asset_events) self.log.info( "Created asset-triggered DagRun for '%s': run_id=%s, consumed %d asset events", diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index 97fd669204456..3612feb011fc3 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -110,7 +110,7 @@ from airflow.serialization.serialized_objects import DagSerialization from airflow.triggers.base import BaseEventTrigger, BaseTrigger, DiscrimatedTriggerEvent, TriggerEvent from airflow.triggers.shared_stream import SharedStreamManager -from airflow.utils.helpers import log_filename_template_renderer +from airflow.utils.helpers import log_filename_template_renderer, prune_dict from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.session import create_session, provide_session @@ -679,7 +679,7 @@ def heartbeat(self): perform_heartbeat(self.job, heartbeat_callback=self.heartbeat_callback, only_if_necessary=True) def heartbeat_callback(self, session: Session | None = None) -> None: - stats.incr("triggerer_heartbeat", 1, 1) + stats.incr("triggerer_heartbeat", 1, 1, tags=prune_dict({"team_name": self.team_name})) def load_triggers(self) -> None: """Assign triggers to this triggerer and update the runner with the IDs it should run.""" @@ -710,7 +710,7 @@ def handle_events(self): if entry.persist_seq is not None: self.persisted_event_seqs.append(entry.persist_seq) # Emit stat event - stats.incr("triggers.succeeded") + stats.incr("triggers.succeeded", tags=prune_dict({"team_name": self.team_name})) def on_trigger_event(self, trigger_id: int, event: TriggerEvent) -> None: """Record that a trigger fired an event.""" @@ -731,7 +731,7 @@ def handle_failed_triggers(self): trigger_id, exc = self.failed_triggers.popleft() self.on_trigger_failure(trigger_id=trigger_id, exc=exc) # Emit stat event - stats.incr("triggers.failed") + stats.incr("triggers.failed", tags=prune_dict({"team_name": self.team_name})) def on_trigger_failure(self, trigger_id: int, exc: list[str] | None) -> None: """Record that a trigger failed.""" @@ -753,7 +753,7 @@ def metric_tags(self) -> dict[str, str]: "TriggerRunnerSupervisor.metric_tags() requires a Job with a hostname; " "subclasses without a metadata-DB Job must override this method." ) - return {"hostname": hostname} + return prune_dict({"hostname": hostname, "team_name": self.team_name}) def emit_metrics(self): tags = self.metric_tags() diff --git a/airflow-core/src/airflow/models/callback.py b/airflow-core/src/airflow/models/callback.py index 78247be9ec018..1e8e758450ded 100644 --- a/airflow-core/src/airflow/models/callback.py +++ b/airflow-core/src/airflow/models/callback.py @@ -39,6 +39,7 @@ from airflow.models import Base from airflow.models.base import StringID from airflow.models.dagbundle import DagBundleModel +from airflow.utils.helpers import prune_dict from airflow.utils.sqlalchemy import ExtendedJSON, UtcDateTime from airflow.utils.state import CallbackState @@ -156,7 +157,7 @@ def __init__(self, priority_weight: int = 1, prefix: str = "", **kwargs): def queue(self, *, session: Session) -> None: self.state = CallbackState.QUEUED - def get_metric_info(self, status: CallbackState, result: Any) -> dict: + def get_metric_info(self, status: CallbackState, result: Any, team_name: str | None = None) -> dict: tags = {"result": result, **self.data} tags.pop("prefix", None) @@ -177,6 +178,10 @@ def get_metric_info(self, status: CallbackState, result: Any) -> dict: for k, v in tags.items() } + # team_name is omitted entirely when not in a multi-team deployment or when the + # callback's bundle is not mapped to a team. + tags = prune_dict({**tags, "team_name": team_name}) + prefix = self.data.get("prefix", "") name = f"{prefix}.callback_{status}" if prefix else f"callback_{status}" @@ -250,9 +255,12 @@ def handle_event(self, event: TriggerEvent, session: Session): if (status := event.payload.get(PAYLOAD_STATUS_KEY)) and status in (ACTIVE_STATES | TERMINAL_STATES): self.state = status if status in TERMINAL_STATES: + team_name: str | None = None + if self.bundle_name and conf.getboolean("core", "multi_team"): + team_name = DagBundleModel.get_team_name(self.bundle_name, session=session) self.trigger = None self.output = event.payload.get(PAYLOAD_BODY_KEY) - stats.incr(**self.get_metric_info(status, self.output)) + stats.incr(**self.get_metric_info(status, self.output, team_name=team_name)) session.add(self) else: diff --git a/airflow-core/src/airflow/models/dagbundle.py b/airflow-core/src/airflow/models/dagbundle.py index abd723a4e6766..f18c548a01930 100644 --- a/airflow-core/src/airflow/models/dagbundle.py +++ b/airflow-core/src/airflow/models/dagbundle.py @@ -30,6 +30,8 @@ from airflow.utils.sqlalchemy import UtcDateTime if TYPE_CHECKING: + from collections.abc import Collection + from sqlalchemy.orm import Session @@ -121,3 +123,18 @@ def get_team_name(bundle_name: str, *, session: Session = NEW_SESSION) -> str | return session.scalar( select(Team.name).join(DagBundleModel.teams).where(DagBundleModel.name == bundle_name) ) + + @staticmethod + @provide_session + def get_team_names( + bundle_names: Collection[str], *, session: Session = NEW_SESSION + ) -> dict[str, str | None]: + """Return a mapping of bundle name to team name (None for bundles not mapped to a team).""" + if not bundle_names: + return {} + rows = session.execute( + select(DagBundleModel.name, Team.name) + .join(DagBundleModel.teams) + .where(DagBundleModel.name.in_(bundle_names)) + ).all() + return {name: team_name for name, team_name in rows} diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index 80e279f149936..348d6ed9813e4 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -1072,6 +1072,16 @@ def _emit_dagrun_span(self, state: DagRunState): ctx = TraceContextTextMapPropagator().extract(self.context_carrier) span = trace.get_current_span(context=ctx) span_context = span.get_span_context() + + # Skip if the run was head-sampled out. Unlike the task spans, this guard is + # required (not just an optimization): the span below is forced to be a root span + # (context=context.Context()), so the configured sampler never sees the carrier's + # flag. A valid-but-unsampled carrier means head-sampled out; an invalid/empty + # carrier (legacy DagRun) recorded no decision, so it falls through and still + # emits — prior behavior we may want to reconsider. + if span_context.is_valid and not span_context.trace_flags.sampled: + return + with override_ids(span_context.trace_id, span_context.span_id): attributes: dict[str, str] = { "airflow.dag_id": str(self.dag_id), @@ -1089,13 +1099,14 @@ def _emit_dagrun_span(self, state: DagRunState): attributes["airflow.dag_run.logical_date"] = str(self.logical_date) if self.partition_key: attributes["airflow.dag_run.partition_key"] = str(self.partition_key) + # TODO: make the empty parent context optional. Default should be to - # nest the dag run span under the currently active parent span (by - # omitting `context` here); only use the empty `context.Context()` to - # force a root span when Airflow itself initiates the run (e.g. dag - # triggered via API, scheduler, or backfill). Today this forces a - # root span unconditionally. - # Tracked at https://github.com/apache/airflow/issues/67210 + # nest the dag run span under the currently active parent span (by + # omitting `context` here); only use the empty `context.Context()` to + # force a root span when Airflow itself initiates the run (e.g. dag + # triggered via API, scheduler, or backfill). Today this forces a + # root span unconditionally. + # Tracked at https://github.com/apache/airflow/issues/67210 span = tracer.start_span( name=f"dag_run.{self.dag_id}", start_time=int((self.queued_at or self.start_date or timezone.utcnow()).timestamp() * 1e9), @@ -1420,6 +1431,7 @@ def execute_dag_callbacks( TaskInstance as TIDataModel, TIRunContext, ) + from airflow.models.dag import DagModel from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance if relevant_ti: @@ -1463,7 +1475,19 @@ def execute_dag_callbacks( callback(context) except Exception: self.log.exception("Callback failed for %s", dag.dag_id) - stats.incr("dag.callback_exceptions", tags={"dag_id": dag.dag_id}) + stats.incr( + "dag.callback_exceptions", + tags=prune_dict( + { + "dag_id": dag.dag_id, + "team_name": ( + DagModel.get_team_name(dag.dag_id) + if airflow_conf.getboolean("core", "multi_team") + else None + ), + } + ), + ) def _get_ready_tis( self, diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 740596f9d69b6..5388b4b6b488b 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -1505,7 +1505,7 @@ def emit_state_change_metric(self, new_state: TaskInstanceState) -> None: stats.timing( f"task.{metric_name}", timing, - tags={"task_id": self.task_id, "dag_id": self.dag_id, "queue": self.queue}, + tags={**self.stats_tags, "queue": self.queue}, ) def clear_next_method_args(self) -> None: @@ -1524,6 +1524,14 @@ def register_asset_changes_in_db( *, session: Session = NEW_SESSION, ) -> None: + # Fast path: a task with no outlets and no outlet events has nothing to + # register. Returning early avoids the AssetModel lookup below (which + # would run with empty IN () clauses) and all downstream work. This is + # the common case -- most tasks declare no outlets -- and it sits on the + # task-success path that gates scheduling the next task. + if not task_outlets and not outlet_events: + return + from airflow.serialization.definitions.assets import ( SerializedAsset, SerializedAssetNameRef, diff --git a/airflow-core/src/airflow/timetables/base.py b/airflow-core/src/airflow/timetables/base.py index b07e5c98c9352..400b63f3d7090 100644 --- a/airflow-core/src/airflow/timetables/base.py +++ b/airflow-core/src/airflow/timetables/base.py @@ -292,6 +292,45 @@ def resolve_day_bound(self, day: datetime.date) -> DateTime: datetime.datetime(day.year, day.month, day.day, tzinfo=datetime.timezone.utc) ) + def resolve_partition_date(self, partition_key: str | None) -> datetime.datetime | None: + """ + Decode *partition_key* into the period-start datetime it represents. + + Returns the timezone-aware datetime that was used to format *partition_key* + when the timetable originally created the run, or ``None`` when no temporal + meaning can be derived from the key. ``None`` is returned without decoding + when *partition_key* is ``None``, when this timetable is not ``partitioned``, + or when it defers partition selection to runtime (``partitioned_at_runtime``). + + Partitioned timetables whose keys carry a temporal structure override + :meth:`_decode_partition_date`: + + - :class:`~airflow.timetables.trigger.CronPartitionTimetable` parses the + key with ``strptime`` using its ``key_format`` and localizes with its + timezone. + - :class:`~airflow.timetables.simple.PartitionedAssetTimetable` delegates + to each asset's partition mapper; when the mappers agree on the same + instant it is returned, otherwise ``None`` is returned. + + :param partition_key: The partition key string to decode, or ``None``. + :returns: The period-start datetime, or ``None`` if not resolvable. + :raises InvalidPartitionKeyError: When *partition_key* is syntactically + invalid for this timetable's key format (e.g. ``strptime`` fails). + """ + if partition_key is None or not self.partitioned or self.partitioned_at_runtime: + return None + return self._decode_partition_date(partition_key) + + def _decode_partition_date(self, partition_key: str) -> datetime.datetime | None: + """ + Decode a non-empty *partition_key* into its period-start datetime. + + Called by :meth:`resolve_partition_date` only after the partitioned-state + guards pass. The default returns ``None``; partitioned timetables whose + keys carry temporal structure override this. + """ + return None + @property def partition_mapper_info(self) -> list[PartitionMapperInfo]: """ diff --git a/airflow-core/src/airflow/timetables/simple.py b/airflow-core/src/airflow/timetables/simple.py index 496ba9da156ed..a8b499dbe50ba 100644 --- a/airflow-core/src/airflow/timetables/simple.py +++ b/airflow-core/src/airflow/timetables/simple.py @@ -17,11 +17,13 @@ from __future__ import annotations from contextlib import suppress +from datetime import datetime from typing import TYPE_CHECKING, Any, TypeAlias import structlog from airflow._shared.timezones import timezone +from airflow.exceptions import InvalidPartitionKeyError from airflow.partition_mappers.identity import IdentityMapper from airflow.serialization.definitions.assets import ( SerializedAsset, @@ -355,6 +357,46 @@ def partition_mapper_info(self) -> list[PartitionMapperInfo]: entries.append(PartitionMapperInfo(uri=s_asset_ref.uri, is_rollup=mapper.is_rollup)) return entries + def _decode_partition_date(self, partition_key: str) -> datetime | None: + """ + Decode *partition_key* into the period-start datetime shared by all asset mappers. + + Iterates every asset (and asset ref) reachable from the asset condition, asks + each mapper for the temporal anchor of *partition_key*, and returns it when all + temporal mappers agree. Returns ``None`` when no mapper is temporal or when the + mappers disagree — consistent with how the scheduler resolves ``partition_date`` + for asset-triggered runs. + """ + anchors: set[datetime] = set() + for unique_key, _ in self.asset_condition.iter_assets(): + mapper = self.get_partition_mapper(name=unique_key.name, uri=unique_key.uri) + try: + anchor = mapper.to_partition_date(partition_key) + except ValueError as exc: + raise InvalidPartitionKeyError( + f"Partition key {partition_key!r} is invalid for this timetable's mappers: {exc}" + ) from exc + if anchor is not None: + anchors.add(anchor) + for s_asset_ref in self.asset_condition.iter_asset_refs(): + if isinstance(s_asset_ref, SerializedAssetNameRef): + mapper = self.get_partition_mapper(name=s_asset_ref.name) + elif isinstance(s_asset_ref, SerializedAssetUriRef): + mapper = self.get_partition_mapper(uri=s_asset_ref.uri) + else: + continue + try: + anchor = mapper.to_partition_date(partition_key) + except ValueError as exc: + raise InvalidPartitionKeyError( + f"Partition key {partition_key!r} is invalid for this timetable's mappers: {exc}" + ) from exc + if anchor is not None: + anchors.add(anchor) + if len(anchors) == 1: + return anchors.pop() + return None + def serialize(self) -> dict[str, Any]: from airflow.serialization.serialized_objects import encode_asset_like diff --git a/airflow-core/src/airflow/timetables/trigger.py b/airflow-core/src/airflow/timetables/trigger.py index b3ac7cbf2e071..a1de05ac76e65 100644 --- a/airflow-core/src/airflow/timetables/trigger.py +++ b/airflow-core/src/airflow/timetables/trigger.py @@ -27,7 +27,8 @@ import structlog -from airflow._shared.timezones.timezone import coerce_datetime, parse_timezone, utcnow +from airflow._shared.timezones.timezone import coerce_datetime, make_aware, parse_timezone, utcnow +from airflow.exceptions import InvalidPartitionKeyError from airflow.timetables._cron import CronMixin from airflow.timetables._delta import DeltaMixin from airflow.timetables.base import DagRunInfo, DataInterval, Timetable @@ -518,6 +519,26 @@ def _format_key(self, partition_date: DateTime) -> str: # midnight partition keys as "...T00:00:00", not the prior UTC day's "...T16:00:00"). return partition_date.in_timezone(self._timezone).strftime(self._key_format) + def _decode_partition_date(self, partition_key: str) -> datetime.datetime: + """ + Decode *partition_key* back to the period-start datetime. + + Parses the key with ``strptime`` using this timetable's ``key_format`` + and localizes with the timetable's timezone, mirroring the forward + direction in :meth:`_format_key`. + + :raises InvalidPartitionKeyError: When *partition_key* does not match + the timetable's ``key_format``. + """ + try: + naive = datetime.datetime.strptime(partition_key, self._key_format) + except ValueError as exc: + raise InvalidPartitionKeyError( + f"Partition key {partition_key!r} does not match the timetable's " + f"key_format {self._key_format!r}: {exc}" + ) from exc + return make_aware(naive, self._timezone) + def next_dagrun_info_v2( self, *, diff --git a/airflow-core/src/airflow/triggers/base.py b/airflow-core/src/airflow/triggers/base.py index 6f6061e3366f8..f587fd8de3ff4 100644 --- a/airflow-core/src/airflow/triggers/base.py +++ b/airflow-core/src/airflow/triggers/base.py @@ -101,8 +101,8 @@ def task(self) -> Operator | None: return None @property - def task_instance(self) -> TaskInstance: - return self._task_instance + def task_instance(self) -> TaskInstance | None: + return getattr(self, "_task_instance", None) @task_instance.setter def task_instance(self, value: TaskInstance | None) -> None: diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/common.json b/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/common.json index aeb55c4a006e3..90debb241c24b 100644 --- a/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/common.json +++ b/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/common.json @@ -179,9 +179,12 @@ "note": { "add": "新增筆記", "dagRun": "Dag 執行筆記", + "edit": "編輯筆記", "label": "筆記", "placeholder": "新增筆記...", - "taskInstance": "任務實例筆記" + "preview": "預覽", + "taskInstance": "任務實例筆記", + "write": "撰寫" }, "overallStatus": "整體狀態", "partitionedDagRun_one": "分區的 Dag 執行", diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/dashboard.json b/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/dashboard.json index 510f3d8bff8e5..a3f6018219b42 100644 --- a/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/dashboard.json +++ b/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/dashboard.json @@ -1,4 +1,16 @@ { + "deadlines": { + "pending": { + "empty": "沒有即將到期的截止期限", + "title": "即將到期" + }, + "recentlyMissed": { + "empty": "沒有已逾期的截止期限", + "title": "已逾期的截止期限" + }, + "showMore": "顯示更多", + "title": "截止期限" + }, "deferredSlotsNotCounted": "未計入配額的延後任務：{{count}}", "deferredSlotsNotCountedTooltip": "長條圖中顯示的延後任務會計入資源池配額。長條圖下方顯示的延後任務來自不將延後任務計入配額的資源池。", "favorite": { diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/hitl.json b/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/hitl.json index 20410b3f7bd89..d529d3179879a 100644 --- a/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/hitl.json +++ b/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/hitl.json @@ -24,6 +24,20 @@ "success": "任務 {{taskId}} 回應成功", "title": "人類參與流程任務實例 - {{taskId}}" }, + "review": { + "detail": { + "selectRequiredAction": "選擇一個待回應的任務實例以查看詳細資訊" + }, + "list": { + "completedRequiredActions": "已回應的任務實例 ({{count}})", + "loadError": "無法載入待回應的任務實例", + "loadingActions": "正在載入待回應的任務實例...", + "pendingRequiredActions": "待回應的任務實例 ({{count}})" + }, + "openReviewDrawer": "開啟回應面板", + "pageLimitHint": "這裡僅顯示前幾項任務實例。請使用「$t(review.viewAll)」以查看全部。", + "viewAll": "查看所有待回應的任務實例" + }, "state": { "approvalReceived": "已核准", "approvalRequired": "需要核准", diff --git a/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/CalendarCell.tsx b/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/CalendarCell.tsx index ea91e34309bb0..6660194329e72 100644 --- a/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/CalendarCell.tsx +++ b/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/CalendarCell.tsx @@ -28,8 +28,8 @@ type Props = { | Record | string | { - actual: string | { _dark: string; _light: string }; - planned: string | { _dark: string; _light: string }; + primary: string | { _dark: string; _light: string }; + secondary: string | { _dark: string; _light: string }; }; readonly cellData: CalendarCellData | undefined; readonly index?: number; @@ -64,7 +64,7 @@ export const CalendarCell = ({ : []; const isMixedState = - typeof backgroundColor === "object" && "planned" in backgroundColor && "actual" in backgroundColor; + typeof backgroundColor === "object" && "secondary" in backgroundColor && "primary" in backgroundColor; const cellBox = isMixedState ? ( + | string + | { primary: Record | string; secondary: Record | string }; + +const LegendIcon = ({ color, cursor }: { readonly color: LegendColorType; readonly cursor?: string }) => { + const isMixedState = typeof color === "object" && "primary" in color && "secondary" in color; + + if (isMixedState) { + return ( + + + + + ); + } + + return ; +}; + export const CalendarLegend = ({ scale, vertical = false, viewMode }: Props) => { - const { t: translate } = useTranslation("dag"); + const { t: translate } = useTranslation(["dag", "common"]); const legendTitle = viewMode === "failed" ? translate("overview.buttons.failedRun_other") : translate("calendar.totalRuns"); @@ -54,7 +94,9 @@ export const CalendarLegend = ({ scale, vertical = false, viewMode }: Props) => {[...scale.legendItems].reverse().map(({ color, label }) => ( - + + + ))} @@ -70,7 +112,9 @@ export const CalendarLegend = ({ scale, vertical = false, viewMode }: Props) => {scale.legendItems.map(({ color, label }) => ( - + + + ))} @@ -83,42 +127,49 @@ export const CalendarLegend = ({ scale, vertical = false, viewMode }: Props) => + {viewMode === "total" && ( + <> + + + + {translate("common:states.success")} + + + + + + {translate("common:states.running")} + + + + )} + + + + + {translate("common:states.failed")} + + + {translate("common:states.planned")} + - - - - + : { _dark: "green.700", _light: "green.400" }, + secondary: PLANNED_COLOR, + }} + /> - {translate("calendar.legend.mixed")} + {translate("dag:calendar.legend.mixed")} diff --git a/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/CalendarTooltip.tsx b/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/CalendarTooltip.tsx index f5a3a50ecf96e..90ae1d8ab3681 100644 --- a/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/CalendarTooltip.tsx +++ b/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/CalendarTooltip.tsx @@ -32,6 +32,7 @@ type Props = { const stateColorMap = { failed: "failed.solid", planned: "stone.solid", + queued: "queued.solid", running: "running.solid", success: "success.solid", }; diff --git a/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/calendarUtils.test.ts b/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/calendarUtils.test.ts index 4cb86deec295c..c35225cebdc8f 100644 --- a/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/calendarUtils.test.ts +++ b/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/calendarUtils.test.ts @@ -27,6 +27,7 @@ const EMPTY_COLOR = { _dark: "gray.700", _light: "gray.100" }; const PLANNED_COLOR = { _dark: "stone.600", _light: "stone.500" }; const DEFAULT_TOTAL_COLOR = { _dark: "green.700", _light: "green.400" }; const DEFAULT_FAILED_COLOR = { _dark: "red.700", _light: "red.400" }; +const DEFAULT_RUNNING_COLOR = { _dark: "cyan.700", _light: "cyan.400" }; const EMPTY_COUNTS: RunCounts = { failed: 0, @@ -158,8 +159,8 @@ describe("createCalendarScale", () => { }); expect(scale.getColor({ ...EMPTY_COUNTS, planned: 1, success: 1, total: 2 })).toEqual({ - actual: DEFAULT_TOTAL_COLOR, - planned: PLANNED_COLOR, + primary: DEFAULT_TOTAL_COLOR, + secondary: PLANNED_COLOR, }); }); @@ -171,8 +172,57 @@ describe("createCalendarScale", () => { }); expect(scale.getColor({ ...EMPTY_COUNTS, queued: 1, success: 1, total: 2 })).toEqual({ - actual: DEFAULT_TOTAL_COLOR, - planned: PLANNED_COLOR, + primary: DEFAULT_TOTAL_COLOR, + secondary: PLANNED_COLOR, + }); + }); + + it("returns the failed color for a failed-only cell in total mode", () => { + const scale = createCalendarScale([run("failed", 1)], { + granularity: "hourly", + timezone: "UTC", + viewMode: "total", + }); + + expect(scale.getColor({ ...EMPTY_COUNTS, failed: 1, total: 1 })).toEqual(DEFAULT_FAILED_COLOR); + }); + + it("returns a mixed red and green color for failed and success runs in total mode", () => { + const scale = createCalendarScale([run("failed", 1), run("success", 1)], { + granularity: "hourly", + timezone: "UTC", + viewMode: "total", + }); + + expect(scale.getColor({ ...EMPTY_COUNTS, failed: 1, success: 1, total: 2 })).toEqual({ + primary: DEFAULT_FAILED_COLOR, + secondary: DEFAULT_TOTAL_COLOR, + }); + }); + + it("returns a mixed cyan and green color for running and success runs in total mode", () => { + const scale = createCalendarScale([run("running", 1), run("success", 1)], { + granularity: "hourly", + timezone: "UTC", + viewMode: "total", + }); + + expect(scale.getColor({ ...EMPTY_COUNTS, running: 1, success: 1, total: 2 })).toEqual({ + primary: DEFAULT_RUNNING_COLOR, + secondary: DEFAULT_TOTAL_COLOR, + }); + }); + + it("returns a mixed cyan and red color for running and failed runs in total mode", () => { + const scale = createCalendarScale([run("running", 1), run("failed", 1)], { + granularity: "hourly", + timezone: "UTC", + viewMode: "total", + }); + + expect(scale.getColor({ ...EMPTY_COUNTS, failed: 1, running: 1, total: 2 })).toEqual({ + primary: DEFAULT_FAILED_COLOR, + secondary: DEFAULT_RUNNING_COLOR, }); }); @@ -195,8 +245,8 @@ describe("createCalendarScale", () => { }); expect(scale.getColor({ ...EMPTY_COUNTS, failed: 1, planned: 1, total: 2 })).toEqual({ - actual: DEFAULT_FAILED_COLOR, - planned: PLANNED_COLOR, + primary: DEFAULT_FAILED_COLOR, + secondary: PLANNED_COLOR, }); }); @@ -218,8 +268,76 @@ describe("createCalendarScale", () => { }); expect(scale.getColor({ ...EMPTY_COUNTS, failed: 1, queued: 1, total: 2 })).toEqual({ - actual: DEFAULT_FAILED_COLOR, - planned: PLANNED_COLOR, + primary: DEFAULT_FAILED_COLOR, + secondary: PLANNED_COLOR, + }); + }); + + it("returns the correct gradient color when runs span across different dates", () => { + const scale = createCalendarScale( + [run("failed", 1, "2026-04-08T10:00:00Z"), run("failed", 5, "2026-04-09T10:00:00Z")], + { + granularity: "hourly", + timezone: "UTC", + viewMode: "total", + }, + ); + + const lowIntensityColor = { _dark: "red.900", _light: "red.200" }; + const highIntensityColor = { _dark: "red.300", _light: "red.800" }; + + expect(scale.getColor({ ...EMPTY_COUNTS, failed: 1, total: 1 })).toEqual(lowIntensityColor); + expect(scale.getColor({ ...EMPTY_COUNTS, failed: 5, total: 5 })).toEqual(highIntensityColor); + }); + + it("prioritizes failed over running over success when multiple actual states coexist with pending", () => { + const scale = createCalendarScale([run("planned", 1), run("failed", 1), run("success", 1)], { + granularity: "hourly", + timezone: "UTC", + viewMode: "total", + }); + + expect(scale.getColor({ ...EMPTY_COUNTS, failed: 1, planned: 1, success: 1, total: 3 })).toEqual({ + primary: DEFAULT_FAILED_COLOR, + secondary: PLANNED_COLOR, }); }); + + it("returns an empty scale when no data is provided", () => { + const scale = createCalendarScale([], { + granularity: "hourly", + timezone: "UTC", + viewMode: "total", + }); + + expect(scale.type).toBe("empty"); + expect(scale.getColor(EMPTY_COUNTS)).toEqual(EMPTY_COLOR); + expect(scale.legendItems).toEqual([{ color: EMPTY_COLOR, label: "0" }]); + }); + + it("prioritizes running and failed colors when failed, running, and success coexist without pending states", () => { + const scale = createCalendarScale([run("failed", 1), run("running", 1), run("success", 1)], { + granularity: "hourly", + timezone: "UTC", + viewMode: "total", + }); + + expect(scale.getColor({ ...EMPTY_COUNTS, failed: 1, running: 1, success: 1, total: 3 })).toEqual({ + primary: DEFAULT_FAILED_COLOR, + secondary: DEFAULT_RUNNING_COLOR, + }); + }); + + it("returns the correct gradient color for failed mode when failed runs span across different dates", () => { + const scale = createCalendarScale( + [run("failed", 1, "2026-04-08T10:00:00Z"), run("failed", 10, "2026-04-09T10:00:00Z")], + { granularity: "hourly", timezone: "UTC", viewMode: "failed" }, + ); + + const lowIntensityFailedColor = { _dark: "red.900", _light: "red.200" }; + const highIntensityFailedColor = { _dark: "red.300", _light: "red.800" }; + + expect(scale.getColor({ ...EMPTY_COUNTS, failed: 1, total: 1 })).toEqual(lowIntensityFailedColor); + expect(scale.getColor({ ...EMPTY_COUNTS, failed: 10, total: 10 })).toEqual(highIntensityFailedColor); + }); }); diff --git a/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/calendarUtils.ts b/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/calendarUtils.ts index 3f57eb35030c5..527ebc6389c26 100644 --- a/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/calendarUtils.ts +++ b/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/calendarUtils.ts @@ -40,6 +40,7 @@ dayjs.extend(tz); // Calendar color constants export const PLANNED_COLOR = { _dark: "stone.600", _light: "stone.500" }; const EMPTY_COLOR = { _dark: "gray.700", _light: "gray.100" }; +const RUNNING_COLOR = { _dark: "cyan.700", _light: "cyan.400" }; const TOTAL_COLOR_INTENSITIES = [ EMPTY_COLOR, // 0 @@ -244,6 +245,75 @@ type ScaleOptions = { viewMode: CalendarColorMode; }; +type ColorValue = string | { _dark: string; _light: string }; + +type ResolveColorParams = { + failedColor: ColorValue; + failedCount: number; + hasPending: boolean; + runningCount: number; + successColor: ColorValue; + successCount: number; +}; + +const resolveCellColor = ({ + failedColor, + failedCount, + hasPending, + runningCount, + successColor, + successCount, +}: ResolveColorParams): ColorValue | { primary: ColorValue; secondary: ColorValue } => { + const hasActual = failedCount > 0 || runningCount > 0 || successCount > 0; + + if (hasPending && hasActual) { + let primaryColor: ColorValue = EMPTY_COLOR; + + if (failedCount > 0) { + primaryColor = failedColor; + } else if (runningCount > 0) { + primaryColor = RUNNING_COLOR; + } else if (successCount > 0) { + primaryColor = successColor; + } + + return { + primary: primaryColor, + secondary: PLANNED_COLOR, + }; + } + + if (hasPending && !hasActual) { + return PLANNED_COLOR; + } + + if (hasActual) { + if (failedCount > 0 && runningCount > 0) { + return { primary: failedColor, secondary: RUNNING_COLOR }; + } + + if (failedCount > 0 && successCount > 0) { + return { primary: failedColor, secondary: successColor }; + } + + if (runningCount > 0 && successCount > 0) { + return { primary: RUNNING_COLOR, secondary: successColor }; + } + + if (failedCount > 0) { + return failedColor; + } + if (runningCount > 0) { + return RUNNING_COLOR; + } + if (successCount > 0) { + return successColor; + } + } + + return EMPTY_COLOR; +}; + export const createCalendarScale = ( data: Array, options: ScaleOptions, @@ -267,22 +337,23 @@ export const createCalendarScale = ( return { getColor: (counts: RunCounts) => { - const actualCount = getActualRunCount(counts, viewMode); - const hasPending = getPendingRunCount(counts) > 0; - const hasActual = actualCount > 0; - - if (hasPending && hasActual) { - return { - actual: singleColor, - planned: PLANNED_COLOR, - }; - } + const failedCount = counts.failed; + const runningCount = viewMode === "total" ? counts.running : 0; + const successCount = viewMode === "total" ? counts.success : 0; - if (hasPending && !hasActual) { - return PLANNED_COLOR; - } + const hasPending = getPendingRunCount(counts) > 0; - return actualCount === 0 ? EMPTY_COLOR : singleColor; + const failedColor = FAILURE_COLOR_INTENSITIES[2] ?? EMPTY_COLOR; + const successColor = TOTAL_COLOR_INTENSITIES[2] ?? EMPTY_COLOR; + + return resolveCellColor({ + failedColor, + failedCount, + hasPending, + runningCount, + successColor, + successCount, + }); }, legendItems: [ { color: EMPTY_COLOR, label: "0" }, @@ -312,54 +383,43 @@ export const createCalendarScale = ( | string | { _dark: string; _light: string } | { - actual: string | { _dark: string; _light: string }; - planned: string | { _dark: string; _light: string }; + primary: string | { _dark: string; _light: string }; + secondary: string | { _dark: string; _light: string }; } => { - const actualCount = getActualRunCount(counts, viewMode); - const hasPending = getPendingRunCount(counts) > 0; - const hasActual = actualCount > 0; + const failedCount = counts.failed; + const runningCount = viewMode === "total" ? counts.running : 0; + const successCount = viewMode === "total" ? counts.success : 0; - if (hasPending && hasActual) { - let actualColor = colorScheme[0] ?? EMPTY_COLOR; + const hasPending = getPendingRunCount(counts) > 0; + const getIntensityColor = (count: number, scheme: Array) => { + if (count === 0) { + return scheme[0] ?? EMPTY_COLOR; + } for (let index = uniqueThresholds.length - 1; index >= 1; index -= 1) { const threshold = uniqueThresholds[index]; - if (threshold !== undefined && actualCount >= threshold) { - actualColor = colorScheme[Math.min(index, colorScheme.length - 1)] ?? EMPTY_COLOR; - break; + if (threshold !== undefined && count >= threshold) { + return scheme[Math.min(index, scheme.length - 1)] ?? EMPTY_COLOR; } } - if (actualCount > 0 && actualColor === colorScheme[0]) { - actualColor = colorScheme[1] ?? EMPTY_COLOR; - } - - return { - actual: actualColor, - planned: PLANNED_COLOR, - }; - } - - if (hasPending && !hasActual) { - return PLANNED_COLOR; - } - - const targetCount = actualCount; - - if (targetCount === 0) { - return colorScheme[0] ?? EMPTY_COLOR; - } - - for (let index = uniqueThresholds.length - 1; index >= 1; index -= 1) { - const threshold = uniqueThresholds[index]; - - if (threshold !== undefined && targetCount >= threshold) { - return colorScheme[Math.min(index, colorScheme.length - 1)] ?? EMPTY_COLOR; - } - } + return scheme[1] ?? EMPTY_COLOR; + }; - return colorScheme[1] ?? EMPTY_COLOR; + const failedColor = + failedCount > 0 ? getIntensityColor(failedCount, FAILURE_COLOR_INTENSITIES) : EMPTY_COLOR; + const successColor = + successCount > 0 ? getIntensityColor(successCount, TOTAL_COLOR_INTENSITIES) : EMPTY_COLOR; + + return resolveCellColor({ + failedColor, + failedCount, + hasPending, + runningCount, + successColor, + successCount, + }); }; const legendItems: Array = []; diff --git a/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/types.ts b/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/types.ts index 8ef78a66af4d3..7e93c63773462 100644 --- a/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/types.ts +++ b/airflow-core/src/airflow/ui/src/pages/Dag/Calendar/types.ts @@ -69,8 +69,8 @@ export type CalendarScale = { | string | { _dark: string; _light: string } | { - actual: string | { _dark: string; _light: string }; - planned: string | { _dark: string; _light: string }; + primary: string | { _dark: string; _light: string }; + secondary: string | { _dark: string; _light: string }; }; readonly legendItems: Array; readonly type: CalendarScaleType; diff --git a/airflow-core/src/airflow/ui/src/pages/TaskInstances/TaskInstancesFilter.tsx b/airflow-core/src/airflow/ui/src/pages/TaskInstances/TaskInstancesFilter.tsx index 40b6c1c36fe6c..d3531224948d1 100644 --- a/airflow-core/src/airflow/ui/src/pages/TaskInstances/TaskInstancesFilter.tsx +++ b/airflow-core/src/airflow/ui/src/pages/TaskInstances/TaskInstancesFilter.tsx @@ -88,14 +88,12 @@ export const TaskInstancesFilter = () => { }); return ( - - - - + + ); }; diff --git a/airflow-core/src/airflow/ui/src/queryClient.ts b/airflow-core/src/airflow/ui/src/queryClient.ts index e465402c3088a..48b47499046df 100644 --- a/airflow-core/src/airflow/ui/src/queryClient.ts +++ b/airflow-core/src/airflow/ui/src/queryClient.ts @@ -29,6 +29,12 @@ if (OpenAPI.BASE.endsWith("/")) { OpenAPI.BASE = OpenAPI.BASE.slice(0, -1); } +// Encode path params as full URI components so values containing "/" (e.g. a variable key like +// "/foo") become "%2Ffoo" rather than a literal "//", which proxies may collapse. The generated +// client otherwise defaults to encodeURI, which leaves "/" untouched. +// The backend automatically decodes path params. +OpenAPI.ENCODE_PATH = encodeURIComponent; + const RETRY_COUNT = 3; const retryFunction = (failureCount: number, error: unknown) => { diff --git a/airflow-core/src/airflow/ui/tests/e2e/components/HITLReviewDrawer.ts b/airflow-core/src/airflow/ui/tests/e2e/components/HITLReviewDrawer.ts index 09990c6cd07e6..6cae368b33d09 100644 --- a/airflow-core/src/airflow/ui/tests/e2e/components/HITLReviewDrawer.ts +++ b/airflow-core/src/airflow/ui/tests/e2e/components/HITLReviewDrawer.ts @@ -25,8 +25,8 @@ export class HITLReviewDrawer { this.root = page.getByRole("dialog", { name: "Required Actions" }); } - public async expectOpenWith(dagId: string): Promise { + public async expectOpenWith(expectedText: string): Promise { await expect(this.root).toBeVisible(); - await expect(this.root).toContainText(dagId); + await expect(this.root).toContainText(expectedText); } } diff --git a/airflow-core/src/airflow/ui/tests/e2e/components/HITLReviewModal.ts b/airflow-core/src/airflow/ui/tests/e2e/components/HITLReviewModal.ts index b9975762c0e54..c86adeb3932bf 100644 --- a/airflow-core/src/airflow/ui/tests/e2e/components/HITLReviewModal.ts +++ b/airflow-core/src/airflow/ui/tests/e2e/components/HITLReviewModal.ts @@ -25,8 +25,8 @@ export class HITLReviewModal { this.root = page.getByRole("dialog", { name: "Required Actions" }); } - public async expectOpenWith(dagId: string): Promise { + public async expectOpenWith(expectedText: string): Promise { await expect(this.root).toBeVisible(); - await expect(this.root).toContainText(dagId, { timeout: 60_000 }); + await expect(this.root).toContainText(expectedText); } } diff --git a/airflow-core/src/airflow/ui/tests/e2e/utils/api/hitl.ts b/airflow-core/src/airflow/ui/tests/e2e/utils/api/hitl.ts index 159dd0b05c780..619fb1179ad19 100644 --- a/airflow-core/src/airflow/ui/tests/e2e/utils/api/hitl.ts +++ b/airflow-core/src/airflow/ui/tests/e2e/utils/api/hitl.ts @@ -200,7 +200,7 @@ export async function setupPendingHITLFlowViaAPI(source: RequestLike, dagId: str const response = await request.patch(`${baseUrl}/api/v2/dags/${dagId}`, { data: { is_paused: false } }); if (!response.ok()) { - throw new Error(`HITL response failed (${response.status()})`); + throw new Error(`Failed to unpause Dag ${dagId} (${response.status()})`); } const { dagRunId } = await apiTriggerDagRun(request, dagId); diff --git a/airflow-core/src/airflow/utils/email.py b/airflow-core/src/airflow/utils/email.py index 7376e14baf860..393841efa3dd3 100644 --- a/airflow-core/src/airflow/utils/email.py +++ b/airflow-core/src/airflow/utils/email.py @@ -257,8 +257,8 @@ def send_mime_email( log.debug("No user/password found for SMTP, so logging in with no authentication.") if not dryrun: - for attempt in range(1, smtp_retry_limit + 1): - log.info("Email alerting: attempt %s", str(attempt)) + for attempt in range(smtp_retry_limit + 1): + log.info("Email alerting: attempt %s", str(attempt + 1)) try: smtp_conn = _get_smtp_connection(smtp_host, smtp_port, smtp_timeout, smtp_ssl) except smtplib.SMTPServerDisconnected: diff --git a/airflow-core/tests/unit/api/common/test_trigger_dag.py b/airflow-core/tests/unit/api/common/test_trigger_dag.py index 49a385865f409..6fc48bd081411 100644 --- a/airflow-core/tests/unit/api/common/test_trigger_dag.py +++ b/airflow-core/tests/unit/api/common/test_trigger_dag.py @@ -16,12 +16,17 @@ # under the License. from __future__ import annotations +from datetime import datetime, timezone + import pytest from sqlalchemy import select from airflow.api.common.trigger_dag import trigger_dag +from airflow.exceptions import InvalidPartitionKeyError from airflow.models import DagModel from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.timetables.simple import PartitionAtRuntime +from airflow.timetables.trigger import CronPartitionTimetable from airflow.utils.types import DagRunTriggeredByType, DagRunType from tests_common.test_utils.db import ( @@ -83,6 +88,69 @@ def test_trigger_dag_with_custom_run_type(dag_maker, session): assert dag_run.run_type == DagRunType.ASSET_MATERIALIZATION +def test_trigger_dag_populates_partition_date_for_cron_partition_timetable(dag_maker, session): + """Manually triggering a CronPartitionTimetable Dag with a partition_key populates partition_date.""" + with dag_maker( + session=session, + dag_id="TEST_CRON_PARTITION_DAG", + schedule=CronPartitionTimetable("0 0 * * *", timezone="UTC"), + ): + EmptyOperator(task_id="mytask") + session.commit() + + dag_run = trigger_dag( + dag_id="TEST_CRON_PARTITION_DAG", + triggered_by=DagRunTriggeredByType.REST_API, + partition_key="2025-06-01T00:00:00", + session=session, + ) + + assert dag_run is not None + assert dag_run.partition_key == "2025-06-01T00:00:00" + assert dag_run.partition_date == datetime(2025, 6, 1, 0, 0, 0, tzinfo=timezone.utc) + + +def test_trigger_dag_raises_invalid_partition_key_for_cron_partition_timetable(dag_maker, session): + """A malformed partition_key for CronPartitionTimetable raises InvalidPartitionKeyError.""" + with dag_maker( + session=session, + dag_id="TEST_CRON_PARTITION_BAD_KEY", + schedule=CronPartitionTimetable("0 0 * * *", timezone="UTC"), + ): + EmptyOperator(task_id="mytask") + session.commit() + + with pytest.raises(InvalidPartitionKeyError, match="does not match"): + trigger_dag( + dag_id="TEST_CRON_PARTITION_BAD_KEY", + triggered_by=DagRunTriggeredByType.REST_API, + partition_key="not-a-valid-date", + session=session, + ) + + +def test_trigger_dag_partition_at_runtime_leaves_partition_date_none(dag_maker, session): + """PartitionAtRuntime Dags accept arbitrary keys; partition_date stays None.""" + with dag_maker( + session=session, + dag_id="TEST_PARTITION_AT_RUNTIME", + schedule=PartitionAtRuntime(), + ): + EmptyOperator(task_id="mytask") + session.commit() + + dag_run = trigger_dag( + dag_id="TEST_PARTITION_AT_RUNTIME", + triggered_by=DagRunTriggeredByType.REST_API, + partition_key="arbitrary-runtime-key", + session=session, + ) + + assert dag_run is not None + assert dag_run.partition_key == "arbitrary-runtime-key" + assert dag_run.partition_date is None + + def test_trigger_dag_operator_denied_when_only_manual_allowed(dag_maker, session): with dag_maker(session=session, dag_id="TEST_DAG_1", schedule="0 * * * *"): EmptyOperator(task_id="mytask") diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 6930039e1abbc..74920b11d51fd 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -1984,6 +1984,40 @@ def test_bulk_rejects_team_name_when_multi_team_is_disabled(self, test_client): expected_error_conn_ids = {err["input"]["connection_id"] for err in detail} assert sorted(expected_error_conn_ids) == ["test_conn_id_2", "test_conn_id_3"] + @conf_vars({("core", "multi_team"): "True"}) + def test_bulk_create_overwrite_preserves_unset_team_name(self, test_client, testing_team, session): + """A bulk create+overwrite that omits ``team_name`` must NOT reset an existing connection's + ``team_name`` to ``None`` (parity with the pools fix). Overwriting with only ``conn_type`` + previously clobbered every unset field via ``model_dump(by_alias=True)`` — silently nulling + the connection's multi-team ownership. ``exclude_unset=True`` preserves omitted fields. + """ + self.create_connection(team_name=testing_team.name) + before = session.scalar(select(Connection).where(Connection.conn_id == TEST_CONN_ID)) + assert before.team_name == testing_team.name + + response = test_client.patch( + "/connections", + json={ + "actions": [ + { + "action": "create", + "action_on_existence": "overwrite", + "entities": [{"connection_id": TEST_CONN_ID, "conn_type": "new_type"}], + } + ] + }, + ) + assert response.status_code == 200 + assert response.json()["create"]["success"] == [TEST_CONN_ID] + + session.expire_all() + after = session.scalar(select(Connection).where(Connection.conn_id == TEST_CONN_ID)) + assert after.conn_type == "new_type" # provided field is applied + assert after.team_name == testing_team.name, ( + "bulk overwrite that omitted team_name must preserve existing ownership, " + f"got team_name={after.team_name!r}" + ) + class TestPostConnectionExtraBackwardCompatibility(TestConnectionEndpoint): def test_post_should_accept_empty_string_as_extra(self, test_client, session): diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py index 0d0478f8d01e6..ad11f764d4839 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py @@ -46,7 +46,7 @@ from airflow.sdk.definitions.param import Param from airflow.settings import _configure_async_session from airflow.timetables.interval import CronDataIntervalTimetable -from airflow.timetables.simple import PartitionAtRuntime +from airflow.timetables.simple import PartitionAtRuntime, PartitionedAssetTimetable from airflow.timetables.trigger import CronPartitionTimetable from airflow.utils.session import provide_session from airflow.utils.state import DagRunState, State @@ -2874,7 +2874,7 @@ def test_should_respond_200_when_exactly_max_length_partition_key(self, dag_make partitioned_dag_id = "test_max_length_partition_key" with dag_maker( dag_id=partitioned_dag_id, - schedule=CronPartitionTimetable("0 * * * *", timezone="UTC"), + schedule=PartitionedAssetTimetable(assets=Asset("test")), start_date=START_DATE1, session=session, serialized=True, @@ -2889,6 +2889,86 @@ def test_should_respond_200_when_exactly_max_length_partition_key(self, dag_make ) assert response.status_code == 200 + @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") + def test_trigger_partitioned_dag_populates_partition_date(self, dag_maker, test_client, session): + """Triggering a CronPartitionTimetable Dag with a valid key populates partition_date on the run. + + Regression guard: before this fix partition_date was NULL for manually triggered runs even + when partition_key was supplied, making partition-date-based filtering (e.g. + ``airflow dags clear --partition-date-*``) silently skip those runs. + """ + partitioned_dag_id = "test_trigger_populates_partition_date" + with dag_maker( + dag_id=partitioned_dag_id, + schedule=CronPartitionTimetable("0 0 * * *", timezone="UTC"), + start_date=START_DATE1, + session=session, + serialized=True, + ): + EmptyOperator(task_id="task") + + session.commit() + + response = test_client.post( + f"/dags/{partitioned_dag_id}/dagRuns", + json={"logical_date": None, "partition_key": "2025-06-01T00:00:00"}, + ) + assert response.status_code == 200 + + dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == partitioned_dag_id)) + assert dag_run is not None + assert dag_run.partition_key == "2025-06-01T00:00:00" + assert dag_run.partition_date == datetime(2025, 6, 1, 0, 0, 0, tzinfo=timezone.utc) + + @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") + def test_trigger_partitioned_dag_invalid_key_returns_400(self, dag_maker, test_client, session): + """An invalid partition_key for a CronPartitionTimetable Dag must return HTTP 400.""" + partitioned_dag_id = "test_trigger_invalid_partition_key" + with dag_maker( + dag_id=partitioned_dag_id, + schedule=CronPartitionTimetable("0 0 * * *", timezone="UTC"), + start_date=START_DATE1, + session=session, + serialized=True, + ): + EmptyOperator(task_id="task") + + session.commit() + + response = test_client.post( + f"/dags/{partitioned_dag_id}/dagRuns", + json={"logical_date": None, "partition_key": "not-a-valid-date"}, + ) + assert response.status_code == 400 + + @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") + def test_trigger_partition_at_runtime_dag_leaves_partition_date_none( + self, dag_maker, test_client, session + ): + """PartitionAtRuntime Dag with an arbitrary key must produce partition_date=None.""" + runtime_dag_id = "test_trigger_partition_at_runtime_none_date" + with dag_maker( + dag_id=runtime_dag_id, + schedule=PartitionAtRuntime(), + start_date=START_DATE1, + session=session, + serialized=True, + ): + EmptyOperator(task_id="task") + + session.commit() + + response = test_client.post( + f"/dags/{runtime_dag_id}/dagRuns", + json={"logical_date": None, "partition_key": "arbitrary-key"}, + ) + assert response.status_code == 200 + + dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == runtime_dag_id)) + assert dag_run is not None + assert dag_run.partition_key == "arbitrary-key" + assert dag_run.partition_date is None + class TestResolveRunOnLatestVersion: @pytest.mark.parametrize("explicit_value", [True, False]) diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py index 6e17598f07ebc..a383713fe8866 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py @@ -1240,3 +1240,68 @@ def test_bulk_rejects_team_name_when_multi_team_is_disabled(self, test_client): expected_error_names = {err["input"]["name"] for err in detail} assert sorted(expected_error_names) == ["pool_2", "pool_3"] + + @conf_vars({("core", "multi_team"): "True"}) + def test_bulk_create_overwrite_preserves_unset_team_name(self, test_client, session): + """A bulk create+overwrite that omits ``team_name`` must NOT reset an existing pool's + ``team_name`` to ``None``. + + ``POOL1_NAME`` is owned by team ``test``. Overwriting it with a body that changes only + ``slots`` (no ``team_name``) previously clobbered every unset field back to its default via + ``model_dump()`` — silently nulling the pool's multi-team ownership. With + ``model_dump(exclude_unset=True)`` the omitted ``team_name`` keeps its current value. + """ + self.create_pools() + before = session.scalar(select(Pool).where(Pool.pool == POOL1_NAME)) + assert before.team_name == "test" + + response = test_client.patch( + "/pools", + json={ + "actions": [ + { + "action": "create", + "action_on_existence": "overwrite", + "entities": [{"name": POOL1_NAME, "slots": 99}], + } + ] + }, + ) + assert response.status_code == 200 + assert response.json()["create"]["success"] == [POOL1_NAME] + + session.expire_all() + after = session.scalar(select(Pool).where(Pool.pool == POOL1_NAME)) + assert after.slots == 99 # the field that WAS provided is applied + assert after.team_name == "test", ( + "bulk overwrite that omitted team_name must preserve existing ownership, " + f"got team_name={after.team_name!r}" + ) + + @conf_vars({("core", "multi_team"): "True"}) + def test_bulk_create_overwrite_applies_explicit_team_name(self, test_client, session): + """An explicitly-provided ``team_name`` on a bulk overwrite is still applied (the fix only + skips *omitted* fields, it must not skip fields the request actually set).""" + _create_team() + session.add(Team(name="other")) + session.commit() + _create_pools() # POOL1 owned by team "test" + + response = test_client.patch( + "/pools", + json={ + "actions": [ + { + "action": "create", + "action_on_existence": "overwrite", + "entities": [{"name": POOL1_NAME, "slots": 7, "team_name": "other"}], + } + ] + }, + ) + assert response.status_code == 200 + + session.expire_all() + after = session.scalar(select(Pool).where(Pool.pool == POOL1_NAME)) + assert after.team_name == "other" + assert after.slots == 7 diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py index 500e1ba4de4c3..5d3c37c9b521d 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py @@ -71,7 +71,7 @@ def test_trigger_dag_run_with_partition_key(self, client, session, dag_maker): dag_id = "test_trigger_dag_run_partition_key" run_id = "test_run_id" logical_date = timezone.datetime(2025, 2, 20) - partition_key = "2025-02-20" + partition_key = "2025-02-20T00:00:00" with dag_maker( dag_id=dag_id, @@ -126,6 +126,31 @@ def test_trigger_dag_run_partition_key_for_non_partitioned_dag(self, client, ses } } + def test_trigger_dag_run_invalid_partition_key(self, client, session, dag_maker): + """partition_key that the timetable cannot decode must return 400.""" + dag_id = "test_trigger_dag_run_invalid_partition_key" + run_id = "test_run_id_invalid_pk" + logical_date = timezone.datetime(2025, 2, 20) + + with dag_maker( + dag_id=dag_id, + schedule=CronPartitionTimetable("0 * * * *", timezone="UTC"), + session=session, + serialized=True, + ): + EmptyOperator(task_id="test_task") + session.commit() + + response = client.post( + f"/execution/dag-runs/{dag_id}/{run_id}", + json={"logical_date": logical_date.isoformat(), "partition_key": "not-a-date"}, + ) + + assert response.status_code == 400 + detail = response.json()["detail"] + assert detail["reason"] == "invalid_partition_key" + assert "does not match the timetable's key_format" in detail["message"] + def test_trigger_dag_run_dag_not_found(self, client): """Test that a DAG that does not exist cannot be triggered.""" dag_id = "dag_not_found" diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 395837c0e61ed..af44599776107 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -4244,3 +4244,19 @@ def test_emit_task_span_skips_if_no_dagrun_carrier(self): _emit_task_span(ti, TaskInstanceState.SUCCESS) assert len(self.exporter.get_finished_spans()) == 0 + + @pytest.mark.parametrize( + ("trace_flag", "expected_spans"), + [ + pytest.param("01", 1, id="sampled-carrier-emits"), + pytest.param("00", 0, id="unsampled-carrier-skips"), + ], + ) + def test_emit_task_span_honors_dagrun_carrier_sampling(self, trace_flag, expected_spans): + """A SAMPLED dag_run carrier (flag 01) emits the task span; an unsampled one (flag 00) is head-sampled out.""" + ti = self._make_ti() + ti.dag_run.context_carrier = { + "traceparent": f"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-{trace_flag}" + } + _emit_task_span(ti, TaskInstanceState.SUCCESS) + assert len(self.exporter.get_finished_spans()) == expected_spans diff --git a/airflow-core/tests/unit/assets/test_manager.py b/airflow-core/tests/unit/assets/test_manager.py index 687f4d178479d..40d9401b9ffc0 100644 --- a/airflow-core/tests/unit/assets/test_manager.py +++ b/airflow-core/tests/unit/assets/test_manager.py @@ -29,6 +29,7 @@ from sqlalchemy.orm import Session from airflow import settings +from airflow._shared.observability.metrics.base_stats_logger import StatsLogger from airflow.assets.manager import AssetManager from airflow.models.asset import ( AssetAliasModel, @@ -40,7 +41,9 @@ DagScheduleAssetReference, ) from airflow.models.dag import DAG, DagModel +from airflow.models.dagbundle import DagBundleModel from airflow.models.log import Log +from airflow.models.team import Team from airflow.partition_mappers.temporal import FanOutMapper, StartOfWeekMapper from airflow.partition_mappers.window import WeekWindow from airflow.providers.standard.operators.empty import EmptyOperator @@ -48,7 +51,11 @@ from airflow.sdk.definitions.timetables.assets import PartitionedAssetTimetable from tests_common.test_utils.config import conf_vars -from tests_common.test_utils.db import clear_db_apdr, clear_db_logs, clear_db_pakl +from tests_common.test_utils.db import ( + clear_db_apdr, + clear_db_logs, + clear_db_pakl, +) from unit.listeners import asset_listener pytestmark = pytest.mark.db_test @@ -670,6 +677,54 @@ def _make_asset_model( return model +class TestAssetMetricsTeamName: + @pytest.mark.parametrize( + ("multi_team", "expect_team_tag"), + [ + pytest.param("true", True, id="with_team"), + pytest.param("false", False, id="without_team"), + ], + ) + @mock.patch("airflow._shared.observability.metrics.stats._get_backend") + def test_asset_updates_respects_team_name( + self, mock_get_backend, multi_team, expect_team_tag, session, dag_maker + ): + mock_stats = mock.MagicMock(spec=StatsLogger) + mock_get_backend.return_value = mock_stats + + suffix = "with_team" if expect_team_tag else "without_team" + + team_name = f"team_asset_upd_{suffix}" + team = Team(name=team_name) + session.add(team) + session.flush() + + bundle_name = f"bundle_asset_upd_{suffix}" + bundle = DagBundleModel(name=bundle_name) + bundle.teams.append(team) + session.add(bundle) + session.flush() + + asset_name = f"metric_asset_{suffix}" + asset = Asset(uri=f"test://{asset_name}", name=asset_name, group="asset") + with dag_maker(dag_id=f"asset_dag_{suffix}", bundle_name=bundle_name, session=session): + EmptyOperator(task_id="task1", outlets=[asset]) + + ti = mock.MagicMock() + ti.dag_id = f"asset_dag_{suffix}" + ti.task_id = "task1" + ti.run_id = "run1" + ti.map_index = -1 + + with conf_vars({("core", "multi_team"): multi_team}): + AssetManager().register_asset_change(task_instance=ti, asset=asset, session=session) + + if expect_team_tag: + mock_stats.incr.assert_any_call("asset.updates", tags={"team_name": team_name}) + else: + mock_stats.incr.assert_any_call("asset.updates") + + class TestFilterDagsByTeam: @conf_vars({("core", "multi_team"): "false"}) def test_multi_team_disabled_returns_all_dags(self): diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 0f6c061fc43a4..4b55b6b2f1ee9 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -2732,6 +2732,15 @@ def test_refresh_dag_bundles_non_versioned_calls_update_bundle_state(self): mock_update.assert_called_once_with("mock_bundle", last_refreshed=mock.ANY, version=None) assert manager._bundle_versions["mock_bundle"] is None + def test_refresh_dag_bundles_clears_team_name_cache(self): + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_name_to_team_name = {"stale_bundle": "old_team"} + bundle = self._make_refresh_bundle(supports_versioning=False) + + self._refresh_with_mocked_state(manager, bundle, BundleState(last_refreshed=None, version=None)) + + assert manager._bundle_name_to_team_name == {} + def test_refresh_dag_bundles_versioned_version_changed_calls_update_bundle_state(self): """Versioned bundle with new version: update_bundle_state called with the new version.""" manager = DagFileProcessorManager(max_runs=1) @@ -2922,3 +2931,341 @@ def test_bundle_version_data_stored_after_refresh(self, session): assert manager._bundle_versions["mock_bundle"] == "newhash" assert manager._bundle_version_data["mock_bundle"] == test_data + + +class TestMultiTeamMetrics: + """Tests for team_name tag on dag processing metrics in multi-team mode.""" + + def mock_processor(self, start_time: float | None = None) -> DagFileProcessorProcess: + proc = MagicMock() + proc.wait.return_value = 0 + read_end, write_end = socketpair() + ret = DagFileProcessorProcess( + process_log=MagicMock(), + id=uuid7(), + pid=1234, + process=proc, + stdin=write_end, + logger_filehandle=MagicMock(), + client=MagicMock(), + bundle_name="testing", + dag_file_rel_path="test_dag.py", + ) + if start_time: + ret.start_time = start_time + ret._open_sockets.clear() + return ret + + @conf_vars({("core", "multi_team"): "true"}) + @mock.patch( + "airflow.dag_processing.manager.DagBundleModel.get_team_names", return_value={"testing": "team_alpha"} + ) + @mock.patch("airflow.dag_processing.manager.stats.gauge") + def test_log_file_processing_stats_includes_team_name(self, mock_gauge, mock_get_team_name): + manager = DagFileProcessorManager(max_runs=1) + dag_file = DagFileInfo( + bundle_name="testing", + rel_path=Path("dag_file.py"), + bundle_path=TEST_DAGS_FOLDER, + ) + manager._file_stats[dag_file] = DagFileStat( + last_finish_time=timezone.utcnow() - timedelta(seconds=5), + ) + + manager._log_file_processing_stats({"testing": {dag_file}}) + + mock_gauge.assert_any_call( + "dag_processing.last_run.seconds_ago", + mock.ANY, + tags={ + "file_path": "dag_file.py", + "bundle_name": "testing", + "file_name": "dag_file", + "team_name": "team_alpha", + }, + ) + + @conf_vars({("core", "multi_team"): "false"}) + @mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_names") + @mock.patch("airflow.dag_processing.manager.stats.gauge") + def test_log_file_processing_stats_omits_team_name_when_not_multi_team( + self, mock_gauge, mock_get_team_name + ): + manager = DagFileProcessorManager(max_runs=1) + dag_file = DagFileInfo( + bundle_name="testing", + rel_path=Path("dag_file.py"), + bundle_path=TEST_DAGS_FOLDER, + ) + manager._file_stats[dag_file] = DagFileStat( + last_finish_time=timezone.utcnow() - timedelta(seconds=5), + ) + + manager._log_file_processing_stats({"testing": {dag_file}}) + + mock_get_team_name.assert_not_called() + mock_gauge.assert_any_call( + "dag_processing.last_run.seconds_ago", + mock.ANY, + tags={ + "file_path": "dag_file.py", + "bundle_name": "testing", + "file_name": "dag_file", + }, + ) + + @conf_vars({("core", "multi_team"): "true"}) + @mock.patch( + "airflow.dag_processing.manager.DagBundleModel.get_team_names", return_value={"testing": "team_alpha"} + ) + @mock.patch("airflow.dag_processing.manager.stats.incr") + def test_start_new_processes_includes_team_name(self, mock_incr, mock_get_team_name): + manager = DagFileProcessorManager(max_runs=1) + dag_file = DagFileInfo( + bundle_name="testing", + rel_path=Path("dag_file.py"), + bundle_path=TEST_DAGS_FOLDER, + ) + manager._file_queue[dag_file] = None + manager._create_process = MagicMock(return_value=self.mock_processor()) + + manager._start_new_processes() + + mock_incr.assert_any_call( + "dag_processing.processes", + tags={"file_path": "dag_file.py", "action": "start", "team_name": "team_alpha"}, + ) + + @conf_vars({("core", "multi_team"): "true"}) + @mock.patch( + "airflow.dag_processing.manager.DagBundleModel.get_team_names", return_value={"testing": "team_alpha"} + ) + @mock.patch("airflow.dag_processing.manager.stats.incr") + @mock.patch("airflow.dag_processing.manager.stats.decr") + def test_kill_timed_out_processors_includes_team_name(self, mock_decr, mock_incr, mock_get_team_name): + manager = DagFileProcessorManager(max_runs=1, processor_timeout=5) + start_time = time.monotonic() - manager.processor_timeout - 1 + processor = self.mock_processor(start_time=start_time) + dag_file = DagFileInfo( + bundle_name="testing", rel_path=Path("dag_file.py"), bundle_path=TEST_DAGS_FOLDER + ) + manager._processors = {dag_file: processor} + + with mock.patch.object(type(processor), "kill"): + manager._kill_timed_out_processors() + + mock_decr.assert_called_once_with( + "dag_processing.processes", + tags={"file_path": "dag_file.py", "action": "timeout", "team_name": "team_alpha"}, + ) + mock_incr.assert_any_call( + "dag_processing.processor_timeouts", + tags={"file_path": "dag_file.py", "team_name": "team_alpha"}, + ) + + @conf_vars({("core", "multi_team"): "true"}) + @mock.patch( + "airflow.dag_processing.manager.DagBundleModel.get_team_names", return_value={"testing": "team_alpha"} + ) + @mock.patch("airflow.dag_processing.manager.stats.timing") + def test_process_parse_results_includes_team_name(self, mock_timing, mock_get_team_name): + from airflow.dag_processing.manager import process_parse_results + + result = DagFileParsingResult(fileloc="/tmp/dag.py", serialized_dags=[]) + process_parse_results( + run_duration=1.5, + finish_time=timezone.utcnow(), + run_count=0, + bundle_name="testing", + parsing_result=result, + relative_fileloc="dag.py", + team_name="team_alpha", + ) + + mock_timing.assert_called_once_with( + "dag_processing.last_duration", + 1.5, + tags={"bundle_name": "testing", "file_name": "dag", "team_name": "team_alpha"}, + ) + + @conf_vars({("core", "multi_team"): "false"}) + @mock.patch("airflow.dag_processing.manager.stats.timing") + def test_process_parse_results_omits_team_name_when_none(self, mock_timing): + from airflow.dag_processing.manager import process_parse_results + + result = DagFileParsingResult(fileloc="/tmp/dag.py", serialized_dags=[]) + process_parse_results( + run_duration=1.5, + finish_time=timezone.utcnow(), + run_count=0, + bundle_name="testing", + parsing_result=result, + relative_fileloc="dag.py", + team_name=None, + ) + + mock_timing.assert_called_once_with( + "dag_processing.last_duration", + 1.5, + tags={"bundle_name": "testing", "file_name": "dag"}, + ) + + @pytest.mark.parametrize( + ("multi_team", "team_name", "expected_tags"), + [ + pytest.param( + True, + "team_alpha", + {"file_path": "dag_file.py", "action": "stop", "team_name": "team_alpha"}, + id="with_team", + ), + pytest.param( + False, + None, + {"file_path": "dag_file.py", "action": "stop"}, + id="without_team", + ), + ], + ) + @mock.patch("airflow.dag_processing.manager.stats.decr") + def test_terminate_orphan_processes_includes_team_name( + self, mock_decr, multi_team, team_name, expected_tags + ): + manager = DagFileProcessorManager(max_runs=1) + manager._multi_team = multi_team + dag_file = DagFileInfo( + bundle_name="testing", rel_path=Path("dag_file.py"), bundle_path=TEST_DAGS_FOLDER + ) + processor = self.mock_processor() + manager._processors = {dag_file: processor} + + with ( + mock.patch( + "airflow.dag_processing.manager.DagBundleModel.get_team_names", + return_value={"testing": team_name}, + ), + mock.patch.object(type(processor), "kill"), + ): + # Empty "present" set means the file is orphaned, so its processor is stopped. + manager.terminate_orphan_processes(present=set()) + + mock_decr.assert_called_once_with("dag_processing.processes", tags=expected_tags) + + @pytest.mark.parametrize( + ("multi_team", "team_name", "expected_tags"), + [ + pytest.param( + True, + "team_alpha", + {"file_path": "dag_file.py", "action": "terminate", "team_name": "team_alpha"}, + id="with_team", + ), + pytest.param( + False, + None, + {"file_path": "dag_file.py", "action": "terminate"}, + id="without_team", + ), + ], + ) + @mock.patch("airflow.dag_processing.manager.stats.decr") + def test_terminate_includes_team_name(self, mock_decr, multi_team, team_name, expected_tags): + manager = DagFileProcessorManager(max_runs=1) + manager._multi_team = multi_team + dag_file = DagFileInfo( + bundle_name="testing", rel_path=Path("dag_file.py"), bundle_path=TEST_DAGS_FOLDER + ) + processor = self.mock_processor() + manager._processors = {dag_file: processor} + + with ( + mock.patch( + "airflow.dag_processing.manager.DagBundleModel.get_team_names", + return_value={"testing": team_name}, + ), + mock.patch.object(type(processor), "kill"), + ): + manager.terminate() + + mock_decr.assert_called_once_with("dag_processing.processes", tags=expected_tags) + + @pytest.mark.parametrize( + ("team_name", "expected_tags"), + [ + pytest.param("team_alpha", {"team_name": "team_alpha"}, id="with_team"), + pytest.param(None, {}, id="without_team"), + ], + ) + @mock.patch("airflow.dag_processing.manager.stats.incr") + def test_process_parse_results_callback_only_count_includes_team_name( + self, mock_incr, team_name, expected_tags + ): + from airflow.dag_processing.manager import process_parse_results + + process_parse_results( + run_duration=1.5, + finish_time=timezone.utcnow(), + run_count=0, + bundle_name="testing", + parsing_result=None, + is_callback_only=True, + relative_fileloc="dag.py", + team_name=team_name, + ) + + mock_incr.assert_called_once_with("dag_processing.callback_only_count", tags=expected_tags) + + @pytest.mark.parametrize( + ("multi_team", "team_name", "expected_tags"), + [ + pytest.param(True, "team_alpha", {"team_name": "team_alpha"}, id="with_team"), + pytest.param(False, None, {}, id="without_team"), + ], + ) + @mock.patch("airflow.dag_processing.manager.stats.incr") + def test_add_callback_to_queue_includes_team_name(self, mock_incr, multi_team, team_name, expected_tags): + manager = DagFileProcessorManager(max_runs=1) + manager._multi_team = multi_team + request = MagicMock(filepath="test_dag.py", bundle_name="testing", bundle_version=None) + bundle = MagicMock(path=TEST_DAGS_FOLDER) + + with ( + mock.patch.object(manager, "prepare_callback_bundle", return_value=bundle), + mock.patch.object(manager, "_add_files_to_queue"), + mock.patch( + "airflow.dag_processing.manager.DagBundleModel.get_team_names", + return_value={"testing": team_name}, + ), + ): + manager._add_callback_to_queue(request) + + mock_incr.assert_called_once_with("dag_processing.other_callback_count", tags=expected_tags) + + @mock.patch( + "airflow.dag_processing.manager.DagBundleModel.get_team_names", + return_value={"bundle_a": "team_alpha"}, + ) + def test_get_team_name_caches_lookup(self, mock_get_team_names): + manager = DagFileProcessorManager(max_runs=1) + manager._multi_team = True + + assert manager._get_team_name("bundle_a") == "team_alpha" + assert manager._get_team_name("bundle_a") == "team_alpha" + assert manager._get_team_name("bundle_a") == "team_alpha" + + mock_get_team_names.assert_called_once() + + @mock.patch( + "airflow.dag_processing.manager.DagBundleModel.get_team_names", + return_value={"bundle_a": "team_alpha", "bundle_b": "team_alpha"}, + ) + def test_get_team_names_batches_and_caches(self, mock_get_team_names): + manager = DagFileProcessorManager(max_runs=1) + manager._multi_team = True + + manager._get_team_names({"bundle_a", "bundle_b"}) + manager._get_team_names({"bundle_a", "bundle_b"}) + + # Two bundles resolved in a single batched query; the repeat call is served from cache. + mock_get_team_names.assert_called_once() + assert manager._bundle_name_to_team_name == {"bundle_a": "team_alpha", "bundle_b": "team_alpha"} diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 2e3d6940cc7eb..599c36d47f695 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -2149,3 +2149,80 @@ def test_handle_request_get_variable_masks_value_with_key(self, proc): "value": "super-secret-value", "type": "VariableResult", } + + +class TestMultiTeamCallbackMetrics: + """Tests for team_name tag on dag.callback_exceptions in multi-team mode.""" + + @conf_vars({("core", "multi_team"): "true"}) + @patch("airflow.dag_processing.processor.DagModel.get_team_name", return_value="team_alpha") + @patch("airflow.dag_processing.processor.stats.incr") + def test_callback_exception_includes_team_name(self, mock_incr, mock_get_team_name, spy_agency): + def failing_callback(context): + raise RuntimeError("boom") + + with DAG(dag_id="test_dag", on_failure_callback=failing_callback) as dag: + BaseOperator(task_id="test_task") + + @spy_agency.spy_for(DagBag.collect_dags, owner=DagBag) + def fake_collect_dags(self, *args, **kwargs): + self.dags[dag.dag_id] = dag + + dagbag = DagBag() + dagbag.collect_dags() + + request = DagCallbackRequest( + filepath="test.py", + dag_id="test_dag", + run_id="test_run", + bundle_name="testing", + bundle_version=None, + is_failure_callback=True, + msg="Test failure", + ) + + log = structlog.get_logger() + _execute_dag_callbacks(dagbag, request, log) + + mock_incr.assert_called_once_with( + "dag.callback_exceptions", + tags={"dag_id": "test_dag", "team_name": "team_alpha"}, + ) + + @conf_vars({("core", "multi_team"): "false"}) + @patch("airflow.dag_processing.processor.DagModel.get_team_name") + @patch("airflow.dag_processing.processor.stats.incr") + def test_callback_exception_omits_team_name_when_not_multi_team( + self, mock_incr, mock_get_team_name, spy_agency + ): + def failing_callback(context): + raise RuntimeError("boom") + + with DAG(dag_id="test_dag", on_failure_callback=failing_callback) as dag: + BaseOperator(task_id="test_task") + + @spy_agency.spy_for(DagBag.collect_dags, owner=DagBag) + def fake_collect_dags(self, *args, **kwargs): + self.dags[dag.dag_id] = dag + + dagbag = DagBag() + dagbag.collect_dags() + + request = DagCallbackRequest( + filepath="test.py", + dag_id="test_dag", + run_id="test_run", + bundle_name="testing", + bundle_version=None, + is_failure_callback=True, + msg="Test failure", + ) + + log = structlog.get_logger() + _execute_dag_callbacks(dagbag, request, log) + + mock_get_team_name.assert_not_called() + mock_incr.assert_called_once_with( + "dag.callback_exceptions", + tags={"dag_id": "test_dag"}, + ) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 540bdc46054a5..b312f9fd1b0a8 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -1425,6 +1425,70 @@ def test_find_executable_task_instances_pool(self, dag_maker): assert tis[3].key in res_keys session.rollback() + @conf_vars({("core", "multi_team"): "true"}) + def test_find_executable_task_instances_pool_team_enforcement(self, dag_maker, session): + """Tasks using a pool owned by another team are not scheduled.""" + clear_db_teams() + clear_db_dag_bundles() + + team_a = Team(name="team_a") + team_b = Team(name="team_b") + session.add_all([team_a, team_b]) + session.flush() + + bundle_a = DagBundleModel(name="bundle_a") + bundle_a.teams.append(team_a) + bundle_b = DagBundleModel(name="bundle_b") + bundle_b.teams.append(team_b) + session.add_all([bundle_a, bundle_b]) + session.flush() + + # Pool owned by team_a + pool_a = Pool(pool="pool_a", slots=10, include_deferred=False, team_name="team_a") + # Shared pool (no team) + pool_shared = Pool(pool="pool_shared", slots=10, include_deferred=False) + session.add_all([pool_a, pool_shared]) + session.flush() + + # DAG in team_a using pool_a (allowed) + with dag_maker(dag_id="dag_a", bundle_name="bundle_a", session=session): + EmptyOperator(task_id="task_a", pool="pool_a") + dr_a = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) + ti_a = dr_a.get_task_instance("task_a", session=session) + ti_a.state = State.SCHEDULED + session.merge(ti_a) + + # DAG in team_b using pool_a (should be blocked) + with dag_maker(dag_id="dag_b_cross", bundle_name="bundle_b", session=session): + EmptyOperator(task_id="task_cross", pool="pool_a") + dr_b = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) + ti_b = dr_b.get_task_instance("task_cross", session=session) + ti_b.state = State.SCHEDULED + session.merge(ti_b) + + # DAG in team_b using shared pool (allowed) + with dag_maker(dag_id="dag_b_shared", bundle_name="bundle_b", session=session): + EmptyOperator(task_id="task_shared", pool="pool_shared") + dr_b2 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) + ti_b2 = dr_b2.get_task_instance("task_shared", session=session) + ti_b2.state = State.SCHEDULED + session.merge(ti_b2) + session.flush() + + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job) + + res = self.job_runner._executable_task_instances_to_queued(max_tis=32, session=session) + queued_keys = {ti.key for ti in res} + + # team_a task using its own pool: allowed + assert ti_a.key in queued_keys + # team_b task using team_a's pool: blocked + assert ti_b.key not in queued_keys + # team_b task using shared pool: allowed + assert ti_b2.key in queued_keys + session.rollback() + @pytest.mark.parametrize( ("state", "total_executed_ti"), [ @@ -5603,6 +5667,67 @@ def dict_from_obj(obj): assert created_run.data_interval_end is None assert created_run.creating_job_id == scheduler_job.id + @pytest.mark.parametrize( + ("multi_team", "expect_team_tag"), + [ + pytest.param("true", True, id="with_team"), + pytest.param("false", False, id="without_team"), + ], + ) + @mock.patch("airflow._shared.observability.metrics.stats._get_backend") + def test_asset_triggered_dagruns_respects_team_name( + self, mock_get_backend, multi_team, expect_team_tag, session, dag_maker + ): + mock_stats = mock.MagicMock(spec=StatsLogger) + mock_get_backend.return_value = mock_stats + + suffix = "with_team" if expect_team_tag else "without_team" + + team_name = f"team_asset_trig_{suffix}" + team = Team(name=team_name) + session.add(team) + session.flush() + + bundle_name = f"bundle_asset_trig_{suffix}" + bundle = DagBundleModel(name=bundle_name) + bundle.teams.append(team) + session.add(bundle) + session.commit() + + asset_name = f"test_team_asset_{suffix}" + asset = Asset(uri=f"test://{asset_name}", name=asset_name, group="test_group") + with dag_maker(dag_id=f"producer_{suffix}", bundle_name=bundle_name, session=session): + BashOperator(task_id="task", bash_command="echo 1", outlets=[asset]) + dr = dag_maker.create_dagrun() + + asset_id = session.scalar(select(AssetModel.id).where(AssetModel.uri == asset.uri)) + event = AssetEvent( + asset_id=asset_id, + source_task_id="task", + source_dag_id=dr.dag_id, + source_run_id=dr.run_id, + source_map_index=-1, + ) + session.add(event) + + with dag_maker( + dag_id=f"consumer_{suffix}", schedule=[asset], bundle_name=bundle_name, session=session + ): + pass + + session.add(AssetDagRunQueue(asset_id=asset_id, target_dag_id=f"consumer_{suffix}")) + session.flush() + + with conf_vars({("core", "multi_team"): multi_team}): + scheduler_job = Job() + self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[self.null_exec]) + self.job_runner._create_dagruns_for_dags(session, session) + + if expect_team_tag: + mock_stats.incr.assert_any_call("asset.triggered_dagruns", tags={"team_name": team_name}) + else: + mock_stats.incr.assert_any_call("asset.triggered_dagruns") + @pytest.mark.need_serialized_dag @pytest.mark.parametrize( ("disable", "enable"), diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index 923779df0034c..2a875fbfbf6c1 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -2460,6 +2460,83 @@ def test_handle_events_does_not_confirm_seq_when_persist_fails(jobless_superviso assert list(jobless_supervisor.persisted_event_seqs) == [] +@pytest.mark.parametrize( + ("team_name", "expected_tags"), + [ + pytest.param("team_alpha", {"team_name": "team_alpha"}, id="with_team"), + pytest.param(None, {}, id="without_team"), + ], +) +def test_handle_events_emits_team_name(jobless_supervisor, team_name, expected_tags): + """triggers.succeeded carries the triggerer's team_name (omitted when the triggerer has none).""" + jobless_supervisor.team_name = team_name + jobless_supervisor.events.append(TriggerEventEntry(1, TriggerEvent(True), 7)) + + with ( + mock.patch.object(TriggerRunnerSupervisor, "on_trigger_event", autospec=True), + mock.patch("airflow.jobs.triggerer_job_runner.stats.incr") as mock_incr, + ): + jobless_supervisor.handle_events() + + mock_incr.assert_called_once_with("triggers.succeeded", tags=expected_tags) + + +@pytest.mark.parametrize( + ("team_name", "expected_tags"), + [ + pytest.param("team_alpha", {"team_name": "team_alpha"}, id="with_team"), + pytest.param(None, {}, id="without_team"), + ], +) +def test_handle_failed_triggers_emits_team_name(jobless_supervisor, team_name, expected_tags): + """triggers.failed carries the triggerer's team_name (omitted when the triggerer has none).""" + jobless_supervisor.team_name = team_name + jobless_supervisor.failed_triggers.append((1, None)) + + with ( + mock.patch.object(TriggerRunnerSupervisor, "on_trigger_failure", autospec=True), + mock.patch("airflow.jobs.triggerer_job_runner.stats.incr") as mock_incr, + ): + jobless_supervisor.handle_failed_triggers() + + mock_incr.assert_called_once_with("triggers.failed", tags=expected_tags) + + +@pytest.mark.parametrize( + ("team_name", "expected_tags"), + [ + pytest.param("team_alpha", {"team_name": "team_alpha"}, id="with_team"), + pytest.param(None, {}, id="without_team"), + ], +) +def test_heartbeat_callback_emits_team_name(jobless_supervisor, team_name, expected_tags): + jobless_supervisor.team_name = team_name + + with mock.patch("airflow.jobs.triggerer_job_runner.stats.incr") as mock_incr: + jobless_supervisor.heartbeat_callback() + + mock_incr.assert_called_once_with("triggerer_heartbeat", 1, 1, tags=expected_tags) + + +@pytest.mark.parametrize( + ("team_name", "expected_extra"), + [ + pytest.param("team_alpha", {"team_name": "team_alpha"}, id="with_team"), + pytest.param(None, {}, id="without_team"), + ], +) +def test_emit_metrics_includes_team_name(supervisor_builder, mocker, team_name, expected_extra): + supervisor = supervisor_builder() + supervisor.team_name = team_name + gauge = mocker.patch("airflow.jobs.triggerer_job_runner.stats.gauge") + + supervisor.emit_metrics() + + expected_tags = {"hostname": supervisor.job.hostname, **expected_extra} + gauge.assert_any_call("triggers.running", mock.ANY, tags=expected_tags) + gauge.assert_any_call("triggerer.capacity_left", mock.ANY, tags=expected_tags) + + def test_state_sync_carries_and_drains_persist_confirmations(jobless_supervisor): """The state-sync response carries pending confirmations once, then None when there are none.""" jobless_supervisor.persisted_event_seqs.extend([3, 9]) diff --git a/airflow-core/tests/unit/models/test_callback.py b/airflow-core/tests/unit/models/test_callback.py index b5296979ed417..b2b13582710fc 100644 --- a/airflow-core/tests/unit/models/test_callback.py +++ b/airflow-core/tests/unit/models/test_callback.py @@ -143,6 +143,21 @@ def test_get_metric_info_dict_values_are_stringified(self): # insertion order collapse to one metric series (no needless cardinality split). assert metric_info["tags"]["result"] == '{"code": 0, "output": [1, 2]}' + @pytest.mark.parametrize( + ("team_name", "expect_tag"), + [ + pytest.param("team_alpha", True, id="with_team"), + pytest.param(None, False, id="without_team"), + ], + ) + def test_get_metric_info_includes_team_name(self, team_name, expect_tag): + callback = TriggererCallback(TEST_ASYNC_CALLBACK, prefix="deadline_alerts", dag_id=TEST_DAG_ID) + metric_info = callback.get_metric_info(CallbackState.SUCCESS, "0", team_name=team_name) + if expect_tag: + assert metric_info["tags"]["team_name"] == team_name + else: + assert "team_name" not in metric_info["tags"] + class TestTriggererCallback: def test_polymorphic_serde(self, session): @@ -249,6 +264,36 @@ def test_handle_event(self, session, event, terminal_state): assert callback.trigger is None assert callback.output == event.payload[PAYLOAD_BODY_KEY] + @pytest.mark.parametrize( + ("multi_team", "team_name", "expect_tag"), + [ + pytest.param("true", "team_alpha", True, id="with_team"), + pytest.param("false", None, False, id="without_team"), + ], + ) + @patch("airflow.models.callback.stats.incr") + def test_handle_event_emits_team_name(self, mock_incr, multi_team, team_name, expect_tag, session): + """On a terminal event, callback_{status} carries team_name resolved from the bundle.""" + callback = TriggererCallback(TEST_ASYNC_CALLBACK, dag_id=TEST_DAG_ID) + callback.bundle_name = "test_bundle" + event = TriggerEvent({PAYLOAD_STATUS_KEY: CallbackState.SUCCESS, PAYLOAD_BODY_KEY: "0"}) + + with ( + conf_vars({("core", "multi_team"): multi_team}), + patch( + "airflow.models.callback.DagBundleModel.get_team_name", return_value=team_name + ) as mock_get_team_name, + ): + callback.handle_event(event, session) + + mock_incr.assert_called_once() + _, kwargs = mock_incr.call_args + if expect_tag: + assert kwargs["tags"]["team_name"] == team_name + else: + mock_get_team_name.assert_not_called() + assert "team_name" not in kwargs["tags"] + class TestExecutorCallback: def test_polymorphic_serde(self, session): diff --git a/airflow-core/tests/unit/models/test_dagbundle.py b/airflow-core/tests/unit/models/test_dagbundle.py index 84223e34a9b31..0c5c5def9d735 100644 --- a/airflow-core/tests/unit/models/test_dagbundle.py +++ b/airflow-core/tests/unit/models/test_dagbundle.py @@ -58,3 +58,20 @@ def test_get_team_name_no_team(self, session: Session): def test_get_team_name_unknown_bundle(self, session: Session): assert DagBundleModel.get_team_name("does_not_exist", session=session) is None + + def test_get_team_names(self, testing_team: Team, session: Session): + mapped = DagBundleModel(name="mapped_bundle") + mapped.teams.append(testing_team) + unmapped = DagBundleModel(name="unmapped_bundle") + session.add_all([mapped, unmapped]) + session.flush() + + result = DagBundleModel.get_team_names( + ["mapped_bundle", "unmapped_bundle", "does_not_exist"], session=session + ) + + # Only bundles actually mapped to a team are returned; callers treat absent keys as None. + assert result == {"mapped_bundle": "testing"} + + def test_get_team_names_empty(self, session: Session): + assert DagBundleModel.get_team_names([], session=session) == {} diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index a8d384d948099..1225cd07eced3 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -1535,7 +1535,6 @@ def test_dagrun_deadline_variable_interval_stable(self, _, mock_get, session, de @mock.patch.object(Deadline, "prune_deadlines") def test_dagrun_deadline_variable_interval_missing_variable_fails(self, _, session, deadline_test_dag): - mock_err = mock.Mock() mock_err.error.value = "MISSING_DEADLINE" mock_err.detail = "missing deadline" @@ -4118,16 +4117,60 @@ def on_failure(context): assert "ti" not in context_received assert context_received["run_id"] == dr.run_id + @pytest.mark.parametrize( + ("multi_team", "team_name", "expected_tags"), + [ + pytest.param( + "true", "team_alpha", {"dag_id": "test_dag", "team_name": "team_alpha"}, id="with_team" + ), + pytest.param("false", None, {"dag_id": "test_dag"}, id="without_team"), + ], + ) + @mock.patch("airflow._shared.observability.metrics.stats.incr") + def test_callback_exception_team_name_tag( + self, mock_incr, multi_team, team_name, expected_tags, dag_maker, session + ): + def failing_callback(context): + raise RuntimeError("boom") + + with dag_maker("test_dag", session=session, on_failure_callback=failing_callback) as dag: + BashOperator(task_id="test_task", bash_command="echo 1") + + dr = dag_maker.create_dagrun() + dag.on_failure_callback = failing_callback + dag.has_on_failure_callback = True + + with ( + conf_vars({("core", "multi_team"): multi_team}), + mock.patch("airflow.models.dag.DagModel.get_team_name", return_value=team_name), + ): + dr.execute_dag_callbacks(dag, success=False) + + mock_incr.assert_any_call("dag.callback_exceptions", tags=expected_tags) + class TestDagRunTracing: """Tests for DagRun OpenTelemetry span behavior.""" @pytest.fixture(autouse=True) def sdk_tracer_provider(self): - """Patch the module-level tracer with one backed by a real SDK provider so spans have valid IDs.""" + """Patch the module-level tracer with one backed by a real SDK provider so spans have valid IDs. + + Also patch the provider that ``new_dagrun_trace_carrier`` consults so the + head-sampling decision is made by a real SDK sampler (default + parentbased_always_on -> SAMPLED) rather than the no-op ProxyTracerProvider + that is otherwise active in the test process (which would honestly produce + an unsampled carrier and suppress emission). + """ provider = TracerProvider() real_tracer = provider.get_tracer("airflow.models.dagrun") - with mock.patch("airflow.models.dagrun.tracer", real_tracer): + with ( + mock.patch("airflow.models.dagrun.tracer", real_tracer), + mock.patch( + "airflow._shared.observability.traces.trace.get_tracer_provider", + return_value=provider, + ), + ): yield def test_context_carrier_set_on_init(self, dag_maker): @@ -4294,6 +4337,34 @@ def test_emit_dagrun_span_with_none_or_empty_carrier(self, dag_maker, session, c else: assert len(spans) == 0 + @pytest.mark.parametrize( + ("dag_id", "trace_flag", "expected_spans"), + [ + pytest.param("test_tracing_sampled", "01", 1, id="sampled-carrier-emits"), + pytest.param("test_tracing_unsampled", "00", 0, id="unsampled-carrier-skips"), + ], + ) + def test_emit_dagrun_span_honors_carrier_sampling( + self, dag_id, trace_flag, expected_spans, dag_maker, session + ): + """A SAMPLED carrier (flag 01) emits the dag_run span; an unsampled carrier (flag 00) is head-sampled out.""" + in_mem_exporter = InMemorySpanExporter() + provider = TracerProvider(id_generator=OverrideableRandomIdGenerator()) + provider.add_span_processor(SimpleSpanProcessor(in_mem_exporter)) + test_tracer = provider.get_tracer("test") + + with dag_maker(dag_id, session=session) as dag: + EmptyOperator(task_id="t1") + dr = dag_maker.create_dagrun(state=DagRunState.RUNNING) + dr.dag = dag + traceparent = f"00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-{trace_flag}" + dr.context_carrier = {"traceparent": traceparent} + + with mock.patch("airflow.models.dagrun.tracer", test_tracer): + dr._emit_dagrun_span(state=DagRunState.SUCCESS) + + assert len(in_mem_exporter.get_finished_spans()) == expected_spans + @pytest.mark.db_test def test_context_carrier_includes_detail_level_from_conf(self, dag_maker): """DagRun created with TASK_SPAN_DETAIL_LEVEL_KEY in conf should encode the level in trace state.""" diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index b9d8c85f61321..97ec0242235b4 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -113,6 +113,7 @@ from airflow.utils.types import DagRunTriggeredByType, DagRunType from tests_common.test_utils import db +from tests_common.test_utils.asserts import assert_queries_count from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import clear_db_runs from tests_common.test_utils.mock_operators import MockOperator @@ -3548,6 +3549,23 @@ def test_find_relevant_relatives_with_non_mapped_task_as_tuple(dag_maker, sessio assert result == {"t1"} +def test_register_asset_changes_in_db_no_outlets_is_a_noop(dag_maker, session): + """A task with no outlets and no outlet events must not issue any queries.""" + with dag_maker(dag_id="no_outlets", schedule=None, session=session): + EmptyOperator(task_id="hi") + dr = dag_maker.create_dagrun(session=session) + [ti] = dr.get_task_instances(session=session) + session.commit() + + with assert_queries_count(0): + TaskInstance.register_asset_changes_in_db( + ti=ti, + task_outlets=[], + outlet_events=[], + session=session, + ) + + def test_when_dag_run_has_partition_then_asset_does(dag_maker, session): asset = Asset(name="hello") with dag_maker(dag_id="asset_event_tester", schedule=PartitionAtRuntime()) as dag: @@ -4126,3 +4144,41 @@ def test_stats_tags_with_none_team_name(self, dag_maker, session): ti._team_name = None tags = ti.stats_tags assert "team_name" not in tags + + @pytest.mark.parametrize( + ("team_name", "expected_tags"), + [ + pytest.param( + "my_team", + {"dag_id": "test_dag", "task_id": "my_task", "team_name": "my_team", "queue": "default"}, + id="with_team", + ), + pytest.param( + None, + {"dag_id": "test_dag", "task_id": "my_task", "queue": "default"}, + id="without_team", + ), + ], + ) + @mock.patch("airflow._shared.observability.metrics.stats.timing") + def test_emit_state_change_metric_includes_team_name( + self, mock_timing, team_name, expected_tags, dag_maker, session + ): + with dag_maker("test_dag"): + EmptyOperator(task_id="my_task") + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("my_task", session=session) + ti.state = TaskInstanceState.SCHEDULED + ti.scheduled_dttm = timezone.utcnow() + if team_name: + ti._team_name = team_name + session.merge(ti) + session.flush() + + ti.emit_state_change_metric(TaskInstanceState.QUEUED) + + mock_timing.assert_called_once_with( + "task.scheduled_duration", + mock.ANY, + tags=expected_tags, + ) diff --git a/airflow-core/tests/unit/timetables/test_partitioned_timetable.py b/airflow-core/tests/unit/timetables/test_partitioned_timetable.py index 23caf06e79d0e..cb5d882179421 100644 --- a/airflow-core/tests/unit/timetables/test_partitioned_timetable.py +++ b/airflow-core/tests/unit/timetables/test_partitioned_timetable.py @@ -27,6 +27,7 @@ import pytest from airflow._shared.module_loading import qualname +from airflow.exceptions import InvalidPartitionKeyError from airflow.partition_mappers.base import RollupMapper from airflow.partition_mappers.identity import IdentityMapper as IdentityMapper from airflow.partition_mappers.temporal import StartOfDayMapper @@ -258,6 +259,28 @@ def test_deserialize(self): assert isinstance(timetable.default_partition_mapper, IdentityMapper) assert isinstance(timetable.partition_mapper_config[ser_asset], IdentityMapper) + @pytest.mark.parametrize( + "asset_like", + [ + Asset(name="daily", uri="s3://bucket/daily"), + Asset.ref(name="daily"), + ], + ) + def test_decode_partition_date_raises_on_invalid_key(self, asset_like): + timetable = PartitionedAssetTimetable( + assets=asset_like, + default_partition_mapper=StartOfDayMapper(), + ) + with pytest.raises(InvalidPartitionKeyError, match="is invalid for this timetable's mappers"): + timetable._decode_partition_date("not-a-date") + + def test_decode_partition_date_returns_period_start_for_valid_key(self): + timetable = PartitionedAssetTimetable( + assets=Asset(name="daily", uri="s3://bucket/daily"), + default_partition_mapper=StartOfDayMapper(), + ) + assert timetable._decode_partition_date("2025-01-01") == pendulum.datetime(2025, 1, 1, tz="UTC") + def test_partitioned_asset_timetable_resolve_day_bound_returns_midnight_utc(self): """PartitionedAssetTimetable has no local timezone; resolve_day_bound uses the base default. diff --git a/airflow-core/tests/unit/triggers/test_base_trigger.py b/airflow-core/tests/unit/triggers/test_base_trigger.py index 5f32cb23bfe41..429e54615293a 100644 --- a/airflow-core/tests/unit/triggers/test_base_trigger.py +++ b/airflow-core/tests/unit/triggers/test_base_trigger.py @@ -140,6 +140,25 @@ def test_render_template_fields_empty_when_no_trigger_kwargs(create_task_instanc assert trigger.name == "Hello {{ name }}" +class _TriggerWithoutSuperInit(BaseTrigger): + """A trigger that does not call super().__init__() — simulates third-party triggers.""" + + def __init__(self, queue_url: str): + self.queue_url = queue_url + + def serialize(self): + return (f"{type(self).__module__}.{type(self).__qualname__}", {"queue_url": self.queue_url}) + + async def run(self): + yield TriggerEvent({"queue_url": self.queue_url}) + + +def test_task_instance_property_works_without_super_init(): + """task_instance property must return None when subclass skips super().__init__().""" + trigger = _TriggerWithoutSuperInit(queue_url="https://sqs.example.com/queue") + assert trigger.task_instance is None + + class _PlainEventTrigger(BaseEventTrigger): """A BaseEventTrigger that does not opt into shared streams.""" diff --git a/airflow-core/tests/unit/utils/test_email.py b/airflow-core/tests/unit/utils/test_email.py index c45e83d6e7714..f73edeb3679a3 100644 --- a/airflow-core/tests/unit/utils/test_email.py +++ b/airflow-core/tests/unit/utils/test_email.py @@ -324,7 +324,7 @@ def test_send_mime_complete_failure(self, mock_smtp: mock.Mock, mock_smtp_ssl: m port=conf.getint("smtp", "SMTP_PORT"), timeout=conf.getint("smtp", "SMTP_TIMEOUT"), ) - assert mock_smtp.call_count == conf.getint("smtp", "SMTP_RETRY_LIMIT") + assert mock_smtp.call_count == conf.getint("smtp", "SMTP_RETRY_LIMIT") + 1 assert not mock_smtp_ssl.called assert not mock_smtp.return_value.starttls.called assert not mock_smtp.return_value.login.called @@ -348,7 +348,7 @@ def test_send_mime_ssl_complete_failure(self, create_default_context, mock_smtp, context=create_default_context.return_value, ) assert create_default_context.called - assert mock_smtp_ssl.call_count == conf.getint("smtp", "SMTP_RETRY_LIMIT") + assert mock_smtp_ssl.call_count == conf.getint("smtp", "SMTP_RETRY_LIMIT") + 1 assert not mock_smtp.called assert not mock_smtp_ssl.return_value.starttls.called assert not mock_smtp_ssl.return_value.login.called @@ -377,7 +377,7 @@ def test_send_mime_custom_timeout_retrylimit(self, mock_smtp, mock_smtp_ssl): host=conf.get("smtp", "SMTP_HOST"), port=conf.getint("smtp", "SMTP_PORT"), timeout=custom_timeout ) assert not mock_smtp_ssl.called - assert mock_smtp.call_count == 10 + assert mock_smtp.call_count == custom_retry_limit + 1 @mock.patch("smtplib.SMTP_SSL") @mock.patch("smtplib.SMTP") diff --git a/dev/breeze/doc/09_release_management_tasks.rst b/dev/breeze/doc/09_release_management_tasks.rst index 6a4d1a86b2109..63d691950875a 100644 --- a/dev/breeze/doc/09_release_management_tasks.rst +++ b/dev/breeze/doc/09_release_management_tasks.rst @@ -447,6 +447,25 @@ You can also add ``--answer yes`` to perform non-interactive build. :width: 100% :alt: Breeze prepare-provider-documentation +Classifying provider changes +"""""""""""""""""""""""""""" + +You can use Breeze to classify each provider's unreleased changes using hard-coded, +high-confidence rules, flagging ambiguous commits as ``needs_llm`` for an agent or skill +to assess. The result is emitted as JSON, providing a deterministic alternative to the +``--non-interactive`` documentation run used purely for change discovery. + +The below example classifies the pending changes for all providers. + +.. code-block:: bash + + breeze release-management classify-provider-changes + +.. image:: ./images/output_release-management_classify-provider-changes.svg + :target: https://raw.githubusercontent.com/apache/airflow/main/dev/breeze/doc/images/output_release-management_classify-provider-changes.svg + :width: 100% + :alt: Breeze classify-provider-changes + Updating provider next version """""""""""""""""""""""""""""" diff --git a/dev/breeze/doc/images/output_release-management.svg b/dev/breeze/doc/images/output_release-management.svg index 2613a5159f574..309ebac51d6d4 100644 --- a/dev/breeze/doc/images/output_release-management.svg +++ b/dev/breeze/doc/images/output_release-management.svg @@ -1,4 +1,4 @@ - +