-
Notifications
You must be signed in to change notification settings - Fork 14.4k
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.
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
Add HTTP retry handling into task SDK api.client #45121
base: main
Are you sure you want to change the base?
Add HTTP retry handling into task SDK api.client #45121
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.
Comments suppressed due to low confidence (1)
task_sdk/tests/api/test_client.py:101
- There is an extra single quote in the mounts dictionary key. It should be mounts={'http://': httpx.MockTransport(mock_handle_request)}.
mounts={'http://': httpx.MockTransport(mock_handle_request)},
bc58ae6
to
5399ace
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Overall looks good, some optimisations.
API_RETRIES = int(os.getenv("AIRFLOW__WORKERS__API_RETRIES", 10)) | ||
API_RETRY_WAIT_MIN = int(os.getenv("AIRFLOW__WORKERS__API_RETRY_WAIT_MIN", 1)) | ||
API_RETRY_WAIT_MAX = int(os.getenv("AIRFLOW__WORKERS__API_RETRY_WAIT_MAX", 90)) | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Two things here:
- Currently, the API_RETRIES, API_RETRY_WAIT_MIN, and API_RETRY_WAIT_MAX are directly cast to integers. This can raise a ValueError if the environment variables are not set correctly. Can we add validation or some fallback defaults (In case of error, if we don't want to fail)?
- Let's also have a check,
API_RETRY_WAIT_MIN < API_RETRY_WAIT_MAX
and raise alarm if not ?
token="", | ||
mounts={"'http://": httpx.MockTransport(mock_handle_request)}, | ||
) | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The section of code has been duplicated multiple times. Should we consider getting them out?
@pytest.fixture
def mock_client():
"""Helper fixture to create a mock client with custom responses."""
def _create_client(responses):
def mock_handle_request(request: httpx.Request) -> httpx.Response:
return responses.pop(0)
return Client(
base_url=None,
dry_run=True,
token="",
mounts={"'http://": httpx.MockTransport(mock_handle_request)},
)
return _create_client
and use client = mock_client(responses)
closes: #44355
Add HTTP retry handling to Task SDK.