Add extra init containers user can add (#424)

cockroachdb · Oct 18, 2024 · 2396411 · 2396411
1 parent 0d4ad59
commit 2396411
Show file tree

Hide file tree

Showing 4 changed files with 185 additions and 0 deletions.
diff --git a/build/templates/values.yaml b/build/templates/values.yaml
@@ -293,6 +293,39 @@ statefulset:
     # Additional serviceAccount annotations (e.g. for attaching AWS IAM roles to pods)
     annotations: {}
 
+  # initContainers allows you to add additional containers to cockroachdb statefulset.
+  initContainers: []
+#  - name: "fetch-metadata"
+#    image: "badouralix/curl-jq"
+#    command:
+#      - "sh"
+#      - "-c"
+#      - "curl -s -H \"Metadata:true\" --noproxy \"*\" \"http://169.254.169.254/metadata/instance?api-version=2021-02-01\" | jq '.' > /metadata/instance_metadata.json"
+#    resources: {}
+#        # requests:
+#        #   cpu: "10m"
+#        #   memory: "128Mi"
+#        # limits:
+#        #   cpu: "10m"
+#        #   memory: "128Mi"
+#    securityContext:
+#      allowPrivilegeEscalation: false
+#      capabilities:
+#        drop:
+#          - ALL
+#      privileged: false
+#      readOnlyRootFilesystem: true
+
+  # volumeMounts are mounted on the same path in the main crdb container and all init containers.
+  volumeMounts: []
+#   - name: metadata
+#     mountPath: /metadata
+
+  # volumes allows you to add additional volumes to cockroachdb statefulset.
+  volumes: []
+#  - name: metadata
+#    emptyDir: {}
+
 service:
   ports:
     # You can set a different external and internal gRPC ports and their name.

diff --git a/cockroachdb/templates/statefulset.yaml b/cockroachdb/templates/statefulset.yaml
@@ -82,6 +82,13 @@ spec:
         {{- with .Values.tls.copyCerts.resources }}
           resources: {{- toYaml . | nindent 12 }}
         {{- end }}
+        {{- range $ic := .Values.statefulset.initContainers }}
+        - {{- toYaml $ic | nindent 10 }}
+          {{ with $.Values.statefulset.volumeMounts}}
+          volumeMounts:
+          {{- toYaml . | nindent 12 }}
+          {{- end }}
+        {{- end }}
       {{- end }}
     {{- if or .Values.statefulset.nodeAffinity .Values.statefulset.podAffinity .Values.statefulset.podAntiAffinity }}
       affinity:
@@ -283,6 +290,9 @@ spec:
               mountPath: /cockroach/log-config
               readOnly: true
           {{- end }}
+          {{ with .Values.statefulset.volumeMounts }}
+            {{ toYaml . | nindent 12 }}
+          {{- end }}
           {{- if .Values.statefulset.customStartupProbe }}
           startupProbe:
             {{ toYaml .Values.statefulset.customStartupProbe | nindent 12 }}
@@ -339,6 +349,9 @@ spec:
         {{- else }}
           emptyDir: {}
         {{- end }}
+        {{ with .Values.statefulset.volumes }}
+          {{ toYaml . | nindent 8 }}
+        {{- end }}
       {{- if .Values.tls.enabled }}
         - name: certs
           emptyDir: {}

diff --git a/cockroachdb/values.yaml b/cockroachdb/values.yaml
@@ -294,6 +294,39 @@ statefulset:
     # Additional serviceAccount annotations (e.g. for attaching AWS IAM roles to pods)
     annotations: {}
 
+  # initContainers allows you to add additional containers to cockroachdb statefulset.
+  initContainers: []
+#  - name: "fetch-metadata"
+#    image: "badouralix/curl-jq"
+#    command:
+#      - "sh"
+#      - "-c"
+#      - "curl -s -H \"Metadata:true\" --noproxy \"*\" \"http://169.254.169.254/metadata/instance?api-version=2021-02-01\" | jq '.' > /metadata/instance_metadata.json"
+#    resources: {}
+#        # requests:
+#        #   cpu: "10m"
+#        #   memory: "128Mi"
+#        # limits:
+#        #   cpu: "10m"
+#        #   memory: "128Mi"
+#    securityContext:
+#      allowPrivilegeEscalation: false
+#      capabilities:
+#        drop:
+#          - ALL
+#      privileged: false
+#      readOnlyRootFilesystem: true
+
+  # volumeMounts are mounted on the same path in the main crdb container and all init containers.
+  volumeMounts: []
+#   - name: metadata
+#     mountPath: /metadata
+
+  # volumes allows you to add additional volumes to cockroachdb statefulset.
+  volumes: []
+#  - name: metadata
+#    emptyDir: {}
+
 service:
   ports:
     # You can set a different external and internal gRPC ports and their name.

diff --git a/tests/template/cockroachdb_helm_template_test.go b/tests/template/cockroachdb_helm_template_test.go
@@ -1283,3 +1283,109 @@ func TestHelmInitJobAnnotations(t *testing.T) {
 		})
 	}
 }
+
+func TestStatefulSetInitContainers(t *testing.T) {
+	t.Parallel()
+	testCases := []struct {
+		name          string
+		initContainer bool
+		volume        bool
+		values        map[string]string
+	}{
+		{
+			"Add extra init container",
+			true,
+			true,
+			map[string]string{
+				"statefulset.initContainers[0].name":       "fetch-metadata",
+				"statefulset.initContainers[0].image":      "busybox",
+				"statefulset.initContainers[0].command[0]": "/bin/bash",
+				"statefulset.initContainers[0].command[1]": "-c",
+				"statefulset.initContainers[0].command[2]": "echo 'Fetching metadata'",
+				"statefulset.volumeMounts[0].name":         "metadata",
+				"statefulset.volumeMounts[0].mountPath":    "/metadata",
+				"statefulset.volumes[0].name":              "metadata",
+				"statefulset.volumes[0].configMap.name":    "log-config",
+			},
+		},
+		{
+			"Add extra volume without init container",
+			false,
+			true,
+			map[string]string{
+				"statefulset.volumeMounts[0].name":      "metadata",
+				"statefulset.volumeMounts[0].mountPath": "/metadata",
+				"statefulset.volumes[0].name":           "metadata",
+				"statefulset.volumes[0].configMap.name": "log-config",
+			},
+		},
+		{
+			"Add extra init container without volume",
+			true,
+			false,
+			map[string]string{
+				"statefulset.initContainers[0].name":       "fetch-metadata",
+				"statefulset.initContainers[0].image":      "busybox",
+				"statefulset.initContainers[0].command[0]": "/bin/bash",
+				"statefulset.initContainers[0].command[1]": "-c",
+				"statefulset.initContainers[0].command[2]": "echo 'Fetching metadata'",
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		// Here, we capture the range variable and force it into the scope of this block. If we don't do this, when the
+		// subtest switches contexts (because of t.Parallel), the testCase value will have been updated by the for loop
+		// and will be the next testCase!
+		testCase := testCase
+		t.Run(testCase.name, func(subT *testing.T) {
+			subT.Parallel()
+
+			options := &helm.Options{
+				KubectlOptions: k8s.NewKubectlOptions("", "", namespaceName),
+				SetValues:      testCase.values,
+			}
+			output, err := helm.RenderTemplateE(t, options, helmChartPath, releaseName, []string{"templates/statefulset.yaml"})
+
+			require.Equal(subT, err, nil)
+
+			var sts appsv1.StatefulSet
+			helm.UnmarshalK8SYaml(t, output, &sts)
+
+			if testCase.initContainer {
+				var fetchMetadata bool
+				for _, c := range sts.Spec.Template.Spec.InitContainers {
+					if c.Name == "fetch-metadata" {
+						fetchMetadata = true
+						require.Equal(subT, "busybox", c.Image)
+						require.Equal(subT, []string{"/bin/bash", "-c", "echo 'Fetching metadata'"}, c.Command)
+						if testCase.volume {
+							require.Equal(subT, []corev1.VolumeMount{{Name: "metadata", MountPath: "/metadata"}}, c.VolumeMounts)
+						}
+						break
+					}
+				}
+
+				if !fetchMetadata {
+					require.Fail(subT, "Init container fetch-metadata not found")
+				}
+			}
+
+			if testCase.volume {
+				var metadataVolume bool
+				for _, v := range sts.Spec.Template.Spec.Volumes {
+					if v.Name == "metadata" {
+						metadataVolume = true
+						require.Equal(subT, v.ConfigMap.Name, "log-config")
+						break
+					}
+				}
+
+				if !metadataVolume {
+					require.Fail(subT, "Volume metadata not found")
+				}
+			}
+		})
+	}
+
+}