Improve mysql and postgres tests tear down.

Author: igorborgest

awswrangler/db.py

@@ -227,6 +227,29 @@ def _convert_params(sql: str, params: Optional[Union[List, Tuple, Dict]]) -> Lis
     return args
 
 
+def _read_parquet_iterator(
+    paths: List[str],
+    keep_files: bool,
+    use_threads: bool,
+    categories: List[str] = None,
+    chunked: Union[bool, int] = True,
+    boto3_session: Optional[boto3.Session] = None,
+    s3_additional_kwargs: Optional[Dict[str, str]] = None,
+) -> Iterator[pd.DataFrame]:
+    dfs: Iterator[pd.DataFrame] = s3.read_parquet(
+        path=paths,
+        categories=categories,
+        chunked=chunked,
+        dataset=False,
+        use_threads=use_threads,
+        boto3_session=boto3_session,
+        s3_additional_kwargs=s3_additional_kwargs,
+    )
+    yield from dfs
+    if keep_files is False:
+        s3.delete_objects(path=paths, use_threads=use_threads, boto3_session=boto3_session)
+
+
 def to_sql(df: pd.DataFrame, con: sqlalchemy.engine.Engine, **pandas_kwargs) -> None:
     """Write records stored in a DataFrame to a SQL database.
 
@@ -1111,29 +1134,6 @@ def unload_redshift(
     )
 
 
-def _read_parquet_iterator(
-    paths: List[str],
-    keep_files: bool,
-    use_threads: bool,
-    categories: List[str] = None,
-    chunked: Union[bool, int] = True,
-    boto3_session: Optional[boto3.Session] = None,
-    s3_additional_kwargs: Optional[Dict[str, str]] = None,
-) -> Iterator[pd.DataFrame]:
-    dfs: Iterator[pd.DataFrame] = s3.read_parquet(
-        path=paths,
-        categories=categories,
-        chunked=chunked,
-        dataset=False,
-        use_threads=use_threads,
-        boto3_session=boto3_session,
-        s3_additional_kwargs=s3_additional_kwargs,
-    )
-    yield from dfs
-    if keep_files is False:
-        s3.delete_objects(path=paths, use_threads=use_threads, boto3_session=boto3_session)
-
-
 def unload_redshift_to_files(
     sql: str,
     path: str,

tests/conftest.py

@@ -199,10 +199,30 @@ def path3(bucket):
 
 
 @pytest.fixture(scope="function")
-def redshift_table(databases_parameters):
+def redshift_table():
     name = f"tbl_{get_time_str_with_random_suffix()}"
     print(f"Table name: {name}")
     yield name
     engine = wr.catalog.get_engine(connection="aws-data-wrangler-redshift")
     with engine.connect() as con:
         con.execute(f"DROP TABLE IF EXISTS public.{name}")
+
+
+@pytest.fixture(scope="function")
+def postgresql_table():
+    name = f"tbl_{get_time_str_with_random_suffix()}"
+    print(f"Table name: {name}")
+    yield name
+    engine = wr.catalog.get_engine(connection="aws-data-wrangler-postgresql")
+    with engine.connect() as con:
+        con.execute(f"DROP TABLE IF EXISTS public.{name}")
+
+
+@pytest.fixture(scope="function")
+def mysql_table():
+    name = f"tbl_{get_time_str_with_random_suffix()}"
+    print(f"Table name: {name}")
+    yield name
+    engine = wr.catalog.get_engine(connection="aws-data-wrangler-mysql")
+    with engine.connect() as con:
+        con.execute(f"DROP TABLE IF EXISTS test.{name}")

tests/test_db.py

@@ -17,7 +17,13 @@
 
 
 @pytest.mark.parametrize("db_type", ["mysql", "redshift", "postgresql"])
-def test_sql(redshift_table, databases_parameters, db_type):
+def test_sql(redshift_table, postgresql_table, mysql_table, databases_parameters, db_type):
+    if db_type == "postgresql":
+        table = postgresql_table
+    elif db_type == "mysql":
+        table = mysql_table
+    else:
+        table = redshift_table
     df = get_df()
     if db_type == "redshift":
         df.drop(["binary"], axis=1, inplace=True)
@@ -26,7 +32,7 @@ def test_sql(redshift_table, databases_parameters, db_type):
     wr.db.to_sql(
         df=df,
         con=engine,
-        name=redshift_table,
+        name=table,
         schema=databases_parameters[db_type]["schema"],
         if_exists="replace",
         index=index,
@@ -36,7 +42,7 @@ def test_sql(redshift_table, databases_parameters, db_type):
         dtype={"iint32": sqlalchemy.types.Integer},
     )
     df = wr.db.read_sql_query(
-        sql=f"SELECT * FROM {databases_parameters[db_type]['schema']}.{redshift_table}", con=engine
+        sql=f"SELECT * FROM {databases_parameters[db_type]['schema']}.{table}", con=engine
     )
     ensure_data_types(df, has_list=False)
     engine = wr.db.get_engine(
@@ -49,7 +55,7 @@ def test_sql(redshift_table, databases_parameters, db_type):
         echo=False,
     )
     dfs = wr.db.read_sql_query(
-        sql=f"SELECT * FROM {databases_parameters[db_type]['schema']}.{redshift_table}",
+        sql=f"SELECT * FROM {databases_parameters[db_type]['schema']}.{table}",
         con=engine,
         chunksize=1,
         dtype={
@@ -76,7 +82,7 @@ def test_sql(redshift_table, databases_parameters, db_type):
         wr.db.to_sql(
             df=pd.DataFrame({"col0": [1, 2, 3]}, dtype="Int32"),
             con=engine,
-            name=redshift_table,
+            name=table,
             schema=databases_parameters[db_type]["schema"],
             if_exists="replace",
             index=True,
@@ -85,7 +91,7 @@ def test_sql(redshift_table, databases_parameters, db_type):
         schema = None
         if db_type == "postgresql":
             schema = databases_parameters[db_type]["schema"]
-        df = wr.db.read_sql_table(con=engine, table=redshift_table, schema=schema, index_col="index")
+        df = wr.db.read_sql_table(con=engine, table=table, schema=schema, index_col="index")
         assert df.shape == (3, 1)
 
 
@@ -373,8 +379,13 @@ def test_redshift_unload_extras(bucket, path, redshift_table, databases_paramete
 
 
 @pytest.mark.parametrize("db_type", ["mysql", "redshift", "postgresql"])
-def test_to_sql_cast(redshift_table, databases_parameters, db_type):
-    table = redshift_table
+def test_to_sql_cast(redshift_table, postgresql_table, mysql_table, databases_parameters, db_type):
+    if db_type == "postgresql":
+        table = postgresql_table
+    elif db_type == "mysql":
+        table = mysql_table
+    else:
+        table = redshift_table
     schema = databases_parameters[db_type]["schema"]
     df = pd.DataFrame(
         {
@@ -403,8 +414,8 @@ def test_to_sql_cast(redshift_table, databases_parameters, db_type):
     assert df.equals(df2)
 
 
-def test_uuid(redshift_table, databases_parameters):
-    table = redshift_table
+def test_uuid(postgresql_table, databases_parameters):
+    table = postgresql_table
     schema = databases_parameters["postgresql"]["schema"]
     engine = wr.catalog.get_engine(connection="aws-data-wrangler-postgresql")
     df = pd.DataFrame(
@@ -436,8 +447,13 @@ def test_uuid(redshift_table, databases_parameters):
 
 
 @pytest.mark.parametrize("db_type", ["mysql", "redshift", "postgresql"])
-def test_null(redshift_table, databases_parameters, db_type):
-    table = redshift_table
+def test_null(redshift_table, postgresql_table, mysql_table, databases_parameters, db_type):
+    if db_type == "postgresql":
+        table = postgresql_table
+    elif db_type == "mysql":
+        table = mysql_table
+    else:
+        table = redshift_table
     schema = databases_parameters[db_type]["schema"]
     engine = wr.catalog.get_engine(connection=f"aws-data-wrangler-{db_type}")
     df = pd.DataFrame({"id": [1, 2, 3], "nothing": [None, None, None]})