Add endpoints urls in the global config. #418

Author: igorborgest

awswrangler/_config.py

@@ -7,7 +7,7 @@
 
 import pandas as pd
 
-from awswrangler import _utils, exceptions
+from awswrangler import exceptions
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -31,15 +31,30 @@ class _ConfigArg(NamedTuple):
     "max_cache_seconds": _ConfigArg(dtype=int, nullable=False),
     "s3_block_size": _ConfigArg(dtype=int, nullable=False, enforced=True),
     "workgroup": _ConfigArg(dtype=str, nullable=False, enforced=True),
+    # Endpoints URLs
+    "s3_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
+    "athena_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
+    "sts_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
+    "glue_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
+    "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),
 }
 
 
-class _Config:
+class _Config:  # pylint: disable=too-many-instance-attributes
     """Wrangler's Configuration class."""
 
     def __init__(self) -> None:
         self._loaded_values: Dict[str, _ConfigValueType] = {}
         name: str
+        self.s3_endpoint_url = None
+        self.athena_endpoint_url = None
+        self.sts_endpoint_url = None
+        self.glue_endpoint_url = None
+        self.redshift_endpoint_url = None
+        self.kms_endpoint_url = None
+        self.emr_endpoint_url = None
         for name in _CONFIG_ARGS:
             self._load_config(name=name)
 
@@ -125,7 +140,10 @@ def __getitem__(self, item: str) -> _ConfigValueType:
 
     def _reset_item(self, item: str) -> None:
         if item in self._loaded_values:
-            del self._loaded_values[item]
+            if item.endswith("_endpoint_url"):
+                self._loaded_values[item] = None
+            else:
+                del self._loaded_values[item]
         self._load_config(name=item)
 
     def _repr_html_(self) -> Any:
@@ -224,6 +242,75 @@ def workgroup(self) -> Optional[str]:
     def workgroup(self, value: Optional[str]) -> None:
         self._set_config_value(key="workgroup", value=value)
 
+    @property
+    def s3_endpoint_url(self) -> Optional[str]:
+        """Property s3_endpoint_url."""
+        return cast(Optional[str], self["s3_endpoint_url"])
+
+    @s3_endpoint_url.setter
+    def s3_endpoint_url(self, value: Optional[str]) -> None:
+        self._set_config_value(key="s3_endpoint_url", value=value)
+
+    @property
+    def athena_endpoint_url(self) -> Optional[str]:
+        """Property athena_endpoint_url."""
+        return cast(Optional[str], self["athena_endpoint_url"])
+
+    @athena_endpoint_url.setter
+    def athena_endpoint_url(self, value: Optional[str]) -> None:
+        self._set_config_value(key="athena_endpoint_url", value=value)
+
+    @property
+    def sts_endpoint_url(self) -> Optional[str]:
+        """Property sts_endpoint_url."""
+        return cast(Optional[str], self["sts_endpoint_url"])
+
+    @sts_endpoint_url.setter
+    def sts_endpoint_url(self, value: Optional[str]) -> None:
+        self._set_config_value(key="sts_endpoint_url", value=value)
+
+    @property
+    def glue_endpoint_url(self) -> Optional[str]:
+        """Property glue_endpoint_url."""
+        return cast(Optional[str], self["glue_endpoint_url"])
+
+    @glue_endpoint_url.setter
+    def glue_endpoint_url(self, value: Optional[str]) -> None:
+        self._set_config_value(key="glue_endpoint_url", value=value)
+
+    @property
+    def redshift_endpoint_url(self) -> Optional[str]:
+        """Property redshift_endpoint_url."""
+        return cast(Optional[str], self["redshift_endpoint_url"])
+
+    @redshift_endpoint_url.setter
+    def redshift_endpoint_url(self, value: Optional[str]) -> None:
+        self._set_config_value(key="redshift_endpoint_url", value=value)
+
+    @property
+    def kms_endpoint_url(self) -> Optional[str]:
+        """Property kms_endpoint_url."""
+        return cast(Optional[str], self["kms_endpoint_url"])
+
+    @kms_endpoint_url.setter
+    def kms_endpoint_url(self, value: Optional[str]) -> None:
+        self._set_config_value(key="kms_endpoint_url", value=value)
+
+    @property
+    def emr_endpoint_url(self) -> Optional[str]:
+        """Property emr_endpoint_url."""
+        return cast(Optional[str], self["emr_endpoint_url"])
+
+    @emr_endpoint_url.setter
+    def emr_endpoint_url(self, value: Optional[str]) -> None:
+        self._set_config_value(key="emr_endpoint_url", value=value)
+
+
+def _insert_str(text: str, token: str, insert: str) -> str:
+    """Insert string into other."""
+    index: int = text.find(token)
+    return text[:index] + insert + text[index:]
+
 
 def _inject_config_doc(doc: Optional[str], available_configs: Tuple[str, ...]) -> str:
     if doc is None:
@@ -244,7 +331,7 @@ def _inject_config_doc(doc: Optional[str], available_configs: Tuple[str, ...]) -
         " for details.\n"
     )
     insertion: str = header + args_block + footer + "\n\n"
-    return _utils.insert_str(text=doc, token="\n    Parameters", insert=insertion)
+    return _insert_str(text=doc, token="\n    Parameters", insert=insertion)
 
 
 def apply_configs(function: Callable[..., Any]) -> Callable[..., Any]:

awswrangler/_utils.py

@@ -16,7 +16,7 @@
 import pandas as pd
 import psycopg2
 
-from awswrangler import exceptions
+from awswrangler import _config, exceptions
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -65,14 +65,39 @@ def botocore_config() -> botocore.config.Config:
     return botocore.config.Config(retries={"max_attempts": 5}, connect_timeout=10, max_pool_connections=10)
 
 
+def _get_endpoint_url(service_name: str) -> Optional[str]:
+    endpoint_url: Optional[str] = None
+    if service_name == "s3" and _config.config.s3_endpoint_url is not None:
+        endpoint_url = _config.config.s3_endpoint_url
+    elif service_name == "athena" and _config.config.athena_endpoint_url is not None:
+        endpoint_url = _config.config.athena_endpoint_url
+    elif service_name == "sts" and _config.config.sts_endpoint_url is not None:
+        endpoint_url = _config.config.sts_endpoint_url
+    elif service_name == "glue" and _config.config.glue_endpoint_url is not None:
+        endpoint_url = _config.config.glue_endpoint_url
+    elif service_name == "redshift" and _config.config.redshift_endpoint_url is not None:
+        endpoint_url = _config.config.redshift_endpoint_url
+    elif service_name == "kms" and _config.config.kms_endpoint_url is not None:
+        endpoint_url = _config.config.kms_endpoint_url
+    elif service_name == "emr" and _config.config.emr_endpoint_url is not None:
+        endpoint_url = _config.config.emr_endpoint_url
+    return endpoint_url
+
+
 def client(service_name: str, session: Optional[boto3.Session] = None) -> boto3.client:
     """Create a valid boto3.client."""
-    return ensure_session(session=session).client(service_name=service_name, use_ssl=True, config=botocore_config())
+    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()
+    )
 
 
 def resource(service_name: str, session: Optional[boto3.Session] = None) -> boto3.resource:
     """Create a valid boto3.resource."""
-    return ensure_session(session=session).resource(service_name=service_name, use_ssl=True, config=botocore_config())
+    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()
+    )
 
 
 def parse_path(path: str) -> Tuple[str, str]:
@@ -238,12 +263,6 @@ def ensure_df_is_mutable(df: pd.DataFrame) -> pd.DataFrame:
     return df
 
 
-def insert_str(text: str, token: str, insert: str) -> str:
-    """Insert string into other."""
-    index: int = text.find(token)
-    return text[:index] + insert + text[index:]
-
-
 def check_duplicated_columns(df: pd.DataFrame) -> Any:
     """Raise an exception if there are duplicated columns names."""
     duplicated: List[str] = df.loc[:, df.columns.duplicated()].columns.to_list()

awswrangler/athena/_utils.py

@@ -305,7 +305,7 @@ def create_athena_bucket(boto3_session: Optional[boto3.Session] = None) -> str:
     account_id: str = sts.get_account_id(boto3_session=session)
     region_name: str = str(session.region_name).lower()
     s3_output = f"s3://aws-athena-query-results-{account_id}-{region_name}/"
-    s3_resource = session.resource("s3")
+    s3_resource = _utils.resource(service_name="s3", session=session)
     s3_resource.Bucket(s3_output)
     return s3_output
 

tests/test_config.py

@@ -1,6 +1,9 @@
 import logging
 import os
+from unittest.mock import patch
 
+import boto3
+import botocore
 import pytest
 
 import awswrangler as wr
@@ -9,6 +12,26 @@
 logging.getLogger("awswrangler").setLevel(logging.DEBUG)
 
 
+def _urls_test(glue_database):
+    original = botocore.client.ClientCreator.create_client
+
+    def wrapper(self, **kwarg):
+        name = kwarg["service_name"]
+        url = kwarg["endpoint_url"]
+        if name == "sts":
+            assert url == wr.config.sts_endpoint_url
+        elif name == "athena":
+            assert url == wr.config.athena_endpoint_url
+        elif name == "s3":
+            assert url == wr.config.s3_endpoint_url
+        elif name == "glue":
+            assert url == wr.config.glue_endpoint_url
+        return original(self, **kwarg)
+
+    with patch("botocore.client.ClientCreator.create_client", new=wrapper):
+        wr.athena.read_sql_query(sql="SELECT 1 as col0", database=glue_database)
+
+
 def test_basics(path, glue_database, glue_table, workgroup0, workgroup1):
     args = {"table": glue_table, "path": "", "columns_types": {"col0": "bigint"}}
 
@@ -71,3 +94,17 @@ def test_basics(path, glue_database, glue_table, workgroup0, workgroup1):
     wr.config.reset()
     df = wr.athena.read_sql_query(sql="SELECT 1 as col0", database=glue_database)
     assert df.query_metadata["WorkGroup"] == workgroup1
+
+    # Endpoints URLs
+    region = boto3.Session().region_name
+    wr.config.sts_endpoint_url = f"https://sts.{region}.amazonaws.com"
+    wr.config.s3_endpoint_url = f"https://s3.{region}.amazonaws.com"
+    wr.config.athena_endpoint_url = f"https://athena.{region}.amazonaws.com"
+    wr.config.glue_endpoint_url = f"https://glue.{region}.amazonaws.com"
+    _urls_test(glue_database)
+    os.environ["WR_STS_ENDPOINT_URL"] = f"https://sts.{region}.amazonaws.com"
+    os.environ["WR_S3_ENDPOINT_URL"] = f"https://s3.{region}.amazonaws.com"
+    os.environ["WR_ATHENA_ENDPOINT_URL"] = f"https://athena.{region}.amazonaws.com"
+    os.environ["WR_GLUE_ENDPOINT_URL"] = f"https://glue.{region}.amazonaws.com"
+    wr.config.reset()
+    _urls_test(glue_database)