Add tests to VARCHAR lengths for Redshift and Aurora tables

Author: igorborgest

awswrangler/redshift.py

@@ -409,12 +409,13 @@ def _get_redshift_schema(dataframe,
                     redshift_type = data_types.pyarrow2redshift(dtype=dtype, varchar_length=varchar_len)
                     schema_built.append((name, redshift_type))
         elif dataframe_type.lower() == "spark":
+            logger.debug(f"cast_columns.keys: {cast_columns.keys()}")
             for name, dtype in dataframe.dtypes:
                 varchar_len = varchar_lengths.get(name, varchar_default_length)
                 if name in cast_columns.keys():
                     redshift_type = data_types.athena2redshift(dtype=cast_columns[name], varchar_length=varchar_len)
                 else:
-                    redshift_type = data_types.spark2redshift(dtype=cast_columns[name], varchar_length=varchar_len)
+                    redshift_type = data_types.spark2redshift(dtype=dtype, varchar_length=varchar_len)
                 schema_built.append((name, redshift_type))
         else:
             raise InvalidDataframeType(

testing/test_awswrangler/test_pandas.py

@@ -2365,3 +2365,68 @@ def test_s3_overall_nan(bucket, database):
                                     ctas_approach=True)
     wr.s3.delete_objects(path=path)
     assert df.equals(df2)
+
+
+def test_aurora_postgres_load_varchar(bucket, postgres_parameters):
+    df = pd.DataFrame({"id": [1, 2, 3], "varchar3": ["foo", "boo", "bar"], "varchar1": ["a", "b", "c"]})
+    path = f"s3://{bucket}/test_aurora_postgres_load_varchar"
+    wr.pandas.to_aurora(dataframe=df,
+                        connection="aws-data-wrangler-postgres",
+                        schema="public",
+                        table="test_aurora_postgres_load_varchar",
+                        mode="overwrite",
+                        temp_s3_path=path,
+                        engine="postgres",
+                        preserve_index=False,
+                        varchar_default_length=3,
+                        varchar_lengths={"varchar1": 1})
+    conn = Aurora.generate_connection(database="postgres",
+                                      host=postgres_parameters["PostgresAddress"],
+                                      port=3306,
+                                      user="test",
+                                      password=postgres_parameters["DatabasesPassword"],
+                                      engine="postgres")
+    with conn.cursor() as cursor:
+        cursor.execute("SELECT * FROM public.test_aurora_postgres_load_varchar")
+        rows = cursor.fetchall()
+        assert len(rows) == len(df.index)
+        assert rows[0][0] == 1
+        assert rows[1][0] == 2
+        assert rows[2][0] == 3
+        assert rows[0][1] == "foo"
+        assert rows[1][1] == "boo"
+        assert rows[2][1] == "bar"
+        assert rows[0][2] == "a"
+        assert rows[1][2] == "b"
+        assert rows[2][2] == "c"
+    conn.close()
+
+
+def test_aurora_mysql_load_varchar(bucket):
+    df = pd.DataFrame({"id": [1, 2, 3], "varchar3": ["foo", "boo", "bar"], "varchar1": ["a", "b", "c"]})
+    path = f"s3://{bucket}/test_aurora_mysql_load_varchar"
+    wr.pandas.to_aurora(dataframe=df,
+                        connection="aws-data-wrangler-mysql",
+                        schema="test",
+                        table="test_aurora_mysql_load_varchar",
+                        mode="overwrite",
+                        temp_s3_path=path,
+                        engine="mysql",
+                        preserve_index=False,
+                        varchar_default_length=3,
+                        varchar_lengths={"varchar1": 1})
+    conn = wr.glue.get_connection("aws-data-wrangler-mysql")
+    with conn.cursor() as cursor:
+        cursor.execute("SELECT * FROM test.test_aurora_mysql_load_varchar")
+        rows = cursor.fetchall()
+        assert len(rows) == len(df.index)
+        assert rows[0][0] == 1
+        assert rows[1][0] == 2
+        assert rows[2][0] == 3
+        assert rows[0][1] == "foo"
+        assert rows[1][1] == "boo"
+        assert rows[2][1] == "bar"
+        assert rows[0][2] == "a"
+        assert rows[1][2] == "b"
+        assert rows[2][2] == "c"
+    conn.close()

testing/test_awswrangler/test_redshift.py

@@ -866,3 +866,49 @@ def test_spectrum_csv(bucket, glue_database, external_schema):
     assert len(rows) == len(df.index)
     for row in rows:
         assert len(row) == len(df.columns)
+
+
+def test_to_redshift_pandas_varchar(bucket, redshift_parameters):
+    df = pd.DataFrame({"id": [1, 2, 3], "varchar3": ["foo", "boo", "bar"], "varchar1": ["a", "b", "c"]})
+    path = f"s3://{bucket}/test_to_redshift_pandas_varchar"
+    wr.pandas.to_redshift(dataframe=df,
+                          path=path,
+                          schema="public",
+                          table="test_to_redshift_pandas_varchar",
+                          connection="aws-data-wrangler-redshift",
+                          iam_role=redshift_parameters.get("RedshiftRole"),
+                          mode="overwrite",
+                          preserve_index=False,
+                          varchar_default_length=3,
+                          varchar_lengths={"varchar1": 1})
+    conn = wr.glue.get_connection("aws-data-wrangler-redshift")
+    with conn.cursor() as cursor:
+        cursor.execute("SELECT * FROM public.test_to_redshift_pandas_varchar")
+        rows = cursor.fetchall()
+        assert len(rows) == len(df.index)
+        for row in rows:
+            assert len(row) == len(df.columns)
+    conn.close()
+
+
+def test_to_redshift_spark_varchar(session, bucket, redshift_parameters):
+    pdf = pd.DataFrame({"id": [1, 2, 3], "varchar3": ["foo", "boo", "bar"], "varchar1": ["a", "b", "c"]})
+    df = session.spark_session.createDataFrame(pdf)
+    path = f"s3://{bucket}/test_to_redshift_spark_varchar"
+    session.spark.to_redshift(dataframe=df,
+                              path=path,
+                              schema="public",
+                              table="test_to_redshift_spark_varchar",
+                              connection="aws-data-wrangler-redshift",
+                              iam_role=redshift_parameters.get("RedshiftRole"),
+                              mode="overwrite",
+                              varchar_default_length=3,
+                              varchar_lengths={"varchar1": 1})
+    conn = wr.glue.get_connection("aws-data-wrangler-redshift")
+    with conn.cursor() as cursor:
+        cursor.execute("SELECT * FROM public.test_to_redshift_spark_varchar")
+        rows = cursor.fetchall()
+        assert len(rows) == len(pdf.index)
+        for row in rows:
+            assert len(row) == len(pdf.columns)
+    conn.close()