coder · mafredri · Jan 27, 2025 · Jan 23, 2025 · Jan 24, 2025 · Jan 24, 2025
diff --git a/cli/provisionerjobs.go b/cli/provisionerjobs.go
@@ -4,9 +4,11 @@ import (
 	"fmt"
 	"slices"
 
+	"github.com/google/uuid"
 	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/v2/cli/cliui"
+	"github.com/coder/coder/v2/coderd/util/ptr"
 	"github.com/coder/coder/v2/coderd/util/slice"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/serpent"
@@ -21,6 +23,7 @@ func (r *RootCmd) provisionerJobs() *serpent.Command {
 		},
 		Aliases: []string{"job"},
 		Children: []*serpent.Command{
+			r.provisionerJobsCancel(),
 			r.provisionerJobsList(),
 		},
 	}
@@ -124,3 +127,58 @@ func (r *RootCmd) provisionerJobsList() *serpent.Command {
 
 	return cmd
 }
+
+func (r *RootCmd) provisionerJobsCancel() *serpent.Command {
+	var (
+		client     = new(codersdk.Client)
+		orgContext = NewOrganizationContext()
+	)
+	cmd := &serpent.Command{
+		Use:   "cancel <job_id>",
+		Short: "Cancel a provisioner job",
+		Middleware: serpent.Chain(
+			serpent.RequireNArgs(1),
+			r.InitClient(client),
+		),
+		Handler: func(inv *serpent.Invocation) error {
+			ctx := inv.Context()
+			org, err := orgContext.Selected(inv, client)
+			if err != nil {
+				return xerrors.Errorf("current organization: %w", err)
+			}
+
+			jobID, err := uuid.Parse(inv.Args[0])
+			if err != nil {
+				return xerrors.Errorf("invalid job ID: %w", err)
+			}
+
+			job, err := client.OrganizationProvisionerJob(ctx, org.ID, jobID)
+			if err != nil {
+				return xerrors.Errorf("get provisioner job: %w", err)
+			}
+
+			switch job.Type {
+			case codersdk.ProvisionerJobTypeTemplateVersionDryRun:
+				_, _ = fmt.Fprintf(inv.Stdout, "Canceling template version dry run job %s...\n", job.ID)
+				err = client.CancelTemplateVersionDryRun(ctx, ptr.NilToEmpty(job.Input.TemplateVersionID), job.ID)
+			case codersdk.ProvisionerJobTypeTemplateVersionImport:
+				_, _ = fmt.Fprintf(inv.Stdout, "Canceling template version import job %s...\n", job.ID)
+				err = client.CancelTemplateVersion(ctx, ptr.NilToEmpty(job.Input.TemplateVersionID))
+			case codersdk.ProvisionerJobTypeWorkspaceBuild:
+				_, _ = fmt.Fprintf(inv.Stdout, "Canceling workspace build job %s...\n", job.ID)
+				err = client.CancelWorkspaceBuild(ctx, ptr.NilToEmpty(job.Input.WorkspaceBuildID))
+			}
+			if err != nil {
+				return xerrors.Errorf("cancel provisioner job: %w", err)
+			}
+
+			_, _ = fmt.Fprintln(inv.Stdout, "Job canceled")
+
+			return nil
+		},
+	}
+
+	orgContext.AttachOptions(cmd)
+
+	return cmd
+}
diff --git a/cli/provisionerjobs_test.go b/cli/provisionerjobs_test.go
@@ -0,0 +1,189 @@
+package cli_test
+
+import (
+	"bytes"
+	"database/sql"
+	"encoding/json"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/aws/smithy-go/ptr"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/cli/clitest"
+	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbgen"
+	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"github.com/coder/coder/v2/coderd/rbac"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/testutil"
+)
+
+func TestProvisionerJobs(t *testing.T) {
+	t.Parallel()
+
+	db, ps := dbtestutil.NewDB(t)
+	client, _, coderdAPI := coderdtest.NewWithAPI(t, &coderdtest.Options{
+		IncludeProvisionerDaemon: false,
+		Database:                 db,
+		Pubsub:                   ps,
+	})
+	owner := coderdtest.CreateFirstUser(t, client)
+	templateAdminClient, templateAdmin := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID, rbac.ScopedRoleOrgTemplateAdmin(owner.OrganizationID))
+	memberClient, member := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
+
+	// Create initial resources with a running provisioner.
+	firstProvisioner := coderdtest.NewTaggedProvisionerDaemon(t, coderdAPI, "default-provisioner", map[string]string{"owner": "", "scope": "organization"})
+	t.Cleanup(func() { _ = firstProvisioner.Close() })
+	version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, completeWithAgent())
+	coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
+	template := coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID, func(req *codersdk.CreateTemplateRequest) {
+		req.AllowUserCancelWorkspaceJobs = ptr.Bool(true)
+	})
+
+	// Stop the provisioner so it doesn't grab any more jobs.
+	firstProvisioner.Close()
+
+	t.Run("Cancel", func(t *testing.T) {
+		t.Parallel()
+
+		// Set up test helpers.
+		type jobInput struct {
+			WorkspaceBuildID  string `json:"workspace_build_id,omitempty"`
+			TemplateVersionID string `json:"template_version_id,omitempty"`
+			DryRun            bool   `json:"dry_run,omitempty"`
+		}
+		prepareJob := func(t *testing.T, input jobInput) database.ProvisionerJob {
+			t.Helper()
+
+			inputBytes, err := json.Marshal(input)
+			require.NoError(t, err)
+
+			var typ database.ProvisionerJobType
+			switch {
+			case input.WorkspaceBuildID != "":
+				typ = database.ProvisionerJobTypeWorkspaceBuild
+			case input.TemplateVersionID != "":
+				if input.DryRun {
+					typ = database.ProvisionerJobTypeTemplateVersionDryRun
+				} else {
+					typ = database.ProvisionerJobTypeTemplateVersionImport
+				}
+			default:
+				t.Fatal("invalid input")
+			}
+
+			var (
+				tags = database.StringMap{"owner": "", "scope": "organization", "foo": uuid.New().String()}
+				_    = dbgen.ProvisionerDaemon(t, db, database.ProvisionerDaemon{Tags: tags})
+				job  = dbgen.ProvisionerJob(t, db, coderdAPI.Pubsub, database.ProvisionerJob{
+					InitiatorID: member.ID,
+					Input:       json.RawMessage(inputBytes),
+					Type:        typ,
+					Tags:        tags,
+					StartedAt:   sql.NullTime{Time: coderdAPI.Clock.Now().Add(-time.Minute), Valid: true},
+				})
+			)
+			return job
+		}
+
+		prepareWorkspaceBuildJob := func(t *testing.T) database.ProvisionerJob {
+			t.Helper()
+			var (
+				wbID = uuid.New()
+				job  = prepareJob(t, jobInput{WorkspaceBuildID: wbID.String()})
+				w    = dbgen.Workspace(t, db, database.WorkspaceTable{
+					OrganizationID: owner.OrganizationID,
+					OwnerID:        member.ID,
+					TemplateID:     template.ID,
+				})
+				_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+					ID:                wbID,
+					InitiatorID:       member.ID,
+					WorkspaceID:       w.ID,
+					TemplateVersionID: version.ID,
+					JobID:             job.ID,
+				})
+			)
+			return job
+		}
+
+		prepareTemplateVersionImportJobBuilder := func(t *testing.T, dryRun bool) database.ProvisionerJob {
+			t.Helper()
+			var (
+				tvID = uuid.New()
+				job  = prepareJob(t, jobInput{TemplateVersionID: tvID.String(), DryRun: dryRun})
+				_    = dbgen.TemplateVersion(t, db, database.TemplateVersion{
+					OrganizationID: owner.OrganizationID,
+					CreatedBy:      templateAdmin.ID,
+					ID:             tvID,
+					TemplateID:     uuid.NullUUID{UUID: template.ID, Valid: true},
+					JobID:          job.ID,
+				})
+			)
+			return job
+		}
+		prepareTemplateVersionImportJob := func(t *testing.T) database.ProvisionerJob {
+			return prepareTemplateVersionImportJobBuilder(t, false)
+		}
+		prepareTemplateVersionImportJobDryRun := func(t *testing.T) database.ProvisionerJob {
+			return prepareTemplateVersionImportJobBuilder(t, true)
+		}
+
+		// Run the cancellation test suite.
+		for _, tt := range []struct {
+			role          string
+			client        *codersdk.Client
+			name          string
+			prepare       func(*testing.T) database.ProvisionerJob
+			wantCancelled bool
+		}{
+			{"Owner", client, "WorkspaceBuild", prepareWorkspaceBuildJob, true},
+			{"Owner", client, "TemplateVersionImport", prepareTemplateVersionImportJob, true},
+			{"Owner", client, "TemplateVersionImportDryRun", prepareTemplateVersionImportJobDryRun, true},
+			{"TemplateAdmin", templateAdminClient, "WorkspaceBuild", prepareWorkspaceBuildJob, false},
+			{"TemplateAdmin", templateAdminClient, "TemplateVersionImport", prepareTemplateVersionImportJob, true},
+			{"TemplateAdmin", templateAdminClient, "TemplateVersionImportDryRun", prepareTemplateVersionImportJobDryRun, false},
+			{"Member", memberClient, "WorkspaceBuild", prepareWorkspaceBuildJob, false},
+			{"Member", memberClient, "TemplateVersionImport", prepareTemplateVersionImportJob, false},
+			{"Member", memberClient, "TemplateVersionImportDryRun", prepareTemplateVersionImportJobDryRun, false},
+		} {
+			tt := tt
+			wantMsg := "OK"
+			if !tt.wantCancelled {
+				wantMsg = "FAIL"
+			}
+			t.Run(fmt.Sprintf("%s/%s/%v", tt.role, tt.name, wantMsg), func(t *testing.T) {
+				t.Parallel()
+
+				job := tt.prepare(t)
+				require.False(t, job.CanceledAt.Valid, "job.CanceledAt.Valid")
+
+				inv, root := clitest.New(t, "provisioner", "jobs", "cancel", job.ID.String())
+				clitest.SetupConfig(t, tt.client, root)
+				var buf bytes.Buffer
+				inv.Stdout = &buf
+				err := inv.Run()
+				if tt.wantCancelled {
+					assert.NoError(t, err)
+				} else {
+					assert.Error(t, err)
+				}
+
+				job, err = db.GetProvisionerJobByID(testutil.Context(t, testutil.WaitShort), job.ID)
+				require.NoError(t, err)
+				assert.Equal(t, tt.wantCancelled, job.CanceledAt.Valid, "job.CanceledAt.Valid")
+				assert.Equal(t, tt.wantCancelled, job.CanceledAt.Time.After(job.StartedAt.Time), "job.CanceledAt.Time")
+				if tt.wantCancelled {
+					assert.Contains(t, buf.String(), "Job canceled")
+				} else {
+					assert.NotContains(t, buf.String(), "Job canceled")
+				}
+			})
+		}
+	})
+}
diff --git a/cli/testdata/coder_provisioner_jobs_--help.golden b/cli/testdata/coder_provisioner_jobs_--help.golden
@@ -8,7 +8,8 @@ USAGE:
   Aliases: job
 
 SUBCOMMANDS:
-    list    List provisioner jobs
+    cancel    Cancel a provisioner job
+    list      List provisioner jobs
 
 ———
 Run `coder --help` for a list of global options.
diff --git a/cli/testdata/coder_provisioner_jobs_cancel_--help.golden b/cli/testdata/coder_provisioner_jobs_cancel_--help.golden
@@ -0,0 +1,13 @@
+coder v0.0.0-devel
+
+USAGE:
+  coder provisioner jobs cancel [flags] <job_id>
+
+  Cancel a provisioner job
+
+OPTIONS:
+  -O, --org string, $CODER_ORGANIZATION
+          Select which organization (uuid or name) to use.
+
+———
+Run `coder --help` for a list of global options.
diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -1011,6 +1011,7 @@ func New(options *Options) *API {
 					r.Get("/", api.provisionerDaemons)
 				})
 				r.Route("/provisionerjobs", func(r chi.Router) {
+					r.Get("/{job}", api.provisionerJob)
 					r.Get("/", api.provisionerJobs)
 				})
 			})