diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 1da539f29a3dc..7165b5e6dfa84 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -899,17 +899,49 @@ def trigger( run_after=run_after, ) + # GET dag-runs/count is available since Airflow 3.0.0. Failures propagate rather + # than falling through to POST — if the count endpoint is unreachable, the POST + # would fail too. None signals dry-run (skip pre-check, no-op POST). + dag_run_count_before = ( + None if self.client._dry_run else self.get_count(dag_id=dag_id, run_ids=[run_id]).count + ) + if dag_run_count_before is not None and dag_run_count_before > 0: + if reset_dag_run: + log.info("Dag Run already exists; Resetting Dag Run.", dag_id=dag_id, run_id=run_id) + # TODO: Make clear() idempotent as a follow-up. + return self.clear(run_id=run_id, dag_id=dag_id) + log.info("Dag Run already exists!", dag_id=dag_id, run_id=run_id) + return ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) + try: - self.client.post( - f"dag-runs/{dag_id}/{run_id}", content=body.model_dump_json(exclude_defaults=True) + self.client._request_without_retry( + "POST", f"dag-runs/{dag_id}/{run_id}", content=body.model_dump_json(exclude_defaults=True) ) + except (httpx.ReadError, httpx.ReadTimeout, httpx.RemoteProtocolError): + if dag_run_count_before == 0 and self.get_count(dag_id=dag_id, run_ids=[run_id]).count > 0: + log.info( + "Dag Run exists after ambiguous trigger response; treating trigger as successful.", + dag_id=dag_id, + run_id=run_id, + ) + return OKResponse(ok=True) + raise except ServerResponseError as e: 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) + log.info( + "Dag Run already exists after trigger attempt; Resetting Dag Run.", + detail=e.detail, + dag_id=dag_id, + run_id=run_id, + ) return self.clear(run_id=run_id, dag_id=dag_id) - - log.info("Dag Run already exists!", detail=e.detail, dag_id=dag_id, run_id=run_id) + log.info( + "Dag Run already exists after trigger attempt.", + detail=e.detail, + dag_id=dag_id, + run_id=run_id, + ) return ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) raise @@ -1116,6 +1148,7 @@ def __init__(self, *, base_url: str | None, dry_run: bool = False, token: str, * if (not base_url) ^ dry_run: raise ValueError(f"Can only specify one of {base_url=} or {dry_run=}") auth = BearerAuth(token) + self._dry_run: bool = dry_run if dry_run: # If dry run is requested, install a no op handler so that simple tasks can "heartbeat" using a @@ -1155,6 +1188,19 @@ def _update_auth(self, response: httpx.Response): log.debug("Execution API issued us a refreshed Task token") self.auth = BearerAuth(new_token) + @staticmethod + def _ensure_json_content_type(kwargs: dict[str, Any]) -> None: + # Set content type as convenience if not already set + if kwargs.get("content", None) is not None and "content-type" not in ( + kwargs.get("headers", {}) or {} + ): + kwargs["headers"] = {"content-type": "application/json"} + + def _request_without_retry(self, *args, **kwargs): + """Implement a convenience for httpx.Client.request without retrying.""" + self._ensure_json_content_type(kwargs) + return super().request(*args, **kwargs) + @retry( retry=retry_if_exception(_should_retry_api_request), stop=stop_after_attempt(API_RETRIES), @@ -1164,12 +1210,7 @@ def _update_auth(self, response: httpx.Response): ) def request(self, *args, **kwargs): """Implement a convenience for httpx.Client.request with a retry layer.""" - # Set content type as convenience if not already set - if kwargs.get("content", None) is not None and "content-type" not in ( - kwargs.get("headers", {}) or {} - ): - kwargs["headers"] = {"content-type": "application/json"} - + self._ensure_json_content_type(kwargs) return super().request(*args, **kwargs) # We "group" or "namespace" operations by what they operate on, rather than a flat namespace with all diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 93e2041878e01..2c3a13364085d 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -1288,8 +1288,15 @@ 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 + requests: list[tuple[str, str]] = [] + def handle_request(request: httpx.Request) -> httpx.Response: + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + assert request.url.params["dag_id"] == "test_trigger" + assert request.url.params["run_ids"] == "test_run_id" + return httpx.Response(status_code=200, json=0) if request.url.path == "/dag-runs/test_trigger/test_run_id": actual_body = json.loads(request.read()) assert actual_body["logical_date"] == "2025-01-01T00:00:00Z" @@ -1309,18 +1316,339 @@ def handle_request(request: httpx.Request) -> httpx.Response: ) assert result == OKResponse(ok=True) + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger/test_run_id"), + ] + + def test_trigger_dry_run_skips_precheck_conflict(self): + client = make_client_w_dry_run() + result = client.dag_runs.trigger(dag_id="test_trigger", run_id="test_run_id") + + assert result == OKResponse(ok=True) + + def test_trigger_pre_existing_dag_run_returns_conflict_without_posting(self): + requests: list[tuple[str, str]] = [] + + def handle_request(request: httpx.Request) -> httpx.Response: + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=1) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id": + return httpx.Response(status_code=500, json={"detail": "POST should not happen"}) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger(dag_id="test_trigger", run_id="test_run_id") + + assert result == ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) + assert requests == [ + ("GET", "/dag-runs/count"), + ] + + def test_trigger_clears_pre_existing_dag_run_without_posting_when_resetting(self): + requests: list[tuple[str, str]] = [] + + def handle_request(request: httpx.Request) -> httpx.Response: + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=1) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id": + return httpx.Response(status_code=500, json={"detail": "POST should not happen"}) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id/clear": + return httpx.Response(status_code=204) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger(dag_id="test_trigger", run_id="test_run_id", reset_dag_run=True) + + assert result == OKResponse(ok=True) + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger/test_run_id/clear"), + ] + + def test_trigger_returns_conflict_from_post_when_run_was_missing_before(self): + requests: list[tuple[str, str]] = [] + + def handle_request(request: httpx.Request) -> httpx.Response: + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=0) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id": + return httpx.Response( + status_code=409, + json={ + "detail": { + "reason": "already_exists", + "message": "A Dag Run already exists for Dag test_trigger with run id test_run_id", + } + }, + ) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger(dag_id="test_trigger", run_id="test_run_id") + + assert result == ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger/test_run_id"), + ] + + def test_trigger_treats_read_error_as_success_when_dag_run_appears_after_missing_precheck(self): + requests: list[tuple[str, str]] = [] + dag_run_exists = False + + def handle_request(request: httpx.Request) -> httpx.Response: + nonlocal dag_run_exists + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=1 if dag_run_exists else 0) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id": + dag_run_exists = True + raise httpx.ReadError("Trigger response was lost", request=request) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger(dag_id="test_trigger", run_id="test_run_id") + + assert result == OKResponse(ok=True) + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger/test_run_id"), + ("GET", "/dag-runs/count"), + ] + + def test_trigger_treats_read_timeout_as_success_when_dag_run_appears_after_missing_precheck(self): + requests: list[tuple[str, str]] = [] + dag_run_exists = False + + def handle_request(request: httpx.Request) -> httpx.Response: + nonlocal dag_run_exists + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=1 if dag_run_exists else 0) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id": + dag_run_exists = True + raise httpx.ReadTimeout("Trigger response timed out", request=request) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger(dag_id="test_trigger", run_id="test_run_id") + + assert result == OKResponse(ok=True) + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger/test_run_id"), + ("GET", "/dag-runs/count"), + ] + + def test_trigger_retries_followup_probe_after_ambiguous_response(self): + requests: list[tuple[str, str]] = [] + dag_run_exists = False + followup_attempts = 0 + + with time_machine.travel("2023-01-01T00:00:00Z", tick=False): + + def handle_request(request: httpx.Request) -> httpx.Response: + nonlocal dag_run_exists, followup_attempts + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + if not dag_run_exists: + return httpx.Response(status_code=200, json=0) + followup_attempts += 1 + if followup_attempts == 1: + return httpx.Response(status_code=500, json={"detail": "Internal Server Error"}) + return httpx.Response(status_code=200, json=1) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id": + dag_run_exists = True + raise httpx.ReadError("Trigger response was lost", request=request) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger(dag_id="test_trigger", run_id="test_run_id") + + assert result == OKResponse(ok=True) + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger/test_run_id"), + ("GET", "/dag-runs/count"), + ("GET", "/dag-runs/count"), + ] + + def test_trigger_treats_read_error_as_success_when_resetting_missing_dag_run_appears(self): + requests: list[tuple[str, str]] = [] + dag_run_exists = False + + def handle_request(request: httpx.Request) -> httpx.Response: + nonlocal dag_run_exists + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=1 if dag_run_exists else 0) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id": + dag_run_exists = True + raise httpx.ReadError("Trigger response was lost", request=request) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger( + dag_id="test_trigger", + run_id="test_run_id", + reset_dag_run=True, + ) + + assert result == OKResponse(ok=True) + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger/test_run_id"), + ("GET", "/dag-runs/count"), + ] + + def test_trigger_treats_remote_protocol_error_as_success_when_dag_run_appears_after_missing_precheck( + self, + ): + requests: list[tuple[str, str]] = [] + dag_run_exists = False + + def handle_request(request: httpx.Request) -> httpx.Response: + nonlocal dag_run_exists + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=1 if dag_run_exists else 0) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id": + dag_run_exists = True + raise httpx.RemoteProtocolError("Trigger response was lost", request=request) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger(dag_id="test_trigger", run_id="test_run_id") + + assert result == OKResponse(ok=True) + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger/test_run_id"), + ("GET", "/dag-runs/count"), + ] + + def test_trigger_conflict_from_post_clears_run_when_resetting(self): + requests: list[tuple[str, str]] = [] + + def handle_request(request: httpx.Request) -> httpx.Response: + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=0) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id": + return httpx.Response( + status_code=409, + json={ + "detail": { + "reason": "already_exists", + "message": "A Dag Run already exists for Dag test_trigger with run id test_run_id", + } + }, + ) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id/clear": + return httpx.Response(status_code=204) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.dag_runs.trigger( + dag_id="test_trigger", + run_id="test_run_id", + reset_dag_run=True, + ) + + assert result == OKResponse(ok=True) + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger/test_run_id"), + ("POST", "/dag-runs/test_trigger/test_run_id/clear"), + ] + + def test_trigger_reraises_read_error_when_dag_run_is_missing_after_precheck(self): + requests: list[tuple[str, str]] = [] + + def handle_request(request: httpx.Request) -> httpx.Response: + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=0) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id": + raise httpx.ReadError("Trigger response was lost", request=request) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + + with pytest.raises(httpx.ReadError, match="Trigger response was lost"): + client.dag_runs.trigger(dag_id="test_trigger", run_id="test_run_id") + + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger/test_run_id"), + ("GET", "/dag-runs/count"), + ] + + def test_trigger_reraises_connect_error_even_if_dag_run_exists_after_precheck(self): + requests: list[tuple[str, str]] = [] + + def handle_request(request: httpx.Request) -> httpx.Response: + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=0) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id": + raise httpx.ConnectError("Could not connect", request=request) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + + with pytest.raises(httpx.ConnectError, match="Could not connect"): + client.dag_runs.trigger(dag_id="test_trigger", run_id="test_run_id") + + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger/test_run_id"), + ] + + def test_trigger_reraises_pool_timeout_even_if_dag_run_exists_after_precheck(self): + requests: list[tuple[str, str]] = [] + + def handle_request(request: httpx.Request) -> httpx.Response: + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=0) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger/test_run_id": + raise httpx.PoolTimeout("Could not get a connection from the pool", request=request) + return httpx.Response(status_code=422) + + client = make_client(transport=httpx.MockTransport(handle_request)) + + with pytest.raises(httpx.PoolTimeout, match="Could not get a connection from the pool"): + client.dag_runs.trigger(dag_id="test_trigger", run_id="test_run_id") + + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger/test_run_id"), + ] def test_trigger_conflict(self): - """Test that if the dag run already exists, the client returns an error when default reset_dag_run=False""" + """Test that if the Dag run already exists, the client returns an error when default reset_dag_run=False""" + + requests: list[tuple[str, str]] = [] def handle_request(request: httpx.Request) -> httpx.Response: - if request.url.path == "/dag-runs/test_trigger_conflict/test_run_id": + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=1) + if request.method == "POST" and request.url.path == "/dag-runs/test_trigger_conflict/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 with run id test_run_id", + "message": ( + "A Dag Run already exists for Dag test_trigger_conflict " + "with run id test_run_id" + ), } }, ) @@ -1330,22 +1658,39 @@ def handle_request(request: httpx.Request) -> httpx.Response: result = client.dag_runs.trigger(dag_id="test_trigger_conflict", run_id="test_run_id") assert result == ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS) + assert requests == [ + ("GET", "/dag-runs/count"), + ] 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 the Dag run already exists and reset_dag_run=True, the client clears the Dag run""" + + requests: list[tuple[str, str]] = [] def handle_request(request: httpx.Request) -> httpx.Response: - if request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id": + requests.append((request.method, request.url.path)) + if request.method == "GET" and request.url.path == "/dag-runs/count": + return httpx.Response(status_code=200, json=1) + if ( + request.method == "POST" + and 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 with run id test_run_id", + "message": ( + "A Dag Run already exists for Dag test_trigger_conflict_reset " + "with run id test_run_id" + ), } }, ) - if request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id/clear": + if ( + request.method == "POST" + and request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id/clear" + ): return httpx.Response(status_code=204) return httpx.Response(status_code=422) @@ -1357,9 +1702,13 @@ def handle_request(request: httpx.Request) -> httpx.Response: ) assert result == OKResponse(ok=True) + assert requests == [ + ("GET", "/dag-runs/count"), + ("POST", "/dag-runs/test_trigger_conflict_reset/test_run_id/clear"), + ] 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":