Add out of box private registry support for dag deploy (#2109)

* fix pre-commit * add tests for dagdeploy private registry * fix ci failures
astronomer · Feb 2, 2024 · c007c04 · c007c04
1 parent 8b541cf
commit c007c04
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 2 deletions.
diff --git a/charts/astronomer/templates/houston/houston-configmap.yaml b/charts/astronomer/templates/houston/houston-configmap.yaml
@@ -136,12 +136,13 @@ data:
 
       {{- if .Values.global.dagOnlyDeployment.enabled }}
       # enables dag only deployment for airflow deployments
+      {{- $dagOnlyDeployment := include "dagOnlyDeployment.image" . }}
       dagDeploy:
         enabled: {{ .Values.global.dagOnlyDeployment.enabled }}
         images:
           dagServer:
-            repository: {{ (splitList ":"  .Values.global.dagOnlyDeployment.image ) | first | quote }}
-            tag: {{ (splitList ":"  .Values.global.dagOnlyDeployment.image ) | last  }}
+            repository: {{ (splitList ":"  $dagOnlyDeployment ) | first  }}
+            tag: {{ (splitList ":"  $dagOnlyDeployment ) | last  }}
         {{- if .Values.global.dagOnlyDeployment.enabled }}
         securityContext: {{- .Values.global.dagOnlyDeployment.securityContext | toYaml | nindent 10 }}
         {{- end }}

diff --git a/templates/_helpers.tpl b/templates/_helpers.tpl
@@ -23,3 +23,11 @@ nginx.ingress.kubernetes.io/auth-url: https://houston.{{ .Values.global.baseDoma
 {{ define "containerd.configToml" -}}
 {{- .Values.global.privateCaCertsAddToHost.containerdConfigToml -}}
 {{- end }}
+
+{{ define "dagOnlyDeployment.image" -}}
+{{- if .Values.global.privateRegistry.enabled -}}
+{{ .Values.global.privateRegistry.repository }}/ap-dag-deploy:{{ (splitList ":"  .Values.global.dagOnlyDeployment.image ) | last  }}
+{{- else -}}
+{{ .Values.global.dagOnlyDeployment.image }}
+{{- end }}
+{{- end }}
diff --git a/tests/chart_tests/test_astronomer_dag_only_deploy.py b/tests/chart_tests/test_astronomer_dag_only_deploy.py
@@ -55,3 +55,30 @@ def test_dagonlydeploy_config_enabled(self, kube_version):
             "securityContext": {"fsGroup": 55555},
             "resources": resources,
         }
+
+    def test_dagonlydeploy_config_enabled_with_private_registry(self, kube_version):
+        """Test dagonlydeploy with private registry."""
+        private_registry = "private-registry.example.com"
+        docs = render_chart(
+            kube_version=kube_version,
+            values={
+                "global": {
+                    "privateRegistry": {
+                        "enabled": True,
+                        "repository": private_registry,
+                    },
+                    "dagOnlyDeployment": {
+                        "enabled": True,
+                    },
+                }
+            },
+            show_only=["charts/astronomer/templates/houston/houston-configmap.yaml"],
+        )
+
+        doc = docs[0]
+        prod = yaml.safe_load(doc["data"]["production.yaml"])
+        assert prod["deployments"]["dagOnlyDeployment"] is True
+
+        assert prod["deployments"]["dagDeploy"]["images"]["dagServer"][
+            "repository"
+        ].startswith(private_registry)