Skip to content

Validation logic for the payload moved to the Serializer for trigger_dag_run #46882

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Lee-W merged 3 commits into from
Feb 20, 2025
Merged

Validation logic for the payload moved to the Serializer for trigger_dag_run #46882

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a022a48
validation logic for the payload moved to the Serializer
sunank200 Feb 19, 2025
603038f
Add type checking
sunank200 Feb 19, 2025
1dfaa23
fix the PR comments
sunank200 Feb 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions airflow/api_fastapi/core_api/datamodels/dag_run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -99,6 +104,36 @@ def check_data_intervals(cls, values):
)
return values

def validate_context(self, dag: DAG) -> dict:
Comment thread
uranusjr marked this conversation as resolved.
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:
Expand Down
37 changes: 6 additions & 31 deletions airflow/api_fastapi/core_api/routes/public/dag_run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand Down