feat(Deployments): support Job Variables (#250)

* feat(Deployments): support Job Variables * Use 'env' key for job_variables test and example The 'env' key is "the only job variable that is configurable for all work pool types" (https://docs.prefect.io/latest/guides/deployment/overriding-job-variables/) so let's use that as the example and as the test to confirm it works as expected. * Address golangci-lint var naming changes --------- Co-authored-by: Mitchell Nielsen <[email protected]>
PrefectHQ · Aug 27, 2024 · 20413da · 20413da
1 parent 0a82c91
commit 20413da
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 9 deletions.
diff --git a/docs/resources/deployment.md b/docs/resources/deployment.md
@@ -32,7 +32,10 @@ resource "prefect_deployment" "deployment" {
   entrypoint               = "hello_world.py:hello_world"
   tags                     = ["test"]
   enforce_parameter_schema = false
-  manifest_path            = "./bar/foo"
+  job_variables = jsonencode({
+    "env" : { "some-key" : "some-value" }
+  })
+  manifest_path = "./bar/foo"
   parameters = jsonencode({
     "some-parameter" : "some-value",
     "some-parameter2" : "some-value2"
@@ -59,6 +62,7 @@ resource "prefect_deployment" "deployment" {
 - `description` (String) A description for the deployment.
 - `enforce_parameter_schema` (Boolean) Whether or not the deployment should enforce the parameter schema.
 - `entrypoint` (String) The path to the entrypoint for the workflow, relative to the path.
+- `job_variables` (String) Overrides for the flow's infrastructure configuration.
 - `manifest_path` (String) The path to the flow's manifest file, relative to the chosen storage.
 - `parameters` (String) Parameters for flow runs scheduled by the deployment.
 - `path` (String) The path to the working directory for the workflow, relative to remote storage or an absolute path.

diff --git a/examples/resources/prefect_deployment/resource.tf b/examples/resources/prefect_deployment/resource.tf
@@ -17,7 +17,10 @@ resource "prefect_deployment" "deployment" {
   entrypoint               = "hello_world.py:hello_world"
   tags                     = ["test"]
   enforce_parameter_schema = false
-  manifest_path            = "./bar/foo"
+  job_variables = jsonencode({
+    "env" : { "some-key" : "some-value" }
+  })
+  manifest_path = "./bar/foo"
   parameters = jsonencode({
     "some-parameter" : "some-value",
     "some-parameter2" : "some-value2"

diff --git a/internal/api/deployments.go b/internal/api/deployments.go
@@ -25,6 +25,7 @@ type Deployment struct {
 	EnforceParameterSchema bool                   `json:"enforce_parameter_schema"`
 	Entrypoint             string                 `json:"entrypoint"`
 	FlowID                 uuid.UUID              `json:"flow_id"`
+	JobVariables           map[string]interface{} `json:"job_variables,omitempty"`
 	ManifestPath           string                 `json:"manifest_path,omitempty"`
 	Name                   string                 `json:"name"`
 	Parameters             map[string]interface{} `json:"parameters,omitempty"`
@@ -42,6 +43,7 @@ type DeploymentCreate struct {
 	EnforceParameterSchema bool                   `json:"enforce_parameter_schema,omitempty"`
 	Entrypoint             string                 `json:"entrypoint,omitempty"`
 	FlowID                 uuid.UUID              `json:"flow_id"`
+	JobVariables           map[string]interface{} `json:"job_variables,omitempty"`
 	ManifestPath           string                 `json:"manifest_path,omitempty"`
 	Name                   string                 `json:"name"`
 	Parameters             map[string]interface{} `json:"parameters,omitempty"`
@@ -58,6 +60,7 @@ type DeploymentUpdate struct {
 	Description            string                 `json:"description,omitempty"`
 	EnforceParameterSchema bool                   `json:"enforce_parameter_schema,omitempty"`
 	Entrypoint             string                 `json:"entrypoint,omitempty"`
+	JobVariables           map[string]interface{} `json:"job_variables,omitempty"`
 	ManifestPath           string                 `json:"manifest_path"`
 	Parameters             map[string]interface{} `json:"parameters,omitempty"`
 	Path                   string                 `json:"path,omitempty"`

diff --git a/internal/provider/resources/deployment.go b/internal/provider/resources/deployment.go
@@ -48,6 +48,7 @@ type DeploymentResourceModel struct {
 	EnforceParameterSchema types.Bool            `tfsdk:"enforce_parameter_schema"`
 	Entrypoint             types.String          `tfsdk:"entrypoint"`
 	FlowID                 customtypes.UUIDValue `tfsdk:"flow_id"`
+	JobVariables           jsontypes.Normalized  `tfsdk:"job_variables"`
 	ManifestPath           types.String          `tfsdk:"manifest_path"`
 	Name                   types.String          `tfsdk:"name"`
 	Parameters             jsontypes.Normalized  `tfsdk:"parameters"`
@@ -163,6 +164,12 @@ func (r *DeploymentResource) Schema(_ context.Context, _ resource.SchemaRequest,
 					stringplanmodifier.UseStateForUnknown(),
 				},
 			},
+			"job_variables": schema.StringAttribute{
+				Description: "Overrides for the flow's infrastructure configuration.",
+				Optional:    true,
+				Computed:    true,
+				CustomType:  jsontypes.NormalizedType{},
+			},
 			"work_queue_name": schema.StringAttribute{
 				Description: "The work queue for the deployment. If no work queue is set, work will not be scheduled.",
 				Optional:    true,
@@ -279,8 +286,14 @@ func (r *DeploymentResource) Create(ctx context.Context, req resource.CreateRequ
 		return
 	}
 
-	var data map[string]interface{}
-	resp.Diagnostics.Append(plan.Parameters.Unmarshal(&data)...)
+	var parameters map[string]interface{}
+	resp.Diagnostics.Append(plan.Parameters.Unmarshal(&parameters)...)
+	if resp.Diagnostics.HasError() {
+		return
+	}
+
+	var jobVariables map[string]interface{}
+	resp.Diagnostics.Append(plan.JobVariables.Unmarshal(&jobVariables)...)
 	if resp.Diagnostics.HasError() {
 		return
 	}
@@ -290,9 +303,10 @@ func (r *DeploymentResource) Create(ctx context.Context, req resource.CreateRequ
 		EnforceParameterSchema: plan.EnforceParameterSchema.ValueBool(),
 		Entrypoint:             plan.Entrypoint.ValueString(),
 		FlowID:                 plan.FlowID.ValueUUID(),
+		JobVariables:           jobVariables,
 		ManifestPath:           plan.ManifestPath.ValueString(),
 		Name:                   plan.Name.ValueString(),
-		Parameters:             data,
+		Parameters:             parameters,
 		Path:                   plan.Path.ValueString(),
 		Paused:                 plan.Paused.ValueBool(),
 		Tags:                   tags,
@@ -371,11 +385,17 @@ func (r *DeploymentResource) Read(ctx context.Context, req resource.ReadRequest,
 		return
 	}
 
-	byteSlice, err := json.Marshal(deployment.Parameters)
+	parametersByteSlice, err := json.Marshal(deployment.Parameters)
 	if err != nil {
 		resp.Diagnostics.Append(helpers.SerializeDataErrorDiagnostic("parameters", "Deployment parameters", err))
 	}
-	model.Parameters = jsontypes.NewNormalizedValue(string(byteSlice))
+	model.Parameters = jsontypes.NewNormalizedValue(string(parametersByteSlice))
+
+	jobVariablesByteSlice, err := json.Marshal(deployment.JobVariables)
+	if err != nil {
+		resp.Diagnostics.Append(helpers.SerializeDataErrorDiagnostic("job_variables", "Deployment job variables", err))
+	}
+	model.JobVariables = jsontypes.NewNormalizedValue(string(jobVariablesByteSlice))
 
 	resp.Diagnostics.Append(resp.State.Set(ctx, &model)...)
 	if resp.Diagnostics.HasError() {
@@ -423,10 +443,17 @@ func (r *DeploymentResource) Update(ctx context.Context, req resource.UpdateRequ
 		return
 	}
 
+	var jobVariables map[string]interface{}
+	resp.Diagnostics.Append(model.JobVariables.Unmarshal(&jobVariables)...)
+	if resp.Diagnostics.HasError() {
+		return
+	}
+
 	payload := api.DeploymentUpdate{
 		Description:            model.Description.ValueString(),
 		EnforceParameterSchema: model.EnforceParameterSchema.ValueBool(),
 		Entrypoint:             model.Entrypoint.ValueString(),
+		JobVariables:           jobVariables,
 		ManifestPath:           model.ManifestPath.ValueString(),
 		Parameters:             parameters,
 		Path:                   model.Path.ValueString(),
@@ -462,13 +489,21 @@ func (r *DeploymentResource) Update(ctx context.Context, req resource.UpdateRequ
 		return
 	}
 
-	byteSlice, err := json.Marshal(deployment.Parameters)
+	parametersByteSlice, err := json.Marshal(deployment.Parameters)
 	if err != nil {
 		resp.Diagnostics.Append(helpers.SerializeDataErrorDiagnostic("parameters", "Deployment parameters", err))
 
 		return
 	}
-	model.Parameters = jsontypes.NewNormalizedValue(string(byteSlice))
+	model.Parameters = jsontypes.NewNormalizedValue(string(parametersByteSlice))
+
+	jobVariablesByteSlice, err := json.Marshal(deployment.JobVariables)
+	if err != nil {
+		resp.Diagnostics.Append(helpers.SerializeDataErrorDiagnostic("job_variables", "Deployment job variables", err))
+
+		return
+	}
+	model.JobVariables = jsontypes.NewNormalizedValue(string(jobVariablesByteSlice))
 
 	resp.Diagnostics.Append(resp.State.Set(ctx, &model)...)
 	if resp.Diagnostics.HasError() {

diff --git a/internal/provider/resources/deployment_test.go b/internal/provider/resources/deployment_test.go
@@ -23,6 +23,7 @@ type deploymentConfig struct {
 	Description            string
 	EnforceParameterSchema bool
 	Entrypoint             string
+	JobVariables           string
 	ManifestPath           string
 	Parameters             string
 	Path                   string
@@ -54,6 +55,9 @@ resource "prefect_deployment" "{{.DeploymentName}}" {
 	enforce_parameter_schema = {{.EnforceParameterSchema}}
 	entrypoint = "{{.Entrypoint}}"
 	flow_id = prefect_flow.{{.FlowName}}.id
+	job_variables = jsonencode(
+		{{.JobVariables}}
+	)
 	manifest_path = "{{.ManifestPath}}"
 	parameters = jsonencode({
 		"some-parameter": "{{.Parameters}}"
@@ -87,6 +91,7 @@ func TestAccResource_deployment(t *testing.T) {
 		Description:            "My deployment description",
 		EnforceParameterSchema: false,
 		Entrypoint:             "hello_world.py:hello_world",
+		JobVariables:           `{"env":{"some-key":"some-value"}}`,
 		ManifestPath:           "some-manifest-path",
 		Parameters:             "some-value1",
 		Path:                   "some-path",
@@ -110,6 +115,7 @@ func TestAccResource_deployment(t *testing.T) {
 		// Configure new values to test the update.
 		Description:   "My deployment description v2",
 		Entrypoint:    "hello_world.py:hello_world2",
+		JobVariables:  `{"env":{"some-key":"some-value2"}}`,
 		ManifestPath:  "some-manifest-path2",
 		Parameters:    "some-value2",
 		Path:          "some-path2",
@@ -150,6 +156,7 @@ func TestAccResource_deployment(t *testing.T) {
 					resource.TestCheckResourceAttr(cfgCreate.DeploymentResourceName, "description", cfgCreate.Description),
 					resource.TestCheckResourceAttr(cfgCreate.DeploymentResourceName, "enforce_parameter_schema", strconv.FormatBool(cfgCreate.EnforceParameterSchema)),
 					resource.TestCheckResourceAttr(cfgCreate.DeploymentResourceName, "entrypoint", cfgCreate.Entrypoint),
+					resource.TestCheckResourceAttr(cfgCreate.DeploymentResourceName, "job_variables", cfgCreate.JobVariables),
 					resource.TestCheckResourceAttr(cfgCreate.DeploymentResourceName, "manifest_path", cfgCreate.ManifestPath),
 					resource.TestCheckResourceAttr(cfgCreate.DeploymentResourceName, "parameters", `{"some-parameter":"some-value1"}`),
 					resource.TestCheckResourceAttr(cfgCreate.DeploymentResourceName, "path", cfgCreate.Path),
@@ -175,6 +182,7 @@ func TestAccResource_deployment(t *testing.T) {
 					resource.TestCheckResourceAttr(cfgUpdate.DeploymentResourceName, "description", cfgUpdate.Description),
 					resource.TestCheckResourceAttr(cfgUpdate.DeploymentResourceName, "enforce_parameter_schema", strconv.FormatBool(cfgUpdate.EnforceParameterSchema)),
 					resource.TestCheckResourceAttr(cfgUpdate.DeploymentResourceName, "entrypoint", cfgUpdate.Entrypoint),
+					resource.TestCheckResourceAttr(cfgCreate.DeploymentResourceName, "job_variables", cfgUpdate.JobVariables),
 					resource.TestCheckResourceAttr(cfgUpdate.DeploymentResourceName, "manifest_path", cfgUpdate.ManifestPath),
 					resource.TestCheckResourceAttr(cfgCreate.DeploymentResourceName, "parameters", `{"some-parameter":"some-value2"}`),
 					resource.TestCheckResourceAttr(cfgUpdate.DeploymentResourceName, "path", cfgUpdate.Path),