You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2025/10/31 14:50:36 UTC
(airflow) 01/02: Migrate retry handler in task SDK API client to use tenacity instead of retryhttp (#56762)
This is an automated email from the ASF dual-hosted git repository.
kaxilnaik pushed a commit to branch v3-1-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 96ab2f0a5176dc8c137d399b4f8751ba9ef58b96
Author: Amogh Desai <am...@gmail.com>
AuthorDate: Thu Oct 23 13:21:48 2025 +0530
Migrate retry handler in task SDK API client to use tenacity instead of retryhttp (#56762)
(cherry picked from commit fa60a7a02f1a0c5485e70e5a3397cb3dac326221)
---
task-sdk/pyproject.toml | 4 -
task-sdk/src/airflow/sdk/api/client.py | 29 ++-
task-sdk/tests/task_sdk/api/test_client.py | 370 ++++++++++++++---------------
3 files changed, 195 insertions(+), 208 deletions(-)
diff --git a/task-sdk/pyproject.toml b/task-sdk/pyproject.toml
index 5089ef2f5c9..9af60cb38d8 100644
--- a/task-sdk/pyproject.toml
+++ b/task-sdk/pyproject.toml
@@ -61,11 +61,7 @@ dependencies = [
"python-dateutil>=2.7.0",
"psutil>=6.1.0",
"structlog>=25.4.0",
- "retryhttp>=1.2.0,!=1.3.0",
"greenback>=1.2.1",
- # Requests is known to introduce breaking changes, so we pin it to a specific range
- "requests>=2.31.0,<3",
- "types-requests>=2.31.0",
"tenacity>=8.3.0",
# Start of shared timezones dependencies
"pendulum>=3.1.0",
diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py
index c4441dbdebd..2a83bb8d7da 100644
--- a/task-sdk/src/airflow/sdk/api/client.py
+++ b/task-sdk/src/airflow/sdk/api/client.py
@@ -30,8 +30,13 @@ import httpx
import msgspec
import structlog
from pydantic import BaseModel
-from retryhttp import retry, wait_retry_after
-from tenacity import before_log, wait_random_exponential
+from tenacity import (
+ before_log,
+ retry,
+ retry_if_exception,
+ stop_after_attempt,
+ wait_random_exponential,
+)
from uuid6 import uuid7
from airflow.configuration import conf
@@ -813,6 +818,14 @@ API_SSL_CERT_PATH = conf.get("api", "ssl_cert")
API_TIMEOUT = conf.getfloat("workers", "execution_api_timeout")
+def _should_retry_api_request(exception: BaseException) -> bool:
+ """Determine if an API request should be retried based on the exception type."""
+ if isinstance(exception, httpx.HTTPStatusError):
+ return exception.response.status_code >= 500
+
+ return isinstance(exception, httpx.RequestError)
+
+
class Client(httpx.Client):
@lru_cache()
@staticmethod
@@ -851,21 +864,17 @@ class Client(httpx.Client):
**kwargs,
)
- _default_wait = wait_random_exponential(min=API_RETRY_WAIT_MIN, max=API_RETRY_WAIT_MAX)
-
def _update_auth(self, response: httpx.Response):
if new_token := response.headers.get("Refreshed-API-Token"):
log.debug("Execution API issued us a refreshed Task token")
self.auth = BearerAuth(new_token)
@retry(
- reraise=True,
- max_attempt_number=API_RETRIES,
- wait_server_errors=_default_wait,
- wait_network_errors=_default_wait,
- wait_timeouts=_default_wait,
- wait_rate_limited=wait_retry_after(fallback=_default_wait), # No infinite timeout on HTTP 429
+ retry=retry_if_exception(_should_retry_api_request),
+ stop=stop_after_attempt(API_RETRIES),
+ wait=wait_random_exponential(min=API_RETRY_WAIT_MIN, max=API_RETRY_WAIT_MAX),
before_sleep=before_log(log, logging.WARNING),
+ reraise=True,
)
def request(self, *args, **kwargs):
"""Implement a convenience for httpx.Client.request with a retry layer."""
diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py
index e5322519fbf..fd52433488a 100644
--- a/task-sdk/tests/task_sdk/api/test_client.py
+++ b/task-sdk/tests/task_sdk/api/test_client.py
@@ -26,6 +26,7 @@ from unittest import mock
import certifi
import httpx
import pytest
+import time_machine
import uuid6
from task_sdk import make_client, make_client_w_dry_run, make_client_w_responses
from uuid6 import uuid7
@@ -172,76 +173,57 @@ class TestClient:
assert unpickled.response.status_code == 404
assert unpickled.request.url == "http://error"
- @mock.patch("time.sleep", return_value=None)
- def test_retry_handling_unrecoverable_error(self, mock_sleep):
- responses: list[httpx.Response] = [
- *[httpx.Response(500, text="Internal Server Error")] * 6,
- httpx.Response(200, json={"detail": "Recovered from error - but will fail before"}),
- httpx.Response(400, json={"detail": "Should not get here"}),
- ]
- client = make_client_w_responses(responses)
-
- with pytest.raises(httpx.HTTPStatusError) as err:
- client.get("http://error")
- assert not isinstance(err.value, ServerResponseError)
- assert len(responses) == 3
- assert mock_sleep.call_count == 4
-
- @mock.patch("time.sleep", return_value=None)
- def test_retry_handling_recovered(self, mock_sleep):
- responses: list[httpx.Response] = [
- *[httpx.Response(500, text="Internal Server Error")] * 2,
- httpx.Response(200, json={"detail": "Recovered from error"}),
- httpx.Response(400, json={"detail": "Should not get here"}),
- ]
- client = make_client_w_responses(responses)
-
- response = client.get("http://error")
- assert response.status_code == 200
- assert len(responses) == 1
- assert mock_sleep.call_count == 2
-
- @mock.patch("time.sleep", return_value=None)
- def test_retry_handling_overload(self, mock_sleep):
- responses: list[httpx.Response] = [
- httpx.Response(429, text="I am really busy atm, please back-off", headers={"Retry-After": "37"}),
- httpx.Response(200, json={"detail": "Recovered from error"}),
- httpx.Response(400, json={"detail": "Should not get here"}),
- ]
- client = make_client_w_responses(responses)
-
- response = client.get("http://error")
- assert response.status_code == 200
- assert len(responses) == 1
- assert mock_sleep.call_count == 1
- assert mock_sleep.call_args[0][0] == 37
-
- @mock.patch("time.sleep", return_value=None)
- def test_retry_handling_non_retry_error(self, mock_sleep):
- responses: list[httpx.Response] = [
- httpx.Response(422, json={"detail": "Somehow this is a bad request"}),
- httpx.Response(400, json={"detail": "Should not get here"}),
- ]
- client = make_client_w_responses(responses)
-
- with pytest.raises(ServerResponseError) as err:
- client.get("http://error")
- assert len(responses) == 1
- assert mock_sleep.call_count == 0
- assert err.value.args == ("Somehow this is a bad request",)
-
- @mock.patch("time.sleep", return_value=None)
- def test_retry_handling_ok(self, mock_sleep):
- responses: list[httpx.Response] = [
- httpx.Response(200, json={"detail": "Recovered from error"}),
- httpx.Response(400, json={"detail": "Should not get here"}),
- ]
- client = make_client_w_responses(responses)
-
- response = client.get("http://error")
- assert response.status_code == 200
- assert len(responses) == 1
- assert mock_sleep.call_count == 0
+ def test_retry_handling_unrecoverable_error(self):
+ with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
+ responses: list[httpx.Response] = [
+ *[httpx.Response(500, text="Internal Server Error")] * 6,
+ httpx.Response(200, json={"detail": "Recovered from error - but will fail before"}),
+ httpx.Response(400, json={"detail": "Should not get here"}),
+ ]
+ client = make_client_w_responses(responses)
+
+ with pytest.raises(httpx.HTTPStatusError) as err:
+ client.get("http://error")
+ assert not isinstance(err.value, ServerResponseError)
+ assert len(responses) == 3
+
+ def test_retry_handling_recovered(self):
+ with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
+ responses: list[httpx.Response] = [
+ *[httpx.Response(500, text="Internal Server Error")] * 2,
+ httpx.Response(200, json={"detail": "Recovered from error"}),
+ httpx.Response(400, json={"detail": "Should not get here"}),
+ ]
+ client = make_client_w_responses(responses)
+
+ response = client.get("http://error")
+ assert response.status_code == 200
+ assert len(responses) == 1
+
+ def test_retry_handling_non_retry_error(self):
+ with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
+ responses: list[httpx.Response] = [
+ httpx.Response(422, json={"detail": "Somehow this is a bad request"}),
+ httpx.Response(400, json={"detail": "Should not get here"}),
+ ]
+ client = make_client_w_responses(responses)
+
+ with pytest.raises(ServerResponseError) as err:
+ client.get("http://error")
+ assert len(responses) == 1
+ assert err.value.args == ("Somehow this is a bad request",)
+
+ def test_retry_handling_ok(self):
+ with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
+ responses: list[httpx.Response] = [
+ httpx.Response(200, json={"detail": "Recovered from error"}),
+ httpx.Response(400, json={"detail": "Should not get here"}),
+ ]
+ client = make_client_w_responses(responses)
+
+ response = client.get("http://error")
+ assert response.status_code == 200
+ assert len(responses) == 1
def test_token_renewal(self):
responses: list[httpx.Response] = [
@@ -287,40 +269,40 @@ class TestTaskInstanceOperations:
response parsing.
"""
- @mock.patch("time.sleep", return_value=None) # To have retries not slowing down tests
- def test_task_instance_start(self, mock_sleep, make_ti_context):
- # Simulate a successful response from the server that starts a task
- ti_id = uuid6.uuid7()
- start_date = "2024-10-31T12:00:00Z"
- ti_context = make_ti_context(
- start_date=start_date,
- logical_date="2024-10-31T12:00:00Z",
- run_type="manual",
- )
-
- # ...including a validation that retry really works
- call_count = 0
+ def test_task_instance_start(self, make_ti_context):
+ with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
+ # Simulate a successful response from the server that starts a task
+ ti_id = uuid6.uuid7()
+ start_date = "2024-10-31T12:00:00Z"
+ ti_context = make_ti_context(
+ start_date=start_date,
+ logical_date="2024-10-31T12:00:00Z",
+ run_type="manual",
+ )
- def handle_request(request: httpx.Request) -> httpx.Response:
- nonlocal call_count
- call_count += 1
- if call_count < 3:
- return httpx.Response(status_code=500, json={"detail": "Internal Server Error"})
- if request.url.path == f"/task-instances/{ti_id}/run":
- actual_body = json.loads(request.read())
- assert actual_body["pid"] == 100
- assert actual_body["start_date"] == start_date
- assert actual_body["state"] == "running"
- return httpx.Response(
- status_code=200,
- json=ti_context.model_dump(mode="json"),
- )
- return httpx.Response(status_code=400, json={"detail": "Bad Request"})
+ # ...including a validation that retry really works
+ call_count = 0
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ nonlocal call_count
+ call_count += 1
+ if call_count < 3:
+ return httpx.Response(status_code=500, json={"detail": "Internal Server Error"})
+ if request.url.path == f"/task-instances/{ti_id}/run":
+ actual_body = json.loads(request.read())
+ assert actual_body["pid"] == 100
+ assert actual_body["start_date"] == start_date
+ assert actual_body["state"] == "running"
+ return httpx.Response(
+ status_code=200,
+ json=ti_context.model_dump(mode="json"),
+ )
+ return httpx.Response(status_code=400, json={"detail": "Bad Request"})
- client = make_client(transport=httpx.MockTransport(handle_request))
- resp = client.task_instances.start(ti_id, 100, start_date)
- assert resp == ti_context
- assert call_count == 3
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ resp = client.task_instances.start(ti_id, 100, start_date)
+ assert resp == ti_context
+ assert call_count == 3
@pytest.mark.parametrize(
"state", [state for state in TerminalTIState if state != TerminalTIState.SUCCESS]
@@ -563,31 +545,31 @@ class TestVariableOperations:
response parsing.
"""
- @mock.patch("time.sleep", return_value=None) # To have retries not slowing down tests
- def test_variable_get_success(self, mock_sleep):
- # Simulate a successful response from the server with a variable
- # ...including a validation that retry really works
- call_count = 0
+ def test_variable_get_success(self):
+ with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
+ # Simulate a successful response from the server with a variable
+ # ...including a validation that retry really works
+ call_count = 0
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ nonlocal call_count
+ call_count += 1
+ if call_count < 2:
+ return httpx.Response(status_code=500, json={"detail": "Internal Server Error"})
+ if request.url.path == "/variables/test_key":
+ return httpx.Response(
+ status_code=200,
+ json={"key": "test_key", "value": "test_value"},
+ )
+ return httpx.Response(status_code=400, json={"detail": "Bad Request"})
- def handle_request(request: httpx.Request) -> httpx.Response:
- nonlocal call_count
- call_count += 1
- if call_count < 2:
- return httpx.Response(status_code=500, json={"detail": "Internal Server Error"})
- if request.url.path == "/variables/test_key":
- return httpx.Response(
- status_code=200,
- json={"key": "test_key", "value": "test_value"},
- )
- return httpx.Response(status_code=400, json={"detail": "Bad Request"})
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ result = client.variables.get(key="test_key")
- client = make_client(transport=httpx.MockTransport(handle_request))
- result = client.variables.get(key="test_key")
-
- assert isinstance(result, VariableResponse)
- assert result.key == "test_key"
- assert result.value == "test_value"
- assert call_count == 2
+ assert isinstance(result, VariableResponse)
+ assert result.key == "test_key"
+ assert result.value == "test_value"
+ assert call_count == 2
def test_variable_not_found(self):
# Simulate a 404 response from the server
@@ -612,26 +594,26 @@ class TestVariableOperations:
assert resp.error == ErrorType.VARIABLE_NOT_FOUND
assert resp.detail == {"key": "non_existent_var"}
- @mock.patch("time.sleep", return_value=None)
- def test_variable_get_500_error(self, mock_sleep):
- # Simulate a response from the server returning a 500 error
- def handle_request(request: httpx.Request) -> httpx.Response:
- if request.url.path == "/variables/test_key":
- return httpx.Response(
- status_code=500,
- headers=[("content-Type", "application/json")],
- json={
- "reason": "internal_server_error",
- "message": "Internal Server Error",
- },
- )
- return httpx.Response(status_code=400, json={"detail": "Bad Request"})
+ def test_variable_get_500_error(self):
+ with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
+ # Simulate a response from the server returning a 500 error
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == "/variables/test_key":
+ return httpx.Response(
+ status_code=500,
+ headers=[("content-Type", "application/json")],
+ json={
+ "reason": "internal_server_error",
+ "message": "Internal Server Error",
+ },
+ )
+ return httpx.Response(status_code=400, json={"detail": "Bad Request"})
- client = make_client(transport=httpx.MockTransport(handle_request))
- with pytest.raises(ServerResponseError):
- client.variables.get(
- key="test_key",
- )
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ with pytest.raises(ServerResponseError):
+ client.variables.get(
+ key="test_key",
+ )
def test_variable_set_success(self):
# Simulate a successful response from the server when putting a variable
@@ -681,35 +663,35 @@ class TestXCOMOperations:
pytest.param({"key": "test_key", "value": {"key2": "value2"}}, id="nested-dict-value"),
],
)
- @mock.patch("time.sleep", return_value=None) # To have retries not slowing down tests
- def test_xcom_get_success(self, mock_sleep, value):
- # Simulate a successful response from the server when getting an xcom
- # ...including a validation that retry really works
- call_count = 0
-
- def handle_request(request: httpx.Request) -> httpx.Response:
- nonlocal call_count
- call_count += 1
- if call_count < 3:
- return httpx.Response(status_code=500, json={"detail": "Internal Server Error"})
- if request.url.path == "/xcoms/dag_id/run_id/task_id/key":
- return httpx.Response(
- status_code=201,
- json={"key": "test_key", "value": value},
- )
- return httpx.Response(status_code=400, json={"detail": "Bad Request"})
+ def test_xcom_get_success(self, value):
+ with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
+ # Simulate a successful response from the server when getting an xcom
+ # ...including a validation that retry really works
+ call_count = 0
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ nonlocal call_count
+ call_count += 1
+ if call_count < 3:
+ return httpx.Response(status_code=500, json={"detail": "Internal Server Error"})
+ if request.url.path == "/xcoms/dag_id/run_id/task_id/key":
+ return httpx.Response(
+ status_code=201,
+ json={"key": "test_key", "value": value},
+ )
+ return httpx.Response(status_code=400, json={"detail": "Bad Request"})
- client = make_client(transport=httpx.MockTransport(handle_request))
- result = client.xcoms.get(
- dag_id="dag_id",
- run_id="run_id",
- task_id="task_id",
- key="key",
- )
- assert isinstance(result, XComResponse)
- assert result.key == "test_key"
- assert result.value == value
- assert call_count == 3
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ result = client.xcoms.get(
+ dag_id="dag_id",
+ run_id="run_id",
+ task_id="task_id",
+ key="key",
+ )
+ assert isinstance(result, XComResponse)
+ assert result.key == "test_key"
+ assert result.value == value
+ assert call_count == 3
def test_xcom_get_success_with_map_index(self):
# Simulate a successful response from the server when getting an xcom with map_index passed
@@ -760,29 +742,29 @@ class TestXCOMOperations:
assert result.key == "test_key"
assert result.value == "test_value"
- @mock.patch("time.sleep", return_value=None)
- def test_xcom_get_500_error(self, mock_sleep):
- # Simulate a successful response from the server returning a 500 error
- def handle_request(request: httpx.Request) -> httpx.Response:
- if request.url.path == "/xcoms/dag_id/run_id/task_id/key":
- return httpx.Response(
- status_code=500,
- headers=[("content-Type", "application/json")],
- json={
- "reason": "invalid_format",
- "message": "XCom value is not a valid JSON",
- },
- )
- return httpx.Response(status_code=400, json={"detail": "Bad Request"})
+ def test_xcom_get_500_error(self):
+ with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
+ # Simulate a successful response from the server returning a 500 error
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == "/xcoms/dag_id/run_id/task_id/key":
+ return httpx.Response(
+ status_code=500,
+ headers=[("content-Type", "application/json")],
+ json={
+ "reason": "invalid_format",
+ "message": "XCom value is not a valid JSON",
+ },
+ )
+ return httpx.Response(status_code=400, json={"detail": "Bad Request"})
- client = make_client(transport=httpx.MockTransport(handle_request))
- with pytest.raises(ServerResponseError):
- client.xcoms.get(
- dag_id="dag_id",
- run_id="run_id",
- task_id="task_id",
- key="key",
- )
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ with pytest.raises(ServerResponseError):
+ client.xcoms.get(
+ dag_id="dag_id",
+ run_id="run_id",
+ task_id="task_id",
+ key="key",
+ )
@pytest.mark.parametrize(
"values",