Add botocore Config to wr.config (#535)

Author: maxispeicher

awswrangler/_config.py

@@ -5,18 +5,19 @@
 import os
 from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple, Type, Union, cast
 
+import botocore.config
 import pandas as pd
 
 from awswrangler import exceptions
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
-_ConfigValueType = Union[str, bool, int, None]
+_ConfigValueType = Union[str, bool, int, botocore.config.Config, None]
 
 
 class _ConfigArg(NamedTuple):
-    dtype: Type[Union[str, bool, int]]
+    dtype: Type[Union[str, bool, int, botocore.config.Config]]
     nullable: bool
     enforced: bool = False
 
@@ -41,6 +42,8 @@ class _ConfigArg(NamedTuple):
     "redshift_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
     "kms_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
     "emr_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
+    # Botocore config
+    "botocore_config": _ConfigArg(dtype=botocore.config.Config, nullable=True),
 }
 
 
@@ -57,6 +60,7 @@ def __init__(self) -> None:
         self.redshift_endpoint_url = None
         self.kms_endpoint_url = None
         self.emr_endpoint_url = None
+        self.botocore_config = None
         for name in _CONFIG_ARGS:
             self._load_config(name=name)
 
@@ -338,6 +342,15 @@ def emr_endpoint_url(self) -> Optional[str]:
     def emr_endpoint_url(self, value: Optional[str]) -> None:
         self._set_config_value(key="emr_endpoint_url", value=value)
 
+    @property
+    def botocore_config(self) -> botocore.config.Config:
+        """Property botocore_config."""
+        return cast(Optional[botocore.config.Config], self["botocore_config"])
+
+    @botocore_config.setter
+    def botocore_config(self, value: Optional[botocore.config.Config]) -> None:
+        self._set_config_value(key="botocore_config", value=value)
+
 
 def _insert_str(text: str, token: str, insert: str) -> str:
     """Insert string into other."""

awswrangler/_utils.py

@@ -16,6 +16,7 @@
 import pandas as pd
 
 from awswrangler import _config, exceptions
+from awswrangler._config import apply_configs
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -59,9 +60,13 @@ def boto3_from_primitives(primitives: Optional[Boto3PrimitivesType] = None) -> b
     return boto3.Session(**args)
 
 
-def botocore_config() -> botocore.config.Config:
+def default_botocore_config() -> botocore.config.Config:
     """Botocore configuration."""
-    return botocore.config.Config(retries={"max_attempts": 5}, connect_timeout=10, max_pool_connections=10)
+    retries_config: Dict[str, Union[str, int]] = {
+        "max_attempts": int(os.getenv("AWS_MAX_ATTEMPTS", "3")),
+        "mode": os.getenv("AWS_RETRY_MODE", "standard"),
+    }
+    return botocore.config.Config(retries=retries_config, connect_timeout=10, max_pool_connections=10)
 
 
 def _get_endpoint_url(service_name: str) -> Optional[str]:
@@ -83,24 +88,31 @@ def _get_endpoint_url(service_name: str) -> Optional[str]:
     return endpoint_url
 
 
+@apply_configs
 def client(
-    service_name: str, session: Optional[boto3.Session] = None, config: Optional[botocore.config.Config] = None
+    service_name: str, session: Optional[boto3.Session] = None, botocore_config: Optional[botocore.config.Config] = None
 ) -> boto3.client:
     """Create a valid boto3.client."""
     endpoint_url: Optional[str] = _get_endpoint_url(service_name=service_name)
     return ensure_session(session=session).client(
         service_name=service_name,
         endpoint_url=endpoint_url,
         use_ssl=True,
-        config=botocore_config() if config is None else config,
+        config=default_botocore_config() if botocore_config is None else botocore_config,
     )
 
 
-def resource(service_name: str, session: Optional[boto3.Session] = None) -> boto3.resource:
+@apply_configs
+def resource(
+    service_name: str, session: Optional[boto3.Session] = None, botocore_config: Optional[botocore.config.Config] = None
+) -> boto3.resource:
     """Create a valid boto3.resource."""
     endpoint_url: Optional[str] = _get_endpoint_url(service_name=service_name)
     return ensure_session(session=session).resource(
-        service_name=service_name, endpoint_url=endpoint_url, use_ssl=True, config=botocore_config()
+        service_name=service_name,
+        endpoint_url=endpoint_url,
+        use_ssl=True,
+        config=default_botocore_config() if botocore_config is None else botocore_config,
     )
 
 

awswrangler/timestream.py

@@ -39,7 +39,7 @@ def _write_batch(
     client: boto3.client = _utils.client(
         service_name="timestream-write",
         session=boto3_session,
-        config=Config(read_timeout=20, max_pool_connections=5000, retries={"max_attempts": 10}),
+        botocore_config=Config(read_timeout=20, max_pool_connections=5000, retries={"max_attempts": 10}),
     )
     try:
         _utils.try_it(
@@ -218,7 +218,7 @@ def query(sql: str, boto3_session: Optional[boto3.Session] = None) -> pd.DataFra
     client: boto3.client = _utils.client(
         service_name="timestream-query",
         session=boto3_session,
-        config=Config(read_timeout=60, retries={"max_attempts": 10}),
+        botocore_config=Config(read_timeout=60, retries={"max_attempts": 10}),
     )
     paginator = client.get_paginator("query")
     rows: List[List[Any]] = []

tests/test_config.py

@@ -4,6 +4,8 @@
 
 import boto3
 import botocore
+import botocore.client
+import botocore.config
 import pytest
 
 import awswrangler as wr
@@ -121,3 +123,60 @@ def test_basics(path, glue_database, glue_table, workgroup0, workgroup1):
 def test_athena_cache_configuration():
     wr.config.max_local_cache_entries = 20
     assert wr.config.max_remote_cache_entries == 20
+
+
+def test_botocore_config(path):
+    original = botocore.client.ClientCreator.create_client
+
+    # Default values for botocore.config.Config
+    expected_max_retries_attempt = 3
+    expected_connect_timeout = 10
+    expected_max_pool_connections = 10
+    expected_retry_mode = "standard"
+
+    def wrapper(self, **kwarg):
+        assert kwarg["client_config"].retries["max_attempts"] == expected_max_retries_attempt
+        assert kwarg["client_config"].connect_timeout == expected_connect_timeout
+        assert kwarg["client_config"].max_pool_connections == expected_max_pool_connections
+        assert kwarg["client_config"].retries["mode"] == expected_retry_mode
+        return original(self, **kwarg)
+
+    # Check for default values
+    with patch("botocore.client.ClientCreator.create_client", new=wrapper):
+        with open_s3_object(path, mode="wb") as s3obj:
+            s3obj.write(b"foo")
+
+    # Update default config with environment variables
+    expected_max_retries_attempt = 20
+    expected_connect_timeout = 10
+    expected_max_pool_connections = 10
+    expected_retry_mode = "adaptive"
+
+    os.environ["AWS_MAX_ATTEMPTS"] = str(expected_max_retries_attempt)
+    os.environ["AWS_RETRY_MODE"] = expected_retry_mode
+
+    with patch("botocore.client.ClientCreator.create_client", new=wrapper):
+        with open_s3_object(path, mode="wb") as s3obj:
+            s3obj.write(b"foo")
+
+    del os.environ["AWS_MAX_ATTEMPTS"]
+    del os.environ["AWS_RETRY_MODE"]
+
+    # Update botocore.config.Config
+    expected_max_retries_attempt = 30
+    expected_connect_timeout = 40
+    expected_max_pool_connections = 50
+    expected_retry_mode = "legacy"
+
+    botocore_config = botocore.config.Config(
+        retries={"max_attempts": expected_max_retries_attempt, "mode": expected_retry_mode},
+        connect_timeout=expected_connect_timeout,
+        max_pool_connections=expected_max_pool_connections,
+    )
+    wr.config.botocore_config = botocore_config
+
+    with patch("botocore.client.ClientCreator.create_client", new=wrapper):
+        with open_s3_object(path, mode="wb") as s3obj:
+            s3obj.write(b"foo")
+
+    wr.config.reset()