From c54cb967c9f4818f8088d76c62dcee55b35fbf75 Mon Sep 17 00:00:00 2001
From: Ryan Br <ryan.broughan@gmail.com>
Date: Fri, 9 Aug 2024 14:18:28 -0700
Subject: [PATCH] chore: add assumed role secrets. (#13466)

---
 .../src/main/kotlin/config/EnvVarConfigBeanFactory.kt | 11 ++++++-----
 .../test/kotlin/config/EnvVarConfigBeanFactoryTest.kt |  4 ++--
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/airbyte-workload-launcher/src/main/kotlin/config/EnvVarConfigBeanFactory.kt b/airbyte-workload-launcher/src/main/kotlin/config/EnvVarConfigBeanFactory.kt
index 925fb06f1c8..a3f04fb5ac6 100644
--- a/airbyte-workload-launcher/src/main/kotlin/config/EnvVarConfigBeanFactory.kt
+++ b/airbyte-workload-launcher/src/main/kotlin/config/EnvVarConfigBeanFactory.kt
@@ -44,6 +44,7 @@ class EnvVarConfigBeanFactory {
     @Named("workloadApiEnvMap") workloadApiEnvMap: Map<String, String>,
     @Named("workloadApiSecretEnv") secretsEnvMap: Map<String, EnvVarSource>,
     @Named("databaseEnvMap") dbEnvMap: Map<String, String>,
+    @Named("awsAssumedRoleSecretEnv") awsAssumedRoleSecretEnv: Map<String, EnvVarSource>,
   ): List<EnvVar> {
     val envMap: MutableMap<String, String> = HashMap()
 
@@ -71,7 +72,7 @@ class EnvVarConfigBeanFactory {
         .toList()
 
     val secretEnvVars =
-      (secretsEnvMap + secretPersistenceSecretsEnvMap)
+      (secretsEnvMap + secretPersistenceSecretsEnvMap + awsAssumedRoleSecretEnv)
         .map { EnvVar(it.key, null, it.value) }
         .toList()
 
@@ -201,8 +202,8 @@ class EnvVarConfigBeanFactory {
    * To be injected into orchestrator pods, so they can start AWS connector pods that use assumed role access.
    */
   @Singleton
-  @Named("orchestratorAwsAssumedRoleSecretEnv")
-  fun orchestratorAwsAssumedRoleSecretEnv(
+  @Named("awsAssumedRoleSecretEnv")
+  fun awsAssumedRoleSecretEnv(
     @Value("\${airbyte.connector.source.credentials.aws.assumed-role.access-key}") awsAssumedRoleAccessKey: String,
     @Value("\${airbyte.connector.source.credentials.aws.assumed-role.secret-key}") awsAssumedRoleSecretKey: String,
     @Value("\${airbyte.connector.source.credentials.aws.assumed-role.secret-name}") awsAssumedRoleSecretName: String,
@@ -241,9 +242,9 @@ class EnvVarConfigBeanFactory {
   @Named("orchestratorSecretsEnvMap")
   fun orchestratorSecretsEnvMap(
     @Named("workloadApiSecretEnv") workloadApiSecretEnv: Map<String, EnvVarSource>,
-    @Named("orchestratorAwsAssumedRoleSecretEnv") orchestratorAwsAssumedRoleSecretEnv: Map<String, EnvVarSource>,
+    @Named("awsAssumedRoleSecretEnv") awsAssumedRoleSecretEnv: Map<String, EnvVarSource>,
   ): Map<String, EnvVarSource> {
-    return workloadApiSecretEnv + orchestratorAwsAssumedRoleSecretEnv
+    return workloadApiSecretEnv + awsAssumedRoleSecretEnv
   }
 
   private fun createEnvVarSource(
diff --git a/airbyte-workload-launcher/src/test/kotlin/config/EnvVarConfigBeanFactoryTest.kt b/airbyte-workload-launcher/src/test/kotlin/config/EnvVarConfigBeanFactoryTest.kt
index 9f1736cdb27..42603cd9a5c 100644
--- a/airbyte-workload-launcher/src/test/kotlin/config/EnvVarConfigBeanFactoryTest.kt
+++ b/airbyte-workload-launcher/src/test/kotlin/config/EnvVarConfigBeanFactoryTest.kt
@@ -54,7 +54,7 @@ class EnvVarConfigBeanFactoryTest {
   fun `orchestrator aws assumed role secret creation`() {
     val factory = EnvVarConfigBeanFactory()
     val envMap =
-      factory.orchestratorAwsAssumedRoleSecretEnv(
+      factory.awsAssumedRoleSecretEnv(
         AWS_ASSUMED_ROLE_ACCESS_KEY,
         AWS_ASSUMED_ROLE_SECRET_KEY,
         AWS_ASSUMED_ROLE_SECRET_NAME,
@@ -87,7 +87,7 @@ class EnvVarConfigBeanFactoryTest {
   fun `orchestrator aws assumed role secret creation with blank names`() {
     val factory = EnvVarConfigBeanFactory()
     val envMap =
-      factory.orchestratorAwsAssumedRoleSecretEnv(
+      factory.awsAssumedRoleSecretEnv(
         AWS_ASSUMED_ROLE_ACCESS_KEY,
         AWS_ASSUMED_ROLE_SECRET_KEY,
         "",