SQLAlchemy with RDS IAM Authentication

sqlalchemy-rds-iam-pgsql.py

'''
  Example of IAM Authentication to Postgres RDS from Python and SQLAlchemy. Uses SQLAlchemy events to
  run a callback just before making a db connection and adding it to the internal pool. The callback
  sets the connection parameters including the password/token.

  https://docs.sqlalchemy.org/en/14/core/events.html#sqlalchemy.events.DialectEvents.do_connect
  https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/UsingWithRDS.IAMDBAuth.html
  https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/UsingWithRDS.IAMDBAuth.Connecting.Python.html
  https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/rds.html#RDS.Client.generate_db_auth_token

  Note: one requirement is that the pg user needs to be granted the rds_iam role
  `GRANT rds_iam TO postgres;`
'''
import boto3
from sqlalchemy import create_engine, event, text

REGION = "us-east-2" # can be assumed, unless you connect from one region to another
DBHostname = "demo.….us-east-2.rds.amazonaws.com"
DBPort = "5432"
DBUsername = "postgres"
DBName = "postgres"

engine = create_engine(f"postgresql:///") # connection params will be set by the event callback

@event.listens_for(engine, "do_connect")
def provide_token(dialect, conn_rec, cargs, cparams):
    client = boto3.client("rds")
    token = client.generate_db_auth_token(DBHostname=DBHostname, Port=DBPort, DBUsername=DBUsername, Region=REGION)
    # set up db connection parameters, alternatively we can get these from boto3 describe_db_instances
    cparams['host'] = DBHostname
    cparams['port'] = DBPort
    cparams['user'] = DBUsername
    cparams['password'] = token
    cparams['database'] = DBName
    print(token)


with engine.connect() as connection:
    print("="*40)
    print("Results:")

    result = connection.execute(text("select now()"))
    for row in result:
        print(row)

    result = connection.execute(text("select 1"))
    for row in result:
        print(row)

sqlalchemy-redshift-iam.py

'''
  aws sts get-caller-identity -- to check current aws role/user

On the redshift cluster:
  GRANT ALL ON  SCHEMA "datalake0" to "IAM:admin";
'''
import boto3
from sqlalchemy import create_engine, event, text

REGION = "us-west-2" # can be assumed, unless you connect from one region to another
DBHost = "11111111111111111.us-west-2.redshift-serverless.amazonaws.com"
DBPort = "5439"
DBUser = "admin"
DBName = "dev"
ClusterIdentifier = "redshift-serverless-default"

engine = create_engine(f"postgresql:///") # connection params will be set by the event callback

@event.listens_for(engine, "do_connect")
def provide_token(dialect, conn_rec, cargs, cparams):
    client = boto3.client("redshift", region_name=REGION)
    creds = client.get_cluster_credentials(
                DbUser=DBUser,
                DbName=DBName,
                AutoCreate=True,
                ClusterIdentifier=ClusterIdentifier,
            )
    # set up db connection parameters, alternatively we can get these from boto3 describe_db_instances
    cparams['host'] = DBHost
    cparams['port'] = DBPort
    cparams['database'] = DBName
    cparams['user'] = creds.get('DbUser')
    cparams['password'] = creds.get('DbPassword')
    print('Got creds!', creds)

with engine.connect() as connection:
    print("="*40)
    print("Results:")

    result = connection.execute(text("select now()"))
    for row in result:
        print(row)
    result = connection.execute(text("select 1"))
    for row in result:
        print(row)
    result = connection.execute(text("select current_user"))
    for row in result:
        print(row)

    result = connection.execute(text('SELECT * from "public"."redshift_rds"'))
    for row in result:
        print(row)