Handle multiple states correctly when provisioning a project

kylecarbs · kylecarbs · commit 4466836a2208 · 2022-02-10T21:18:05.000Z
diff --git a/coderd/parameter/compute.go b/coderd/parameter/compute.go
@@ -11,6 +11,11 @@ import (
 	"github.com/coder/coder/database"
 )
 
+const (
+	CoderUsername            = "coder_username"
+	CoderWorkspaceTransition = "coder_workspace_transition"
+)
+
 // ComputeScope targets identifiers to pull parameters from.
 type ComputeScope struct {
 	ProjectImportJobID uuid.UUID
diff --git a/provisionerd/provisionerd.go b/provisionerd/provisionerd.go
@@ -19,11 +19,23 @@ import (
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
+	"github.com/coder/coder/coderd/parameter"
 	"github.com/coder/coder/provisionerd/proto"
 	sdkproto "github.com/coder/coder/provisionersdk/proto"
 	"github.com/coder/retry"
 )
 
+const (
+	missingParameterErrorText = "missing parameter"
+)
+
+// IsMissingParameterError returns whether the error message provided
+// is a missing parameter error. This can indicate to consumers that
+// they should check parameters.
+func IsMissingParameterError(err string) bool {
+	return strings.Contains(err, missingParameterErrorText)
+}
+
 // Dialer represents the function to create a daemon client connection.
 type Dialer func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error)
 
@@ -352,32 +364,53 @@ func (p *provisionerDaemon) runProjectImport(ctx context.Context, provisioner sd
 		valueByName[parameterValue.Name] = parameterValue
 	}
 	for _, parameterSchema := range parameterSchemas {
+		if parameterSchema.Name == parameter.CoderWorkspaceTransition {
+			// Hardcode the workspace transition variable. We'll
+			// make it do stuff later!
+			continue
+		}
 		_, ok := valueByName[parameterSchema.Name]
 		if !ok {
-			p.cancelActiveJobf("missing parameter: %s", parameterSchema.Name)
+			p.cancelActiveJobf("%s: %s", missingParameterErrorText, parameterSchema.Name)
 			return
 		}
 	}
+	// Checks if the schema has defined a workspace transition variable.
+	// If not, we don't need to check for resources provisioned in a stopped state.
+	hasWorkspaceTransition := false
+	for _, parameterSchema := range parameterSchemas {
+		if parameterSchema.Name != parameter.CoderWorkspaceTransition {
+			continue
+		}
+		hasWorkspaceTransition = true
+		break
+	}
 
-	startResources, err := p.runProjectImportProvision(ctx, provisioner, job, append(updateResponse.ParameterValues, &sdkproto.ParameterValue{
-		DestinationScheme: sdkproto.ParameterDestination_PROVISIONER_VARIABLE,
-		// TODO: Make this a constant higher-up in the stack.
-		Name:  "coder_workspace_transition",
-		Value: "start",
-	}))
-	if err != nil {
-		p.cancelActiveJobf("project import provision for start: %s", err)
-		return
+	startParameters := updateResponse.ParameterValues
+	if hasWorkspaceTransition {
+		startParameters = append(updateResponse.ParameterValues, &sdkproto.ParameterValue{
+			DestinationScheme: sdkproto.ParameterDestination_PROVISIONER_VARIABLE,
+			Name:              parameter.CoderWorkspaceTransition,
+			Value:             "start",
+		})
 	}
-	stopResources, err := p.runProjectImportProvision(ctx, provisioner, job, append(updateResponse.ParameterValues, &sdkproto.ParameterValue{
-		DestinationScheme: sdkproto.ParameterDestination_PROVISIONER_VARIABLE,
-		Name:              "coder_workspace_transition",
-		Value:             "stop",
-	}))
+	startResources, err := p.runProjectImportProvision(ctx, provisioner, job, startParameters)
 	if err != nil {
 		p.cancelActiveJobf("project import provision for start: %s", err)
 		return
 	}
+	stopResources := startResources
+	if hasWorkspaceTransition {
+		stopResources, err = p.runProjectImportProvision(ctx, provisioner, job, append(updateResponse.ParameterValues, &sdkproto.ParameterValue{
+			DestinationScheme: sdkproto.ParameterDestination_PROVISIONER_VARIABLE,
+			Name:              "coder_workspace_transition",
+			Value:             "stop",
+		}))
+		if err != nil {
+			p.cancelActiveJobf("project import provision for start: %s", err)
+			return
+		}
+	}
 
 	_, err = p.client.CompleteJob(ctx, &proto.CompletedJob{
 		JobId: job.JobId,
diff --git a/provisionerd/provisionerd_test.go b/provisionerd/provisionerd_test.go
@@ -21,6 +21,7 @@ import (
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
 
+	"github.com/coder/coder/coderd/parameter"
 	"github.com/coder/coder/provisionerd"
 	"github.com/coder/coder/provisionerd/proto"
 	"github.com/coder/coder/provisionersdk"
@@ -195,6 +196,7 @@ func TestProvisionerd(t *testing.T) {
 			didComplete   atomic.Bool
 			didLog        atomic.Bool
 			didAcquireJob atomic.Bool
+			didDryRun     atomic.Bool
 		)
 		completeChan := make(chan struct{})
 		closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
@@ -247,14 +249,19 @@ func TestProvisionerd(t *testing.T) {
 					err = stream.Send(&sdkproto.Parse_Response{
 						Type: &sdkproto.Parse_Response_Complete{
 							Complete: &sdkproto.Parse_Complete{
-								ParameterSchemas: []*sdkproto.ParameterSchema{},
+								ParameterSchemas: []*sdkproto.ParameterSchema{{
+									Name: parameter.CoderWorkspaceTransition,
+								}},
 							},
 						},
 					})
 					require.NoError(t, err)
 					return nil
 				},
 				provision: func(request *sdkproto.Provision_Request, stream sdkproto.DRPCProvisioner_ProvisionStream) error {
+					if request.DryRun {
+						didDryRun.Store(true)
+					}
 					err := stream.Send(&sdkproto.Provision_Response{
 						Type: &sdkproto.Provision_Response_Log{
 							Log: &sdkproto.Log{
@@ -280,6 +287,7 @@ func TestProvisionerd(t *testing.T) {
 		<-completeChan
 		require.True(t, didLog.Load())
 		require.True(t, didComplete.Load())
+		require.True(t, didDryRun.Load())
 		require.NoError(t, closer.Close())
 	})
 
