feat: add prefect_worker_metadata datasource, for base job configs (#106

)

* feat: add prefect_worker_metadata datasource, for base job configs

* Generate Terraform Docs

* Generate Terraform Docs

* remove count check bc of parallel tests

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

docs/data-sources/work_pools.md

@@ -35,13 +35,11 @@ data "prefect_work_pools" "all_pools" {}
 
 Optional:
 
-- `account_id` (String) Account UUID, defaults to the account set in the provider
 - `concurrency_limit` (Number) The concurrency limit applied to this work pool
 - `default_queue_id` (String) The UUID of the default queue associated with this work pool
 - `description` (String) Description of the work pool
 - `id` (String) Work pool UUID
 - `name` (String) Name of the work pool
-- `workspace_id` (String) Workspace UUID, defaults to the workspace set in the provider
 
 Read-Only:
 

docs/data-sources/worker_metadata.md

@@ -0,0 +1,68 @@
+---
+# generated by https://github.com/hashicorp/terraform-plugin-docs
+page_title: "prefect_worker_metadata Data Source - prefect"
+subcategory: ""
+description: |-
+  Data Source for metdata for common Worker types
+---
+
+# prefect_worker_metadata (Data Source)
+
+Data Source for metdata for common Worker types
+
+## Example Usage
+
+```terraform
+# Use the prefect_worker_metadata datasource
+# to fetch a set of default base job configurations
+# to be used with several common worker types.
+data "prefect_worker_metadata" "d" {}
+
+resource "prefect_work_pool" "kubernetes" {
+  name              = "test-k8s-pool"
+  type              = "kubernetes"
+  workspace_id      = data.prefect_workspace.prd.id
+  paused            = false
+  base_job_template = data.prefect_worker_metadata.d.base_job_configs.kubernetes
+}
+
+resource "prefect_work_pool" "ecs" {
+  name              = "test-ecs-pool"
+  type              = "ecs"
+  workspace_id      = data.prefect_workspace.prd.id
+  paused            = false
+  base_job_template = data.prefect_worker_metadata.d.base_job_configs.ecs
+}
+
+resource "prefect_work_pool" "process" {
+  name              = "test-process-pool"
+  type              = "cloud-run:push"
+  workspace_id      = data.prefect_workspace.prd.id
+  paused            = false
+  base_job_template = data.prefect_worker_metadata.d.base_job_configs.cloud_run_push
+}
+```
+
+<!-- schema generated by tfplugindocs -->
+## Schema
+
+### Read-Only
+
+- `base_job_configs` (Attributes) A map of default base job configurations (JSON) for each of the primary worker types (see [below for nested schema](#nestedatt--base_job_configs))
+
+<a id="nestedatt--base_job_configs"></a>
+### Nested Schema for `base_job_configs`
+
+Read-Only:
+
+- `azure_container_instances` (String) Default base job configuration for Azure Container Instances workers
+- `azure_container_instances_push` (String) Default base job configuration for Azure Container Instances Push workers
+- `cloud_run` (String) Default base job configuration for Cloud Run workers
+- `cloud_run_push` (String) Default base job configuration for Cloud Run Push workers
+- `docker` (String) Default base job configuration for Docker workers
+- `ecs` (String) Default base job configuration for ECS workers
+- `ecs_push` (String) Default base job configuration for ECS Push workers
+- `kubernetes` (String) Default base job configuration for Kubernetes workers
+- `prefect_agent` (String) Default base job configuration for Prefect Agent workers
+- `process` (String) Default base job configuration for Process workers
+- `vertex_ai` (String) Default base job configuration for Vertex AI workers

docs/resources/work_pool.md

@@ -14,9 +14,31 @@ Resource representing a Prefect work pool
 
 ```terraform
 resource "prefect_work_pool" "example" {
-  name   = "My Work Pool"
-  type   = "Kubernetes"
-  paused = false
+  name         = "my-work-pool"
+  type         = "kubernetes"
+  paused       = false
+  workspace_id = "my-workspace-id"
+}
+
+# Use a JSON file to load a base job configuration
+resource "prefect_work_pool" "example" {
+  name              = "test-k8s-pool"
+  type              = "kubernetes"
+  workspace_id      = data.prefect_workspace.prd.id
+  paused            = false
+  base_job_template = file("./base-job-template.json")
+}
+
+# Or use the prefect_worker_metadata datasource
+# to load a default base job configuration
+data "prefect_worker_metadata" "d" {}
+
+resource "prefect_work_pool" "example" {
+  name              = "test-k8s-pool"
+  type              = "kubernetes"
+  workspace_id      = data.prefect_workspace.prd.id
+  paused            = false
+  base_job_template = data.prefect_worker_metadata.d.base_job_configs.kubernetes
 }
 ```
 
@@ -32,7 +54,6 @@ resource "prefect_work_pool" "example" {
 - `account_id` (String) Account UUID, defaults to the account set in the provider
 - `base_job_template` (String) The base job template for the work pool, as a JSON string
 - `concurrency_limit` (Number) The concurrency limit applied to this work pool
-- `default_queue_id` (String) The UUID of the default queue associated with this work pool
 - `description` (String) Description of the work pool
 - `paused` (Boolean) Whether this work pool is paused
 - `type` (String) Type of the work pool
@@ -41,6 +62,7 @@ resource "prefect_work_pool" "example" {
 ### Read-Only
 
 - `created` (String) Date and time of the work pool creation in RFC 3339 format
+- `default_queue_id` (String) The UUID of the default queue associated with this work pool
 - `id` (String) Work pool UUID
 - `updated` (String) Date and time that the work pool was last updated in RFC 3339 format
 

examples/data-sources/prefect_worker_metadata/data-source.tf

@@ -0,0 +1,28 @@
+# Use the prefect_worker_metadata datasource
+# to fetch a set of default base job configurations
+# to be used with several common worker types.
+data "prefect_worker_metadata" "d" {}
+
+resource "prefect_work_pool" "kubernetes" {
+  name              = "test-k8s-pool"
+  type              = "kubernetes"
+  workspace_id      = data.prefect_workspace.prd.id
+  paused            = false
+  base_job_template = data.prefect_worker_metadata.d.base_job_configs.kubernetes
+}
+
+resource "prefect_work_pool" "ecs" {
+  name              = "test-ecs-pool"
+  type              = "ecs"
+  workspace_id      = data.prefect_workspace.prd.id
+  paused            = false
+  base_job_template = data.prefect_worker_metadata.d.base_job_configs.ecs
+}
+
+resource "prefect_work_pool" "process" {
+  name              = "test-process-pool"
+  type              = "cloud-run:push"
+  workspace_id      = data.prefect_workspace.prd.id
+  paused            = false
+  base_job_template = data.prefect_worker_metadata.d.base_job_configs.cloud_run_push
+}

examples/resources/prefect_work_pool/resource.tf

@@ -1,5 +1,27 @@
 resource "prefect_work_pool" "example" {
-  name   = "My Work Pool"
-  type   = "Kubernetes"
-  paused = false
+  name         = "my-work-pool"
+  type         = "kubernetes"
+  paused       = false
+  workspace_id = "my-workspace-id"
+}
+
+# Use a JSON file to load a base job configuration
+resource "prefect_work_pool" "example" {
+  name              = "test-k8s-pool"
+  type              = "kubernetes"
+  workspace_id      = data.prefect_workspace.prd.id
+  paused            = false
+  base_job_template = file("./base-job-template.json")
+}
+
+# Or use the prefect_worker_metadata datasource
+# to load a default base job configuration
+data "prefect_worker_metadata" "d" {}
+
+resource "prefect_work_pool" "example" {
+  name              = "test-k8s-pool"
+  type              = "kubernetes"
+  workspace_id      = data.prefect_workspace.prd.id
+  paused            = false
+  base_job_template = data.prefect_worker_metadata.d.base_job_configs.kubernetes
 }

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/google/uuid v1.3.1
 	github.com/hashicorp/terraform-plugin-docs v0.16.0
 	github.com/hashicorp/terraform-plugin-framework v1.4.2
+	github.com/hashicorp/terraform-plugin-framework-jsontypes v0.1.0
 	github.com/hashicorp/terraform-plugin-framework-validators v0.12.0
 	github.com/hashicorp/terraform-plugin-go v0.19.0
 	github.com/hashicorp/terraform-plugin-testing v1.5.1

internal/api/client.go

@@ -7,6 +7,7 @@ type PrefectClient interface {
 	Accounts() (AccountsClient, error)
 	AccountMemberships(accountID uuid.UUID) (AccountMembershipsClient, error)
 	AccountRoles(accountID uuid.UUID) (AccountRolesClient, error)
+	Collections() (CollectionsClient, error)
 	Workspaces(accountID uuid.UUID) (WorkspacesClient, error)
 	WorkspaceAccess(accountID uuid.UUID, workspaceID uuid.UUID) (WorkspaceAccessClient, error)
 	WorkspaceRoles(accountID uuid.UUID) (WorkspaceRolesClient, error)

internal/api/collections.go

@@ -0,0 +1,25 @@
+package api
+
+import (
+	"context"
+	"encoding/json"
+)
+
+type CollectionsClient interface {
+	GetWorkerMetadataViews(ctx context.Context) (WorkerTypeByPackage, error)
+}
+
+// { "prefect": {...}, "prefect-aws": {...} }.
+type WorkerTypeByPackage map[string]MetadataByWorkerType
+
+// { "ecs": {...} }.
+type MetadataByWorkerType map[string]WorkerMetadata
+
+type WorkerMetadata struct {
+	Type                        string          `json:"type"`
+	DocumentationURL            string          `json:"documentation_url"`
+	DisplayName                 string          `json:"display_name"`
+	InstallCommand              string          `json:"install_command"`
+	Description                 string          `json:"description"`
+	DefaultBaseJobConfiguration json.RawMessage `json:"default_base_job_configuration"`
+}

internal/client/collections.go

@@ -0,0 +1,59 @@
+package client
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+
+	"github.com/prefecthq/terraform-provider-prefect/internal/api"
+)
+
+var _ = api.CollectionsClient(&CollectionsClient{})
+
+type CollectionsClient struct {
+	hc          *http.Client
+	apiKey      string
+	routePrefix string
+}
+
+// Collections returns an CollectionsClient.
+//
+//nolint:ireturn // required to support PrefectClient mocking
+func (c *Client) Collections() (api.CollectionsClient, error) {
+	return &CollectionsClient{
+		hc:          c.hc,
+		apiKey:      c.apiKey,
+		routePrefix: fmt.Sprintf("%s/collections", c.endpoint),
+	}, nil
+}
+
+// GetWorkerMetadataViews returns a map of worker metadata views by prefect package name.
+// This endpoint serves base job configurations for the primary worker types.
+func (c *CollectionsClient) GetWorkerMetadataViews(ctx context.Context) (api.WorkerTypeByPackage, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("%s/views/aggregate-worker-metadata", c.routePrefix), http.NoBody)
+	if err != nil {
+		return nil, fmt.Errorf("error creating request: %w", err)
+	}
+	setDefaultHeaders(req, c.apiKey)
+
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("http error: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		errorBody, _ := io.ReadAll(resp.Body)
+
+		return nil, fmt.Errorf("status code %s, error=%s", resp.Status, errorBody)
+	}
+
+	var workerTypeByPackage api.WorkerTypeByPackage
+	if err := json.NewDecoder(resp.Body).Decode(&workerTypeByPackage); err != nil {
+		return nil, fmt.Errorf("failed to decode response: %w", err)
+	}
+
+	return workerTypeByPackage, nil
+}

internal/provider/datasources/work_pool_test.go

@@ -56,7 +56,6 @@ func TestAccDatasource_work_pool(t *testing.T) {
 				// Check that we can query multiple work pools
 				Config: fixtureAccMultipleWorkPools(),
 				Check: resource.ComposeAggregateTestCheckFunc(
-					resource.TestCheckResourceAttr(multipleWorkPoolDatasourceName, "work_pools.#", "1"),
 					resource.TestCheckResourceAttr(multipleWorkPoolDatasourceName, "work_pools.0.name", "evergreen-pool"),
 					resource.TestCheckResourceAttrSet(multipleWorkPoolDatasourceName, "work_pools.0.id"),
 					resource.TestCheckResourceAttrSet(multipleWorkPoolDatasourceName, "work_pools.0.created"),