Use DataFrame column names in INSERT statement for UPSERT operation (#1317)

* Use DataFrame column names in INSERT statement for UPSERT operation

* Make column_names as an optional parameter

* Fix input parameter type to Optional[List[str]]

* Fix Pylint issue. Replace unnecessary for-comprehension to list(df.columns)

* Reformat code

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

Author: Ildar Almakaev

awswrangler/redshift.py

@@ -174,6 +174,7 @@ def _upsert(
     schema: str,
     primary_keys: Optional[List[str]] = None,
     precombine_key: Optional[str] = None,
+    column_names: Optional[List[str]] = None,
 ) -> None:
     if not primary_keys:
         primary_keys = _get_primary_keys(cursor=cursor, schema=schema, table=table)
@@ -199,7 +200,11 @@ def _upsert(
         sql: str = f'DELETE FROM "{schema}"."{table}" USING {temp_table} WHERE {join_clause}'
         _logger.debug(sql)
         cursor.execute(sql)
-    insert_sql = f"INSERT INTO {schema}.{table} SELECT * FROM {temp_table}"
+    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}"
+    else:
+        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)
@@ -903,11 +908,12 @@ def to_sql(  # pylint: disable=too-many-locals
             )
             if index:
                 df.reset_index(level=df.index.names, inplace=True)
-            column_placeholders: str = ", ".join(["%s"] * len(df.columns))
+            column_names = list(df.columns)
+            column_placeholders: str = ", ".join(["%s"] * len(column_names))
             schema_str = f'"{created_schema}".' if created_schema else ""
             insertion_columns = ""
             if use_column_names:
-                insertion_columns = f"({', '.join(df.columns)})"
+                insertion_columns = f"({', '.join(column_names)})"
             placeholder_parameter_pair_generator = _db_utils.generate_placeholder_parameter_pairs(
                 df=df, column_placeholders=column_placeholders, chunksize=chunksize
             )
@@ -923,6 +929,7 @@ def to_sql(  # pylint: disable=too-many-locals
                     temp_table=created_table,
                     primary_keys=primary_keys,
                     precombine_key=precombine_key,
+                    column_names=column_names,
                 )
             if commit_transaction:
                 con.commit()