Add Redshift tests

Author: igorborgest

testing/test_awswrangler/test_db.py

@@ -76,6 +76,11 @@ def external_schema(cloudformation_outputs, parameters, glue_database):
     yield "aws_data_wrangler_external"
 
 
+@pytest.fixture(scope="module")
+def kms_key_id(cloudformation_outputs):
+    yield cloudformation_outputs["KmsKeyArn"].split("/", 1)[1]
+
+
 @pytest.mark.parametrize("db_type", ["mysql", "redshift", "postgresql"])
 def test_sql(parameters, db_type):
     df = get_df()
@@ -386,3 +391,72 @@ def test_redshift_category(bucket, parameters):
     for df2 in dfs:
         ensure_data_types_category(df2)
     wr.s3.delete_objects(path=path)
+
+
+def test_redshift_unload_extras(bucket, parameters, kms_key_id):
+    table = "test_redshift_unload_extras"
+    schema = parameters["redshift"]["schema"]
+    path = f"s3://{bucket}/{table}/"
+    wr.s3.delete_objects(path=path)
+    engine = wr.catalog.get_engine(connection=f"aws-data-wrangler-redshift")
+    df = pd.DataFrame({"id": [1, 2], "name": ["foo", "boo"]})
+    wr.db.to_sql(df=df, con=engine, name=table, schema=schema, if_exists="replace", index=False)
+    paths = wr.db.unload_redshift_to_files(
+        sql=f"SELECT * FROM {schema}.{table}",
+        path=path,
+        con=engine,
+        iam_role=parameters["redshift"]["role"],
+        region=wr.s3.get_bucket_region(bucket),
+        max_file_size=5.0,
+        kms_key_id=kms_key_id,
+        partition_cols=["name"],
+    )
+    wr.s3.wait_objects_exist(paths=paths)
+    df = wr.s3.read_parquet(path=path, dataset=True)
+    assert len(df.index) == 2
+    assert len(df.columns) == 2
+    wr.s3.delete_objects(path=path)
+    df = wr.db.unload_redshift(
+        sql=f"SELECT * FROM {schema}.{table}",
+        con=engine,
+        iam_role=parameters["redshift"]["role"],
+        path=path,
+        keep_files=False,
+        region=wr.s3.get_bucket_region(bucket),
+        max_file_size=5.0,
+        kms_key_id=kms_key_id,
+    )
+    assert len(df.index) == 2
+    assert len(df.columns) == 2
+    wr.s3.delete_objects(path=path)
+
+
+@pytest.mark.parametrize("db_type", ["mysql", "redshift", "postgresql"])
+def test_to_sql_cast(parameters, db_type):
+    table = "test_to_sql_cast"
+    schema = parameters[db_type]["schema"]
+    df = pd.DataFrame(
+        {
+            "col": [
+                "".join([str(i)[-1] for i in range(1_024)]),
+                "".join([str(i)[-1] for i in range(1_024)]),
+                "".join([str(i)[-1] for i in range(1_024)]),
+            ]
+        },
+        dtype="string",
+    )
+    engine = wr.catalog.get_engine(connection=f"aws-data-wrangler-{db_type}")
+    wr.db.to_sql(
+        df=df,
+        con=engine,
+        name=table,
+        schema=schema,
+        if_exists="replace",
+        index=False,
+        index_label=None,
+        chunksize=None,
+        method=None,
+        dtype={"col": sqlalchemy.types.VARCHAR(length=1_024)},
+    )
+    df2 = wr.db.read_sql_query(sql=f"SELECT * FROM {schema}.{table}", con=engine)
+    assert df.equals(df2)