fix: data_api uuid typehint (#2961)

* fix data_api uuid typehint * tests: add uuid case to unit test * fix: sort deps correctly --------- Co-authored-by: jaidisido <[email protected]>
aws · Sep 30, 2024 · 5da8bde · 5da8bde
1 parent 4bb8afc
commit 5da8bde
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/awswrangler/data_api/rds.py b/awswrangler/data_api/rds.py
@@ -359,6 +359,9 @@ def _create_value_dict(  # noqa: PLR0911
     if isinstance(value, Decimal):
         return {"stringValue": str(value)}, "DECIMAL"
 
+    if isinstance(value, uuid.UUID):
+        return {"stringValue": str(value)}, "UUID"
+
     raise exceptions.InvalidArgumentType(f"Value {value} not supported.")
 
 

diff --git a/tests/unit/test_data_api.py b/tests/unit/test_data_api.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import uuid
 from typing import Any, Iterator
 
 import boto3
@@ -284,7 +285,8 @@ def test_data_api_mysql_ansi(mysql_serverless_connector: "RdsDataApi", mysql_ser
 
 def test_data_api_postgresql(postgresql_serverless_connector: "RdsDataApi", postgresql_serverless_table: str) -> None:
     database = "test"
-    frame = pd.DataFrame([[42, "test"]], columns=["id", "name"])
+    _id = uuid.uuid4()
+    frame = pd.DataFrame([[_id, "test"]], columns=["id", "name"])
 
     wr.data_api.rds.to_sql(
         df=frame,
@@ -295,7 +297,7 @@ def test_data_api_postgresql(postgresql_serverless_connector: "RdsDataApi", post
     )
 
     out_frame = wr.data_api.rds.read_sql_query(
-        f"SELECT name FROM {postgresql_serverless_table} WHERE id = 42", con=postgresql_serverless_connector
+        f"SELECT name FROM {postgresql_serverless_table} WHERE id = '{_id}'", con=postgresql_serverless_connector
     )
     expected_dataframe = pd.DataFrame([["test"]], columns=["name"])
     assert_pandas_equals(out_frame, expected_dataframe)