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
2 changes: 2 additions & 0 deletions mkdocs/docs/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -478,6 +478,8 @@ Legacy OAuth2 Properties will be removed in PyIceberg 1.0 in place of pluggable
| rest.signing-region | us-east-1 | The region to use when SigV4 signing a request |
| rest.signing-name | execute-api | The service signing name to use when SigV4 signing a request |

The `rest.client.connection-timeout-ms` and `rest.client.socket-timeout-ms` timeouts also apply to SigV4-signed requests. SigV4 retries are controlled by `rest.sigv4.max-retries`, not `rest.client.max-retries`.

##### Pluggable Authentication via AuthManager

The RESTCatalog supports pluggable authentication via the `auth` configuration block. This allows you to specify which how the access token will be fetched and managed for use with the HTTP requests to the RESTCatalog server. The authentication method is selected by setting the `auth.type` property, and additional configuration can be provided as needed for each method.
Expand Down
40 changes: 26 additions & 14 deletions pyiceberg/catalog/rest/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -499,11 +499,10 @@ def send(
return super().send(request, stream=stream, timeout=timeout, verify=verify, cert=cert, proxies=proxies)


def _create_connection_adapter(properties: Properties) -> _RetryTimeoutHTTPAdapter | None:
"""Build a connection adapter from the optional `rest.client.*` properties.
def _connection_timeout(properties: Properties) -> float | None:
"""Compute the default request timeout in seconds from the optional `rest.client.*` timeout properties.

Returns None when no connection properties are supplied, leaving the default
Session behavior unchanged. Raises ValueError on invalid input.
Raises ValueError on invalid input.
"""
connection_timeout_ms = property_as_int(properties, REST_CLIENT_CONNECTION_TIMEOUT_MS)
if connection_timeout_ms is not None and connection_timeout_ms <= 0:
Expand All @@ -513,6 +512,19 @@ def _create_connection_adapter(properties: Properties) -> _RetryTimeoutHTTPAdapt
if socket_timeout_ms is not None and socket_timeout_ms <= 0:
raise ValueError(f"`{REST_CLIENT_SOCKET_TIMEOUT_MS}` must be a positive number, got: {socket_timeout_ms}")

# 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.
if connection_timeout_ms is None and socket_timeout_ms is None:
return None
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.

Returns None when no connection properties are supplied, leaving the default
Session behavior unchanged. Raises ValueError on invalid input.
"""
retries = property_as_int(properties, REST_CLIENT_MAX_RETRIES)
if retries is not None and retries < 0:
raise ValueError(f"`{REST_CLIENT_MAX_RETRIES}` must be non-negative, got: {retries}")
Expand All @@ -521,14 +533,10 @@ def _create_connection_adapter(properties: Properties) -> _RetryTimeoutHTTPAdapt
if backoff_factor is not None and backoff_factor < 0:
raise ValueError(f"`{REST_CLIENT_RETRY_BACKOFF_FACTOR}` must be non-negative, got: {backoff_factor}")

if all(value is None for value in (connection_timeout_ms, socket_timeout_ms, retries, backoff_factor)):
return None
timeout = _connection_timeout(properties)

# 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
if all(value is None for value in (timeout, retries, backoff_factor)):
return None

return _RetryTimeoutHTTPAdapter(
timeout=timeout,
Expand Down Expand Up @@ -575,7 +583,8 @@ 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 still wins for that host,
# so it applies the same timeout itself.
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 +1107,14 @@ 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)
super().__init__(
timeout=_connection_timeout(self._properties),
max_retries=SIGV4_MAX_RETRIES_DEFAULT if max_retries is None else 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
21 changes: 21 additions & 0 deletions tests/catalog/test_rest.py
Original file line number Diff line number Diff line change
Expand Up @@ -695,6 +695,27 @@ def test_sigv4_adapter_override_retry_config(rest_mock: Mocker) -> None:
assert adapter.max_retries.total == 3


def test_sigv4_adapter_applies_client_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": "1000",
"rest.client.socket-timeout-ms": "2000",
},
)

adapter = catalog._session.adapters[catalog.uri]
assert isinstance(adapter, _RetryTimeoutHTTPAdapter)
assert adapter._timeout == 3
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