feat: expose is_prebuild_claim attribute (#396)

dannykopping · web-flow · commit 50ed1c120bb8 · 2025-05-09T18:44:08.000+02:00
* feat: expose is_prebuild_claim attribute

Signed-off-by: Danny Kopping &lt;dannykopping@gmail.com&gt;

* fix: string comparison hardening

Signed-off-by: Danny Kopping &lt;dannykopping@gmail.com&gt;

---------

Signed-off-by: Danny Kopping &lt;dannykopping@gmail.com&gt;
diff --git a/docs/data-sources/workspace.md b/docs/data-sources/workspace.md
@@ -70,6 +70,7 @@ resource "docker_container" "workspace" {
 - `access_url` (String) The access URL of the Coder deployment provisioning this workspace.
 - `id` (String) UUID of the workspace.
 - `is_prebuild` (Boolean) Similar to `prebuild_count`, but a boolean value instead of a count. This is set to true if the workspace is a currently unassigned prebuild. Once the workspace is assigned, this value will be false.
+- `is_prebuild_claim` (Boolean) Indicates whether a prebuilt workspace has just been claimed and this is the first `apply` after that occurrence.
 - `name` (String) Name of the workspace.
 - `prebuild_count` (Number) A computed count, equal to 1 if the workspace is a currently unassigned prebuild. Use this to conditionally act on the status of a prebuild. Actions that do not require user identity can be taken when this value is set to 1. Actions that should only be taken once the workspace has been assigned to a user may be taken when this value is set to 0.
 - `start_count` (Number) A computed count based on `transition` state. If `start`, count will equal 1.
diff --git a/provider/workspace.go b/provider/workspace.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"reflect"
 	"strconv"
+	"strings"
 
 	"github.com/google/uuid"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
@@ -30,10 +31,23 @@ func workspaceDataSource() *schema.Resource {
 			if isPrebuiltWorkspace() {
 				_ = rd.Set("prebuild_count", 1)
 				_ = rd.Set("is_prebuild", true)
+
+				// A claim can only take place AFTER a prebuild, so it's not logically consistent to have this set to any other value.
+				_ = rd.Set("is_prebuild_claim", false)
 			} else {
 				_ = rd.Set("prebuild_count", 0)
 				_ = rd.Set("is_prebuild", false)
 			}
+			if isPrebuiltWorkspaceClaim() {
+				// Indicate that a prebuild claim has taken place.
+				_ = rd.Set("is_prebuild_claim", true)
+
+				// A claim can only take place AFTER a prebuild, so it's not logically consistent to have these set to any other values.
+				_ = rd.Set("prebuild_count", 0)
+				_ = rd.Set("is_prebuild", false)
+			} else {
+				_ = rd.Set("is_prebuild_claim", false)
+			}
 
 			name := helpers.OptionalEnvOrDefault("CODER_WORKSPACE_NAME", "default")
 			rd.Set("name", name)
@@ -116,6 +130,11 @@ func workspaceDataSource() *schema.Resource {
 				Computed:    true,
 				Description: "Similar to `prebuild_count`, but a boolean value instead of a count. This is set to true if the workspace is a currently unassigned prebuild. Once the workspace is assigned, this value will be false.",
 			},
+			"is_prebuild_claim": {
+				Type:        schema.TypeBool,
+				Computed:    true,
+				Description: "Indicates whether a prebuilt workspace has just been claimed and this is the first `apply` after that occurrence.",
+			},
 			"name": {
 				Type:        schema.TypeString,
 				Computed:    true,
@@ -142,7 +161,12 @@ func workspaceDataSource() *schema.Resource {
 
 // isPrebuiltWorkspace returns true if the workspace is an unclaimed prebuilt workspace.
 func isPrebuiltWorkspace() bool {
-	return helpers.OptionalEnv(IsPrebuildEnvironmentVariable()) == "true"
+	return strings.EqualFold(helpers.OptionalEnv(IsPrebuildEnvironmentVariable()), "true")
+}
+
+// isPrebuiltWorkspaceClaim returns true if the workspace is a prebuilt workspace which has just been claimed.
+func isPrebuiltWorkspaceClaim() bool {
+	return strings.EqualFold(helpers.OptionalEnv(IsPrebuildClaimEnvironmentVariable()), "true")
 }
 
 // IsPrebuildEnvironmentVariable returns the name of the environment variable that
@@ -161,3 +185,21 @@ func isPrebuiltWorkspace() bool {
 func IsPrebuildEnvironmentVariable() string {
 	return "CODER_WORKSPACE_IS_PREBUILD"
 }
+
+// IsPrebuildClaimEnvironmentVariable returns the name of the environment variable that
+// indicates whether the workspace is a prebuilt workspace which has just been claimed, and this is the first Terraform
+// apply after that occurrence.
+//
+// Knowing whether the workspace is a claimed prebuilt workspace allows template
+// authors to conditionally execute code in the template based on whether the workspace
+// has been assigned to a user or not. This allows identity specific configuration to
+// be applied only after the workspace is claimed, while the rest of the workspace can
+// be pre-configured.
+//
+// The value of this environment variable should be set to "true" if the workspace is prebuilt
+// and it has just been claimed by a user. Any other values, including "false"
+// and "" will be interpreted to mean that the workspace is not prebuilt, or was
+// prebuilt but has not been claimed by a user.
+func IsPrebuildClaimEnvironmentVariable() string {
+	return "CODER_WORKSPACE_IS_PREBUILD_CLAIM"
+}
diff --git a/provider/workspace_test.go b/provider/workspace_test.go
@@ -4,6 +4,7 @@ import (
 	"regexp"
 	"testing"
 
+	"github.com/coder/terraform-provider-coder/v2/provider"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
 	"github.com/stretchr/testify/assert"
@@ -102,3 +103,122 @@ func TestWorkspace_MissingTemplateName(t *testing.T) {
 		}},
 	})
 }
+
+// TestWorkspace_PrebuildEnv validates that our handling of input environment variables is correct.
+func TestWorkspace_PrebuildEnv(t *testing.T) {
+	cases := []struct {
+		name  string
+		envs  map[string]string
+		check func(state *terraform.State, resource *terraform.ResourceState) error
+	}{
+		{
+			name: "unused",
+			envs: map[string]string{},
+			check: func(state *terraform.State, resource *terraform.ResourceState) error {
+				attribs := resource.Primary.Attributes
+				assert.Equal(t, "false", attribs["is_prebuild"])
+				assert.Equal(t, "0", attribs["prebuild_count"])
+				assert.Equal(t, "false", attribs["is_prebuild_claim"])
+				return nil
+			},
+		},
+		{
+			name: "prebuild=true",
+			envs: map[string]string{
+				provider.IsPrebuildEnvironmentVariable(): "true",
+			},
+			check: func(state *terraform.State, resource *terraform.ResourceState) error {
+				attribs := resource.Primary.Attributes
+				assert.Equal(t, "true", attribs["is_prebuild"])
+				assert.Equal(t, "1", attribs["prebuild_count"])
+				assert.Equal(t, "false", attribs["is_prebuild_claim"])
+				return nil
+			},
+		},
+		{
+			name: "prebuild=false",
+			envs: map[string]string{
+				provider.IsPrebuildEnvironmentVariable(): "false",
+			},
+			check: func(state *terraform.State, resource *terraform.ResourceState) error {
+				attribs := resource.Primary.Attributes
+				assert.Equal(t, "false", attribs["is_prebuild"])
+				assert.Equal(t, "0", attribs["prebuild_count"])
+				assert.Equal(t, "false", attribs["is_prebuild_claim"])
+				return nil
+			},
+		},
+		{
+			name: "prebuild_claim=true",
+			envs: map[string]string{
+				provider.IsPrebuildClaimEnvironmentVariable(): "true",
+			},
+			check: func(state *terraform.State, resource *terraform.ResourceState) error {
+				attribs := resource.Primary.Attributes
+				assert.Equal(t, "false", attribs["is_prebuild"])
+				assert.Equal(t, "0", attribs["prebuild_count"])
+				assert.Equal(t, "true", attribs["is_prebuild_claim"])
+				return nil
+			},
+		},
+		{
+			name: "prebuild_claim=false",
+			envs: map[string]string{
+				provider.IsPrebuildClaimEnvironmentVariable(): "false",
+			},
+			check: func(state *terraform.State, resource *terraform.ResourceState) error {
+				attribs := resource.Primary.Attributes
+				assert.Equal(t, "false", attribs["is_prebuild"])
+				assert.Equal(t, "0", attribs["prebuild_count"])
+				assert.Equal(t, "false", attribs["is_prebuild_claim"])
+				return nil
+			},
+		},
+		{
+			// Should not ever happen, but let's ensure our defensive check is activated. We can't ever have both flags
+			// being true.
+			name: "prebuild=true,prebuild_claim=true",
+			envs: map[string]string{
+				provider.IsPrebuildEnvironmentVariable():      "true",
+				provider.IsPrebuildClaimEnvironmentVariable(): "true",
+			},
+			check: func(state *terraform.State, resource *terraform.ResourceState) error {
+				attribs := resource.Primary.Attributes
+				assert.Equal(t, "false", attribs["is_prebuild"])
+				assert.Equal(t, "0", attribs["prebuild_count"])
+				assert.Equal(t, "true", attribs["is_prebuild_claim"])
+				return nil
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			for k, v := range tc.envs {
+				t.Setenv(k, v)
+			}
+
+			resource.Test(t, resource.TestCase{
+				ProviderFactories: coderFactory(),
+				IsUnitTest:        true,
+				Steps: []resource.TestStep{{
+					Config: `
+provider "coder" {
+	url = "https://example.com:8080"
+}
+data "coder_workspace" "me" {
+}`,
+					Check: func(state *terraform.State) error {
+						// Baseline checks
+						require.Len(t, state.Modules, 1)
+						require.Len(t, state.Modules[0].Resources, 1)
+						resource := state.Modules[0].Resources["data.coder_workspace.me"]
+						require.NotNil(t, resource)
+
+						return tc.check(state, resource)
+					},
+				}},
+			})
+		})
+	}
+}
