Skip to content

fix(task-sdk): handle missing DAG in TriggerDagRunOperator gracefully #63207

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

Closed
YoannAbriel wants to merge 2 commits into from
Closed

fix(task-sdk): handle missing DAG in TriggerDagRunOperator gracefully #63207

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cfcc2d1
fix(task-sdk): handle missing DAG in TriggerDagRunOperator gracefully
YoannAbriel Mar 9, 2026
aa9068e
fix: ruff format test file
YoannAbriel Mar 10, 2026
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
3 changes: 3 additions & 0 deletions task-sdk/src/airflow/sdk/api/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -724,6 +724,9 @@ def trigger(
f"dag-runs/{dag_id}/{run_id}", content=body.model_dump_json(exclude_defaults=True)
)
except ServerResponseError as e:
if e.response.status_code == HTTPStatus.NOT_FOUND:
log.error("Dag not found.", dag_id=dag_id)
return ErrorResponse(error=ErrorType.DAG_NOT_FOUND)
Comment on lines +728 to +729

Copilot AI Apr 10, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new 404 handling logs only a generic "Dag not found" message and doesn’t include available response context (e.detail, status_code). Consider logging those fields (and/or returning them in ErrorResponse.detail) for easier debugging, consistent with other client methods’ 404 handling.

Suggested change
log.error("Dag not found.", dag_id=dag_id)
return ErrorResponse(error=ErrorType.DAG_NOT_FOUND)
log.error(
"Dag not found.",
dag_id=dag_id,
status_code=e.response.status_code,
detail=e.detail,
)
return ErrorResponse(error=ErrorType.DAG_NOT_FOUND, detail=e.detail)

Copilot uses AI. Check for mistakes.
if e.response.status_code == HTTPStatus.CONFLICT:
if reset_dag_run:
log.info("Dag Run already exists; Resetting Dag Run.", dag_id=dag_id, run_id=run_id)
Expand Down
1 change: 1 addition & 0 deletions task-sdk/src/airflow/sdk/exceptions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@ class ErrorType(enum.Enum):
VARIABLE_NOT_FOUND = "VARIABLE_NOT_FOUND"
XCOM_NOT_FOUND = "XCOM_NOT_FOUND"
ASSET_NOT_FOUND = "ASSET_NOT_FOUND"
DAG_NOT_FOUND = "DAG_NOT_FOUND"
DAGRUN_ALREADY_EXISTS = "DAGRUN_ALREADY_EXISTS"
GENERIC_ERROR = "GENERIC_ERROR"
API_SERVER_ERROR = "API_SERVER_ERROR"
Expand Down
11 changes: 10 additions & 1 deletion task-sdk/src/airflow/sdk/execution_time/task_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1240,7 +1240,7 @@ def _on_term(signum, frame):
signal.signal(signal.SIGTERM, _on_term)

msg: ToSupervisor | None = None
state: TaskInstanceState
state: TaskInstanceState = TaskInstanceState.FAILED
error: BaseException | None = None

stats_tags = {"dag_id": ti.dag_id, "task_id": ti.task_id}
Expand Down Expand Up @@ -1441,6 +1441,15 @@ def _handle_trigger_dag_run(
),
)

if isinstance(comms_msg, ErrorResponse) and comms_msg.error == ErrorType.DAG_NOT_FOUND:
log.error("Dag not found, marking task as failed.", dag_id=drte.trigger_dag_id)
msg = TaskState(
state=TaskInstanceState.FAILED,
end_date=datetime.now(tz=timezone.utc),
rendered_map_index=ti.rendered_map_index,
)
return msg, TaskInstanceState.FAILED

if isinstance(comms_msg, ErrorResponse) and comms_msg.error == ErrorType.DAGRUN_ALREADY_EXISTS:
if drte.skip_when_already_exists:
log.info(
Expand Down
21 changes: 21 additions & 0 deletions task-sdk/tests/task_sdk/api/test_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1264,6 +1264,27 @@ def handle_request(request: httpx.Request) -> httpx.Response:

assert result == ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS)

def test_trigger_dag_not_found(self):
"""Test that if the target dag does not exist, the client returns a DAG_NOT_FOUND error."""

def handle_request(request: httpx.Request) -> httpx.Response:
if request.url.path == "/dag-runs/nonexistent_dag/test_run_id":
return httpx.Response(
status_code=404,
json={
"detail": {
"reason": "not_found",
"message": "Dag with dag_id: 'nonexistent_dag' not found",
}
},
)
return httpx.Response(status_code=422)

client = make_client(transport=httpx.MockTransport(handle_request))
result = client.dag_runs.trigger(dag_id="nonexistent_dag", run_id="test_run_id")

assert result == ErrorResponse(error=ErrorType.DAG_NOT_FOUND)

def test_trigger_conflict_reset_dag_run(self):
"""Test that if dag run already exists and reset_dag_run=True, the client clears the dag run"""

Expand Down
19 changes: 19 additions & 0 deletions task-sdk/tests/task_sdk/execution_time/test_task_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4448,6 +4448,25 @@ def test_handle_trigger_dag_run_conflict(
]
mock_supervisor_comms.assert_has_calls(expected_calls)

@time_machine.travel("2025-01-01 00:00:00", tick=False)
def test_handle_trigger_dag_run_dag_not_found(self, create_runtime_ti, mock_supervisor_comms):
"""Test that TriggerDagRunOperator fails gracefully when the target DAG doesn't exist."""
from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator

task = TriggerDagRunOperator(
task_id="test_task",
trigger_dag_id="nonexistent_dag",
trigger_run_id="test_run_id",
)
ti = create_runtime_ti(dag_id="test_handle_trigger_dag_run_not_found", run_id="test_run", task=task)

log = mock.MagicMock()
mock_supervisor_comms.send.return_value = ErrorResponse(error=ErrorType.DAG_NOT_FOUND)
state, msg, _ = run(ti, ti.get_template_context(), log)
Comment on lines +4462 to +4465

Copilot AI Apr 10, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

New test introduces an unspecced mock.MagicMock() for log. Airflow’s testing guidelines prefer mocks with spec/autospec to avoid silently accepting unexpected attributes; consider using a real structlog.get_logger(...) or mock.Mock(spec=...) for the logger instead.

Copilot generated this review using guidance from repository custom instructions.

assert state == TaskInstanceState.FAILED
assert msg.state == TaskInstanceState.FAILED

@pytest.mark.parametrize(
("allowed_states", "failed_states", "target_dr_state", "expected_task_state"),
[
Expand Down
Loading