From e42711cd545500e3f1ee732543c2da2333b90b42 Mon Sep 17 00:00:00 2001
From: Bryce Groff <bgroff@hawaii.edu>
Date: Thu, 1 Aug 2024 17:12:22 -0700
Subject: [PATCH] fix: allow the pg container used to create the keycloak
 database to be overriden (#13385)

---
 .../templates/statefulset.yaml                |  4 +++
 charts/airbyte/values.yaml                    |  4 +++
 .../tests/enterprise_config_test.go           | 25 ++++++++++++++-----
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/charts/airbyte-keycloak/templates/statefulset.yaml b/charts/airbyte-keycloak/templates/statefulset.yaml
index 40f2b5f3219..1aabfc835c4 100644
--- a/charts/airbyte-keycloak/templates/statefulset.yaml
+++ b/charts/airbyte-keycloak/templates/statefulset.yaml
@@ -50,7 +50,11 @@ spec:
         {{- if and (eq .Values.global.deploymentMode "oss") (ne .Values.bypassInit true)}}
         # This init container will only executed if the deployment mode is "oss" and the bypassInit is not true.
         - name: init-db
+          {{- if .Values.initContainers.initDb.image }}
+          image: {{ .Values.initContainers.initDb.image }}
+          {{- else }}
           image: postgres:13-alpine
+          {{- end }}
           command: [ "sh", "-c" ]
           args:
             - >
diff --git a/charts/airbyte/values.yaml b/charts/airbyte/values.yaml
index 3b420aba9b8..127da7e254a 100644
--- a/charts/airbyte/values.yaml
+++ b/charts/airbyte/values.yaml
@@ -1731,6 +1731,10 @@ keycloak:
     # gid=0(root)
     fsGroup: 0
 
+  initContainers:
+    initDb:
+      image: {}
+
   initContainerSecurityContext:
     allowPrivilegeEscalation: false 
     runAsNonRoot: true
diff --git a/charts/helm-tests/tests/enterprise_config_test.go b/charts/helm-tests/tests/enterprise_config_test.go
index c31c81be4e0..f50939cb73c 100644
--- a/charts/helm-tests/tests/enterprise_config_test.go
+++ b/charts/helm-tests/tests/enterprise_config_test.go
@@ -242,7 +242,7 @@ func TestBasicEnterpriseConfigWithHelmValues(t *testing.T) {
 	})
 }
 
-func TestKeycloakSetupInitContainerOverride(t *testing.T) {
+func TestKeycloakInitContainerOverride(t *testing.T) {
 	t.Run("default keycloak readiness image is curlimages/curl", func(t *testing.T) {
 		helmOpts := baseHelmOptionsForEnterpriseWithValues()
 		helmOpts.SetValues["global.auth.instanceAdmin.firstName"] = "Octavia"
@@ -260,8 +260,14 @@ func TestKeycloakSetupInitContainerOverride(t *testing.T) {
 		assert.NoError(t, err)
 		assert.NotNil(t, keycloakSetupJob)
 
-		initContainers := keycloakSetupJob.Spec.Template.Spec.InitContainers
-		assert.Equal(t, "curlimages/curl:8.1.1", initContainers[0].Image)
+		keycloakStatefulSet, err := getStatefulSet(chartYaml, "airbyte-keycloak")
+		assert.NoError(t, err)
+		assert.NotNil(t, keycloakStatefulSet)
+
+		setupInitContainers := keycloakSetupJob.Spec.Template.Spec.InitContainers
+		keycloakInitContainers := keycloakStatefulSet.Spec.Template.Spec.InitContainers
+		assert.Equal(t, "curlimages/curl:8.1.1", setupInitContainers[0].Image)
+		assert.Equal(t, "postgres:13-alpine", keycloakInitContainers[0].Image)
 	})
 
 	t.Run("override init container image ", func(t *testing.T) {
@@ -275,14 +281,21 @@ func TestKeycloakSetupInitContainerOverride(t *testing.T) {
 		helmOpts.SetValues["global.auth.identityProvider.oidc.clientIdSecretKey"] = "client-id"
 		helmOpts.SetValues["global.auth.identityProvider.oidc.clientSecretSecretKey"] = "client-secret"
 		helmOpts.SetValues["keycloak-setup.initContainers.keycloakReadinessCheck.image"] = "airbyte/custom-curl-image"
+		helmOpts.SetValues["keycloak.initContainers.initDb.image"] = "airbyte/custom-postgres-image"
 		chartYaml, err := helm.RenderTemplateE(t, helmOpts, chartPath, "airbyte", nil)
 		assert.NoError(t, err)
 
 		keycloakSetupJob, err := getJob(chartYaml, "airbyte-keycloak-setup")
-		assert.NoError(t, err)
 		assert.NotNil(t, keycloakSetupJob)
+		assert.NoError(t, err)
+
+		keycloakStatefulSet, err := getStatefulSet(chartYaml, "airbyte-keycloak")
+		assert.NotNil(t, keycloakStatefulSet)
+		assert.NoError(t, err)
 
-		initContainers := keycloakSetupJob.Spec.Template.Spec.InitContainers
-		assert.Equal(t, "airbyte/custom-curl-image", initContainers[0].Image)
+		setupInitContainers := keycloakSetupJob.Spec.Template.Spec.InitContainers
+		keycloakInitContainers := keycloakStatefulSet.Spec.Template.Spec.InitContainers
+		assert.Equal(t, "airbyte/custom-curl-image", setupInitContainers[0].Image)
+		assert.Equal(t, "airbyte/custom-postgres-image", keycloakInitContainers[0].Image)
 	})
 }