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.
PrefectHQ · Oct 1, 2024 · 167e5d4 · 167e5d4
1 parent 38593d8
commit 167e5d4
Show file tree

Hide file tree

Showing 8 changed files with 134 additions and 4 deletions.
diff --git 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
@@ -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
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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",
 				}))
 			})