feat(coderd): use task data model when creating a new task (#20275)

Updates coder/internal#976
This commit is contained in:
Mathias Fredriksson
2025-10-23 19:12:09 +03:00
committed by GitHub
parent 0f342ecc04
commit 5c802c2627
27 changed files with 1006 additions and 124 deletions
+24
View File
@@ -5009,6 +5009,30 @@ func (q *querier) UpdateTailnetPeerStatusByCoordinator(ctx context.Context, arg
return q.db.UpdateTailnetPeerStatusByCoordinator(ctx, arg)
}
func (q *querier) UpdateTaskWorkspaceID(ctx context.Context, arg database.UpdateTaskWorkspaceIDParams) (database.TaskTable, error) {
// An actor is allowed to update the workspace ID of a task if they are the
// owner of the task and workspace or have the appropriate permissions.
task, err := q.db.GetTaskByID(ctx, arg.ID)
if err != nil {
return database.TaskTable{}, err
}
if err := q.authorizeContext(ctx, policy.ActionUpdate, task.RBACObject()); err != nil {
return database.TaskTable{}, err
}
ws, err := q.db.GetWorkspaceByID(ctx, arg.WorkspaceID.UUID)
if err != nil {
return database.TaskTable{}, err
}
if err := q.authorizeContext(ctx, policy.ActionUpdate, ws.RBACObject()); err != nil {
return database.TaskTable{}, err
}
return q.db.UpdateTaskWorkspaceID(ctx, arg)
}
func (q *querier) UpdateTemplateACLByID(ctx context.Context, arg database.UpdateTemplateACLByIDParams) error {
fetch := func(ctx context.Context, arg database.UpdateTemplateACLByIDParams) (database.Template, error) {
return q.db.GetTemplateByID(ctx, arg.ID)
+14
View File
@@ -2395,6 +2395,20 @@ func (s *MethodTestSuite) TestTasks() {
check.Args(arg).Asserts(task, policy.ActionUpdate).Returns(database.TaskWorkspaceApp{})
}))
s.Run("UpdateTaskWorkspaceID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
task := testutil.Fake(s.T(), faker, database.Task{})
ws := testutil.Fake(s.T(), faker, database.Workspace{})
arg := database.UpdateTaskWorkspaceIDParams{
ID: task.ID,
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
}
dbm.EXPECT().GetTaskByID(gomock.Any(), task.ID).Return(task, nil).AnyTimes()
dbm.EXPECT().GetWorkspaceByID(gomock.Any(), ws.ID).Return(ws, nil).AnyTimes()
dbm.EXPECT().UpdateTaskWorkspaceID(gomock.Any(), arg).Return(database.TaskTable{}, nil).AnyTimes()
check.Args(arg).Asserts(task, policy.ActionUpdate, ws, policy.ActionUpdate).Returns(database.TaskTable{})
}))
s.Run("GetTaskByWorkspaceID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
task := testutil.Fake(s.T(), faker, database.Task{})
task.WorkspaceID = uuid.NullUUID{UUID: uuid.New(), Valid: true}
@@ -3064,6 +3064,13 @@ func (m queryMetricsStore) UpdateTailnetPeerStatusByCoordinator(ctx context.Cont
return r0
}
func (m queryMetricsStore) UpdateTaskWorkspaceID(ctx context.Context, arg database.UpdateTaskWorkspaceIDParams) (database.TaskTable, error) {
start := time.Now()
r0, r1 := m.s.UpdateTaskWorkspaceID(ctx, arg)
m.queryLatencies.WithLabelValues("UpdateTaskWorkspaceID").Observe(time.Since(start).Seconds())
return r0, r1
}
func (m queryMetricsStore) UpdateTemplateACLByID(ctx context.Context, arg database.UpdateTemplateACLByIDParams) error {
start := time.Now()
err := m.s.UpdateTemplateACLByID(ctx, arg)
+15
View File
@@ -6578,6 +6578,21 @@ func (mr *MockStoreMockRecorder) UpdateTailnetPeerStatusByCoordinator(ctx, arg a
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTailnetPeerStatusByCoordinator", reflect.TypeOf((*MockStore)(nil).UpdateTailnetPeerStatusByCoordinator), ctx, arg)
}
// UpdateTaskWorkspaceID mocks base method.
func (m *MockStore) UpdateTaskWorkspaceID(ctx context.Context, arg database.UpdateTaskWorkspaceIDParams) (database.TaskTable, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "UpdateTaskWorkspaceID", ctx, arg)
ret0, _ := ret[0].(database.TaskTable)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// UpdateTaskWorkspaceID indicates an expected call of UpdateTaskWorkspaceID.
func (mr *MockStoreMockRecorder) UpdateTaskWorkspaceID(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateTaskWorkspaceID", reflect.TypeOf((*MockStore)(nil).UpdateTaskWorkspaceID), ctx, arg)
}
// UpdateTemplateACLByID mocks base method.
func (m *MockStore) UpdateTemplateACLByID(ctx context.Context, arg database.UpdateTemplateACLByIDParams) error {
m.ctrl.T.Helper()
+1
View File
@@ -657,6 +657,7 @@ type sqlcQuerier interface {
UpdateProvisionerJobWithCompleteWithStartedAtByID(ctx context.Context, arg UpdateProvisionerJobWithCompleteWithStartedAtByIDParams) error
UpdateReplica(ctx context.Context, arg UpdateReplicaParams) (Replica, error)
UpdateTailnetPeerStatusByCoordinator(ctx context.Context, arg UpdateTailnetPeerStatusByCoordinatorParams) error
UpdateTaskWorkspaceID(ctx context.Context, arg UpdateTaskWorkspaceIDParams) (TaskTable, error)
UpdateTemplateACLByID(ctx context.Context, arg UpdateTemplateACLByIDParams) error
UpdateTemplateAccessControlByID(ctx context.Context, arg UpdateTemplateAccessControlByIDParams) error
UpdateTemplateActiveVersionByID(ctx context.Context, arg UpdateTemplateActiveVersionByIDParams) error
+184
View File
@@ -7540,3 +7540,187 @@ func TestListTasks(t *testing.T) {
})
}
}
func TestUpdateTaskWorkspaceID(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
// Create organization, users, template, and template version.
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
template := dbgen.Template(t, db, database.Template{
OrganizationID: org.ID,
CreatedBy: user.ID,
})
templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
OrganizationID: org.ID,
TemplateID: uuid.NullUUID{Valid: true, UUID: template.ID},
CreatedBy: user.ID,
})
// Create another template for mismatch test.
template2 := dbgen.Template(t, db, database.Template{
OrganizationID: org.ID,
CreatedBy: user.ID,
})
tests := []struct {
name string
setupTask func(t *testing.T) database.Task
setupWS func(t *testing.T) database.WorkspaceTable
wantErr bool
wantNoRow bool
}{
{
name: "successful update with matching template",
setupTask: func(t *testing.T) database.Task {
return dbgen.Task(t, db, database.TaskTable{
OrganizationID: org.ID,
OwnerID: user.ID,
Name: testutil.GetRandomName(t),
WorkspaceID: uuid.NullUUID{},
TemplateVersionID: templateVersion.ID,
Prompt: "Test prompt",
})
},
setupWS: func(t *testing.T) database.WorkspaceTable {
return dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
OwnerID: user.ID,
TemplateID: template.ID,
})
},
wantErr: false,
wantNoRow: false,
},
{
name: "task already has workspace_id",
setupTask: func(t *testing.T) database.Task {
existingWS := dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
OwnerID: user.ID,
TemplateID: template.ID,
})
return dbgen.Task(t, db, database.TaskTable{
OrganizationID: org.ID,
OwnerID: user.ID,
Name: testutil.GetRandomName(t),
WorkspaceID: uuid.NullUUID{Valid: true, UUID: existingWS.ID},
TemplateVersionID: templateVersion.ID,
Prompt: "Test prompt",
})
},
setupWS: func(t *testing.T) database.WorkspaceTable {
return dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
OwnerID: user.ID,
TemplateID: template.ID,
})
},
wantErr: false,
wantNoRow: true, // No row should be returned because WHERE condition fails.
},
{
name: "template mismatch between task and workspace",
setupTask: func(t *testing.T) database.Task {
return dbgen.Task(t, db, database.TaskTable{
OrganizationID: org.ID,
OwnerID: user.ID,
Name: testutil.GetRandomName(t),
WorkspaceID: uuid.NullUUID{}, // NULL workspace_id
TemplateVersionID: templateVersion.ID,
Prompt: "Test prompt",
})
},
setupWS: func(t *testing.T) database.WorkspaceTable {
return dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
OwnerID: user.ID,
TemplateID: template2.ID, // Different template, JOIN will fail.
})
},
wantErr: false,
wantNoRow: true, // No row should be returned because JOIN condition fails.
},
{
name: "task does not exist",
setupTask: func(t *testing.T) database.Task {
return database.Task{
ID: uuid.New(), // Non-existent task ID.
}
},
setupWS: func(t *testing.T) database.WorkspaceTable {
return dbgen.Workspace(t, db, database.WorkspaceTable{
OrganizationID: org.ID,
OwnerID: user.ID,
TemplateID: template.ID,
})
},
wantErr: false,
wantNoRow: true,
},
{
name: "workspace does not exist",
setupTask: func(t *testing.T) database.Task {
return dbgen.Task(t, db, database.TaskTable{
OrganizationID: org.ID,
OwnerID: user.ID,
Name: testutil.GetRandomName(t),
WorkspaceID: uuid.NullUUID{},
TemplateVersionID: templateVersion.ID,
Prompt: "Test prompt",
})
},
setupWS: func(t *testing.T) database.WorkspaceTable {
return database.WorkspaceTable{
ID: uuid.New(), // Non-existent workspace ID.
}
},
wantErr: false,
wantNoRow: true,
},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
task := tt.setupTask(t)
workspace := tt.setupWS(t)
updatedTask, err := db.UpdateTaskWorkspaceID(ctx, database.UpdateTaskWorkspaceIDParams{
ID: task.ID,
WorkspaceID: uuid.NullUUID{Valid: true, UUID: workspace.ID},
})
if tt.wantErr {
require.Error(t, err)
return
}
if tt.wantNoRow {
require.ErrorIs(t, err, sql.ErrNoRows)
return
}
require.NoError(t, err)
require.Equal(t, task.ID, updatedTask.ID)
require.True(t, updatedTask.WorkspaceID.Valid)
require.Equal(t, workspace.ID, updatedTask.WorkspaceID.UUID)
require.Equal(t, task.OrganizationID, updatedTask.OrganizationID)
require.Equal(t, task.OwnerID, updatedTask.OwnerID)
require.Equal(t, task.Name, updatedTask.Name)
require.Equal(t, task.TemplateVersionID, updatedTask.TemplateVersionID)
// Verify the update persisted by fetching the task again.
fetchedTask, err := db.GetTaskByID(ctx, task.ID)
require.NoError(t, err)
require.True(t, fetchedTask.WorkspaceID.Valid)
require.Equal(t, workspace.ID, fetchedTask.WorkspaceID.UUID)
})
}
}
+43
View File
@@ -12730,6 +12730,49 @@ func (q *sqlQuerier) ListTasks(ctx context.Context, arg ListTasksParams) ([]Task
return items, nil
}
const updateTaskWorkspaceID = `-- name: UpdateTaskWorkspaceID :one
UPDATE
tasks
SET
workspace_id = $2
FROM
workspaces w
JOIN
template_versions tv
ON
tv.template_id = w.template_id
WHERE
tasks.id = $1
AND tasks.workspace_id IS NULL
AND w.id = $2
AND tv.id = tasks.template_version_id
RETURNING
tasks.id, tasks.organization_id, tasks.owner_id, tasks.name, tasks.workspace_id, tasks.template_version_id, tasks.template_parameters, tasks.prompt, tasks.created_at, tasks.deleted_at
`
type UpdateTaskWorkspaceIDParams struct {
ID uuid.UUID `db:"id" json:"id"`
WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"`
}
func (q *sqlQuerier) UpdateTaskWorkspaceID(ctx context.Context, arg UpdateTaskWorkspaceIDParams) (TaskTable, error) {
row := q.db.QueryRowContext(ctx, updateTaskWorkspaceID, arg.ID, arg.WorkspaceID)
var i TaskTable
err := row.Scan(
&i.ID,
&i.OrganizationID,
&i.OwnerID,
&i.Name,
&i.WorkspaceID,
&i.TemplateVersionID,
&i.TemplateParameters,
&i.Prompt,
&i.CreatedAt,
&i.DeletedAt,
)
return i, err
}
const upsertTaskWorkspaceApp = `-- name: UpsertTaskWorkspaceApp :one
INSERT INTO task_workspace_apps
(task_id, workspace_build_number, workspace_agent_id, workspace_app_id)
+19
View File
@@ -5,6 +5,25 @@ VALUES
(gen_random_uuid(), $1, $2, $3, $4, $5, $6, $7, $8)
RETURNING *;
-- name: UpdateTaskWorkspaceID :one
UPDATE
tasks
SET
workspace_id = $2
FROM
workspaces w
JOIN
template_versions tv
ON
tv.template_id = w.template_id
WHERE
tasks.id = $1
AND tasks.workspace_id IS NULL
AND w.id = $2
AND tv.id = tasks.template_version_id
RETURNING
tasks.*;
-- name: UpsertTaskWorkspaceApp :one
INSERT INTO task_workspace_apps
(task_id, workspace_build_number, workspace_agent_id, workspace_app_id)