Skip to content

Commit e1ad75b

Browse files
dannykoppingdannykopping
committed
feat: proto & db changes to persist coder_ai_task state
Signed-off-by: Danny Kopping <dannykopping@gmail.com>
1 parent 6d8f5b1 commit e1ad75b

File tree

16 files changed

+917
-486
lines changed

16 files changed

+917
-486
lines changed

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4489,6 +4489,28 @@ func (q *querier) UpdateTemplateScheduleByID(ctx context.Context, arg database.U
44894489
return update(q.log, q.auth, fetch, q.db.UpdateTemplateScheduleByID)(ctx, arg)
44904490
}
44914491

4492+
func (q *querier) UpdateTemplateVersionAITaskByJobID(ctx context.Context, arg database.UpdateTemplateVersionAITaskByJobIDParams) error {
4493+
// An actor is allowed to update the template version AI task flag if they are authorized to update the template.
4494+
tv, err := q.db.GetTemplateVersionByJobID(ctx, arg.JobID)
4495+
if err != nil {
4496+
return err
4497+
}
4498+
var obj rbac.Objecter
4499+
if !tv.TemplateID.Valid {
4500+
obj = rbac.ResourceTemplate.InOrg(tv.OrganizationID)
4501+
} else {
4502+
tpl, err := q.db.GetTemplateByID(ctx, tv.TemplateID.UUID)
4503+
if err != nil {
4504+
return err
4505+
}
4506+
obj = tpl
4507+
}
4508+
if err := q.authorizeContext(ctx, policy.ActionUpdate, obj); err != nil {
4509+
return err
4510+
}
4511+
return q.db.UpdateTemplateVersionAITaskByJobID(ctx, arg)
4512+
}
4513+
44924514
func (q *querier) UpdateTemplateVersionByID(ctx context.Context, arg database.UpdateTemplateVersionByIDParams) error {
44934515
// An actor is allowed to update the template version if they are authorized to update the template.
44944516
tv, err := q.db.GetTemplateVersionByID(ctx, arg.ID)
@@ -4845,6 +4867,13 @@ func (q *querier) UpdateWorkspaceAutostart(ctx context.Context, arg database.Upd
48454867
return update(q.log, q.auth, fetch, q.db.UpdateWorkspaceAutostart)(ctx, arg)
48464868
}
48474869

4870+
func (q *querier) UpdateWorkspaceBuildAITaskByID(ctx context.Context, arg database.UpdateWorkspaceBuildAITaskByIDParams) error {
4871+
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
4872+
return err
4873+
}
4874+
return q.db.UpdateWorkspaceBuildAITaskByID(ctx, arg)
4875+
}
4876+
48484877
// UpdateWorkspaceBuildCostByID is used by the provisioning system to update the cost of a workspace build.
48494878
func (q *querier) UpdateWorkspaceBuildCostByID(ctx context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) error {
48504879
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {

‎coderd/database/dbgen/dbgen.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -369,7 +369,6 @@ func WorkspaceBuild(t testing.TB, db database.Store, orig database.WorkspaceBuil
369369
UUID: uuid.UUID{},
370370
Valid: false,
371371
}),
372-
HasAITask: orig.HasAITask,
373372
})
374373
if err != nil {
375374
return err
@@ -944,7 +943,6 @@ func TemplateVersion(t testing.TB, db database.Store, orig database.TemplateVers
944943
JobID: takeFirst(orig.JobID, uuid.New()),
945944
CreatedBy: takeFirst(orig.CreatedBy, uuid.New()),
946945
SourceExampleID: takeFirst(orig.SourceExampleID, sql.NullString{}),
947-
HasAITask: orig.HasAITask,
948946
})
949947
if err != nil {
950948
return err

‎coderd/database/dbmem/dbmem.go

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9399,7 +9399,6 @@ func (q *FakeQuerier) InsertTemplateVersion(_ context.Context, arg database.Inse
93999399
JobID: arg.JobID,
94009400
CreatedBy: arg.CreatedBy,
94019401
SourceExampleID: arg.SourceExampleID,
9402-
HasAITask: arg.HasAITask,
94039402
}
94049403
q.templateVersions = append(q.templateVersions, version)
94059404
return nil
@@ -10079,7 +10078,6 @@ func (q *FakeQuerier) InsertWorkspaceBuild(_ context.Context, arg database.Inser
1007910078
MaxDeadline: arg.MaxDeadline,
1008010079
Reason: arg.Reason,
1008110080
TemplateVersionPresetID: arg.TemplateVersionPresetID,
10082-
HasAITask: arg.HasAITask,
1008310081
}
1008410082
q.workspaceBuilds = append(q.workspaceBuilds, workspaceBuild)
1008510083
return nil
@@ -11285,6 +11283,26 @@ func (q *FakeQuerier) UpdateTemplateScheduleByID(_ context.Context, arg database
1128511283
return sql.ErrNoRows
1128611284
}
1128711285

11286+
func (q *FakeQuerier) UpdateTemplateVersionAITaskByJobID(_ context.Context, arg database.UpdateTemplateVersionAITaskByJobIDParams) error {
11287+
if err := validateDatabaseType(arg); err != nil {
11288+
return err
11289+
}
11290+
11291+
q.mutex.Lock()
11292+
defer q.mutex.Unlock()
11293+
11294+
for index, templateVersion := range q.templateVersions {
11295+
if templateVersion.JobID != arg.JobID {
11296+
continue
11297+
}
11298+
templateVersion.HasAITask = arg.HasAITask
11299+
templateVersion.UpdatedAt = arg.UpdatedAt
11300+
q.templateVersions[index] = templateVersion
11301+
return nil
11302+
}
11303+
return sql.ErrNoRows
11304+
}
11305+
1128811306
func (q *FakeQuerier) UpdateTemplateVersionByID(_ context.Context, arg database.UpdateTemplateVersionByIDParams) error {
1128911307
if err := validateDatabaseType(arg); err != nil {
1129011308
return err
@@ -11980,6 +11998,27 @@ func (q *FakeQuerier) UpdateWorkspaceAutostart(_ context.Context, arg database.U
1198011998
return sql.ErrNoRows
1198111999
}
1198212000

12001+
func (q *FakeQuerier) UpdateWorkspaceBuildAITaskByID(_ context.Context, arg database.UpdateWorkspaceBuildAITaskByIDParams) error {
12002+
err := validateDatabaseType(arg)
12003+
if err != nil {
12004+
return err
12005+
}
12006+
12007+
q.mutex.Lock()
12008+
defer q.mutex.Unlock()
12009+
12010+
for index, workspaceBuild := range q.workspaceBuilds {
12011+
if workspaceBuild.ID != arg.ID {
12012+
continue
12013+
}
12014+
workspaceBuild.HasAITask = arg.HasAITask
12015+
workspaceBuild.UpdatedAt = dbtime.Now()
12016+
q.workspaceBuilds[index] = workspaceBuild
12017+
return nil
12018+
}
12019+
return sql.ErrNoRows
12020+
}
12021+
1198312022
func (q *FakeQuerier) UpdateWorkspaceBuildCostByID(_ context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) error {
1198412023
if err := validateDatabaseType(arg); err != nil {
1198512024
return err

‎coderd/database/dbmetrics/querymetrics.go

Lines changed: 14 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go

Lines changed: 28 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier.go

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go

Lines changed: 51 additions & 9 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/templateversions.sql

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -88,11 +88,10 @@ INSERT INTO
8888
readme,
8989
job_id,
9090
created_by,
91-
source_example_id,
92-
has_ai_task
91+
source_example_id
9392
)
9493
VALUES
95-
($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12);
94+
($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11);
9695

9796
-- name: UpdateTemplateVersionByID :exec
9897
UPDATE
@@ -123,6 +122,15 @@ SET
123122
WHERE
124123
job_id = $1;
125124

125+
-- name: UpdateTemplateVersionAITaskByJobID :exec
126+
UPDATE
127+
template_versions
128+
SET
129+
has_ai_task = $2,
130+
updated_at = $3
131+
WHERE
132+
job_id = $1;
133+
126134
-- name: GetPreviousTemplateVersion :one
127135
SELECT
128136
*

‎coderd/database/queries/workspacebuilds.sql

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -121,11 +121,10 @@ INSERT INTO
121121
deadline,
122122
max_deadline,
123123
reason,
124-
template_version_preset_id,
125-
has_ai_task
124+
template_version_preset_id
126125
)
127126
VALUES
128-
($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15);
127+
($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14);
129128

130129
-- name: UpdateWorkspaceBuildCostByID :exec
131130
UPDATE
@@ -152,6 +151,15 @@ SET
152151
updated_at = @updated_at::timestamptz
153152
WHERE id = @id::uuid;
154153

154+
-- name: UpdateWorkspaceBuildAITaskByID :exec
155+
UPDATE
156+
workspace_builds
157+
SET
158+
has_ai_task = @has_ai_task,
159+
ai_tasks_sidebar_app_id = @sidebar_app_id,
160+
updated_at = @updated_at::timestamptz
161+
WHERE id = @id::uuid;
162+
155163
-- name: GetActiveWorkspaceBuildsByTemplateID :many
156164
SELECT wb.*
157165
FROM (

0 commit comments

Comments
 (0)