Skip to content

Commit 4b7ed1a

Browse files
kacpersawkacpersaw
committed
feat(database): add has_external_agent flag to template_versions and workspace_builds
1 parent a8c89a1 commit 4b7ed1a

File tree

19 files changed

+570
-259
lines changed

19 files changed

+570
-259
lines changed

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 26 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -4733,9 +4733,9 @@ func (q *querier) UpdateTemplateScheduleByID(ctx context.Context, arg database.U
47334733
return update(q.log, q.auth, fetch, q.db.UpdateTemplateScheduleByID)(ctx, arg)
47344734
}
47354735

4736-
func (q *querier) UpdateTemplateVersionAITaskByJobID(ctx context.Context, arg database.UpdateTemplateVersionAITaskByJobIDParams) error {
4737-
// An actor is allowed to update the template version AI task flag if they are authorized to update the template.
4738-
tv, err := q.db.GetTemplateVersionByJobID(ctx, arg.JobID)
4736+
func (q *querier) UpdateTemplateVersionByID(ctx context.Context, arg database.UpdateTemplateVersionByIDParams) error {
4737+
// An actor is allowed to update the template version if they are authorized to update the template.
4738+
tv, err := q.db.GetTemplateVersionByID(ctx, arg.ID)
47394739
if err != nil {
47404740
return err
47414741
}
@@ -4752,12 +4752,12 @@ func (q *querier) UpdateTemplateVersionAITaskByJobID(ctx context.Context, arg da
47524752
if err := q.authorizeContext(ctx, policy.ActionUpdate, obj); err != nil {
47534753
return err
47544754
}
4755-
return q.db.UpdateTemplateVersionAITaskByJobID(ctx, arg)
4755+
return q.db.UpdateTemplateVersionByID(ctx, arg)
47564756
}
47574757

4758-
func (q *querier) UpdateTemplateVersionByID(ctx context.Context, arg database.UpdateTemplateVersionByIDParams) error {
4759-
// An actor is allowed to update the template version if they are authorized to update the template.
4760-
tv, err := q.db.GetTemplateVersionByID(ctx, arg.ID)
4758+
func (q *querier) UpdateTemplateVersionDescriptionByJobID(ctx context.Context, arg database.UpdateTemplateVersionDescriptionByJobIDParams) error {
4759+
// An actor is allowed to update the template version description if they are authorized to update the template.
4760+
tv, err := q.db.GetTemplateVersionByJobID(ctx, arg.JobID)
47614761
if err != nil {
47624762
return err
47634763
}
@@ -4774,11 +4774,11 @@ func (q *querier) UpdateTemplateVersionByID(ctx context.Context, arg database.Up
47744774
if err := q.authorizeContext(ctx, policy.ActionUpdate, obj); err != nil {
47754775
return err
47764776
}
4777-
return q.db.UpdateTemplateVersionByID(ctx, arg)
4777+
return q.db.UpdateTemplateVersionDescriptionByJobID(ctx, arg)
47784778
}
47794779

4780-
func (q *querier) UpdateTemplateVersionDescriptionByJobID(ctx context.Context, arg database.UpdateTemplateVersionDescriptionByJobIDParams) error {
4781-
// An actor is allowed to update the template version description if they are authorized to update the template.
4780+
func (q *querier) UpdateTemplateVersionExternalAuthProvidersByJobID(ctx context.Context, arg database.UpdateTemplateVersionExternalAuthProvidersByJobIDParams) error {
4781+
// An actor is allowed to update the template version external auth providers if they are authorized to update the template.
47824782
tv, err := q.db.GetTemplateVersionByJobID(ctx, arg.JobID)
47834783
if err != nil {
47844784
return err
@@ -4796,11 +4796,11 @@ func (q *querier) UpdateTemplateVersionDescriptionByJobID(ctx context.Context, a
47964796
if err := q.authorizeContext(ctx, policy.ActionUpdate, obj); err != nil {
47974797
return err
47984798
}
4799-
return q.db.UpdateTemplateVersionDescriptionByJobID(ctx, arg)
4799+
return q.db.UpdateTemplateVersionExternalAuthProvidersByJobID(ctx, arg)
48004800
}
48014801

4802-
func (q *querier) UpdateTemplateVersionExternalAuthProvidersByJobID(ctx context.Context, arg database.UpdateTemplateVersionExternalAuthProvidersByJobIDParams) error {
4803-
// An actor is allowed to update the template version external auth providers if they are authorized to update the template.
4802+
func (q *querier) UpdateTemplateVersionFlagsByJobID(ctx context.Context, arg database.UpdateTemplateVersionFlagsByJobIDParams) error {
4803+
// An actor is allowed to update the template version ai task and external agent flag if they are authorized to update the template.
48044804
tv, err := q.db.GetTemplateVersionByJobID(ctx, arg.JobID)
48054805
if err != nil {
48064806
return err
@@ -4818,7 +4818,7 @@ func (q *querier) UpdateTemplateVersionExternalAuthProvidersByJobID(ctx context.
48184818
if err := q.authorizeContext(ctx, policy.ActionUpdate, obj); err != nil {
48194819
return err
48204820
}
4821-
return q.db.UpdateTemplateVersionExternalAuthProvidersByJobID(ctx, arg)
4821+
return q.db.UpdateTemplateVersionFlagsByJobID(ctx, arg)
48224822
}
48234823

48244824
func (q *querier) UpdateTemplateWorkspacesLastUsedAt(ctx context.Context, arg database.UpdateTemplateWorkspacesLastUsedAtParams) error {
@@ -5143,7 +5143,15 @@ func (q *querier) UpdateWorkspaceAutostart(ctx context.Context, arg database.Upd
51435143
return update(q.log, q.auth, fetch, q.db.UpdateWorkspaceAutostart)(ctx, arg)
51445144
}
51455145

5146-
func (q *querier) UpdateWorkspaceBuildAITaskByID(ctx context.Context, arg database.UpdateWorkspaceBuildAITaskByIDParams) error {
5146+
// UpdateWorkspaceBuildCostByID is used by the provisioning system to update the cost of a workspace build.
5147+
func (q *querier) UpdateWorkspaceBuildCostByID(ctx context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) error {
5148+
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
5149+
return err
5150+
}
5151+
return q.db.UpdateWorkspaceBuildCostByID(ctx, arg)
5152+
}
5153+
5154+
func (q *querier) UpdateWorkspaceBuildDeadlineByID(ctx context.Context, arg database.UpdateWorkspaceBuildDeadlineByIDParams) error {
51475155
build, err := q.db.GetWorkspaceBuildByID(ctx, arg.ID)
51485156
if err != nil {
51495157
return err
@@ -5158,18 +5166,10 @@ func (q *querier) UpdateWorkspaceBuildAITaskByID(ctx context.Context, arg databa
51585166
if err != nil {
51595167
return err
51605168
}
5161-
return q.db.UpdateWorkspaceBuildAITaskByID(ctx, arg)
5162-
}
5163-
5164-
// UpdateWorkspaceBuildCostByID is used by the provisioning system to update the cost of a workspace build.
5165-
func (q *querier) UpdateWorkspaceBuildCostByID(ctx context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) error {
5166-
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
5167-
return err
5168-
}
5169-
return q.db.UpdateWorkspaceBuildCostByID(ctx, arg)
5169+
return q.db.UpdateWorkspaceBuildDeadlineByID(ctx, arg)
51705170
}
51715171

5172-
func (q *querier) UpdateWorkspaceBuildDeadlineByID(ctx context.Context, arg database.UpdateWorkspaceBuildDeadlineByIDParams) error {
5172+
func (q *querier) UpdateWorkspaceBuildFlagsByID(ctx context.Context, arg database.UpdateWorkspaceBuildFlagsByIDParams) error {
51735173
build, err := q.db.GetWorkspaceBuildByID(ctx, arg.ID)
51745174
if err != nil {
51755175
return err
@@ -5184,7 +5184,7 @@ func (q *querier) UpdateWorkspaceBuildDeadlineByID(ctx context.Context, arg data
51845184
if err != nil {
51855185
return err
51865186
}
5187-
return q.db.UpdateWorkspaceBuildDeadlineByID(ctx, arg)
5187+
return q.db.UpdateWorkspaceBuildFlagsByID(ctx, arg)
51885188
}
51895189

51905190
func (q *querier) UpdateWorkspaceBuildProvisionerStateByID(ctx context.Context, arg database.UpdateWorkspaceBuildProvisionerStateByIDParams) error {

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1355,13 +1355,13 @@ func (s *MethodTestSuite) TestTemplate() {
13551355
dbm.EXPECT().UpdateTemplateScheduleByID(gomock.Any(), arg).Return(nil).AnyTimes()
13561356
check.Args(arg).Asserts(t1, policy.ActionUpdate)
13571357
}))
1358-
s.Run("UpdateTemplateVersionAITaskByJobID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
1358+
s.Run("UpdateTemplateVersionFlagsByJobID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
13591359
t := testutil.Fake(s.T(), faker, database.Template{})
13601360
tv := testutil.Fake(s.T(), faker, database.TemplateVersion{TemplateID: uuid.NullUUID{UUID: t.ID, Valid: true}})
1361-
arg := database.UpdateTemplateVersionAITaskByJobIDParams{JobID: tv.JobID, HasAITask: sql.NullBool{Bool: true, Valid: true}}
1361+
arg := database.UpdateTemplateVersionFlagsByJobIDParams{JobID: tv.JobID, HasAITask: sql.NullBool{Bool: true, Valid: true}, HasExternalAgent: sql.NullBool{Bool: true, Valid: true}}
13621362
dbm.EXPECT().GetTemplateVersionByJobID(gomock.Any(), tv.JobID).Return(tv, nil).AnyTimes()
13631363
dbm.EXPECT().GetTemplateByID(gomock.Any(), t.ID).Return(t, nil).AnyTimes()
1364-
dbm.EXPECT().UpdateTemplateVersionAITaskByJobID(gomock.Any(), arg).Return(nil).AnyTimes()
1364+
dbm.EXPECT().UpdateTemplateVersionFlagsByJobID(gomock.Any(), arg).Return(nil).AnyTimes()
13651365
check.Args(arg).Asserts(t, policy.ActionUpdate)
13661366
}))
13671367
s.Run("UpdateTemplateWorkspacesLastUsedAt", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
@@ -2955,7 +2955,7 @@ func (s *MethodTestSuite) TestWorkspace() {
29552955
Deadline: b.Deadline,
29562956
}).Asserts(w, policy.ActionUpdate)
29572957
}))
2958-
s.Run("UpdateWorkspaceBuildAITaskByID", s.Subtest(func(db database.Store, check *expects) {
2958+
s.Run("UpdateWorkspaceBuildFlagsByID", s.Subtest(func(db database.Store, check *expects) {
29592959
u := dbgen.User(s.T(), db, database.User{})
29602960
o := dbgen.Organization(s.T(), db, database.Organization{})
29612961
tpl := dbgen.Template(s.T(), db, database.Template{
@@ -2983,10 +2983,12 @@ func (s *MethodTestSuite) TestWorkspace() {
29832983
res := dbgen.WorkspaceResource(s.T(), db, database.WorkspaceResource{JobID: b.JobID})
29842984
agt := dbgen.WorkspaceAgent(s.T(), db, database.WorkspaceAgent{ResourceID: res.ID})
29852985
app := dbgen.WorkspaceApp(s.T(), db, database.WorkspaceApp{AgentID: agt.ID})
2986-
check.Args(database.UpdateWorkspaceBuildAITaskByIDParams{
2987-
HasAITask: sql.NullBool{Bool: true, Valid: true},
2988-
SidebarAppID: uuid.NullUUID{UUID: app.ID, Valid: true},
2989-
ID: b.ID,
2986+
check.Args(database.UpdateWorkspaceBuildFlagsByIDParams{
2987+
ID: b.ID,
2988+
HasAITask: sql.NullBool{Bool: true, Valid: true},
2989+
HasExternalAgent: sql.NullBool{Bool: true, Valid: true},
2990+
SidebarAppID: uuid.NullUUID{UUID: app.ID, Valid: true},
2991+
UpdatedAt: b.UpdatedAt,
29902992
}).Asserts(w, policy.ActionUpdate)
29912993
}))
29922994
s.Run("SoftDeleteWorkspaceByID", s.Subtest(func(db database.Store, check *expects) {

‎coderd/database/dbgen/dbgen.go

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -437,6 +437,7 @@ func WorkspaceBuild(t testing.TB, db database.Store, orig database.WorkspaceBuil
437437
jobID := takeFirst(orig.JobID, uuid.New())
438438
hasAITask := takeFirst(orig.HasAITask, sql.NullBool{})
439439
sidebarAppID := takeFirst(orig.AITaskSidebarAppID, uuid.NullUUID{})
440+
hasExternalAgent := takeFirst(orig.HasExternalAgent, sql.NullBool{})
440441
var build database.WorkspaceBuild
441442
err := db.InTx(func(db database.Store) error {
442443
err := db.InsertWorkspaceBuild(genCtx, database.InsertWorkspaceBuildParams{
@@ -470,12 +471,13 @@ func WorkspaceBuild(t testing.TB, db database.Store, orig database.WorkspaceBuil
470471
require.NoError(t, err)
471472
}
472473

473-
if hasAITask.Valid {
474-
require.NoError(t, db.UpdateWorkspaceBuildAITaskByID(genCtx, database.UpdateWorkspaceBuildAITaskByIDParams{
475-
HasAITask: hasAITask,
476-
SidebarAppID: sidebarAppID,
477-
UpdatedAt: dbtime.Now(),
478-
ID: buildID,
474+
if hasAITask.Valid || hasExternalAgent.Valid {
475+
require.NoError(t, db.UpdateWorkspaceBuildFlagsByID(genCtx, database.UpdateWorkspaceBuildFlagsByIDParams{
476+
ID: buildID,
477+
HasAITask: hasAITask,
478+
HasExternalAgent: hasExternalAgent,
479+
SidebarAppID: sidebarAppID,
480+
UpdatedAt: dbtime.Now(),
479481
}))
480482
}
481483

@@ -1028,6 +1030,7 @@ func ExternalAuthLink(t testing.TB, db database.Store, orig database.ExternalAut
10281030
func TemplateVersion(t testing.TB, db database.Store, orig database.TemplateVersion) database.TemplateVersion {
10291031
var version database.TemplateVersion
10301032
hasAITask := takeFirst(orig.HasAITask, sql.NullBool{})
1033+
hasExternalAgent := takeFirst(orig.HasExternalAgent, sql.NullBool{})
10311034
jobID := takeFirst(orig.JobID, uuid.New())
10321035
err := db.InTx(func(db database.Store) error {
10331036
versionID := takeFirst(orig.ID, uuid.New())
@@ -1048,11 +1051,12 @@ func TemplateVersion(t testing.TB, db database.Store, orig database.TemplateVers
10481051
return err
10491052
}
10501053

1051-
if hasAITask.Valid {
1052-
require.NoError(t, db.UpdateTemplateVersionAITaskByJobID(genCtx, database.UpdateTemplateVersionAITaskByJobIDParams{
1053-
JobID: jobID,
1054-
HasAITask: hasAITask,
1055-
UpdatedAt: dbtime.Now(),
1054+
if hasAITask.Valid || hasExternalAgent.Valid {
1055+
require.NoError(t, db.UpdateTemplateVersionFlagsByJobID(genCtx, database.UpdateTemplateVersionFlagsByJobIDParams{
1056+
JobID: jobID,
1057+
HasAITask: hasAITask,
1058+
HasExternalAgent: hasExternalAgent,
1059+
UpdatedAt: dbtime.Now(),
10561060
}))
10571061
}
10581062

‎coderd/database/dbmetrics/querymetrics.go

Lines changed: 14 additions & 14 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 & 28 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)