diff --git a/pkg/controller/mongodb/mongodb_controller.go b/pkg/controller/mongodb/mongodb_controller.go index 7cc70dad7..d397ebf9c 100644 --- a/pkg/controller/mongodb/mongodb_controller.go +++ b/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), diff --git a/pkg/kube/container/container_test.go b/pkg/kube/container/container_test.go new file mode 100644 index 000000000..23c3c611e --- /dev/null +++ b/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 +} diff --git a/pkg/kube/container/containers.go b/pkg/kube/container/containers.go index 09b061a35..604cc486b 100644 --- a/pkg/kube/container/containers.go +++ b/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 } } diff --git a/pkg/kube/lifecycle/lifecyle.go b/pkg/kube/lifecycle/lifecyle.go new file mode 100644 index 000000000..3905a1659 --- /dev/null +++ b/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 + } +} diff --git a/pkg/kube/persistentvolumeclaim/pvc.go b/pkg/kube/persistentvolumeclaim/pvc.go index bf3715bdd..c719527b4 100644 --- a/pkg/kube/persistentvolumeclaim/pvc.go +++ b/pkg/kube/persistentvolumeclaim/pvc.go @@ -2,10 +2,12 @@ package persistentvolumeclaim import ( corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ) type Modification func(claim *corev1.PersistentVolumeClaim) +// Apply returns a function which applies a series of Modification functions to a *corev1.PersistentVolumeClaim func Apply(funcs ...Modification) Modification { return func(claim *corev1.PersistentVolumeClaim) { for _, f := range funcs { @@ -14,20 +16,42 @@ func Apply(funcs ...Modification) Modification { } } +// NOOP is a valid Modification which applies no changes +func NOOP() Modification { + return func(claim *corev1.PersistentVolumeClaim) {} +} + +// WithName sets the PersistentVolumeClaim's name func WithName(name string) Modification { return func(claim *corev1.PersistentVolumeClaim) { claim.Name = name } } +// WithAccessModes sets the PersistentVolumeClaim's AccessModes func WithAccessModes(accessMode corev1.PersistentVolumeAccessMode) Modification { return func(claim *corev1.PersistentVolumeClaim) { claim.Spec.AccessModes = []corev1.PersistentVolumeAccessMode{accessMode} } } +// WithResourceRequests sets the PersistentVolumeClaim's Resource Requests func WithResourceRequests(requests corev1.ResourceList) Modification { return func(claim *corev1.PersistentVolumeClaim) { claim.Spec.Resources.Requests = requests } } + +// WithLabelSelector sets the PersistentVolumeClaim's LevelSelector +func WithLabelSelector(selector *metav1.LabelSelector) Modification { + return func(claim *corev1.PersistentVolumeClaim) { + claim.Spec.Selector = selector + } +} + +// WithStorageClassName sets the PersistentVolumeClaim's storage class name +func WithStorageClassName(storageClassName string) Modification { + return func(claim *corev1.PersistentVolumeClaim) { + claim.Spec.StorageClassName = &storageClassName + } +} diff --git a/pkg/kube/podtemplatespec/podsoec_template_test.go b/pkg/kube/podtemplatespec/podsoec_template_test.go new file mode 100644 index 000000000..c30e90ec6 --- /dev/null +++ b/pkg/kube/podtemplatespec/podsoec_template_test.go @@ -0,0 +1,81 @@ +package podtemplatespec + +import ( + "testing" + + "github.com/mongodb/mongodb-kubernetes-operator/pkg/kube/container" + "github.com/stretchr/testify/assert" + corev1 "k8s.io/api/core/v1" +) + +func TestPodTemplateSpec(t *testing.T) { + p := New( + WithVolume(corev1.Volume{ + Name: "vol-1", + }), + WithVolume(corev1.Volume{ + Name: "vol-2", + }), + WithFsGroup(100), + WithImagePullSecrets("pull-secrets"), + WithInitContainerByIndex(0, container.Apply( + container.WithName("init-container-0"), + container.WithImage("init-image"), + )), + WithContainerByIndex(0, container.Apply( + container.WithName("container-0"), + container.WithImage("image"), + )), + WithContainerByIndex(1, container.Apply( + container.WithName("container-1"), + container.WithImage("image"), + )), + ) + + assert.Len(t, p.Spec.Volumes, 2) + assert.Equal(t, p.Spec.Volumes[0].Name, "vol-1") + assert.Equal(t, p.Spec.Volumes[1].Name, "vol-2") + + expected := int64(100) + assert.Equal(t, &expected, p.Spec.SecurityContext.FSGroup) + + assert.Len(t, p.Spec.ImagePullSecrets, 1) + assert.Equal(t, "pull-secrets", p.Spec.ImagePullSecrets[0].Name) + + assert.Len(t, p.Spec.InitContainers, 1) + assert.Equal(t, "init-container-0", p.Spec.InitContainers[0].Name) + assert.Equal(t, "init-image", p.Spec.InitContainers[0].Image) + + assert.Len(t, p.Spec.Containers, 2) + assert.Equal(t, "container-0", p.Spec.Containers[0].Name) + assert.Equal(t, "image", p.Spec.Containers[0].Image) + assert.Equal(t, "container-1", p.Spec.Containers[1].Name) + assert.Equal(t, "image", p.Spec.Containers[1].Image) +} + +func TestPodTemplateSpec_MultipleEditsToContainer(t *testing.T) { + p := New( + WithContainerByIndex(0, + container.Apply( + container.WithName("container-0"), + )), + WithContainerByIndex(0, + container.Apply( + container.WithImage("image"), + )), + WithContainerByIndex(0, + container.Apply( + container.WithImagePullPolicy(corev1.PullAlways), + )), + WithContainer("container-0", container.Apply( + container.WithCommand([]string{"cmd"}), + )), + ) + + assert.Len(t, p.Spec.Containers, 1) + c := p.Spec.Containers[0] + assert.Equal(t, "container-0", c.Name) + assert.Equal(t, "image", c.Image) + assert.Equal(t, corev1.PullAlways, c.ImagePullPolicy) + assert.Equal(t, "cmd", c.Command[0]) +} diff --git a/pkg/kube/podtemplatespec/podspec_template.go b/pkg/kube/podtemplatespec/podspec_template.go index f4d6550d8..ad32960be 100644 --- a/pkg/kube/podtemplatespec/podspec_template.go +++ b/pkg/kube/podtemplatespec/podspec_template.go @@ -2,6 +2,7 @@ package podtemplatespec import ( corev1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ) type Modification func(*corev1.PodTemplateSpec) @@ -10,6 +11,7 @@ const ( notFound = -1 ) +// Apply returns a function which applies a series of Modification functions to a *corev1.PodTemplateSpec func Apply(templateMods ...Modification) Modification { return func(template *corev1.PodTemplateSpec) { for _, f := range templateMods { @@ -17,6 +19,23 @@ func Apply(templateMods ...Modification) Modification { } } } + +// New returns a concrete corev1.PodTemplateSpec instance which has been modified based on the provided +// modifications +func New(templateMods ...Modification) corev1.PodTemplateSpec { + podTemplateSpec := corev1.PodTemplateSpec{} + for _, templateMod := range templateMods { + templateMod(&podTemplateSpec) + } + return podTemplateSpec +} + +// NOOP is a valid Modification which applies no changes +func NOOP() Modification { + return func(spec *corev1.PodTemplateSpec) {} +} + +// WithContainer applies the modifications to the container with the provided name func WithContainer(name string, containerfunc func(*corev1.Container)) Modification { return func(podTemplateSpec *corev1.PodTemplateSpec) { idx := findIndexByName(name, podTemplateSpec.Spec.Containers) @@ -30,6 +49,21 @@ func WithContainer(name string, containerfunc func(*corev1.Container)) Modificat } } +// WithContainerByIndex applies the modifications to the container with the provided index +// if the index is out of range, a new container is added to accept these changes. +func WithContainerByIndex(index int, funcs ...func(container *corev1.Container)) func(podTemplateSpec *corev1.PodTemplateSpec) { + return func(podTemplateSpec *corev1.PodTemplateSpec) { + if index >= len(podTemplateSpec.Spec.Containers) { + podTemplateSpec.Spec.Containers = append(podTemplateSpec.Spec.Containers, corev1.Container{}) + } + c := &podTemplateSpec.Spec.Containers[index] + for _, f := range funcs { + f(c) + } + } +} + +// WithInitContainer applies the modifications to the init container with the provided name func WithInitContainer(name string, containerfunc func(*corev1.Container)) Modification { return func(podTemplateSpec *corev1.PodTemplateSpec) { idx := findIndexByName(name, podTemplateSpec.Spec.InitContainers) @@ -43,6 +77,21 @@ func WithInitContainer(name string, containerfunc func(*corev1.Container)) Modif } } +// WithInitContainerByIndex applies the modifications to the container with the provided index +// if the index is out of range, a new container is added to accept these changes. +func WithInitContainerByIndex(index int, funcs ...func(container *corev1.Container)) func(podTemplateSpec *corev1.PodTemplateSpec) { + return func(podTemplateSpec *corev1.PodTemplateSpec) { + if index >= len(podTemplateSpec.Spec.Containers) { + podTemplateSpec.Spec.InitContainers = append(podTemplateSpec.Spec.InitContainers, corev1.Container{}) + } + c := &podTemplateSpec.Spec.InitContainers[index] + for _, f := range funcs { + f(c) + } + } +} + +// WithPodLabels sets the PodTemplateSpec's Labels func WithPodLabels(labels map[string]string) Modification { if labels == nil { labels = map[string]string{} @@ -52,12 +101,14 @@ func WithPodLabels(labels map[string]string) Modification { } } +// WithServiceAccount sets the PodTemplateSpec's ServiceAccount name func WithServiceAccount(serviceAccountName string) Modification { return func(podTemplateSpec *corev1.PodTemplateSpec) { podTemplateSpec.Spec.ServiceAccountName = serviceAccountName } } +// WithVolume ensures the given volume exists func WithVolume(volume corev1.Volume) Modification { return func(template *corev1.PodTemplateSpec) { for _, v := range template.Spec.Volumes { @@ -77,3 +128,86 @@ func findIndexByName(name string, containers []corev1.Container) int { } return notFound } + +// WithTerminationGracePeriodSeconds sets the PodTemplateSpec's termination grace period seconds +func WithTerminationGracePeriodSeconds(seconds int) Modification { + s := int64(seconds) + return func(podTemplateSpec *corev1.PodTemplateSpec) { + podTemplateSpec.Spec.TerminationGracePeriodSeconds = &s + } +} + +// WithFsGroup sets the PodTemplateSpec's fs group +func WithFsGroup(fsGroup int) Modification { + return func(podTemplateSpec *corev1.PodTemplateSpec) { + spec := &podTemplateSpec.Spec + fsGroup64 := int64(fsGroup) + spec.SecurityContext = &corev1.PodSecurityContext{ + FSGroup: &fsGroup64, + } + } +} + +// WithImagePullSecrets adds an ImagePullSecrets local reference with the given name +func WithImagePullSecrets(name string) Modification { + return func(podTemplateSpec *corev1.PodTemplateSpec) { + podTemplateSpec.Spec.ImagePullSecrets = append(podTemplateSpec.Spec.ImagePullSecrets, corev1.LocalObjectReference{ + Name: name, + }) + } +} + +// WithTopologyKey sets the PodTemplateSpec's topology at a given index +func WithTopologyKey(topologyKey string, idx int) Modification { + return func(podTemplateSpec *corev1.PodTemplateSpec) { + podTemplateSpec.Spec.Affinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution[idx].PodAffinityTerm.TopologyKey = topologyKey + } +} + +// WithAffinity updates the name, antiAffinityLabelKey and weight of the PodTemplateSpec's Affinity +func WithAffinity(stsName, antiAffinityLabelKey string, weight int) Modification { + return func(podTemplateSpec *corev1.PodTemplateSpec) { + podTemplateSpec.Spec.Affinity = + &corev1.Affinity{ + PodAntiAffinity: &corev1.PodAntiAffinity{ + PreferredDuringSchedulingIgnoredDuringExecution: []corev1.WeightedPodAffinityTerm{{ + Weight: int32(weight), + PodAffinityTerm: corev1.PodAffinityTerm{ + LabelSelector: &metav1.LabelSelector{MatchLabels: map[string]string{antiAffinityLabelKey: stsName}}, + }, + }}, + }, + } + } +} + +// WithNodeAffinity sets the PodTemplateSpec's node affinity +func WithNodeAffinity(nodeAffinity *corev1.NodeAffinity) Modification { + return func(podTemplateSpec *corev1.PodTemplateSpec) { + podTemplateSpec.Spec.Affinity.NodeAffinity = nodeAffinity + } +} + +// WithPodAffinity sets the PodTemplateSpec's pod affinity +func WithPodAffinity(podAffinity *corev1.PodAffinity) Modification { + return func(podTemplateSpec *corev1.PodTemplateSpec) { + podTemplateSpec.Spec.Affinity.PodAffinity = podAffinity + } +} + +// WithTolerations sets the PodTemplateSpec's tolerations +func WithTolerations(tolerations []corev1.Toleration) Modification { + return func(podTemplateSpec *corev1.PodTemplateSpec) { + podTemplateSpec.Spec.Tolerations = tolerations + } +} + +// WithAnnotations sets the PodTemplateSpec's annotations +func WithAnnotations(annotations map[string]string) Modification { + if annotations == nil { + annotations = map[string]string{} + } + return func(podTemplateSpec *corev1.PodTemplateSpec) { + podTemplateSpec.Annotations = annotations + } +} diff --git a/pkg/kube/probes/readiness.go b/pkg/kube/probes/probes.go similarity index 60% rename from pkg/kube/probes/readiness.go rename to pkg/kube/probes/probes.go index 71440bfe5..44b86b6b2 100644 --- a/pkg/kube/probes/readiness.go +++ b/pkg/kube/probes/probes.go @@ -4,6 +4,22 @@ import corev1 "k8s.io/api/core/v1" type Modification func(*corev1.Probe) +func Apply(funcs ...Modification) Modification { + return func(probe *corev1.Probe) { + for _, f := range funcs { + f(probe) + } + } +} + +func New(funcs ...Modification) corev1.Probe { + probe := corev1.Probe{} + for _, f := range funcs { + f(&probe) + } + return probe +} + func WithExecCommand(cmd []string) Modification { return func(probe *corev1.Probe) { if probe.Handler.Exec == nil { @@ -24,19 +40,24 @@ func WithInitialDelaySeconds(initialDelaySeconds int) Modification { probe.InitialDelaySeconds = int32(initialDelaySeconds) } } - -func Apply(funcs ...Modification) Modification { +func WithSuccessThreshold(successThreshold int) Modification { return func(probe *corev1.Probe) { - for _, f := range funcs { - f(probe) - } + probe.SuccessThreshold = int32(successThreshold) + } +} +func WithPeriodSeconds(periodSeconds int) Modification { + return func(probe *corev1.Probe) { + probe.PeriodSeconds = int32(periodSeconds) + } +} +func WithTimeoutSeconds(initialDelaySeconds int) Modification { + return func(probe *corev1.Probe) { + probe.InitialDelaySeconds = int32(initialDelaySeconds) } } -func New(funcs ...Modification) corev1.Probe { - probe := corev1.Probe{} - for _, f := range funcs { - f(&probe) +func WithHandler(handler corev1.Handler) Modification { + return func(probe *corev1.Probe) { + probe.Handler = handler } - return probe } diff --git a/pkg/kube/statefulset/statefulset.go b/pkg/kube/statefulset/statefulset.go index e9c0c7aa2..d313bdd0e 100644 --- a/pkg/kube/statefulset/statefulset.go +++ b/pkg/kube/statefulset/statefulset.go @@ -86,6 +86,22 @@ func IsReady(sts appsv1.StatefulSet, expectedReplicas int) bool { type Modification func(*appsv1.StatefulSet) +func New(mods ...Modification) appsv1.StatefulSet { + sts := appsv1.StatefulSet{} + for _, mod := range mods { + mod(&sts) + } + return sts +} + +func Apply(funcs ...Modification) func(*appsv1.StatefulSet) { + return func(sts *appsv1.StatefulSet) { + for _, f := range funcs { + f(sts) + } + } +} + func WithName(name string) Modification { return func(sts *appsv1.StatefulSet) { sts.Name = name @@ -167,11 +183,3 @@ func findVolumeClaimIndexByName(name string, pvcs []corev1.PersistentVolumeClaim } return notFound } - -func Apply(funcs ...Modification) func(*appsv1.StatefulSet) { - return func(sts *appsv1.StatefulSet) { - for _, f := range funcs { - f(sts) - } - } -}