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", })) })