From 167e5d49fc018d3dfb4dc64ffb8bcc1cf9df1ed5 Mon Sep 17 00:00:00 2001 From: Chris Guidry Date: Tue, 1 Oct 2024 13:35:00 -0400 Subject: [PATCH] Allows for adding labels to the generated deployments, services, and jobs (#102) This addition gives users fine-grained control over adding extra labels to any of the deployment, services, or jobs produced by the operator. This gives folks control over labels they may need to add, like suppressing Istio sidecars for specific pods, or adding additional custom labels for application use. --- api/v1/prefectserver_types.go | 35 ++++++++++++++++++- api/v1/prefectworkpool_types.go | 11 +++++- api/v1/zz_generated.deepcopy.go | 28 +++++++++++++++ .../crds/prefect.io_prefectservers.yaml | 18 ++++++++++ .../crds/prefect.io_prefectworkpools.yaml | 6 ++++ .../controller/prefectserver_controller.go | 4 +-- .../prefectserver_controller_test.go | 25 +++++++++++++ .../prefectworkpool_controller_test.go | 11 ++++++ 8 files changed, 134 insertions(+), 4 deletions(-) diff --git a/api/v1/prefectserver_types.go b/api/v1/prefectserver_types.go index 9da49d0..0a51076 100644 --- a/api/v1/prefectserver_types.go +++ b/api/v1/prefectserver_types.go @@ -51,6 +51,15 @@ type PrefectServerSpec struct { // A list of environment variables to set on the Prefect Server Settings []corev1.EnvVar `json:"settings,omitempty"` + + // DeploymentLabels defines additional labels to add to the server Deployment + DeploymentLabels map[string]string `json:"deploymentLabels,omitempty"` + + // ServiceLabels defines additional labels to add to the server Service + ServiceLabels map[string]string `json:"serviceLabels,omitempty"` + + // MigrationJobLabels defines additional labels to add to the migration Job + MigrationJobLabels map[string]string `json:"migrationJobLabels,omitempty"` } type EphemeralConfiguration struct { @@ -220,9 +229,33 @@ type PrefectServer struct { } func (s *PrefectServer) ServerLabels() map[string]string { - return map[string]string{ + labels := map[string]string{ + "prefect.io/server": s.Name, + } + for k, v := range s.Spec.DeploymentLabels { + labels[k] = v + } + return labels +} + +func (s *PrefectServer) ServiceLabels() map[string]string { + labels := map[string]string{ "prefect.io/server": s.Name, } + for k, v := range s.Spec.ServiceLabels { + labels[k] = v + } + return labels +} + +func (s *PrefectServer) MigrationJobLabels() map[string]string { + labels := map[string]string{ + "prefect.io/server": s.Name, + } + for k, v := range s.Spec.MigrationJobLabels { + labels[k] = v + } + return labels } func (s *PrefectServer) Image() string { diff --git a/api/v1/prefectworkpool_types.go b/api/v1/prefectworkpool_types.go index d4ad96d..ff2f46a 100644 --- a/api/v1/prefectworkpool_types.go +++ b/api/v1/prefectworkpool_types.go @@ -47,6 +47,9 @@ type PrefectWorkPoolSpec struct { // A list of environment variables to set on the Prefect Worker Settings []corev1.EnvVar `json:"settings,omitempty"` + + // DeploymentLabels defines additional labels to add to the server deployment + DeploymentLabels map[string]string `json:"deploymentLabels,omitempty"` } type PrefectServerReference struct { @@ -111,9 +114,15 @@ type PrefectWorkPool struct { } func (s *PrefectWorkPool) WorkerLabels() map[string]string { - return map[string]string{ + labels := map[string]string{ "prefect.io/worker": s.Name, } + + for k, v := range s.Spec.DeploymentLabels { + labels[k] = v + } + + return labels } func (s *PrefectWorkPool) Image() string { diff --git a/api/v1/zz_generated.deepcopy.go b/api/v1/zz_generated.deepcopy.go index 9fc7293..2339a65 100644 --- a/api/v1/zz_generated.deepcopy.go +++ b/api/v1/zz_generated.deepcopy.go @@ -261,6 +261,27 @@ func (in *PrefectServerSpec) DeepCopyInto(out *PrefectServerSpec) { (*in)[i].DeepCopyInto(&(*out)[i]) } } + if in.DeploymentLabels != nil { + in, out := &in.DeploymentLabels, &out.DeploymentLabels + *out = make(map[string]string, len(*in)) + for key, val := range *in { + (*out)[key] = val + } + } + if in.ServiceLabels != nil { + in, out := &in.ServiceLabels, &out.ServiceLabels + *out = make(map[string]string, len(*in)) + for key, val := range *in { + (*out)[key] = val + } + } + if in.MigrationJobLabels != nil { + in, out := &in.MigrationJobLabels, &out.MigrationJobLabels + *out = make(map[string]string, len(*in)) + for key, val := range *in { + (*out)[key] = val + } + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new PrefectServerSpec. @@ -376,6 +397,13 @@ func (in *PrefectWorkPoolSpec) DeepCopyInto(out *PrefectWorkPoolSpec) { (*in)[i].DeepCopyInto(&(*out)[i]) } } + if in.DeploymentLabels != nil { + in, out := &in.DeploymentLabels, &out.DeploymentLabels + *out = make(map[string]string, len(*in)) + for key, val := range *in { + (*out)[key] = val + } + } } // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new PrefectWorkPoolSpec. diff --git a/deploy/charts/prefect-operator/crds/prefect.io_prefectservers.yaml b/deploy/charts/prefect-operator/crds/prefect.io_prefectservers.yaml index 63617f6..430f270 100644 --- a/deploy/charts/prefect-operator/crds/prefect.io_prefectservers.yaml +++ b/deploy/charts/prefect-operator/crds/prefect.io_prefectservers.yaml @@ -48,6 +48,12 @@ spec: spec: description: PrefectServerSpec defines the desired state of a PrefectServer properties: + deploymentLabels: + additionalProperties: + type: string + description: DeploymentLabels defines additional labels to add to + the server Deployment + type: object ephemeral: description: Ephemeral defines whether the server will be deployed with an ephemeral storage backend @@ -56,6 +62,12 @@ spec: description: Image defines the exact image to deploy for the Prefect Server, overriding Version type: string + migrationJobLabels: + additionalProperties: + type: string + description: MigrationJobLabels defines additional labels to add to + the migration Job + type: object postgres: description: |- Postgres defines whether the server will be deployed with a PostgreSQL backend connecting to the @@ -563,6 +575,12 @@ spec: More info: https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/ type: object type: object + serviceLabels: + additionalProperties: + type: string + description: ServiceLabels defines additional labels to add to the + server Service + type: object settings: description: A list of environment variables to set on the Prefect Server diff --git a/deploy/charts/prefect-operator/crds/prefect.io_prefectworkpools.yaml b/deploy/charts/prefect-operator/crds/prefect.io_prefectworkpools.yaml index 3588a57..0326c7f 100644 --- a/deploy/charts/prefect-operator/crds/prefect.io_prefectworkpools.yaml +++ b/deploy/charts/prefect-operator/crds/prefect.io_prefectworkpools.yaml @@ -60,6 +60,12 @@ spec: spec: description: PrefectWorkPoolSpec defines the desired state of PrefectWorkPool properties: + deploymentLabels: + additionalProperties: + type: string + description: DeploymentLabels defines additional labels to add to + the server deployment + type: object image: description: Image defines the exact image to deploy for the Prefect Server, overriding Version diff --git a/internal/controller/prefectserver_controller.go b/internal/controller/prefectserver_controller.go index edd82b0..1d56479 100644 --- a/internal/controller/prefectserver_controller.go +++ b/internal/controller/prefectserver_controller.go @@ -526,7 +526,7 @@ func (r *PrefectServerReconciler) postgresMigrationJob(server *prefectiov1.Prefe TTLSecondsAfterFinished: ptr.To(int32(7 * 24 * 60 * 60)), // 7 days Template: corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{ - Labels: server.ServerLabels(), + Labels: server.MigrationJobLabels(), }, Spec: corev1.PodSpec{ InitContainers: []corev1.Container{ @@ -569,7 +569,7 @@ func (r *PrefectServerReconciler) prefectServerService(server *prefectiov1.Prefe Name: server.Name, }, Spec: corev1.ServiceSpec{ - Selector: server.ServerLabels(), + Selector: server.ServiceLabels(), Ports: []corev1.ServicePort{ { Name: "api", diff --git a/internal/controller/prefectserver_controller_test.go b/internal/controller/prefectserver_controller_test.go index f0c3135..bf029e5 100644 --- a/internal/controller/prefectserver_controller_test.go +++ b/internal/controller/prefectserver_controller_test.go @@ -182,6 +182,10 @@ var _ = Describe("PrefectServer controller", func() { corev1.ResourceMemory: resource.MustParse("512Mi"), }, }, + DeploymentLabels: map[string]string{ + "some": "additional-label", + "another": "extra-label", + }, }, } Expect(k8sClient.Create(ctx, prefectserver)).To(Succeed()) @@ -265,6 +269,13 @@ var _ = Describe("PrefectServer controller", func() { It("should have appropriate labels", func() { Expect(deployment.Spec.Selector.MatchLabels).To(Equal(map[string]string{ "prefect.io/server": "prefect-on-anything", + "some": "additional-label", + "another": "extra-label", + })) + Expect(deployment.Spec.Template.Labels).To(Equal(map[string]string{ + "prefect.io/server": "prefect-on-anything", + "some": "additional-label", + "another": "extra-label", })) }) @@ -1093,6 +1104,14 @@ var _ = Describe("PrefectServer controller", func() { Password: ptr.To("this-is-a-bad-idea"), Database: ptr.To("some-prefect"), }, + DeploymentLabels: map[string]string{ + "some": "additional-label", + "another": "extra-label", + }, + MigrationJobLabels: map[string]string{ + "some": "additional-label-for-migrations", + "another": "extra-label-for-migrations", + }, }, } Expect(k8sClient.Create(ctx, prefectserver)).To(Succeed()) @@ -1211,6 +1230,12 @@ var _ = Describe("PrefectServer controller", func() { )) }) + It("should have the correct labels", func() { + Expect(migrateJob.Labels).To(HaveKeyWithValue("prefect.io/server", "prefect-on-postgres")) + Expect(migrateJob.Labels).To(HaveKeyWithValue("some", "additional-label-for-migrations")) + Expect(migrateJob.Labels).To(HaveKeyWithValue("another", "extra-label-for-migrations")) + }) + It("should have an environment pointing to the PostgreSQL database", func() { Expect(migrateJob.Spec.Template.Spec.Containers).To(HaveLen(1)) container := migrateJob.Spec.Template.Spec.Containers[0] diff --git a/internal/controller/prefectworkpool_controller_test.go b/internal/controller/prefectworkpool_controller_test.go index aa4dd62..f1afe0d 100644 --- a/internal/controller/prefectworkpool_controller_test.go +++ b/internal/controller/prefectworkpool_controller_test.go @@ -181,6 +181,10 @@ var _ = Describe("PrefectWorkPool Controller", func() { corev1.ResourceMemory: resource.MustParse("512Mi"), }, }, + DeploymentLabels: map[string]string{ + "some": "additional-label", + "another": "extra-label", + }, }, } Expect(k8sClient.Create(ctx, prefectworkpool)).To(Succeed()) @@ -232,6 +236,13 @@ var _ = Describe("PrefectWorkPool Controller", func() { It("should have appropriate labels", func() { Expect(deployment.Spec.Selector.MatchLabels).To(Equal(map[string]string{ "prefect.io/worker": "example-work-pool", + "some": "additional-label", + "another": "extra-label", + })) + Expect(deployment.Spec.Template.Labels).To(Equal(map[string]string{ + "prefect.io/worker": "example-work-pool", + "some": "additional-label", + "another": "extra-label", })) })