Skip to content

feat: allow setting or unsetting the boto retry configuration #1271

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

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
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
15 changes: 15 additions & 0 deletions docs/settings.rst
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,21 @@ Default: automatically constructed by boto to account for region
The URL endpoint for DynamoDB. This can be used to use a local implementation of DynamoDB such as DynamoDB Local or dynalite.


retry_configuration
-------------------

Default: ``"LEGACY"``

This controls the PynamoDB retry behavior. The default of ``"LEGACY"`` keeps the
existing PynamoDB retry behavior. If set to ``None``, this will use botocore's default
retry configuration discovery mechanism as documented
`in boto3 <https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html#retries>`_
and
`in the AWS SDK docs <https://docs.aws.amazon.com/sdkref/latest/guide/feature-retry-behavior.html>`_.
If set to a retry configuration dictionary as described
`here <https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html#defining-a-retry-configuration-in-a-config-object-for-your-boto3-client>`_
it will be used directly in the botocore client configuration.

Overriding settings
~~~~~~~~~~~~~~~~~~~

Expand Down
42 changes: 37 additions & 5 deletions pynamodb/connection/base.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,17 @@
"""
Lowest level connection
"""
import sys
import logging
import uuid
from threading import local
from typing import Any, Dict, List, Mapping, Optional, Sequence, cast
from typing import Any, Dict, List, Mapping, Optional, Sequence, Union, cast
if sys.version_info >= (3, 8):
from typing import Literal
else:
from typing_extensions import Literal

import botocore.config
import botocore.client
import botocore.exceptions
from botocore.client import ClientError
Expand Down Expand Up @@ -247,6 +253,13 @@ def __init__(self,
read_timeout_seconds: Optional[float] = None,
connect_timeout_seconds: Optional[float] = None,
max_retry_attempts: Optional[int] = None,
retry_configuration: Optional[
Union[
Literal["LEGACY"],
Literal["UNSET"],
"botocore.config._RetryDict",
]
] = None,
max_pool_connections: Optional[int] = None,
extra_headers: Optional[Mapping[str, str]] = None,
aws_access_key_id: Optional[str] = None,
Expand Down Expand Up @@ -277,6 +290,18 @@ def __init__(self,
else:
self._max_retry_attempts_exception = get_settings_value('max_retry_attempts')

# Since we have the pattern of using `None` to indicate "read from the
# settings", we use a literal of "UNSET" to indicate we want the
# `_retry_configuration` attribute set to `None` so botocore will use its own
# retry configuration discovery logic. This was required so direct users of the
# `Connection` class can still leave the retry configuration unset.
if retry_configuration is not None and retry_configuration == "UNSET":
self._retry_configuration = None
elif retry_configuration is not None:
self._retry_configuration = retry_configuration
else:
self._retry_configuration = get_settings_value('retry_configuration')

if max_pool_connections is not None:
self._max_pool_connections = max_pool_connections
else:
Expand Down Expand Up @@ -399,15 +424,22 @@ def client(self) -> BotocoreBaseClientPrivate:
# if the client does not have credentials, we create a new client
# otherwise the client is permanently poisoned in the case of metadata service flakiness when using IAM roles
if not self._client or (self._client._request_signer and not self._client._request_signer._credentials):
# Check if we are using the "LEGACY" retry mode to keep previous PynamoDB
# retry behavior, or if we are using the new retry configuration settings.
if self._retry_configuration != "LEGACY":
retries = self._retry_configuration
else:
retries = {
'total_max_attempts': 1 + self._max_retry_attempts_exception,
'mode': 'standard',
}

config = botocore.client.Config(
parameter_validation=False, # Disable unnecessary validation for performance
connect_timeout=self._connect_timeout_seconds,
read_timeout=self._read_timeout_seconds,
max_pool_connections=self._max_pool_connections,
retries={
'total_max_attempts': 1 + self._max_retry_attempts_exception,
'mode': 'standard',
}
retries=retries
)
self._client = cast(BotocoreBaseClientPrivate, self.session.create_client(SERVICE_NAME, self.region, endpoint_url=self.host, config=config))

Expand Down
1 change: 1 addition & 0 deletions pynamodb/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
'region': None,
'max_pool_connections': 10,
'extra_headers': None,
'retry_configuration': 'LEGACY'
}

OVERRIDE_SETTINGS_PATH = getenv('PYNAMODB_CONFIG', '/etc/pynamodb/global_default_settings.py')
Expand Down
66 changes: 66 additions & 0 deletions tests/test_base_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -1632,3 +1632,69 @@ def test_connection_update_time_to_live__fail():
req.side_effect = BotoCoreError
with pytest.raises(TableError):
conn.update_time_to_live('test table', 'my_ttl')

@pytest.mark.parametrize(
"retry_configuration, expected_retries",
(
(None, None),
(
"LEGACY",
{
'total_max_attempts': 4,
'mode': 'standard',
}
),
(
{"max_attempts": 10, "mode": "adaptive"},
{"max_attempts": 10, "mode": "adaptive"},
)
)
)
def test_connection_client_retry_configuration(
retry_configuration, expected_retries, mocker
):
"""Test that the client respects the retry configuration setting."""
mock_client_config = mocker.patch(target="botocore.client.Config", autospec=True)
mock_session_property = mocker.patch.object(
target=Connection, attribute="session", autospec=True
)

unit_under_test = Connection()
unit_under_test._retry_configuration = retry_configuration
unit_under_test.client

# Ensure the configuration was called correctly, and used the appropriate retry
# configuration.
mock_client_config.assert_called_once_with(
parameter_validation=False,
connect_timeout=unit_under_test._connect_timeout_seconds,
read_timeout=unit_under_test._read_timeout_seconds,
max_pool_connections=unit_under_test._max_pool_connections,
retries=expected_retries
)
# Ensure the session was created correctly.
mock_session_property.create_client.assert_called_once_with(
"dynamodb",
unit_under_test.region,
endpoint_url=unit_under_test.host,
config=mock_client_config.return_value,
)

@pytest.mark.parametrize(
"retry_configuration, expected_retry_configuration",
(
(None, "LEGACY"),
("LEGACY","LEGACY"),
("UNSET", None),
(
{"max_attempts": 10, "mode": "adaptive"},
{"max_attempts": 10, "mode": "adaptive"},
)
)
)
def test_connection_client_retry_configuration__init__(
retry_configuration, expected_retry_configuration
):
"""Test that the __init__ properly sets the `_retry_configuration` attribute."""
unit_under_test = Connection(retry_configuration=retry_configuration)
assert unit_under_test._retry_configuration == expected_retry_configuration
15 changes: 14 additions & 1 deletion tests/test_settings.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import sys
from unittest.mock import patch

import pytest
Expand All @@ -20,3 +19,17 @@ def test_override_old_attributes(settings_str, tmpdir):
reload(pynamodb.settings)
assert len(warns) == 1
assert 'options are no longer supported' in str(warns[0].message)

def test_default_settings():
"""Ensure that the default settings are what we expect. This is mainly done to catch
any potentially breaking changes to default settings.
"""
assert pynamodb.settings.default_settings_dict == {
'connect_timeout_seconds': 15,
'read_timeout_seconds': 30,
'max_retry_attempts': 3,
'region': None,
'max_pool_connections': 10,
'extra_headers': None,
'retry_configuration': 'LEGACY'
}