Add Amazon Timestream support.

Author: igorborgest

awswrangler/__init__.py

@@ -19,6 +19,7 @@
     redshift,
     s3,
     sts,
+    timestream,
 )
 from awswrangler.__metadata__ import __description__, __license__, __title__, __version__  # noqa
 from awswrangler._config import config  # noqa
@@ -36,6 +37,7 @@
     "mysql",
     "postgresql",
     "config",
+    "timestream",
     "__description__",
     "__license__",
     "__title__",

awswrangler/_data_types.py

@@ -166,6 +166,29 @@ def pyarrow2postgresql(  # pylint: disable=too-many-branches,too-many-return-sta
     raise exceptions.UnsupportedType(f"Unsupported PostgreSQL type: {dtype}")
 
 
+def pyarrow2timestream(dtype: pa.DataType) -> str:  # pylint: disable=too-many-branches,too-many-return-statements
+    """Pyarrow to Amazon Timestream data types conversion."""
+    if pa.types.is_int8(dtype):
+        return "BIGINT"
+    if pa.types.is_int16(dtype) or pa.types.is_uint8(dtype):
+        return "BIGINT"
+    if pa.types.is_int32(dtype) or pa.types.is_uint16(dtype):
+        return "BIGINT"
+    if pa.types.is_int64(dtype) or pa.types.is_uint32(dtype):
+        return "BIGINT"
+    if pa.types.is_uint64(dtype):
+        return "BIGINT"
+    if pa.types.is_float32(dtype):
+        return "DOUBLE"
+    if pa.types.is_float64(dtype):
+        return "DOUBLE"
+    if pa.types.is_boolean(dtype):
+        return "BOOLEAN"
+    if pa.types.is_string(dtype):
+        return "VARCHAR"
+    raise exceptions.UnsupportedType(f"Unsupported Amazon Timestream measure type: {dtype}")
+
+
 def athena2pyarrow(dtype: str) -> pa.DataType:  # pylint: disable=too-many-return-statements
     """Athena to PyArrow data types conversion."""
     dtype = dtype.lower().replace(" ", "")
@@ -587,3 +610,13 @@ def database_types_from_pandas(
             database_types[col_name] = converter_func(col_dtype, string_type)
     _logger.debug("database_types: %s", database_types)
     return database_types
+
+
+def timestream_type_from_pandas(df: pd.DataFrame) -> str:
+    """Extract Amazon Timestream types from a Pandas DataFrame."""
+    pyarrow_types: Dict[str, Optional[pa.DataType]] = pyarrow_types_from_pandas(df=df, index=False, ignore_cols=[])
+    if len(pyarrow_types) != 1 or list(pyarrow_types.values())[0] is None:
+        raise RuntimeError(f"Invalid pyarrow_types: {pyarrow_types}")
+    pyarrow_type: pa.DataType = list(pyarrow_types.values())[0]
+    _logger.debug("pyarrow_type: %s", pyarrow_type)
+    return pyarrow2timestream(dtype=pyarrow_type)

awswrangler/_databases.py

@@ -133,3 +133,15 @@ def read_sql_query(
         return _iterate_cursor(
             cursor=cursor, chunksize=chunksize, cols_names=cols_names, index=index_col, dtype=dtype, safe=safe
         )
+
+
+def extract_parameters(df: pd.DataFrame) -> List[List[Any]]:
+    """Extract Parameters."""
+    parameters: List[List[Any]] = df.values.tolist()
+    for i, row in enumerate(parameters):
+        for j, value in enumerate(row):
+            if pd.isna(value):
+                parameters[i][j] = None
+            elif hasattr(value, "to_pydatetime"):
+                parameters[i][j] = value.to_pydatetime()
+    return parameters

awswrangler/_utils.py

@@ -83,11 +83,16 @@ def _get_endpoint_url(service_name: str) -> Optional[str]:
     return endpoint_url
 
 
-def client(service_name: str, session: Optional[boto3.Session] = None) -> boto3.client:
+def client(
+    service_name: str, session: Optional[boto3.Session] = None, 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()
+        service_name=service_name,
+        endpoint_url=endpoint_url,
+        use_ssl=True,
+        config=botocore_config() if config is None else config,
     )
 
 

awswrangler/mysql.py

@@ -66,17 +66,6 @@ def _create_table(
     cursor.execute(sql)
 
 
-def _extract_parameters(df: pd.DataFrame) -> List[List[Any]]:
-    parameters: List[List[Any]] = df.values.tolist()
-    for i, row in enumerate(parameters):
-        for j, value in enumerate(row):
-            if pd.isna(value):
-                parameters[i][j] = None
-            elif hasattr(value, "to_pydatetime"):
-                parameters[i][j] = value.to_pydatetime()
-    return parameters
-
-
 def connect(
     connection: str,
     catalog_id: Optional[str] = None,
@@ -339,7 +328,7 @@ def to_sql(
             placeholders: str = ", ".join(["%s"] * len(df.columns))
             sql: str = f"INSERT INTO {schema}.{table} VALUES ({placeholders})"
             _logger.debug("sql: %s", sql)
-            parameters: List[List[Any]] = _extract_parameters(df=df)
+            parameters: List[List[Any]] = _db_utils.extract_parameters(df=df)
             cursor.executemany(sql, parameters)
             con.commit()  # type: ignore
     except Exception as ex:

awswrangler/postgresql.py

@@ -70,17 +70,6 @@ def _create_table(
     cursor.execute(sql)
 
 
-def _extract_parameters(df: pd.DataFrame) -> List[List[Any]]:
-    parameters: List[List[Any]] = df.values.tolist()
-    for i, row in enumerate(parameters):
-        for j, value in enumerate(row):
-            if pd.isna(value):
-                parameters[i][j] = None
-            elif hasattr(value, "to_pydatetime"):
-                parameters[i][j] = value.to_pydatetime()
-    return parameters
-
-
 def connect(
     connection: str,
     catalog_id: Optional[str] = None,
@@ -343,7 +332,7 @@ def to_sql(
             placeholders: str = ", ".join(["%s"] * len(df.columns))
             sql: str = f"INSERT INTO {schema}.{table} VALUES ({placeholders})"
             _logger.debug("sql: %s", sql)
-            parameters: List[List[Any]] = _extract_parameters(df=df)
+            parameters: List[List[Any]] = _db_utils.extract_parameters(df=df)
             cursor.executemany(sql, parameters)
             con.commit()
     except Exception as ex:

awswrangler/redshift.py

@@ -229,17 +229,6 @@ def _create_table(
     return table, schema
 
 
-def _extract_parameters(df: pd.DataFrame) -> List[List[Any]]:
-    parameters: List[List[Any]] = df.values.tolist()
-    for i, row in enumerate(parameters):
-        for j, value in enumerate(row):
-            if pd.isna(value):
-                parameters[i][j] = None
-            elif hasattr(value, "to_pydatetime"):
-                parameters[i][j] = value.to_pydatetime()
-    return parameters
-
-
 def _read_parquet_iterator(
     path: str,
     keep_files: bool,
@@ -664,7 +653,7 @@ def to_sql(
             schema_str = f"{created_schema}." if created_schema else ""
             sql: str = f"INSERT INTO {schema_str}{created_table} VALUES ({placeholders})"
             _logger.debug("sql: %s", sql)
-            parameters: List[List[Any]] = _extract_parameters(df=df)
+            parameters: List[List[Any]] = _db_utils.extract_parameters(df=df)
             cursor.executemany(sql, parameters)
             if table != created_table:  # upsert
                 _upsert(cursor=cursor, schema=schema, table=table, temp_table=created_table, primary_keys=primary_keys)