Skip to content

Commit 763855e

Browse files
kacpersawkacpersaw
committed
feat(database): add has_external_agent flag to template_versions and workspace_builds
1 parent fb8036a commit 763855e

File tree

19 files changed

+572
-260
lines changed

19 files changed

+572
-260
lines changed

‎coderd/database/dbauthz/dbauthz.go

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

4694-
func (q *querier) UpdateTemplateVersionAITaskByJobID(ctx context.Context, arg database.UpdateTemplateVersionAITaskByJobIDParams) error {
4695-
// An actor is allowed to update the template version AI task flag if they are authorized to update the template.
4696-
tv, err := q.db.GetTemplateVersionByJobID(ctx, arg.JobID)
4694+
func (q *querier) UpdateTemplateVersionByID(ctx context.Context, arg database.UpdateTemplateVersionByIDParams) error {
4695+
// An actor is allowed to update the template version if they are authorized to update the template.
4696+
tv, err := q.db.GetTemplateVersionByID(ctx, arg.ID)
46974697
if err != nil {
46984698
return err
46994699
}
@@ -4710,12 +4710,12 @@ func (q *querier) UpdateTemplateVersionAITaskByJobID(ctx context.Context, arg da
47104710
if err := q.authorizeContext(ctx, policy.ActionUpdate, obj); err != nil {
47114711
return err
47124712
}
4713-
return q.db.UpdateTemplateVersionAITaskByJobID(ctx, arg)
4713+
return q.db.UpdateTemplateVersionByID(ctx, arg)
47144714
}
47154715

4716-
func (q *querier) UpdateTemplateVersionByID(ctx context.Context, arg database.UpdateTemplateVersionByIDParams) error {
4717-
// An actor is allowed to update the template version if they are authorized to update the template.
4718-
tv, err := q.db.GetTemplateVersionByID(ctx, arg.ID)
4716+
func (q *querier) UpdateTemplateVersionDescriptionByJobID(ctx context.Context, arg database.UpdateTemplateVersionDescriptionByJobIDParams) error {
4717+
// An actor is allowed to update the template version description if they are authorized to update the template.
4718+
tv, err := q.db.GetTemplateVersionByJobID(ctx, arg.JobID)
47194719
if err != nil {
47204720
return err
47214721
}
@@ -4732,11 +4732,11 @@ func (q *querier) UpdateTemplateVersionByID(ctx context.Context, arg database.Up
47324732
if err := q.authorizeContext(ctx, policy.ActionUpdate, obj); err != nil {
47334733
return err
47344734
}
4735-
return q.db.UpdateTemplateVersionByID(ctx, arg)
4735+
return q.db.UpdateTemplateVersionDescriptionByJobID(ctx, arg)
47364736
}
47374737

4738-
func (q *querier) UpdateTemplateVersionDescriptionByJobID(ctx context.Context, arg database.UpdateTemplateVersionDescriptionByJobIDParams) error {
4739-
// An actor is allowed to update the template version description if they are authorized to update the template.
4738+
func (q *querier) UpdateTemplateVersionExternalAuthProvidersByJobID(ctx context.Context, arg database.UpdateTemplateVersionExternalAuthProvidersByJobIDParams) error {
4739+
// An actor is allowed to update the template version external auth providers if they are authorized to update the template.
47404740
tv, err := q.db.GetTemplateVersionByJobID(ctx, arg.JobID)
47414741
if err != nil {
47424742
return err
@@ -4754,11 +4754,11 @@ func (q *querier) UpdateTemplateVersionDescriptionByJobID(ctx context.Context, a
47544754
if err := q.authorizeContext(ctx, policy.ActionUpdate, obj); err != nil {
47554755
return err
47564756
}
4757-
return q.db.UpdateTemplateVersionDescriptionByJobID(ctx, arg)
4757+
return q.db.UpdateTemplateVersionExternalAuthProvidersByJobID(ctx, arg)
47584758
}
47594759

4760-
func (q *querier) UpdateTemplateVersionExternalAuthProvidersByJobID(ctx context.Context, arg database.UpdateTemplateVersionExternalAuthProvidersByJobIDParams) error {
4761-
// An actor is allowed to update the template version external auth providers if they are authorized to update the template.
4760+
func (q *querier) UpdateTemplateVersionFlagsByJobID(ctx context.Context, arg database.UpdateTemplateVersionFlagsByJobIDParams) error {
4761+
// An actor is allowed to update the template version ai task and external agent flag if they are authorized to update the template.
47624762
tv, err := q.db.GetTemplateVersionByJobID(ctx, arg.JobID)
47634763
if err != nil {
47644764
return err
@@ -4776,7 +4776,7 @@ func (q *querier) UpdateTemplateVersionExternalAuthProvidersByJobID(ctx context.
47764776
if err := q.authorizeContext(ctx, policy.ActionUpdate, obj); err != nil {
47774777
return err
47784778
}
4779-
return q.db.UpdateTemplateVersionExternalAuthProvidersByJobID(ctx, arg)
4779+
return q.db.UpdateTemplateVersionFlagsByJobID(ctx, arg)
47804780
}
47814781

47824782
func (q *querier) UpdateTemplateWorkspacesLastUsedAt(ctx context.Context, arg database.UpdateTemplateWorkspacesLastUsedAtParams) error {
@@ -5094,7 +5094,15 @@ func (q *querier) UpdateWorkspaceAutostart(ctx context.Context, arg database.Upd
50945094
return update(q.log, q.auth, fetch, q.db.UpdateWorkspaceAutostart)(ctx, arg)
50955095
}
50965096

5097-
func (q *querier) UpdateWorkspaceBuildAITaskByID(ctx context.Context, arg database.UpdateWorkspaceBuildAITaskByIDParams) error {
5097+
// UpdateWorkspaceBuildCostByID is used by the provisioning system to update the cost of a workspace build.
5098+
func (q *querier) UpdateWorkspaceBuildCostByID(ctx context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) error {
5099+
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
5100+
return err
5101+
}
5102+
return q.db.UpdateWorkspaceBuildCostByID(ctx, arg)
5103+
}
5104+
5105+
func (q *querier) UpdateWorkspaceBuildDeadlineByID(ctx context.Context, arg database.UpdateWorkspaceBuildDeadlineByIDParams) error {
50985106
build, err := q.db.GetWorkspaceBuildByID(ctx, arg.ID)
50995107
if err != nil {
51005108
return err
@@ -5109,18 +5117,10 @@ func (q *querier) UpdateWorkspaceBuildAITaskByID(ctx context.Context, arg databa
51095117
if err != nil {
51105118
return err
51115119
}
5112-
return q.db.UpdateWorkspaceBuildAITaskByID(ctx, arg)
5113-
}
5114-
5115-
// UpdateWorkspaceBuildCostByID is used by the provisioning system to update the cost of a workspace build.
5116-
func (q *querier) UpdateWorkspaceBuildCostByID(ctx context.Context, arg database.UpdateWorkspaceBuildCostByIDParams) error {
5117-
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
5118-
return err
5119-
}
5120-
return q.db.UpdateWorkspaceBuildCostByID(ctx, arg)
5120+
return q.db.UpdateWorkspaceBuildDeadlineByID(ctx, arg)
51215121
}
51225122

5123-
func (q *querier) UpdateWorkspaceBuildDeadlineByID(ctx context.Context, arg database.UpdateWorkspaceBuildDeadlineByIDParams) error {
5123+
func (q *querier) UpdateWorkspaceBuildFlagsByID(ctx context.Context, arg database.UpdateWorkspaceBuildFlagsByIDParams) error {
51245124
build, err := q.db.GetWorkspaceBuildByID(ctx, arg.ID)
51255125
if err != nil {
51265126
return err
@@ -5135,7 +5135,7 @@ func (q *querier) UpdateWorkspaceBuildDeadlineByID(ctx context.Context, arg data
51355135
if err != nil {
51365136
return err
51375137
}
5138-
return q.db.UpdateWorkspaceBuildDeadlineByID(ctx, arg)
5138+
return q.db.UpdateWorkspaceBuildFlagsByID(ctx, arg)
51395139
}
51405140

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

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1537,7 +1537,7 @@ func (s *MethodTestSuite) TestTemplate() {
15371537
ID: t1.ID,
15381538
}).Asserts(t1, policy.ActionUpdate)
15391539
}))
1540-
s.Run("UpdateTemplateVersionAITaskByJobID", s.Subtest(func(db database.Store, check *expects) {
1540+
s.Run("UpdateTemplateVersionFlagsByJobID", s.Subtest(func(db database.Store, check *expects) {
15411541
dbtestutil.DisableForeignKeysAndTriggers(s.T(), db)
15421542
o := dbgen.Organization(s.T(), db, database.Organization{})
15431543
u := dbgen.User(s.T(), db, database.User{})
@@ -1550,9 +1550,10 @@ func (s *MethodTestSuite) TestTemplate() {
15501550
JobID: job.ID,
15511551
TemplateID: uuid.NullUUID{UUID: t.ID, Valid: true},
15521552
})
1553-
check.Args(database.UpdateTemplateVersionAITaskByJobIDParams{
1554-
JobID: job.ID,
1555-
HasAITask: sql.NullBool{Bool: true, Valid: true},
1553+
check.Args(database.UpdateTemplateVersionFlagsByJobIDParams{
1554+
JobID: job.ID,
1555+
HasAITask: sql.NullBool{Bool: true, Valid: true},
1556+
HasExternalAgent: sql.NullBool{Bool: true, Valid: true},
15561557
}).Asserts(t, policy.ActionUpdate)
15571558
}))
15581559
s.Run("UpdateTemplateWorkspacesLastUsedAt", s.Subtest(func(db database.Store, check *expects) {
@@ -3161,7 +3162,7 @@ func (s *MethodTestSuite) TestWorkspace() {
31613162
Deadline: b.Deadline,
31623163
}).Asserts(w, policy.ActionUpdate)
31633164
}))
3164-
s.Run("UpdateWorkspaceBuildAITaskByID", s.Subtest(func(db database.Store, check *expects) {
3165+
s.Run("UpdateWorkspaceBuildFlagsByID", s.Subtest(func(db database.Store, check *expects) {
31653166
u := dbgen.User(s.T(), db, database.User{})
31663167
o := dbgen.Organization(s.T(), db, database.Organization{})
31673168
tpl := dbgen.Template(s.T(), db, database.Template{
@@ -3189,10 +3190,12 @@ func (s *MethodTestSuite) TestWorkspace() {
31893190
res := dbgen.WorkspaceResource(s.T(), db, database.WorkspaceResource{JobID: b.JobID})
31903191
agt := dbgen.WorkspaceAgent(s.T(), db, database.WorkspaceAgent{ResourceID: res.ID})
31913192
app := dbgen.WorkspaceApp(s.T(), db, database.WorkspaceApp{AgentID: agt.ID})
3192-
check.Args(database.UpdateWorkspaceBuildAITaskByIDParams{
3193-
HasAITask: sql.NullBool{Bool: true, Valid: true},
3194-
SidebarAppID: uuid.NullUUID{UUID: app.ID, Valid: true},
3195-
ID: b.ID,
3193+
check.Args(database.UpdateWorkspaceBuildFlagsByIDParams{
3194+
ID: b.ID,
3195+
HasAITask: sql.NullBool{Bool: true, Valid: true},
3196+
HasExternalAgent: sql.NullBool{Bool: true, Valid: true},
3197+
SidebarAppID: uuid.NullUUID{UUID: app.ID, Valid: true},
3198+
UpdatedAt: b.UpdatedAt,
31963199
}).Asserts(w, policy.ActionUpdate)
31973200
}))
31983201
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)