Skip to content

Commit 31dcf69

Browse files
kacpersawkacpersaw
committed
feat(database): add has_external_agent flag to template_versions and workspace_builds
1 parent b7e0266 commit 31dcf69

File tree

19 files changed

+498
-245
lines changed

19 files changed

+498
-245
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
@@ -1523,7 +1523,7 @@ func (s *MethodTestSuite) TestTemplate() {
15231523
ID: t1.ID,
15241524
}).Asserts(t1, policy.ActionUpdate)
15251525
}))
1526-
s.Run("UpdateTemplateVersionAITaskByJobID", s.Subtest(func(db database.Store, check *expects) {
1526+
s.Run("UpdateTemplateVersionFlagsByJobID", s.Subtest(func(db database.Store, check *expects) {
15271527
dbtestutil.DisableForeignKeysAndTriggers(s.T(), db)
15281528
o := dbgen.Organization(s.T(), db, database.Organization{})
15291529
u := dbgen.User(s.T(), db, database.User{})
@@ -1536,9 +1536,10 @@ func (s *MethodTestSuite) TestTemplate() {
15361536
JobID: job.ID,
15371537
TemplateID: uuid.NullUUID{UUID: t.ID, Valid: true},
15381538
})
1539-
check.Args(database.UpdateTemplateVersionAITaskByJobIDParams{
1540-
JobID: job.ID,
1541-
HasAITask: sql.NullBool{Bool: true, Valid: true},
1539+
check.Args(database.UpdateTemplateVersionFlagsByJobIDParams{
1540+
JobID: job.ID,
1541+
HasAITask: sql.NullBool{Bool: true, Valid: true},
1542+
HasExternalAgent: sql.NullBool{Bool: true, Valid: true},
15421543
}).Asserts(t, policy.ActionUpdate)
15431544
}))
15441545
s.Run("UpdateTemplateWorkspacesLastUsedAt", s.Subtest(func(db database.Store, check *expects) {
@@ -3229,7 +3230,7 @@ func (s *MethodTestSuite) TestWorkspace() {
32293230
Deadline: b.Deadline,
32303231
}).Asserts(w, policy.ActionUpdate)
32313232
}))
3232-
s.Run("UpdateWorkspaceBuildAITaskByID", s.Subtest(func(db database.Store, check *expects) {
3233+
s.Run("UpdateWorkspaceBuildFlagsByID", s.Subtest(func(db database.Store, check *expects) {
32333234
u := dbgen.User(s.T(), db, database.User{})
32343235
o := dbgen.Organization(s.T(), db, database.Organization{})
32353236
tpl := dbgen.Template(s.T(), db, database.Template{
@@ -3257,10 +3258,12 @@ func (s *MethodTestSuite) TestWorkspace() {
32573258
res := dbgen.WorkspaceResource(s.T(), db, database.WorkspaceResource{JobID: b.JobID})
32583259
agt := dbgen.WorkspaceAgent(s.T(), db, database.WorkspaceAgent{ResourceID: res.ID})
32593260
app := dbgen.WorkspaceApp(s.T(), db, database.WorkspaceApp{AgentID: agt.ID})
3260-
check.Args(database.UpdateWorkspaceBuildAITaskByIDParams{
3261-
HasAITask: sql.NullBool{Bool: true, Valid: true},
3262-
SidebarAppID: uuid.NullUUID{UUID: app.ID, Valid: true},
3263-
ID: b.ID,
3261+
check.Args(database.UpdateWorkspaceBuildFlagsByIDParams{
3262+
ID: b.ID,
3263+
HasAITask: sql.NullBool{Bool: true, Valid: true},
3264+
HasExternalAgent: sql.NullBool{Bool: true, Valid: true},
3265+
SidebarAppID: uuid.NullUUID{UUID: app.ID, Valid: true},
3266+
UpdatedAt: b.UpdatedAt,
32643267
}).Asserts(w, policy.ActionUpdate)
32653268
}))
32663269
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)