Allows for adding labels to the generated deployments, services, and jobs

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 {

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 {

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

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

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

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]

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