diff --git a/redash/models/base.py b/redash/models/base.py index 2ed95c38fb..74210ff6e1 100644 --- a/redash/models/base.py +++ b/redash/models/base.py @@ -1,7 +1,10 @@ import functools +import boto3 from flask_sqlalchemy import BaseQuery, SQLAlchemy from sqlalchemy.dialects.postgresql import UUID +from sqlalchemy.engine import Engine +from sqlalchemy.event import listens_for from sqlalchemy.orm import object_session from sqlalchemy.pool import NullPool from sqlalchemy_searchable import SearchQueryMixin, make_searchable, vectorizer @@ -11,6 +14,21 @@ class RedashSQLAlchemy(SQLAlchemy): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + if settings.REDASH_DATABASE_IAM_AUTH: + + @listens_for(Engine, "do_connect") + def db_connect_hook(dialect, conn_rec, cargs, cparams): + rds_client = boto3.client("rds") + auth_token = rds_client.generate_db_auth_token( + DBHostname=cparams["host"], + Port=cparams["port"], + DBUsername=cparams["user"], + ) + cparams["password"] = auth_token + def apply_driver_hacks(self, app, info, options): options.update(json_serializer=json_dumps) if settings.SQLALCHEMY_ENABLE_POOL_PRE_PING: diff --git a/redash/settings/__init__.py b/redash/settings/__init__.py index 973beffc55..2005123775 100644 --- a/redash/settings/__init__.py +++ b/redash/settings/__init__.py @@ -459,3 +459,6 @@ def email_server_is_configured(): # Email blocked domains, use delimiter comma to separated multiple domains BLOCKED_DOMAINS = set_from_string(os.environ.get("REDASH_BLOCKED_DOMAINS", "qq.com")) + +# AWS +REDASH_DATABASE_IAM_AUTH = parse_boolean(os.environ.get("REDASH_DATABASE_IAM_AUTH", "false"))