Merge pull request #108 from awslabs/aurora

Handling null values for Pandas.to_aurora() and eventual consistency for postgres load

Author: igorborgest

awswrangler/aurora.py

@@ -1,12 +1,14 @@
 from typing import TYPE_CHECKING, Union, List, Dict, Tuple, Any
-from logging import getLogger, Logger
+from logging import getLogger, Logger, INFO
 import json
 import warnings
 
 import pg8000  # type: ignore
+from pg8000 import ProgrammingError  # type: ignore
 import pymysql  # type: ignore
 import pandas as pd  # type: ignore
 from boto3 import client  # type: ignore
+import tenacity  # type: ignore
 
 from awswrangler import data_types
 from awswrangler.exceptions import InvalidEngine, InvalidDataframeType, AuroraLoadError
@@ -134,7 +136,7 @@ def load_table(dataframe: pd.DataFrame,
                    schema_name: str,
                    table_name: str,
                    connection: Any,
-                   num_files,
+                   num_files: int,
                    mode: str = "append",
                    preserve_index: bool = False,
                    engine: str = "mysql",
@@ -156,6 +158,54 @@ def load_table(dataframe: pd.DataFrame,
         :param region: AWS S3 bucket region (Required only for postgres engine)
         :return: None
         """
+        if "postgres" in engine.lower():
+            Aurora.load_table_postgres(dataframe=dataframe,
+                                       dataframe_type=dataframe_type,
+                                       load_paths=load_paths,
+                                       schema_name=schema_name,
+                                       table_name=table_name,
+                                       connection=connection,
+                                       mode=mode,
+                                       preserve_index=preserve_index,
+                                       region=region)
+        elif "mysql" in engine.lower():
+            Aurora.load_table_mysql(dataframe=dataframe,
+                                    dataframe_type=dataframe_type,
+                                    manifest_path=load_paths[0],
+                                    schema_name=schema_name,
+                                    table_name=table_name,
+                                    connection=connection,
+                                    mode=mode,
+                                    preserve_index=preserve_index,
+                                    num_files=num_files)
+        else:
+            raise InvalidEngine(f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!")
+
+    @staticmethod
+    def load_table_postgres(dataframe: pd.DataFrame,
+                            dataframe_type: str,
+                            load_paths: List[str],
+                            schema_name: str,
+                            table_name: str,
+                            connection: Any,
+                            mode: str = "append",
+                            preserve_index: bool = False,
+                            region: str = "us-east-1"):
+        """
+        Load text/CSV files into a Aurora table using a manifest file.
+        Creates the table if necessary.
+
+        :param dataframe: Pandas or Spark Dataframe
+        :param dataframe_type: "pandas" or "spark"
+        :param load_paths: S3 paths to be loaded (E.g. S3://...)
+        :param schema_name: Aurora schema
+        :param table_name: Aurora table name
+        :param connection: A PEP 249 compatible connection (Can be generated with Aurora.generate_connection())
+        :param mode: append or overwrite
+        :param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
+        :param region: AWS S3 bucket region (Required only for postgres engine)
+        :return: None
+        """
         with connection.cursor() as cursor:
             if mode == "overwrite":
                 Aurora._create_table(cursor=cursor,
@@ -164,30 +214,94 @@ def load_table(dataframe: pd.DataFrame,
                                      schema_name=schema_name,
                                      table_name=table_name,
                                      preserve_index=preserve_index,
-                                     engine=engine)
-            for path in load_paths:
-                sql = Aurora._get_load_sql(path=path,
-                                           schema_name=schema_name,
-                                           table_name=table_name,
-                                           engine=engine,
-                                           region=region)
-                logger.debug(sql)
+                                     engine="postgres")
+                connection.commit()
+                logger.debug("CREATE TABLE committed.")
+        for path in load_paths:
+            Aurora._load_object_postgres_with_retry(connection=connection,
+                                                    schema_name=schema_name,
+                                                    table_name=table_name,
+                                                    path=path,
+                                                    region=region)
+
+    @staticmethod
+    @tenacity.retry(retry=tenacity.retry_if_exception_type(exception_types=ProgrammingError),
+                    wait=tenacity.wait_random_exponential(multiplier=0.5),
+                    stop=tenacity.stop_after_attempt(max_attempt_number=5),
+                    reraise=True,
+                    after=tenacity.after_log(logger, INFO))
+    def _load_object_postgres_with_retry(connection: Any, schema_name: str, table_name: str, path: str,
+                                         region: str) -> None:
+        with connection.cursor() as cursor:
+            sql = Aurora._get_load_sql(path=path,
+                                       schema_name=schema_name,
+                                       table_name=table_name,
+                                       engine="postgres",
+                                       region=region)
+            logger.debug(sql)
+            try:
                 cursor.execute(sql)
+            except ProgrammingError as ex:
+                if "The file has been modified" in str(ex):
+                    connection.rollback()
+                    raise ex
+            connection.commit()
+            logger.debug(f"Load committed for: {path}.")
 
-        connection.commit()
-        logger.debug("Load committed.")
+    @staticmethod
+    def load_table_mysql(dataframe: pd.DataFrame,
+                         dataframe_type: str,
+                         manifest_path: str,
+                         schema_name: str,
+                         table_name: str,
+                         connection: Any,
+                         num_files: int,
+                         mode: str = "append",
+                         preserve_index: bool = False):
+        """
+        Load text/CSV files into a Aurora table using a manifest file.
+        Creates the table if necessary.
 
-        if "mysql" in engine.lower():
-            with connection.cursor() as cursor:
-                sql = ("-- AWS DATA WRANGLER\n"
-                       f"SELECT COUNT(*) as num_files_loaded FROM mysql.aurora_s3_load_history "
-                       f"WHERE load_prefix = '{path}'")
-                logger.debug(sql)
-                cursor.execute(sql)
-                num_files_loaded = cursor.fetchall()[0][0]
-                if num_files_loaded != (num_files + 1):
-                    raise AuroraLoadError(
-                        f"Missing files to load. {num_files_loaded} files counted. {num_files + 1} expected.")
+        :param dataframe: Pandas or Spark Dataframe
+        :param dataframe_type: "pandas" or "spark"
+        :param manifest_path: S3 manifest path to be loaded (E.g. S3://...)
+        :param schema_name: Aurora schema
+        :param table_name: Aurora table name
+        :param connection: A PEP 249 compatible connection (Can be generated with Aurora.generate_connection())
+        :param num_files: Number of files to be loaded
+        :param mode: append or overwrite
+        :param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
+        :return: None
+        """
+        with connection.cursor() as cursor:
+            if mode == "overwrite":
+                Aurora._create_table(cursor=cursor,
+                                     dataframe=dataframe,
+                                     dataframe_type=dataframe_type,
+                                     schema_name=schema_name,
+                                     table_name=table_name,
+                                     preserve_index=preserve_index,
+                                     engine="mysql")
+            sql = Aurora._get_load_sql(path=manifest_path,
+                                       schema_name=schema_name,
+                                       table_name=table_name,
+                                       engine="mysql")
+            logger.debug(sql)
+            cursor.execute(sql)
+            logger.debug(f"Load done for: {manifest_path}")
+            connection.commit()
+            logger.debug("Load committed.")
+
+        with connection.cursor() as cursor:
+            sql = ("-- AWS DATA WRANGLER\n"
+                   f"SELECT COUNT(*) as num_files_loaded FROM mysql.aurora_s3_load_history "
+                   f"WHERE load_prefix = '{manifest_path}'")
+            logger.debug(sql)
+            cursor.execute(sql)
+            num_files_loaded = cursor.fetchall()[0][0]
+            if num_files_loaded != (num_files + 1):
+                raise AuroraLoadError(
+                    f"Missing files to load. {num_files_loaded} files counted. {num_files + 1} expected.")
 
     @staticmethod
     def _parse_path(path):

awswrangler/pandas.py

@@ -644,9 +644,11 @@ def _apply_dates_to_generator(generator, parse_dates):
     def to_csv(self,
                dataframe: pd.DataFrame,
                path: str,
-               sep: str = ",",
+               sep: Optional[str] = None,
+               na_rep: Optional[str] = None,
+               quoting: Optional[int] = None,
                escapechar: Optional[str] = None,
-               serde: str = "OpenCSVSerDe",
+               serde: Optional[str] = "OpenCSVSerDe",
                database: Optional[str] = None,
                table: Optional[str] = None,
                partition_cols: Optional[List[str]] = None,
@@ -665,8 +667,10 @@ def to_csv(self,
         :param dataframe: Pandas Dataframe
         :param path: AWS S3 path (E.g. s3://bucket-name/folder_name/
         :param sep: Same as pandas.to_csv()
+        :param na_rep: Same as pandas.to_csv()
+        :param quoting: Same as pandas.to_csv()
         :param escapechar: Same as pandas.to_csv()
-        :param serde: SerDe library name (e.g. OpenCSVSerDe, LazySimpleSerDe)
+        :param serde: SerDe library name (e.g. OpenCSVSerDe, LazySimpleSerDe) (For Athena/Glue Catalog only)
         :param database: AWS Glue Database name
         :param table: AWS Glue table name
         :param partition_cols: List of columns names that will be partitions on S3
@@ -680,9 +684,17 @@ def to_csv(self,
         :param columns_comments: Columns names and the related comments (Optional[Dict[str, str]])
         :return: List of objects written on S3
         """
-        if serde not in Pandas.VALID_CSV_SERDES:
+        if (serde is not None) and (serde not in Pandas.VALID_CSV_SERDES):
             raise InvalidSerDe(f"{serde} in not in the valid SerDe list ({Pandas.VALID_CSV_SERDES})")
-        extra_args: Dict[str, Optional[str]] = {"sep": sep, "serde": serde, "escapechar": escapechar}
+        if (database is not None) and (serde is None):
+            raise InvalidParameters(f"It is not possible write to a Glue Database without a SerDe.")
+        extra_args: Dict[str, Optional[Union[str, int]]] = {
+            "sep": sep,
+            "na_rep": na_rep,
+            "serde": serde,
+            "escapechar": escapechar,
+            "quoting": quoting
+        }
         return self.to_s3(dataframe=dataframe,
                           path=path,
                           file_format="csv",
@@ -767,7 +779,7 @@ def to_s3(self,
               procs_cpu_bound=None,
               procs_io_bound=None,
               cast_columns=None,
-              extra_args: Optional[Dict[str, Optional[str]]] = None,
+              extra_args: Optional[Dict[str, Optional[Union[str, int]]]] = None,
               inplace: bool = True,
               description: Optional[str] = None,
               parameters: Optional[Dict[str, str]] = None,
@@ -1053,17 +1065,24 @@ def write_csv_dataframe(dataframe, path, preserve_index, compression, fs, extra_
 
         serde = extra_args.get("serde")
         if serde is None:
-            escapechar = extra_args.get("escapechar")
+            escapechar: Optional[str] = extra_args.get("escapechar")
             if escapechar is not None:
                 csv_extra_args["escapechar"] = escapechar
+            quoting: Optional[str] = extra_args.get("quoting")
+            if escapechar is not None:
+                csv_extra_args["quoting"] = quoting
+            na_rep: Optional[str] = extra_args.get("na_rep")
+            if na_rep is not None:
+                csv_extra_args["na_rep"] = na_rep
         else:
             if serde == "OpenCSVSerDe":
                 csv_extra_args["quoting"] = csv.QUOTE_ALL
                 csv_extra_args["escapechar"] = "\\"
             elif serde == "LazySimpleSerDe":
                 csv_extra_args["quoting"] = csv.QUOTE_NONE
                 csv_extra_args["escapechar"] = "\\"
-        csv_buffer = bytes(
+        logger.debug(f"csv_extra_args: {csv_extra_args}")
+        csv_buffer: bytes = bytes(
             dataframe.to_csv(None, header=False, index=preserve_index, compression=compression, **csv_extra_args),
             "utf-8")
         Pandas._write_csv_to_s3_retrying(fs=fs, path=path, buffer=csv_buffer)
@@ -1554,9 +1573,13 @@ def to_aurora(self,
                     temp_s3_path = self._session.athena.create_athena_bucket() + temp_directory + "/"
             temp_s3_path = temp_s3_path if temp_s3_path[-1] == "/" else temp_s3_path + "/"
             logger.debug(f"temp_s3_path: {temp_s3_path}")
+            na_rep: str = "NULL" if "mysql" in engine.lower() else ""
             paths: List[str] = self.to_csv(dataframe=dataframe,
                                            path=temp_s3_path,
+                                           serde=None,
                                            sep=",",
+                                           na_rep=na_rep,
+                                           quoting=csv.QUOTE_MINIMAL,
                                            escapechar="\"",
                                            preserve_index=preserve_index,
                                            mode="overwrite",

awswrangler/s3.py

@@ -308,8 +308,13 @@ def get_objects_sizes(self, objects_paths: List[str], procs_io_bound: Optional[i
             receive_pipes[i].close()
         return objects_sizes
 
-    def copy_listed_objects(self, objects_paths, source_path, target_path, mode="append", procs_io_bound=None):
-        if not procs_io_bound:
+    def copy_listed_objects(self,
+                            objects_paths: List[str],
+                            source_path: str,
+                            target_path: str,
+                            mode: str = "append",
+                            procs_io_bound: Optional[int] = None):
+        if procs_io_bound is None:
             procs_io_bound = self._session.procs_io_bound
         logger.debug(f"procs_io_bound: {procs_io_bound}")
         logger.debug(f"len(objects_paths): {len(objects_paths)}")