Add support for completing a job

kylecarbs · kylecarbs · commit bc8c0e000b7e · 2022-01-29T18:50:54.000Z
diff --git a/coderd/provisionerd.go b/coderd/provisionerd.go
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"reflect"
 
 	"golang.org/x/xerrors"
 	"storj.io/drpc/drpcmux"
@@ -267,6 +268,194 @@ func (s *provisionerdServer) CancelJob(ctx context.Context, cancelJob *proto.Can
 	return &proto.Empty{}, nil
 }
 
+// CompleteJob is triggered by a provision daemon to mark a provisioner job as completed.
 func (s *provisionerdServer) CompleteJob(ctx context.Context, completed *proto.CompletedJob) (*proto.Empty, error) {
-	return nil, nil
+	jobID, err := uuid.Parse(completed.JobId)
+	if err != nil {
+		return nil, xerrors.Errorf("parse job id: %w", err)
+	}
+	job, err := s.Database.GetProvisionerJobByID(ctx, jobID)
+	if err != nil {
+		return nil, xerrors.Errorf("get job by id: %w", err)
+	}
+	// TODO: Check if the worker ID matches!
+	// If it doesn't, a provisioner daemon could be impersonating another job!
+
+	switch jobType := completed.Type.(type) {
+	case *proto.CompletedJob_ProjectImport_:
+		var input projectImportJob
+		err = json.Unmarshal(job.Input, &input)
+		if err != nil {
+			return nil, xerrors.Errorf("unmarshal job data: %w", err)
+		}
+
+		// Validate that all parameters send from the provisioner daemon
+		// follow the protocol.
+		projectParameters := make([]database.InsertProjectParameterParams, 0, len(jobType.ProjectImport.ParameterSchemas))
+		for _, protoParameter := range jobType.ProjectImport.ParameterSchemas {
+			validationTypeSystem, err := convertValidationTypeSystem(protoParameter.ValidationTypeSystem)
+			if err != nil {
+				return nil, xerrors.Errorf("convert validation type system for %q: %w", protoParameter.Name, err)
+			}
+
+			projectParameter := database.InsertProjectParameterParams{
+				ID:                   uuid.New(),
+				CreatedAt:            database.Now(),
+				ProjectHistoryID:     input.ProjectHistoryID,
+				Name:                 protoParameter.Name,
+				Description:          protoParameter.Description,
+				RedisplayValue:       protoParameter.RedisplayValue,
+				ValidationError:      protoParameter.ValidationError,
+				ValidationCondition:  protoParameter.ValidationCondition,
+				ValidationValueType:  protoParameter.ValidationValueType,
+				ValidationTypeSystem: validationTypeSystem,
+
+				AllowOverrideDestination: protoParameter.AllowOverrideDestination,
+				AllowOverrideSource:      protoParameter.AllowOverrideSource,
+			}
+
+			// It's possible a parameter doesn't define a default source!
+			if protoParameter.DefaultSource != nil {
+				parameterSourceScheme, err := convertParameterSourceScheme(protoParameter.DefaultSource.Scheme)
+				if err != nil {
+					return nil, xerrors.Errorf("convert parameter source scheme: %w", err)
+				}
+				projectParameter.DefaultSourceScheme = parameterSourceScheme
+				projectParameter.DefaultSourceValue = sql.NullString{
+					String: protoParameter.DefaultSource.Value,
+					Valid:  protoParameter.DefaultSource.Value != "",
+				}
+			}
+
+			// It's possible a parameter doesn't define a default destination!
+			if protoParameter.DefaultDestination != nil {
+				parameterDestinationScheme, err := convertParameterDestinationScheme(protoParameter.DefaultDestination.Scheme)
+				if err != nil {
+					return nil, xerrors.Errorf("convert parameter destination scheme: %w", err)
+				}
+				projectParameter.DefaultDestinationScheme = parameterDestinationScheme
+				projectParameter.DefaultDestinationValue = sql.NullString{
+					String: protoParameter.DefaultDestination.Value,
+					Valid:  protoParameter.DefaultDestination.Value != "",
+				}
+			}
+
+			projectParameters = append(projectParameters, projectParameter)
+		}
+
+		// This must occur in a transaction in case of failure.
+		err = s.Database.InTx(func(db database.Store) error {
+			err = db.UpdateProvisionerJobByID(ctx, database.UpdateProvisionerJobByIDParams{
+				ID:        jobID,
+				UpdatedAt: database.Now(),
+				CompletedAt: sql.NullTime{
+					Time:  database.Now(),
+					Valid: true,
+				},
+			})
+			if err != nil {
+				return xerrors.Errorf("update provisioner job: %w", err)
+			}
+			for _, projectParameter := range projectParameters {
+				_, err = db.InsertProjectParameter(ctx, projectParameter)
+				if err != nil {
+					return xerrors.Errorf("insert project parameter %q: %w", projectParameter.Name, err)
+				}
+			}
+			return nil
+		})
+		if err != nil {
+			return nil, xerrors.Errorf("complete job: %w", err)
+		}
+	case *proto.CompletedJob_WorkspaceProvision_:
+		var input workspaceProvisionJob
+		err = json.Unmarshal(job.Input, &input)
+		if err != nil {
+			return nil, xerrors.Errorf("unmarshal job data: %w", err)
+		}
+
+		workspaceHistory, err := s.Database.GetWorkspaceHistoryByID(ctx, input.WorkspaceHistoryID)
+		if err != nil {
+			return nil, xerrors.Errorf("get workspace history: %w", err)
+		}
+
+		err = s.Database.InTx(func(db database.Store) error {
+			err = db.UpdateProvisionerJobByID(ctx, database.UpdateProvisionerJobByIDParams{
+				ID:        jobID,
+				UpdatedAt: database.Now(),
+				CompletedAt: sql.NullTime{
+					Time:  database.Now(),
+					Valid: true,
+				},
+			})
+			if err != nil {
+				return xerrors.Errorf("update provisioner job: %w", err)
+			}
+			err = db.UpdateWorkspaceHistoryByID(ctx, database.UpdateWorkspaceHistoryByIDParams{
+				ID:               workspaceHistory.ID,
+				UpdatedAt:        database.Now(),
+				ProvisionerState: jobType.WorkspaceProvision.State,
+				CompletedAt: sql.NullTime{
+					Time:  database.Now(),
+					Valid: true,
+				},
+			})
+			if err != nil {
+				return xerrors.Errorf("update workspace history: %w", err)
+			}
+			for _, protoResource := range jobType.WorkspaceProvision.Resources {
+				_, err = db.InsertWorkspaceResource(ctx, database.InsertWorkspaceResourceParams{
+					ID:                 uuid.New(),
+					CreatedAt:          database.Now(),
+					WorkspaceHistoryID: input.WorkspaceHistoryID,
+					Type:               protoResource.Type,
+					Name:               protoResource.Name,
+					// TODO: Generate this at the variable validation phase.
+					// Set the value in `default_source`, and disallow overwrite.
+					WorkspaceAgentToken: uuid.NewString(),
+				})
+				if err != nil {
+					return xerrors.Errorf("insert workspace resource %q: %w", protoResource.Name, err)
+				}
+			}
+			return nil
+		})
+		if err != nil {
+			return nil, xerrors.Errorf("complete job: %w", err)
+		}
+	default:
+		return nil, xerrors.Errorf("unknown job type %q; ensure coderd and provisionerd versions match",
+			reflect.TypeOf(completed.Type).String())
+	}
+
+	return &proto.Empty{}, nil
+}
+
+func convertValidationTypeSystem(typeSystem sdkproto.ParameterSchema_TypeSystem) (database.ParameterTypeSystem, error) {
+	switch typeSystem {
+	case sdkproto.ParameterSchema_HCL:
+		return database.ParameterTypeSystemHCL, nil
+	default:
+		return database.ParameterTypeSystem(""), xerrors.Errorf("unknown type system: %d", typeSystem)
+	}
+}
+
+func convertParameterSourceScheme(sourceScheme sdkproto.ParameterSource_Scheme) (database.ParameterSourceScheme, error) {
+	switch sourceScheme {
+	case sdkproto.ParameterSource_DATA:
+		return database.ParameterSourceSchemeData, nil
+	default:
+		return database.ParameterSourceScheme(""), xerrors.Errorf("unknown parameter source scheme: %d", sourceScheme)
+	}
+}
+
+func convertParameterDestinationScheme(destinationScheme sdkproto.ParameterDestination_Scheme) (database.ParameterDestinationScheme, error) {
+	switch destinationScheme {
+	case sdkproto.ParameterDestination_ENVIRONMENT_VARIABLE:
+		return database.ParameterDestinationSchemeEnvironmentVariable, nil
+	case sdkproto.ParameterDestination_PROVISIONER_VARIABLE:
+		return database.ParameterDestinationSchemeProvisionerVariable, nil
+	default:
+		return database.ParameterDestinationScheme(""), xerrors.Errorf("unknown parameter destination scheme: %d", destinationScheme)
+	}
 }
diff --git a/database/databasefake/databasefake.go b/database/databasefake/databasefake.go
@@ -659,7 +659,9 @@ func (q *fakeQuerier) UpdateWorkspaceHistoryByID(_ context.Context, arg database
 			continue
 		}
 		workspaceHistory.UpdatedAt = arg.UpdatedAt
+		workspaceHistory.CompletedAt = arg.CompletedAt
 		workspaceHistory.AfterID = arg.AfterID
+		workspaceHistory.ProvisionerState = arg.ProvisionerState
 		q.workspaceHistory[index] = workspaceHistory
 		return nil
 	}
diff --git a/database/query.sql b/database/query.sql
@@ -543,6 +543,8 @@ UPDATE
   workspace_history
 SET
   updated_at = $2,
-  after_id = $3
+  completed_at = $3,
+  after_id = $4,
+  provisioner_state = $5
 WHERE
   id = $1;
diff --git a/database/query.sql.go b/database/query.sql.go
diff --git a/provisionerd/provisionerd.go b/provisionerd/provisionerd.go
@@ -168,6 +168,7 @@ func (a *API) acquireJob() {
 			a.cancelActiveJob(fmt.Sprintf("remove all from %q directory: %s", a.opts.WorkDirectory, err))
 			return
 		}
+		a.opts.Logger.Debug(context.Background(), "cleaned up work directory")
 	}()
 
 	err = os.MkdirAll(a.opts.WorkDirectory, 0600)
@@ -233,20 +234,67 @@ func (a *API) acquireJob() {
 		a.opts.Logger.Debug(context.Background(), "acquired job is project import",
 			slog.F("project_history_name", jobType.ProjectImport.ProjectHistoryName),
 		)
+
+		response, err := provisioner.Parse(a.closeContext, &provisionersdkproto.Parse_Request{
+			Directory: a.opts.WorkDirectory,
+		})
+		if err != nil {
+			a.cancelActiveJob(fmt.Sprintf("parse source: %s", err))
+			return
+		}
+		_, err = a.client.CompleteJob(a.closeContext, &proto.CompletedJob{
+			JobId: a.activeJob.JobId,
+			Type: &proto.CompletedJob_ProjectImport_{
+				ProjectImport: &proto.CompletedJob_ProjectImport{
+					ParameterSchemas: response.ParameterSchemas,
+				},
+			},
+		})
+		if err != nil {
+			a.cancelActiveJob(fmt.Sprintf("complete job: %s", err))
+			return
+		}
 	case *proto.AcquiredJob_WorkspaceProvision_:
 		a.opts.Logger.Debug(context.Background(), "acquired job is workspace provision",
 			slog.F("workspace_name", jobType.WorkspaceProvision.WorkspaceName),
 			slog.F("state_length", len(jobType.WorkspaceProvision.State)),
 			slog.F("parameters", jobType.WorkspaceProvision.ParameterValues),
 		)
 
+		response, err := provisioner.Provision(a.closeContext, &provisionersdkproto.Provision_Request{
+			Directory:       a.opts.WorkDirectory,
+			ParameterValues: jobType.WorkspaceProvision.ParameterValues,
+			State:           jobType.WorkspaceProvision.State,
+		})
+		if err != nil {
+			a.cancelActiveJob(fmt.Sprintf("provision: %s", err))
+			return
+		}
+		a.opts.Logger.Debug(context.Background(), "provision successful; marking job as complete",
+			slog.F("resource_count", len(response.Resources)),
+			slog.F("resources", response.Resources),
+			slog.F("state_length", len(response.State)),
+		)
+
+		// Complete job may need to be async if we disconnected...
+		// When we reconnect we can flush any of these cached values.
+		_, err = a.client.CompleteJob(a.closeContext, &proto.CompletedJob{
+			JobId: a.activeJob.JobId,
+			Type: &proto.CompletedJob_WorkspaceProvision_{
+				WorkspaceProvision: &proto.CompletedJob_WorkspaceProvision{
+					State:     response.State,
+					Resources: response.Resources,
+				},
+			},
+		})
+		if err != nil {
+			a.cancelActiveJob(fmt.Sprintf("complete job: %s", err))
+			return
+		}
 	default:
 		a.cancelActiveJob(fmt.Sprintf("unknown job type %q; ensure your provisioner daemon is up-to-date", reflect.TypeOf(a.activeJob.Type).String()))
 		return
 	}
-
-	fmt.Printf("Provisioner: %s\n", provisioner)
-	// Work!
 }
 
 func (a *API) cancelActiveJob(errMsg string) {
@@ -257,6 +305,9 @@ func (a *API) cancelActiveJob(errMsg string) {
 		a.activeJob = nil
 		return
 	}
+	if a.activeJob == nil {
+		return
+	}
 
 	a.opts.Logger.Info(context.Background(), "canceling active job",
 		slog.F("error_message", errMsg),
diff --git a/provisionerd/provisionerd_test.go b/provisionerd/provisionerd_test.go
@@ -42,12 +42,13 @@ func TestProvisionerd(t *testing.T) {
 	setupProjectVersion := func(t *testing.T, client *codersdk.Client, user coderd.CreateInitialUserRequest, project coderd.Project) coderd.ProjectHistory {
 		var buffer bytes.Buffer
 		writer := tar.NewWriter(&buffer)
+		content := `resource "null_resource" "hi" {}`
 		err := writer.WriteHeader(&tar.Header{
-			Name: "file",
-			Size: 1 << 10,
+			Name: "main.tf",
+			Size: int64(len(content)),
 		})
 		require.NoError(t, err)
-		_, err = writer.Write(make([]byte, 1<<10))
+		_, err = writer.Write([]byte(content))
 		require.NoError(t, err)
 		projectHistory, err := client.CreateProjectHistory(context.Background(), user.Organization, project.Name, coderd.CreateProjectVersionRequest{
 			StorageMethod: database.ProjectStorageMethodInlineArchive,
@@ -102,6 +103,6 @@ func TestProvisionerd(t *testing.T) {
 			WorkDirectory:   t.TempDir(),
 		})
 		defer api.Close()
-		time.Sleep(time.Millisecond * 500)
+		time.Sleep(time.Millisecond * 1500)
 	})
 }

Original file line number	Diff line number	Diff line change
`@@ -659,7 +659,9 @@ func (q *fakeQuerier) UpdateWorkspaceHistoryByID(_ context.Context, arg database`
`659`	`659`	`continue`
`660`	`660`	`}`
`661`	`661`	`workspaceHistory.UpdatedAt = arg.UpdatedAt`
	`662`	`+ workspaceHistory.CompletedAt = arg.CompletedAt`
`662`	`663`	`workspaceHistory.AfterID = arg.AfterID`
	`664`	`+ workspaceHistory.ProvisionerState = arg.ProvisionerState`
`663`	`665`	`q.workspaceHistory[index] = workspaceHistory`
`664`	`666`	`return nil`
`665`	`667`	`}`