awswrangler/redshift/_connect.py

"""Amazon Redshift Connect Module (PRIVATE)."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

import boto3

from awswrangler import _databases as _db_utils
from awswrangler import _utils, exceptions

if TYPE_CHECKING:
    try:
        import redshift_connector
    except ImportError:
        pass
else:
    redshift_connector = _utils.import_optional_dependency("redshift_connector")


def _validate_connection(con: "redshift_connector.Connection") -> None:
    if not isinstance(con, redshift_connector.Connection):
        raise exceptions.InvalidConnection(
            "Invalid 'conn' argument, please pass a "
            "redshift_connector.Connection object. Use redshift_connector.connect() to use "
            "credentials directly or wr.redshift.connect() to fetch it from the Glue Catalog."
        )


@_utils.check_optional_dependency(redshift_connector, "redshift_connector")
def connect(
    connection: str | None = None,
    secret_id: str | None = None,
    catalog_id: str | None = None,
    dbname: str | None = None,
    boto3_session: boto3.Session | None = None,
    ssl: bool = True,
    timeout: int | None = None,
    max_prepared_statements: int = 1000,
    tcp_keepalive: bool = True,
    **kwargs: Any,
) -> "redshift_connector.Connection":
    """Return a redshift_connector connection from a Glue Catalog or Secret Manager.

    Note
    ----
    You MUST pass a `connection` OR `secret_id`.
    Here is an example of the secret structure in Secrets Manager:
    {
    "host":"my-host.us-east-1.redshift.amazonaws.com",
    "username":"test",
    "password":"test",
    "engine":"redshift",
    "port":"5439",
    "dbname": "mydb"
    }


    https://github.com/aws/amazon-redshift-python-driver

    Parameters
    ----------
    connection
        Glue Catalog Connection name.
    secret_id
        Specifies the secret containing the connection details that you want to retrieve.
        You can specify either the Amazon Resource Name (ARN) or the friendly name of the secret.
    catalog_id
        The ID of the Data Catalog.
        If none is provided, the AWS account ID is used by default.
    dbname
        Optional database name to overwrite the stored one.
    boto3_session
        The default boto3 session will be used if **boto3_session** is ``None``.
    ssl
        This governs SSL encryption for TCP/IP sockets.
        This parameter is forward to redshift_connector.
        https://github.com/aws/amazon-redshift-python-driver
    timeout
        This is the time in seconds before the connection to the server will time out.
        The default is None which means no timeout.
        This parameter is forward to redshift_connector.
        https://github.com/aws/amazon-redshift-python-driver
    max_prepared_statements
        This parameter is forward to redshift_connector.
        https://github.com/aws/amazon-redshift-python-driver
    tcp_keepalive
        If True then use TCP keepalive. The default is True.
        This parameter is forward to redshift_connector.
        https://github.com/aws/amazon-redshift-python-driver
    **kwargs
        Forwarded to redshift_connector.connect.
        e.g. ``is_serverless=True, serverless_acct_id='...', serverless_work_group='...'``

    Returns
    -------
        ``redshift_connector`` connection.

    Examples
    --------
    Fetching Redshift connection from Glue Catalog

    >>> import awswrangler as wr
    >>> with wr.redshift.connect("MY_GLUE_CONNECTION") as con:
    ...     with con.cursor() as cursor:
    ...         cursor.execute("SELECT 1")
    ...         print(cursor.fetchall())

    Fetching Redshift connection from Secrets Manager

    >>> import awswrangler as wr
    >>> with wr.redshift.connect(secret_id="MY_SECRET") as con:
    ...     with con.cursor() as cursor:
    ...         cursor.execute("SELECT 1")
    ...         print(cursor.fetchall())

    """
    attrs: _db_utils.ConnectionAttributes = _db_utils.get_connection_attributes(
        connection=connection, secret_id=secret_id, catalog_id=catalog_id, dbname=dbname, boto3_session=boto3_session
    )
    if attrs.kind != "redshift":
        raise exceptions.InvalidDatabaseType(
            f"Invalid connection type ({attrs.kind}. It must be a redshift connection.)"
        )
    return redshift_connector.connect(
        user=attrs.user,
        database=attrs.database,
        password=attrs.password,
        port=int(attrs.port),
        host=attrs.host,
        ssl=ssl,
        timeout=timeout,
        max_prepared_statements=max_prepared_statements,
        tcp_keepalive=tcp_keepalive,
        **kwargs,
    )


@_utils.check_optional_dependency(redshift_connector, "redshift_connector")
def connect_temp(
    cluster_identifier: str,
    user: str,
    database: str | None = None,
    duration: int = 900,
    auto_create: bool = True,
    db_groups: list[str] | None = None,
    boto3_session: boto3.Session | None = None,
    ssl: bool = True,
    timeout: int | None = None,
    max_prepared_statements: int = 1000,
    tcp_keepalive: bool = True,
    **kwargs: Any,
) -> "redshift_connector.Connection":
    """Return a redshift_connector temporary connection (No password required).

    https://github.com/aws/amazon-redshift-python-driver

    Parameters
    ----------
    cluster_identifier
        The unique identifier of a cluster.
        This parameter is case sensitive.
    user
        The name of a database user.
    database
        Database name. If None, the default Database is used.
    duration
        The number of seconds until the returned temporary password expires.
        Constraint: minimum 900, maximum 3600.
        Default: 900
    auto_create
        Create a database user with the name specified for the user named in user if one does not exist.
    db_groups
        A list of the names of existing database groups that the user named in user will join for the current session,
        in addition to any group memberships for an existing user. If not specified, a new user is added only to PUBLIC.
    boto3_session
        The default boto3 session will be used if **boto3_session** is ``None``.
    ssl
        This governs SSL encryption for TCP/IP sockets.
        This parameter is forward to redshift_connector.
        https://github.com/aws/amazon-redshift-python-driver
    timeout
        This is the time in seconds before the connection to the server will time out.
        The default is None which means no timeout.
        This parameter is forward to redshift_connector.
        https://github.com/aws/amazon-redshift-python-driver
    max_prepared_statements
        This parameter is forward to redshift_connector.
        https://github.com/aws/amazon-redshift-python-driver
    tcp_keepalive
        If True then use TCP keepalive. The default is True.
        This parameter is forward to redshift_connector.
        https://github.com/aws/amazon-redshift-python-driver
    **kwargs
        Forwarded to redshift_connector.connect.
        e.g. is_serverless=True, serverless_acct_id='...', serverless_work_group='...'

    Returns
    -------
        ``redshift_connector`` connection.

    Examples
    --------
    >>> import awswrangler as wr
    >>> with wr.redshift.connect_temp(cluster_identifier="my-cluster", user="test") as con:
    ...     with con.cursor() as cursor:
    ...         cursor.execute("SELECT 1")
    ...         print(cursor.fetchall())

    """
    client_redshift = _utils.client(service_name="redshift", session=boto3_session)
    args: dict[str, Any] = {
        "DbUser": user,
        "ClusterIdentifier": cluster_identifier,
        "DurationSeconds": duration,
        "AutoCreate": auto_create,
    }
    if db_groups is not None:
        args["DbGroups"] = db_groups
    else:
        db_groups = []
    res = client_redshift.get_cluster_credentials(**args)
    cluster = client_redshift.describe_clusters(ClusterIdentifier=cluster_identifier)["Clusters"][0]
    return redshift_connector.connect(
        user=res["DbUser"],
        database=database if database else cluster["DBName"],
        password=res["DbPassword"],
        port=cluster["Endpoint"]["Port"],
        host=cluster["Endpoint"]["Address"],
        ssl=ssl,
        timeout=timeout,
        max_prepared_statements=max_prepared_statements,
        tcp_keepalive=tcp_keepalive,
        db_groups=db_groups,
        **kwargs,
    )