fix: escape column names in PRIMARY KEY statement in SQL query (#2351)

* escape column names in PRIMARY KEY statement in SQL query

* add test for writing to redshift with hyphenated primary key

---------

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

Author: 3 people

awswrangler/redshift/_utils.py

@@ -349,7 +349,9 @@ def _create_table(  # pylint: disable=too-many-locals,too-many-arguments,too-man
         primary_keys=primary_keys,
     )
     cols_str: str = "".join([f'"{k}" {v},\n' for k, v in redshift_types.items()])[:-2]
-    primary_keys_str: str = f",\nPRIMARY KEY ({', '.join(primary_keys)})" if primary_keys else ""
+    primary_keys_str: str = (
+        ",\nPRIMARY KEY ({})".format(", ".join('"' + pk + '"' for pk in primary_keys)) if primary_keys else ""
+    )
     distkey_str: str = f"\nDISTKEY({distkey})" if distkey and diststyle == "KEY" else ""
     sortkey_str: str = f"\n{sortstyle} SORTKEY({','.join(sortkey)})" if sortkey else ""
     sql = (

tests/unit/test_redshift.py

@@ -17,6 +17,7 @@
 from awswrangler import _utils
 
 from .._utils import (
+    assert_pandas_equals,
     dt,
     ensure_data_types,
     ensure_data_types_category,
@@ -63,6 +64,21 @@ def test_to_sql_simple(redshift_table: str, redshift_con: redshift_connector.Con
     wr.redshift.to_sql(df, redshift_con, redshift_table, "public", "overwrite", overwrite_method, True)
 
 
+def test_to_sql_with_hyphenated_primary_key(
+    redshift_table: str,
+    redshift_con: redshift_connector.Connection,
+) -> None:
+    schema = "public"
+    df = pd.DataFrame({"id-col": [1, 2, 3], "other-col": ["foo", "boo", "bar"]})
+    df["id-col"] = df["id-col"].astype("Int64")
+    df["other-col"] = df["other-col"].astype("string")
+    wr.redshift.to_sql(
+        df=df, con=redshift_con, table=redshift_table, schema=schema, mode="overwrite", primary_keys=["id-col"]
+    )
+    df_out = wr.redshift.read_sql_table(table=redshift_table, con=redshift_con, schema=schema)
+    assert_pandas_equals(df, df_out)
+
+
 def test_empty_table(redshift_table: str, redshift_con: redshift_connector.Connection) -> None:
     with redshift_con.cursor() as cursor:
         cursor.execute(f"CREATE TABLE public.{redshift_table}(c0 integer not null, c1 integer, primary key(c0));")