Support AWS RDS IAM Authentication for Redash database

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:

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"))