Support AWS RDS IAM Authentication for Redash database

Author: winebarrel

Makefile

@@ -4,11 +4,11 @@ compose_build: .env
 	COMPOSE_DOCKER_CLI_BUILD=1 DOCKER_BUILDKIT=1 docker compose build
 
 up:
-	docker compose up -d redis postgres --remove-orphans
-	docker compose exec -u postgres postgres psql postgres --csv \
-		-1tqc "SELECT table_name FROM information_schema.tables WHERE table_name = 'organizations'" 2> /dev/null \
-		| grep -q "organizations" || make create_database
-	COMPOSE_DOCKER_CLI_BUILD=1 DOCKER_BUILDKIT=1 docker compose up -d --build --remove-orphans
+	# docker compose up -d redis postgres --remove-orphans
+	# docker compose exec -u postgres postgres psql postgres --csv \
+	# 	-1tqc "SELECT table_name FROM information_schema.tables WHERE table_name = 'organizations'" 2> /dev/null \
+	# 	| grep -q "organizations" || make create_database
+	COMPOSE_DOCKER_CLI_BUILD=1 DOCKER_BUILDKIT=1 docker compose up --build --remove-orphans
 
 test_db:
 	@for i in `seq 1 5`; do \

compose.yaml

@@ -13,7 +13,7 @@ x-redash-environment: &redash-environment
   REDASH_HOST: http://localhost:5001
   REDASH_LOG_LEVEL: "INFO"
   REDASH_REDIS_URL: "redis://redis:6379/0"
-  REDASH_DATABASE_URL: "postgresql://postgres@postgres/postgres"
+  # REDASH_DATABASE_URL: "postgresql://postgres@postgres/postgres"
   REDASH_RATELIMIT_ENABLED: "false"
   REDASH_MAIL_DEFAULT_SENDER: "[email protected]"
   REDASH_MAIL_SERVER: "email"
@@ -26,7 +26,7 @@ services:
     <<: *redash-service
     command: dev_server
     depends_on:
-      - postgres
+      # - postgres
       - redis
     ports:
       - "5001:5000"
@@ -52,17 +52,17 @@ services:
   redis:
     image: redis:7-alpine
     restart: unless-stopped
-  postgres:
-    image: pgautoupgrade/pgautoupgrade:latest
-    ports:
-      - "15432:5432"
-    # The following turns the DB into less durable, but gains significant performance improvements for the tests run (x3
-    # improvement on my personal machine). We should consider moving this into a dedicated Docker Compose configuration for
-    # tests.
-    command: "postgres -c fsync=off -c full_page_writes=off -c synchronous_commit=OFF"
-    restart: unless-stopped
-    environment:
-      POSTGRES_HOST_AUTH_METHOD: "trust"
+  # postgres:
+  #   image: pgautoupgrade/pgautoupgrade:latest
+  #   ports:
+  #     - "15432:5432"
+  #   # The following turns the DB into less durable, but gains significant performance improvements for the tests run (x3
+  #   # improvement on my personal machine). We should consider moving this into a dedicated Docker Compose configuration for
+  #   # tests.
+  #   command: "postgres -c fsync=off -c full_page_writes=off -c synchronous_commit=OFF"
+  #   restart: unless-stopped
+  #   environment:
+  #     POSTGRES_HOST_AUTH_METHOD: "trust"
   email:
     image: maildev/maildev
     ports:

redash/models/base.py

@@ -1,7 +1,10 @@
+import boto3
 import functools
 
 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
@@ -41,6 +44,18 @@ def apply_pool_defaults(self, app, options):
 # and indexes for the full text search
 make_searchable(db.metadata, options={"regconfig": "pg_catalog.simple"})
 
+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):
     """

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