Add formatter for PartiQL

Author: LeonLuttenberger

awswrangler/_sql_formatter.py

@@ -9,6 +9,7 @@
 class _EngineType(Enum):
     PRESTO = "presto"
     HIVE = "hive"
+    PARTIQL = "partiql"
 
     def __str__(self) -> str:
         return self.value
@@ -30,14 +31,17 @@ def __str__(self) -> str:
 
 class _NullType(_AbstractType[_NoneType]):
     def __str__(self) -> str:
+        if self.engine == _EngineType.PARTIQL:
+            return "null"
+
         return "NULL"
 
 
 class _StringType(_AbstractType[str]):
     supported_formats = {"s", "i"}
 
     def __str__(self) -> str:
-        if self.engine == _EngineType.PRESTO:
+        if self.engine in [_EngineType.PRESTO, _EngineType.PARTIQL]:
             return f"""'{self.data.replace("'", "''")}'"""
 
         if self.engine == _EngineType.HIVE:
@@ -54,6 +58,9 @@ def __str__(self) -> str:
 
 class _BooleanType(_AbstractType[bool]):
     def __str__(self) -> str:
+        if self.engine == _EngineType.PARTIQL:
+            return "1" if self.data else "0"
+
         return str(self.data).upper()
 
 
@@ -69,28 +76,44 @@ def __str__(self) -> str:
 
 class _DecimalType(_AbstractType[decimal.Decimal]):
     def __str__(self) -> str:
+        if self.engine == _EngineType.PARTIQL:
+            return f"'{self.data}'"
+
         return f"DECIMAL '{self.data:f}'"
 
 
 class _TimestampType(_AbstractType[datetime.datetime]):
     def __str__(self) -> str:
         if self.data.tzinfo is not None:
             raise TypeError(f"Supports only timezone aware datatype, got {self.data}.")
+
+        if self.engine == _EngineType.PARTIQL:
+            return f"'{self.data.isoformat()}'"
+
         return f"TIMESTAMP '{self.data.isoformat(sep=' ', timespec='milliseconds')}'"
 
 
 class _DateType(_AbstractType[datetime.date]):
     def __str__(self) -> str:
+        if self.engine == _EngineType.PARTIQL:
+            return f"'{self.data.isoformat()}'"
+
         return f"DATE '{self.data.isoformat()}'"
 
 
 class _ArrayType(_AbstractType[Sequence[_PythonType]]):
     def __str__(self) -> str:
+        if self.engine == _EngineType.PARTIQL:
+            super().__str__()
+
         return f"ARRAY [{', '.join(map(str, self.data))}]"
 
 
 class _MapType(_AbstractType[Dict[_PythonType, _PythonTypeMapValue]]):
     def __str__(self) -> str:
+        if self.engine == _EngineType.PARTIQL:
+            super().__str__()
+
         if not self.data:
             return "MAP()"
 

awswrangler/lakeformation/_read.py

@@ -10,7 +10,7 @@
 from awswrangler import _data_types, _utils, catalog
 from awswrangler._config import apply_configs
 from awswrangler._distributed import engine
-from awswrangler._sql_formatter import _process_sql_params
+from awswrangler._sql_formatter import _EngineType, _process_sql_params
 from awswrangler._threading import _get_executor
 from awswrangler.catalog._utils import _catalog_id, _transaction_id
 from awswrangler.distributed.ray import RayLogger
@@ -168,7 +168,7 @@ def read_sql_query(
     client_lakeformation: boto3.client = _utils.client(service_name="lakeformation", session=session)
     commit_trans: bool = False
 
-    sql = _process_sql_params(sql, params)
+    sql = _process_sql_params(sql, params, engine=_EngineType.PARTIQL)
 
     if not any([transaction_id, query_as_of_time]):
         _logger.debug("Neither `transaction_id` nor `query_as_of_time` were specified, starting transaction")

tests/_utils.py

@@ -105,7 +105,7 @@ def get_df_list(governed=False):
     df["category"] = df["category"].astype("category")
 
     if governed:
-        df = (df.drop(["iint8", "binary"], axis=1),)  # tinyint & binary currently not supported
+        df = df.drop(["iint8", "binary"], axis=1)  # tinyint & binary currently not supported
     return df
 
 

tests/unit/test_lakeformation.py

@@ -1,13 +1,15 @@
 import calendar
+import datetime as dt
 import logging
 import time
+from decimal import Decimal
 
 import pytest
 
 import awswrangler as wr
 from awswrangler._distributed import EngineEnum, MemoryFormatEnum
 
-from .._utils import ensure_data_types, ensure_data_types_csv, get_df, get_df_csv
+from .._utils import ensure_data_types, ensure_data_types_csv, get_df, get_df_csv, get_df_list
 
 if wr.engine.get() == EngineEnum.RAY and wr.memory_format.get() == MemoryFormatEnum.MODIN:
     import modin.pandas as pd
@@ -145,3 +147,71 @@ def test_lakeformation_multi_transaction(path, path2, glue_database, glue_table,
 
     assert df2.shape == df4.shape
     assert df2.c1.sum() == df4.c1.sum()
+
+
+@pytest.mark.parametrize(
+    "col_name,col_value",
+    [
+        ("date", dt.date(2020, 1, 1)),
+        ("timestamp", dt.datetime(2020, 1, 1)),
+        ("bool", True),
+        ("decimal", Decimal(("1.99"))),
+        ("float", 0.0),
+        ("iint16", 1),
+    ],
+)
+def test_lakeformation_partiql_formatting(path, path2, glue_database, glue_table, glue_table2, col_name, col_value):
+    wr.catalog.delete_table_if_exists(database=glue_database, table=glue_table)
+    wr.catalog.delete_table_if_exists(database=glue_database, table=glue_table2)
+
+    wr.s3.to_parquet(
+        df=get_df_list(governed=True),
+        path=path,
+        index=False,
+        boto3_session=None,
+        s3_additional_kwargs=None,
+        dataset=True,
+        partition_cols=["par0", "par1"],
+        mode="overwrite",
+        table=glue_table,
+        table_type="GOVERNED",
+        database=glue_database,
+    )
+
+    # Filter query
+    df = wr.lakeformation.read_sql_query(
+        sql=f'SELECT * FROM {glue_table} WHERE "{col_name}" = :col_value',
+        database=glue_database,
+        params={"col_value": col_value},
+    )
+    assert len(df) == 1
+
+
+def test_lakeformation_partiql_formatting_escape_string(path, path2, glue_database, glue_table, glue_table2):
+    df = pd.DataFrame(
+        {
+            "id": [1, 2, 3],
+            "string": ["normal string", "'weird' string", "another normal string"],
+        }
+    )
+
+    wr.s3.to_parquet(
+        df=df,
+        path=path,
+        index=False,
+        boto3_session=None,
+        s3_additional_kwargs=None,
+        dataset=True,
+        mode="overwrite",
+        table=glue_table,
+        table_type="GOVERNED",
+        database=glue_database,
+    )
+
+    # Filter query
+    df = wr.lakeformation.read_sql_query(
+        sql=f'SELECT * FROM {glue_table} WHERE "string" = :col_value',
+        database=glue_database,
+        params={"col_value": "'weird' string"},
+    )
+    assert len(df) == 1