diff --git a/airflow/api_fastapi/core_api/datamodels/dag_run.py b/airflow/api_fastapi/core_api/datamodels/dag_run.py index 77727406bcea9..a9bc630c8368b 100644 --- a/airflow/api_fastapi/core_api/datamodels/dag_run.py +++ b/airflow/api_fastapi/core_api/datamodels/dag_run.py @@ -19,15 +19,20 @@ from datetime import datetime from enum import Enum +from typing import TYPE_CHECKING from pydantic import AwareDatetime, Field, NonNegativeInt, model_validator from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel from airflow.models import DagRun +from airflow.timetables.base import DataInterval from airflow.utils import timezone from airflow.utils.state import DagRunState from airflow.utils.types import DagRunTriggeredByType, DagRunType +if TYPE_CHECKING: + from airflow.models import DAG + class DAGRunPatchStates(str, Enum): """Enum for DAG Run states when updating a DAG Run.""" @@ -99,6 +104,36 @@ def check_data_intervals(cls, values): ) return values + def validate_context(self, dag: DAG) -> dict: + coerced_logical_date = timezone.coerce_datetime(self.logical_date) + run_after = self.run_after + data_interval = None + if coerced_logical_date: + if self.data_interval_start and self.data_interval_end: + data_interval = DataInterval( + start=timezone.coerce_datetime(self.data_interval_start), + end=timezone.coerce_datetime(self.data_interval_end), + ) + else: + data_interval = dag.timetable.infer_manual_data_interval( + run_after=coerced_logical_date or timezone.coerce_datetime(self.run_after) + ) + run_after = data_interval.end + + run_id = self.dag_run_id or DagRun.generate_run_id( + run_type=DagRunType.SCHEDULED, + logical_date=coerced_logical_date, + run_after=self.run_after, + ) + return { + "run_id": run_id, + "logical_date": coerced_logical_date, + "data_interval": data_interval, + "run_after": run_after, + "conf": self.conf, + "note": self.note, + } + @model_validator(mode="after") def validate_dag_run_id(self): if not self.dag_run_id: diff --git a/airflow/api_fastapi/core_api/routes/public/dag_run.py b/airflow/api_fastapi/core_api/routes/public/dag_run.py index fc4397447da12..ad556b00da84d 100644 --- a/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -19,7 +19,6 @@ from typing import Annotated, Literal, cast -import pendulum from fastapi import Depends, HTTPException, Query, Request, status from fastapi.exceptions import RequestValidationError from pydantic import ValidationError @@ -65,8 +64,6 @@ from airflow.listeners.listener import get_listener_manager from airflow.models import DAG, DagModel, DagRun from airflow.models.dag_version import DagVersion -from airflow.timetables.base import DataInterval -from airflow.utils import timezone from airflow.utils.state import DagRunState from airflow.utils.types import DagRunTriggeredByType, DagRunType @@ -358,38 +355,16 @@ def trigger_dag_run( f"DAG with dag_id: '{dag_id}' has import errors and cannot be triggered", ) - logical_date = timezone.coerce_datetime(body.logical_date) - coerced_logical_date = timezone.coerce_datetime(logical_date) - run_after = timezone.coerce_datetime(body.run_after) - try: dag: DAG = request.app.state.dag_bag.get_dag(dag_id) - data_interval = None - if body.logical_date: - if body.data_interval_start and body.data_interval_end: - data_interval = DataInterval( - start=pendulum.instance(body.data_interval_start), - end=pendulum.instance(body.data_interval_end), - ) - else: - data_interval = dag.timetable.infer_manual_data_interval( - run_after=coerced_logical_date or run_after - ) - run_after = data_interval.end - - if body.dag_run_id: - run_id = body.dag_run_id - else: - run_id = DagRun.generate_run_id( - run_type=DagRunType.SCHEDULED, logical_date=coerced_logical_date, run_after=run_after - ) + params = body.validate_context(dag) dag_run = dag.create_dagrun( - run_id=run_id, - logical_date=coerced_logical_date, - data_interval=data_interval, - run_after=run_after, - conf=body.conf, + run_id=params["run_id"], + logical_date=params["logical_date"], + data_interval=params["data_interval"], + run_after=params["run_after"], + conf=params["conf"], run_type=DagRunType.MANUAL, triggered_by=DagRunTriggeredByType.REST_API, external_trigger=True,