diff --git a/pynamodb/connection/base.py b/pynamodb/connection/base.py index 301ee171..c8a5c3bb 100644 --- a/pynamodb/connection/base.py +++ b/pynamodb/connection/base.py @@ -264,7 +264,8 @@ def __init__(self, extra_headers: Optional[Mapping[str, str]] = None, aws_access_key_id: Optional[str] = None, aws_secret_access_key: Optional[str] = None, - aws_session_token: Optional[str] = None): + aws_session_token: Optional[str] = None, + tcp_keepalive: Optional[bool] = None): self._tables: Dict[str, MetaTable] = {} self.host = host self._local = local() @@ -312,6 +313,12 @@ def __init__(self, else: self._extra_headers = get_settings_value('extra_headers') + if tcp_keepalive is not None: + self._tcp_keepalive = tcp_keepalive + else: + self._tcp_keepalive = get_settings_value('tcp_keepalive') + + self._aws_access_key_id = aws_access_key_id self._aws_secret_access_key = aws_secret_access_key self._aws_session_token = aws_session_token @@ -439,6 +446,7 @@ def client(self) -> BotocoreBaseClientPrivate: connect_timeout=self._connect_timeout_seconds, read_timeout=self._read_timeout_seconds, max_pool_connections=self._max_pool_connections, + tcp_keepalive=self._tcp_keepalive, retries=retries, ) self._client = cast(BotocoreBaseClientPrivate, self.session.create_client(SERVICE_NAME, self.region, endpoint_url=self.host, config=config)) diff --git a/pynamodb/connection/table.py b/pynamodb/connection/table.py index 5e70ba5c..b8a4526b 100644 --- a/pynamodb/connection/table.py +++ b/pynamodb/connection/table.py @@ -29,6 +29,7 @@ def __init__( aws_access_key_id: Optional[str] = None, aws_secret_access_key: Optional[str] = None, aws_session_token: Optional[str] = None, + tcp_keepalive: Optional[bool] = None, *, meta_table: Optional[MetaTable] = None, ) -> None: @@ -42,7 +43,8 @@ def __init__( extra_headers=extra_headers, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token) + aws_session_token=aws_session_token, + tcp_keepalive=tcp_keepalive) if meta_table is not None: self.connection.add_meta_table(meta_table) diff --git a/pynamodb/models.py b/pynamodb/models.py index c75e99b3..15ed737b 100644 --- a/pynamodb/models.py +++ b/pynamodb/models.py @@ -183,6 +183,7 @@ class MetaProtocol(Protocol): read_timeout_seconds: int max_retry_attempts: int max_pool_connections: int + tcp_keepalive: bool extra_headers: Mapping[str, str] aws_access_key_id: Optional[str] aws_secret_access_key: Optional[str] @@ -237,6 +238,8 @@ def __init__(self, name, bases, namespace, discriminator=None) -> None: warnings.warn("The `session_cls` and `request_timeout_second` options are no longer supported") if not hasattr(attr_obj, 'connect_timeout_seconds'): setattr(attr_obj, 'connect_timeout_seconds', get_settings_value('connect_timeout_seconds')) + if not hasattr(attr_obj, 'tcp_keepalive'): + setattr(attr_obj, 'tcp_keepalive', get_settings_value('tcp_keepalive')) if not hasattr(attr_obj, 'read_timeout_seconds'): setattr(attr_obj, 'read_timeout_seconds', get_settings_value('read_timeout_seconds')) if not hasattr(attr_obj, 'max_retry_attempts'): @@ -1085,6 +1088,7 @@ def _get_connection(cls) -> TableConnection: read_timeout_seconds=cls.Meta.read_timeout_seconds, max_retry_attempts=cls.Meta.max_retry_attempts, max_pool_connections=cls.Meta.max_pool_connections, + tcp_keepalive=cls.Meta.tcp_keepalive, extra_headers=cls.Meta.extra_headers, aws_access_key_id=cls.Meta.aws_access_key_id, aws_secret_access_key=cls.Meta.aws_secret_access_key, diff --git a/pynamodb/settings.py b/pynamodb/settings.py index 503f060d..ff8283dc 100644 --- a/pynamodb/settings.py +++ b/pynamodb/settings.py @@ -15,7 +15,8 @@ 'region': None, 'max_pool_connections': 10, 'extra_headers': None, - 'retry_configuration': 'LEGACY' + 'retry_configuration': 'LEGACY', + 'tcp_keepalive': False } OVERRIDE_SETTINGS_PATH = getenv('PYNAMODB_CONFIG', '/etc/pynamodb/global_default_settings.py') diff --git a/tests/test_base_connection.py b/tests/test_base_connection.py index c90acd27..0a3987fd 100644 --- a/tests/test_base_connection.py +++ b/tests/test_base_connection.py @@ -1517,6 +1517,11 @@ def test_connection__botocore_config(): assert c.client._client_config.connect_timeout == 5 assert c.client._client_config.read_timeout == 10 assert c.client._client_config.max_pool_connections == 20 + assert c.client._client_config.tcp_keepalive is False + +def test_connection__botocore_config_tcp_keepalive(): + c = Connection(tcp_keepalive=True) + assert c.client._client_config.tcp_keepalive is True @freeze_time() @@ -1670,6 +1675,7 @@ def test_connection_client_retry_configuration( 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, + tcp_keepalive=False, retries=expected_retries ) # Ensure the session was created correctly. diff --git a/tests/test_model.py b/tests/test_model.py index da54303b..2987d337 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -519,11 +519,13 @@ def fake_dynamodb(*args): assert UserModel.Meta.read_timeout_seconds == 30 assert UserModel.Meta.max_retry_attempts == 3 assert UserModel.Meta.max_pool_connections == 10 + assert UserModel.Meta.tcp_keepalive == False assert UserModel._connection.connection._connect_timeout_seconds == 15 assert UserModel._connection.connection._read_timeout_seconds == 30 assert UserModel._connection.connection._max_retry_attempts_exception == 3 assert UserModel._connection.connection._max_pool_connections == 10 + assert UserModel._connection.connection._tcp_keepalive == False with patch(PATCH_METHOD) as req: req.return_value = MODEL_TABLE_DATA diff --git a/tests/test_settings.py b/tests/test_settings.py index 88b3a598..9881c095 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -31,6 +31,7 @@ def test_default_settings(): 'max_retry_attempts': 3, 'region': None, 'max_pool_connections': 10, + 'tcp_keepalive': False, 'extra_headers': None, 'retry_configuration': 'LEGACY' }