Add secretmanager module and support for databases connections. #402

Author: igorborgest

README.md

@@ -12,7 +12,7 @@
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 
 [![Checked with mypy](http://www.mypy-lang.org/static/mypy_badge.svg)](http://mypy-lang.org/)
-[![Coverage](https://img.shields.io/badge/coverage-92%25-brightgreen.svg)](https://pypi.org/project/awswrangler/)
+[![Coverage](https://img.shields.io/badge/coverage-91%25-brightgreen.svg)](https://pypi.org/project/awswrangler/)
 ![Static Checking](https://github.com/awslabs/aws-data-wrangler/workflows/Static%20Checking/badge.svg?branch=master)
 [![Documentation Status](https://readthedocs.org/projects/aws-data-wrangler/badge/?version=latest)](https://aws-data-wrangler.readthedocs.io/?badge=latest)
 
@@ -138,6 +138,7 @@ FROM "sampleDB"."sampleTable" ORDER BY time DESC LIMIT 3
   - [Amazon CloudWatch Logs](https://aws-data-wrangler.readthedocs.io/en/stable/api.html#amazon-cloudwatch-logs)
   - [Amazon QuickSight](https://aws-data-wrangler.readthedocs.io/en/stable/api.html#amazon-quicksight)
   - [AWS STS](https://aws-data-wrangler.readthedocs.io/en/stable/api.html#aws-sts)
+  - [AWS Secrets Manager](https://aws-data-wrangler.readthedocs.io/en/stable/api.html#aws-secrets-manager)
 - [**License**](https://github.com/awslabs/aws-data-wrangler/blob/master/LICENSE.txt)
 - [**Contributing**](https://github.com/awslabs/aws-data-wrangler/blob/master/CONTRIBUTING.md)
 - [**Legacy Docs** (pre-1.0.0)](https://aws-data-wrangler.readthedocs.io/en/0.3.3/)
@@ -202,6 +203,6 @@ Please [send a Pull Request](https://github.com/awslabs/aws-data-wrangler/edit/m
 
 **Amazon SageMaker Data Wrangler** is a new SageMaker Studio feature that has a similar name but has a different purpose than the **AWS Data Wrangler** open source project.
 
-* **AWS Data Wrangler** is open source, runs anywhere, and is focused on code.
+- **AWS Data Wrangler** is open source, runs anywhere, and is focused on code.
 
-* **Amazon SageMaker Data Wrangler** is specific for the SageMaker Studio environment and is focused on a visual interface.
+- **Amazon SageMaker Data Wrangler** is specific for the SageMaker Studio environment and is focused on a visual interface.

awswrangler/__init__.py

@@ -18,6 +18,7 @@
     quicksight,
     redshift,
     s3,
+    secretsmanager,
     sts,
     timestream,
 )
@@ -36,6 +37,7 @@
     "redshift",
     "mysql",
     "postgresql",
+    "secretsmanager",
     "config",
     "timestream",
     "__description__",

awswrangler/_databases.py

@@ -8,7 +8,7 @@
 import pandas as pd
 import pyarrow as pa
 
-from awswrangler import _data_types
+from awswrangler import _data_types, _utils, exceptions, secretsmanager
 from awswrangler.catalog import get_connection
 
 _logger: logging.Logger = logging.getLogger(__name__)
@@ -25,12 +25,15 @@ class ConnectionAttributes(NamedTuple):
     database: str
 
 
-def get_connection_attributes(
-    connection: str,
-    catalog_id: Optional[str] = None,
-    boto3_session: Optional[boto3.Session] = None,
+def _get_dbname(cluster_id: str, boto3_session: Optional[boto3.Session] = None) -> str:
+    client_redshift: boto3.client = _utils.client(service_name="redshift", session=boto3_session)
+    res: Dict[str, Any] = client_redshift.describe_clusters(ClusterIdentifier=cluster_id)["Clusters"][0]
+    return cast(str, res["DBName"])
+
+
+def _get_connection_attributes_from_catalog(
+    connection: str, catalog_id: Optional[str], dbname: Optional[str], boto3_session: Optional[boto3.Session]
 ) -> ConnectionAttributes:
-    """Get Connection Attributes."""
     details: Dict[str, Any] = get_connection(name=connection, catalog_id=catalog_id, boto3_session=boto3_session)[
         "ConnectionProperties"
     ]
@@ -41,7 +44,51 @@ def get_connection_attributes(
         password=quote_plus(details["PASSWORD"]),
         host=details["JDBC_CONNECTION_URL"].split(":")[2].replace("/", ""),
         port=int(port),
-        database=database,
+        database=dbname if dbname is not None else database,
+    )
+
+
+def _get_connection_attributes_from_secrets_manager(
+    secret_id: str, dbname: Optional[str], boto3_session: Optional[boto3.Session]
+) -> ConnectionAttributes:
+    secret_value: Dict[str, Any] = secretsmanager.get_secret_json(name=secret_id, boto3_session=boto3_session)
+    kind: str = secret_value["engine"]
+    if dbname is not None:
+        _dbname: str = dbname
+    elif "dbname" in secret_value:
+        _dbname = secret_value["dbname"]
+    else:
+        if kind != "redshift":
+            raise exceptions.InvalidConnection(f"The secret {secret_id} MUST have a dbname property.")
+        _dbname = _get_dbname(cluster_id=secret_value["dbClusterIdentifier"], boto3_session=boto3_session)
+    return ConnectionAttributes(
+        kind=kind,
+        user=quote_plus(secret_value["username"]),
+        password=quote_plus(secret_value["password"]),
+        host=secret_value["host"],
+        port=secret_value["port"],
+        database=_dbname,
+    )
+
+
+def get_connection_attributes(
+    connection: Optional[str] = None,
+    secret_id: Optional[str] = None,
+    catalog_id: Optional[str] = None,
+    dbname: Optional[str] = None,
+    boto3_session: Optional[boto3.Session] = None,
+) -> ConnectionAttributes:
+    """Get Connection Attributes."""
+    if connection is None and secret_id is None:
+        raise exceptions.InvalidArgumentCombination(
+            "Failed attempt to connect. You MUST pass a connection name (Glue Catalog) OR a secret_id as argument."
+        )
+    if connection is not None:
+        return _get_connection_attributes_from_catalog(
+            connection=connection, catalog_id=catalog_id, dbname=dbname, boto3_session=boto3_session
+        )
+    return _get_connection_attributes_from_secrets_manager(
+        secret_id=cast(str, secret_id), dbname=dbname, boto3_session=boto3_session
     )
 
 

awswrangler/mysql.py

@@ -67,8 +67,10 @@ def _create_table(
 
 
 def connect(
-    connection: str,
+    connection: Optional[str] = None,
+    secret_id: Optional[str] = None,
     catalog_id: Optional[str] = None,
+    dbname: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
     read_timeout: Optional[int] = None,
     write_timeout: Optional[int] = None,
@@ -82,9 +84,14 @@ def connect(
     ----------
     connection : str
         Glue Catalog Connection name.
+    secret_id: Optional[str]:
+        Specifies the secret containing the version that you want to retrieve.
+        You can specify either the Amazon Resource Name (ARN) or the friendly name of the secret.
     catalog_id : str, optional
         The ID of the Data Catalog.
         If none is provided, the AWS account ID is used by default.
+    dbname: Optional[str]
+        Optional database name to overwrite the stored one.
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
     read_timeout: Optional[int]
@@ -117,7 +124,7 @@ def connect(
 
     """
     attrs: _db_utils.ConnectionAttributes = _db_utils.get_connection_attributes(
-        connection=connection, catalog_id=catalog_id, boto3_session=boto3_session
+        connection=connection, secret_id=secret_id, catalog_id=catalog_id, dbname=dbname, boto3_session=boto3_session
     )
     if attrs.kind != "mysql":
         exceptions.InvalidDatabaseType(f"Invalid connection type ({attrs.kind}. It must be a MySQL connection.)")

awswrangler/postgresql.py

@@ -71,8 +71,10 @@ def _create_table(
 
 
 def connect(
-    connection: str,
+    connection: Optional[str] = None,
+    secret_id: Optional[str] = None,
     catalog_id: Optional[str] = None,
+    dbname: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
     ssl_context: Optional[Dict[Any, Any]] = None,
     timeout: Optional[int] = None,
@@ -84,11 +86,16 @@ def connect(
 
     Parameters
     ----------
-    connection : str
+    connection : Optional[str]
         Glue Catalog Connection name.
+    secret_id: Optional[str]:
+        Specifies the secret containing the version that you want to retrieve.
+        You can specify either the Amazon Resource Name (ARN) or the friendly name of the secret.
     catalog_id : str, optional
         The ID of the Data Catalog.
         If none is provided, the AWS account ID is used by default.
+    dbname: Optional[str]
+        Optional database name to overwrite the stored one.
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
     ssl_context: Optional[Dict]
@@ -121,7 +128,7 @@ def connect(
 
     """
     attrs: _db_utils.ConnectionAttributes = _db_utils.get_connection_attributes(
-        connection=connection, catalog_id=catalog_id, boto3_session=boto3_session
+        connection=connection, secret_id=secret_id, catalog_id=catalog_id, dbname=dbname, boto3_session=boto3_session
     )
     if attrs.kind != "postgresql":
         exceptions.InvalidDatabaseType(f"Invalid connection type ({attrs.kind}. It must be a postgresql connection.)")

awswrangler/redshift.py

@@ -262,25 +262,37 @@ def _read_parquet_iterator(
 
 
 def connect(
-    connection: str,
+    connection: Optional[str] = None,
+    secret_id: Optional[str] = None,
     catalog_id: Optional[str] = None,
+    dbname: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
     ssl: bool = True,
     timeout: Optional[int] = None,
     max_prepared_statements: int = 1000,
     tcp_keepalive: bool = True,
 ) -> redshift_connector.Connection:
-    """Return a redshift_connector connection from a Glue Catalog Connection.
+    """Return a redshift_connector connection from a Glue Catalog or Secret Manager.
+
+    Note
+    ----
+    You MUST pass a `connection` OR `secret_id`
+
 
     https://github.com/aws/amazon-redshift-python-driver
 
     Parameters
     ----------
-    connection : str
+    connection : Optional[str]
         Glue Catalog Connection name.
+    secret_id: Optional[str]:
+        Specifies the secret containing the version that you want to retrieve.
+        You can specify either the Amazon Resource Name (ARN) or the friendly name of the secret.
     catalog_id : str, optional
         The ID of the Data Catalog.
         If none is provided, the AWS account ID is used by default.
+    dbname: Optional[str]
+        Optional database name to overwrite the stored one.
     boto3_session : boto3.Session(), optional
         Boto3 Session. The default boto3 session will be used if boto3_session receive None.
     ssl: bool
@@ -307,16 +319,27 @@ def connect(
 
     Examples
     --------
+    Fetching Redshit connection from Glue Catalog
+
     >>> import awswrangler as wr
     >>> con = wr.redshift.connect("MY_GLUE_CONNECTION")
     >>> with con.cursor() as cursor:
     >>>     cursor.execute("SELECT 1")
     >>>     print(cursor.fetchall())
     >>> con.close()
 
+    Fetching Redshit connection from Secrets Manager
+
+    >>> import awswrangler as wr
+    >>> con = wr.redshift.connect(secret_id="MY_SECRET")
+    >>> with con.cursor() as cursor:
+    >>>     cursor.execute("SELECT 1")
+    >>>     print(cursor.fetchall())
+    >>> con.close()
+
     """
     attrs: _db_utils.ConnectionAttributes = _db_utils.get_connection_attributes(
-        connection=connection, catalog_id=catalog_id, boto3_session=boto3_session
+        connection=connection, secret_id=secret_id, catalog_id=catalog_id, dbname=dbname, boto3_session=boto3_session
     )
     if attrs.kind != "redshift":
         exceptions.InvalidDatabaseType(f"Invalid connection type ({attrs.kind}. It must be a redshift connection.)")

awswrangler/secretsmanager.py

@@ -0,0 +1,68 @@
+"""Secrets Manager module."""
+
+import base64
+import json
+import logging
+from typing import Any, Dict, Optional, Union, cast
+
+import boto3
+
+from awswrangler import _utils
+
+_logger: logging.Logger = logging.getLogger(__name__)
+
+
+def get_secret(name: str, boto3_session: Optional[boto3.Session] = None) -> Union[str, bytes]:
+    """Get secret value.
+
+    Parameters
+    ----------
+    name: str:
+        Specifies the secret containing the version that you want to retrieve.
+        You can specify either the Amazon Resource Name (ARN) or the friendly name of the secret.
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+
+    Returns
+    -------
+    Union[str, bytes]
+        Secret value.
+
+    Examples
+    --------
+    >>> import awswrangler as wr
+    >>> value = wr.secretsmanager.get_secret("my-secret")
+
+    """
+    session: boto3.Session = _utils.ensure_session(session=boto3_session)
+    client: boto3.client = _utils.client(service_name="secretsmanager", session=session)
+    response: Dict[str, Any] = client.get_secret_value(SecretId=name)
+    if "SecretString" in response:
+        return cast(str, response["SecretString"])
+    return base64.b64decode(response["SecretBinary"])
+
+
+def get_secret_json(name: str, boto3_session: Optional[boto3.Session] = None) -> Dict[str, Any]:
+    """Get JSON secret value.
+
+    Parameters
+    ----------
+    name: str:
+        Specifies the secret containing the version that you want to retrieve.
+        You can specify either the Amazon Resource Name (ARN) or the friendly name of the secret.
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+
+    Returns
+    -------
+    Dict[str, Any]
+        Secret JSON value parsed as a dictionary.
+
+    Examples
+    --------
+    >>> import awswrangler as wr
+    >>> value = wr.secretsmanager.get_secret_json("my-secret-with-json-content")
+
+    """
+    value: Union[str, bytes] = get_secret(name=name, boto3_session=boto3_session)
+    return cast(Dict[str, Any], json.loads(value))