mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat(coderd): use task data model when creating a new task (#20275)
Updates coder/internal#976
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user