Password-less database authentication with Apache Superset and AWS Aurora Serverless V2 (PostgreSQL)

Liam Baker
Sage Ai
Published in
5 min readNov 29, 2023

--

Fewer static passwords, less risk. Migrating applications to Aurora Serverless V2.

Introduction

AWS has long supported IAM Database Authentication. This was a great opportunity as part of a larger migration to Aurora Serverless V2 and PostgreSQL 15 being pursued for scaling and feature improvements (namely the pgvector extension), to rethink how we provided storage for use by modern workloads.

Planning for the migration

Changing a database authentication pattern is challenging for any application. At Sage AI, all our applications are deployed to Kubernetes via Helm Charts. To ensure backwards compatibility for tenants with services that could not yet migrate to this authentication strategy, an optional feature flag was used in our core application configuration (a custom Helm Chart). Apache Superset being the flexible OSS project that it is, supports the use of password retrieval using a custom function. If the feature flag is enabled the custom function is patched into the Apache Superset (henceforth Superset) configuration file. This disables the legacy method, which required the ExternalSecrets Operator to fetch the password stored in AWS SSM Parameter Store.

Our cloud infrastructure is managed using Terraform, so when provisioning the database using the PostgreSQL provider, it was necessary to grant an additional, special role named rds_iam to the database users within our module. This ensures password authentication is disabled for users granted the role. A benefit of doing this is that it would be quick to revoke the grant in the event the authentication strategy was broken, in a “break glass” scenario.

To add to the complexity, the major version of the database engine had recently been updated courtesy of the migration to Aurora Serverless V2, so it was critical to ensure recent changes were accounted for in the target version (v15.x) from our current version (v13.9).

Implementation

IAM Policies

AWS provides a special reference for this action as described here, rds-db:connect, which allows the specified principal to retrieve temporary tokens in order make connections to the database cluster with the specified database user. Since Terraform is used to configure IAM policies and provision the database clusters, it was necessary to add an additional output to the aurora module maintained by Sage AI which the child module didn’t create for us, shown below:

output "cluster_resource_arn" {
description = "The computed ARN of the created cluster, using cluster ID"
value = local.cluster_resource_arn
}
locals {
...
cluster_resource_arn = "arn:aws:rds-db:${var.tf_config.aws_region}:${var.tf_config.aws_account}:dbuser:${module.aurora.cluster_resource_id}"
...
}

Included below is an excerpt from the Terraform module that computes the policy statement that the IAM Role for Superset requires.

output "statements" {
description = "Computed Aurora RDS Connect policy statements"
value = length(var.accessible_clusters) == 0 ? {} : {
RDSConnectAccess = {
actions = ["rds-db:connect"]
resources = flatten([
for value in var.accessible_clusters : [
format("%s/%s",
value.cluster_resource_arn,
value.username,
)
]
])
}
}
}

This produces the following IAM policy statement:

{
"Version": "2012–10–17",
"Statement": [
{
"Action": "rds-db:connect",
"Effect": "Allow",
"Resource": "arn:aws:rds-db:eu-west-1:000000000000:dbuser:cluster-XXXXXXXXXX/bob",
"Sid": "RDSConnectAccess"
}
]
}

Custom Password Store Function

With IAM policies appropriately setup, there needs to be a mechanism to signal to Superset that it needs to perform additional actions when substituting password values into the connection string.

Superset can be configured with a Python file, which provides a variable SQLALCHEMY_CUSTOM_PASSWORD_STORE that can be assigned as a regular Python function.

Using Helm, a feature flag was added to the custom Chart maintained by the Infrastructure team, when specified a small template helper patches the function. This ensured backwards compatibility.

The body of the function can be seen below (with enclosing templating logic):

{{/*
RDS IAM Authentication code for custom password store extension
*/}}
{{- define "superset.rdsIamAuthCustomFunc" -}}
import boto3
import time
from urllib import parse
DATABASE_AUTH_ROLE_SESSION_NAME = "superset-rds-iam-auth-"
SAGEAI_ACCOUNTS = {
"development": "<devAccountID>",
"pre-production": "<preProdAccountId>",
"production": "<ProdAccountId>",
}
def env_var_password_formatter(url):
if not url.password:
return None
account = ""
if "DEVELOPMENT" in url.password:
account = SAGEAI_ACCOUNTS.get("development")
elif "PRE_PRODUCTION" in url.password:
account = SAGEAI_ACCOUNTS.get("pre-production")
else:
account = SAGEAI_ACCOUNTS.get("production")
formatted = url.password.format(**os.environ) if "generate" in formatted:
role_name = os.getenv("DATABASE_AUTH_ROLE_NAME")
region = os.getenv("AWS_REGION")
role_arn = f"arn:aws:iam::{account}:role/{role_name}@{region}"
SINCE_EPOCH = int(time.time())
role_session_name = f"{DATABASE_AUTH_ROLE_SESSION_NAME}-{SINCE_EPOCH}"
credentials = assume_role(role_arn, role_session_name)
rds_token = get_db_auth_token(credentials, url, region)
return rds_token
return formatted
def assume_role(role_arn, role_session_name):
session = boto3.Session()
sts_client = boto3.client("sts")
response = sts_client.assume_role(
RoleArn=role_arn, RoleSessionName=role_session_name
)
return response["Credentials"]
def get_db_auth_token(credentials, url, region) -> str:
rds_client = boto3.client(
"rds",
aws_access_key_id=credentials["AccessKeyId"],
aws_secret_access_key=credentials["SecretAccessKey"],
aws_session_token=credentials["SessionToken"],
region_name=region,
)
port = os.getenv("POSTGRES_PORT", "5432") rds_token = rds_client.generate_db_auth_token(
DBHostname=url.host,
Port=port,
DBUsername=url.username,
Region=region,
)
return parse.quote_plus(rds_token)
{{- end -}}

Helm Release Configuration

Within the HelmRelease object created for the tenant-specific deployment of Superset, there is a special placeholder of “generate” in the connection string where the password would usually be declared. This signals to the patched function above that it should retrieve the token from AWS, instead of the environment. It operates in the “legacy” or normal mode if this keyword is not encountered in the connection string.

PostgreSQL Role Grant

When a database cluster is provisioned through Terraform in our cloud environment, PostgreSQL database users are also created along with it. As the documentation explains, an Aurora PostgreSQL role named rds_iam must be granted to the user wishing to use IAM database authentication, this was achieved by appending this role to the list of roles provided to the postgresql_role resource within the Terraform configuration.

Gotchas

When trying to allow a cross-account IAM role to connect to the target database cluster, to reduce the number of IAM roles involved, database connections would fail with a PAM authentication error. It was later clear that an IAM role in the source account would need to assume an IAM role within the same account that the database cluster was located in order for authentication to succeed.

Assuming a Python setup, the token returned must be sanitised by the function parse.unquote_plus (or a similar function in other languages) otherwise SQL client libraries will incorrectly parse the password portion of the connection URL, again resulting in PAM authentication failures.

Conclusion

At a high level we’ve covered how we planned and introduced a new feature to enable a seamless database authentication experience. In detail, we’ve shown that with a combination of a custom python function and some basic helm templating, we can customise Superset to support additional authentication mechanisms. This has effectively reduced the risk associated with static password exposure. After all, data protection remains a critical and top priority.

--

--