From 77f560c605a8f7e2323621b4d70bf1296fe2cbf0 Mon Sep 17 00:00:00 2001
From: Mina Slater <slater.mina@gmail.com>
Date: Wed, 9 Nov 2022 15:51:13 -0600
Subject: [PATCH] Adjust admin-login module to handle replica

Saves the replica database url to the same secrets so its value can be exposed as env variable when needed.
---
 rds-postgres/admin-login/main.tf              | 26 +++++++++++++------
 .../admin-login/rotation/lambda_function.py   | 17 +++++++++---
 rds-postgres/admin-login/variables.tf         |  6 +++++
 3 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/rds-postgres/admin-login/main.tf b/rds-postgres/admin-login/main.tf
index 4ad1b59..8fe2168 100644
--- a/rds-postgres/admin-login/main.tf
+++ b/rds-postgres/admin-login/main.tf
@@ -8,14 +8,7 @@ module "secret" {
   resource_tags    = var.tags
   trust_tags       = var.trust_tags
 
-  initial_value = jsonencode({
-    dbname   = var.database_name
-    engine   = data.aws_db_instance.this.engine
-    host     = data.aws_db_instance.this.address
-    password = var.initial_password
-    port     = tostring(data.aws_db_instance.this.port)
-    username = var.username
-  })
+  initial_value = jsonencode(local.initial_secret_value)
 }
 
 module "rotation" {
@@ -77,6 +70,23 @@ data "aws_db_instance" "this" {
   db_instance_identifier = var.identifier
 }
 
+data "aws_db_instance" "replica" {
+  count = var.replica_identifier ? 1 : 0
+
+  db_instance_identifier = var.replica_identifier
+}
+
 locals {
   full_name = join("-", ["rds-postgres", var.identifier])
+
+  base_value = {
+    dbname       = var.database_name
+    engine       = data.aws_db_instance.this.engine
+    host         = data.aws_db_instance.this.address
+    password     = var.initial_password
+    port         = tostring(data.aws_db_instance.this.port)
+    username     = var.username
+  }
+
+  initial_secret_value = var.replica_identifier ? merge(local.base_value, { replica_host = data.aws_db_instance.replica[0].address }) : local.base_value
 }
diff --git a/rds-postgres/admin-login/rotation/lambda_function.py b/rds-postgres/admin-login/rotation/lambda_function.py
index 76e621f..4c4e36c 100644
--- a/rds-postgres/admin-login/rotation/lambda_function.py
+++ b/rds-postgres/admin-login/rotation/lambda_function.py
@@ -31,7 +31,8 @@ def lambda_handler(event, context):
         'username': <required: username>,
         'password': <required: password>,
         'dbname': <optional: database name, default to 'postgres'>,
-        'port': <optional: if not specified, default port 5432 will be used>
+        'port': <optional: if not specified, default port 5432 will be used>,
+        'replica_host': <optional: host address of replica DB>
     }
 
     Args:
@@ -126,7 +127,11 @@ def create_secret(service_client, arn, token):
         current_dict['password'] = passwd['RandomPassword']
 
         # Add DATABASE_URL to secret
-        current_dict['DATABASE_URL'] = dict_to_url(current_dict)
+        current_dict['DATABASE_URL'] = dict_to_url(current_dict, false)
+
+        if secret['replica_host']:
+            # Add DATABASE_REPLICA_URL to secret
+            current_dict['DATABASE_REPLICA_URL'] = dict_to_url(current_dict, true)
 
         # Put the secret
         service_client.put_secret_value(SecretId=arn, ClientRequestToken=token, SecretString=json.dumps(current_dict), VersionStages=['AWSPENDING'])
@@ -278,7 +283,7 @@ def finish_secret(service_client, arn, token):
     service_client.update_secret_version_stage(SecretId=arn, VersionStage="AWSCURRENT", MoveToVersionId=token, RemoveFromVersionId=current_version)
     logger.info("finishSecret: Successfully set AWSCURRENT stage to version %s for secret %s." % (token, arn))
 
-def dict_to_url(secret):
+def dict_to_url(secret, replica):
     """Reformats connection details as a URL string
 
     Generate a Heroku-style DATABASE_URL with connection details
@@ -289,9 +294,13 @@ def dict_to_url(secret):
     Returns:
         url: DATABASE_URL-style string
     """
+    if replica:
+        host = secret['host']
+    else:
+        host = secret['replica_host']
 
     return "postgres://%s:%s@%s:%s/%s" % (secret['username'],
-            secret['password'], secret['host'], secret['port'],
+            secret['password'], host, secret['port'],
             secret['dbname'])
 
 def get_connection(secret_dict):
diff --git a/rds-postgres/admin-login/variables.tf b/rds-postgres/admin-login/variables.tf
index ea01d32..a2c4c14 100644
--- a/rds-postgres/admin-login/variables.tf
+++ b/rds-postgres/admin-login/variables.tf
@@ -31,6 +31,12 @@ variable "read_principals" {
   default     = null
 }
 
+variable "replica_identifier" {
+  description = "Identifier of the database replica"
+  type        = string
+  default     = null
+}
+
 variable "secret_name" {
   description = "Override the name for this secret"
   type        = string