Add generate_connection() to Aurora

Author: igorborgest

awswrangler/aurora.py

@@ -1,8 +1,103 @@
+from typing import Union
 import logging
 
+import pg8000  # type: ignore
+import pymysql  # type: ignore
+
+from awswrangler.exceptions import InvalidEngine
+
 logger = logging.getLogger(__name__)
 
 
 class Aurora:
     def __init__(self, session):
         self._session = session
+
+    @staticmethod
+    def _validate_connection(database: str,
+                             host: str,
+                             port: Union[str, int],
+                             user: str,
+                             password: str,
+                             engine: str = "mysql",
+                             tcp_keepalive: bool = True,
+                             application_name: str = "aws-data-wrangler-validation",
+                             validation_timeout: int = 5) -> None:
+        if "postgres" in engine.lower():
+            conn = pg8000.connect(database=database,
+                                  host=host,
+                                  port=int(port),
+                                  user=user,
+                                  password=password,
+                                  ssl=True,
+                                  application_name=application_name,
+                                  tcp_keepalive=tcp_keepalive,
+                                  timeout=validation_timeout)
+        elif "mysql" in engine.lower():
+            conn = pymysql.connect(database=database,
+                                   host=host,
+                                   port=int(port),
+                                   user=user,
+                                   password=password,
+                                   program_name=application_name,
+                                   connect_timeout=validation_timeout)
+        else:
+            raise InvalidEngine(f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!")
+        conn.close()
+
+    @staticmethod
+    def generate_connection(database: str,
+                            host: str,
+                            port: Union[str, int],
+                            user: str,
+                            password: str,
+                            engine: str = "mysql",
+                            tcp_keepalive: bool = True,
+                            application_name: str = "aws-data-wrangler",
+                            connection_timeout: int = 1_200_000,
+                            validation_timeout: int = 5):
+        """
+        Generates a valid connection object
+
+        :param database: The name of the database instance to connect with.
+        :param host: The hostname of the Aurora server to connect with.
+        :param port: The TCP/IP port of the Aurora server instance.
+        :param user: The username to connect to the Aurora database with.
+        :param password: The user password to connect to the server with.
+        :param engine: "mysql" or "postgres"
+        :param tcp_keepalive: If True then use TCP keepalive
+        :param application_name: Application name
+        :param connection_timeout: Connection Timeout
+        :param validation_timeout: Timeout to try to validate the connection
+        :return: PEP 249 compatible connection
+        """
+        Aurora._validate_connection(database=database,
+                                    host=host,
+                                    port=port,
+                                    user=user,
+                                    password=password,
+                                    engine=engine,
+                                    tcp_keepalive=tcp_keepalive,
+                                    application_name=application_name,
+                                    validation_timeout=validation_timeout)
+        if "postgres" in engine.lower():
+            conn = pg8000.connect(database=database,
+                                  host=host,
+                                  port=int(port),
+                                  user=user,
+                                  password=password,
+                                  ssl=True,
+                                  application_name=application_name,
+                                  tcp_keepalive=tcp_keepalive,
+                                  timeout=connection_timeout)
+        elif "mysql" in engine.lower():
+            conn = pymysql.connect(database=database,
+                                   host=host,
+                                   port=int(port),
+                                   user=user,
+                                   password=password,
+                                   program_name=application_name,
+                                   connect_timeout=connection_timeout)
+        else:
+            raise InvalidEngine(f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!")
+        return conn

awswrangler/exceptions.py

@@ -96,3 +96,7 @@ class InvalidParameters(Exception):
 
 class AWSCredentialsNotFound(Exception):
     pass
+
+
+class InvalidEngine(Exception):
+    pass

awswrangler/pandas.py

@@ -1433,3 +1433,43 @@ def read_sql_redshift(self,
             else:
                 self._session.s3.delete_objects(path=temp_s3_path)
             raise e
+
+    def read_sql_aurora(self,
+                        sql: str,
+                        iam_role: str,
+                        connection: Any,
+                        temp_s3_path: Optional[str] = None) -> pd.DataFrame:
+        """
+        Convert a query result in a Pandas Dataframe.
+
+        :param sql: SQL Query
+        :param iam_role: AWS IAM role with the related permissions
+        :param connection: A PEP 249 compatible connection (Can be generated with Aurora.generate_connection())
+        :param temp_s3_path: AWS S3 path to write temporary data (e.g. s3://...) (Default uses the Athena's results bucket)
+        """
+        guid: str = pa.compat.guid()
+        name: str = f"temp_redshift_{guid}"
+        if temp_s3_path is None:
+            if self._session.athena_s3_output is not None:
+                temp_s3_path = self._session.redshift_temp_s3_path
+            else:
+                temp_s3_path = self._session.athena.create_athena_bucket()
+        temp_s3_path = temp_s3_path[:-1] if temp_s3_path[-1] == "/" else temp_s3_path
+        temp_s3_path = f"{temp_s3_path}/{name}"
+        logger.debug(f"temp_s3_path: {temp_s3_path}")
+        paths: Optional[List[str]] = None
+        try:
+            paths = self._session.redshift.to_parquet(sql=sql,
+                                                      path=temp_s3_path,
+                                                      iam_role=iam_role,
+                                                      connection=connection)
+            logger.debug(f"paths: {paths}")
+            df: pd.DataFrame = self.read_parquet(path=paths)  # type: ignore
+            self._session.s3.delete_listed_objects(objects_paths=paths)
+            return df
+        except Exception as e:
+            if paths is not None:
+                self._session.s3.delete_listed_objects(objects_paths=paths)
+            else:
+                self._session.s3.delete_objects(path=temp_s3_path)
+            raise e

awswrangler/redshift.py

@@ -38,19 +38,16 @@ def _validate_connection(database,
                              tcp_keepalive=True,
                              application_name="aws-data-wrangler-validation",
                              validation_timeout=5):
-        try:
-            conn = pg8000.connect(database=database,
-                                  host=host,
-                                  port=int(port),
-                                  user=user,
-                                  password=password,
-                                  ssl=True,
-                                  application_name=application_name,
-                                  tcp_keepalive=tcp_keepalive,
-                                  timeout=validation_timeout)
-            conn.close()
-        except pg8000.core.InterfaceError as e:
-            raise e
+        conn = pg8000.connect(database=database,
+                              host=host,
+                              port=int(port),
+                              user=user,
+                              password=password,
+                              ssl=True,
+                              application_name=application_name,
+                              tcp_keepalive=tcp_keepalive,
+                              timeout=validation_timeout)
+        conn.close()
 
     @staticmethod
     def generate_connection(database,
@@ -86,8 +83,6 @@ def generate_connection(database,
                                       tcp_keepalive=tcp_keepalive,
                                       application_name=application_name,
                                       validation_timeout=validation_timeout)
-        if isinstance(type(port), str) or isinstance(type(port), float):
-            port = int(port)
         conn = pg8000.connect(database=database,
                               host=host,
                               port=int(port),

awswrangler/sagemaker.py

@@ -12,7 +12,9 @@ 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", use_ssl=True, config=session.botocore_config)
+        self._client_sagemaker = session.boto3_session.client(service_name="sagemaker",
+                                                              use_ssl=True,
+                                                              config=session.botocore_config)
 
     @staticmethod
     def _parse_path(path):

testing/parameters.properties

@@ -1,4 +1,5 @@
 VpcId=VPC_ID
 SubnetId=SUBNET_ID
+SubnetId2=SUBNET_ID2
 Password=REDSHIFT_PASSWORD
 TestUser=AWS_USER_THAT_WILL_RUN_THE_TESTS_ON_CLI

testing/template.yaml

@@ -234,9 +234,9 @@ Outputs:
   RedshiftPort:
     Value: !GetAtt Redshift.Endpoint.Port
     Description: Redshift Endpoint Port.
-  RedshiftPassword:
+  Password:
     Value: !Ref Password
-    Description: Redshift Password.
+    Description: Password.
   RedshiftRole:
     Value: !GetAtt RedshiftRole.Arn
     Description: Redshift IAM role.

testing/test_awswrangler/test_athena.py

@@ -15,7 +15,7 @@
 
 @pytest.fixture(scope="module")
 def cloudformation_outputs():
-    response = boto3.client("cloudformation").describe_stacks(StackName="aws-data-wrangler-test-arena")
+    response = boto3.client("cloudformation").describe_stacks(StackName="aws-data-wrangler-test")
     outputs = {}
     for output in response.get("Stacks")[0].get("Outputs"):
         outputs[output.get("OutputKey")] = output.get("OutputValue")

testing/test_awswrangler/test_aurora.py

@@ -0,0 +1,89 @@
+import logging
+
+import pytest
+import boto3
+
+from awswrangler import Aurora
+from awswrangler.exceptions import InvalidEngine
+
+logging.basicConfig(level=logging.INFO, format="[%(asctime)s][%(levelname)s][%(name)s][%(funcName)s] %(message)s")
+logging.getLogger("awswrangler").setLevel(logging.DEBUG)
+
+
+@pytest.fixture(scope="module")
+def cloudformation_outputs():
+    response = boto3.client("cloudformation").describe_stacks(StackName="aws-data-wrangler-test")
+    outputs = {}
+    for output in response.get("Stacks")[0].get("Outputs"):
+        outputs[output.get("OutputKey")] = output.get("OutputValue")
+    yield outputs
+
+
+@pytest.fixture(scope="module")
+def postgres_parameters(cloudformation_outputs):
+    postgres_parameters = {}
+    if "PostgresAddress" in cloudformation_outputs:
+        postgres_parameters["PostgresAddress"] = cloudformation_outputs.get("PostgresAddress")
+    else:
+        raise Exception("You must deploy the test infrastructure using SAM!")
+    if "Password" in cloudformation_outputs:
+        postgres_parameters["Password"] = cloudformation_outputs.get("Password")
+    else:
+        raise Exception("You must deploy the test infrastructure using SAM!")
+    yield postgres_parameters
+
+
+@pytest.fixture(scope="module")
+def mysql_parameters(cloudformation_outputs):
+    mysql_parameters = {}
+    if "MysqlAddress" in cloudformation_outputs:
+        mysql_parameters["MysqlAddress"] = cloudformation_outputs.get("MysqlAddress")
+    else:
+        raise Exception("You must deploy the test infrastructure using SAM!")
+    if "Password" in cloudformation_outputs:
+        mysql_parameters["Password"] = cloudformation_outputs.get("Password")
+    else:
+        raise Exception("You must deploy the test infrastructure using SAM!")
+    yield mysql_parameters
+
+
+def test_postgres_connection(postgres_parameters):
+    conn = Aurora.generate_connection(database="postgres",
+                                      host=postgres_parameters["PostgresAddress"],
+                                      port=3306,
+                                      user="test",
+                                      password=postgres_parameters["Password"],
+                                      engine="postgres")
+    cursor = conn.cursor()
+    cursor.execute("SELECT 1 + 2, 3 + 4")
+    first_row = cursor.fetchall()[0]
+    assert first_row[0] == 3
+    assert first_row[1] == 7
+    cursor.close()
+    conn.close()
+
+
+def test_mysql_connection(mysql_parameters):
+    conn = Aurora.generate_connection(database="mysql",
+                                      host=mysql_parameters["MysqlAddress"],
+                                      port=3306,
+                                      user="test",
+                                      password=mysql_parameters["Password"],
+                                      engine="mysql")
+    cursor = conn.cursor()
+    cursor.execute("SELECT 1 + 2, 3 + 4")
+    first_row = cursor.fetchall()[0]
+    assert first_row[0] == 3
+    assert first_row[1] == 7
+    cursor.close()
+    conn.close()
+
+
+def test_invalid_engine(mysql_parameters):
+    with pytest.raises(InvalidEngine):
+        Aurora.generate_connection(database="mysql",
+                                   host=mysql_parameters["MysqlAddress"],
+                                   port=3306,
+                                   user="test",
+                                   password=mysql_parameters["Password"],
+                                   engine="foo")

testing/test_awswrangler/test_cloudwatchlogs.py

@@ -13,7 +13,7 @@
 
 @pytest.fixture(scope="module")
 def cloudformation_outputs():
-    response = boto3.client("cloudformation").describe_stacks(StackName="aws-data-wrangler-test-arena")
+    response = boto3.client("cloudformation").describe_stacks(StackName="aws-data-wrangler-test")
     outputs = {}
     for output in response.get("Stacks")[0].get("Outputs"):
         outputs[output.get("OutputKey")] = output.get("OutputValue")