Improving databases internals.

Author: igorborgest

awswrangler/_databases.py

@@ -116,23 +116,28 @@ def read_sql_query(
 ) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
     """Read SQL Query (generic)."""
     args = _convert_params(sql, params)
-    with con.cursor() as cursor:
-        cursor.execute(*args)
-        cols_names: List[str] = [
-            col[0].decode("utf-8") if isinstance(col[0], bytes) else col[0] for col in cursor.description
-        ]
-        _logger.debug("cols_names: %s", cols_names)
-        if chunksize is None:
-            return _records2df(
-                records=cast(List[Tuple[Any]], cursor.fetchall()),
-                cols_names=cols_names,
-                index=index_col,
-                dtype=dtype,
-                safe=safe,
+    try:
+        with con.cursor() as cursor:
+            cursor.execute(*args)
+            cols_names: List[str] = [
+                col[0].decode("utf-8") if isinstance(col[0], bytes) else col[0] for col in cursor.description
+            ]
+            _logger.debug("cols_names: %s", cols_names)
+            if chunksize is None:
+                return _records2df(
+                    records=cast(List[Tuple[Any]], cursor.fetchall()),
+                    cols_names=cols_names,
+                    index=index_col,
+                    dtype=dtype,
+                    safe=safe,
+                )
+            return _iterate_cursor(
+                cursor=cursor, chunksize=chunksize, cols_names=cols_names, index=index_col, dtype=dtype, safe=safe
             )
-        return _iterate_cursor(
-            cursor=cursor, chunksize=chunksize, cols_names=cols_names, index=index_col, dtype=dtype, safe=safe
-        )
+    except Exception as ex:
+        con.rollback()
+        _logger.error(ex)
+        raise
 
 
 def extract_parameters(df: pd.DataFrame) -> List[List[Any]]:

awswrangler/mysql.py

@@ -26,8 +26,8 @@ def _validate_connection(con: pymysql.connections.Connection) -> None:
 
 
 def _drop_table(cursor: Cursor, schema: Optional[str], table: str) -> None:
-    schema_str = f"{schema}." if schema else ""
-    sql = f"DROP TABLE IF EXISTS {schema_str}{table}"
+    schema_str = f"`{schema}`." if schema else ""
+    sql = f"DROP TABLE IF EXISTS {schema_str}`{table}`"
     _logger.debug("Drop table query:\n%s", sql)
     cursor.execute(sql)
 
@@ -61,7 +61,7 @@ def _create_table(
         converter_func=_data_types.pyarrow2mysql,
     )
     cols_str: str = "".join([f"`{k}` {v},\n" for k, v in mysql_types.items()])[:-2]
-    sql = f"CREATE TABLE IF NOT EXISTS {schema}.{table} (\n" f"{cols_str}" f")"
+    sql = f"CREATE TABLE IF NOT EXISTS `{schema}`.`{table}` (\n{cols_str})"
     _logger.debug("Create table query:\n%s", sql)
     cursor.execute(sql)
 
@@ -246,7 +246,7 @@ def read_sql_table(
     >>> con.close()
 
     """
-    sql: str = f"SELECT * FROM {table}" if schema is None else f"SELECT * FROM {schema}.{table}"
+    sql: str = f"SELECT * FROM `{table}`" if schema is None else f"SELECT * FROM `{schema}`.`{table}`"
     return read_sql_query(
         sql=sql, con=con, index_col=index_col, params=params, chunksize=chunksize, dtype=dtype, safe=safe
     )
@@ -310,7 +310,6 @@ def to_sql(
     if df.empty is True:
         raise exceptions.EmptyDataFrame()
     _validate_connection(con=con)
-    con.autocommit(True)  # type: ignore
     try:
         with con.cursor() as cursor:
             _create_table(
@@ -326,7 +325,7 @@ def to_sql(
             if index:
                 df.reset_index(level=df.index.names, inplace=True)
             placeholders: str = ", ".join(["%s"] * len(df.columns))
-            sql: str = f"INSERT INTO {schema}.{table} VALUES ({placeholders})"
+            sql: str = f"INSERT INTO `{schema}`.`{table}` VALUES ({placeholders})"
             _logger.debug("sql: %s", sql)
             parameters: List[List[Any]] = _db_utils.extract_parameters(df=df)
             cursor.executemany(sql, parameters)

awswrangler/postgresql.py

@@ -25,8 +25,8 @@ def _validate_connection(con: pg8000.Connection) -> None:
 
 
 def _drop_table(cursor: pg8000.Cursor, schema: Optional[str], table: str) -> None:
-    schema_str = f"{schema}." if schema else ""
-    sql = f"DROP TABLE IF EXISTS {schema_str}{table}"
+    schema_str = f'"{schema}".' if schema else ""
+    sql = f'DROP TABLE IF EXISTS {schema_str}"{table}"'
     _logger.debug("Drop table query:\n%s", sql)
     cursor.execute(sql)
 
@@ -65,7 +65,7 @@ def _create_table(
         converter_func=_data_types.pyarrow2postgresql,
     )
     cols_str: str = "".join([f"{k} {v},\n" for k, v in postgresql_types.items()])[:-2]
-    sql = f"CREATE TABLE IF NOT EXISTS {schema}.{table} (\n" f"{cols_str}" f")"
+    sql = f'CREATE TABLE IF NOT EXISTS "{schema}"."{table}" (\n{cols_str})'
     _logger.debug("Create table query:\n%s", sql)
     cursor.execute(sql)
 
@@ -250,7 +250,7 @@ def read_sql_table(
     >>> con.close()
 
     """
-    sql: str = f"SELECT * FROM {table}" if schema is None else f"SELECT * FROM {schema}.{table}"
+    sql: str = f'SELECT * FROM "{table}"' if schema is None else f'SELECT * FROM "{schema}"."{table}"'
     return read_sql_query(
         sql=sql, con=con, index_col=index_col, params=params, chunksize=chunksize, dtype=dtype, safe=safe
     )
@@ -314,7 +314,6 @@ def to_sql(
     if df.empty is True:
         raise exceptions.EmptyDataFrame()
     _validate_connection(con=con)
-    con.autocommit = True
     try:
         with con.cursor() as cursor:
             _create_table(
@@ -330,7 +329,7 @@ def to_sql(
             if index:
                 df.reset_index(level=df.index.names, inplace=True)
             placeholders: str = ", ".join(["%s"] * len(df.columns))
-            sql: str = f"INSERT INTO {schema}.{table} VALUES ({placeholders})"
+            sql: str = f'INSERT INTO "{schema}"."{table}" VALUES ({placeholders})'
             _logger.debug("sql: %s", sql)
             parameters: List[List[Any]] = _db_utils.extract_parameters(df=df)
             cursor.executemany(sql, parameters)

awswrangler/redshift.py

@@ -29,8 +29,8 @@ def _validate_connection(con: redshift_connector.Connection) -> None:
 
 
 def _drop_table(cursor: redshift_connector.Cursor, schema: Optional[str], table: str) -> None:
-    schema_str = f"{schema}." if schema else ""
-    sql = f"DROP TABLE IF EXISTS {schema_str}{table}"
+    schema_str = f'"{schema}".' if schema else ""
+    sql = f'DROP TABLE IF EXISTS {schema_str}"{table}"'
     _logger.debug("Drop table query:\n%s", sql)
     cursor.execute(sql)
 
@@ -62,9 +62,9 @@ def _copy(
     schema: Optional[str] = None,
 ) -> None:
     if schema is None:
-        table_name: str = table
+        table_name: str = f'"{table}"'
     else:
-        table_name = f"{schema}.{table}"
+        table_name = f'"{schema}"."{table}"'
     sql: str = f"COPY {table_name} FROM '{path}'\nIAM_ROLE '{iam_role}'\nFORMAT AS PARQUET"
     _logger.debug("copy query:\n%s", sql)
     cursor.execute(sql)
@@ -84,7 +84,7 @@ def _upsert(
         raise exceptions.InvalidRedshiftPrimaryKeys()
     equals_clause: str = f"{table}.%s = {temp_table}.%s"
     join_clause: str = " AND ".join([equals_clause % (pk, pk) for pk in primary_keys])
-    sql: str = f"DELETE FROM {schema}.{table} USING {temp_table} WHERE {join_clause}"
+    sql: str = f'DELETE FROM "{schema}"."{table}" USING {temp_table} WHERE {join_clause}'
     _logger.debug(sql)
     cursor.execute(sql)
     sql = f"INSERT INTO {schema}.{table} SELECT * FROM {temp_table}"
@@ -177,7 +177,7 @@ def _create_table(
         if mode == "upsert":
             guid: str = uuid.uuid4().hex
             temp_table: str = f"temp_redshift_{guid}"
-            sql: str = f"CREATE TEMPORARY TABLE {temp_table} (LIKE {schema}.{table})"
+            sql: str = f'CREATE TEMPORARY TABLE {temp_table} (LIKE "{schema}"."{table}")'
             _logger.debug(sql)
             cursor.execute(sql)
             return temp_table, None
@@ -217,7 +217,7 @@ def _create_table(
     distkey_str: str = f"\nDISTKEY({distkey})" if distkey and diststyle == "KEY" else ""
     sortkey_str: str = f"\n{sortstyle} SORTKEY({','.join(sortkey)})" if sortkey else ""
     sql = (
-        f"CREATE TABLE IF NOT EXISTS {schema}.{table} (\n"
+        f'CREATE TABLE IF NOT EXISTS "{schema}"."{table}" (\n'
         f"{cols_str}"
         f"{primary_keys_str}"
         f")\nDISTSTYLE {diststyle}"
@@ -538,7 +538,7 @@ def read_sql_table(
     >>> con.close()
 
     """
-    sql: str = f"SELECT * FROM {table}" if schema is None else f"SELECT * FROM {schema}.{table}"
+    sql: str = f'SELECT * FROM "{table}"' if schema is None else f'SELECT * FROM "{schema}"."{table}"'
     return read_sql_query(
         sql=sql, con=con, index_col=index_col, params=params, chunksize=chunksize, dtype=dtype, safe=safe
     )
@@ -627,6 +627,7 @@ def to_sql(
     if df.empty is True:
         raise exceptions.EmptyDataFrame()
     _validate_connection(con=con)
+    autocommit_temp: bool = con.autocommit
     con.autocommit = False
     try:
         with con.cursor() as cursor:
@@ -650,8 +651,8 @@ def to_sql(
             if index:
                 df.reset_index(level=df.index.names, inplace=True)
             placeholders: str = ", ".join(["%s"] * len(df.columns))
-            schema_str = f"{created_schema}." if created_schema else ""
-            sql: str = f"INSERT INTO {schema_str}{created_table} VALUES ({placeholders})"
+            schema_str = f'"{created_schema}".' if created_schema else ""
+            sql: str = f'INSERT INTO {schema_str}"{created_table}" VALUES ({placeholders})'
             _logger.debug("sql: %s", sql)
             parameters: List[List[Any]] = _db_utils.extract_parameters(df=df)
             cursor.executemany(sql, parameters)
@@ -662,6 +663,8 @@ def to_sql(
         con.rollback()
         _logger.error(ex)
         raise
+    finally:
+        con.autocommit = autocommit_temp
 
 
 def unload_to_files(
@@ -1009,6 +1012,7 @@ def copy_from_files(  # pylint: disable=too-many-locals,too-many-arguments
     >>> con.close()
 
     """
+    autocommit_temp: bool = con.autocommit
     con.autocommit = False
     try:
         with con.cursor() as cursor:
@@ -1047,6 +1051,8 @@ def copy_from_files(  # pylint: disable=too-many-locals,too-many-arguments
         con.rollback()
         _logger.error(ex)
         raise
+    finally:
+        con.autocommit = autocommit_temp
 
 
 def copy(  # pylint: disable=too-many-arguments

tests/conftest.py

@@ -208,6 +208,7 @@ def redshift_table():
     con = wr.redshift.connect("aws-data-wrangler-redshift")
     with con.cursor() as cursor:
         cursor.execute(f"DROP TABLE IF EXISTS public.{name}")
+    con.commit()
     con.close()
 
 
@@ -219,6 +220,7 @@ def postgresql_table():
     con = wr.postgresql.connect("aws-data-wrangler-postgresql")
     with con.cursor() as cursor:
         cursor.execute(f"DROP TABLE IF EXISTS public.{name}")
+    con.commit()
     con.close()
 
 
@@ -230,6 +232,7 @@ def mysql_table():
     con = wr.mysql.connect("aws-data-wrangler-mysql")
     with con.cursor() as cursor:
         cursor.execute(f"DROP TABLE IF EXISTS test.{name}")
+    con.commit()
     con.close()
 
 

tests/test_mysql.py

@@ -148,3 +148,26 @@ def test_decimal_cast(mysql_table):
     assert 3.88 <= df2.col1.sum() <= 3.89
     assert df2.col2.sum() == 2
     con.close()
+
+
+def test_read_retry():
+    con = wr.mysql.connect(connection="aws-data-wrangler-mysql")
+    try:
+        wr.mysql.read_sql_query("ERROR", con)
+    except:  # noqa
+        pass
+    df = wr.mysql.read_sql_query("SELECT 1", con)
+    assert df.shape == (1, 1)
+    con.close()
+
+
+def test_table_name():
+    df = pd.DataFrame({"col0": [1]})
+    con = wr.mysql.connect(connection="aws-data-wrangler-mysql")
+    wr.mysql.to_sql(df, con, "Test Name", "test")
+    df = wr.mysql.read_sql_table(schema="test", table="Test Name")
+    assert df.shape == (1, 1)
+    with con.cursor() as cursor:
+        cursor.execute("DROP TABLE `Test Name`")
+    con.commit()
+    con.close()

tests/test_postgresql.py

@@ -148,3 +148,26 @@ def test_decimal_cast(postgresql_table):
     assert 3.88 <= df2.col1.sum() <= 3.89
     assert df2.col2.sum() == 2
     con.close()
+
+
+def test_read_retry():
+    con = wr.postgresql.connect(connection="aws-data-wrangler-postgresql")
+    try:
+        wr.postgresql.read_sql_query("ERROR", con)
+    except:  # noqa
+        pass
+    df = wr.postgresql.read_sql_query("SELECT 1", con)
+    assert df.shape == (1, 1)
+    con.close()
+
+
+def test_table_name():
+    df = pd.DataFrame({"col0": [1]})
+    con = wr.postgresql.connect(connection="aws-data-wrangler-postgresql")
+    wr.postgresql.to_sql(df, con, "Test Name", "public", mode="overwrite")
+    df = wr.postgresql.read_sql_table(schema="public", con=con, table="Test Name")
+    assert df.shape == (1, 1)
+    with con.cursor() as cursor:
+        cursor.execute('DROP TABLE "Test Name"')
+    con.commit()
+    con.close()

tests/test_redshift.py

@@ -667,3 +667,26 @@ def test_upsert(redshift_table):
     assert len(df.columns) == len(df4.columns)
 
     con.close()
+
+
+def test_read_retry():
+    con = wr.redshift.connect(connection="aws-data-wrangler-redshift")
+    try:
+        wr.redshift.read_sql_query("ERROR", con)
+    except:  # noqa
+        pass
+    df = wr.redshift.read_sql_query("SELECT 1", con)
+    assert df.shape == (1, 1)
+    con.close()
+
+
+def test_table_name():
+    df = pd.DataFrame({"col0": [1]})
+    con = wr.redshift.connect(connection="aws-data-wrangler-redshift")
+    wr.redshift.to_sql(df, con, "Test Name", "public", mode="overwrite")
+    df = wr.redshift.read_sql_table(schema="public", con=con, table="Test Name")
+    assert df.shape == (1, 1)
+    with con.cursor() as cursor:
+        cursor.execute('DROP TABLE "Test Name"')
+    con.commit()
+    con.close()