Skip to content

Commit 5e4aa79

Browse files
kacpersawkacpersaw
authored
feat(coderd): add has_external_agent flag to template_versions and workspace_builds (#19285)
This pull request introduces support for external workspace management, allowing users to register and manage workspaces that are provisioned and managed outside of the Coder. * Added has_external_agent field to workspace builds and template versions
1 parent e67f0f6 commit 5e4aa79

File tree

19 files changed

+508
-253
lines changed

19 files changed

+508
-253
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: 24 additions & 18 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,38 +2955,44 @@ 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) {
2959-
u := dbgen.User(s.T(), db, database.User{})
2960-
o := dbgen.Organization(s.T(), db, database.Organization{})
2961-
tpl := dbgen.Template(s.T(), db, database.Template{
2958+
s.Run("UpdateWorkspaceBuildFlagsByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
2959+
u := testutil.Fake(s.T(), faker, database.User{})
2960+
o := testutil.Fake(s.T(), faker, database.Organization{})
2961+
tpl := testutil.Fake(s.T(), faker, database.Template{
29622962
OrganizationID: o.ID,
29632963
CreatedBy: u.ID,
29642964
})
2965-
tv := dbgen.TemplateVersion(s.T(), db, database.TemplateVersion{
2965+
tv := testutil.Fake(s.T(), faker, database.TemplateVersion{
29662966
TemplateID: uuid.NullUUID{UUID: tpl.ID, Valid: true},
29672967
OrganizationID: o.ID,
29682968
CreatedBy: u.ID,
29692969
})
2970-
w := dbgen.Workspace(s.T(), db, database.WorkspaceTable{
2970+
w := testutil.Fake(s.T(), faker, database.Workspace{
29712971
TemplateID: tpl.ID,
29722972
OrganizationID: o.ID,
29732973
OwnerID: u.ID,
29742974
})
2975-
j := dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{
2975+
j := testutil.Fake(s.T(), faker, database.ProvisionerJob{
29762976
Type: database.ProvisionerJobTypeWorkspaceBuild,
29772977
})
2978-
b := dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{
2978+
b := testutil.Fake(s.T(), faker, database.WorkspaceBuild{
29792979
JobID: j.ID,
29802980
WorkspaceID: w.ID,
29812981
TemplateVersionID: tv.ID,
29822982
})
2983-
res := dbgen.WorkspaceResource(s.T(), db, database.WorkspaceResource{JobID: b.JobID})
2984-
agt := dbgen.WorkspaceAgent(s.T(), db, database.WorkspaceAgent{ResourceID: res.ID})
2985-
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,
2983+
res := testutil.Fake(s.T(), faker, database.WorkspaceResource{JobID: b.JobID})
2984+
agt := testutil.Fake(s.T(), faker, database.WorkspaceAgent{ResourceID: res.ID})
2985+
app := testutil.Fake(s.T(), faker, database.WorkspaceApp{AgentID: agt.ID})
2986+
2987+
dbm.EXPECT().GetWorkspaceByID(gomock.Any(), w.ID).Return(w, nil).AnyTimes()
2988+
dbm.EXPECT().GetWorkspaceBuildByID(gomock.Any(), b.ID).Return(b, nil).AnyTimes()
2989+
dbm.EXPECT().UpdateWorkspaceBuildFlagsByID(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
2990+
check.Args(database.UpdateWorkspaceBuildFlagsByIDParams{
2991+
ID: b.ID,
2992+
HasAITask: sql.NullBool{Bool: true, Valid: true},
2993+
HasExternalAgent: sql.NullBool{Bool: true, Valid: true},
2994+
SidebarAppID: uuid.NullUUID{UUID: app.ID, Valid: true},
2995+
UpdatedAt: b.UpdatedAt,
29902996
}).Asserts(w, policy.ActionUpdate)
29912997
}))
29922998
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.

0 commit comments

Comments
 (0)