Bug: Add correct service names for timestream boto3 clients (#1716)

* add correct service names for timestream boto3 clients

* format

* pylint

* account for write and query endpoints

Co-authored-by: kukushking <[email protected]>

Author: malachi-constant

awswrangler/_config.py

@@ -46,7 +46,8 @@ class _ConfigArg(NamedTuple):
     "lakeformation_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
     "dynamodb_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
     "secretsmanager_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
-    "timestream_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
+    "timestream_query_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
+    "timestream_write_endpoint_url": _ConfigArg(dtype=str, nullable=True, enforced=True),
     # Botocore config
     "botocore_config": _ConfigArg(dtype=botocore.config.Config, nullable=True),
     "verify": _ConfigArg(dtype=str, nullable=True),
@@ -69,7 +70,8 @@ def __init__(self) -> None:
         self.lakeformation_endpoint_url = None
         self.dynamodb_endpoint_url = None
         self.secretsmanager_endpoint_url = None
-        self.timestream_endpoint_url = None
+        self.timestream_write_endpoint_url = None
+        self.timestream_query_endpoint_url = None
         self.botocore_config = None
         self.verify = None
         for name in _CONFIG_ARGS:
@@ -390,13 +392,22 @@ def secretsmanager_endpoint_url(self, value: Optional[str]) -> None:
         self._set_config_value(key="secretsmanager_endpoint_url", value=value)
 
     @property
-    def timestream_endpoint_url(self) -> Optional[str]:
-        """Property timestream_endpoint_url."""
-        return cast(Optional[str], self["timestream_endpoint_url"])
+    def timestream_query_endpoint_url(self) -> Optional[str]:
+        """Property timestream_query_endpoint_url."""
+        return cast(Optional[str], self["timestream_query_endpoint_url"])
 
-    @timestream_endpoint_url.setter
-    def timestream_endpoint_url(self, value: Optional[str]) -> None:
-        self._set_config_value(key="timestream_endpoint_url", value=value)
+    @timestream_query_endpoint_url.setter
+    def timestream_query_endpoint_url(self, value: Optional[str]) -> None:
+        self._set_config_value(key="timestream_query_endpoint_url", value=value)
+
+    @property
+    def timestream_write_endpoint_url(self) -> Optional[str]:
+        """Property timestream_write_endpoint_url."""
+        return cast(Optional[str], self["timestream_write_endpoint_url"])
+
+    @timestream_write_endpoint_url.setter
+    def timestream_write_endpoint_url(self, value: Optional[str]) -> None:
+        self._set_config_value(key="timestream_write_endpoint_url", value=value)
 
     @property
     def botocore_config(self) -> botocore.config.Config:

awswrangler/_utils.py

@@ -99,8 +99,11 @@ def _get_endpoint_url(service_name: str) -> Optional[str]:
         endpoint_url = _config.config.dynamodb_endpoint_url
     elif service_name == "secretsmanager" and _config.config.secretsmanager_endpoint_url is not None:
         endpoint_url = _config.config.secretsmanager_endpoint_url
-    elif service_name == "timestream" and _config.config.timestream_endpoint_url is not None:
-        endpoint_url = _config.config.timestream_endpoint_url
+    elif service_name == "timestream-write" and _config.config.timestream_write_endpoint_url is not None:
+        endpoint_url = _config.config.timestream_write_endpoint_url
+    elif service_name == "timestream-query" and _config.config.timestream_query_endpoint_url is not None:
+        endpoint_url = _config.config.timestream_query_endpoint_url
+
     return endpoint_url
 
 

tests/test_config.py

@@ -31,8 +31,10 @@ def wrapper(self, **kwarg):
             assert url == wr.config.glue_endpoint_url
         elif name == "secretsmanager":
             assert url == wr.config.secretsmanager_endpoint_url
-        elif name == "timestream":
-            assert url == wr.config.timestream_endpoint_url
+        elif name == "timestream-write":
+            assert url == wr.config.timestream_write_endpoint_url
+        elif name == "timestream-query":
+            assert url == wr.config.timestream_query_endpoint_url
         return original(self, **kwarg)
 
     with patch("botocore.client.ClientCreator.create_client", new=wrapper):
@@ -117,14 +119,16 @@ def test_basics(path, glue_database, glue_table, workgroup0, workgroup1):
     wr.config.athena_endpoint_url = f"https://athena.{region}.amazonaws.com"
     wr.config.glue_endpoint_url = f"https://glue.{region}.amazonaws.com"
     wr.config.secretsmanager_endpoint_url = f"https://secretsmanager.{region}.amazonaws.com"
-    wr.config.timestream_endpoint_url = f"https://timestream.{region}.amazonaws.com"
+    wr.config.timestream_write_endpoint_url = f"https://ingest.timestream.{region}.amazonaws.com"
+    wr.config.timestream_query_endpoint_url = f"https://query.timestream.{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"
     os.environ["WR_SECRETSMANAGER_ENDPOINT_URL"] = f"https://secretsmanager.{region}.amazonaws.com"
-    os.environ["WR_TIMESTREAM_ENDPOINT_URL"] = f"https://timestream.{region}.amazonaws.com"
+    os.environ["WR_TIMESTREAM_WRITE_ENDPOINT_URL"] = f"https://ingest.timestream.{region}.amazonaws.com"
+    os.environ["WR_TIMESTREAM_QUERY_ENDPOINT_URL"] = f"https://query.timestream.{region}.amazonaws.com"
     wr.config.reset()
     _urls_test(glue_database)