Fix Redshift unload breaking with hyphenated table names (#1762)

Author: LeonLuttenberger

awswrangler/redshift.py

@@ -184,30 +184,32 @@ def _upsert(
     _logger.debug("primary_keys: %s", primary_keys)
     if not primary_keys:
         raise exceptions.InvalidRedshiftPrimaryKeys()
-    equals_clause: str = f"{table}.%s = {temp_table}.%s"
+    equals_clause: str = f'"{table}".%s = "{temp_table}".%s'
     join_clause: str = " AND ".join([equals_clause % (pk, pk) for pk in primary_keys])
     if precombine_key:
         delete_from_target_filter: str = f"AND {table}.{precombine_key} <= {temp_table}.{precombine_key}"
         delete_from_temp_filter: str = f"AND {table}.{precombine_key} > {temp_table}.{precombine_key}"
         target_del_sql: str = (
-            f'DELETE FROM "{schema}"."{table}" USING {temp_table} WHERE {join_clause} {delete_from_target_filter}'
+            f'DELETE FROM "{schema}"."{table}" USING "{temp_table}" WHERE {join_clause} {delete_from_target_filter}'
         )
         _logger.debug(target_del_sql)
         cursor.execute(target_del_sql)
         source_del_sql: str = (
-            f'DELETE FROM {temp_table} USING "{schema}"."{table}" WHERE {join_clause} {delete_from_temp_filter}'
+            f'DELETE FROM "{temp_table}" USING "{schema}"."{table}" WHERE {join_clause} {delete_from_temp_filter}'
         )
         _logger.debug(source_del_sql)
         cursor.execute(source_del_sql)
     else:
-        sql: str = f'DELETE FROM "{schema}"."{table}" USING {temp_table} WHERE {join_clause}'
+        sql: str = f'DELETE FROM "{schema}"."{table}" USING "{temp_table}" WHERE {join_clause}'
         _logger.debug(sql)
         cursor.execute(sql)
     if column_names:
         column_names_str = ",".join(column_names)
-        insert_sql = f'INSERT INTO "{schema}"."{table}"({column_names_str}) SELECT {column_names_str} FROM {temp_table}'
+        insert_sql = (
+            f'INSERT INTO "{schema}"."{table}"({column_names_str}) SELECT {column_names_str} FROM "{temp_table}"'
+        )
     else:
-        insert_sql = f'INSERT INTO "{schema}"."{table}" SELECT * FROM {temp_table}'
+        insert_sql = f'INSERT INTO "{schema}"."{table}" SELECT * FROM "{temp_table}"'
     _logger.debug(insert_sql)
     cursor.execute(insert_sql)
     _drop_table(cursor=cursor, schema=schema, table=temp_table)

tests/conftest.py

@@ -251,6 +251,18 @@ def redshift_table():
     con.close()
 
 
+@pytest.fixture(scope="function")
+def redshift_table_with_hyphenated_name():
+    name = f"tbl-{get_time_str_with_random_suffix()}"
+    print(f"Table name: {name}")
+    yield name
+    con = wr.redshift.connect("aws-sdk-pandas-redshift")
+    with con.cursor() as cursor:
+        cursor.execute(f'DROP TABLE IF EXISTS public."{name}"')
+    con.commit()
+    con.close()
+
+
 @pytest.fixture(scope="function")
 def postgresql_table():
     name = f"tbl_{get_time_str_with_random_suffix()}"

tests/test_redshift.py

@@ -3,6 +3,7 @@
 import random
 import string
 from decimal import Decimal
+from typing import Any, Dict
 
 import boto3
 import numpy as np
@@ -164,7 +165,12 @@ def test_copy_unload(path, redshift_table, redshift_con, databases_parameters):
         ensure_data_types(df=chunk, has_list=False)
 
 
-def test_copy_upsert(path, redshift_table, redshift_con, databases_parameters):
+def generic_test_copy_upsert(
+    path: str,
+    redshift_table: str,
+    redshift_con: redshift_connector.Connection,
+    databases_parameters: Dict[str, Any],
+) -> None:
     df = pd.DataFrame({"id": list((range(1_000))), "val": list(["foo" if i % 2 == 0 else "boo" for i in range(1_000)])})
     df3 = pd.DataFrame(
         {"id": list((range(1_000, 1_500))), "val": list(["foo" if i % 2 == 0 else "boo" for i in range(500)])}
@@ -185,7 +191,7 @@ def test_copy_upsert(path, redshift_table, redshift_con, databases_parameters):
     )
     path = f"{path}upsert/test_redshift_copy_upsert2/"
     df2 = wr.redshift.unload(
-        sql=f"SELECT * FROM public.{redshift_table}",
+        sql=f'SELECT * FROM public."{redshift_table}"',
         con=redshift_con,
         iam_role=databases_parameters["redshift"]["role"],
         path=path,
@@ -209,7 +215,7 @@ def test_copy_upsert(path, redshift_table, redshift_con, databases_parameters):
     )
     path = f"{path}upsert/test_redshift_copy_upsert4/"
     df4 = wr.redshift.unload(
-        sql=f"SELECT * FROM public.{redshift_table}",
+        sql=f'SELECT * FROM public."{redshift_table}"',
         con=redshift_con,
         iam_role=databases_parameters["redshift"]["role"],
         path=path,
@@ -232,7 +238,7 @@ def test_copy_upsert(path, redshift_table, redshift_con, databases_parameters):
     )
     path = f"{path}upsert/test_redshift_copy_upsert4/"
     df4 = wr.redshift.unload(
-        sql=f"SELECT * FROM public.{redshift_table}",
+        sql=f'SELECT * FROM public."{redshift_table}"',
         con=redshift_con,
         iam_role=databases_parameters["redshift"]["role"],
         path=path,
@@ -242,6 +248,14 @@ def test_copy_upsert(path, redshift_table, redshift_con, databases_parameters):
     assert len(df.columns) == len(df4.columns)
 
 
+def test_copy_upsert(path, redshift_table, redshift_con, databases_parameters):
+    generic_test_copy_upsert(path, redshift_table, redshift_con, databases_parameters)
+
+
+def test_copy_upsert_hyphenated_name(path, redshift_table_with_hyphenated_name, redshift_con, databases_parameters):
+    generic_test_copy_upsert(path, redshift_table_with_hyphenated_name, redshift_con, databases_parameters)
+
+
 @pytest.mark.parametrize(
     "diststyle,distkey,exc,sortstyle,sortkey",
     [