Support additional sql params in Redshift COPY operation (#1210)

* adding argument 'sql_copy_extra_params'

* checking var before iterating

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

Author: malachi-constant

awswrangler/redshift.py

@@ -130,6 +130,7 @@ def _copy(
     boto3_session: Optional[str] = None,
     schema: Optional[str] = None,
     manifest: Optional[bool] = False,
+    sql_copy_extra_params: Optional[List[str]] = None,
 ) -> None:
     if schema is None:
         table_name: str = f'"{table}"'
@@ -147,6 +148,9 @@ def _copy(
     sql: str = f"COPY {table_name}\nFROM '{path}' {auth_str}\nFORMAT AS PARQUET{ser_json_str}"
     if manifest:
         sql += "\nMANIFEST"
+    if sql_copy_extra_params:
+        for param in sql_copy_extra_params:
+            sql += f"\n{param}"
     _logger.debug("copy query:\n%s", sql)
     cursor.execute(sql)
 
@@ -1199,6 +1203,7 @@ def copy_from_files(  # pylint: disable=too-many-locals,too-many-arguments
     lock: bool = False,
     commit_transaction: bool = True,
     manifest: Optional[bool] = False,
+    sql_copy_extra_params: Optional[List[str]] = None,
     boto3_session: Optional[boto3.Session] = None,
     s3_additional_kwargs: Optional[Dict[str, str]] = None,
 ) -> None:
@@ -1361,6 +1366,7 @@ def copy_from_files(  # pylint: disable=too-many-locals,too-many-arguments
                 aws_session_token=aws_session_token,
                 boto3_session=boto3_session,
                 serialize_to_json=serialize_to_json,
+                sql_copy_extra_params=sql_copy_extra_params,
                 manifest=manifest,
             )
             if table != created_table:  # upsert

tests/test_redshift.py

@@ -732,6 +732,24 @@ def test_copy_from_files(path, redshift_table, redshift_con, databases_parameter
     assert df2["counter"].iloc[0] == 3
 
 
+def test_copy_from_files_extra_params(path, redshift_table, redshift_con, databases_parameters):
+    df = get_df_category().drop(["binary"], axis=1, inplace=False)
+    wr.s3.to_parquet(df, f"{path}test.parquet")
+    bucket, key = wr._utils.parse_path(f"{path}test.csv")
+    boto3.client("s3").put_object(Body=b"", Bucket=bucket, Key=key)
+    wr.redshift.copy_from_files(
+        path=path,
+        path_suffix=[".parquet"],
+        con=redshift_con,
+        table=redshift_table,
+        schema="public",
+        iam_role=databases_parameters["redshift"]["role"],
+        sql_copy_extra_params=["STATUPDATE ON"],
+    )
+    df2 = wr.redshift.read_sql_query(sql=f"SELECT count(*) AS counter FROM public.{redshift_table}", con=redshift_con)
+    assert df2["counter"].iloc[0] == 3
+
+
 def test_get_paths_from_manifest(path):
     manifest_content = {
         "entries": [