Skip to content

Commit ecf6987

Browse files
committed
feat: allow setting or unsetting the boto retry configuration
This adds the ability to directly set the boto retry configuration dictionary, or to leave it unset and allow botocore to automatically discover the configuration from the environment or `~/.aws/config` files. The default is to use the previous PynamoDB behavior for configuring retries so as to not make this a breaking change.
1 parent f0bc917 commit ecf6987

File tree

5 files changed

+132
-6
lines changed

5 files changed

+132
-6
lines changed

docs/settings.rst

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,20 @@ Default: automatically constructed by boto to account for region
6969
The URL endpoint for DynamoDB. This can be used to use a local implementation of DynamoDB such as DynamoDB Local or dynalite.
7070

7171

72+
retry_configuration
73+
-------------------
74+
75+
Default: ``"LEGACY"``
76+
77+
This controls the PynamoDB retry behavior. The default of ``"LEGACY"`` keeps the
78+
existing PynamoDB retry behavior. If set to ``None``, this will use botocore's default
79+
retry configuration discovery mechanism as documented
80+
`here <https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html#retries>`_
81+
and `here <https://docs.aws.amazon.com/sdkref/latest/guide/feature-retry-behavior.html>`_.
82+
If set to a retry configuration dictionary as described
83+
`here <https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html#defining-a-retry-configuration-in-a-config-object-for-your-boto3-client>`_
84+
it will be used directly in the botocore client configuration.
85+
7286
Overriding settings
7387
~~~~~~~~~~~~~~~~~~~
7488

pynamodb/connection/base.py

Lines changed: 37 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,17 @@
11
"""
22
Lowest level connection
33
"""
4+
import sys
45
import logging
56
import uuid
67
from threading import local
7-
from typing import Any, Dict, List, Mapping, Optional, Sequence, cast
8+
from typing import Any, Dict, List, Mapping, Optional, Sequence, Union, cast
9+
if sys.version_info >= (3, 8):
10+
from typing import Literal
11+
else:
12+
from typing_extensions import Literal
813

14+
import botocore.config
915
import botocore.client
1016
import botocore.exceptions
1117
from botocore.client import ClientError
@@ -247,6 +253,13 @@ def __init__(self,
247253
read_timeout_seconds: Optional[float] = None,
248254
connect_timeout_seconds: Optional[float] = None,
249255
max_retry_attempts: Optional[int] = None,
256+
retry_configuration: Optional[
257+
Union[
258+
Literal["LEGACY"],
259+
Literal["UNSET"],
260+
"botocore.config._RetryDict",
261+
]
262+
] = None,
250263
max_pool_connections: Optional[int] = None,
251264
extra_headers: Optional[Mapping[str, str]] = None,
252265
aws_access_key_id: Optional[str] = None,
@@ -277,6 +290,18 @@ def __init__(self,
277290
else:
278291
self._max_retry_attempts_exception = get_settings_value('max_retry_attempts')
279292

293+
# Since we have the pattern of using `None` to indicate "read from the
294+
# settings", we use a literal of "UNSET" to indicate we want the
295+
# `_retry_configuration` attribute set to `None` so botocore will use its own
296+
# retry configuration discovery logic. This was required so direct users of the
297+
# `Connection` class can still leave the retry configuration unset.
298+
if retry_configuration is not None and retry_configuration == "UNSET":
299+
self._retry_configuration = None
300+
elif retry_configuration is not None:
301+
self._retry_configuration = retry_configuration
302+
else:
303+
self._retry_configuration = get_settings_value('retry_configuration')
304+
280305
if max_pool_connections is not None:
281306
self._max_pool_connections = max_pool_connections
282307
else:
@@ -399,15 +424,22 @@ def client(self) -> BotocoreBaseClientPrivate:
399424
# if the client does not have credentials, we create a new client
400425
# otherwise the client is permanently poisoned in the case of metadata service flakiness when using IAM roles
401426
if not self._client or (self._client._request_signer and not self._client._request_signer._credentials):
427+
# Check if we are using the "LEGACY" retry mode to keep previous PynamoDB
428+
# retry behavior, or if we are using the new retry configuration settings.
429+
if self._retry_configuration != "LEGACY":
430+
retries = self._retry_configuration
431+
else:
432+
retries = {
433+
'total_max_attempts': 1 + self._max_retry_attempts_exception,
434+
'mode': 'standard',
435+
}
436+
402437
config = botocore.client.Config(
403438
parameter_validation=False, # Disable unnecessary validation for performance
404439
connect_timeout=self._connect_timeout_seconds,
405440
read_timeout=self._read_timeout_seconds,
406441
max_pool_connections=self._max_pool_connections,
407-
retries={
408-
'total_max_attempts': 1 + self._max_retry_attempts_exception,
409-
'mode': 'standard',
410-
}
442+
retries=retries
411443
)
412444
self._client = cast(BotocoreBaseClientPrivate, self.session.create_client(SERVICE_NAME, self.region, endpoint_url=self.host, config=config))
413445

pynamodb/settings.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
'region': None,
1616
'max_pool_connections': 10,
1717
'extra_headers': None,
18+
'retry_configuration': 'LEGACY'
1819
}
1920

2021
OVERRIDE_SETTINGS_PATH = getenv('PYNAMODB_CONFIG', '/etc/pynamodb/global_default_settings.py')

tests/test_base_connection.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1632,3 +1632,69 @@ def test_connection_update_time_to_live__fail():
16321632
req.side_effect = BotoCoreError
16331633
with pytest.raises(TableError):
16341634
conn.update_time_to_live('test table', 'my_ttl')
1635+
1636+
@pytest.mark.parametrize(
1637+
"retry_configuration, expected_retries",
1638+
(
1639+
(None, None),
1640+
(
1641+
"LEGACY",
1642+
{
1643+
'total_max_attempts': 4,
1644+
'mode': 'standard',
1645+
}
1646+
),
1647+
(
1648+
{"max_attempts": 10, "mode": "adaptive"},
1649+
{"max_attempts": 10, "mode": "adaptive"},
1650+
)
1651+
)
1652+
)
1653+
def test_connection_client_retry_configuration(
1654+
retry_configuration, expected_retries, mocker
1655+
):
1656+
"""Test that the client respects the retry configuration setting."""
1657+
mock_client_config = mocker.patch(target="botocore.client.Config", autospec=True)
1658+
mock_session_property = mocker.patch.object(
1659+
target=Connection, attribute="session", autospec=True
1660+
)
1661+
1662+
unit_under_test = Connection()
1663+
unit_under_test._retry_configuration = retry_configuration
1664+
unit_under_test.client
1665+
1666+
# Ensure the configuration was called correctly, and used the appropriate retry
1667+
# configuration.
1668+
mock_client_config.assert_called_once_with(
1669+
parameter_validation=False,
1670+
connect_timeout=unit_under_test._connect_timeout_seconds,
1671+
read_timeout=unit_under_test._read_timeout_seconds,
1672+
max_pool_connections=unit_under_test._max_pool_connections,
1673+
retries=expected_retries
1674+
)
1675+
# Ensure the session was created correctly.
1676+
mock_session_property.create_client.assert_called_once_with(
1677+
"dynamodb",
1678+
unit_under_test.region,
1679+
endpoint_url=unit_under_test.host,
1680+
config=mock_client_config.return_value,
1681+
)
1682+
1683+
@pytest.mark.parametrize(
1684+
"retry_configuration, expected_retry_configuration",
1685+
(
1686+
(None, "LEGACY"),
1687+
("LEGACY","LEGACY"),
1688+
("UNSET", None),
1689+
(
1690+
{"max_attempts": 10, "mode": "adaptive"},
1691+
{"max_attempts": 10, "mode": "adaptive"},
1692+
)
1693+
)
1694+
)
1695+
def test_connection_client_retry_configuration__init__(
1696+
retry_configuration, expected_retry_configuration
1697+
):
1698+
"""Test that the __init__ properly sets the `_retry_configuration` attribute."""
1699+
unit_under_test = Connection(retry_configuration=retry_configuration)
1700+
assert unit_under_test._retry_configuration == expected_retry_configuration

tests/test_settings.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import sys
21
from unittest.mock import patch
32

43
import pytest
@@ -20,3 +19,17 @@ def test_override_old_attributes(settings_str, tmpdir):
2019
reload(pynamodb.settings)
2120
assert len(warns) == 1
2221
assert 'options are no longer supported' in str(warns[0].message)
22+
23+
def test_default_settings():
24+
"""Ensure that the default settings are what we expect. This is mainly done to catch
25+
any potentially breaking changes to default settings.
26+
"""
27+
assert pynamodb.settings.default_settings_dict == {
28+
'connect_timeout_seconds': 15,
29+
'read_timeout_seconds': 30,
30+
'max_retry_attempts': 3,
31+
'region': None,
32+
'max_pool_connections': 10,
33+
'extra_headers': None,
34+
'retry_configuration': 'LEGACY'
35+
}

0 commit comments

Comments
 (0)