-
Notifications
You must be signed in to change notification settings - Fork 576
REST: Add retry and timeout configuration for REST catalog #3418
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
9d85042
6fb87ff
afb1f51
47a5382
9a799c6
36053c7
afa6685
6c52148
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -19,17 +19,21 @@ | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| import logging | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| import time | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from collections import deque | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from collections.abc import Callable, Mapping | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from enum import Enum | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from typing import ( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| TYPE_CHECKING, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| Any, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| TypeVar, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from urllib.parse import quote, unquote | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from pydantic import ConfigDict, Field, TypeAdapter, field_validator | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from requests import HTTPError, Session | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from requests import HTTPError, PreparedRequest, Response, Session | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from requests.adapters import DEFAULT_RETRIES, HTTPAdapter | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from tenacity import RetryCallState, retry, retry_if_exception_type, stop_after_attempt | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from typing_extensions import override | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from urllib3.util.retry import Retry | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from pyiceberg import __version__ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from pyiceberg.catalog import BOTOCORE_SESSION, TOKEN, URI, WAREHOUSE_LOCATION, Catalog, PropertiesUpdateSummary | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
@@ -274,6 +278,13 @@ class ScanPlanningMode(Enum): | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| SIGV4_SERVICE = "rest.signing-name" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| SIGV4_MAX_RETRIES = "rest.sigv4.max-retries" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| SIGV4_MAX_RETRIES_DEFAULT = 10 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| REST_CLIENT_REQUEST_TIMEOUT = "rest.client.request-timeout" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| REST_CLIENT_MAX_RETRIES = "rest.client.max-retries" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| REST_CLIENT_RETRY_BACKOFF_FACTOR = "rest.client.retry-backoff-factor" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| # Hard-coded internally so users cannot misconfigure the retry policy | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| # (e.g. setting raise_on_status=False would swallow 4xx errors silently). | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| _CONNECTION_RETRY_STATUS_FORCELIST = (429, 500, 502, 503, 504) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| _CONNECTION_RETRY_ALLOWED_METHODS = frozenset({"GET", "HEAD", "OPTIONS", "PUT", "DELETE"}) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| EMPTY_BODY_SHA256: str = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| OAUTH2_SERVER_URI = "oauth2-server-uri" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| SNAPSHOT_LOADING_MODE = "snapshot-loading-mode" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
@@ -442,6 +453,111 @@ class ListViewsResponse(IcebergBaseModel): | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| _PLANNING_RESPONSE_ADAPTER = TypeAdapter(PlanningResponse) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| _T = TypeVar("_T", int, float) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def _parse_connection_property( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| properties: Properties, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| property_name: str, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| converter: Callable[[Any], _T], | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| type_description: str, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| is_invalid: Callable[[_T], bool], | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| range_description: str, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) -> _T | None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| raw_value = properties.get(property_name) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| if raw_value is None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| return None | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| try: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| value = converter(raw_value) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| except (TypeError, ValueError) as e: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| raise ValueError(f"`{property_name}` must be {type_description}, got: {raw_value!r}") from e | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| if is_invalid(value): | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| raise ValueError(f"`{property_name}` must be {range_description}, got: {value}") | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| return value | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| class _RetryTimeoutHTTPAdapter(HTTPAdapter): | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| """HTTPAdapter that applies a default per-request timeout. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| requests does not provide a way to set a default timeout on a Session; | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| without this adapter, every call would have to thread `timeout=` through. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| The adapter applies `self._timeout` whenever a per-call timeout is not set. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| """ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def __init__(self, timeout: float | None = None, max_retries: Retry | int = DEFAULT_RETRIES) -> None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| self._timeout = timeout | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| super().__init__(max_retries=max_retries) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def send( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| self, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| request: PreparedRequest, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| stream: bool = False, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| timeout: None | float | tuple[float, float] | tuple[float, None] = None, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| verify: bool | str = True, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| cert: None | bytes | str | tuple[bytes | str, bytes | str] = None, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| proxies: Mapping[str, str] | None = None, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) -> Response: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| if timeout is None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| timeout = self._timeout | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 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. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| Returns None when no connection properties are supplied, leaving the default | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| Session behavior unchanged. Raises ValueError on invalid input. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| """ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| if not any( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| property_name in properties | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| for property_name in (REST_CLIENT_REQUEST_TIMEOUT, REST_CLIENT_MAX_RETRIES, REST_CLIENT_RETRY_BACKOFF_FACTOR) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ): | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| return None | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| timeout = _parse_connection_property( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| properties, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| REST_CLIENT_REQUEST_TIMEOUT, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| float, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| "a number", | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| lambda value: value <= 0, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| "a positive number", | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| retries = _parse_connection_property( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| properties, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| REST_CLIENT_MAX_RETRIES, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| int, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| "an integer", | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| lambda value: value < 0, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| "non-negative", | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| backoff_factor = _parse_connection_property( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| properties, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| REST_CLIENT_RETRY_BACKOFF_FACTOR, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| float, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| "a number", | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| lambda value: value < 0, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| "non-negative", | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
Comment on lines
+517
to
+541
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We can compact this into by reusing the methods in
Suggested change
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| return _RetryTimeoutHTTPAdapter( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think there's a pretty significant issue here. The
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Good catch — fixed in 9a799c6. Added
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Great catch, thanks for the extra review! |
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| timeout=timeout, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 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. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| total=retries if retries is not None else DEFAULT_RETRIES, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| backoff_factor=backoff_factor if backoff_factor is not None else 0.0, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| status_forcelist=list(_CONNECTION_RETRY_STATUS_FORCELIST), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| allowed_methods=_CONNECTION_RETRY_ALLOWED_METHODS, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| # Return the final response on retry exhaustion (instead of raising MaxRetryError) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| # so `_handle_non_200_response` can map the 5xx status to a typed exception | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| # (ServiceUnavailableError, etc.). 4xx codes are not in status_forcelist and are | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| # never retried, so they reach the same mapping unchanged. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| raise_on_status=False, | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ), | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| class RestCatalog(Catalog): | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| uri: str | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| _session: Session | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
@@ -468,6 +584,12 @@ def _create_session(self) -> Session: | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| """Create a request session with provided catalog configuration.""" | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| 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. | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| if (connection_adapter := _create_connection_adapter(self.properties)) is not None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| session.mount("http://", connection_adapter) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| session.mount("https://", connection_adapter) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| # Set HTTP headers | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| self._config_headers(session) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
@@ -980,8 +1102,6 @@ def _init_sigv4(self, session: Session) -> None: | |||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| import boto3 | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from botocore.auth import SigV4Auth | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from botocore.awsrequest import AWSRequest | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from requests import PreparedRequest | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| from requests.adapters import HTTPAdapter | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| class SigV4Adapter(HTTPAdapter): | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| def __init__(self, **properties: str): | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -19,7 +19,8 @@ | |
|
|
||
| import base64 | ||
| import os | ||
| from collections.abc import Callable | ||
| from collections.abc import Callable, Iterator | ||
| from contextlib import contextmanager | ||
| from typing import Any, cast | ||
| from unittest import mock | ||
|
|
||
|
|
@@ -36,6 +37,9 @@ | |
| EMPTY_BODY_SHA256, | ||
| OAUTH2_SERVER_URI, | ||
| PAGE_SIZE, | ||
| REST_CLIENT_MAX_RETRIES, | ||
| REST_CLIENT_REQUEST_TIMEOUT, | ||
| REST_CLIENT_RETRY_BACKOFF_FACTOR, | ||
| SIGV4_MAX_RETRIES, | ||
| SIGV4_MAX_RETRIES_DEFAULT, | ||
| SNAPSHOT_LOADING_MODE, | ||
|
|
@@ -44,6 +48,7 @@ | |
| HttpMethod, | ||
| RestCatalog, | ||
| ScanPlanningMode, | ||
| _RetryTimeoutHTTPAdapter, | ||
| ) | ||
| from pyiceberg.exceptions import ( | ||
| AuthorizationExpiredError, | ||
|
|
@@ -55,6 +60,7 @@ | |
| NoSuchViewError, | ||
| OAuthError, | ||
| ServerError, | ||
| ServiceUnavailableError, | ||
| TableAlreadyExistsError, | ||
| ViewAlreadyExistsError, | ||
| ) | ||
|
|
@@ -2284,6 +2290,158 @@ def test_request_session_with_ssl_client_cert() -> None: | |
| assert "Could not find the TLS certificate file, invalid path: path_to_client_cert" in str(e.value) | ||
|
|
||
|
|
||
| def test_session_without_connection_config_uses_default_adapter(rest_mock: Mocker) -> None: | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can we get a test where we set the retry logic and then see the retries occur? We should be able to simulate this with mock HTTP calls and then see that X number of HTTP calls were made afterwards.
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Added |
||
| catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN) | ||
| for adapter in catalog._session.adapters.values(): | ||
| assert not isinstance(adapter, _RetryTimeoutHTTPAdapter) | ||
|
|
||
|
|
||
| def test_session_with_connection_timeout_and_retries(rest_mock: Mocker) -> None: | ||
| catalog_properties = { | ||
| "uri": TEST_URI, | ||
| "token": TEST_TOKEN, | ||
| REST_CLIENT_REQUEST_TIMEOUT: 60, | ||
| REST_CLIENT_MAX_RETRIES: 5, | ||
| REST_CLIENT_RETRY_BACKOFF_FACTOR: 1.0, | ||
| } | ||
| catalog = RestCatalog("rest", **catalog_properties) # type: ignore | ||
|
|
||
| https_adapter = catalog._session.adapters["https://"] | ||
| http_adapter = catalog._session.adapters["http://"] | ||
| assert isinstance(https_adapter, _RetryTimeoutHTTPAdapter) | ||
| assert https_adapter is http_adapter | ||
| assert https_adapter._timeout == 60.0 | ||
| assert https_adapter.max_retries.total == 5 | ||
| assert https_adapter.max_retries.backoff_factor == 1.0 | ||
| # Internal retry policy: transient codes and idempotent methods only. | ||
| assert https_adapter.max_retries.status_forcelist == [429, 500, 502, 503, 504] | ||
| allowed_methods = https_adapter.max_retries.allowed_methods or frozenset() | ||
| assert set(allowed_methods) == {"GET", "HEAD", "OPTIONS", "PUT", "DELETE"} | ||
|
|
||
|
|
||
| def test_session_with_connection_timeout_only(rest_mock: Mocker) -> None: | ||
| catalog_properties = { | ||
| "uri": TEST_URI, | ||
| "token": TEST_TOKEN, | ||
| REST_CLIENT_REQUEST_TIMEOUT: "30", | ||
| } | ||
| catalog = RestCatalog("rest", **catalog_properties) | ||
| adapter = catalog._session.adapters["https://"] | ||
| assert isinstance(adapter, _RetryTimeoutHTTPAdapter) | ||
| assert adapter._timeout == 30.0 | ||
| # Default retry policy (total=0) is a no-op when only timeout is configured. | ||
| assert adapter.max_retries.total == 0 | ||
|
|
||
|
|
||
| @contextmanager | ||
| def _local_rest_server_503_then_200(num_failures: int) -> Iterator[dict[str, Any]]: | ||
| """Stand up a loopback HTTP server that returns `num_failures` 503s for `/v1/namespaces` then a 200. | ||
|
|
||
| Used in place of `requests_mock`, which replaces the HTTPAdapter and would bypass the retry logic. | ||
|
|
||
| Yields a dict with `port` and `namespace_calls` keys (the latter is updated in-place as requests arrive). | ||
| """ | ||
| import json | ||
| import threading | ||
| from http.server import BaseHTTPRequestHandler, HTTPServer | ||
|
|
||
| state: dict[str, Any] = {"namespace_calls": 0} | ||
| config_body = json.dumps( | ||
| {"defaults": {}, "overrides": {}, "endpoints": [str(endpoint) for endpoint in TEST_SUPPORTED_ENDPOINTS]} | ||
| ).encode() | ||
|
|
||
| class _Handler(BaseHTTPRequestHandler): | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can you add the server setup into a different method? That way, we can easily see what this is actually testing + less about the test setup.
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done in 47a5382 — extracted the handler + threading setup into a |
||
| def do_GET(self) -> None: | ||
| if self.path.endswith("/v1/config"): | ||
| self._respond(200, config_body) | ||
| elif self.path.endswith("/v1/namespaces"): | ||
| state["namespace_calls"] += 1 | ||
| if state["namespace_calls"] <= num_failures: | ||
| self._respond(503, b"") | ||
| else: | ||
| self._respond(200, json.dumps({"namespaces": [["foo"]]}).encode()) | ||
| else: | ||
| self._respond(404, b"") | ||
|
|
||
| def _respond(self, status: int, body: bytes) -> None: | ||
| self.send_response(status) | ||
| self.send_header("Content-Type", "application/json") | ||
| self.send_header("Content-Length", str(len(body))) | ||
| self.end_headers() | ||
| if body: | ||
| self.wfile.write(body) | ||
|
|
||
| def log_message(self, format: str, *args: Any) -> None: # silence default access logs | ||
| pass | ||
|
|
||
| server = HTTPServer(("127.0.0.1", 0), _Handler) | ||
| state["port"] = server.server_address[1] | ||
| server_thread = threading.Thread(target=server.serve_forever, daemon=True) | ||
| server_thread.start() | ||
| try: | ||
| yield state | ||
| finally: | ||
| server.shutdown() | ||
| server.server_close() | ||
|
|
||
|
|
||
| def test_session_retries_on_transient_5xx_then_succeeds() -> None: | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We should add a test that throws
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Added |
||
| """The catalog should retry on transient 5xx and succeed once the server stabilizes.""" | ||
| with _local_rest_server_503_then_200(num_failures=3) as server: | ||
| catalog = RestCatalog( | ||
| "rest", | ||
| **{ # type: ignore | ||
| "uri": f"http://127.0.0.1:{server['port']}/", | ||
| "token": TEST_TOKEN, | ||
| # backoff-factor=0 keeps the test fast; retries=3 covers three 503s + the eventual 200. | ||
| REST_CLIENT_MAX_RETRIES: 3, | ||
| REST_CLIENT_RETRY_BACKOFF_FACTOR: 0, | ||
| }, | ||
| ) | ||
| assert catalog.list_namespaces() == [("foo",)] | ||
| assert server["namespace_calls"] == 4 | ||
|
|
||
|
|
||
| def test_session_exhausted_retries_surfaces_typed_exception() -> None: | ||
| """When retries are exhausted, the typed exception from `_handle_non_200_response` should be raised | ||
| (e.g. `ServiceUnavailableError` for 503), not the urllib3 `MaxRetryError` / `RetryError`.""" | ||
| # `num_failures` greater than `retries + 1` guarantees the server never returns success. | ||
| with _local_rest_server_503_then_200(num_failures=100) as server: | ||
| catalog = RestCatalog( | ||
| "rest", | ||
| **{ # type: ignore | ||
| "uri": f"http://127.0.0.1:{server['port']}/", | ||
| "token": TEST_TOKEN, | ||
| REST_CLIENT_MAX_RETRIES: 2, | ||
| REST_CLIENT_RETRY_BACKOFF_FACTOR: 0, | ||
| }, | ||
| ) | ||
| with pytest.raises(ServiceUnavailableError): | ||
| catalog.list_namespaces() | ||
| # retries=2 means 1 initial attempt + 2 retries = 3 calls | ||
| assert server["namespace_calls"] == 3 | ||
|
|
||
|
|
||
| def test_session_with_invalid_connection_timeout_raises(rest_mock: Mocker) -> None: | ||
| catalog_properties = { | ||
| "uri": TEST_URI, | ||
| "token": TEST_TOKEN, | ||
| REST_CLIENT_REQUEST_TIMEOUT: -1, | ||
| } | ||
| with pytest.raises(ValueError, match="`rest.client.request-timeout` must be a positive number"): | ||
| RestCatalog("rest", **catalog_properties) # type: ignore | ||
|
|
||
|
|
||
| def test_session_with_invalid_connection_retries_raises(rest_mock: Mocker) -> None: | ||
| catalog_properties = { | ||
| "uri": TEST_URI, | ||
| "token": TEST_TOKEN, | ||
| REST_CLIENT_MAX_RETRIES: -1, | ||
| } | ||
| with pytest.raises(ValueError, match="`rest.client.max-retries` must be non-negative"): | ||
| RestCatalog("rest", **catalog_properties) # type: ignore | ||
|
|
||
|
|
||
| def test_rest_catalog_with_basic_auth_type(rest_mock: Mocker) -> None: | ||
| # Given | ||
| rest_mock.get( | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sorry for not flagging this earlier, but we can drop this, and re-use the methods in
pyiceberg/utils/properties.py