diff --git a/changelog/+transient-retry.added.md b/changelog/+transient-retry.added.md
new file mode 100644
index 000000000..17627ed81
--- /dev/null
+++ b/changelog/+transient-retry.added.md
@@ -0,0 +1 @@
+With `retry_on_failure` enabled, the client now retries every transient failure, not only connection errors: dropped or timed-out connections, HTTP `500`/`502`/`503`/`504` responses, and GraphQL errors the server flags with one of those statuses, on every request path. Retries use exponential backoff with jitter (`retry_delay` up to the new `retry_max_delay`), `retry_status_codes` tunes what counts as transient, and `max_retry_duration=0` retries indefinitely.
diff --git a/docs/docs/python-sdk/guides/client.mdx b/docs/docs/python-sdk/guides/client.mdx
index 90872a0f1..bcbb56ce9 100644
--- a/docs/docs/python-sdk/guides/client.mdx
+++ b/docs/docs/python-sdk/guides/client.mdx
@@ -432,6 +432,52 @@ export INFRAHUB_PROXY_MOUNTS_HTTPS=http://https-proxy.example.com:8080
The `proxy` and `proxy_mounts` configurations are mutually exclusive and cannot be used together. Specifying both will cause a `ValueError` to be raised when the client is initialized.
:::
+### Retry transient failures
+
+Long-running scripts and generators can fail halfway through when Infrahub is temporarily unavailable: a database restart, an overloaded backend or a short network outage. With `retry_on_failure` enabled, the client retries a request when it fails with a transient error and raises immediately for every other error. A failure is transient when it is one of:
+
+- A connection error, or a connection dropped before the response arrived, such as when a load balancer restarts.
+- A timeout, while connecting, sending the request or reading the response.
+- An HTTP response with a status listed in `retry_status_codes` (`500`, `502`, `503` and `504` by default).
+- A GraphQL response in which every error is flagged by the server with one of those statuses.
+
+Retries use exponential backoff with jitter, starting at `retry_delay` seconds and capped at `retry_max_delay`. They stop after `max_retry_duration` seconds (5 minutes by default). Set `max_retry_duration` to `0` to keep retrying until the request succeeds, which lets a generator survive an outage instead of aborting.
+
+
+
+
+ ```python
+ from infrahub_sdk import Config, InfrahubClient
+ config = Config(retry_on_failure=True, max_retry_duration=0)
+ client = InfrahubClient(config=config)
+ ```
+
+
+
+
+ ```bash
+ export INFRAHUB_RETRY_ON_FAILURE=true
+ export INFRAHUB_MAX_RETRY_DURATION=0
+ ```
+
+ ```python
+ from infrahub_sdk import InfrahubClient
+ client = InfrahubClient() # retry settings are read from the environment variables
+ ```
+
+
+
+
+Every retry is logged at `WARNING` level with the attempt number and the time spent so far. After 5 minutes of retrying the log level switches to `ERROR`, so an operation that keeps retrying stays visible.
+
+:::warning Retrying mutations
+A mutation that timed out may have been applied by the server before the retry. Saving with `allow_upsert=True` is safe to retry. A plain create that had already succeeded fails on retry with a non-transient error, which is raised.
+:::
+
+:::note Unclassified server errors
+Infrahub reports some transient database failures as HTTP `500` without further classification, so `500` is part of the default `retry_status_codes`. The same status also covers genuine bugs, which are therefore retried until `max_retry_duration` expires. Remove `500` from `retry_status_codes` to fail fast on them instead.
+:::
+
## Next steps
Now that you have a fully configured Infrahub client, you're ready to start working with your infrastructure data. Here's what you can explore next:
diff --git a/docs/docs/python-sdk/reference/config.mdx b/docs/docs/python-sdk/reference/config.mdx
index b8750fe0a..5ffe9c056 100644
--- a/docs/docs/python-sdk/reference/config.mdx
+++ b/docs/docs/python-sdk/reference/config.mdx
@@ -144,21 +144,39 @@ The following settings can be defined in the `Config` class
## retry_delay
-**Description**: Number of seconds to wait until attempting a retry.
+**Description**: Base delay in seconds before retrying a request that failed with a transient error. The delay doubles after every attempt, with jitter, up to the maximum retry delay.
**Type**: `integer`
**Default value**: 5
**Environment variable**: `INFRAHUB_RETRY_DELAY`
+## retry_max_delay
+
+
+**Description**: Maximum delay in seconds between two retries of a request that failed with a transient error.
+**Type**: `integer`
+**Default value**: 60
+**Environment variable**: `INFRAHUB_RETRY_MAX_DELAY`
+
+
## retry_on_failure
-**Description**: Retry operation in case of failure
+**Description**: Retry requests that fail with a transient error: connection error, timeout, an HTTP status listed in the retry status codes, or a GraphQL error the server flags with one of those statuses. Other errors are never retried. The maximum retry duration controls how long to keep retrying.
**Type**: `boolean`
**Default value**: False
**Environment variable**: `INFRAHUB_RETRY_ON_FAILURE`
+## retry_status_codes
+
+
+**Description**: HTTP status codes treated as transient when retrying on failure is enabled. Also matched against the HTTP status reported in GraphQL error extensions. 500 is included because Infrahub reports some transient database errors without further classification; remove it to fail fast on them.
+**Type**: `array`
+**Default value**: [500, 502, 503, 504]
+**Environment variable**: `INFRAHUB_RETRY_STATUS_CODES`
+
+
## rate_limit_retry_enabled
@@ -198,7 +216,7 @@ The following settings can be defined in the `Config` class
## max_retry_duration
-**Description**: Maximum duration until we stop attempting to retry if enabled.
+**Description**: Maximum number of seconds to keep retrying a request that fails with transient errors when retrying on failure is enabled. Set to 0 to retry indefinitely.
**Type**: `integer`
**Default value**: 300
**Environment variable**: `INFRAHUB_MAX_RETRY_DURATION`
diff --git a/docs/docs/python-sdk/sdk_ref/infrahub_sdk/client.mdx b/docs/docs/python-sdk/sdk_ref/infrahub_sdk/client.mdx
index 1f783d46c..b7fcffacc 100644
--- a/docs/docs/python-sdk/sdk_ref/infrahub_sdk/client.mdx
+++ b/docs/docs/python-sdk/sdk_ref/infrahub_sdk/client.mdx
@@ -344,7 +344,9 @@ execute_graphql(self, query: str, variables: dict | None = None, branch_name: st
Execute a GraphQL query (or mutation).
-If retry_on_failure is True, the query will retry until the server becomes reachable.
+If retry_on_failure is True, transient failures (connection errors, timeouts, transient HTTP statuses and
+GraphQL errors the server flags as transient) are retried until max_retry_duration is exhausted, or
+indefinitely when max_retry_duration is 0.
**Args:**
@@ -364,11 +366,12 @@ client-wide default for this request only. When None, the client default (if any
**Raises:**
-- `GraphQLError`: When the GraphQL response contains errors.
-- `ServerNotReachableError`: If the server is not reachable after exhausting retries.
+- `GraphQLError`: When the GraphQL response contains errors that are not transient, or transient ones
+once the retry budget is exhausted.
+- `ServerNotReachableError`: If the server is not reachable, after exhausting retries when enabled.
+- `ServerNotResponsiveError`: If the server does not answer before the timeout, after exhausting retries.
- `AuthenticationError`: If the server returns a 401 or 403 response.
- `URLNotFoundError`: If the server returns a 404 response.
-- `Error`: If the response is unexpectedly missing.
#### `refresh_login`
@@ -661,7 +664,9 @@ execute_graphql(self, query: str, variables: dict | None = None, branch_name: st
Execute a GraphQL query (or mutation).
-If retry_on_failure is True, the query will retry until the server becomes reachable.
+If retry_on_failure is True, transient failures (connection errors, timeouts, transient HTTP statuses and
+GraphQL errors the server flags as transient) are retried until max_retry_duration is exhausted, or
+indefinitely when max_retry_duration is 0.
**Args:**
@@ -681,11 +686,12 @@ client-wide default for this request only. When None, the client default (if any
**Raises:**
-- `GraphQLError`: When the GraphQL response contains errors.
-- `ServerNotReachableError`: If the server is not reachable after exhausting retries.
+- `GraphQLError`: When the GraphQL response contains errors that are not transient, or transient ones
+once the retry budget is exhausted.
+- `ServerNotReachableError`: If the server is not reachable, after exhausting retries when enabled.
+- `ServerNotResponsiveError`: If the server does not answer before the timeout, after exhausting retries.
- `AuthenticationError`: If the server returns a 401 or 403 response.
- `URLNotFoundError`: If the server returns a 404 response.
-- `Error`: If the response is unexpectedly missing.
#### `count`
@@ -1076,6 +1082,34 @@ Base class for InfrahubClient and InfrahubClientSync.
**Methods:**
+#### `retry_on_failure`
+
+```python
+retry_on_failure(self) -> bool
+```
+
+Whether transient failures are retried. Can be toggled at runtime, e.g. by a long-running generator.
+
+#### `retry_on_failure`
+
+```python
+retry_on_failure(self, value: bool) -> None
+```
+
+#### `retry_delay`
+
+```python
+retry_delay(self) -> float
+```
+
+Base delay in seconds between retries of a transient failure; doubles per attempt up to retry_max_delay.
+
+#### `retry_delay`
+
+```python
+retry_delay(self, value: float) -> None
+```
+
#### `request_context`
```python
diff --git a/infrahub_sdk/client.py b/infrahub_sdk/client.py
index d600ea2fc..51d3eca66 100644
--- a/infrahub_sdk/client.py
+++ b/infrahub_sdk/client.py
@@ -3,14 +3,14 @@
import asyncio
import copy
import logging
-import time
+import shutil
+import tempfile
from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping, MutableMapping
from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager, suppress
from datetime import datetime
from enum import Enum
from functools import wraps
-from time import sleep
-from typing import TYPE_CHECKING, Any, BinaryIO, Literal, TypedDict, TypeVar, overload
+from typing import TYPE_CHECKING, Any, BinaryIO, Literal, TypedDict, TypeVar, cast, overload
from urllib.parse import quote, urlencode
import httpx
@@ -50,6 +50,7 @@
from .queries import QUERY_USER, get_commit_update_mutation
from .query_groups import InfrahubGroupContext, InfrahubGroupContextSync
from .rate_limit import RateLimitRetryHandler
+from .retry import CONNECTION_LOST_EXCEPTIONS, RetryState, TransientRetryHandler
from .schema import InfrahubSchema, InfrahubSchemaSync, NodeSchemaAPI
from .store import NodeStore, NodeStoreSync
from .task.manager import InfrahubTaskManager, InfrahubTaskManagerSync
@@ -80,6 +81,78 @@ class ProxyConfig(TypedDict):
mounts: Mapping[str, AsyncBaseTransport | None] | None
+def _is_rewindable(file_content: BinaryIO) -> bool:
+ """Whether ``file_content`` can be rewound, so a retried upload can re-send it from the start."""
+ seekable = getattr(file_content, "seekable", None)
+ return callable(seekable) and bool(seekable())
+
+
+@contextmanager
+def _seekable_upload(file_content: BinaryIO | None) -> Iterator[BinaryIO | None]:
+ """Yield ``file_content`` when it can be rewound, otherwise a seekable copy that is closed afterwards.
+
+ A retried upload re-sends the file from the start, which a pipe, a socket or another non-seekable
+ stream cannot do: ``_rewind_multipart_files`` would leave it consumed and the retry would carry an
+ empty body. Such a stream is copied once, to a temporary file so the size does not matter, before
+ the first attempt.
+ """
+ if file_content is None or _is_rewindable(file_content):
+ yield file_content
+ return
+ with tempfile.TemporaryFile() as buffer:
+ shutil.copyfileobj(file_content, buffer)
+ buffer.seek(0)
+ # TemporaryFile is typed as IO[bytes]; it offers every call the multipart sender makes.
+ yield cast("BinaryIO", buffer)
+
+
+@asynccontextmanager
+async def _aseekable_upload(file_content: BinaryIO | None) -> AsyncIterator[BinaryIO | None]:
+ """Async counterpart of :func:`_seekable_upload`.
+
+ The copy runs in a worker thread, so draining a slow or large stream does not block the event loop
+ for the other tasks sharing it. A thread cannot be interrupted: a cancellation arriving during the
+ copy is honoured at once, and the worker keeps the temporary file until its pending ``read`` returns,
+ then closes it, so the file is never closed under a worker still writing to it.
+
+ Yields:
+ BinaryIO | None: ``file_content`` itself when it can be rewound, otherwise its seekable copy.
+
+ Raises:
+ asyncio.CancelledError: If the caller is cancelled during the copy.
+
+ """
+ if file_content is None or _is_rewindable(file_content):
+ yield file_content
+ return
+ source: BinaryIO = file_content
+ buffer = tempfile.TemporaryFile() # noqa: SIM115 # the worker closes it if the caller is cancelled mid-copy
+
+ def copy() -> None:
+ shutil.copyfileobj(source, buffer)
+
+ def close_after_copy(task: asyncio.Future[None]) -> None:
+ if not task.cancelled():
+ task.exception() # nobody is waiting any more; mark a copy error retrieved so asyncio does not log it
+ buffer.close()
+
+ copy_task = asyncio.ensure_future(asyncio.to_thread(copy))
+ try:
+ await asyncio.shield(copy_task)
+ except BaseException:
+ if copy_task.done():
+ buffer.close()
+ else:
+ copy_task.add_done_callback(close_after_copy)
+ raise
+ try:
+ buffer.seek(0)
+ # TemporaryFile is typed as IO[bytes]; it offers every call the multipart sender makes.
+ yield cast("BinaryIO", buffer)
+ finally:
+ buffer.close()
+
+
def _rewind_multipart_files(files: dict[str, Any]) -> None:
"""Rewind seekable file objects in a multipart ``files`` payload to position 0.
@@ -211,8 +284,14 @@ def __init__(
self.address = self.config.address
self.mode = self.config.mode
self.pagination_size = self.config.pagination_size
- self.retry_delay = self.config.retry_delay
- self.retry_on_failure = self.config.retry_on_failure
+ self._retry_handler = TransientRetryHandler(
+ enabled=self.config.retry_on_failure,
+ base_delay=self.config.retry_delay,
+ max_delay=self.config.retry_max_delay,
+ max_duration=self.config.max_retry_duration,
+ status_codes=self.config.retry_status_codes,
+ log=self.log,
+ )
if self.config.api_token:
self.headers["X-INFRAHUB-KEY"] = self.config.api_token
@@ -229,6 +308,34 @@ def __init__(
self._request_context: RequestContext | None = None
_ = self.config.tls_context # Early load of the TLS context to catch errors
+ @property
+ def retry_on_failure(self) -> bool:
+ """Whether transient failures are retried. Can be toggled at runtime, e.g. by a long-running generator."""
+ return self._retry_handler.enabled
+
+ @retry_on_failure.setter
+ def retry_on_failure(self, value: bool) -> None:
+ self._retry_handler.enabled = value
+
+ @property
+ def retry_delay(self) -> float:
+ """Base delay in seconds between retries of a transient failure; doubles per attempt up to retry_max_delay."""
+ return self._retry_handler.base_delay
+
+ @retry_delay.setter
+ def retry_delay(self, value: float) -> None:
+ self._retry_handler.base_delay = value
+
+ def _is_retried_stream_status(self, status_code: int, retry_state: RetryState) -> bool:
+ """Whether a streamed response with this status will be retried, so its body is read and closed first.
+
+ A transient status that the handler will not retry, because retries are disabled or the budget is
+ spent, is handed to the caller open like any other response.
+ """
+ if status_code == 429:
+ return True
+ return self._retry_handler.is_transient_status(status_code) and self._retry_handler.should_retry(retry_state)
+
def _initialize(self) -> None:
"""Sets the properties for each version of the client."""
@@ -1286,7 +1393,9 @@ async def execute_graphql(
) -> dict:
"""Execute a GraphQL query (or mutation).
- If retry_on_failure is True, the query will retry until the server becomes reachable.
+ If retry_on_failure is True, transient failures (connection errors, timeouts, transient HTTP statuses and
+ GraphQL errors the server flags as transient) are retried until max_retry_duration is exhausted, or
+ indefinitely when max_retry_duration is 0.
Args:
query (_type_): GraphQL Query to execute, can be a query or a mutation
@@ -1303,11 +1412,12 @@ async def execute_graphql(
dict: The GraphQL data payload (response["data"]).
Raises:
- GraphQLError: When the GraphQL response contains errors.
- ServerNotReachableError: If the server is not reachable after exhausting retries.
+ GraphQLError: When the GraphQL response contains errors that are not transient, or transient ones
+ once the retry budget is exhausted.
+ ServerNotReachableError: If the server is not reachable, after exhausting retries when enabled.
+ ServerNotResponsiveError: If the server does not answer before the timeout, after exhausting retries.
AuthenticationError: If the server returns a 401 or 403 response.
URLNotFoundError: If the server returns a 404 response.
- Error: If the response is unexpectedly missing.
"""
branch_name = branch_name or self.default_branch
@@ -1323,25 +1433,11 @@ async def execute_graphql(
self._echo(url=url, query=query, variables=variables)
- retry = True
- resp = None
- start_time = time.time()
- while retry and time.time() - start_time < self.config.max_retry_duration:
- retry = self.retry_on_failure
+ retry_state = self._retry_handler.new_state()
+ while True:
+ resp = await self._post(url=url, payload=payload, headers=headers, timeout=timeout, retry_state=retry_state)
try:
- resp = await self._post(url=url, payload=payload, headers=headers, timeout=timeout)
resp.raise_for_status()
-
- retry = False
- except ServerNotReachableError:
- if retry:
- self.log.warning(
- f"Unable to connect to {self.address}, will retry in {self.retry_delay} seconds .."
- )
- await asyncio.sleep(delay=self.retry_delay)
- else:
- self.log.error(f"Unable to connect to {self.address} .. ")
- raise
except httpx.HTTPStatusError as exc:
if exc.response.status_code in {401, 403}:
response = decode_json(response=exc.response)
@@ -1350,16 +1446,22 @@ async def execute_graphql(
raise AuthenticationError(" | ".join(messages)) from exc
if exc.response.status_code == 404:
raise URLNotFoundError(url=url) from exc
+ # Any other status falls through: the body is expected to carry a GraphQL error envelope.
- if not resp:
- raise Error("Unexpected situation, resp hasn't been initialized.")
+ response = decode_json(response=resp)
- response = decode_json(response=resp)
-
- if "errors" in response:
- raise GraphQLError(errors=response["errors"], query=query, variables=variables)
+ if "errors" in response:
+ errors = response["errors"]
+ if self._retry_handler.is_transient_graphql_errors(errors) and self._retry_handler.should_retry(
+ retry_state
+ ):
+ await self._retry_handler.asleep_before_retry(
+ state=retry_state, url=url, reason=self._retry_handler.describe_graphql_errors(errors)
+ )
+ continue
+ raise GraphQLError(errors=errors, query=query, variables=variables)
- return response["data"]
+ return response["data"]
# TODO add a special method to execute mutation that will check if the method returned OK
@@ -1378,7 +1480,10 @@ async def _execute_graphql_with_file(
"""Execute a GraphQL mutation with a file upload using multipart/form-data.
This method follows the GraphQL Multipart Request Spec for file uploads.
- The file is attached to the 'file' variable in the mutation.
+ The file is attached to the 'file' variable in the mutation. Transient failures, including GraphQL
+ errors the server flags as transient, are retried like in ``execute_graphql`` when retry_on_failure is
+ enabled; the file is rewound before every attempt, and a non-seekable stream is copied once up front so
+ every attempt carries the full body.
Args:
query: GraphQL mutation query that includes a $file variable of type Upload!
@@ -1411,24 +1516,36 @@ async def _execute_graphql_with_file(
self._echo(url=url, query=query, variables=variables)
- resp = await self._post_multipart(
- url=url,
- query=query,
- variables=variables,
- file_content=file_content,
- file_name=file_name or "upload",
- headers=headers,
- timeout=timeout,
- operation_name=operation_name,
- )
+ retry_state = self._retry_handler.new_state()
+ async with _aseekable_upload(file_content) as upload:
+ while True:
+ resp = await self._post_multipart(
+ url=url,
+ query=query,
+ variables=variables,
+ file_content=upload,
+ file_name=file_name or "upload",
+ headers=headers,
+ timeout=timeout,
+ operation_name=operation_name,
+ retry_state=retry_state,
+ )
- resp.raise_for_status()
- response = decode_json(response=resp)
+ resp.raise_for_status()
+ response = decode_json(response=resp)
- if "errors" in response:
- raise GraphQLError(errors=response["errors"], query=query, variables=variables)
+ if "errors" in response:
+ errors = response["errors"]
+ if self._retry_handler.is_transient_graphql_errors(errors) and self._retry_handler.should_retry(
+ retry_state
+ ):
+ await self._retry_handler.asleep_before_retry(
+ state=retry_state, url=url, reason=self._retry_handler.describe_graphql_errors(errors)
+ )
+ continue
+ raise GraphQLError(errors=errors, query=query, variables=variables)
- return response["data"]
+ return response["data"]
@handle_relogin
async def _post_multipart(
@@ -1441,6 +1558,7 @@ async def _post_multipart(
headers: dict | None = None,
timeout: int | None = None,
operation_name: str | None = None,
+ retry_state: RetryState | None = None,
) -> httpx.Response:
"""Execute a HTTP POST with multipart/form-data for GraphQL file uploads.
@@ -1462,7 +1580,7 @@ async def _post_multipart(
)
return await self._request_multipart(
- url=url, headers=headers, timeout=timeout or self.default_timeout, files=files
+ url=url, headers=headers, timeout=timeout or self.default_timeout, files=files, retry_state=retry_state
)
def _build_proxy_config(self) -> ProxyConfig:
@@ -1478,7 +1596,12 @@ def _build_proxy_config(self) -> ProxyConfig:
return proxy_config
async def _request_multipart(
- self, url: str, headers: dict[str, Any], timeout: int, files: dict[str, Any]
+ self,
+ url: str,
+ headers: dict[str, Any],
+ timeout: int,
+ files: dict[str, Any],
+ retry_state: RetryState | None = None,
) -> httpx.Response:
"""Execute a multipart HTTP POST request.
@@ -1493,12 +1616,15 @@ async def send() -> httpx.Response:
async with httpx.AsyncClient(**self._build_proxy_config(), verify=self.config.tls_context) as client:
try:
return await client.post(url=url, headers=headers, timeout=timeout, files=files)
- except httpx.NetworkError as exc:
+ except CONNECTION_LOST_EXCEPTIONS as exc:
raise ServerNotReachableError(address=self.address) from exc
- except httpx.ReadTimeout as exc:
+ except httpx.TimeoutException as exc:
raise ServerNotResponsiveError(url=url, timeout=timeout) from exc
- response = await self._rate_limit_handler.asend(send=send, url=url)
+ async def send_with_rate_limit() -> httpx.Response:
+ return await self._rate_limit_handler.asend(send=send, url=url)
+
+ response = await self._retry_handler.asend(send=send_with_rate_limit, url=url, state=retry_state)
self._record(response)
return response
@@ -1509,9 +1635,13 @@ async def _post(
payload: dict,
headers: dict | None = None,
timeout: int | None = None,
+ retry_state: RetryState | None = None,
) -> httpx.Response:
"""Execute a HTTP POST with HTTPX.
+ ``retry_state`` lets a caller that retries on its own (``execute_graphql``) share one retry budget
+ with the transport-level retries performed by ``_request``.
+
Raises:
ServerNotReachableError: If we are not able to connect to the server.
ServerNotResponsiveError: If the server didn't respond before the timeout expired.
@@ -1527,6 +1657,7 @@ async def _post(
headers=headers,
timeout=timeout or self.default_timeout,
payload=payload,
+ retry_state=retry_state,
)
@handle_relogin
@@ -1551,19 +1682,28 @@ async def _get(self, url: str, headers: dict | None = None, timeout: int | None
@asynccontextmanager
async def _get_streaming(
- self, url: str, headers: dict | None = None, timeout: int | None = None
+ self,
+ url: str,
+ headers: dict | None = None,
+ timeout: int | None = None,
+ retry_state: RetryState | None = None,
) -> AsyncIterator[httpx.Response]:
"""Execute a streaming HTTP GET with HTTPX.
Returns an async context manager that yields the streaming response.
Use this for downloading large files without loading into memory.
+ Only the stream initiation is retried here: a transient failure while the caller reads the body is
+ raised as ``ServerNotResponsiveError``. Pass ``retry_state`` from an outer loop that restarts the whole
+ transfer, such as the file handler's download to disk, so both share one retry budget.
+
Yields:
httpx.Response: The streaming HTTP response.
Raises:
ServerNotReachableError: If we are not able to connect to the server.
- ServerNotResponsiveError: If the server didn't respond before the timeout expired.
+ ServerNotResponsiveError: If the server didn't respond before the timeout expired, or the
+ connection was lost while the caller was reading the body.
"""
await self.login()
@@ -1573,16 +1713,22 @@ async def _get_streaming(
request_timeout = timeout or self.default_timeout
async with httpx.AsyncClient(**self._build_proxy_config(), verify=self.config.tls_context) as client:
open_stream: dict[str, AsyncExitStack] = {}
+ retry_state = retry_state or self._retry_handler.new_state()
async def send() -> httpx.Response:
- # Retry stream initiation only (a 429 arrives in the headers before the body): a
- # failed attempt is read and closed here, a successful stream is left open for
- # the caller and closed afterwards.
+ # Retry stream initiation only (a 429 or a transient status arrives in the headers
+ # before the body): a failed attempt is read and closed here, a successful stream is
+ # left open for the caller and closed afterwards.
stack = AsyncExitStack()
- response = await stack.enter_async_context(
- client.stream(method="GET", url=url, headers=headers, timeout=request_timeout)
- )
- if response.status_code == 429:
+ try:
+ response = await stack.enter_async_context(
+ client.stream(method="GET", url=url, headers=headers, timeout=request_timeout)
+ )
+ except CONNECTION_LOST_EXCEPTIONS as exc:
+ raise ServerNotReachableError(address=self.address) from exc
+ except httpx.TimeoutException as exc:
+ raise ServerNotResponsiveError(url=url, timeout=request_timeout) from exc
+ if self._is_retried_stream_status(response.status_code, retry_state):
try:
await response.aread()
finally:
@@ -1591,17 +1737,23 @@ async def send() -> httpx.Response:
open_stream["stack"] = stack
return response
+ async def send_with_rate_limit() -> httpx.Response:
+ return await self._rate_limit_handler.asend(send=send, url=url)
+
try:
- response = await self._rate_limit_handler.asend(send=send, url=url)
+ response = await self._retry_handler.asend(send=send_with_rate_limit, url=url, state=retry_state)
try:
yield response
finally:
stack = open_stream.get("stack")
if stack is not None:
await stack.aclose()
- except httpx.NetworkError as exc:
- raise ServerNotReachableError(address=self.address) from exc
- except httpx.ReadTimeout as exc:
+ except CONNECTION_LOST_EXCEPTIONS as exc:
+ # Stream initiation maps its own failures inside send(); this is the body read breaking off.
+ raise ServerNotResponsiveError(
+ url=url, message=f"Connection to '{url}' was lost while reading the response."
+ ) from exc
+ except httpx.TimeoutException as exc:
raise ServerNotResponsiveError(url=url, timeout=request_timeout) from exc
async def _request(
@@ -1611,11 +1763,15 @@ async def _request(
headers: dict[str, Any],
timeout: int,
payload: dict | None = None,
+ retry_state: RetryState | None = None,
) -> httpx.Response:
async def send() -> httpx.Response:
return await self._request_method(url=url, method=method, headers=headers, timeout=timeout, payload=payload)
- response = await self._rate_limit_handler.asend(send=send, url=url)
+ async def send_with_rate_limit() -> httpx.Response:
+ return await self._rate_limit_handler.asend(send=send, url=url)
+
+ response = await self._retry_handler.asend(send=send_with_rate_limit, url=url, state=retry_state)
self._record(response)
return response
@@ -1640,9 +1796,9 @@ async def _default_request_method(
timeout=timeout,
**params,
)
- except httpx.NetworkError as exc:
+ except CONNECTION_LOST_EXCEPTIONS as exc:
raise ServerNotReachableError(address=self.address) from exc
- except httpx.ReadTimeout as exc:
+ except httpx.TimeoutException as exc:
raise ServerNotResponsiveError(url=url, timeout=timeout) from exc
return response
@@ -2274,7 +2430,9 @@ def execute_graphql(
) -> dict:
"""Execute a GraphQL query (or mutation).
- If retry_on_failure is True, the query will retry until the server becomes reachable.
+ If retry_on_failure is True, transient failures (connection errors, timeouts, transient HTTP statuses and
+ GraphQL errors the server flags as transient) are retried until max_retry_duration is exhausted, or
+ indefinitely when max_retry_duration is 0.
Args:
query (str): GraphQL Query to execute, can be a query or a mutation
@@ -2291,11 +2449,12 @@ def execute_graphql(
dict: The GraphQL data payload (`response["data"]`).
Raises:
- GraphQLError: When the GraphQL response contains errors.
- ServerNotReachableError: If the server is not reachable after exhausting retries.
+ GraphQLError: When the GraphQL response contains errors that are not transient, or transient ones
+ once the retry budget is exhausted.
+ ServerNotReachableError: If the server is not reachable, after exhausting retries when enabled.
+ ServerNotResponsiveError: If the server does not answer before the timeout, after exhausting retries.
AuthenticationError: If the server returns a 401 or 403 response.
URLNotFoundError: If the server returns a 404 response.
- Error: If the response is unexpectedly missing.
"""
branch_name = branch_name or self.default_branch
@@ -2311,25 +2470,11 @@ def execute_graphql(
self._echo(url=url, query=query, variables=variables)
- retry = True
- resp = None
- start_time = time.time()
- while retry and time.time() - start_time < self.config.max_retry_duration:
- retry = self.retry_on_failure
+ retry_state = self._retry_handler.new_state()
+ while True:
+ resp = self._post(url=url, payload=payload, headers=headers, timeout=timeout, retry_state=retry_state)
try:
- resp = self._post(url=url, payload=payload, headers=headers, timeout=timeout)
resp.raise_for_status()
-
- retry = False
- except ServerNotReachableError:
- if retry:
- self.log.warning(
- f"Unable to connect to {self.address}, will retry in {self.retry_delay} seconds .."
- )
- sleep(self.retry_delay)
- else:
- self.log.error(f"Unable to connect to {self.address} .. ")
- raise
except httpx.HTTPStatusError as exc:
if exc.response.status_code in {401, 403}:
response = decode_json(response=exc.response)
@@ -2338,16 +2483,22 @@ def execute_graphql(
raise AuthenticationError(" | ".join(messages)) from exc
if exc.response.status_code == 404:
raise URLNotFoundError(url=url) from exc
+ # Any other status falls through: the body is expected to carry a GraphQL error envelope.
- if not resp:
- raise Error("Unexpected situation, resp hasn't been initialized.")
-
- response = decode_json(response=resp)
+ response = decode_json(response=resp)
- if "errors" in response:
- raise GraphQLError(errors=response["errors"], query=query, variables=variables)
+ if "errors" in response:
+ errors = response["errors"]
+ if self._retry_handler.is_transient_graphql_errors(errors) and self._retry_handler.should_retry(
+ retry_state
+ ):
+ self._retry_handler.sleep_before_retry(
+ state=retry_state, url=url, reason=self._retry_handler.describe_graphql_errors(errors)
+ )
+ continue
+ raise GraphQLError(errors=errors, query=query, variables=variables)
- return response["data"]
+ return response["data"]
# TODO add a special method to execute mutation that will check if the method returned OK
@@ -2366,7 +2517,10 @@ def _execute_graphql_with_file(
"""Execute a GraphQL mutation with a file upload using multipart/form-data.
This method follows the GraphQL Multipart Request Spec for file uploads.
- The file is attached to the 'file' variable in the mutation.
+ The file is attached to the 'file' variable in the mutation. Transient failures, including GraphQL
+ errors the server flags as transient, are retried like in ``execute_graphql`` when retry_on_failure is
+ enabled; the file is rewound before every attempt, and a non-seekable stream is copied once up front so
+ every attempt carries the full body.
Args:
query: GraphQL mutation query that includes a $file variable of type Upload!
@@ -2399,24 +2553,36 @@ def _execute_graphql_with_file(
self._echo(url=url, query=query, variables=variables)
- resp = self._post_multipart(
- url=url,
- query=query,
- variables=variables,
- file_content=file_content,
- file_name=file_name or "upload",
- headers=headers,
- timeout=timeout,
- operation_name=operation_name,
- )
+ retry_state = self._retry_handler.new_state()
+ with _seekable_upload(file_content) as upload:
+ while True:
+ resp = self._post_multipart(
+ url=url,
+ query=query,
+ variables=variables,
+ file_content=upload,
+ file_name=file_name or "upload",
+ headers=headers,
+ timeout=timeout,
+ operation_name=operation_name,
+ retry_state=retry_state,
+ )
- resp.raise_for_status()
- response = decode_json(response=resp)
+ resp.raise_for_status()
+ response = decode_json(response=resp)
- if "errors" in response:
- raise GraphQLError(errors=response["errors"], query=query, variables=variables)
+ if "errors" in response:
+ errors = response["errors"]
+ if self._retry_handler.is_transient_graphql_errors(errors) and self._retry_handler.should_retry(
+ retry_state
+ ):
+ self._retry_handler.sleep_before_retry(
+ state=retry_state, url=url, reason=self._retry_handler.describe_graphql_errors(errors)
+ )
+ continue
+ raise GraphQLError(errors=errors, query=query, variables=variables)
- return response["data"]
+ return response["data"]
@handle_relogin_sync
def _post_multipart(
@@ -2429,6 +2595,7 @@ def _post_multipart(
headers: dict | None = None,
timeout: int | None = None,
operation_name: str | None = None,
+ retry_state: RetryState | None = None,
) -> httpx.Response:
"""Execute a HTTP POST with multipart/form-data for GraphQL file uploads.
@@ -2449,7 +2616,9 @@ def _post_multipart(
operation_name=operation_name,
)
- return self._request_multipart(url=url, headers=headers, timeout=timeout or self.default_timeout, files=files)
+ return self._request_multipart(
+ url=url, headers=headers, timeout=timeout or self.default_timeout, files=files, retry_state=retry_state
+ )
def _build_proxy_config(self) -> ProxyConfigSync:
"""Build proxy configuration for httpx Client."""
@@ -2464,7 +2633,12 @@ def _build_proxy_config(self) -> ProxyConfigSync:
return proxy_config
def _request_multipart(
- self, url: str, headers: dict[str, Any], timeout: int, files: dict[str, Any]
+ self,
+ url: str,
+ headers: dict[str, Any],
+ timeout: int,
+ files: dict[str, Any],
+ retry_state: RetryState | None = None,
) -> httpx.Response:
"""Execute a multipart HTTP POST request.
@@ -2479,12 +2653,15 @@ def send() -> httpx.Response:
with httpx.Client(**self._build_proxy_config(), verify=self.config.tls_context) as client:
try:
return client.post(url=url, headers=headers, timeout=timeout, files=files)
- except httpx.NetworkError as exc:
+ except CONNECTION_LOST_EXCEPTIONS as exc:
raise ServerNotReachableError(address=self.address) from exc
- except httpx.ReadTimeout as exc:
+ except httpx.TimeoutException as exc:
raise ServerNotResponsiveError(url=url, timeout=timeout) from exc
- response = self._rate_limit_handler.send(send=send, url=url)
+ def send_with_rate_limit() -> httpx.Response:
+ return self._rate_limit_handler.send(send=send, url=url)
+
+ response = self._retry_handler.send(send=send_with_rate_limit, url=url, state=retry_state)
self._record(response)
return response
@@ -3693,19 +3870,28 @@ def _get(self, url: str, headers: dict | None = None, timeout: int | None = None
@contextmanager
def _get_streaming(
- self, url: str, headers: dict | None = None, timeout: int | None = None
+ self,
+ url: str,
+ headers: dict | None = None,
+ timeout: int | None = None,
+ retry_state: RetryState | None = None,
) -> Iterator[httpx.Response]:
"""Execute a streaming HTTP GET with HTTPX.
Returns a context manager that yields the streaming response.
Use this for downloading large files without loading into memory.
+ Only the stream initiation is retried here: a transient failure while the caller reads the body is
+ raised as ``ServerNotResponsiveError``. Pass ``retry_state`` from an outer loop that restarts the whole
+ transfer, such as the file handler's download to disk, so both share one retry budget.
+
Yields:
httpx.Response: The streaming HTTP response.
Raises:
ServerNotReachableError: If we are not able to connect to the server.
- ServerNotResponsiveError: If the server didn't respond before the timeout expired.
+ ServerNotResponsiveError: If the server didn't respond before the timeout expired, or the
+ connection was lost while the caller was reading the body.
"""
self.login()
@@ -3715,16 +3901,22 @@ def _get_streaming(
request_timeout = timeout or self.default_timeout
with httpx.Client(**self._build_proxy_config(), verify=self.config.tls_context) as client:
open_stream: dict[str, ExitStack] = {}
+ retry_state = retry_state or self._retry_handler.new_state()
def send() -> httpx.Response:
- # Retry stream initiation only (a 429 arrives in the headers before the body): a
- # failed attempt is read and closed here, a successful stream is left open for
- # the caller and closed afterwards.
+ # Retry stream initiation only (a 429 or a transient status arrives in the headers
+ # before the body): a failed attempt is read and closed here, a successful stream is
+ # left open for the caller and closed afterwards.
stack = ExitStack()
- response = stack.enter_context(
- client.stream(method="GET", url=url, headers=headers, timeout=request_timeout)
- )
- if response.status_code == 429:
+ try:
+ response = stack.enter_context(
+ client.stream(method="GET", url=url, headers=headers, timeout=request_timeout)
+ )
+ except CONNECTION_LOST_EXCEPTIONS as exc:
+ raise ServerNotReachableError(address=self.address) from exc
+ except httpx.TimeoutException as exc:
+ raise ServerNotResponsiveError(url=url, timeout=request_timeout) from exc
+ if self._is_retried_stream_status(response.status_code, retry_state):
try:
response.read()
finally:
@@ -3733,17 +3925,23 @@ def send() -> httpx.Response:
open_stream["stack"] = stack
return response
+ def send_with_rate_limit() -> httpx.Response:
+ return self._rate_limit_handler.send(send=send, url=url)
+
try:
- response = self._rate_limit_handler.send(send=send, url=url)
+ response = self._retry_handler.send(send=send_with_rate_limit, url=url, state=retry_state)
try:
yield response
finally:
stack = open_stream.get("stack")
if stack is not None:
stack.close()
- except httpx.NetworkError as exc:
- raise ServerNotReachableError(address=self.address) from exc
- except httpx.ReadTimeout as exc:
+ except CONNECTION_LOST_EXCEPTIONS as exc:
+ # Stream initiation maps its own failures inside send(); this is the body read breaking off.
+ raise ServerNotResponsiveError(
+ url=url, message=f"Connection to '{url}' was lost while reading the response."
+ ) from exc
+ except httpx.TimeoutException as exc:
raise ServerNotResponsiveError(url=url, timeout=request_timeout) from exc
@handle_relogin_sync
@@ -3753,9 +3951,13 @@ def _post(
payload: dict,
headers: dict | None = None,
timeout: int | None = None,
+ retry_state: RetryState | None = None,
) -> httpx.Response:
"""Execute a HTTP POST with HTTPX.
+ ``retry_state`` lets a caller that retries on its own (``execute_graphql``) share one retry budget
+ with the transport-level retries performed by ``_request``.
+
Raises:
ServerNotReachableError: If we are not able to connect to the server.
ServerNotResponsiveError: If the server didn't respond before the timeout expired.
@@ -3771,6 +3973,7 @@ def _post(
payload=payload,
headers=headers,
timeout=timeout or self.default_timeout,
+ retry_state=retry_state,
)
def _request(
@@ -3780,11 +3983,15 @@ def _request(
headers: dict[str, Any],
timeout: int,
payload: dict | None = None,
+ retry_state: RetryState | None = None,
) -> httpx.Response:
def send() -> httpx.Response:
return self._request_method(url=url, method=method, headers=headers, timeout=timeout, payload=payload)
- response = self._rate_limit_handler.send(send=send, url=url)
+ def send_with_rate_limit() -> httpx.Response:
+ return self._rate_limit_handler.send(send=send, url=url)
+
+ response = self._retry_handler.send(send=send_with_rate_limit, url=url, state=retry_state)
self._record(response)
return response
@@ -3809,9 +4016,9 @@ def _default_request_method(
timeout=timeout,
**params,
)
- except httpx.NetworkError as exc:
+ except CONNECTION_LOST_EXCEPTIONS as exc:
raise ServerNotReachableError(address=self.address) from exc
- except httpx.ReadTimeout as exc:
+ except httpx.TimeoutException as exc:
raise ServerNotResponsiveError(url=url, timeout=timeout) from exc
return response
diff --git a/infrahub_sdk/config.py b/infrahub_sdk/config.py
index f81fc911e..b5c5e1aef 100644
--- a/infrahub_sdk/config.py
+++ b/infrahub_sdk/config.py
@@ -63,8 +63,35 @@ class ConfigBase(BaseSettings):
"one of high|medium|low (case-insensitive). When unset, no header is sent."
),
)
- retry_delay: int = Field(default=5, description="Number of seconds to wait until attempting a retry.")
- retry_on_failure: bool = Field(default=False, description="Retry operation in case of failure")
+ retry_delay: int = Field(
+ default=5,
+ gt=0,
+ description=(
+ "Base delay in seconds before retrying a request that failed with a transient error. "
+ "The delay doubles after every attempt, with jitter, up to the maximum retry delay."
+ ),
+ )
+ retry_max_delay: int = Field(
+ default=60,
+ gt=0,
+ description="Maximum delay in seconds between two retries of a request that failed with a transient error.",
+ )
+ retry_on_failure: bool = Field(
+ default=False,
+ description=(
+ "Retry requests that fail with a transient error: connection error, timeout, an HTTP status listed in "
+ "the retry status codes, or a GraphQL error the server flags with one of those statuses. "
+ "Other errors are never retried. The maximum retry duration controls how long to keep retrying."
+ ),
+ )
+ retry_status_codes: list[int] = Field(
+ default=[500, 502, 503, 504],
+ description=(
+ "HTTP status codes treated as transient when retrying on failure is enabled. Also matched against the "
+ "HTTP status reported in GraphQL error extensions. 500 is included because Infrahub reports some "
+ "transient database errors without further classification; remove it to fail fast on them."
+ ),
+ )
rate_limit_retry_enabled: bool = Field(
default=True,
description="Retry requests that receive HTTP 429 using backoff. Set False to disable.",
@@ -85,7 +112,12 @@ class ConfigBase(BaseSettings):
description="Maximum wait in seconds for any single 429 retry (also clamps Retry-After).",
)
max_retry_duration: int = Field(
- default=300, description="Maximum duration until we stop attempting to retry if enabled."
+ default=300,
+ ge=0,
+ description=(
+ "Maximum number of seconds to keep retrying a request that fails with transient errors when "
+ "retrying on failure is enabled. Set to 0 to retry indefinitely."
+ ),
)
schema_converge_timeout: int = Field(
default=60, description="Number of seconds to wait for schema to have converged"
diff --git a/infrahub_sdk/file_handler.py b/infrahub_sdk/file_handler.py
index 56a7f106c..3c0af3139 100644
--- a/infrahub_sdk/file_handler.py
+++ b/infrahub_sdk/file_handler.py
@@ -9,10 +9,11 @@
import anyio
import httpx
-from .exceptions import AuthenticationError, NodeNotFoundError, ServerNotReachableError
+from .exceptions import AuthenticationError, NodeNotFoundError, ServerNotReachableError, ServerNotResponsiveError
if TYPE_CHECKING:
from .client import InfrahubClient, InfrahubClientSync
+ from .retry import RetryState
_SHA1_CHUNK_BYTES = 64 * 1024
@@ -245,6 +246,7 @@ async def download(self, node_id: str, branch: str | None, dest: Path | None = N
Raises:
ServerNotReachableError: If the server is not reachable.
+ ServerNotResponsiveError: If the server timed out or the connection was lost mid-download.
AuthenticationError: If authentication fails.
NodeNotFoundError: If the node/file is not found.
@@ -266,6 +268,10 @@ async def download(self, node_id: str, branch: str | None, dest: Path | None = N
async def _stream_to_file(self, url: str, dest: Path) -> int:
"""Stream download directly to a file without loading into memory.
+ A transfer interrupted by a transient failure (the connection dropped or timed out mid-body) is
+ restarted from the beginning when the client retries on failure, on the same time budget as the
+ stream initiation; ``dest`` is removed after a failed attempt so it never holds a partial body.
+
Args:
url: The URL to download from.
dest: The destination path to write to.
@@ -275,28 +281,44 @@ async def _stream_to_file(self, url: str, dest: Path) -> int:
Raises:
ServerNotReachableError: If the server is not reachable.
+ ServerNotResponsiveError: If the server timed out or the connection was lost mid-download.
AuthenticationError: If authentication fails.
NodeNotFoundError: If the file is not found.
"""
- try:
- async with self._client._get_streaming(url=url) as resp:
- try:
- resp.raise_for_status()
- except httpx.HTTPStatusError as exc:
- # Need to read the response body for error details
- await resp.aread()
- self.handle_error_response(exc=exc)
-
- bytes_written = 0
+ retry_handler = self._client._retry_handler
+ retry_state = retry_handler.new_state()
+ while True:
+ try:
+ return await self._download_to_file(url=url, dest=dest, retry_state=retry_state)
+ except (ServerNotReachableError, ServerNotResponsiveError) as exc:
+ if retry_handler.should_retry(retry_state):
+ await retry_handler.asleep_before_retry(state=retry_state, url=url, reason=str(exc))
+ continue
+ if isinstance(exc, ServerNotReachableError):
+ self._client.log.error(f"Unable to connect to {self._client.address}")
+ raise
+
+ async def _download_to_file(self, url: str, dest: Path, retry_state: RetryState) -> int:
+ """One attempt at streaming ``url`` into ``dest``; the file is removed again if the transfer fails."""
+ async with self._client._get_streaming(url=url, retry_state=retry_state) as resp:
+ try:
+ resp.raise_for_status()
+ except httpx.HTTPStatusError as exc:
+ # Need to read the response body for error details
+ await resp.aread()
+ self.handle_error_response(exc=exc)
+
+ bytes_written = 0
+ try:
async with await anyio.Path(dest).open("wb") as f:
async for chunk in resp.aiter_bytes(chunk_size=65536):
await f.write(chunk)
bytes_written += len(chunk)
- return bytes_written
- except ServerNotReachableError:
- self._client.log.error(f"Unable to connect to {self._client.address}")
- raise
+ except BaseException:
+ await anyio.Path(dest).unlink(missing_ok=True)
+ raise
+ return bytes_written
class FileHandlerSync(FileHandlerBase):
@@ -354,6 +376,7 @@ def download(self, node_id: str, branch: str | None, dest: Path | None = None) -
Raises:
ServerNotReachableError: If the server is not reachable.
+ ServerNotResponsiveError: If the server timed out or the connection was lost mid-download.
AuthenticationError: If authentication fails.
NodeNotFoundError: If the node/file is not found.
@@ -375,6 +398,10 @@ def download(self, node_id: str, branch: str | None, dest: Path | None = None) -
def _stream_to_file(self, url: str, dest: Path) -> int:
"""Stream download directly to a file without loading into memory.
+ A transfer interrupted by a transient failure (the connection dropped or timed out mid-body) is
+ restarted from the beginning when the client retries on failure, on the same time budget as the
+ stream initiation; ``dest`` is removed after a failed attempt so it never holds a partial body.
+
Args:
url: The URL to download from.
dest: The destination path to write to.
@@ -384,25 +411,41 @@ def _stream_to_file(self, url: str, dest: Path) -> int:
Raises:
ServerNotReachableError: If the server is not reachable.
+ ServerNotResponsiveError: If the server timed out or the connection was lost mid-download.
AuthenticationError: If authentication fails.
NodeNotFoundError: If the file is not found.
"""
- try:
- with self._client._get_streaming(url=url) as resp:
- try:
- resp.raise_for_status()
- except httpx.HTTPStatusError as exc:
- # Need to read the response body for error details
- resp.read()
- self.handle_error_response(exc=exc)
-
- bytes_written = 0
+ retry_handler = self._client._retry_handler
+ retry_state = retry_handler.new_state()
+ while True:
+ try:
+ return self._download_to_file(url=url, dest=dest, retry_state=retry_state)
+ except (ServerNotReachableError, ServerNotResponsiveError) as exc:
+ if retry_handler.should_retry(retry_state):
+ retry_handler.sleep_before_retry(state=retry_state, url=url, reason=str(exc))
+ continue
+ if isinstance(exc, ServerNotReachableError):
+ self._client.log.error(f"Unable to connect to {self._client.address}")
+ raise
+
+ def _download_to_file(self, url: str, dest: Path, retry_state: RetryState) -> int:
+ """One attempt at streaming ``url`` into ``dest``; the file is removed again if the transfer fails."""
+ with self._client._get_streaming(url=url, retry_state=retry_state) as resp:
+ try:
+ resp.raise_for_status()
+ except httpx.HTTPStatusError as exc:
+ # Need to read the response body for error details
+ resp.read()
+ self.handle_error_response(exc=exc)
+
+ bytes_written = 0
+ try:
with dest.open("wb") as f:
for chunk in resp.iter_bytes(chunk_size=65536):
f.write(chunk)
bytes_written += len(chunk)
- return bytes_written
- except ServerNotReachableError:
- self._client.log.error(f"Unable to connect to {self._client.address}")
- raise
+ except BaseException:
+ dest.unlink(missing_ok=True)
+ raise
+ return bytes_written
diff --git a/infrahub_sdk/retry.py b/infrahub_sdk/retry.py
new file mode 100644
index 000000000..f287a6d29
--- /dev/null
+++ b/infrahub_sdk/retry.py
@@ -0,0 +1,277 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import random
+import time
+from collections.abc import Callable, Coroutine, Iterable
+from typing import TYPE_CHECKING, Any
+
+import httpx
+
+from .exceptions import ServerNotReachableError, ServerNotResponsiveError
+
+if TYPE_CHECKING:
+ from .types import InfrahubLoggers
+
+LOGGER = logging.getLogger("infrahub_sdk")
+
+DEFAULT_RETRY_STATUS_CODES: frozenset[int] = frozenset({500, 502, 503, 504})
+"""HTTP status codes treated as transient unless ``retry_status_codes`` says otherwise."""
+
+ESCALATE_AFTER_SECONDS = 300.0
+"""Once an operation has been retrying for this long, retry log lines switch from WARNING to ERROR."""
+
+TRANSIENT_EXCEPTIONS = (ServerNotReachableError, ServerNotResponsiveError)
+"""Client-side failures (connection error, timeout) that are always considered transient."""
+
+CONNECTION_LOST_EXCEPTIONS = (httpx.NetworkError, httpx.ConnectTimeout, httpx.RemoteProtocolError)
+"""httpx failures meaning the connection was lost or never established, mapped to ``ServerNotReachableError``.
+
+``ConnectTimeout`` is included because a load balancer that drops packets during a failover leaves the TCP
+handshake hanging instead of refusing it, and httpx classifies it as a timeout rather than a network error.
+``RemoteProtocolError`` is included because a server or load balancer that goes away mid-request closes the
+socket before answering. httpx reports that as a protocol error rather than a network error, and it is the
+shape a failover takes on a request that was already in flight. The other timeouts (read, write while sending
+a request body, connection pool) are mapped to ``ServerNotResponsiveError`` by the client.
+"""
+
+
+class RetryState:
+ """Retry bookkeeping for one logical operation.
+
+ ``execute_graphql`` creates a state and hands it down to ``_request`` so the transport-level
+ retries (network errors, timeouts, transient HTTP statuses) and the GraphQL-envelope retries
+ (transient ``errors`` inside a 200 response) share a single time budget and attempt counter.
+ """
+
+ __slots__ = ("attempts", "started")
+
+ def __init__(self, started: float) -> None:
+ self.started = started
+ self.attempts = 0
+
+
+class TransientRetryHandler:
+ """Retry policy for transient failures, enabled by ``retry_on_failure``.
+
+ A failure is transient when it is a connection error, a timeout, an HTTP response whose
+ status is listed in ``status_codes``, or a GraphQL error envelope in which every error carries
+ one of those statuses in its ``extensions`` (Infrahub sets ``extensions.http_status`` on the
+ GraphQL endpoint and an integer ``extensions.code`` on the REST endpoints). Anything else is
+ raised immediately, so a bad query or a schema error still fails fast even when the budget is
+ unlimited.
+
+ Retries are spaced with exponential backoff (``base_delay * 2**(attempt - 1)``, capped at
+ ``max_delay``) with equal jitter, and stop once ``max_duration`` seconds have elapsed since the
+ operation started. ``max_duration == 0`` means retry indefinitely.
+
+ The decision methods are pure so a single handler can be shared across concurrent requests;
+ ``send``/``asend`` are the sync and async drivers.
+ """
+
+ def __init__(
+ self,
+ *,
+ enabled: bool,
+ base_delay: float,
+ max_delay: float,
+ max_duration: float,
+ status_codes: Iterable[int] = DEFAULT_RETRY_STATUS_CODES,
+ log: InfrahubLoggers | None = None,
+ clock: Callable[[], float] = time.monotonic,
+ ) -> None:
+ self.enabled = enabled
+ self.base_delay = base_delay
+ self.max_delay = max_delay
+ self.max_duration = max_duration
+ self.status_codes = frozenset(status_codes)
+ self.log = log or LOGGER
+ self.clock = clock
+
+ @property
+ def unlimited(self) -> bool:
+ """``True`` when there is no time limit on retries (``max_duration == 0``)."""
+ return self.max_duration == 0
+
+ def new_state(self) -> RetryState:
+ """Start the retry bookkeeping for a new logical operation."""
+ return RetryState(started=self.clock())
+
+ # --- Classification -------------------------------------------------------------------------
+
+ def is_transient_status(self, status_code: int) -> bool:
+ """Return ``True`` when ``status_code`` is one of the retryable HTTP statuses."""
+ return status_code in self.status_codes
+
+ @staticmethod
+ def graphql_error_status(error: Any) -> int | None:
+ """Return the HTTP status Infrahub attached to a formatted GraphQL error, if any.
+
+ Looks at ``extensions.http_status`` first (GraphQL endpoint) and falls back to an integer
+ ``extensions.code`` (REST endpoints). Booleans are rejected since ``bool`` subclasses ``int``.
+ """
+ if not isinstance(error, dict):
+ return None
+ extensions = error.get("extensions")
+ if not isinstance(extensions, dict):
+ return None
+ for key in ("http_status", "code"):
+ value = extensions.get(key)
+ if isinstance(value, int) and not isinstance(value, bool):
+ return value
+ return None
+
+ def is_transient_graphql_errors(self, errors: Any) -> bool:
+ """Return ``True`` when ``errors`` is a non-empty list whose every entry carries a transient status.
+
+ A single error without a status, or with a non-transient one, makes the whole response
+ non-transient: it is safer to surface an unclassified failure than to retry it forever.
+ """
+ if not isinstance(errors, list) or not errors:
+ return False
+ for error in errors:
+ status = self.graphql_error_status(error)
+ if status is None or not self.is_transient_status(status):
+ return False
+ return True
+
+ @staticmethod
+ def describe_graphql_errors(errors: list[dict[str, Any]]) -> str:
+ """One-line description of a transient GraphQL error envelope for the retry log."""
+ first = errors[0]
+ status = TransientRetryHandler.graphql_error_status(first)
+ message = first.get("message", "")
+ suffix = f" (+{len(errors) - 1} more)" if len(errors) > 1 else ""
+ return f"GraphQL error with HTTP status {status}: {message}{suffix}"
+
+ # --- Budget ---------------------------------------------------------------------------------
+
+ def elapsed(self, state: RetryState) -> float:
+ """Seconds since the operation tracked by ``state`` started."""
+ return self.clock() - state.started
+
+ def remaining(self, state: RetryState) -> float | None:
+ """Seconds left in the budget, or ``None`` when retries are unlimited."""
+ if self.unlimited:
+ return None
+ return max(0.0, self.max_duration - self.elapsed(state))
+
+ def should_retry(self, state: RetryState) -> bool:
+ """Return ``True`` while retrying is enabled and the time budget is not exhausted."""
+ if not self.enabled:
+ return False
+ if self.unlimited:
+ return True
+ return self.elapsed(state) < self.max_duration
+
+ # --- Delay ----------------------------------------------------------------------------------
+
+ def compute_backoff(self, attempt: int) -> float:
+ """Backoff ceiling for retry number ``attempt`` (1-based): ``min(max_delay, base_delay * 2**(attempt - 1))``.
+
+ The exponent is capped so a very long-running unlimited retry cannot overflow ``float``.
+ """
+ exponent = min(max(attempt - 1, 0), 63)
+ return min(self.max_delay, self.base_delay * (2**exponent))
+
+ def jittered_delay(self, ceiling: float) -> float:
+ """Equal-jitter delay drawn from ``[ceiling / 2, ceiling]``."""
+ return random.uniform(ceiling / 2, ceiling)
+
+ def next_delay(self, state: RetryState) -> float:
+ """Delay before retry number ``state.attempts``, never sleeping past the remaining budget."""
+ delay = self.jittered_delay(self.compute_backoff(state.attempts))
+ remaining = self.remaining(state)
+ if remaining is not None:
+ delay = min(delay, remaining)
+ return delay
+
+ # --- Drivers --------------------------------------------------------------------------------
+
+ def prepare_retry(self, state: RetryState, url: str, reason: str) -> float:
+ """Record one more attempt, log it, and return how long to sleep before it.
+
+ Log lines are emitted at WARNING and escalate to ERROR once the operation has been
+ retrying for ``ESCALATE_AFTER_SECONDS``, so an indefinitely retrying job stays visible.
+ """
+ state.attempts += 1
+ delay = self.next_delay(state)
+ elapsed = self.elapsed(state)
+ budget = "no time limit" if self.unlimited else f"{self.max_duration:.0f}s budget"
+ message = (
+ f"Transient failure on {url}: {reason}. "
+ f"Retry {state.attempts} in {delay:.1f}s ({elapsed:.0f}s elapsed, {budget})"
+ )
+ if elapsed >= ESCALATE_AFTER_SECONDS:
+ self.log.error(message)
+ else:
+ self.log.warning(message)
+ return delay
+
+ async def asleep_before_retry(self, state: RetryState, url: str, reason: str) -> None:
+ """Async: record the attempt, log it, and wait for the computed delay."""
+ await asyncio.sleep(self.prepare_retry(state=state, url=url, reason=reason))
+
+ def sleep_before_retry(self, state: RetryState, url: str, reason: str) -> None:
+ """Sync: record the attempt, log it, and wait for the computed delay."""
+ time.sleep(self.prepare_retry(state=state, url=url, reason=reason))
+
+ async def asend(
+ self,
+ send: Callable[[], Coroutine[Any, Any, httpx.Response]],
+ url: str,
+ state: RetryState | None = None,
+ ) -> httpx.Response:
+ """Send via ``send``, retrying transient failures until success or budget exhaustion.
+
+ ``send`` performs one HTTP send per call and is re-invoked per attempt. Pass ``state`` to
+ share the budget with an outer retry loop. When the budget runs out the last transient
+ exception is re-raised, or the last transient response is returned for the caller to
+ handle as it would without retries.
+
+ Raises:
+ ServerNotReachableError: If connection errors persist past the budget.
+ ServerNotResponsiveError: If read timeouts persist past the budget.
+
+ """
+ if not self.enabled:
+ return await send()
+
+ state = state or self.new_state()
+ while True:
+ try:
+ response = await send()
+ except TRANSIENT_EXCEPTIONS as exc:
+ if not self.should_retry(state):
+ raise
+ await self.asleep_before_retry(state=state, url=url, reason=str(exc))
+ continue
+
+ if not self.is_transient_status(response.status_code) or not self.should_retry(state):
+ return response
+ await self.asleep_before_retry(state=state, url=url, reason=f"HTTP {response.status_code}")
+
+ def send(
+ self,
+ send: Callable[[], httpx.Response],
+ url: str,
+ state: RetryState | None = None,
+ ) -> httpx.Response:
+ """Synchronous counterpart of :meth:`asend`; see it for the full contract."""
+ if not self.enabled:
+ return send()
+
+ state = state or self.new_state()
+ while True:
+ try:
+ response = send()
+ except TRANSIENT_EXCEPTIONS as exc:
+ if not self.should_retry(state):
+ raise
+ self.sleep_before_retry(state=state, url=url, reason=str(exc))
+ continue
+
+ if not self.is_transient_status(response.status_code) or not self.should_retry(state):
+ return response
+ self.sleep_before_retry(state=state, url=url, reason=f"HTTP {response.status_code}")
diff --git a/tests/AGENTS.md b/tests/AGENTS.md
index cce67364c..7efb5e597 100644
--- a/tests/AGENTS.md
+++ b/tests/AGENTS.md
@@ -7,6 +7,7 @@ pytest with async auto-mode enabled.
```bash
uv run pytest tests/unit/ # Unit tests (fast, mocked)
uv run pytest tests/integration/ # Integration tests (real Infrahub)
+INFRAHUB_TESTING_FAILOVER=1 uv run pytest tests/integration/test_retry_on_failover.py # Opt-in failover tests, skipped otherwise
uv run pytest -n 4 # Parallel execution
uv run pytest --cov infrahub_sdk # With coverage
uv run pytest tests/unit/test_client.py # Single file
diff --git a/tests/integration/test_retry_on_failover.py b/tests/integration/test_retry_on_failover.py
new file mode 100644
index 000000000..82a703501
--- /dev/null
+++ b/tests/integration/test_retry_on_failover.py
@@ -0,0 +1,327 @@
+"""Retry behaviour when Infrahub goes away in the middle of a request.
+
+These tests reproduce a failover against a real Infrahub deployment. The API is told to sit on
+every GraphQL request for a few seconds, an upsert mutation is started, and containers are killed
+while that mutation is still in flight, in each of the two shapes a failover takes:
+
+- The HAProxy load balancer in front of the API servers is restarted, so the connection is dropped
+ before a single byte of the response arrives.
+- The API servers behind it are restarted while HAProxy stays up, so the connection to the client
+ holds and HAProxy answers for them with a transient HTTP status.
+
+The mutation is an upsert on purpose: the server keeps processing a request whose client has gone
+away, so the first attempt may well have been applied by the time the retry is sent. Retrying is
+at-least-once, and only an idempotent mutation can be replayed safely.
+
+The module is opt-in and skipped by default, so CI never runs it: the tests restart containers,
+pin a host port and take several minutes. Run them with::
+
+ INFRAHUB_TESTING_FAILOVER=1 uv run pytest tests/integration/test_retry_on_failover.py
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import socket
+import subprocess # noqa: S404
+import threading
+import time
+from collections.abc import Callable, Generator, Sequence
+from pathlib import Path
+from typing import TypeVar
+
+import httpx
+import pytest
+from infrahub_testcontainers.container import PROJECT_ENV_VARIABLES, InfrahubDockerCompose
+
+from infrahub_sdk import Config, InfrahubClient, InfrahubClientSync
+from infrahub_sdk.exceptions import ServerNotReachableError
+from infrahub_sdk.testing.docker import TestInfrahubDockerClient
+
+RESPONSE_DELAY = 10
+"""Seconds the API waits before handling each GraphQL request, to widen the in-flight window."""
+
+RESTART_AFTER = 3.0
+"""Seconds into the mutation at which the containers are restarted."""
+
+RETRY_TIMEOUT = 180.0
+"""Upper bound on a retrying call, so an unlimited retry budget cannot hang the suite."""
+
+LOAD_BALANCER = "infrahub-server-lb"
+API_SERVER = "infrahub-server"
+
+ADMIN_TOKEN = PROJECT_ENV_VARIABLES["INFRAHUB_TESTING_INITIAL_ADMIN_TOKEN"]
+
+FAILOVER_TESTS_ENV = "INFRAHUB_TESTING_FAILOVER"
+"""Environment variable that opts into these tests; unset, the whole module is skipped."""
+
+pytestmark = pytest.mark.skipif(
+ os.environ.get(FAILOVER_TESTS_ENV) != "1",
+ reason=f"failover tests restart containers and take minutes; opt in with {FAILOVER_TESTS_ENV}=1",
+)
+
+T = TypeVar("T")
+
+
+def reserve_host_port() -> int:
+ """Return a free TCP port on the host."""
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+ sock.bind(("127.0.0.1", 0))
+ return int(sock.getsockname()[1])
+
+
+def restart_containers(container_names: Sequence[str]) -> None:
+ """Kill and restart containers, dropping every connection they were carrying."""
+ subprocess.run( # noqa: S603
+ ["docker", "restart", "--time", "0", *container_names], # noqa: S607
+ check=True,
+ capture_output=True,
+ )
+
+
+async def restart_after(container_names: Sequence[str], delay: float) -> None:
+ """Restart containers ``delay`` seconds from now, while the caller has a request in flight."""
+ await asyncio.sleep(delay)
+ await asyncio.to_thread(restart_containers, container_names)
+
+
+def run_with_timeout(func: Callable[[], T], timeout: float, on_timeout: Callable[[], None], grace: float = 30.0) -> T:
+ """Run ``func`` in a thread and fail the test if it has not returned after ``timeout`` seconds.
+
+ The synchronous client has no counterpart to ``asyncio.wait_for``, and an unlimited retry budget
+ against a load balancer that never comes back would otherwise hang the whole suite. On timeout,
+ ``on_timeout`` is called to make the call give up, so the thread does not keep retrying against
+ the deployment behind the next test, and it gets ``grace`` seconds to end before the test fails.
+ The thread is a daemon so, should it not end even then, it cannot block interpreter exit.
+ """
+ outcome: list[T] = []
+ failure: list[BaseException] = []
+
+ def target() -> None:
+ try:
+ outcome.append(func())
+ except BaseException as exc: # re-raised in the calling thread below
+ failure.append(exc)
+
+ worker = threading.Thread(target=target, daemon=True)
+ worker.start()
+ worker.join(timeout)
+ if worker.is_alive():
+ on_timeout()
+ worker.join(grace)
+ still_running = " and did not stop within the grace period" if worker.is_alive() else ""
+ pytest.fail(f"the call was still running after {timeout} seconds{still_running}")
+ if failure:
+ raise failure[0]
+ return outcome[0]
+
+
+def wait_until_reachable(address: str, timeout: float = 120.0) -> None:
+ """Block until the load balancer answers again after a restart.
+
+ Raises:
+ httpx.HTTPError: If the load balancer is still not answering after ``timeout`` seconds.
+
+ """
+ deadline = time.monotonic() + timeout
+ while True:
+ try:
+ httpx.get(f"{address}/api/config", timeout=10).raise_for_status()
+ return
+ except httpx.HTTPError:
+ if time.monotonic() > deadline:
+ raise
+ time.sleep(1)
+
+
+class TestRetryOnFailover(TestInfrahubDockerClient):
+ @pytest.fixture(scope="class")
+ def infrahub_compose(
+ self,
+ tmp_directory: Path,
+ remote_repos_dir: Path, # initialize repository before running docker compose to fix permissions issues
+ remote_backups_dir: Path,
+ infrahub_version: str,
+ deployment_type: str | None,
+ ) -> Generator[InfrahubDockerCompose, None, None]:
+ """Publish the load balancer on a fixed host port.
+
+ Docker picks a new host port every time a container starts, so a restarted load balancer
+ would come back at a different address and the client would retry against a dead port.
+ Pinning the port keeps the address stable across the restart, the way it is in a real
+ failover. The variable has to stay set for as long as compose is driven, so the whole
+ fixture lifetime runs inside the patched environment.
+
+ Yields:
+ InfrahubDockerCompose: the compose project, with the server port pinned.
+
+ """
+ with pytest.MonkeyPatch.context() as monkeypatch:
+ monkeypatch.setenv("INFRAHUB_TESTING_SERVER_PORT", str(reserve_host_port()))
+ yield InfrahubDockerCompose.init(
+ directory=tmp_directory,
+ version=infrahub_version,
+ deployment_type=deployment_type,
+ )
+
+ @pytest.fixture(scope="class")
+ def address(self, infrahub_port: int) -> str:
+ return f"http://localhost:{infrahub_port}"
+
+ @pytest.fixture(scope="class")
+ def load_balancer_container(self, infrahub_compose: InfrahubDockerCompose, infrahub_port: int) -> str:
+ return str(infrahub_compose.get_container(service_name=LOAD_BALANCER).Name)
+
+ @pytest.fixture(scope="class")
+ def api_server_containers(self, infrahub_compose: InfrahubDockerCompose, infrahub_port: int) -> list[str]:
+ return [
+ str(container.Name) for container in infrahub_compose.get_containers() if container.Service == API_SERVER
+ ]
+
+ @pytest.fixture(scope="class")
+ def slow_api(
+ self, infrahub_compose: InfrahubDockerCompose, address: str, infrahub_port: int
+ ) -> Generator[None, None, None]:
+ """Make every GraphQL request take RESPONSE_DELAY seconds, on every API worker.
+
+ Yields:
+ None: once the delay is active on every worker process.
+
+ """
+ infrahub_compose.set_server_response_delay(RESPONSE_DELAY)
+ yield
+ wait_until_reachable(address)
+ infrahub_compose.set_server_response_delay(0)
+
+ @pytest.fixture
+ def slow_api_across_api_restart(
+ self, infrahub_compose: InfrahubDockerCompose, address: str, slow_api: None
+ ) -> Generator[None, None, None]:
+ """Re-apply the response delay for a test that restarts the API servers and so clears it.
+
+ The delay lives in the memory of each API worker process, set by a broadcast on the
+ message bus, so restarting those processes drops it. Restoring it here keeps the class
+ fixture's promise and leaves this test independent of the order tests run in. The
+ broadcast only reaches workers that are up, so the teardown first waits for the restarted
+ servers to answer again.
+
+ Yields:
+ None: with the delay active, as ``slow_api`` leaves it.
+
+ """
+ yield
+ wait_until_reachable(address)
+ infrahub_compose.set_server_response_delay(RESPONSE_DELAY)
+
+ def build_client(self, address: str, *, retry_on_failure: bool) -> InfrahubClient:
+ return InfrahubClient(config=self.build_config(address, retry_on_failure=retry_on_failure))
+
+ def build_config(self, address: str, *, retry_on_failure: bool) -> Config:
+ return Config(
+ address=address,
+ api_token=ADMIN_TOKEN,
+ retry_on_failure=retry_on_failure,
+ max_retry_duration=0, # retry for as long as it takes
+ retry_delay=1,
+ retry_max_delay=5,
+ )
+
+ async def test_failover_aborts_the_mutation_without_retry(
+ self, address: str, load_balancer_container: str, slow_api: None
+ ) -> None:
+ """Without retries, losing the load balancer mid-mutation surfaces as a hard failure."""
+ client = self.build_client(address, retry_on_failure=False)
+ await client.schema.all() # warm the schema cache so only the mutation is in flight
+
+ node = await client.create(kind="BuiltinTag", name="failover-no-retry")
+
+ restart = asyncio.create_task(restart_after([load_balancer_container], RESTART_AFTER))
+ with pytest.raises(ServerNotReachableError):
+ await node.save(allow_upsert=True)
+ await restart
+
+ wait_until_reachable(address)
+
+ async def test_failover_is_survived_with_unlimited_retries(
+ self, address: str, load_balancer_container: str, slow_api: None, caplog: pytest.LogCaptureFixture
+ ) -> None:
+ """With retries enabled the same failover only delays the mutation."""
+ caplog.set_level("WARNING", logger="infrahub_sdk")
+ client = self.build_client(address, retry_on_failure=True)
+ await client.schema.all()
+
+ node = await client.create(kind="BuiltinTag", name="failover-async")
+
+ restart = asyncio.create_task(restart_after([load_balancer_container], RESTART_AFTER))
+ await asyncio.wait_for(node.save(allow_upsert=True), timeout=RETRY_TIMEOUT)
+ await restart
+
+ assert node.id, "the upsert should have returned the id of the saved node"
+ assert any("Transient failure" in record.message for record in caplog.records), (
+ "the mutation should have been retried, not served on the first attempt"
+ )
+
+ saved = await client.get(kind="BuiltinTag", name__value="failover-async")
+ assert saved.id == node.id
+
+ def test_failover_is_survived_with_unlimited_retries_sync(
+ self, address: str, load_balancer_container: str, slow_api: None, caplog: pytest.LogCaptureFixture
+ ) -> None:
+ """The synchronous client retries the same way the asynchronous one does."""
+ caplog.set_level("WARNING", logger="infrahub_sdk")
+ client = InfrahubClientSync(config=self.build_config(address, retry_on_failure=True))
+ client.schema.all()
+
+ node = client.create(kind="BuiltinTag", name="failover-sync")
+
+ def stop_retrying() -> None:
+ # With retries off, the next failed attempt raises instead of sleeping again.
+ client.retry_on_failure = False
+
+ restart = threading.Timer(RESTART_AFTER, restart_containers, args=([load_balancer_container],))
+ restart.start()
+ try:
+ run_with_timeout(lambda: node.save(allow_upsert=True), timeout=RETRY_TIMEOUT, on_timeout=stop_retrying)
+ finally:
+ restart.join()
+
+ assert node.id, "the upsert should have returned the id of the saved node"
+ assert any("Transient failure" in record.message for record in caplog.records), (
+ "the mutation should have been retried, not served on the first attempt"
+ )
+
+ saved = client.get(kind="BuiltinTag", name__value="failover-sync")
+ assert saved.id == node.id
+
+ async def test_api_server_restart_is_survived_with_unlimited_retries(
+ self,
+ address: str,
+ api_server_containers: list[str],
+ slow_api_across_api_restart: None,
+ caplog: pytest.LogCaptureFixture,
+ ) -> None:
+ """The other shape of a failover: the load balancer survives, the API servers behind it do not.
+
+ Nothing breaks the connection to the client here, so instead of a dropped socket HAProxy
+ answers the in-flight request itself with 502, and answers whatever arrives while the
+ servers boot with 503. Both are in retry_status_codes, so the mutation only has to wait.
+ """
+ caplog.set_level("WARNING", logger="infrahub_sdk")
+ client = self.build_client(address, retry_on_failure=True)
+ await client.schema.all()
+
+ node = await client.create(kind="BuiltinTag", name="failover-api-servers")
+
+ restart = asyncio.create_task(restart_after(api_server_containers, RESTART_AFTER))
+ await asyncio.wait_for(node.save(allow_upsert=True), timeout=RETRY_TIMEOUT)
+ await restart
+
+ assert node.id, "the upsert should have returned the id of the saved node"
+ retries = [record.message for record in caplog.records if "Transient failure" in record.message]
+ assert any("HTTP 502" in message for message in retries), (
+ f"the request in flight when the API servers died should have been retried on a 502: {retries}"
+ )
+
+ saved = await client.get(kind="BuiltinTag", name__value="failover-api-servers")
+ assert saved.id == node.id
diff --git a/tests/unit/sdk/test_client.py b/tests/unit/sdk/test_client.py
index c20227093..d63c8c564 100644
--- a/tests/unit/sdk/test_client.py
+++ b/tests/unit/sdk/test_client.py
@@ -23,7 +23,7 @@
pytestmark = pytest.mark.httpx_mock(can_send_already_matched_responses=True)
-excluded_methods = ["request_context"]
+excluded_methods = ["request_context", "retry_delay", "retry_on_failure"]
async_client_methods = [
method for method in dir(InfrahubClient) if not method.startswith("_") and method not in excluded_methods
diff --git a/tests/unit/sdk/test_rate_limit_retry.py b/tests/unit/sdk/test_rate_limit_retry.py
index 8306afb15..8d52d2797 100644
--- a/tests/unit/sdk/test_rate_limit_retry.py
+++ b/tests/unit/sdk/test_rate_limit_retry.py
@@ -7,9 +7,11 @@
from __future__ import annotations
+import asyncio
import io
import logging
import re
+import time
from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
@@ -82,8 +84,8 @@ async def fake_async_sleep(delay: float) -> None:
def fake_sync_sleep(delay: float) -> None:
recorded.append(delay)
- monkeypatch.setattr(client_module.asyncio, "sleep", fake_async_sleep)
- monkeypatch.setattr(client_module.time, "sleep", fake_sync_sleep)
+ monkeypatch.setattr(asyncio, "sleep", fake_async_sleep)
+ monkeypatch.setattr(time, "sleep", fake_sync_sleep)
return recorded
diff --git a/tests/unit/sdk/test_retry.py b/tests/unit/sdk/test_retry.py
new file mode 100644
index 000000000..ddf58477d
--- /dev/null
+++ b/tests/unit/sdk/test_retry.py
@@ -0,0 +1,1108 @@
+"""Tests for transient-failure retries (``retry_on_failure``) on the async and sync clients.
+
+Covers the ``TransientRetryHandler`` decision logic (classification, backoff, budget), the
+transport-level driver in ``_request`` (connection errors, timeouts, transient HTTP statuses),
+the GraphQL-envelope retries in ``execute_graphql`` sharing one budget with the transport layer,
+the opt-in default, budget exhaustion, unlimited mode, log escalation and configuration plumbing.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import io
+import logging
+import tempfile
+import threading
+import time
+from collections.abc import Callable, Iterator, Sequence
+from dataclasses import dataclass
+from typing import IO, TYPE_CHECKING, Any, BinaryIO
+
+import httpx
+import pytest
+from pydantic import ValidationError
+from pytest_httpx import IteratorStream
+
+from infrahub_sdk import InfrahubClient, InfrahubClientSync
+from infrahub_sdk import client as client_module
+from infrahub_sdk.config import Config
+from infrahub_sdk.exceptions import GraphQLError, ServerNotReachableError, ServerNotResponsiveError
+from infrahub_sdk.file_handler import FileHandler, FileHandlerSync
+from infrahub_sdk.retry import ESCALATE_AFTER_SECONDS, TransientRetryHandler
+from infrahub_sdk.types import HTTPMethod
+
+if TYPE_CHECKING:
+ from pathlib import Path
+
+ from pytest_httpx import HTTPXMock
+
+CLIENT_TYPES = ["standard", "sync"]
+GRAPHQL_URL = "http://mock/graphql/main"
+QUERY = "query { InfraDevice { edges { node { id } } } }"
+LOGGER_NAME = "infrahub_sdk"
+
+
+class FakeClock:
+ """Deterministic monotonic clock advanced explicitly or by the patched sleeps."""
+
+ def __init__(self, now: float = 1000.0, step: float = 0.0) -> None:
+ self.now = now
+ self.step = step
+
+ def __call__(self) -> float:
+ value = self.now
+ self.now += self.step
+ return value
+
+ def advance(self, seconds: float) -> None:
+ self.now += seconds
+
+
+class ScriptedRequester:
+ """A pluggable ``requester``/``sync_requester`` replaying scripted responses or exceptions.
+
+ Each send returns the next ``httpx.Response``, or raises the next exception, and increments
+ ``call_count`` so a test can assert exactly how many HTTP sends were performed.
+ """
+
+ def __init__(self, steps: Sequence[httpx.Response | Exception]) -> None:
+ self._steps = list(steps)
+ self.call_count = 0
+
+ def _next(self) -> httpx.Response:
+ step = self._steps[self.call_count]
+ self.call_count += 1
+ if isinstance(step, Exception):
+ raise step
+ return step
+
+ def sync_request(
+ self,
+ url: str,
+ method: HTTPMethod,
+ headers: dict[str, Any],
+ timeout: int,
+ payload: dict | None = None,
+ ) -> httpx.Response:
+ return self._next()
+
+ async def async_request(
+ self,
+ url: str,
+ method: HTTPMethod,
+ headers: dict[str, Any],
+ timeout: int,
+ payload: dict | None = None,
+ ) -> httpx.Response:
+ return self._next()
+
+
+def _patch_sleep(monkeypatch: pytest.MonkeyPatch, clock: FakeClock | None = None) -> list[float]:
+ """Replace async/sync sleep with recorders that never wait; advance ``clock`` by each delay."""
+ recorded: list[float] = []
+
+ def record(delay: float) -> None:
+ recorded.append(delay)
+ if clock is not None:
+ clock.advance(delay)
+
+ async def fake_async_sleep(delay: float) -> None:
+ record(delay)
+
+ def fake_sync_sleep(delay: float) -> None:
+ record(delay)
+
+ monkeypatch.setattr(asyncio, "sleep", fake_async_sleep)
+ monkeypatch.setattr(time, "sleep", fake_sync_sleep)
+ return recorded
+
+
+def _no_jitter(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Make delays deterministic: the jittered delay becomes the backoff ceiling itself."""
+ monkeypatch.setattr(TransientRetryHandler, "jittered_delay", lambda _self, ceiling: ceiling)
+
+
+def _build_client(
+ client_type: str,
+ requester: ScriptedRequester,
+ clock: FakeClock | None = None,
+ **overrides: bool | int | list[int],
+) -> InfrahubClient | InfrahubClientSync:
+ config_kwargs: dict[str, Any] = dict(overrides)
+ client: InfrahubClient | InfrahubClientSync
+ if client_type == "standard":
+ client = InfrahubClient(
+ config=Config(address="http://mock", requester=requester.async_request, **config_kwargs)
+ )
+ else:
+ client = InfrahubClientSync(
+ config=Config(address="http://mock", sync_requester=requester.sync_request, **config_kwargs)
+ )
+ if clock is not None:
+ client._retry_handler.clock = clock
+ return client
+
+
+def _build_client_over_httpx(client_type: str, **overrides: bool | int) -> InfrahubClient | InfrahubClientSync:
+ """Build a client without a custom requester, so requests go through httpx and its exceptions."""
+ config_kwargs: dict[str, Any] = dict(overrides)
+ if client_type == "standard":
+ return InfrahubClient(config=Config(address="http://mock", **config_kwargs))
+ return InfrahubClientSync(config=Config(address="http://mock", **config_kwargs))
+
+
+async def _execute_graphql(client: InfrahubClient | InfrahubClientSync) -> dict:
+ if isinstance(client, InfrahubClient):
+ return await client.execute_graphql(query=QUERY)
+ return client.execute_graphql(query=QUERY)
+
+
+async def _request(client: InfrahubClient | InfrahubClientSync, url: str = GRAPHQL_URL) -> httpx.Response:
+ if isinstance(client, InfrahubClient):
+ return await client._request(url=url, method=HTTPMethod.POST, headers={}, timeout=10, payload={})
+ return client._request(url=url, method=HTTPMethod.POST, headers={}, timeout=10, payload={})
+
+
+def _response(status_code: int, json: dict | None = None) -> httpx.Response:
+ """Build a response with a request attached so ``raise_for_status`` works on it."""
+ return httpx.Response(status_code=status_code, json=json, request=httpx.Request("POST", GRAPHQL_URL))
+
+
+def _ok(data: dict | None = None) -> httpx.Response:
+ return _response(200, json={"data": data or {"result": "ok"}})
+
+
+def _graphql_errors(*errors: dict, status_code: int = 200) -> httpx.Response:
+ return _response(status_code, json={"data": None, "errors": list(errors)})
+
+
+def _transient_error(status: int = 503, message: str = "Unable to connect to the database") -> dict:
+ return {"message": message, "extensions": {"code": "DATABASE_UNAVAILABLE", "http_status": status}}
+
+
+def _connection_error() -> ServerNotReachableError:
+ return ServerNotReachableError(address="http://mock")
+
+
+def _timeout() -> ServerNotResponsiveError:
+ return ServerNotResponsiveError(url=GRAPHQL_URL, timeout=10)
+
+
+def _handler(**overrides: bool | float | list[int] | FakeClock) -> TransientRetryHandler:
+ params: dict[str, Any] = {"enabled": True, "base_delay": 5, "max_delay": 60, "max_duration": 300}
+ params.update(overrides)
+ return TransientRetryHandler(**params)
+
+
+# --- Handler decision logic -------------------------------------------------------------------------
+
+
+@dataclass
+class StatusCase:
+ name: str
+ status: int
+ expected: bool
+
+
+STATUS_CASES = [
+ StatusCase(name="502", status=502, expected=True),
+ StatusCase(name="503", status=503, expected=True),
+ StatusCase(name="504", status=504, expected=True),
+ StatusCase(name="500-unclassified-server-error", status=500, expected=True),
+ StatusCase(name="429-has-its-own-handler", status=429, expected=False),
+ StatusCase(name="400", status=400, expected=False),
+ StatusCase(name="200", status=200, expected=False),
+]
+
+
+@pytest.mark.parametrize("case", [pytest.param(tc, id=tc.name) for tc in STATUS_CASES])
+def test_handler_default_transient_statuses(case: StatusCase) -> None:
+ assert _handler().is_transient_status(case.status) is case.expected
+
+
+def test_handler_custom_status_codes_replace_the_default_set() -> None:
+ handler = _handler(status_codes=[500, 503])
+ assert handler.is_transient_status(500)
+ assert handler.is_transient_status(503)
+ assert not handler.is_transient_status(502)
+
+
+@dataclass
+class GraphQLErrorsCase:
+ name: str
+ errors: Any
+ expected: bool
+
+
+GRAPHQL_ERRORS_CASES = [
+ GraphQLErrorsCase(name="graphql-http-status", errors=[_transient_error(503)], expected=True),
+ GraphQLErrorsCase(name="all-transient", errors=[_transient_error(502), _transient_error(504)], expected=True),
+ GraphQLErrorsCase(
+ name="rest-integer-code", errors=[{"message": "db down", "extensions": {"code": 503}}], expected=True
+ ),
+ GraphQLErrorsCase(
+ name="one-unclassified", errors=[_transient_error(503), {"message": "Unknown field"}], expected=False
+ ),
+ GraphQLErrorsCase(
+ name="non-transient-status",
+ errors=[{"message": "not found", "extensions": {"code": "NODE_NOT_FOUND", "http_status": 404}}],
+ expected=False,
+ ),
+ GraphQLErrorsCase(
+ name="bool-is-not-a-status", errors=[{"message": "x", "extensions": {"code": True}}], expected=False
+ ),
+ GraphQLErrorsCase(name="malformed-extensions", errors=[{"message": "x", "extensions": "oops"}], expected=False),
+ GraphQLErrorsCase(name="malformed-error", errors=["not a dict"], expected=False),
+ GraphQLErrorsCase(name="empty", errors=[], expected=False),
+ GraphQLErrorsCase(name="none", errors=None, expected=False),
+ GraphQLErrorsCase(name="not-a-list", errors="errors", expected=False),
+]
+
+
+@pytest.mark.parametrize("case", [pytest.param(tc, id=tc.name) for tc in GRAPHQL_ERRORS_CASES])
+def test_handler_graphql_error_classification(case: GraphQLErrorsCase) -> None:
+ assert _handler().is_transient_graphql_errors(case.errors) is case.expected
+
+
+def test_handler_graphql_error_status_prefers_http_status_over_code() -> None:
+ error = {"message": "x", "extensions": {"code": 200, "http_status": 503}}
+ assert TransientRetryHandler.graphql_error_status(error) == 503
+
+
+def test_handler_describe_graphql_errors_mentions_status_message_and_count() -> None:
+ description = TransientRetryHandler.describe_graphql_errors(
+ [_transient_error(503, "db down"), _transient_error(502)]
+ )
+ assert description == "GraphQL error with HTTP status 503: db down (+1 more)"
+
+
+def test_handler_backoff_doubles_from_base_delay_and_clamps() -> None:
+ handler = _handler(base_delay=5, max_delay=60)
+ assert [handler.compute_backoff(attempt) for attempt in range(1, 7)] == [5, 10, 20, 40, 60, 60]
+ assert handler.compute_backoff(0) == 5
+ assert handler.compute_backoff(10_000) == 60
+
+
+def test_handler_jitter_stays_within_half_to_full_ceiling() -> None:
+ handler = _handler()
+ for _ in range(200):
+ assert 30 <= handler.jittered_delay(60) <= 60
+
+
+def test_handler_budget_is_time_based() -> None:
+ clock = FakeClock()
+ handler = _handler(max_duration=10, clock=clock)
+ state = handler.new_state()
+
+ assert not handler.unlimited
+ assert handler.should_retry(state)
+ assert handler.remaining(state) == 10
+ clock.advance(9.9)
+ assert handler.should_retry(state)
+ clock.advance(0.1)
+ assert not handler.should_retry(state)
+ assert handler.remaining(state) == 0
+
+
+def test_handler_zero_duration_means_unlimited() -> None:
+ clock = FakeClock()
+ handler = _handler(max_duration=0, clock=clock)
+ state = handler.new_state()
+
+ assert handler.unlimited
+ assert handler.remaining(state) is None
+ clock.advance(10**6)
+ assert handler.should_retry(state)
+
+
+def test_handler_disabled_never_retries() -> None:
+ handler = _handler(enabled=False)
+ assert not handler.should_retry(handler.new_state())
+
+
+def test_handler_next_delay_never_sleeps_past_the_budget(monkeypatch: pytest.MonkeyPatch) -> None:
+ _no_jitter(monkeypatch)
+ clock = FakeClock()
+ handler = _handler(max_duration=12, clock=clock)
+ state = handler.new_state()
+
+ state.attempts = 3 # backoff ceiling would be 20s
+ assert handler.next_delay(state) == 12
+ clock.advance(10)
+ assert handler.next_delay(state) == 2
+
+
+def test_handler_retry_log_escalates_to_error_after_threshold(
+ monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+ _no_jitter(monkeypatch)
+ caplog.set_level(logging.WARNING, logger=LOGGER_NAME)
+ clock = FakeClock()
+ handler = _handler(max_duration=0, clock=clock)
+ state = handler.new_state()
+
+ handler.prepare_retry(state=state, url=GRAPHQL_URL, reason="HTTP 503")
+ clock.advance(ESCALATE_AFTER_SECONDS)
+ handler.prepare_retry(state=state, url=GRAPHQL_URL, reason="HTTP 503")
+
+ assert [record.levelno for record in caplog.records] == [logging.WARNING, logging.ERROR]
+ assert "Retry 1 in 5.0s (0s elapsed, no time limit)" in caplog.records[0].message
+ assert "Retry 2 in 10.0s (300s elapsed, no time limit)" in caplog.records[1].message
+
+
+# --- Opt-in default ---------------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_disabled_by_default_surfaces_connection_error_immediately(
+ client_type: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ requester = ScriptedRequester([_connection_error(), _ok()])
+ client = _build_client(client_type, requester)
+
+ assert client.retry_on_failure is False
+ with pytest.raises(ServerNotReachableError, match="Unable to connect to 'http://mock'"):
+ await _execute_graphql(client)
+ assert requester.call_count == 1
+ assert recorded_sleeps == []
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_disabled_by_default_returns_transient_status_untouched(
+ client_type: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ requester = ScriptedRequester([_response(503), _ok()])
+ client = _build_client(client_type, requester)
+
+ response = await _request(client)
+
+ assert response.status_code == 503
+ assert requester.call_count == 1
+ assert recorded_sleeps == []
+
+
+# --- Transport-level retries ------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_retries_connection_error_and_timeout_then_succeeds(
+ client_type: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ _no_jitter(monkeypatch)
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ requester = ScriptedRequester([_connection_error(), _timeout(), _ok()])
+ client = _build_client(client_type, requester, retry_on_failure=True)
+
+ data = await _execute_graphql(client)
+
+ assert data == {"result": "ok"}
+ assert requester.call_count == 3
+ assert recorded_sleeps == [5, 10]
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_retries_transient_http_statuses_with_growing_backoff(
+ client_type: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ _no_jitter(monkeypatch)
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ requester = ScriptedRequester(
+ [
+ _response(503),
+ _response(502),
+ _response(504),
+ _ok(),
+ ]
+ )
+ client = _build_client(client_type, requester, retry_on_failure=True)
+
+ response = await _request(client)
+
+ assert response.status_code == 200
+ assert requester.call_count == 4
+ assert recorded_sleeps == [5, 10, 20]
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_custom_status_codes_control_what_is_transient(client_type: str, monkeypatch: pytest.MonkeyPatch) -> None:
+ _patch_sleep(monkeypatch)
+
+ default_requester = ScriptedRequester([_response(500), _ok()])
+ default_client = _build_client(client_type, default_requester, retry_on_failure=True)
+ assert (await _request(default_client)).status_code == 200
+ assert default_requester.call_count == 2
+
+ custom_requester = ScriptedRequester([_response(500), _ok()])
+ custom_client = _build_client(client_type, custom_requester, retry_on_failure=True, retry_status_codes=[502, 503])
+ assert (await _request(custom_client)).status_code == 500
+ assert custom_requester.call_count == 1
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_rest_endpoints_are_retried_too(client_type: str, monkeypatch: pytest.MonkeyPatch) -> None:
+ """``query_gql_query`` (used by generators to collect their data) goes through the same retry path."""
+ _patch_sleep(monkeypatch)
+ payload: dict[str, Any] = {"data": {"InfraDevice": {"edges": []}}}
+ requester = ScriptedRequester([_connection_error(), _response(200, json=payload)])
+ client = _build_client(client_type, requester, retry_on_failure=True)
+
+ if isinstance(client, InfrahubClient):
+ result = await client.query_gql_query(name="my_query")
+ else:
+ result = client.query_gql_query(name="my_query")
+
+ assert result == payload
+ assert requester.call_count == 2
+
+
+# --- Lost connections and timeouts, whichever way httpx reports them --------------------------------
+
+
+@dataclass
+class TransportErrorCase:
+ name: str
+ exception: Exception
+
+
+LOST_CONNECTION_CASES = [
+ TransportErrorCase(name="refused", exception=httpx.ConnectError("connection refused")),
+ TransportErrorCase(name="handshake-timed-out", exception=httpx.ConnectTimeout("timed out")),
+ TransportErrorCase(name="reset", exception=httpx.ReadError("connection reset by peer")),
+ TransportErrorCase(
+ name="disconnected-mid-request",
+ exception=httpx.RemoteProtocolError("Server disconnected without sending a response."),
+ ),
+]
+
+TIMEOUT_CASES = [
+ TransportErrorCase(name="read", exception=httpx.ReadTimeout("timed out")),
+ TransportErrorCase(name="write", exception=httpx.WriteTimeout("timed out")),
+ TransportErrorCase(name="pool", exception=httpx.PoolTimeout("timed out")),
+]
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+@pytest.mark.parametrize("case", [pytest.param(tc, id=tc.name) for tc in LOST_CONNECTION_CASES])
+async def test_lost_connection_is_transient_whichever_httpx_error_reports_it(
+ client_type: str, case: TransportErrorCase, httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ """Losing the connection before a response arrives is transient however httpx classifies it.
+
+ ``RemoteProtocolError`` is the shape a failover takes on a request that is already in flight:
+ the server or the load balancer in front of it closes the socket before sending a single byte
+ of the response. httpx reports that as a protocol error rather than a network error, so it
+ only reaches the retry handler if the client maps it too. ``ConnectTimeout`` is a handshake left
+ hanging by a load balancer dropping packets mid-failover, which httpx files under timeouts.
+ """
+ _no_jitter(monkeypatch)
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ httpx_mock.add_exception(case.exception, url=GRAPHQL_URL)
+ httpx_mock.add_response(url=GRAPHQL_URL, json={"data": {"result": "ok"}})
+ client = _build_client_over_httpx(client_type, retry_on_failure=True)
+
+ assert await _execute_graphql(client) == {"result": "ok"}
+ assert recorded_sleeps == [5]
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+@pytest.mark.parametrize("case", [pytest.param(tc, id=tc.name) for tc in LOST_CONNECTION_CASES])
+async def test_lost_connection_surfaces_as_unreachable_when_retries_are_disabled(
+ client_type: str, case: TransportErrorCase, httpx_mock: HTTPXMock
+) -> None:
+ httpx_mock.add_exception(case.exception, url=GRAPHQL_URL)
+ client = _build_client_over_httpx(client_type, retry_on_failure=False)
+
+ with pytest.raises(ServerNotReachableError, match="Unable to connect to 'http://mock'"):
+ await _execute_graphql(client)
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+@pytest.mark.parametrize("case", [pytest.param(tc, id=tc.name) for tc in TIMEOUT_CASES])
+async def test_timeout_is_transient_whichever_phase_of_the_request_it_hits(
+ client_type: str, case: TransportErrorCase, httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ """Every httpx timeout is transient, not only ``ReadTimeout``.
+
+ A write timeout is an upload the server stopped draining; a pool timeout is a connection that
+ never became available. Both leave the request unprocessed, so retrying it is safe.
+ """
+ _no_jitter(monkeypatch)
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ httpx_mock.add_exception(case.exception, url=GRAPHQL_URL)
+ httpx_mock.add_response(url=GRAPHQL_URL, json={"data": {"result": "ok"}})
+ client = _build_client_over_httpx(client_type, retry_on_failure=True)
+
+ assert await _execute_graphql(client) == {"result": "ok"}
+ assert recorded_sleeps == [5]
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+@pytest.mark.parametrize("case", [pytest.param(tc, id=tc.name) for tc in TIMEOUT_CASES])
+async def test_timeout_surfaces_as_not_responsive_when_retries_are_disabled(
+ client_type: str, case: TransportErrorCase, httpx_mock: HTTPXMock
+) -> None:
+ httpx_mock.add_exception(case.exception, url=GRAPHQL_URL)
+ client = _build_client_over_httpx(client_type, retry_on_failure=False)
+
+ with pytest.raises(ServerNotResponsiveError, match="Unable to read from"):
+ await _execute_graphql(client)
+
+
+# --- GraphQL-envelope retries -----------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_retries_transient_graphql_errors(
+ client_type: str, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+ _no_jitter(monkeypatch)
+ caplog.set_level(logging.WARNING, logger=LOGGER_NAME)
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ requester = ScriptedRequester(
+ [
+ _graphql_errors(_transient_error(503, "Unable to connect to the database")),
+ _graphql_errors(_transient_error(502), _transient_error(504)),
+ _ok(),
+ ]
+ )
+ client = _build_client(client_type, requester, retry_on_failure=True)
+
+ data = await _execute_graphql(client)
+
+ assert data == {"result": "ok"}
+ assert requester.call_count == 3
+ assert recorded_sleeps == [5, 10]
+ assert "GraphQL error with HTTP status 503: Unable to connect to the database" in caplog.text
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_graphql_and_transport_retries_share_one_attempt_counter(
+ client_type: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ """Backoff keeps growing across layers: the envelope retry does not restart from the base delay."""
+ _no_jitter(monkeypatch)
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ requester = ScriptedRequester(
+ [
+ _response(503),
+ _graphql_errors(_transient_error(503)),
+ _connection_error(),
+ _ok(),
+ ]
+ )
+ client = _build_client(client_type, requester, retry_on_failure=True)
+
+ data = await _execute_graphql(client)
+
+ assert data == {"result": "ok"}
+ assert requester.call_count == 4
+ assert recorded_sleeps == [5, 10, 20]
+
+
+@dataclass
+class NonTransientCase:
+ name: str
+ errors: list[dict]
+
+
+NON_TRANSIENT_CASES = [
+ NonTransientCase(
+ name="non-transient-status",
+ errors=[{"message": "Unknown field", "extensions": {"code": "GRAPHQL_VALIDATION", "http_status": 400}}],
+ ),
+ NonTransientCase(name="mixed", errors=[_transient_error(503), {"message": "Unknown field"}]),
+ NonTransientCase(name="unclassified", errors=[{"message": "legacy error without extensions"}]),
+]
+
+
+@pytest.mark.parametrize("case", [pytest.param(tc, id=tc.name) for tc in NON_TRANSIENT_CASES])
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_non_transient_graphql_errors_raise_immediately(
+ client_type: str, case: NonTransientCase, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ requester = ScriptedRequester([_graphql_errors(*case.errors), _ok()])
+ client = _build_client(client_type, requester, retry_on_failure=True, max_retry_duration=0)
+
+ with pytest.raises(GraphQLError, match="An error occurred while executing the GraphQL Query") as exc:
+ await _execute_graphql(client)
+
+ assert exc.value.errors == case.errors
+ assert requester.call_count == 1
+ assert recorded_sleeps == []
+
+
+# --- Budget exhaustion and unlimited mode -----------------------------------------------------------
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_budget_exhaustion_reraises_the_original_exception(
+ client_type: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ _no_jitter(monkeypatch)
+ clock = FakeClock()
+ recorded_sleeps = _patch_sleep(monkeypatch, clock=clock)
+ requester = ScriptedRequester([_timeout() for _ in range(10)])
+ client = _build_client(client_type, requester, clock=clock, retry_on_failure=True, max_retry_duration=12)
+
+ with pytest.raises(ServerNotResponsiveError, match="Unable to read from"):
+ await _execute_graphql(client)
+
+ # 5s, then 10s clamped to the 7s left in the budget, then the budget is spent.
+ assert recorded_sleeps == [5, 7]
+ assert requester.call_count == 3
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_budget_exhaustion_on_transient_status_returns_last_response(
+ client_type: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ _no_jitter(monkeypatch)
+ clock = FakeClock()
+ _patch_sleep(monkeypatch, clock=clock)
+ requester = ScriptedRequester([_response(503) for _ in range(10)])
+ client = _build_client(client_type, requester, clock=clock, retry_on_failure=True, max_retry_duration=12)
+
+ response = await _request(client)
+
+ assert response.status_code == 503
+ assert requester.call_count == 3
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_budget_exhaustion_on_5xx_graphql_envelope_raises_graphql_error(
+ client_type: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ """Once the shared budget is spent, ``execute_graphql`` does not start a second one for the envelope."""
+ _no_jitter(monkeypatch)
+ clock = FakeClock()
+ recorded_sleeps = _patch_sleep(monkeypatch, clock=clock)
+ envelope = {"message": "Service unavailable", "extensions": {"code": 503}}
+ requester = ScriptedRequester([_graphql_errors(envelope, status_code=503) for _ in range(10)])
+ client = _build_client(client_type, requester, clock=clock, retry_on_failure=True, max_retry_duration=12)
+
+ with pytest.raises(GraphQLError, match="An error occurred while executing the GraphQL Query") as exc:
+ await _execute_graphql(client)
+
+ assert exc.value.errors == [envelope]
+ assert requester.call_count == 3
+ assert recorded_sleeps == [5, 7]
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_unlimited_budget_keeps_retrying_and_escalates_logging(
+ client_type: str, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+ _no_jitter(monkeypatch)
+ caplog.set_level(logging.WARNING, logger=LOGGER_NAME)
+ clock = FakeClock()
+ recorded_sleeps = _patch_sleep(monkeypatch, clock=clock)
+ failures = 40
+ requester = ScriptedRequester([_connection_error() for _ in range(failures)] + [_ok()])
+ client = _build_client(client_type, requester, clock=clock, retry_on_failure=True, max_retry_duration=0)
+
+ data = await _execute_graphql(client)
+
+ assert data == {"result": "ok"}
+ assert requester.call_count == failures + 1
+ assert len(recorded_sleeps) == failures
+ assert recorded_sleeps[:5] == [5, 10, 20, 40, 60]
+ assert max(recorded_sleeps) == 60
+ assert clock.now - 1000.0 > ESCALATE_AFTER_SECONDS
+ levels = {record.levelno for record in caplog.records}
+ assert levels == {logging.WARNING, logging.ERROR}
+ assert "no time limit" in caplog.records[-1].message
+
+
+# --- Runtime toggling and configuration -------------------------------------------------------------
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_retry_settings_can_be_toggled_at_runtime(client_type: str, monkeypatch: pytest.MonkeyPatch) -> None:
+ _no_jitter(monkeypatch)
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ requester = ScriptedRequester([_connection_error(), _ok()])
+ client = _build_client(client_type, requester)
+ assert client.retry_delay == 5
+
+ client.retry_on_failure = True
+ client.retry_delay = 1
+
+ data = await _execute_graphql(client)
+
+ assert data == {"result": "ok"}
+ assert recorded_sleeps == [1]
+
+
+def test_config_wires_the_retry_handler() -> None:
+ config = Config(
+ address="http://mock",
+ retry_on_failure=True,
+ retry_delay=2,
+ retry_max_delay=30,
+ max_retry_duration=0,
+ retry_status_codes=[500, 503],
+ )
+ handler = InfrahubClient(config=config)._retry_handler
+
+ assert handler.enabled is True
+ assert handler.base_delay == 2
+ assert handler.max_delay == 30
+ assert handler.unlimited
+ assert handler.status_codes == frozenset({500, 503})
+
+
+def test_config_defaults_keep_retries_opt_in() -> None:
+ config = Config(address="http://mock")
+
+ assert config.retry_on_failure is False
+ assert config.retry_delay == 5
+ assert config.retry_max_delay == 60
+ assert config.max_retry_duration == 300
+ assert config.retry_status_codes == [500, 502, 503, 504]
+
+
+@pytest.mark.parametrize(
+ ("field", "value", "message"),
+ [
+ pytest.param("retry_delay", 0, "greater than 0", id="retry_delay-zero"),
+ pytest.param("retry_max_delay", 0, "greater than 0", id="retry_max_delay-zero"),
+ pytest.param("max_retry_duration", -1, "greater than or equal to 0", id="max_retry_duration-negative"),
+ ],
+)
+def test_config_rejects_retry_settings_that_would_spin_or_run_backwards(field: str, value: int, message: str) -> None:
+ """A zero base delay or ceiling would make every retry sleep 0s; only the duration may be 0 (unlimited)."""
+ overrides: dict[str, Any] = {field: value}
+ with pytest.raises(ValidationError, match=message):
+ Config(address="http://mock", **overrides)
+
+
+def test_config_reads_retry_settings_from_environment(monkeypatch: pytest.MonkeyPatch) -> None:
+ monkeypatch.setenv("INFRAHUB_RETRY_ON_FAILURE", "true")
+ monkeypatch.setenv("INFRAHUB_MAX_RETRY_DURATION", "0")
+ monkeypatch.setenv("INFRAHUB_RETRY_STATUS_CODES", "[503, 504]")
+
+ config = Config(address="http://mock")
+
+ assert config.retry_on_failure is True
+ assert config.max_retry_duration == 0
+ assert config.retry_status_codes == [503, 504]
+
+
+# --- Every request path: multipart uploads and streamed downloads ------------------------------------
+
+ALL_REQUEST_PATHS = ["regular", "multipart", "streaming"]
+MULTIPART_FILE_CONTENT = b"multipart file body that must survive a transient retry\n" * 8
+UPLOAD_MUTATION = "mutation ($file: Upload!) { InfrahubObjectUpload(data: { file: $file }) { ok } }"
+
+
+class NonSeekableStream(io.BytesIO):
+ """A read-only stream that cannot be rewound, like a pipe or an HTTP body."""
+
+ def seekable(self) -> bool:
+ return False
+
+ def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
+ raise io.UnsupportedOperation("seek")
+
+ def tell(self) -> int:
+ raise io.UnsupportedOperation("tell")
+
+
+def _multipart_files() -> dict[str, Any]:
+ return {"file": ("upload.bin", io.BytesIO(MULTIPART_FILE_CONTENT), "application/octet-stream")}
+
+
+async def _drive_path(client: InfrahubClient | InfrahubClientSync, path: str) -> int:
+ """Drive one request path and return the final status; a streamed body is read inside its context."""
+ if path == "regular":
+ return (await _request(client)).status_code
+ if path == "multipart":
+ if isinstance(client, InfrahubClient):
+ response = await client._request_multipart(
+ url=GRAPHQL_URL, headers={}, timeout=10, files=_multipart_files()
+ )
+ else:
+ response = client._request_multipart(url=GRAPHQL_URL, headers={}, timeout=10, files=_multipart_files())
+ return response.status_code
+ if isinstance(client, InfrahubClient):
+ async with client._get_streaming(url=GRAPHQL_URL) as response:
+ await response.aread()
+ return response.status_code
+ with client._get_streaming(url=GRAPHQL_URL) as response:
+ response.read()
+ return response.status_code
+
+
+async def _upload(client: InfrahubClient | InfrahubClientSync, file_content: BinaryIO) -> dict:
+ if isinstance(client, InfrahubClient):
+ return await client._execute_graphql_with_file(
+ query=UPLOAD_MUTATION, variables={}, file_content=file_content, file_name="upload.bin"
+ )
+ return client._execute_graphql_with_file(
+ query=UPLOAD_MUTATION, variables={}, file_content=file_content, file_name="upload.bin"
+ )
+
+
+@pytest.mark.parametrize("path", ALL_REQUEST_PATHS)
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_all_request_paths_retry_transient_status_then_succeed(
+ client_type: str, path: str, httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ httpx_mock.add_response(status_code=503)
+ httpx_mock.add_response(status_code=200, json={"data": {"result": "success"}})
+ client = _build_client_over_httpx(client_type, retry_on_failure=True)
+
+ status = await _drive_path(client, path)
+
+ assert status == 200
+ assert len(httpx_mock.get_requests()) == 2
+ assert len(recorded_sleeps) == 1
+
+
+@pytest.mark.parametrize("path", ALL_REQUEST_PATHS)
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_all_request_paths_retry_lost_connection_then_succeed(
+ client_type: str, path: str, httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ httpx_mock.add_exception(httpx.ConnectError("connection refused"))
+ httpx_mock.add_response(status_code=200, json={"data": {"result": "success"}})
+ client = _build_client_over_httpx(client_type, retry_on_failure=True)
+
+ status = await _drive_path(client, path)
+
+ assert status == 200
+ assert len(httpx_mock.get_requests()) == 2
+ assert len(recorded_sleeps) == 1
+
+
+@pytest.mark.parametrize("path", ALL_REQUEST_PATHS)
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_all_request_paths_pass_transient_status_through_when_disabled(
+ client_type: str, path: str, httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ httpx_mock.add_response(status_code=503, content=b"upstream unavailable")
+ client = _build_client_over_httpx(client_type)
+
+ status = await _drive_path(client, path)
+
+ assert status == 503
+ assert len(httpx_mock.get_requests()) == 1
+ assert recorded_sleeps == []
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_multipart_mutation_retries_transient_graphql_errors_on_the_shared_budget(
+ client_type: str, httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ """A transient status and then a transient GraphQL envelope are retried on one attempt counter, re-sending the file."""
+ _no_jitter(monkeypatch)
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ httpx_mock.add_response(status_code=503)
+ httpx_mock.add_response(status_code=200, json={"data": None, "errors": [_transient_error(503)]})
+ httpx_mock.add_response(status_code=200, json={"data": {"InfrahubObjectUpload": {"ok": True}}})
+ client = _build_client_over_httpx(client_type, retry_on_failure=True)
+
+ data = await _upload(client, io.BytesIO(MULTIPART_FILE_CONTENT))
+
+ assert data == {"InfrahubObjectUpload": {"ok": True}}
+ requests = httpx_mock.get_requests()
+ assert len(requests) == 3
+ assert recorded_sleeps == [5, 10]
+ assert all(MULTIPART_FILE_CONTENT in request.content for request in requests)
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_multipart_mutation_raises_non_transient_graphql_errors_immediately(
+ client_type: str, httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ errors = [{"message": "Unknown field", "extensions": {"code": "GRAPHQL_VALIDATION", "http_status": 400}}]
+ httpx_mock.add_response(status_code=200, json={"data": None, "errors": errors})
+ client = _build_client_over_httpx(client_type, retry_on_failure=True, max_retry_duration=0)
+
+ with pytest.raises(GraphQLError, match="An error occurred while executing the GraphQL Query") as exc:
+ await _upload(client, io.BytesIO(MULTIPART_FILE_CONTENT))
+
+ assert exc.value.errors == errors
+ assert len(httpx_mock.get_requests()) == 1
+ assert recorded_sleeps == []
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_multipart_retry_re_sends_a_non_seekable_upload_in_full(
+ client_type: str, httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ """A stream that cannot be rewound is copied once, so the transport and envelope retries still carry the body."""
+ _patch_sleep(monkeypatch)
+ httpx_mock.add_response(status_code=503)
+ httpx_mock.add_response(status_code=200, json={"data": None, "errors": [_transient_error(503)]})
+ httpx_mock.add_response(status_code=200, json={"data": {"InfrahubObjectUpload": {"ok": True}}})
+ client = _build_client_over_httpx(client_type, retry_on_failure=True)
+
+ data = await _upload(client, NonSeekableStream(MULTIPART_FILE_CONTENT))
+
+ assert data == {"InfrahubObjectUpload": {"ok": True}}
+ requests = httpx_mock.get_requests()
+ assert len(requests) == 3
+ assert all(MULTIPART_FILE_CONTENT in request.content for request in requests)
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_streaming_hands_over_an_open_transient_response_once_the_budget_is_spent(
+ client_type: str, httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ """A transient status the handler will not retry is not pre-read: the caller gets the stream untouched."""
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ httpx_mock.add_response(status_code=503, content=b"upstream unavailable")
+ client = _build_client_over_httpx(client_type, retry_on_failure=True, max_retry_duration=1)
+ client._retry_handler.clock = FakeClock(step=1.0) # every look at the clock moves it past the budget
+
+ if isinstance(client, InfrahubClient):
+ async with client._get_streaming(url=GRAPHQL_URL) as response:
+ assert response.status_code == 503
+ assert not response.is_stream_consumed
+ assert await response.aread() == b"upstream unavailable"
+ else:
+ with client._get_streaming(url=GRAPHQL_URL) as sync_response:
+ assert sync_response.status_code == 503
+ assert not sync_response.is_stream_consumed
+ assert sync_response.read() == b"upstream unavailable"
+
+ assert len(httpx_mock.get_requests()) == 1
+ assert recorded_sleeps == []
+
+
+DOWNLOAD_URL = "http://mock/api/storage/files/file-1?branch=main"
+DOWNLOAD_BODY = b"streamed body that must be complete on disk\n" * 64
+
+
+def _truncated_body() -> IteratorStream:
+ """A body whose connection drops after the first chunk, the way httpx reports a truncated transfer."""
+
+ def chunks() -> Iterator[bytes]:
+ yield DOWNLOAD_BODY[:100]
+ raise httpx.RemoteProtocolError("peer closed connection without sending complete message body")
+
+ return IteratorStream(chunks())
+
+
+async def _download_to(client: InfrahubClient | InfrahubClientSync, dest: Path) -> int:
+ if isinstance(client, InfrahubClient):
+ return await FileHandler(client=client).download(node_id="file-1", branch="main", dest=dest)
+ return FileHandlerSync(client=client).download(node_id="file-1", branch="main", dest=dest)
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_download_to_file_restarts_when_the_connection_drops_mid_body(
+ client_type: str, httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+ """Losing the connection while the body streams is transient too: the download starts over on the shared budget."""
+ _no_jitter(monkeypatch)
+ recorded_sleeps = _patch_sleep(monkeypatch)
+ httpx_mock.add_response(url=DOWNLOAD_URL, stream=_truncated_body())
+ httpx_mock.add_response(url=DOWNLOAD_URL, content=DOWNLOAD_BODY)
+ client = _build_client_over_httpx(client_type, retry_on_failure=True)
+ dest = tmp_path / "download.bin"
+
+ assert await _download_to(client, dest) == len(DOWNLOAD_BODY)
+
+ assert dest.read_bytes() == DOWNLOAD_BODY
+ assert len(httpx_mock.get_requests()) == 2
+ assert recorded_sleeps == [5]
+
+
+@pytest.mark.parametrize("client_type", CLIENT_TYPES)
+async def test_download_interrupted_mid_body_reports_the_lost_connection_and_leaves_no_partial_file(
+ client_type: str, httpx_mock: HTTPXMock, tmp_path: Path
+) -> None:
+ """Without retries the caller learns the transfer broke off, not that the server was unreachable."""
+ httpx_mock.add_response(url=DOWNLOAD_URL, stream=_truncated_body())
+ client = _build_client_over_httpx(client_type, retry_on_failure=False)
+ dest = tmp_path / "download.bin"
+
+ with pytest.raises(ServerNotResponsiveError, match="was lost while reading the response"):
+ await _download_to(client, dest)
+
+ assert not dest.exists()
+ assert len(httpx_mock.get_requests()) == 1
+
+
+async def test_async_upload_copies_a_non_seekable_stream_off_the_event_loop(
+ httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ """Draining a stream that cannot be rewound is blocking I/O, so the async client hands it to a worker thread."""
+ offloaded: list[str] = []
+ real_to_thread = asyncio.to_thread
+
+ async def recording_to_thread(func: Callable[..., object], /, *args: object, **kwargs: object) -> object:
+ offloaded.append(getattr(func, "__name__", repr(func)))
+ return await real_to_thread(func, *args, **kwargs)
+
+ monkeypatch.setattr(asyncio, "to_thread", recording_to_thread)
+ httpx_mock.add_response(status_code=200, json={"data": {"InfrahubObjectUpload": {"ok": True}}})
+ client = _build_client_over_httpx("standard")
+
+ data = await _upload(client, NonSeekableStream(MULTIPART_FILE_CONTENT))
+
+ assert data == {"InfrahubObjectUpload": {"ok": True}}
+ assert len(offloaded) == 1, f"expected exactly the stream copy to be offloaded, got {offloaded}"
+ assert MULTIPART_FILE_CONTENT in httpx_mock.get_requests()[0].content
+
+
+async def test_cancelled_async_upload_returns_at_once_and_leaves_the_buffer_to_the_worker(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """Cancelling mid-copy surfaces immediately and leaves the temporary file to the worker.
+
+ The worker cannot be interrupted, so it closes the file itself once its pending read returns;
+ the file is never closed under it.
+ """
+ started = threading.Event()
+ release = threading.Event()
+ reads: list[bytes] = []
+ buffers: list[IO[bytes]] = []
+ real_temporary_file = tempfile.TemporaryFile
+
+ def recording_temporary_file() -> IO[bytes]:
+ buffer = real_temporary_file()
+ buffers.append(buffer)
+ return buffer
+
+ monkeypatch.setattr(tempfile, "TemporaryFile", recording_temporary_file)
+
+ class BlockingStream(NonSeekableStream):
+ def read(self, size: int | None = -1) -> bytes:
+ started.set()
+ assert release.wait(timeout=5), "the test never released the copy"
+ chunk = super().read(size)
+ reads.append(chunk)
+ return chunk
+
+ async def upload() -> None:
+ async with client_module._aseekable_upload(BlockingStream(MULTIPART_FILE_CONTENT)):
+ pytest.fail("the upload body should never run once cancelled during the copy")
+
+ task = asyncio.create_task(upload())
+ assert await asyncio.to_thread(started.wait, 5), "the copy never started" # the worker is now blocked mid-copy
+
+ task.cancel()
+ done, _ = await asyncio.wait({task}, timeout=5)
+ assert task in done, "the cancellation did not surface while the worker thread was still blocked"
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ [buffer] = buffers
+ assert not buffer.closed, "the temporary file was closed under the worker still copying into it"
+
+ release.set()
+ for _ in range(100):
+ if buffer.closed:
+ break
+ await asyncio.sleep(0.05)
+ assert buffer.closed, "the worker did not close the temporary file once its copy ended"
+ assert reads[-1] == b"", "the copy should have drained the stream to EOF before closing the buffer"