ephraimbuddy commented on code in PR #74790: URL: https://github.com/apache/airflow/pull/74790#discussion_r4422278505
########## 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: Yes, that's the intent. `failed` is now `body.relative_fileloc` in `body.import_errors`, and the comment says only container members get the grace period. -- 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]
