Add functionality for enterprise

pkg/controller/mongodb/mongodb_controller.go

@@ -419,12 +419,10 @@ func mongodbAgentContainer(volumeMounts []corev1.VolumeMount) container.Modifica
 			"-serveStatusPort=5000",
 		},
 		),
-		container.WithEnv(
-			[]corev1.EnvVar{
-				{
-					Name:  agentHealthStatusFilePathEnv,
-					Value: agentHealthStatusFilePathValue,
-				},
+		container.WithEnvs(
+			corev1.EnvVar{
+				Name:  agentHealthStatusFilePathEnv,
+				Value: agentHealthStatusFilePathValue,
 			},
 		),
 	)
@@ -461,16 +459,14 @@ mongod -f /data/automation-mongod.conf ;
 		container.WithImage(fmt.Sprintf("mongo:%s", version)),
 		container.WithResourceRequirements(resourcerequirements.Defaults()),
 		container.WithCommand(mongoDbCommand),
-		container.WithEnv(
-			[]corev1.EnvVar{
-				{
-					Name:  agentHealthStatusFilePathEnv,
-					Value: "/healthstatus/agent-health-status.json",
-				},
-				{
-					Name:  preStopHookLogFilePathEnv,
-					Value: "/hooks/pre-stop-hook.log",
-				},
+		container.WithEnvs(
+			corev1.EnvVar{
+				Name:  agentHealthStatusFilePathEnv,
+				Value: "/healthstatus/agent-health-status.json",
+			},
+			corev1.EnvVar{
+				Name:  preStopHookLogFilePathEnv,
+				Value: "/hooks/pre-stop-hook.log",
 			},
 		),
 		container.WithVolumeMounts(volumeMounts),

pkg/kube/container/container_test.go

@@ -0,0 +1,144 @@
+package container
+
+import (
+	"testing"
+
+	"github.com/mongodb/mongodb-kubernetes-operator/pkg/kube/lifecycle"
+	"github.com/mongodb/mongodb-kubernetes-operator/pkg/kube/probes"
+	"github.com/stretchr/testify/assert"
+	corev1 "k8s.io/api/core/v1"
+)
+
+func TestContainer(t *testing.T) {
+	c := New(
+		WithName("name"),
+		WithImage("image"),
+		WithImagePullPolicy(corev1.PullAlways),
+		WithPorts([]corev1.ContainerPort{{Name: "port-1", ContainerPort: int32(1000)}}),
+		WithSecurityContext(corev1.SecurityContext{
+			RunAsGroup:   int64Ref(100),
+			RunAsNonRoot: boolRef(true),
+		}),
+		WithLifecycle(lifecycle.Apply(
+			lifecycle.WithPrestopCommand([]string{"pre-stop-command"}),
+		)),
+		WithReadinessProbe(probes.Apply(
+			probes.WithExecCommand([]string{"exec"}),
+			probes.WithFailureThreshold(10),
+			probes.WithPeriodSeconds(5),
+		)),
+		WithEnvs(
+			[]corev1.EnvVar{
+				{
+					Name:  "env-1",
+					Value: "env-1-value",
+				},
+			}...,
+		),
+	)
+
+	assert.Equal(t, "name", c.Name)
+	assert.Equal(t, "image", c.Image)
+	assert.Equal(t, corev1.PullAlways, c.ImagePullPolicy)
+
+	assert.Len(t, c.Ports, 1)
+	assert.Equal(t, int32(1000), c.Ports[0].ContainerPort)
+	assert.Equal(t, "port-1", c.Ports[0].Name)
+
+	securityContext := c.SecurityContext
+	assert.Equal(t, int64Ref(100), securityContext.RunAsGroup)
+	assert.Equal(t, boolRef(true), securityContext.RunAsNonRoot)
+
+	readinessProbe := c.ReadinessProbe
+	assert.Equal(t, int32(10), readinessProbe.FailureThreshold)
+	assert.Equal(t, int32(5), readinessProbe.PeriodSeconds)
+	assert.Equal(t, "exec", readinessProbe.Exec.Command[0])
+
+	lifeCycle := c.Lifecycle
+	assert.NotNil(t, lifeCycle)
+	assert.NotNil(t, lifeCycle.PreStop)
+	assert.NotNil(t, lifeCycle.PreStop.Exec)
+	assert.Equal(t, "pre-stop-command", lifeCycle.PreStop.Exec.Command[0])
+
+	assert.Len(t, c.Env, 1)
+	assert.Equal(t, "env-1", c.Env[0].Name)
+	assert.Equal(t, "env-1-value", c.Env[0].Value)
+}
+
+func TestMergeEnvs(t *testing.T) {
+	existing := []corev1.EnvVar{
+		{
+			Name:  "C_env",
+			Value: "C_value",
+		},
+		{
+			Name:  "B_env",
+			Value: "B_value",
+		},
+		{
+			Name:  "A_env",
+			Value: "A_value",
+		},
+		{
+			Name: "F_env",
+			ValueFrom: &corev1.EnvVarSource{
+				SecretKeyRef: &corev1.SecretKeySelector{
+					Key: "f_key",
+				},
+			},
+		},
+	}
+
+	desired := []corev1.EnvVar{
+		{
+			Name:  "D_env",
+			Value: "D_value",
+		},
+		{
+			Name:  "E_env",
+			Value: "E_value",
+		},
+		{
+			Name:  "C_env",
+			Value: "C_value_new",
+		},
+		{
+			Name:  "B_env",
+			Value: "B_value_new",
+		},
+		{
+			Name:  "A_env",
+			Value: "A_value",
+		},
+	}
+
+	merged := mergeEnvs(existing, desired)
+
+	t.Run("EnvVars should be sorted", func(t *testing.T) {
+		assert.Equal(t, "A_env", merged[0].Name)
+		assert.Equal(t, "B_env", merged[1].Name)
+		assert.Equal(t, "C_env", merged[2].Name)
+		assert.Equal(t, "D_env", merged[3].Name)
+		assert.Equal(t, "E_env", merged[4].Name)
+		assert.Equal(t, "F_env", merged[5].Name)
+	})
+
+	t.Run("EnvVars of same name are updated", func(t *testing.T) {
+		assert.Equal(t, "B_env", merged[1].Name)
+		assert.Equal(t, "B_value_new", merged[1].Value)
+	})
+
+	t.Run("Existing EnvVars are not touched", func(t *testing.T) {
+		envVar := merged[5]
+		assert.NotNil(t, envVar.ValueFrom)
+		assert.Equal(t, "f_key", envVar.ValueFrom.SecretKeyRef.Key)
+	})
+}
+
+func boolRef(b bool) *bool {
+	return &b
+}
+
+func int64Ref(i int64) *int64 {
+	return &i
+}

pkg/kube/container/containers.go

@@ -1,27 +1,60 @@
 package container
 
-import corev1 "k8s.io/api/core/v1"
+import (
+	"sort"
+
+	"github.com/mongodb/mongodb-kubernetes-operator/pkg/kube/lifecycle"
+	corev1 "k8s.io/api/core/v1"
+)
 
 type Modification func(*corev1.Container)
 
+// Apply returns a function which applies a series of Modification functions to a *corev1.Container
+func Apply(modifications ...Modification) Modification {
+	return func(container *corev1.Container) {
+		for _, mod := range modifications {
+			mod(container)
+		}
+	}
+}
+
+// New returns a concrete corev1.Container instance which has been modified based on the provided
+// modifications
+func New(mods ...Modification) corev1.Container {
+	c := corev1.Container{}
+	for _, mod := range mods {
+		mod(&c)
+	}
+	return c
+}
+
+// NOOP is a valid Modification which applies no changes
+func NOOP() Modification {
+	return func(container *corev1.Container) {}
+}
+
+// WithName sets the container name
 func WithName(name string) Modification {
 	return func(container *corev1.Container) {
 		container.Name = name
 	}
 }
 
+// WithImage sets the container image
 func WithImage(image string) Modification {
 	return func(container *corev1.Container) {
 		container.Image = image
 	}
 }
 
+// WithImagePullPolicy sets the container pullPolicy
 func WithImagePullPolicy(pullPolicy corev1.PullPolicy) Modification {
 	return func(container *corev1.Container) {
 		container.ImagePullPolicy = pullPolicy
 	}
 }
 
+// WithReadinessProbe modifies the container's Readiness Probe
 func WithReadinessProbe(probeFunc func(*corev1.Probe)) Modification {
 	return func(container *corev1.Container) {
 		if container.ReadinessProbe == nil {
@@ -31,24 +64,70 @@ func WithReadinessProbe(probeFunc func(*corev1.Probe)) Modification {
 	}
 }
 
+// WithLivenessProbe modifies the container's Liveness Probe
+func WithLivenessProbe(readinessProbeFunc func(*corev1.Probe)) Modification {
+	return func(container *corev1.Container) {
+		if container.LivenessProbe == nil {
+			container.LivenessProbe = &corev1.Probe{}
+		}
+		readinessProbeFunc(container.LivenessProbe)
+	}
+}
+
+// WithResourceRequirements sets the container's Resources
 func WithResourceRequirements(resources corev1.ResourceRequirements) Modification {
 	return func(container *corev1.Container) {
 		container.Resources = resources
 	}
 }
 
+// WithCommand sets the containers Command
 func WithCommand(cmd []string) Modification {
 	return func(container *corev1.Container) {
 		container.Command = cmd
 	}
 }
 
-func WithEnv(envs []corev1.EnvVar) Modification {
+// WithLifecycle applies the lifecycle Modification to this container's
+// Lifecycle
+func WithLifecycle(lifeCycleMod lifecycle.Modification) Modification {
+	return func(container *corev1.Container) {
+		if container.Lifecycle == nil {
+			container.Lifecycle = &corev1.Lifecycle{}
+		}
+		lifeCycleMod(container.Lifecycle)
+	}
+}
+
+// WithEnvs ensures all of the provided envs exist in the container
+func WithEnvs(envs ...corev1.EnvVar) Modification {
 	return func(container *corev1.Container) {
-		container.Env = envs
+		container.Env = mergeEnvs(container.Env, envs)
+	}
+}
+
+func mergeEnvs(existing, desired []corev1.EnvVar) []corev1.EnvVar {
+	envMap := make(map[string]corev1.EnvVar)
+	for _, env := range existing {
+		envMap[env.Name] = env
+	}
+
+	for _, env := range desired {
+		envMap[env.Name] = env
+	}
+
+	var mergedEnv []corev1.EnvVar
+	for _, env := range envMap {
+		mergedEnv = append(mergedEnv, env)
 	}
+
+	sort.SliceStable(mergedEnv, func(i, j int) bool {
+		return mergedEnv[i].Name < mergedEnv[j].Name
+	})
+	return mergedEnv
 }
 
+// WithVolumeMounts sets the VolumeMounts
 func WithVolumeMounts(volumeMounts []corev1.VolumeMount) Modification {
 	volumesMountsCopy := make([]corev1.VolumeMount, len(volumeMounts))
 	copy(volumesMountsCopy, volumeMounts)
@@ -57,10 +136,16 @@ func WithVolumeMounts(volumeMounts []corev1.VolumeMount) Modification {
 	}
 }
 
-func Apply(modifications ...Modification) Modification {
+// WithPorts sets the container's Ports
+func WithPorts(ports []corev1.ContainerPort) Modification {
 	return func(container *corev1.Container) {
-		for _, mod := range modifications {
-			mod(container)
-		}
+		container.Ports = ports
+	}
+}
+
+// WithSecurityContext sets teh container's SecurityContext
+func WithSecurityContext(context corev1.SecurityContext) Modification {
+	return func(container *corev1.Container) {
+		container.SecurityContext = &context
 	}
 }

pkg/kube/lifecycle/lifecyle.go

@@ -0,0 +1,27 @@
+package lifecycle
+
+import corev1 "k8s.io/api/core/v1"
+
+type Modification func(lifecycle *corev1.Lifecycle)
+
+// Apply returns a function which applies a series of Modification functions to a *corev1.Lifecycle
+func Apply(modifications ...Modification) Modification {
+	return func(lifecycle *corev1.Lifecycle) {
+		for _, mod := range modifications {
+			mod(lifecycle)
+		}
+	}
+}
+
+// WithPrestopCommand sets the LifeCycles PreStop Exec Command
+func WithPrestopCommand(preStopCmd []string) Modification {
+	return func(lc *corev1.Lifecycle) {
+		if lc.PreStop == nil {
+			lc.PreStop = &corev1.Handler{}
+		}
+		if lc.PreStop.Exec == nil {
+			lc.PreStop.Exec = &corev1.ExecAction{}
+		}
+		lc.PreStop.Exec.Command = preStopCmd
+	}
+}