diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index d0362ce12aa61..20d68de781a60 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -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) 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) diff --git a/task-sdk/src/airflow/sdk/exceptions.py b/task-sdk/src/airflow/sdk/exceptions.py index b69abe6226575..065bf4f399af5 100644 --- a/task-sdk/src/airflow/sdk/exceptions.py +++ b/task-sdk/src/airflow/sdk/exceptions.py @@ -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" diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 73e8fd0cbf62e..f94503ab5b163 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -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} @@ -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( diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 0df8839c55f30..205a1f9152de2 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -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""" diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 5aeb009bd33da..d50a48acc27d3 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -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) + + assert state == TaskInstanceState.FAILED + assert msg.state == TaskInstanceState.FAILED + @pytest.mark.parametrize( ("allowed_states", "failed_states", "target_dr_state", "expected_task_state"), [