diff --git a/cli/exp_task.go b/cli/exp_task.go index 81316d155000d..860f7b954f47f 100644 --- a/cli/exp_task.go +++ b/cli/exp_task.go @@ -14,6 +14,7 @@ func (r *RootCmd) tasksCommand() *serpent.Command { }, Children: []*serpent.Command{ r.taskList(), + r.taskCreate(), }, } return cmd diff --git a/cli/exp_taskcreate.go b/cli/exp_taskcreate.go new file mode 100644 index 0000000000000..b23da632a12c2 --- /dev/null +++ b/cli/exp_taskcreate.go @@ -0,0 +1,127 @@ +package cli + +import ( + "fmt" + "strings" + "time" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/cli/cliui" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/serpent" +) + +func (r *RootCmd) taskCreate() *serpent.Command { + var ( + orgContext = NewOrganizationContext() + client = new(codersdk.Client) + + templateName string + templateVersionName string + presetName string + taskInput string + ) + + return &serpent.Command{ + Use: "create [template]", + Short: "Create an experimental task", + Middleware: serpent.Chain( + serpent.RequireRangeArgs(0, 1), + r.InitClient(client), + ), + Options: serpent.OptionSet{ + { + Flag: "input", + Env: "CODER_TASK_INPUT", + Value: serpent.StringOf(&taskInput), + Required: true, + }, + { + Env: "CODER_TASK_TEMPLATE_NAME", + Value: serpent.StringOf(&templateName), + }, + { + Env: "CODER_TASK_TEMPLATE_VERSION", + Value: serpent.StringOf(&templateVersionName), + }, + { + Flag: "preset", + Env: "CODER_TASK_PRESET_NAME", + Value: serpent.StringOf(&presetName), + Default: PresetNone, + }, + }, + Handler: func(inv *serpent.Invocation) error { + var ( + ctx = inv.Context() + expClient = codersdk.NewExperimentalClient(client) + + templateVersionID uuid.UUID + templateVersionPresetID uuid.UUID + ) + + organization, err := orgContext.Selected(inv, client) + if err != nil { + return xerrors.Errorf("get current organization: %w", err) + } + + if len(inv.Args) > 0 { + templateName, templateVersionName, _ = strings.Cut(inv.Args[0], "@") + } + + if templateName == "" { + return xerrors.Errorf("template name not provided") + } + + if templateVersionName != "" { + templateVersion, err := client.TemplateVersionByOrganizationAndName(ctx, organization.ID, templateName, templateVersionName) + if err != nil { + return xerrors.Errorf("get template version: %w", err) + } + + templateVersionID = templateVersion.ID + } else { + template, err := client.TemplateByName(ctx, organization.ID, templateName) + if err != nil { + return xerrors.Errorf("get template: %w", err) + } + + templateVersionID = template.ActiveVersionID + } + + if presetName != PresetNone { + templatePresets, err := client.TemplateVersionPresets(ctx, templateVersionID) + if err != nil { + return xerrors.Errorf("get template presets: %w", err) + } + + preset, err := resolvePreset(templatePresets, presetName) + if err != nil { + return xerrors.Errorf("resolve preset: %w", err) + } + + templateVersionPresetID = preset.ID + } + + workspace, err := expClient.CreateTask(ctx, codersdk.Me, codersdk.CreateTaskRequest{ + TemplateVersionID: templateVersionID, + TemplateVersionPresetID: templateVersionPresetID, + Prompt: taskInput, + }) + if err != nil { + return xerrors.Errorf("create task: %w", err) + } + + _, _ = fmt.Fprintf( + inv.Stdout, + "The task %s has been created at %s!\n", + cliui.Keyword(workspace.Name), + cliui.Timestamp(time.Now()), + ) + + return nil + }, + } +} diff --git a/cli/exp_taskcreate_test.go b/cli/exp_taskcreate_test.go new file mode 100644 index 0000000000000..7a4a4bfb5a43e --- /dev/null +++ b/cli/exp_taskcreate_test.go @@ -0,0 +1,227 @@ +package cli_test + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + "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/cli/cliui" + "github.com/coder/coder/v2/coderd/httpapi" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" + "github.com/coder/serpent" +) + +func TestTaskCreate(t *testing.T) { + t.Parallel() + + var ( + organizationID = uuid.New() + templateID = uuid.New() + templateVersionID = uuid.New() + templateVersionPresetID = uuid.New() + ) + + templateAndVersionFoundHandler := func(t *testing.T, ctx context.Context, templateName, templateVersionName, presetName, prompt string) http.HandlerFunc { + t.Helper() + + return func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/v2/users/me/organizations": + httpapi.Write(ctx, w, http.StatusOK, []codersdk.Organization{ + {MinimalOrganization: codersdk.MinimalOrganization{ + ID: organizationID, + }}, + }) + case fmt.Sprintf("/api/v2/organizations/%s/templates/my-template/versions/my-template-version", organizationID): + httpapi.Write(ctx, w, http.StatusOK, codersdk.TemplateVersion{ + ID: templateVersionID, + }) + case fmt.Sprintf("/api/v2/organizations/%s/templates/my-template", organizationID): + httpapi.Write(ctx, w, http.StatusOK, codersdk.Template{ + ID: templateID, + ActiveVersionID: templateVersionID, + }) + case fmt.Sprintf("/api/v2/templateversions/%s/presets", templateVersionID): + httpapi.Write(ctx, w, http.StatusOK, []codersdk.Preset{ + { + ID: templateVersionPresetID, + Name: presetName, + }, + }) + case "/api/experimental/tasks/me": + var req codersdk.CreateTaskRequest + if !httpapi.Read(ctx, w, r, &req) { + return + } + + assert.Equal(t, prompt, req.Prompt, "prompt mismatch") + assert.Equal(t, templateVersionID, req.TemplateVersionID, "template version mismatch") + + if presetName == "" { + assert.Equal(t, uuid.Nil, req.TemplateVersionPresetID, "expected no template preset id") + } else { + assert.Equal(t, templateVersionPresetID, req.TemplateVersionPresetID, "template version preset id mismatch") + } + + httpapi.Write(ctx, w, http.StatusCreated, codersdk.Workspace{ + Name: "task-wild-goldfish-27", + }) + default: + t.Errorf("unexpected path: %s", r.URL.Path) + } + } + } + + tests := []struct { + args []string + env []string + expectError string + expectOutput string + handler func(t *testing.T, ctx context.Context) http.HandlerFunc + }{ + { + args: []string{"my-template@my-template-version", "--input", "my custom prompt"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "my-template-version", "", "my custom prompt") + }, + }, + { + args: []string{"my-template", "--input", "my custom prompt"}, + env: []string{"CODER_TASK_TEMPLATE_VERSION=my-template-version"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "my-template-version", "", "my custom prompt") + }, + }, + { + args: []string{"--input", "my custom prompt"}, + env: []string{"CODER_TASK_TEMPLATE_NAME=my-template", "CODER_TASK_TEMPLATE_VERSION=my-template-version"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "my-template-version", "", "my custom prompt") + }, + }, + { + env: []string{"CODER_TASK_TEMPLATE_NAME=my-template", "CODER_TASK_TEMPLATE_VERSION=my-template-version", "CODER_TASK_INPUT=my custom prompt"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "my-template-version", "", "my custom prompt") + }, + }, + { + args: []string{"my-template", "--input", "my custom prompt"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "", "", "my custom prompt") + }, + }, + { + args: []string{"my-template", "--input", "my custom prompt", "--preset", "my-preset"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "", "my-preset", "my custom prompt") + }, + }, + { + args: []string{"my-template", "--input", "my custom prompt"}, + env: []string{"CODER_TASK_PRESET_NAME=my-preset"}, + expectOutput: fmt.Sprintf("The task %s has been created", cliui.Keyword("task-wild-goldfish-27")), + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "", "my-preset", "my custom prompt") + }, + }, + { + args: []string{"my-template", "--input", "my custom prompt", "--preset", "not-real-preset"}, + expectError: `preset "not-real-preset" not found`, + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return templateAndVersionFoundHandler(t, ctx, "my-template", "", "my-preset", "my custom prompt") + }, + }, + { + args: []string{"my-template@not-real-template-version", "--input", "my custom prompt"}, + expectError: httpapi.ResourceNotFoundResponse.Message, + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/v2/users/me/organizations": + httpapi.Write(ctx, w, http.StatusOK, []codersdk.Organization{ + {MinimalOrganization: codersdk.MinimalOrganization{ + ID: organizationID, + }}, + }) + case fmt.Sprintf("/api/v2/organizations/%s/templates/my-template/versions/not-real-template-version", organizationID): + httpapi.ResourceNotFound(w) + default: + t.Errorf("unexpected path: %s", r.URL.Path) + } + } + }, + }, + { + args: []string{"not-real-template", "--input", "my custom prompt"}, + expectError: httpapi.ResourceNotFoundResponse.Message, + handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/v2/users/me/organizations": + httpapi.Write(ctx, w, http.StatusOK, []codersdk.Organization{ + {MinimalOrganization: codersdk.MinimalOrganization{ + ID: organizationID, + }}, + }) + case fmt.Sprintf("/api/v2/organizations/%s/templates/not-real-template", organizationID): + httpapi.ResourceNotFound(w) + default: + t.Errorf("unexpected path: %s", r.URL.Path) + } + } + }, + }, + } + + for _, tt := range tests { + t.Run(strings.Join(tt.args, ","), func(t *testing.T) { + t.Parallel() + + var ( + ctx = testutil.Context(t, testutil.WaitShort) + srv = httptest.NewServer(tt.handler(t, ctx)) + client = new(codersdk.Client) + args = []string{"exp", "task", "create"} + sb strings.Builder + err error + ) + + t.Cleanup(srv.Close) + + client.URL, err = url.Parse(srv.URL) + require.NoError(t, err) + + inv, root := clitest.New(t, append(args, tt.args...)...) + inv.Environ = serpent.ParseEnviron(tt.env, "") + inv.Stdout = &sb + inv.Stderr = &sb + clitest.SetupConfig(t, client, root) + + err = inv.WithContext(ctx).Run() + if tt.expectError == "" { + assert.NoError(t, err) + } else { + assert.ErrorContains(t, err, tt.expectError) + } + + assert.Contains(t, sb.String(), tt.expectOutput) + }) + } +} diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 18c10d6388f37..a8b3c186edd8b 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -32,6 +32,7 @@ import ( "github.com/coder/coder/v2/coderd/provisionerdserver" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/rbac/policy" + "github.com/coder/coder/v2/coderd/util/slice" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/provisionersdk" "github.com/coder/coder/v2/testutil" @@ -6579,3 +6580,75 @@ func TestWorkspaceBuildDeadlineConstraint(t *testing.T) { } } } + +// TestGetLatestWorkspaceBuildsByWorkspaceIDs populates the database with +// workspaces and builds. It then tests that +// GetLatestWorkspaceBuildsByWorkspaceIDs returns the latest build for some +// subset of the workspaces. +func TestGetLatestWorkspaceBuildsByWorkspaceIDs(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + + org := dbgen.Organization(t, db, database.Organization{}) + admin := dbgen.User(t, db, database.User{}) + + tv := dbfake.TemplateVersion(t, db). + Seed(database.TemplateVersion{ + OrganizationID: org.ID, + CreatedBy: admin.ID, + }). + Do() + + users := make([]database.User, 5) + wrks := make([][]database.WorkspaceTable, len(users)) + exp := make(map[uuid.UUID]database.WorkspaceBuild) + for i := range users { + users[i] = dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: users[i].ID, + OrganizationID: org.ID, + }) + + // Each user gets 2 workspaces. + wrks[i] = make([]database.WorkspaceTable, 2) + for wi := range wrks[i] { + wrks[i][wi] = dbgen.Workspace(t, db, database.WorkspaceTable{ + TemplateID: tv.Template.ID, + OwnerID: users[i].ID, + }) + + // Choose a deterministic number of builds per workspace + // No more than 5 builds though, that would be excessive. + for j := int32(1); int(j) <= (i+wi)%5; j++ { + wb := dbfake.WorkspaceBuild(t, db, wrks[i][wi]). + Seed(database.WorkspaceBuild{ + WorkspaceID: wrks[i][wi].ID, + BuildNumber: j + 1, + }). + Do() + + exp[wrks[i][wi].ID] = wb.Build // Save the final workspace build + } + } + } + + // Only take half the users. And only take 1 workspace per user for the test. + // The others are just noice. This just queries a subset of workspaces and builds + // to make sure the noise doesn't interfere with the results. + assertWrks := wrks[:len(users)/2] + ctx := testutil.Context(t, testutil.WaitLong) + ids := slice.Convert[[]database.WorkspaceTable, uuid.UUID](assertWrks, func(pair []database.WorkspaceTable) uuid.UUID { + return pair[0].ID + }) + + require.Greater(t, len(ids), 0, "expected some workspace ids for test") + builds, err := db.GetLatestWorkspaceBuildsByWorkspaceIDs(ctx, ids) + require.NoError(t, err) + for _, b := range builds { + expB, ok := exp[b.WorkspaceID] + require.Truef(t, ok, "unexpected workspace build for workspace id %s", b.WorkspaceID) + require.Equalf(t, expB.ID, b.ID, "unexpected workspace build id for workspace id %s", b.WorkspaceID) + require.Equal(t, expB.BuildNumber, b.BuildNumber, "unexpected build number") + } +} diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 2f56b422f350b..014c433cab690 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -18983,20 +18983,15 @@ func (q *sqlQuerier) GetLatestWorkspaceBuildByWorkspaceID(ctx context.Context, w } const getLatestWorkspaceBuildsByWorkspaceIDs = `-- name: GetLatestWorkspaceBuildsByWorkspaceIDs :many -SELECT wb.id, wb.created_at, wb.updated_at, wb.workspace_id, wb.template_version_id, wb.build_number, wb.transition, wb.initiator_id, wb.provisioner_state, wb.job_id, wb.deadline, wb.reason, wb.daily_cost, wb.max_deadline, wb.template_version_preset_id, wb.has_ai_task, wb.ai_task_sidebar_app_id, wb.has_external_agent, wb.initiator_by_avatar_url, wb.initiator_by_username, wb.initiator_by_name -FROM ( - SELECT - workspace_id, MAX(build_number) as max_build_number - FROM - workspace_build_with_user AS workspace_builds - WHERE - workspace_id = ANY($1 :: uuid [ ]) - GROUP BY - workspace_id -) m -JOIN - workspace_build_with_user AS wb -ON m.workspace_id = wb.workspace_id AND m.max_build_number = wb.build_number +SELECT + DISTINCT ON (workspace_id) + id, created_at, updated_at, workspace_id, template_version_id, build_number, transition, initiator_id, provisioner_state, job_id, deadline, reason, daily_cost, max_deadline, template_version_preset_id, has_ai_task, ai_task_sidebar_app_id, has_external_agent, initiator_by_avatar_url, initiator_by_username, initiator_by_name +FROM + workspace_build_with_user AS workspace_builds +WHERE + workspace_id = ANY($1 :: uuid [ ]) +ORDER BY + workspace_id, build_number DESC -- latest first ` func (q *sqlQuerier) GetLatestWorkspaceBuildsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceBuild, error) { diff --git a/coderd/database/queries/workspacebuilds.sql b/coderd/database/queries/workspacebuilds.sql index 6c020f5a97f50..0736c5514b3f7 100644 --- a/coderd/database/queries/workspacebuilds.sql +++ b/coderd/database/queries/workspacebuilds.sql @@ -76,20 +76,16 @@ LIMIT 1; -- name: GetLatestWorkspaceBuildsByWorkspaceIDs :many -SELECT wb.* -FROM ( - SELECT - workspace_id, MAX(build_number) as max_build_number - FROM - workspace_build_with_user AS workspace_builds - WHERE - workspace_id = ANY(@ids :: uuid [ ]) - GROUP BY - workspace_id -) m -JOIN - workspace_build_with_user AS wb -ON m.workspace_id = wb.workspace_id AND m.max_build_number = wb.build_number; +SELECT + DISTINCT ON (workspace_id) + * +FROM + workspace_build_with_user AS workspace_builds +WHERE + workspace_id = ANY(@ids :: uuid [ ]) +ORDER BY + workspace_id, build_number DESC -- latest first +; -- name: InsertWorkspaceBuild :exec INSERT INTO