Support AWS RDS IAM Authentication for Redash database

Author: winebarrel

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
@@ -42,6 +45,21 @@ def apply_pool_defaults(self, app, options):
 make_searchable(db.metadata, options={"regconfig": "pg_catalog.simple"})
 
 
+# IAM database authentication for AWS RDS
+# See https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/UsingWithRDS.IAMDBAuth.html
+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
+
+
 class SearchBaseQuery(BaseQuery, SearchQueryMixin):
     """
     The SQA query class to use when full text search is wanted.

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