Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 16 additions & 11 deletions pyiceberg/catalog/rest/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -499,6 +499,17 @@ def send(
return super().send(request, stream=stream, timeout=timeout, verify=verify, cert=cert, proxies=proxies)


def _connection_timeout(properties: Properties) -> float | None:
"""Return the request timeout in seconds from the `rest.client.*` timeout properties, or None when unset."""
connection_timeout_ms = property_as_int(properties, REST_CLIENT_CONNECTION_TIMEOUT_MS)
socket_timeout_ms = property_as_int(properties, REST_CLIENT_SOCKET_TIMEOUT_MS)
if connection_timeout_ms is None and socket_timeout_ms is None:
return None
# requests uses a single timeout and cannot split connect vs socket, so follow the Java client
# and sum the two (milliseconds), flooring to whole seconds.
return ((connection_timeout_ms or 0) + (socket_timeout_ms or 0)) // 1000


def _create_connection_adapter(properties: Properties) -> _RetryTimeoutHTTPAdapter | None:
"""Build a connection adapter from the optional `rest.client.*` properties.

Expand All @@ -524,14 +535,8 @@ def _create_connection_adapter(properties: Properties) -> _RetryTimeoutHTTPAdapt
if all(value is None for value in (connection_timeout_ms, socket_timeout_ms, retries, backoff_factor)):
return None

# requests uses a single timeout and cannot split connect vs socket, so follow the Java client
# and sum the two (milliseconds), flooring to whole seconds.
timeout: float | None = None
if connection_timeout_ms is not None or socket_timeout_ms is not None:
timeout = ((connection_timeout_ms or 0) + (socket_timeout_ms or 0)) // 1000

return _RetryTimeoutHTTPAdapter(
timeout=timeout,
timeout=_connection_timeout(properties),
max_retries=Retry(
# `retries` and `backoff_factor` fall back to a no-op Retry when unset, so a user can
# configure only one without having to specify the rest of the policy.
Expand Down Expand Up @@ -575,7 +580,7 @@ def _create_session(self) -> Session:
session = Session()

# Mount the retry/timeout adapter when `connection.*` properties are set.
# SigV4's adapter mounted below at `self.uri` is a longer prefix and still wins for that host.
# SigV4's adapter mounted below at `self.uri` is a longer prefix and wins for that host; it applies the same timeout.
if (connection_adapter := _create_connection_adapter(self.properties)) is not None:
session.mount("http://", connection_adapter)
session.mount("https://", connection_adapter)
Expand Down Expand Up @@ -1098,11 +1103,11 @@ def _init_sigv4(self, session: Session) -> None:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest

class SigV4Adapter(HTTPAdapter):
class SigV4Adapter(_RetryTimeoutHTTPAdapter):
def __init__(self, **properties: str):
self._properties = properties
max_retries = property_as_int(self._properties, SIGV4_MAX_RETRIES, SIGV4_MAX_RETRIES_DEFAULT)
super().__init__(max_retries=max_retries)
max_retries: int = property_as_int(self._properties, SIGV4_MAX_RETRIES, SIGV4_MAX_RETRIES_DEFAULT) # type: ignore # The default is set with non-None value.
super().__init__(timeout=_connection_timeout(self._properties), max_retries=max_retries)
self._boto_session = boto3.Session(
profile_name=get_first_property_value(self._properties, AWS_PROFILE_NAME),
region_name=get_first_property_value(self._properties, AWS_REGION),
Expand Down
22 changes: 22 additions & 0 deletions tests/catalog/test_rest.py
Original file line number Diff line number Diff line change
Expand Up @@ -695,6 +695,28 @@ def test_sigv4_adapter_override_retry_config(rest_mock: Mocker) -> None:
assert adapter.max_retries.total == 3


def test_sigv4_adapter_applies_connection_timeout(rest_mock: Mocker) -> None:
catalog = RestCatalog(
"rest",
**{
"uri": TEST_URI,
"token": TEST_TOKEN,
"rest.sigv4-enabled": "true",
"rest.signing-region": "us-west-2",
"client.access-key-id": "id",
"client.secret-access-key": "secret",
REST_CLIENT_CONNECTION_TIMEOUT_MS: "5000",
REST_CLIENT_SOCKET_TIMEOUT_MS: "60000",
},
)

# The SigV4 adapter is mounted on the catalog URI, so it handles catalog requests instead of the `https://` adapter.
adapter = catalog._session.get_adapter(catalog.uri)
assert isinstance(adapter, _RetryTimeoutHTTPAdapter)
assert adapter._timeout == 65 # (5000 + 60000) ms floored to whole seconds
assert adapter.max_retries.total == SIGV4_MAX_RETRIES_DEFAULT


def test_sigv4_uses_client_profile_name(rest_mock: Mocker) -> None:
with mock.patch("boto3.Session") as mock_session:
RestCatalog(
Expand Down
Loading