Add Aurora module

Author: igorborgest

awswrangler/__init__.py

@@ -9,6 +9,7 @@
 from awswrangler.cloudwatchlogs import CloudWatchLogs  # noqa
 from awswrangler.glue import Glue  # noqa
 from awswrangler.redshift import Redshift  # noqa
+from awswrangler.aurora import Aurora  # noqa
 from awswrangler.emr import EMR  # noqa
 from awswrangler.sagemaker import SageMaker  # noqa
 import awswrangler.utils  # noqa
@@ -38,6 +39,7 @@ def __getattr__(self, name):
 pandas = DynamicInstantiate("pandas")
 athena = DynamicInstantiate("athena")
 redshift = DynamicInstantiate("redshift")
+aurora = DynamicInstantiate("aurora")
 sagemaker = DynamicInstantiate("sagemaker")
 cloudwatchlogs = DynamicInstantiate("cloudwatchlogs")
 

awswrangler/aurora.py

@@ -0,0 +1,8 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class Aurora:
+    def __init__(self, session):
+        self._session = session

awswrangler/sagemaker.py

@@ -1,8 +1,8 @@
+from typing import Any
 import pickle
 import tarfile
 import logging
 
-from typing import Any
 from awswrangler.exceptions import InvalidParameters
 
 logger = logging.getLogger(__name__)
@@ -12,7 +12,7 @@ class SageMaker:
     def __init__(self, session):
         self._session = session
         self._client_s3 = session.boto3_session.client(service_name="s3", use_ssl=True, config=session.botocore_config)
-        self._client_sagemaker = session.boto3_session.client(service_name="sagemaker")
+        self._client_sagemaker = session.boto3_session.client(service_name="sagemaker", use_ssl=True, config=session.botocore_config)
 
     @staticmethod
     def _parse_path(path):

awswrangler/session.py

@@ -12,6 +12,7 @@
 from awswrangler.pandas import Pandas
 from awswrangler.glue import Glue
 from awswrangler.redshift import Redshift
+from awswrangler.aurora import Aurora
 from awswrangler.emr import EMR
 from awswrangler.sagemaker import SageMaker
 from awswrangler.exceptions import AWSCredentialsNotFound
@@ -52,7 +53,8 @@ def __init__(self,
                  athena_kms_key: Optional[str] = None,
                  athena_database: str = "default",
                  athena_ctas_approach: bool = False,
-                 redshift_temp_s3_path: Optional[str] = None):
+                 redshift_temp_s3_path: Optional[str] = None,
+                 aurora_temp_s3_path: Optional[str] = None):
         """
         Most parameters inherit from Boto3 or Pyspark.
         https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html
@@ -76,7 +78,8 @@ def __init__(self,
         :param athena_s3_output: AWS S3 path
         :param athena_encryption: None|'SSE_S3'|'SSE_KMS'|'CSE_KMS'
         :param athena_kms_key: For SSE-KMS and CSE-KMS , this is the KMS key ARN or ID.
-        :param redshift_temp_s3_path: redshift_temp_s3_path: AWS S3 path to write temporary data (e.g. s3://...)
+        :param redshift_temp_s3_path: AWS S3 path to write temporary data (e.g. s3://...)
+        :param aurora_temp_s3_path: AWS S3 path to write temporary data (e.g. s3://...)
         """
         self._profile_name: Optional[str] = (boto3_session.profile_name if boto3_session else profile_name)
         self._aws_access_key_id: Optional[str] = (boto3_session.get_credentials().access_key
@@ -100,6 +103,7 @@ def __init__(self,
         self._athena_database: str = athena_database
         self._athena_ctas_approach: bool = athena_ctas_approach
         self._redshift_temp_s3_path: Optional[str] = redshift_temp_s3_path
+        self._aurora_temp_s3_path: Optional[str] = aurora_temp_s3_path
         self._primitives = None
         self._load_new_primitives()
         if boto3_session:
@@ -113,6 +117,7 @@ def __init__(self,
         self._pandas = None
         self._glue = None
         self._redshift = None
+        self._aurora = None
         self._spark = None
         self._sagemaker = None
 
@@ -160,7 +165,8 @@ def _load_new_primitives(self):
                                              athena_kms_key=self._athena_kms_key,
                                              athena_database=self._athena_database,
                                              athena_ctas_approach=self._athena_ctas_approach,
-                                             redshift_temp_s3_path=self._redshift_temp_s3_path)
+                                             redshift_temp_s3_path=self._redshift_temp_s3_path,
+                                             aurora_temp_s3_path=self._aurora_temp_s3_path)
 
     @property
     def profile_name(self):
@@ -238,6 +244,10 @@ def athena_ctas_approach(self) -> bool:
     def redshift_temp_s3_path(self) -> Optional[str]:
         return self._redshift_temp_s3_path
 
+    @property
+    def aurora_temp_s3_path(self) -> Optional[str]:
+        return self._aurora_temp_s3_path
+
     @property
     def boto3_session(self):
         return self._boto3_session
@@ -288,6 +298,12 @@ def redshift(self):
             self._redshift = Redshift(session=self)
         return self._redshift
 
+    @property
+    def aurora(self):
+        if not self._aurora:
+            self._aurora = Aurora(session=self)
+        return self._aurora
+
     @property
     def sagemaker(self):
         if not self._sagemaker:
@@ -326,7 +342,8 @@ def __init__(self,
                  athena_kms_key: Optional[str] = None,
                  athena_database: Optional[str] = None,
                  athena_ctas_approach: bool = False,
-                 redshift_temp_s3_path: Optional[str] = None):
+                 redshift_temp_s3_path: Optional[str] = None,
+                 aurora_temp_s3_path: Optional[str] = None):
         """
         Most parameters inherit from Boto3.
         https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html
@@ -348,6 +365,7 @@ def __init__(self,
         :param athena_encryption: None|'SSE_S3'|'SSE_KMS'|'CSE_KMS'
         :param athena_kms_key: For SSE-KMS and CSE-KMS , this is the KMS key ARN or ID.
         :param redshift_temp_s3_path: AWS S3 path to write temporary data (e.g. s3://...)
+        :param aurora_temp_s3_path: AWS S3 path to write temporary data (e.g. s3://...)
         """
         self._profile_name: Optional[str] = profile_name
         self._aws_access_key_id: Optional[str] = aws_access_key_id
@@ -366,6 +384,7 @@ def __init__(self,
         self._athena_database: Optional[str] = athena_database
         self._athena_ctas_approach: bool = athena_ctas_approach
         self._redshift_temp_s3_path: Optional[str] = redshift_temp_s3_path
+        self._aurora_temp_s3_path: Optional[str] = aurora_temp_s3_path
 
     @property
     def profile_name(self):
@@ -435,6 +454,10 @@ def athena_ctas_approach(self) -> bool:
     def redshift_temp_s3_path(self) -> Optional[str]:
         return self._redshift_temp_s3_path
 
+    @property
+    def aurora_temp_s3_path(self) -> Optional[str]:
+        return self._aurora_temp_s3_path
+
     @property
     def session(self):
         """
@@ -456,4 +479,5 @@ def session(self):
                        athena_kms_key=self._athena_kms_key,
                        athena_database=self._athena_database,
                        athena_ctas_approach=self._athena_ctas_approach,
-                       redshift_temp_s3_path=self._redshift_temp_s3_path)
+                       redshift_temp_s3_path=self._redshift_temp_s3_path,
+                       aurora_temp_s3_path=self._aurora_temp_s3_path)

requirements.txt

@@ -5,4 +5,5 @@ botocore~=1.13.45
 boto3~=1.10.45
 s3fs~=0.4.0
 tenacity~=6.0.0
-pg8000~=1.13.2
+pg8000~=1.13.2
+pymysql~=0.9.3