Add basic support for Redshift and RDS Data APIs (#828)

* Add basic outline of redshift data api connector

* Add outline of support for pagination

* Restructure data_api module and add basic RDS support

* Remove unused variable

* Add secret ARNs to CDK outputs and databases_parameters fixture

* Add RDS resource identifiers to CDK outputs and databases_parameters fixture

* Pylint fixes

* Add basic outline for full infrastructure data api tests

* Fix CDK typo

* Add Aurora serverless to test resources to allow RDS Data API testing

* Update data api integration tests based on new serverless Aurora instance

* Fix issues based on initial integration tests

* Finish first working integration tests for RDS and Redshift data api

* Refactor code and tests, and add logging

* Remove unused mysql CDK exports

* Fix documentation style

* Remove use of postponed type evaluations

Co-authored-by: kukushking <[email protected]>

Author: pwithams

awswrangler/__init__.py

@@ -12,6 +12,7 @@
     catalog,
     chime,
     cloudwatch,
+    data_api,
     dynamodb,
     emr,
     exceptions,
@@ -34,6 +35,7 @@
     "chime",
     "cloudwatch",
     "emr",
+    "data_api",
     "dynamodb",
     "exceptions",
     "quicksight",

awswrangler/data_api/__init__.py

@@ -0,0 +1,7 @@
+"""Data API Service Module for RDS and Redshift."""
+from awswrangler.data_api import rds, redshift
+
+__all__ = [
+    "redshift",
+    "rds",
+]

awswrangler/data_api/connector.py

@@ -0,0 +1,71 @@
+"""Data API Connector base class."""
+import logging
+from typing import Any, Dict, Optional
+
+import pandas as pd
+
+
+class DataApiConnector:
+    """Base class for Data API (RDS, Redshift, etc.) connectors."""
+
+    def __init__(self, client: Any, logger: logging.Logger):
+        self.client = client
+        self.logger: logging.Logger = logger
+
+    def execute(self, sql: str, database: Optional[str] = None) -> pd.DataFrame:
+        """Execute SQL statement against a Data API Service.
+
+        Parameters
+        ----------
+        sql: str
+            SQL statement to execute.
+
+        Returns
+        -------
+        A Pandas DataFrame containing the execution results.
+        """
+        request_id: str = self._execute_statement(sql, database=database)
+        return self._get_statement_result(request_id)
+
+    def _execute_statement(self, sql: str, database: Optional[str] = None) -> str:
+        raise NotImplementedError()
+
+    def _get_statement_result(self, request_id: str) -> pd.DataFrame:
+        raise NotImplementedError()
+
+    @staticmethod
+    def _get_column_value(column_value: Dict[str, Any]) -> Any:
+        """Return the first non-null key value for a given dictionary.
+
+        The key names for a given record depend on the column type: stringValue, longValue, etc.
+
+        Therefore, a record in the response does not have consistent key names. The ColumnMetadata
+        typeName information could be used to infer the key, but there is no direct mapping here
+        that could be easily parsed with creating a static dictionary:
+            varchar -> stringValue
+            int2 -> longValue
+            timestamp -> stringValue
+
+        What has been observed is that each record appears to have a single key, so this function
+        iterates over the keys and returns the first non-null value. If none are found, None is
+        returned.
+
+        Documentation:
+            https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/redshift-data.html#RedshiftDataAPIService.Client.get_statement_result
+            https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/rds-data.html#RDSDataService.Client.execute_statement
+        """
+        for key in column_value:
+            if column_value[key] is not None:
+                if key == "arrayValue":
+                    raise ValueError(f"arrayValue not supported yet - could not extract {column_value[key]}")
+                return column_value[key]
+        return None
+
+
+class WaitConfig:
+    """Holds standard wait configuration values."""
+
+    def __init__(self, sleep: float, backoff: float, retries: int) -> None:
+        self.sleep = sleep
+        self.backoff = backoff
+        self.retries = retries

awswrangler/data_api/rds.py

@@ -0,0 +1,149 @@
+"""RDS Data API Connector."""
+import logging
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+import boto3
+import pandas as pd
+
+from awswrangler.data_api import connector
+
+
+class RdsDataApi(connector.DataApiConnector):
+    """Provides access to the RDS Data API.
+
+    Parameters
+    ----------
+    resource_arn: str
+        ARN for the RDS resource.
+    database: str
+        Target database name.
+    secret_arn: str
+        The ARN for the secret to be used for authentication.
+    sleep: float
+        Number of seconds to sleep between connection attempts to paused clusters - defaults to 0.5.
+    backoff: float
+        Factor by which to increase the sleep between connection attempts to paused clusters - defaults to 1.0.
+    retries: int
+        Maximum number of connection attempts to paused clusters - defaults to 10.
+    """
+
+    def __init__(
+        self,
+        resource_arn: str,
+        database: str,
+        secret_arn: str = "",
+        sleep: float = 0.5,
+        backoff: float = 1.0,
+        retries: int = 30,
+    ) -> None:
+        self.resource_arn = resource_arn
+        self.database = database
+        self.secret_arn = secret_arn
+        self.wait_config = connector.WaitConfig(sleep, backoff, retries)
+        self.client = boto3.client("rds-data")
+        self.results: Dict[str, Dict[str, Any]] = {}
+        logger: logging.Logger = logging.getLogger(__name__)
+        super().__init__(self.client, logger)
+
+    def _execute_statement(self, sql: str, database: Optional[str] = None) -> str:
+        if database is None:
+            database = self.database
+
+        sleep: float = self.wait_config.sleep
+        total_tries: int = 0
+        total_sleep: float = 0
+        response: Optional[Dict[str, Any]] = None
+        last_exception: Optional[Exception] = None
+        while total_tries < self.wait_config.retries:
+            try:
+                response = self.client.execute_statement(
+                    resourceArn=self.resource_arn,
+                    database=database,
+                    sql=sql,
+                    secretArn=self.secret_arn,
+                    includeResultMetadata=True,
+                )
+                self.logger.debug(
+                    "Response received after %s tries and sleeping for a total of %s seconds", total_tries, total_sleep
+                )
+                break
+            except self.client.exceptions.BadRequestException as exception:
+                last_exception = exception
+                total_sleep += sleep
+                self.logger.debug("BadRequestException occurred: %s", exception)
+                self.logger.debug(
+                    "Cluster may be paused - sleeping for %s seconds for a total of %s before retrying",
+                    sleep,
+                    total_sleep,
+                )
+                time.sleep(sleep)
+                total_tries += 1
+                sleep *= self.wait_config.backoff
+
+        if response is None:
+            self.logger.exception("Maximum BadRequestException retries reached for query %s", sql)
+            raise self.client.exceptions.BadRequestException(
+                f"Query failed - BadRequestException received after {total_tries} tries and sleeping {total_sleep}s"
+            ) from last_exception
+
+        request_id: str = uuid.uuid4().hex
+        self.results[request_id] = response
+        return request_id
+
+    def _get_statement_result(self, request_id: str) -> pd.DataFrame:
+        try:
+            result = self.results.pop(request_id)
+        except KeyError as exception:
+            raise KeyError(f"Request {request_id} not found in results {self.results}") from exception
+
+        if "records" not in result:
+            return pd.DataFrame()
+
+        rows: List[List[Any]] = []
+        for record in result["records"]:
+            row: List[Any] = [connector.DataApiConnector._get_column_value(column) for column in record]
+            rows.append(row)
+
+        column_names: List[str] = [column["name"] for column in result["columnMetadata"]]
+        dataframe = pd.DataFrame(rows, columns=column_names)
+        return dataframe
+
+
+def connect(resource_arn: str, database: str, secret_arn: str = "", **kwargs: Any) -> RdsDataApi:
+    """Create a RDS Data API connection.
+
+    Parameters
+    ----------
+    resource_arn: str
+        ARN for the RDS resource.
+    database: str
+        Target database name.
+    secret_arn: str
+        The ARN for the secret to be used for authentication.
+    **kwargs
+        Any additional kwargs are passed to the underlying RdsDataApi class.
+
+    Returns
+    -------
+    A RdsDataApi connection instance that can be used with `wr.rds.data_api.read_sql_query`.
+    """
+    return RdsDataApi(resource_arn, database, secret_arn=secret_arn, **kwargs)
+
+
+def read_sql_query(sql: str, con: RdsDataApi, database: Optional[str] = None) -> pd.DataFrame:
+    """Run an SQL query on an RdsDataApi connection and return the result as a dataframe.
+
+    Parameters
+    ----------
+    sql: str
+        SQL query to run.
+    database: str
+        Database to run query on - defaults to the database specified by `con`.
+
+    Returns
+    -------
+    A Pandas dataframe containing the query results.
+    """
+    return con.execute(sql, database=database)