kaxil commented on code in PR #74790:
URL: https://github.com/apache/airflow/pull/74790#discussion_r4402517914
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/v2027_02_28/test_jobs.py:
##########
@@ -62,6 +64,10 @@ def test_head_version_routes_to_endpoint(self, client):
assert response.status_code == 404
assert response.json()["detail"]["reason"] == "not_found"
+ def test_old_version_has_no_bundle_catalog(self, client):
+ client.headers["Airflow-API-Version"] = "2026-10-30"
+ assert client.get("/execution/jobs/0/bundles").status_code == 404
Review Comment:
The fixture token has `job_id=1` and this calls `/jobs/0/bundles`, so the
head route's `_validate_job` returns 404 as well, and this passes even without
the `didnt_exist` entry. Asserting `response.json() == {"detail": "Not Found"}`
like `test_old_version_returns_404` does (or adding it to that parametrize)
would pin the version gate.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_parsing.py:
##########
@@ -0,0 +1,413 @@
+# 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.
+
+from __future__ import annotations
+
+import hashlib
+import json
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Request, Security, status
+from sqlalchemy import delete, or_, select
+from sqlalchemy.orm import selectinload
+from starlette.requests import ClientDisconnect
+from uuid6 import uuid7
+
+from airflow._shared.timezones import timezone
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
+from airflow.api_fastapi.execution_api.datamodels.dag_parsing import (
+ MAX_PUBLICATION_BYTES,
+ DagBundleInventoryBody,
+ DagBundleInventoryResponse,
+ DagBundleStateResponse,
+ DagParseResultBody,
+ DagParseResultResponse,
+)
+from airflow.api_fastapi.execution_api.security import CurrentExecutionToken,
ExecutionAPIRoute, require_auth
+from airflow.configuration import conf
+from airflow.dag_processing.collection import (
+ DagBundleOwnershipError,
+ DagPublicationValidationError,
+ _reject_other_teams_plugin_classes,
+ is_source_discovery_failed,
+ update_dag_parsing_results_in_db,
+ validate_dag_bundle_ownership,
+)
+from airflow.exceptions import DeserializationError
+from airflow.jobs.job import Job
+from airflow.models.asset import remove_references_to_deleted_dags
+from airflow.models.dag import DagModel
+from airflow.models.dag_parse_checkpoint import DagParseCheckpoint
+from airflow.models.dagbundle import DagBundleModel
+from airflow.models.dagwarning import DagWarning
+from airflow.models.errors import ParseImportError
+from airflow.sdk.importers.base import DagSourceCode # noqa: SDK001
+from airflow.serialization.serialized_objects import DagSerialization,
LazyDeserializedDAG
+
+if TYPE_CHECKING:
+ from collections.abc import Callable, Coroutine
+
+ from pydantic import BaseModel
+ from sqlalchemy.orm import Session
+ from starlette.responses import Response
+ from starlette.types import Message
+
+
+class DagParsingRoute(ExecutionAPIRoute):
+ """Refuse request bodies over the publication limit before they are read
into memory."""
+
+ def get_route_handler(self) -> Callable[[Request], Coroutine[Any, Any,
Response]]:
+ handler = super().get_route_handler()
+
+ async def handle_bounded_request(request: Request) -> Response:
+ too_large = HTTPException(
+ status.HTTP_413_CONTENT_TOO_LARGE,
+ detail=f"Request body exceeds {MAX_PUBLICATION_BYTES // (1024
* 1024)} MiB",
+ )
+ declared = request.headers.get("content-length", "")
+ if declared.isdigit() and int(declared) > MAX_PUBLICATION_BYTES:
+ raise too_large
+ body = bytearray()
+ try:
+ async for chunk in request.stream():
+ body += chunk
+ if len(body) > MAX_PUBLICATION_BYTES:
+ raise too_large
+ except ClientDisconnect as error:
+ # FastAPI reports a body it could not read as a bad request
rather than a server error.
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST, detail="The client
disconnected while sending the body"
+ ) from error
+ replayed = False
+
+ async def receive() -> Message:
+ nonlocal replayed
+ if replayed:
+ return await request.receive()
+ replayed = True
+ return {"type": "http.request", "body": bytes(body),
"more_body": False}
+
+ return await handler(Request(request.scope, receive))
+
+ return handle_bounded_request
+
+
+router = VersionedAPIRouter(route_class=DagParsingRoute)
+
+
+def _compute_payload_hash(body: BaseModel) -> str:
+ """Hash a request body independently of key order; a non-JSON float raises
ValueError."""
+ payload = json.dumps(body.model_dump(mode="json"), sort_keys=True,
allow_nan=False)
+ return hashlib.sha256(payload.encode()).hexdigest()
+
+
+def _validate_job(job_id: int, token, *, session: Session, for_update: bool =
True) -> None:
+ if job_id != token.claims.job_id:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, detail={"reason":
"not_found"})
+ query = select(Job).where(Job.id == job_id)
+ # Publications lock the Job to serialize checkpoint creation and fence
completion or replacement.
+ job = session.scalar(query.with_for_update() if for_update else query)
+ if job is None or job.session_id != token.id or job.end_date is not None:
+ raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason":
"job_closed"})
+
+
[email protected](
+ "/{job_id}/parse-results",
+ responses=create_openapi_http_exception_doc(
+ [
+ (status.HTTP_403_FORBIDDEN, "The bundle or a Dag's plugin is not
granted, or the Job has ended"),
+ status.HTTP_404_NOT_FOUND,
+ status.HTTP_409_CONFLICT,
+ status.HTTP_413_CONTENT_TOO_LARGE,
+ status.HTTP_422_UNPROCESSABLE_CONTENT,
+ ]
+ ),
+ dependencies=[Security(require_auth, scopes=["token:dag_processor"])],
+)
+def publish_parse_result(
+ job_id: int, body: DagParseResultBody, *, session: SessionDep,
token=CurrentExecutionToken
+) -> DagParseResultResponse:
+ """Atomically publish a file's results and a receipt, without reading the
processor's filesystem."""
+ if body.bundle_name not in token.claims.dag_bundles:
+ raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason":
"bundle_not_granted"})
+ try:
+ payload_hash = _compute_payload_hash(body)
+ except ValueError:
+ raise HTTPException(status.HTTP_422_UNPROCESSABLE_CONTENT,
detail="Invalid JSON value")
+ source_key = hashlib.sha256(
+ json.dumps(["parse", body.bundle_name, body.relative_fileloc]).encode()
+ ).hexdigest()
+
+ _validate_job(job_id, token, session=session)
+ checkpoint = session.get(DagParseCheckpoint, (job_id, source_key))
+ if checkpoint is not None:
+ if checkpoint.attempt_id == body.attempt_id:
+ if checkpoint.payload_hash != payload_hash:
+ raise HTTPException(status.HTTP_409_CONFLICT,
detail={"reason": "publication_conflict"})
+ return DagParseResultResponse(
+ attempt_id=checkpoint.attempt_id,
accepted_at=checkpoint.accepted_at
+ )
+ if body.dispatch_sequence <= checkpoint.dispatch_sequence:
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"publication_superseded"})
+
+ bundle = session.scalar(
+ select(DagBundleModel).where(DagBundleModel.name ==
body.bundle_name).with_for_update(key_share=True)
+ )
+ if bundle is None or not bundle.active:
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"bundle_inactive"})
+ if body.bundle_revision != bundle.parse_revision or (
+ bundle.parse_revision is not None and body.bundle_version !=
bundle.version
+ ):
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"source_changed"})
+ dags = [LazyDeserializedDAG(data=document) for document in
body.serialized_dags]
+ dag_ids = [dag.dag_id for dag in dags]
+ existing = session.scalars(
+ select(DagModel)
+ .where(DagModel.dag_id.in_(dag_ids))
+ .order_by(DagModel.dag_id)
+ .with_for_update(key_share=True)
+ ).all()
+ try:
+ validate_dag_bundle_ownership(existing, body.bundle_name,
session=session)
+ except DagBundleOwnershipError as error:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT, detail={"reason":
"dag_owned_by_another_bundle"}
+ ) from error
+ if len(_reject_other_teams_plugin_classes(body.bundle_name, dags, {},
session=session)) != len(dags):
+ raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason":
"plugin_not_granted"})
+ try:
+ for dag in dags:
+ DagSerialization.validate_serialized_dag(dag.data)
+ except (DeserializationError, ValueError, TypeError, KeyError) as error:
+ raise HTTPException(status.HTTP_422_UNPROCESSABLE_CONTENT,
detail="Invalid serialized Dag") from error
+
+ # A complete accepted import proves absence; an import error only ages out
prior definitions.
+ cutoff = timezone.utcnow() -
timedelta(seconds=conf.getint("dag_processor", "stale_dag_threshold"))
+ source_filter = or_(
+ DagModel.relative_fileloc == body.relative_fileloc,
+ DagModel.relative_fileloc.startswith(body.relative_fileloc + "/",
autoescape=True),
+ )
+ missing = session.scalars(
+ select(DagModel)
+ .where(DagModel.bundle_name == body.bundle_name, source_filter,
DagModel.dag_id.not_in(dag_ids))
+ .order_by(DagModel.dag_id)
+ .with_for_update(key_share=True)
+ )
+ for stored_dag in missing:
+ failed = any(
+ path == body.relative_fileloc or path ==
stored_dag.relative_fileloc
Review Comment:
`_update_import_errors` runs later in this transaction and sets
`is_stale=True` on every Dag whose `relative_fileloc` matches an errored path,
whatever its `last_parsed_time`. So `path == stored_dag.relative_fileloc` never
changes the outcome, and the only Dags this branch keeps are members of a
container whose error is on the container path. Is that the intent? If so, the
comment on line 201 reads broader than what happens, and this could be
`body.relative_fileloc in body.import_errors`.
##########
airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dag_parsing.py:
##########
@@ -0,0 +1,145 @@
+# 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.
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any
+from uuid import UUID
+
+from pydantic import Field, model_validator
+
+from airflow.api_fastapi.core_api.base import StrictBaseModel
+from airflow.api_fastapi.execution_api.datamodels.job import DagParseTokenBody
+from airflow.models.dag import TAG_MAX_LEN
+from airflow.models.dagwarning import DagWarningTypeValue
+from airflow.utils.helpers import validate_key
+
+MAX_PUBLICATION_BYTES = 16 * 1024 * 1024
+"""Largest request body the Dag parsing endpoints read, measured as the bytes
sent."""
+
+
+class ParseWarning(StrictBaseModel):
+ """A warning attached to a Dag returned by this import."""
+
+ dag_id: str = Field(min_length=1, max_length=250)
+ warning_type: DagWarningTypeValue
+ message: str
+
+
+class ParseSourceCode(StrictBaseModel):
+ """Source captured by the importer; null explicitly means unavailable."""
+
+ source_code: str | None
+ language: str = Field(min_length=1, max_length=64)
+
+
+class DagParseResultBody(DagParseTokenBody):
+ """One completed file or container import, including an empty result."""
+
+ dispatch_sequence: int = Field(ge=1, le=2**53 - 1)
+ bundle_revision: UUID | None = None
+ bundle_version: str | None = Field(default=None, max_length=200)
+ version_data: dict[str, Any] | None = None
+ parse_duration: float = Field(ge=0, allow_inf_nan=False)
+ serialized_dags: list[dict[str, Any]]
+ import_errors: dict[str, str] = Field(default_factory=dict)
+ warnings: list[ParseWarning] = Field(default_factory=list)
+ parsed_definitions: list[str] = Field(default_factory=list)
+ source_codes: dict[str, ParseSourceCode]
Review Comment:
`PythonImporter.get_source_code` ignores `dag_id` and returns the whole
file, so a factory file of S bytes that builds N Dags sends N x S bytes of
source in one body, all under the 16 MiB `MAX_PUBLICATION_BYTES` cap. A 40 KB
file building 400 Dags is ~15 MiB of source before any serialized Dag, and at
the top of the stack a 413 turns into an import error for the file. Since
2027-02-28 hasn't shipped, would it be worth deduplicating now? Something like
`sources: dict[str, ParseSourceCode]` keyed by content hash plus `dag_sources:
dict[str, str]` mapping dag_id to that key would still let an importer return
different source per Dag.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_jobs.py:
##########
@@ -410,6 +428,7 @@ def test_records_the_final_state(self, client, session,
authenticate, time_machi
assert response.status_code == 204, response.text
session.refresh(job)
assert (job.state, job.end_date) == (JobState.SUCCESS, NOW +
timedelta(minutes=5))
+ assert session.scalar(select(DagParseCheckpoint)) is None
Review Comment:
With a single checkpoint in the table, this still passes if the delete loses
its `job_id` filter, which would wipe another live processor's receipts.
Seeding a checkpoint for a second open Job and asserting it survives would
cover the scoping (same for the assertion in `test_replaces_a_stopped_job`).
##########
airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dag_parsing.py:
##########
@@ -0,0 +1,145 @@
+# 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.
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any
+from uuid import UUID
+
+from pydantic import Field, model_validator
+
+from airflow.api_fastapi.core_api.base import StrictBaseModel
+from airflow.api_fastapi.execution_api.datamodels.job import DagParseTokenBody
+from airflow.models.dag import TAG_MAX_LEN
+from airflow.models.dagwarning import DagWarningTypeValue
+from airflow.utils.helpers import validate_key
+
+MAX_PUBLICATION_BYTES = 16 * 1024 * 1024
+"""Largest request body the Dag parsing endpoints read, measured as the bytes
sent."""
+
+
+class ParseWarning(StrictBaseModel):
+ """A warning attached to a Dag returned by this import."""
+
+ dag_id: str = Field(min_length=1, max_length=250)
+ warning_type: DagWarningTypeValue
+ message: str
+
+
+class ParseSourceCode(StrictBaseModel):
+ """Source captured by the importer; null explicitly means unavailable."""
+
+ source_code: str | None
+ language: str = Field(min_length=1, max_length=64)
+
+
+class DagParseResultBody(DagParseTokenBody):
+ """One completed file or container import, including an empty result."""
+
+ dispatch_sequence: int = Field(ge=1, le=2**53 - 1)
+ bundle_revision: UUID | None = None
+ bundle_version: str | None = Field(default=None, max_length=200)
+ version_data: dict[str, Any] | None = None
+ parse_duration: float = Field(ge=0, allow_inf_nan=False)
+ serialized_dags: list[dict[str, Any]]
+ import_errors: dict[str, str] = Field(default_factory=dict)
+ warnings: list[ParseWarning] = Field(default_factory=list)
+ parsed_definitions: list[str] = Field(default_factory=list)
+ source_codes: dict[str, ParseSourceCode]
+
+ @model_validator(mode="after")
+ def validate_file_scope(self) -> DagParseResultBody:
+ paths = [*self.parsed_definitions, *self.import_errors]
+ dag_ids: set[str] = set()
+ for document in self.serialized_dags:
+ dag = document.get("dag")
+ if not isinstance(dag, dict) or not isinstance(dag.get("dag_id"),
str):
+ raise ValueError("A serialized Dag must contain a dag_id")
+ # The processor runs Dag code that can bypass the SDK's checks;
these columns are bounded.
+ validate_key(dag["dag_id"])
+ tags = dag.get("tags") or []
+ if not isinstance(tags, list) or any(
+ not isinstance(tag, str) or len(tag) > TAG_MAX_LEN for tag in
tags
+ ):
+ raise ValueError(f"Dag tags must be strings of at most
{TAG_MAX_LEN} characters")
+ if dag["dag_id"] in dag_ids:
+ raise ValueError("Duplicate Dag id in publication")
+ dag_ids.add(dag["dag_id"])
+ relative_fileloc = dag.get("relative_fileloc")
+ if not isinstance(relative_fileloc, str):
+ raise ValueError("Every definition must have a bundle-relative
location")
+ paths.append(relative_fileloc)
+ fileloc = dag.get("fileloc")
+ if not isinstance(fileloc, str) or len(fileloc) > 2000:
+ raise ValueError("A serialized Dag must contain a fileloc of
at most 2000 characters")
+ for path in paths:
+ if len(path) > 2000:
Review Comment:
Import-error keys go through this 2000-character check too, but
`ParseImportError.filename` is `String(1024)`. On Postgres a 1025-2000
character key fails at flush with `DataError`, which comes back as a 500 the
processor treats as retryable instead of a rejected result. Bounding the
`import_errors` keys at 1024 would make it a validation error (and a named
constant would help, the 2000 is repeated on lines 88, 91 and 134).
##########
airflow-core/src/airflow/listeners/spec/importerrors.py:
##########
@@ -25,9 +25,17 @@
@hookspec
def on_new_dag_import_error(filename, stacktrace):
- """Execute when new dag import error appears."""
+ """
+ Execute when new dag import error appears.
+
+ ``filename`` is bundle-relative when the Dag processor publishes through
the Execution API.
Review Comment:
With API publication `filename` becomes bundle-relative with no bundle name
next to it, so a listener can't tell `a.py` in one bundle from `a.py` in
another, or which form of path it was given. Pluggy lets a hookimpl accept a
subset of the hookspec's arguments, so adding `bundle_name` and
`relative_fileloc` to both specs wouldn't break existing listeners. Would that
be better than `filename` meaning different things per deployment mode? It's
also the `todo` already sitting in `_update_import_errors`.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_parsing.py:
##########
@@ -0,0 +1,1116 @@
+# 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.
+
+from __future__ import annotations
+
+import asyncio
+import json
+from concurrent.futures import ThreadPoolExecutor
+from copy import deepcopy
+from threading import Barrier, Event
+from unittest import mock
+from uuid import UUID
+
+import pytest
+from fastapi import APIRouter, FastAPI, Request
+from sqlalchemy import delete, event, select
+from sqlalchemy.exc import OperationalError
+from sqlalchemy.orm import Session
+from uuid6 import uuid7
+
+from airflow import settings
+from airflow._shared.timezones import timezone
+from airflow.api_fastapi.execution_api.datamodels.token import
DagProcessorClaims, DagProcessorToken
+from airflow.api_fastapi.execution_api.routes.dag_parsing import
DagParsingRoute, get_bundles
+from airflow.api_fastapi.execution_api.security import require_auth
+from airflow.dag_processing.collection import DagModelOperation,
update_dag_parsing_results_in_db
+from airflow.jobs.job import Job, JobState
+from airflow.models.asset import DagScheduleAssetReference
+from airflow.models.dag import DagModel
+from airflow.models.dag_parse_checkpoint import DagParseCheckpoint
+from airflow.models.dag_version import DagVersion
+from airflow.models.dagbundle import DagBundleModel
+from airflow.models.dagcode import DagCode
+from airflow.models.errors import ParseImportError
+from airflow.models.serialized_dag import SerializedDagModel
+from airflow.models.team import Team
+from airflow.sdk import DAG, Asset
+from airflow.serialization.serialized_objects import LazyDeserializedDAG
+from airflow.utils.session import create_session
+
+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_assets,
+ clear_db_dag_bundles,
+ clear_db_dags,
+ clear_db_import_errors,
+ clear_db_jobs,
+)
+
+pytestmark = pytest.mark.db_test
+SESSION_ID = UUID("00000000-0000-0000-0000-0000000000aa")
+
+
[email protected](autouse=True)
+def clean_db(time_machine, session):
+ time_machine.move_to("2026-10-05T12:00:00Z", tick=False)
+ clear_db_jobs()
+ clear_db_dags()
+ clear_db_import_errors()
+ clear_db_dag_bundles()
+ yield
+ clear_db_jobs()
+ clear_db_dags()
+ clear_db_import_errors()
+ clear_db_dag_bundles()
+
+
[email protected]
+def job(session, exec_app):
+ job = Job(job_type="DagProcessorJob", state=JobState.RUNNING)
+ job.session_id = SESSION_ID
+ job.registration_id = uuid7()
+ job.bundle_names = ["bundle"]
+ session.add_all([job, DagBundleModel(name="bundle"),
DagBundleModel(name="other")])
+ session.commit()
+
+ async def authenticate(request: Request):
+ return DagProcessorToken(
+ id=SESSION_ID,
+ claims=DagProcessorClaims(exp=1, job_id=job.id,
dag_bundles=frozenset({"bundle"})),
+ )
+
+ exec_app.dependency_overrides[require_auth] = authenticate
+ return job
+
+
[email protected]
+def body():
+ dag = DAG("published", schedule=None)
+ dag.fileloc = "/worker-only/dags/a.py"
+ dag.relative_fileloc = "a.py"
+ return {
+ "attempt_id": str(uuid7()),
+ "dispatch_sequence": 1,
+ "bundle_name": "bundle",
+ "relative_fileloc": "a.py",
+ "bundle_version": "v1",
+ "version_data": {"ref": "v1"},
+ "parse_duration": 0.5,
+ "serialized_dags": [LazyDeserializedDAG.from_dag(dag).data],
+ "source_codes": {dag.dag_id: {"source_code": "# captured on worker",
"language": "python"}},
+ }
+
+
[email protected](
+ ParseImportError, "full_file_path", autospec=True,
side_effect=AssertionError("Server constructed bundle")
+)
[email protected]("airflow.listeners.listener.get_listener_manager", autospec=True)
+def test_import_error_listeners_get_bundle_relative_paths(
+ get_listener_manager, full_file_path, client, job, body
+):
+ hook = get_listener_manager.return_value.hook
+ body.update(serialized_dags=[], source_codes={}, import_errors={"a.py":
"broken"})
+ assert client.post(f"/execution/jobs/{job.id}/parse-results",
json=body).status_code == 200
+ body.update(attempt_id=str(uuid7()), dispatch_sequence=2,
import_errors={"a.py": "still broken"})
+ assert client.post(f"/execution/jobs/{job.id}/parse-results",
json=body).status_code == 200
+
+ hook.on_new_dag_import_error.assert_called_once_with(filename="a.py",
stacktrace="broken")
+ hook.on_existing_dag_import_error.assert_called_once_with(filename="a.py",
stacktrace="still broken")
+ full_file_path.assert_not_called()
+
+
[email protected]("source", ["# captured on worker", "", None])
[email protected](DagCode, "code", autospec=True,
side_effect=AssertionError("Server read source"))
+def test_publishes_worker_source_without_filesystem_access(mock_code, client,
session, job, body, source):
+ body["source_codes"]["published"]["source_code"] = source
+ response = client.post(f"/execution/jobs/{job.id}/parse-results",
json=body)
+ assert response.status_code == 200, response.json()
+ assert session.get(DagModel, "published").relative_fileloc == "a.py"
+ assert session.scalar(select(SerializedDagModel)).dag_id == "published"
+ assert session.scalar(select(DagCode)).source_code == (
+ source if source is not None else "Source unavailable"
+ )
+ assert session.scalar(select(DagParseCheckpoint)).attempt_id ==
UUID(body["attempt_id"])
+ mock_code.assert_not_called()
+
+
+def test_replay_returns_receipt_without_rewriting_metadata(client, session,
job, body, time_machine):
+ first = client.post(f"/execution/jobs/{job.id}/parse-results", json=body)
+ assert first.status_code == 200, first.json()
+ parsed_at = session.get(DagModel, "published").last_parsed_time
+ time_machine.shift(10)
+ replay = client.post(f"/execution/jobs/{job.id}/parse-results", json=body)
+ assert replay.json() == first.json()
+ session.expire_all()
+ assert session.get(DagModel, "published").last_parsed_time == parsed_at
+ assert len(session.scalars(select(DagParseCheckpoint)).all()) == 1
+
+
[email protected]("first", ["parse", "inventory"])
+def test_parse_and_inventory_checkpoints_have_separate_namespaces(client,
session, exec_app, job, first):
+ session.add_all([DagBundleModel(name="inventory"),
DagBundleModel(name="a.py")])
+ job.bundle_names = ["inventory", "a.py"]
+ session.commit()
+
+ async def authenticate(request: Request):
+ return DagProcessorToken(
+ id=job.session_id,
+ claims=DagProcessorClaims(exp=1, job_id=job.id,
dag_bundles=frozenset(job.bundle_names)),
+ )
+
+ exec_app.dependency_overrides[require_auth] = authenticate
+ publications = {
+ "parse": (
+ f"/execution/jobs/{job.id}/parse-results",
+ {
+ "attempt_id": str(uuid7()),
+ "dispatch_sequence": 2,
+ "bundle_name": "inventory",
+ "relative_fileloc": "a.py",
+ "parse_duration": 0.5,
+ "serialized_dags": [],
+ "source_codes": {},
+ "import_errors": {"a.py": "Import failed"},
+ },
+ ),
+ "inventory": (
+ f"/execution/jobs/{job.id}/bundles/a.py/inventory",
+ {
+ "attempt_id": str(uuid7()),
+ "dispatch_sequence": 2,
+ "expected_revision": None,
+ "files": [],
+ },
+ ),
+ }
+ publications[first][1]["dispatch_sequence"] = 3
+ second = "inventory" if first == "parse" else "parse"
+ receipts = {}
+ for kind in (first, second, first, second):
+ url, payload = publications[kind]
+ response = client.post(url, json=payload)
+ assert response.status_code == 200, response.json()
+ if kind in receipts:
+ assert response.json() == receipts[kind]
+ else:
+ receipts[kind] = response.json()
+
+ checkpoints = session.scalars(select(DagParseCheckpoint)).all()
+ assert len(checkpoints) == 2
+ assert {checkpoint.attempt_id for checkpoint in checkpoints} == {
+ UUID(payload["attempt_id"]) for _, payload in publications.values()
+ }
+ error = session.scalar(select(ParseImportError))
+ assert (error.bundle_name, error.filename) == ("inventory", "a.py")
+
+
[email protected]("paused", [True, False])
+def
test_publication_preserves_pause_on_creation_and_existing_pause_state(client,
session, job, body, paused):
+ body["serialized_dags"][0]["dag"]["is_paused_upon_creation"] = paused
+ assert client.post(f"/execution/jobs/{job.id}/parse-results",
json=body).status_code == 200
+ assert session.get(DagModel, "published").is_paused is paused
+ body.update(attempt_id=str(uuid7()), dispatch_sequence=2)
+ body["serialized_dags"][0]["dag"]["is_paused_upon_creation"] = not paused
+ assert client.post(f"/execution/jobs/{job.id}/parse-results",
json=body).status_code == 200
+ session.expire_all()
+ assert session.get(DagModel, "published").is_paused is paused
+
+
[email protected](
+ ("change", "reason"), [("payload", "publication_conflict"), ("attempt",
"publication_superseded")]
+)
+def test_rejects_conflicting_or_superseded_publication(client, session, job,
body, change, reason):
+ assert client.post(f"/execution/jobs/{job.id}/parse-results",
json=body).status_code == 200
+ if change == "payload":
+ body["parse_duration"] = 99
+ else:
+ body["attempt_id"] = str(uuid7())
+ response = client.post(f"/execution/jobs/{job.id}/parse-results",
json=body)
+ assert response.status_code == 409
+ assert response.json()["detail"]["reason"] == reason
+ assert session.get(DagModel, "published").last_parse_duration == 0.5
+
+
+@conf_vars({("core", "min_serialized_dag_update_interval"): "0"})
+def test_newer_result_supersedes_old_attempt(client, session, job, body):
+ original = deepcopy(body)
+ assert client.post(f"/execution/jobs/{job.id}/parse-results",
json=body).status_code == 200
+ body.update(attempt_id=str(uuid7()), dispatch_sequence=2,
parse_duration=2.0)
+ assert client.post(f"/execution/jobs/{job.id}/parse-results",
json=body).status_code == 200
+ assert client.post(f"/execution/jobs/{job.id}/parse-results",
json=original).status_code == 409
+ session.expire_all()
+ assert session.get(DagModel, "published").last_parse_duration == 2.0
+ assert len(session.scalars(select(DagParseCheckpoint)).all()) == 1
+
+
[email protected]("same_bundle", [True, False])
+def test_allows_file_relocation_but_not_cross_bundle_takeover(client, session,
job, body, same_bundle):
+ session.add(
+ DagModel(
+ dag_id="published", bundle_name="bundle" if same_bundle else
"other", relative_fileloc="old.py"
+ )
+ )
+ session.commit()
+ response = client.post(f"/execution/jobs/{job.id}/parse-results",
json=body)
+ assert response.status_code == (200 if same_bundle else 409),
response.json()
+ session.expire_all()
+ assert session.get(DagModel, "published").relative_fileloc == ("a.py" if
same_bundle else "old.py")
+
+
[email protected](
+ "failure", ["bundle", "job", "closed", "replaced", "inactive", "deleted",
"missing_bundle"]
+)
+def test_rejects_unauthorized_publication(client, session, job, body, failure):
+ job_id = job.id
+ if failure == "bundle":
+ body["bundle_name"] = "other"
+ elif failure == "job":
+ job_id += 1000
+ elif failure == "closed":
+ job.end_date = job.start_date
+ elif failure == "replaced":
+ job.session_id = None
+ elif failure == "deleted":
+ session.delete(job)
+ elif failure == "missing_bundle":
+ session.delete(session.get(DagBundleModel, "bundle"))
+ else:
+ session.get(DagBundleModel, "bundle").active = False
+ session.commit()
+ response = client.post(f"/execution/jobs/{job_id}/parse-results",
json=body)
+ assert response.status_code == {"job": 404, "inactive": 409,
"missing_bundle": 409}.get(failure, 403)
Review Comment:
This only checks the status code, so swapping `job_closed` and
`bundle_not_granted` still passes, and `DagProcessorAPIClient` in #74794
branches on `job_closed`. Could the parametrize carry the expected
`detail["reason"]` as well?
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_parsing.py:
##########
@@ -0,0 +1,413 @@
+# 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.
+
+from __future__ import annotations
+
+import hashlib
+import json
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Request, Security, status
+from sqlalchemy import delete, or_, select
+from sqlalchemy.orm import selectinload
+from starlette.requests import ClientDisconnect
+from uuid6 import uuid7
+
+from airflow._shared.timezones import timezone
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
+from airflow.api_fastapi.execution_api.datamodels.dag_parsing import (
+ MAX_PUBLICATION_BYTES,
+ DagBundleInventoryBody,
+ DagBundleInventoryResponse,
+ DagBundleStateResponse,
+ DagParseResultBody,
+ DagParseResultResponse,
+)
+from airflow.api_fastapi.execution_api.security import CurrentExecutionToken,
ExecutionAPIRoute, require_auth
+from airflow.configuration import conf
+from airflow.dag_processing.collection import (
+ DagBundleOwnershipError,
+ DagPublicationValidationError,
+ _reject_other_teams_plugin_classes,
+ is_source_discovery_failed,
+ update_dag_parsing_results_in_db,
+ validate_dag_bundle_ownership,
+)
+from airflow.exceptions import DeserializationError
+from airflow.jobs.job import Job
+from airflow.models.asset import remove_references_to_deleted_dags
+from airflow.models.dag import DagModel
+from airflow.models.dag_parse_checkpoint import DagParseCheckpoint
+from airflow.models.dagbundle import DagBundleModel
+from airflow.models.dagwarning import DagWarning
+from airflow.models.errors import ParseImportError
+from airflow.sdk.importers.base import DagSourceCode # noqa: SDK001
+from airflow.serialization.serialized_objects import DagSerialization,
LazyDeserializedDAG
+
+if TYPE_CHECKING:
+ from collections.abc import Callable, Coroutine
+
+ from pydantic import BaseModel
+ from sqlalchemy.orm import Session
+ from starlette.responses import Response
+ from starlette.types import Message
+
+
+class DagParsingRoute(ExecutionAPIRoute):
+ """Refuse request bodies over the publication limit before they are read
into memory."""
+
+ def get_route_handler(self) -> Callable[[Request], Coroutine[Any, Any,
Response]]:
+ handler = super().get_route_handler()
+
+ async def handle_bounded_request(request: Request) -> Response:
+ too_large = HTTPException(
+ status.HTTP_413_CONTENT_TOO_LARGE,
+ detail=f"Request body exceeds {MAX_PUBLICATION_BYTES // (1024
* 1024)} MiB",
+ )
+ declared = request.headers.get("content-length", "")
+ if declared.isdigit() and int(declared) > MAX_PUBLICATION_BYTES:
+ raise too_large
+ body = bytearray()
+ try:
+ async for chunk in request.stream():
+ body += chunk
+ if len(body) > MAX_PUBLICATION_BYTES:
+ raise too_large
+ except ClientDisconnect as error:
+ # FastAPI reports a body it could not read as a bad request
rather than a server error.
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST, detail="The client
disconnected while sending the body"
+ ) from error
+ replayed = False
+
+ async def receive() -> Message:
+ nonlocal replayed
+ if replayed:
+ return await request.receive()
+ replayed = True
+ return {"type": "http.request", "body": bytes(body),
"more_body": False}
+
+ return await handler(Request(request.scope, receive))
+
+ return handle_bounded_request
+
+
+router = VersionedAPIRouter(route_class=DagParsingRoute)
+
+
+def _compute_payload_hash(body: BaseModel) -> str:
+ """Hash a request body independently of key order; a non-JSON float raises
ValueError."""
+ payload = json.dumps(body.model_dump(mode="json"), sort_keys=True,
allow_nan=False)
+ return hashlib.sha256(payload.encode()).hexdigest()
+
+
+def _validate_job(job_id: int, token, *, session: Session, for_update: bool =
True) -> None:
+ if job_id != token.claims.job_id:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, detail={"reason":
"not_found"})
+ query = select(Job).where(Job.id == job_id)
+ # Publications lock the Job to serialize checkpoint creation and fence
completion or replacement.
+ job = session.scalar(query.with_for_update() if for_update else query)
+ if job is None or job.session_id != token.id or job.end_date is not None:
+ raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason":
"job_closed"})
+
+
[email protected](
+ "/{job_id}/parse-results",
+ responses=create_openapi_http_exception_doc(
+ [
+ (status.HTTP_403_FORBIDDEN, "The bundle or a Dag's plugin is not
granted, or the Job has ended"),
+ status.HTTP_404_NOT_FOUND,
+ status.HTTP_409_CONFLICT,
+ status.HTTP_413_CONTENT_TOO_LARGE,
+ status.HTTP_422_UNPROCESSABLE_CONTENT,
+ ]
+ ),
+ dependencies=[Security(require_auth, scopes=["token:dag_processor"])],
+)
+def publish_parse_result(
+ job_id: int, body: DagParseResultBody, *, session: SessionDep,
token=CurrentExecutionToken
+) -> DagParseResultResponse:
+ """Atomically publish a file's results and a receipt, without reading the
processor's filesystem."""
+ if body.bundle_name not in token.claims.dag_bundles:
+ raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason":
"bundle_not_granted"})
+ try:
+ payload_hash = _compute_payload_hash(body)
+ except ValueError:
+ raise HTTPException(status.HTTP_422_UNPROCESSABLE_CONTENT,
detail="Invalid JSON value")
+ source_key = hashlib.sha256(
+ json.dumps(["parse", body.bundle_name, body.relative_fileloc]).encode()
+ ).hexdigest()
+
+ _validate_job(job_id, token, session=session)
+ checkpoint = session.get(DagParseCheckpoint, (job_id, source_key))
+ if checkpoint is not None:
+ if checkpoint.attempt_id == body.attempt_id:
+ if checkpoint.payload_hash != payload_hash:
+ raise HTTPException(status.HTTP_409_CONFLICT,
detail={"reason": "publication_conflict"})
+ return DagParseResultResponse(
+ attempt_id=checkpoint.attempt_id,
accepted_at=checkpoint.accepted_at
+ )
+ if body.dispatch_sequence <= checkpoint.dispatch_sequence:
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"publication_superseded"})
+
+ bundle = session.scalar(
+ select(DagBundleModel).where(DagBundleModel.name ==
body.bundle_name).with_for_update(key_share=True)
+ )
+ if bundle is None or not bundle.active:
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"bundle_inactive"})
+ if body.bundle_revision != bundle.parse_revision or (
+ bundle.parse_revision is not None and body.bundle_version !=
bundle.version
+ ):
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"source_changed"})
+ dags = [LazyDeserializedDAG(data=document) for document in
body.serialized_dags]
+ dag_ids = [dag.dag_id for dag in dags]
+ existing = session.scalars(
+ select(DagModel)
+ .where(DagModel.dag_id.in_(dag_ids))
+ .order_by(DagModel.dag_id)
+ .with_for_update(key_share=True)
+ ).all()
+ try:
+ validate_dag_bundle_ownership(existing, body.bundle_name,
session=session)
+ except DagBundleOwnershipError as error:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT, detail={"reason":
"dag_owned_by_another_bundle"}
+ ) from error
+ if len(_reject_other_teams_plugin_classes(body.bundle_name, dags, {},
session=session)) != len(dags):
+ raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason":
"plugin_not_granted"})
+ try:
+ for dag in dags:
+ DagSerialization.validate_serialized_dag(dag.data)
+ except (DeserializationError, ValueError, TypeError, KeyError) as error:
+ raise HTTPException(status.HTTP_422_UNPROCESSABLE_CONTENT,
detail="Invalid serialized Dag") from error
+
+ # A complete accepted import proves absence; an import error only ages out
prior definitions.
+ cutoff = timezone.utcnow() -
timedelta(seconds=conf.getint("dag_processor", "stale_dag_threshold"))
+ source_filter = or_(
+ DagModel.relative_fileloc == body.relative_fileloc,
+ DagModel.relative_fileloc.startswith(body.relative_fileloc + "/",
autoescape=True),
+ )
+ missing = session.scalars(
+ select(DagModel)
+ .where(DagModel.bundle_name == body.bundle_name, source_filter,
DagModel.dag_id.not_in(dag_ids))
+ .order_by(DagModel.dag_id)
+ .with_for_update(key_share=True)
+ )
+ for stored_dag in missing:
+ failed = any(
+ path == body.relative_fileloc or path ==
stored_dag.relative_fileloc
+ for path in body.import_errors
+ )
+ if not failed or (stored_dag.last_parsed_time is not None and
stored_dag.last_parsed_time < cutoff):
+ stored_dag.is_stale = True
+
+ import_errors = {(body.bundle_name, path): error for path, error in
body.import_errors.items()}
+ files_parsed = {(body.bundle_name, path) for path in
[body.relative_fileloc, *body.parsed_definitions]}
+ files_parsed.update(import_errors)
+ try:
+ update_dag_parsing_results_in_db(
+ bundle_name=body.bundle_name,
+ bundle_version=body.bundle_version,
+ version_data=body.version_data,
+ dags=dags,
+ import_errors=import_errors,
+ parse_duration=body.parse_duration,
+ warnings={DagWarning(**warning.model_dump()) for warning in
body.warnings},
+ files_parsed=files_parsed,
+ dag_source_codes={
+ dag_id: DagSourceCode(
+ source_code=source.source_code
+ if source.source_code is not None
+ else "Source unavailable",
+ language=source.language,
+ )
+ for dag_id, source in body.source_codes.items()
+ },
+ atomic=True,
+ enforce_bundle_ownership=True,
+ resolve_bundle_paths=False,
+ session=session,
+ )
+ except DagPublicationValidationError as error:
+ raise HTTPException(
+ status.HTTP_422_UNPROCESSABLE_CONTENT,
+ detail={"reason": "invalid_dag", "message": str(error)},
+ ) from error
+ except DagBundleOwnershipError as error:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT, detail={"reason":
"dag_owned_by_another_bundle"}
+ ) from error
+ if checkpoint is None:
+ checkpoint = DagParseCheckpoint(job_id=job_id, source_key=source_key)
+ session.add(checkpoint)
+ checkpoint.attempt_id = body.attempt_id
+ checkpoint.dispatch_sequence = body.dispatch_sequence
+ checkpoint.payload_hash = payload_hash
+ checkpoint.accepted_at = timezone.utcnow()
+ checkpoint.bundle_revision = body.bundle_revision
+ session.flush()
+ return DagParseResultResponse(attempt_id=checkpoint.attempt_id,
accepted_at=checkpoint.accepted_at)
+
+
[email protected](
+ "/{job_id}/bundles",
+ responses=create_openapi_http_exception_doc(
+ [(status.HTTP_403_FORBIDDEN, "The Job has ended or was replaced"),
status.HTTP_404_NOT_FOUND]
+ ),
+ dependencies=[Security(require_auth, scopes=["token:dag_processor"])],
+)
+def get_bundles(
+ job_id: int, *, session: SessionDep, token=CurrentExecutionToken
+) -> list[DagBundleStateResponse]:
+ """Read the provisioned catalog without constructing providers on the API
server."""
+ _validate_job(job_id, token, session=session, for_update=False)
+ bundles = session.scalars(
+ select(DagBundleModel)
+ .options(selectinload(DagBundleModel.teams))
+ .where(DagBundleModel.name.in_(token.claims.dag_bundles),
DagBundleModel.active.is_(True))
+ .order_by(DagBundleModel.name)
+ )
+ return [
+ DagBundleStateResponse(
+ name=bundle.name,
+ version=bundle.version,
+ last_refreshed=bundle.last_refreshed,
+ revision=bundle.parse_revision,
+ team_name=bundle.team_name,
+ )
+ for bundle in bundles
+ ]
+
+
[email protected](
+ "/{job_id}/bundles/{bundle_name:path}/inventory",
+ responses=create_openapi_http_exception_doc(
+ [
+ (status.HTTP_403_FORBIDDEN, "The bundle is not granted, or the Job
has ended"),
+ status.HTTP_404_NOT_FOUND,
+ status.HTTP_409_CONFLICT,
+ status.HTTP_413_CONTENT_TOO_LARGE,
+ ]
+ ),
+ dependencies=[Security(require_auth, scopes=["token:dag_processor"])],
+)
+def publish_inventory(
+ job_id: int,
+ bundle_name: str,
+ body: DagBundleInventoryBody,
+ *,
+ session: SessionDep,
+ token=CurrentExecutionToken,
+) -> DagBundleInventoryResponse:
+ """Reconcile observed sources in one bundle, preserving sources whose
discovery failed."""
+ if bundle_name not in token.claims.dag_bundles:
+ raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason":
"bundle_not_granted"})
+ source_key = hashlib.sha256(json.dumps(["inventory",
bundle_name]).encode()).hexdigest()
+ payload_hash = _compute_payload_hash(body)
+ _validate_job(job_id, token, session=session)
+ checkpoint = session.get(DagParseCheckpoint, (job_id, source_key))
+ if checkpoint is not None:
+ if checkpoint.attempt_id == body.attempt_id:
+ if checkpoint.payload_hash != payload_hash or
checkpoint.bundle_revision is None:
+ raise HTTPException(status.HTTP_409_CONFLICT,
detail={"reason": "publication_conflict"})
+ return DagBundleInventoryResponse(
+ attempt_id=checkpoint.attempt_id,
+ accepted_at=checkpoint.accepted_at,
+ revision=checkpoint.bundle_revision,
+ )
+ if body.dispatch_sequence <= checkpoint.dispatch_sequence:
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"publication_superseded"})
+ bundle = session.scalar(
+ select(DagBundleModel).where(DagBundleModel.name ==
bundle_name).with_for_update(key_share=True)
+ )
+ if bundle is None or not bundle.active:
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"bundle_inactive"})
+ if body.expected_revision != bundle.parse_revision:
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"source_changed"})
+ inventory_hash = hashlib.sha256(
+ json.dumps([body.version, sorted(body.files),
sorted(body.failed_sources)]).encode()
+ ).hexdigest()
+ if bundle.inventory_hash != inventory_hash or bundle.parse_revision is
None:
+ bundle.parse_revision = uuid7()
+ bundle.inventory_hash = inventory_hash
+ bundle.version = body.version
+ bundle.last_refreshed = timezone.utcnow()
+ present = set(body.files)
+ failed_sources = set(body.failed_sources)
+
+ # Like DagModel.deactivate_deleted_dags, a NULL location (a 2.x row never
parsed since) is not present.
+ def is_removed(relative_fileloc: str | None) -> bool:
+ return relative_fileloc not in present and not
is_source_discovery_failed(
+ relative_fileloc, failed_sources
+ )
+
+ # Lock only the Dags being deactivated, so the rest of the bundle stays
available to the scheduler.
+ removed = sorted(
+ dag_id
+ for dag_id, relative_fileloc in session.execute(
+ select(DagModel.dag_id, DagModel.relative_fileloc).where(
+ DagModel.bundle_name == bundle_name,
DagModel.is_stale.is_(False)
+ )
+ )
+ if is_removed(relative_fileloc)
+ )
+ deactivated = False
+ for start in range(0, len(removed), 500):
+ for dag in session.scalars(
+ select(DagModel)
+ .where(
+ DagModel.dag_id.in_(removed[start : start + 500]),
+ DagModel.bundle_name == bundle_name,
+ DagModel.is_stale.is_(False),
+ )
+ .order_by(DagModel.dag_id)
+ .with_for_update(key_share=True)
+ ):
+ if is_removed(dag.relative_fileloc):
+ dag.is_stale = True
+ deactivated = True
+ if deactivated:
+ session.flush()
+ remove_references_to_deleted_dags(session=session)
+ error_ids = [
+ error.id
+ for error in session.scalars(
+ select(ParseImportError).where(ParseImportError.bundle_name ==
bundle_name)
Review Comment:
This loads whole `ParseImportError` rows, stacktraces included, on every
inventory publication, but only reads `id` and `filename`.
`select(ParseImportError.id, ParseImportError.filename)` would skip that. The
DB-mode cleanup in `manager.py` uses `load_only(ParseImportError.filename)` for
the same reason.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_parsing.py:
##########
@@ -0,0 +1,413 @@
+# 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.
+
+from __future__ import annotations
+
+import hashlib
+import json
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any
+
+from cadwyn import VersionedAPIRouter
+from fastapi import HTTPException, Request, Security, status
+from sqlalchemy import delete, or_, select
+from sqlalchemy.orm import selectinload
+from starlette.requests import ClientDisconnect
+from uuid6 import uuid7
+
+from airflow._shared.timezones import timezone
+from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
+from airflow.api_fastapi.execution_api.datamodels.dag_parsing import (
+ MAX_PUBLICATION_BYTES,
+ DagBundleInventoryBody,
+ DagBundleInventoryResponse,
+ DagBundleStateResponse,
+ DagParseResultBody,
+ DagParseResultResponse,
+)
+from airflow.api_fastapi.execution_api.security import CurrentExecutionToken,
ExecutionAPIRoute, require_auth
+from airflow.configuration import conf
+from airflow.dag_processing.collection import (
+ DagBundleOwnershipError,
+ DagPublicationValidationError,
+ _reject_other_teams_plugin_classes,
+ is_source_discovery_failed,
+ update_dag_parsing_results_in_db,
+ validate_dag_bundle_ownership,
+)
+from airflow.exceptions import DeserializationError
+from airflow.jobs.job import Job
+from airflow.models.asset import remove_references_to_deleted_dags
+from airflow.models.dag import DagModel
+from airflow.models.dag_parse_checkpoint import DagParseCheckpoint
+from airflow.models.dagbundle import DagBundleModel
+from airflow.models.dagwarning import DagWarning
+from airflow.models.errors import ParseImportError
+from airflow.sdk.importers.base import DagSourceCode # noqa: SDK001
+from airflow.serialization.serialized_objects import DagSerialization,
LazyDeserializedDAG
+
+if TYPE_CHECKING:
+ from collections.abc import Callable, Coroutine
+
+ from pydantic import BaseModel
+ from sqlalchemy.orm import Session
+ from starlette.responses import Response
+ from starlette.types import Message
+
+
+class DagParsingRoute(ExecutionAPIRoute):
+ """Refuse request bodies over the publication limit before they are read
into memory."""
+
+ def get_route_handler(self) -> Callable[[Request], Coroutine[Any, Any,
Response]]:
+ handler = super().get_route_handler()
+
+ async def handle_bounded_request(request: Request) -> Response:
+ too_large = HTTPException(
+ status.HTTP_413_CONTENT_TOO_LARGE,
+ detail=f"Request body exceeds {MAX_PUBLICATION_BYTES // (1024
* 1024)} MiB",
+ )
+ declared = request.headers.get("content-length", "")
+ if declared.isdigit() and int(declared) > MAX_PUBLICATION_BYTES:
+ raise too_large
+ body = bytearray()
+ try:
+ async for chunk in request.stream():
+ body += chunk
+ if len(body) > MAX_PUBLICATION_BYTES:
+ raise too_large
+ except ClientDisconnect as error:
+ # FastAPI reports a body it could not read as a bad request
rather than a server error.
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST, detail="The client
disconnected while sending the body"
+ ) from error
+ replayed = False
+
+ async def receive() -> Message:
+ nonlocal replayed
+ if replayed:
+ return await request.receive()
+ replayed = True
+ return {"type": "http.request", "body": bytes(body),
"more_body": False}
+
+ return await handler(Request(request.scope, receive))
+
+ return handle_bounded_request
+
+
+router = VersionedAPIRouter(route_class=DagParsingRoute)
+
+
+def _compute_payload_hash(body: BaseModel) -> str:
+ """Hash a request body independently of key order; a non-JSON float raises
ValueError."""
+ payload = json.dumps(body.model_dump(mode="json"), sort_keys=True,
allow_nan=False)
+ return hashlib.sha256(payload.encode()).hexdigest()
+
+
+def _validate_job(job_id: int, token, *, session: Session, for_update: bool =
True) -> None:
+ if job_id != token.claims.job_id:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, detail={"reason":
"not_found"})
+ query = select(Job).where(Job.id == job_id)
+ # Publications lock the Job to serialize checkpoint creation and fence
completion or replacement.
+ job = session.scalar(query.with_for_update() if for_update else query)
+ if job is None or job.session_id != token.id or job.end_date is not None:
+ raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason":
"job_closed"})
+
+
[email protected](
+ "/{job_id}/parse-results",
+ responses=create_openapi_http_exception_doc(
+ [
+ (status.HTTP_403_FORBIDDEN, "The bundle or a Dag's plugin is not
granted, or the Job has ended"),
+ status.HTTP_404_NOT_FOUND,
+ status.HTTP_409_CONFLICT,
+ status.HTTP_413_CONTENT_TOO_LARGE,
+ status.HTTP_422_UNPROCESSABLE_CONTENT,
+ ]
+ ),
+ dependencies=[Security(require_auth, scopes=["token:dag_processor"])],
+)
+def publish_parse_result(
+ job_id: int, body: DagParseResultBody, *, session: SessionDep,
token=CurrentExecutionToken
+) -> DagParseResultResponse:
+ """Atomically publish a file's results and a receipt, without reading the
processor's filesystem."""
+ if body.bundle_name not in token.claims.dag_bundles:
+ raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason":
"bundle_not_granted"})
+ try:
+ payload_hash = _compute_payload_hash(body)
+ except ValueError:
+ raise HTTPException(status.HTTP_422_UNPROCESSABLE_CONTENT,
detail="Invalid JSON value")
+ source_key = hashlib.sha256(
+ json.dumps(["parse", body.bundle_name, body.relative_fileloc]).encode()
+ ).hexdigest()
+
+ _validate_job(job_id, token, session=session)
+ checkpoint = session.get(DagParseCheckpoint, (job_id, source_key))
+ if checkpoint is not None:
+ if checkpoint.attempt_id == body.attempt_id:
+ if checkpoint.payload_hash != payload_hash:
+ raise HTTPException(status.HTTP_409_CONFLICT,
detail={"reason": "publication_conflict"})
+ return DagParseResultResponse(
+ attempt_id=checkpoint.attempt_id,
accepted_at=checkpoint.accepted_at
+ )
+ if body.dispatch_sequence <= checkpoint.dispatch_sequence:
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"publication_superseded"})
+
+ bundle = session.scalar(
+ select(DagBundleModel).where(DagBundleModel.name ==
body.bundle_name).with_for_update(key_share=True)
+ )
+ if bundle is None or not bundle.active:
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"bundle_inactive"})
+ if body.bundle_revision != bundle.parse_revision or (
+ bundle.parse_revision is not None and body.bundle_version !=
bundle.version
+ ):
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"source_changed"})
+ dags = [LazyDeserializedDAG(data=document) for document in
body.serialized_dags]
+ dag_ids = [dag.dag_id for dag in dags]
+ existing = session.scalars(
+ select(DagModel)
+ .where(DagModel.dag_id.in_(dag_ids))
+ .order_by(DagModel.dag_id)
+ .with_for_update(key_share=True)
+ ).all()
+ try:
+ validate_dag_bundle_ownership(existing, body.bundle_name,
session=session)
+ except DagBundleOwnershipError as error:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT, detail={"reason":
"dag_owned_by_another_bundle"}
+ ) from error
+ if len(_reject_other_teams_plugin_classes(body.bundle_name, dags, {},
session=session)) != len(dags):
+ raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason":
"plugin_not_granted"})
+ try:
+ for dag in dags:
+ DagSerialization.validate_serialized_dag(dag.data)
+ except (DeserializationError, ValueError, TypeError, KeyError) as error:
+ raise HTTPException(status.HTTP_422_UNPROCESSABLE_CONTENT,
detail="Invalid serialized Dag") from error
Review Comment:
When this fails, `str(error)` already names the Dag and the cause (a cycle,
a schema path, a timetable that raises on load), but the detail here is a bare
string. So the processor's `get_error_reason` finds no reason, and the import
error it records in #74794 is just "the API server rejected it (HTTP 422)".
Could this return `detail={"reason": "invalid_dag", "message": str(error)}`
like the `DagPublicationValidationError` branch below? The processor already
appends `message` for that reason. It also means a custom timetable that raises
on load is rejected here, before `_reject_dags_with_unloadable_timetables` gets
to record its per-Dag traceback.
##########
airflow-core/src/airflow/models/dag_parse_checkpoint.py:
##########
@@ -0,0 +1,46 @@
+# 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.
+
+from __future__ import annotations
+
+from datetime import datetime
+from uuid import UUID
+
+from sqlalchemy import BigInteger, ForeignKey, String, Uuid
+from sqlalchemy.orm import Mapped, mapped_column
+
+from airflow.models.base import Base
+from airflow.utils.sqlalchemy import UtcDateTime
+
+
+class DagParseCheckpoint(Base):
+ """
+ Latest accepted publication per Job and source.
+
+ Publications need an open Job, so its checkpoints are deleted when the Job
completes or is replaced.
+ The cascade covers a Job purged while still open.
+ """
+
+ __tablename__ = "dag_parse_checkpoint"
Review Comment:
This table also holds the inventory receipts (the `["inventory",
bundle_name]` source key), so `dag_parse_checkpoint` describes only half of it.
Worth renaming while the migration is still unreleased? Something like
`dag_processor_publication_receipt`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
