Skip to content

Commit 067e0dc

Browse files
author
Mrutunjay Kinagi
committed
feat(rest): add sigv4 retry configuration defaults (#3008)
1 parent 9de7deb commit 067e0dc

File tree

2 files changed

+60
-2
lines changed

2 files changed

+60
-2
lines changed

pyiceberg/catalog/rest/__init__.py

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@
7777
from pyiceberg.typedef import EMPTY_DICT, UTF8, IcebergBaseModel, Identifier, Properties
7878
from pyiceberg.types import transform_dict_value_to_str
7979
from pyiceberg.utils.deprecated import deprecation_message
80-
from pyiceberg.utils.properties import get_first_property_value, get_header_properties, property_as_bool
80+
from pyiceberg.utils.properties import get_first_property_value, get_header_properties, property_as_bool, property_as_int
8181

8282
if TYPE_CHECKING:
8383
import pyarrow as pa
@@ -223,6 +223,10 @@ class IdentifierKind(Enum):
223223
SIGV4 = "rest.sigv4-enabled"
224224
SIGV4_REGION = "rest.signing-region"
225225
SIGV4_SERVICE = "rest.signing-name"
226+
SIGV4_MAX_RETRIES = "rest.sigv4.max-retries"
227+
SIGV4_DEFAULT_MAX_RETRIES = 10
228+
SIGV4_RETRY_BACKOFF_FACTOR = 0.5
229+
SIGV4_RETRY_STATUS_CODES = (429, 500, 502, 503, 504)
226230
EMPTY_BODY_SHA256: str = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
227231
OAUTH2_SERVER_URI = "oauth2-server-uri"
228232
SNAPSHOT_LOADING_MODE = "snapshot-loading-mode"
@@ -680,11 +684,25 @@ def _init_sigv4(self, session: Session) -> None:
680684
from botocore.awsrequest import AWSRequest
681685
from requests import PreparedRequest
682686
from requests.adapters import HTTPAdapter
687+
from urllib3.util.retry import Retry
683688

684689
class SigV4Adapter(HTTPAdapter):
685690
def __init__(self, **properties: str):
686-
super().__init__()
687691
self._properties = properties
692+
max_retries = property_as_int(self._properties, SIGV4_MAX_RETRIES, SIGV4_DEFAULT_MAX_RETRIES)
693+
super().__init__(
694+
max_retries=Retry(
695+
total=max_retries,
696+
status=max_retries,
697+
connect=max_retries,
698+
read=max_retries,
699+
# Keep retries conservative for idempotent calls.
700+
allowed_methods=frozenset({"GET", "HEAD", "OPTIONS"}),
701+
status_forcelist=SIGV4_RETRY_STATUS_CODES,
702+
backoff_factor=SIGV4_RETRY_BACKOFF_FACTOR,
703+
respect_retry_after_header=True,
704+
)
705+
)
688706
self._boto_session = boto3.Session(
689707
region_name=get_first_property_value(self._properties, AWS_REGION),
690708
botocore_session=self._properties.get(BOTOCORE_SESSION),

tests/catalog/test_rest.py

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@
3333
DEFAULT_ENDPOINTS,
3434
EMPTY_BODY_SHA256,
3535
OAUTH2_SERVER_URI,
36+
SIGV4_DEFAULT_MAX_RETRIES,
37+
SIGV4_MAX_RETRIES,
3638
SNAPSHOT_LOADING_MODE,
3739
Capability,
3840
RestCatalog,
@@ -527,6 +529,44 @@ def test_sigv4_sign_request_with_body(rest_mock: Mocker) -> None:
527529
assert prepared.headers.get("x-amz-content-sha256") != EMPTY_BODY_SHA256
528530

529531

532+
def test_sigv4_adapter_default_retry_config(rest_mock: Mocker) -> None:
533+
catalog = RestCatalog(
534+
"rest",
535+
**{
536+
"uri": TEST_URI,
537+
"token": TEST_TOKEN,
538+
"rest.sigv4-enabled": "true",
539+
"rest.signing-region": "us-west-2",
540+
"client.access-key-id": "id",
541+
"client.secret-access-key": "secret",
542+
},
543+
)
544+
545+
adapter = catalog._session.adapters[catalog.uri]
546+
assert isinstance(adapter, HTTPAdapter)
547+
assert adapter.max_retries.total == SIGV4_DEFAULT_MAX_RETRIES
548+
assert 429 in adapter.max_retries.status_forcelist
549+
550+
551+
def test_sigv4_adapter_override_retry_config(rest_mock: Mocker) -> None:
552+
catalog = RestCatalog(
553+
"rest",
554+
**{
555+
"uri": TEST_URI,
556+
"token": TEST_TOKEN,
557+
"rest.sigv4-enabled": "true",
558+
"rest.signing-region": "us-west-2",
559+
"client.access-key-id": "id",
560+
"client.secret-access-key": "secret",
561+
SIGV4_MAX_RETRIES: "3",
562+
},
563+
)
564+
565+
adapter = catalog._session.adapters[catalog.uri]
566+
assert isinstance(adapter, HTTPAdapter)
567+
assert adapter.max_retries.total == 3
568+
569+
530570
def test_list_tables_404(rest_mock: Mocker) -> None:
531571
namespace = "examples"
532572
rest_mock.get(

0 commit comments

Comments
 (0)