Now to_redshift() will created the table if it does not exists yet even for append or upsert mode.

Author: igorborgest

awswrangler/redshift.py

@@ -242,16 +242,32 @@ def load_table(dataframe,
                                        varchar_default_length=varchar_default_length,
                                        varchar_lengths=varchar_lengths)
                 table_name = f"{schema_name}.{table_name}"
-            elif mode == "upsert":
-                guid: str = pa.compat.guid()
-                temp_table_name = f"temp_redshift_{guid}"
-                final_table_name = table_name
-                table_name = temp_table_name
-                sql: str = f"CREATE TEMPORARY TABLE {temp_table_name} (LIKE {schema_name}.{final_table_name})"
-                logger.debug(sql)
-                cursor.execute(sql)
             else:
-                table_name = f"{schema_name}.{table_name}"
+                if Redshift.does_table_exists(cursor=cursor, schema=schema_name, table=table_name) is False:
+                    Redshift._create_table(cursor=cursor,
+                                           dataframe=dataframe,
+                                           dataframe_type=dataframe_type,
+                                           schema_name=schema_name,
+                                           table_name=table_name,
+                                           diststyle=diststyle,
+                                           distkey=distkey,
+                                           sortstyle=sortstyle,
+                                           sortkey=sortkey,
+                                           primary_keys=primary_keys,
+                                           preserve_index=preserve_index,
+                                           cast_columns=cast_columns,
+                                           varchar_default_length=varchar_default_length,
+                                           varchar_lengths=varchar_lengths)
+                if mode == "upsert":
+                    guid: str = pa.compat.guid()
+                    temp_table_name = f"temp_redshift_{guid}"
+                    final_table_name = table_name
+                    table_name = temp_table_name
+                    sql: str = f"CREATE TEMPORARY TABLE {temp_table_name} (LIKE {schema_name}.{final_table_name})"
+                    logger.debug(sql)
+                    cursor.execute(sql)
+                else:
+                    table_name = f"{schema_name}.{table_name}"
 
             sql = ("-- AWS DATA WRANGLER\n"
                    f"COPY {table_name} FROM '{manifest_path}'\n"
@@ -513,3 +529,24 @@ def to_parquet(self,
                 logger.debug(f"Waiting path: {p}")
                 self._session.s3.wait_object_exists(path=p, timeout=30.0)
         return paths
+
+    @staticmethod
+    def does_table_exists(cursor, schema: str, table: str) -> bool:
+        """
+        Check if the table exists.
+
+        :param cursor: A PEP 249 compatible cursor (A valid connection be generated with Redshift.generate_connection())
+        :param schema: Schema name
+        :param table: Table name
+        :return: True or False
+        """
+        cursor.execute(f"SELECT true WHERE EXISTS ("
+                       f"SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE "
+                       f"TABLE_SCHEMA = '{schema}' AND TABLE_NAME = '{table}'"
+                       f");")
+        rows: List = cursor.fetchall()
+        if len(rows) > 0:
+            res: bool = rows[0][0]
+        else:
+            res = False
+        return res

testing/test_awswrangler/test_redshift.py

@@ -947,3 +947,54 @@ def test_to_redshift_int_na(bucket, redshift_parameters):
         rows = cursor.fetchall()
         assert rows[0][0] == 3
     conn.close()
+
+
+def test_does_table_exists():
+    conn = wr.glue.get_connection("aws-data-wrangler-redshift")
+    with conn.cursor() as cursor:
+        res = wr.redshift.does_table_exists(cursor=cursor, schema="information_schema", table="tables")
+        assert res is True
+        res = wr.redshift.does_table_exists(cursor=cursor, schema="information_schema", table="adjwbdjawbjd")
+        assert res is False
+    conn.close()
+
+
+def test_append_create_table(bucket, redshift_parameters):
+    df = pd.DataFrame({
+        "id": [1, 2, 3],
+        "decimal_2": [Decimal((0, (1, 9, 9), -2)), None, Decimal((0, (1, 9, 0), -2))],
+        "decimal_5": [Decimal((0, (1, 9, 9, 9, 9, 9), -5)), None,
+                      Decimal((0, (1, 9, 0, 0, 0, 0), -5))],
+    })
+    con = wr.redshift.get_connection("aws-data-wrangler-redshift")
+    cursor = con.cursor()
+    cursor.execute("DROP TABLE IF EXISTS public.test_append_create_table")
+    cursor.close()
+    path = f"s3://{bucket}/test_append_create_table/"
+    wr.pandas.to_redshift(
+        dataframe=df,
+        path=path,
+        schema="public",
+        table="test_append_create_table",
+        connection=con,
+        iam_role=redshift_parameters.get("RedshiftRole"),
+        mode="append",
+        preserve_index=False,
+    )
+    cursor = con.cursor()
+    cursor.execute("SELECT * from public.test_append_create_table")
+    rows = cursor.fetchall()
+    cursor.close()
+    con.close()
+    assert len(df.index) == len(rows)
+    assert len(list(df.columns)) == len(list(rows[0]))
+    for row in rows:
+        if row[0] == 1:
+            assert row[1] == Decimal((0, (1, 9, 9), -2))
+            assert row[2] == Decimal((0, (1, 9, 9, 9, 9, 9), -5))
+        elif row[1] == 2:
+            assert row[1] is None
+            assert row[2] is None
+        elif row[2] == 3:
+            assert row[1] == Decimal((0, (1, 9, 0), -2))
+            assert row[2] == Decimal((0, (1, 9, 0, 0, 0, 0), -5))