From a7d946595c321ba8ba881037ba26e0c0e273fccc Mon Sep 17 00:00:00 2001
From: Mitchell Nielsen <mitchnielsen@users.noreply.github.com>
Date: Tue, 15 Oct 2024 21:40:55 -0500
Subject: [PATCH] Require workspace ID for work pools in Cloud (#282)

* Require workspace ID for work pools in Cloud

In Cloud, workspace_id is required. It's marked as optional to support
OSS.

Related to https://github.com/PrefectHQ/terraform-provider-prefect/issues/281

Related to https://linear.app/prefect/issue/PLA-399/prefect-terraform-provider-errors-with-404-when-trying-to-create-a-new

* Add test case
---
 docs/resources/work_pool.md                   |  2 +-
 internal/client/work_pools.go                 |  6 ++++++
 internal/provider/resources/work_pool.go      |  2 +-
 internal/provider/resources/work_pool_test.go | 13 +++++++++++++
 4 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/docs/resources/work_pool.md b/docs/resources/work_pool.md
index e84c06bd..d6deaf34 100644
--- a/docs/resources/work_pool.md
+++ b/docs/resources/work_pool.md
@@ -60,7 +60,7 @@ resource "prefect_work_pool" "example" {
 - `description` (String) Description of the work pool
 - `paused` (Boolean) Whether this work pool is paused
 - `type` (String) Type of the work pool, eg. kubernetes, ecs, process, etc.
-- `workspace_id` (String) Workspace ID (UUID), defaults to the workspace set in the provider
+- `workspace_id` (String) Workspace ID (UUID), defaults to the workspace set in the provider. In Prefect Cloud, either the `work_pool` resource or the provider's `workspace_id` must be set.
 
 ### Read-Only
 
diff --git a/internal/client/work_pools.go b/internal/client/work_pools.go
index b4d4b1af..204f466a 100644
--- a/internal/client/work_pools.go
+++ b/internal/client/work_pools.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/prefecthq/terraform-provider-prefect/internal/api"
+	"github.com/prefecthq/terraform-provider-prefect/internal/provider/helpers"
 )
 
 var _ = api.WorkPoolsClient(&WorkPoolsClient{})
@@ -28,10 +29,15 @@ func (c *Client) WorkPools(accountID uuid.UUID, workspaceID uuid.UUID) (api.Work
 	if accountID == uuid.Nil {
 		accountID = c.defaultAccountID
 	}
+
 	if workspaceID == uuid.Nil {
 		workspaceID = c.defaultWorkspaceID
 	}
 
+	if helpers.IsCloudEndpoint(c.endpoint) && (accountID == uuid.Nil || workspaceID == uuid.Nil) {
+		return nil, fmt.Errorf("prefect Cloud endpoints require an account_id and workspace_id to be set on either the provider or the resource")
+	}
+
 	return &WorkPoolsClient{
 		hc:          c.hc,
 		apiKey:      c.apiKey,
diff --git a/internal/provider/resources/work_pool.go b/internal/provider/resources/work_pool.go
index 1a2f9bd6..190b80d1 100644
--- a/internal/provider/resources/work_pool.go
+++ b/internal/provider/resources/work_pool.go
@@ -115,7 +115,7 @@ func (r *WorkPoolResource) Schema(_ context.Context, _ resource.SchemaRequest, r
 			},
 			"workspace_id": schema.StringAttribute{
 				CustomType:  customtypes.UUIDType{},
-				Description: "Workspace ID (UUID), defaults to the workspace set in the provider",
+				Description: "Workspace ID (UUID), defaults to the workspace set in the provider. In Prefect Cloud, either the `work_pool` resource or the provider's `workspace_id` must be set.",
 				Optional:    true,
 			},
 			"name": schema.StringAttribute{
diff --git a/internal/provider/resources/work_pool_test.go b/internal/provider/resources/work_pool_test.go
index 8eec9c4f..a2cb3155 100644
--- a/internal/provider/resources/work_pool_test.go
+++ b/internal/provider/resources/work_pool_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"regexp"
 	"strings"
 	"testing"
 
@@ -15,6 +16,13 @@ import (
 	"github.com/prefecthq/terraform-provider-prefect/internal/testutils"
 )
 
+const workPoolWithoutWorkspaceID = `
+resource "prefect_work_pool" "invalid_work_pool" {
+	name = "invalid-work-pool"
+	type = "kubernetes"
+}
+`
+
 func fixtureAccWorkPoolCreate(workspace, workspaceName, name, poolType, baseJobTemplate string, paused bool) string {
 	return fmt.Sprintf(`
 %s
@@ -59,6 +67,11 @@ func TestAccResource_work_pool(t *testing.T) {
 		ProtoV6ProviderFactories: testutils.TestAccProtoV6ProviderFactories,
 		PreCheck:                 func() { testutils.AccTestPreCheck(t) },
 		Steps: []resource.TestStep{
+			{
+				// Check that workspace_id missing causes a failure
+				Config:      workPoolWithoutWorkspaceID,
+				ExpectError: regexp.MustCompile(".*require an account_id and workspace_id to be set.*"),
+			},
 			{
 				// Check creation + existence of the work pool resource
 				Config: fixtureAccWorkPoolCreate(workspace, workspaceName, randomName, poolType, baseJobTemplate, true),