diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index ea220c68fde71..4952e285888c8 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -941,10 +941,35 @@ def trigger( return OKResponse(ok=True) - def clear(self, dag_id: str, run_id: str) -> OKResponse: + def clear(self, dag_id: str, run_id: str) -> OKResponse | ErrorResponse: """Clear a Dag run via the API server.""" - self.client.post(f"dag-runs/{dag_id}/{run_id}/clear") - # TODO: Error handling + try: + self.client.post(f"dag-runs/{dag_id}/{run_id}/clear") + except ServerResponseError as e: + detail = { + "dag_id": dag_id, + "run_id": run_id, + "status_code": e.response.status_code, + "server_detail": e.detail, + } + if e.response.status_code == HTTPStatus.NOT_FOUND: + error = ErrorType.DAGRUN_CLEAR_FAILED + if isinstance(e.detail, dict): + message = str(e.detail.get("message", "")) + if message.startswith("Dag run "): + error = ErrorType.DAGRUN_NOT_FOUND + elif message.startswith("Dag with dag_id"): + error = ErrorType.DAG_NOT_FOUND + log.info("Dag Run clear target not found.", detail=e.detail, dag_id=dag_id, run_id=run_id) + return ErrorResponse(error=error, detail=detail) + if e.response.status_code in ( + HTTPStatus.BAD_REQUEST, + HTTPStatus.UNPROCESSABLE_ENTITY, + ): + log.info("Dag Run clear failed.", detail=e.detail, dag_id=dag_id, run_id=run_id) + return ErrorResponse(error=ErrorType.DAGRUN_CLEAR_FAILED, detail=detail) + raise + return OKResponse(ok=True) def get_detail(self, dag_id: str, run_id: str) -> DagRun: diff --git a/task-sdk/src/airflow/sdk/exceptions.py b/task-sdk/src/airflow/sdk/exceptions.py index 6f43d5421ecf2..fe7fcfbcaaf86 100644 --- a/task-sdk/src/airflow/sdk/exceptions.py +++ b/task-sdk/src/airflow/sdk/exceptions.py @@ -94,6 +94,9 @@ class ErrorType(enum.Enum): TASK_STORE_NOT_FOUND = "TASK_STORE_NOT_FOUND" ASSET_STORE_NOT_FOUND = "ASSET_STORE_NOT_FOUND" DAGRUN_ALREADY_EXISTS = "DAGRUN_ALREADY_EXISTS" + DAG_NOT_FOUND = "DAG_NOT_FOUND" + DAGRUN_NOT_FOUND = "DAGRUN_NOT_FOUND" + DAGRUN_CLEAR_FAILED = "DAGRUN_CLEAR_FAILED" # Distinct from API_SERVER_ERROR: signals an explicit 401/403 from the # Execution API. Callers like ExecutionAPISecretsBackend treat this as # a deny rather than a "not found" so the secrets-backend dispatcher diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 85f8b7cfa4e2b..a542e3720c993 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1865,6 +1865,9 @@ "TASK_STORE_NOT_FOUND", "ASSET_STORE_NOT_FOUND", "DAGRUN_ALREADY_EXISTS", + "DAG_NOT_FOUND", + "DAGRUN_NOT_FOUND", + "DAGRUN_CLEAR_FAILED", "PERMISSION_DENIED", "GENERIC_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 eeb2788062248..6106b378c4713 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1979,6 +1979,9 @@ def _handle_trigger_dag_run( return msg, state + if isinstance(comms_msg, ErrorResponse): + raise AirflowRuntimeError(comms_msg) + log.info("Dag Run triggered successfully.", trigger_dag_id=drte.trigger_dag_id) # Store the run id from the dag run (either created or found above) to diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 684d4e7dcbb17..6d80947434188 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -1507,7 +1507,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: class TestDagRunOperations: def test_trigger(self): - # Simulate a successful response from the server when triggering a dag run + # Simulate a successful response from the server when triggering a Dag run def handle_request(request: httpx.Request) -> httpx.Response: if request.url.path == "/dag-runs/test_trigger/test_run_id": actual_body = json.loads(request.read()) @@ -1551,7 +1551,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert result == ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) 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""" + """Test that if Dag run already exists and reset_dag_run=True, the client clears the Dag run""" def handle_request(request: httpx.Request) -> httpx.Response: if request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id": @@ -1577,8 +1577,54 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert result == OKResponse(ok=True) + def test_trigger_conflict_reset_dag_run_clear_failure(self): + """Test that reset_dag_run=True returns the clear failure response.""" + + def handle_request(request: httpx.Request) -> httpx.Response: + if request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id": + return httpx.Response( + status_code=409, + json={ + "detail": { + "reason": "already_exists", + "message": "A Dag Run already exists for Dag test_trigger_conflict", + } + }, + ) + if request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id/clear": + return httpx.Response( + status_code=404, + json={ + "detail": { + "reason": "not_found", + "message": "Dag run with run_id: 'test_run_id' not found", + } + }, + ) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger( + dag_id="test_trigger_conflict_reset", + run_id="test_run_id", + reset_dag_run=True, + ) + + assert result == ErrorResponse( + error=ErrorType.DAGRUN_NOT_FOUND, + detail={ + "dag_id": "test_trigger_conflict_reset", + "run_id": "test_run_id", + "status_code": 404, + "server_detail": { + "reason": "not_found", + "message": "Dag run with run_id: 'test_run_id' not found", + }, + }, + ) + def test_clear(self): - """Test that the client can clear a dag run""" + """Test that the client can clear a Dag run""" def handle_request(request: httpx.Request) -> httpx.Response: if request.url.path == "/dag-runs/test_clear/test_run_id/clear": @@ -1590,8 +1636,58 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert result == OKResponse(ok=True) + @pytest.mark.parametrize( + ("status_code", "server_detail", "expected_error"), + [ + pytest.param( + 404, + {"reason": "not_found", "message": "Dag run with run_id: 'test_run_id' not found"}, + ErrorType.DAGRUN_NOT_FOUND, + id="dag-run-not-found", + ), + pytest.param( + 404, + {"reason": "not_found", "message": "Dag with dag_id: 'test_clear' not found"}, + ErrorType.DAG_NOT_FOUND, + id="dag-not-found", + ), + pytest.param( + 400, + {"reason": "import_errors", "message": "Dag with dag_id 'test_clear' has import errors"}, + ErrorType.DAGRUN_CLEAR_FAILED, + id="bad-request", + ), + pytest.param( + 422, + {"reason": "validation_error", "message": "Invalid payload"}, + ErrorType.DAGRUN_CLEAR_FAILED, + id="validation-error", + ), + ], + ) + def test_clear_error_response(self, status_code, server_detail, expected_error): + """Test that clear returns semantic errors for known clear failures.""" + + def handle_request(request: httpx.Request) -> httpx.Response: + if request.url.path == "/dag-runs/test_clear/test_run_id/clear": + return httpx.Response(status_code=status_code, json={"detail": server_detail}) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.clear(dag_id="test_clear", run_id="test_run_id") + + assert result == ErrorResponse( + error=expected_error, + detail={ + "dag_id": "test_clear", + "run_id": "test_run_id", + "status_code": status_code, + "server_detail": server_detail, + }, + ) + def test_get_state(self): - """Test that the client can get the state of a dag run""" + """Test that the client can get the state of a Dag run""" def handle_request(request: httpx.Request) -> httpx.Response: if request.url.path == "/dag-runs/test_state/test_run_id/state": 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 cc9fb77e08921..2772c05f3e5c7 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 @@ -5406,6 +5406,37 @@ def _send(msg=None, **kwargs): assert state == TaskInstanceState.FAILED assert isinstance(error, _TriggerSendError) + @time_machine.travel("2025-01-01 00:00:00", tick=False) + def test_handle_trigger_dag_run_error_response_fails_task(self, create_runtime_ti, mock_supervisor_comms): + """Test that semantic trigger errors do not continue down the success path.""" + from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator + + task = TriggerDagRunOperator( + task_id="test_task", + trigger_dag_id="test_dag", + trigger_run_id="test_run_id", + reset_dag_run=True, + ) + ti = create_runtime_ti( + dag_id="test_handle_trigger_dag_run_error_response", + run_id="test_run", + task=task, + ) + + mock_supervisor_comms.send.return_value = ErrorResponse( + error=ErrorType.DAGRUN_CLEAR_FAILED, + detail={"dag_id": "test_dag", "run_id": "test_run_id", "status_code": 400}, + ) + + log = mock.MagicMock(spec=structlog.typing.FilteringBoundLogger) + state, msg, error = run(ti, ti.get_template_context(), log) + + assert state == TaskInstanceState.FAILED + assert isinstance(msg, TaskState) + assert msg.state == TaskInstanceState.FAILED + assert isinstance(error, AirflowRuntimeError) + assert error.error.error == ErrorType.DAGRUN_CLEAR_FAILED + @pytest.mark.parametrize( ("should_retry", "expected_state"), [