From 63e01b6b5fe4f8cc190b3b83dd740c9dfddc1ac7 Mon Sep 17 00:00:00 2001
From: Mitchell Nielsen <inbox@mitchn.me>
Date: Fri, 1 Nov 2024 15:08:57 -0500
Subject: [PATCH] Support setting extra args in Deployment command

Supports setting extra args in the Deployment command.

This will allow us to set `--no-ui` in cases where we don't need to
serve the UI.

Related to https://linear.app/prefect/issue/PLA-452/cycle-5-catch-all
---
 api/v1/prefectserver_types.go                 |  8 ++++-
 api/v1/zz_generated.deepcopy.go               |  5 ++++
 .../crds/prefect.io_prefectservers.yaml       |  6 ++++
 .../prefectserver_controller_test.go          | 29 +++++++++++++++++++
 4 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/api/v1/prefectserver_types.go b/api/v1/prefectserver_types.go
index 889aedd..38192fe 100644
--- a/api/v1/prefectserver_types.go
+++ b/api/v1/prefectserver_types.go
@@ -45,6 +45,9 @@ type PrefectServerSpec struct {
 	// ExtraServicePorts defines additional ports to expose on the Prefect Server Service
 	ExtraServicePorts []corev1.ServicePort `json:"extraServicePorts,omitempty"`
 
+	// ExtraArgs defines additional arguments to pass to the Prefect Server Deployment
+	ExtraArgs []string `json:"extraArgs,omitempty"`
+
 	// Ephemeral defines whether the Prefect Server will be deployed with an ephemeral storage backend
 	Ephemeral *EphemeralConfiguration `json:"ephemeral,omitempty"`
 
@@ -281,7 +284,10 @@ func (s *PrefectServer) Image() string {
 }
 
 func (s *PrefectServer) Command() []string {
-	return []string{"prefect", "server", "start", "--host", "0.0.0.0"}
+	command := []string{"prefect", "server", "start", "--host", "0.0.0.0"}
+	command = append(command, s.Spec.ExtraArgs...)
+
+	return command
 }
 
 func (s *PrefectServer) ToEnvVars() []corev1.EnvVar {
diff --git a/api/v1/zz_generated.deepcopy.go b/api/v1/zz_generated.deepcopy.go
index db79d76..5b13ccb 100644
--- a/api/v1/zz_generated.deepcopy.go
+++ b/api/v1/zz_generated.deepcopy.go
@@ -253,6 +253,11 @@ func (in *PrefectServerSpec) DeepCopyInto(out *PrefectServerSpec) {
 			(*in)[i].DeepCopyInto(&(*out)[i])
 		}
 	}
+	if in.ExtraArgs != nil {
+		in, out := &in.ExtraArgs, &out.ExtraArgs
+		*out = make([]string, len(*in))
+		copy(*out, *in)
+	}
 	if in.Ephemeral != nil {
 		in, out := &in.Ephemeral, &out.Ephemeral
 		*out = new(EphemeralConfiguration)
diff --git a/deploy/charts/prefect-operator/crds/prefect.io_prefectservers.yaml b/deploy/charts/prefect-operator/crds/prefect.io_prefectservers.yaml
index d2dcbcf..518a9ae 100644
--- a/deploy/charts/prefect-operator/crds/prefect.io_prefectservers.yaml
+++ b/deploy/charts/prefect-operator/crds/prefect.io_prefectservers.yaml
@@ -60,6 +60,12 @@ spec:
                 description: Ephemeral defines whether the Prefect Server will be
                   deployed with an ephemeral storage backend
                 type: object
+              extraArgs:
+                description: ExtraArgs defines additional arguments to pass to the
+                  Prefect Server Deployment
+                items:
+                  type: string
+                type: array
               extraContainers:
                 description: ExtraContainers defines additional containers to add
                   to the Prefect Server Deployment
diff --git a/internal/controller/prefectserver_controller_test.go b/internal/controller/prefectserver_controller_test.go
index 95eb4ed..4be9ffa 100644
--- a/internal/controller/prefectserver_controller_test.go
+++ b/internal/controller/prefectserver_controller_test.go
@@ -602,6 +602,35 @@ var _ = Describe("PrefectServer controller", func() {
 				port := service.Spec.Ports[1]
 				Expect(port.Name).To(Equal("extra-port"))
 			})
+
+			It("should update the Command with the extra args", func() {
+				// Update the PrefectServer with an extra args
+				Expect(k8sClient.Get(ctx, name, prefectserver)).To(Succeed())
+				prefectserver.Spec.ExtraArgs = []string{"--some-arg", "some-value"}
+
+				Expect(k8sClient.Update(ctx, prefectserver)).To(Succeed())
+				// Reconcile to apply the changes
+				controllerReconciler := &PrefectServerReconciler{
+					Client: k8sClient,
+					Scheme: k8sClient.Scheme(),
+				}
+				_, err := controllerReconciler.Reconcile(ctx, reconcile.Request{
+					NamespacedName: name,
+				})
+				Expect(err).NotTo(HaveOccurred())
+
+				// Check if the Deployment was updated with the extra args
+				deployment := &appsv1.Deployment{}
+				Eventually(func() error {
+					return k8sClient.Get(ctx, types.NamespacedName{
+						Namespace: namespaceName,
+						Name:      "prefect-on-anything",
+					}, deployment)
+				}).Should(Succeed())
+
+				container := deployment.Spec.Template.Spec.Containers[0]
+				Expect(container.Command).To(Equal([]string{"prefect", "server", "start", "--host", "0.0.0.0", "--some-arg", "some-value"}))
+			})
 		})
 
 		Context("When evaluating changes with any server", func() {