Skip to content

Commit 81e30c9

Browse files
ethanndicksonethanndickson
committed
chore: add db query to retrieve workspaces & their agents
1 parent ba90bb0 commit 81e30c9

File tree

11 files changed

+390
-0
lines changed

11 files changed

+390
-0
lines changed

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2712,6 +2712,14 @@ func (q *querier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesP
27122712
return q.db.GetAuthorizedWorkspaces(ctx, arg, prep)
27132713
}
27142714

2715+
func (q *querier) GetWorkspacesAndAgents(ctx context.Context) ([]database.GetWorkspacesAndAgentsRow, error) {
2716+
prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceWorkspace.Type)
2717+
if err != nil {
2718+
return nil, xerrors.Errorf("(dev error) prepare sql filter: %w", err)
2719+
}
2720+
return q.db.GetAuthorizedWorkspacesAndAgents(ctx, prep)
2721+
}
2722+
27152723
func (q *querier) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]database.Workspace, error) {
27162724
return q.db.GetWorkspacesEligibleForTransition(ctx, now)
27172725
}
@@ -4149,6 +4157,10 @@ func (q *querier) GetAuthorizedWorkspaces(ctx context.Context, arg database.GetW
41494157
return q.GetWorkspaces(ctx, arg)
41504158
}
41514159

4160+
func (q *querier) GetAuthorizedWorkspacesAndAgents(ctx context.Context, _ rbac.PreparedAuthorized) ([]database.GetWorkspacesAndAgentsRow, error) {
4161+
return q.GetWorkspacesAndAgents(ctx)
4162+
}
4163+
41524164
// GetAuthorizedUsers is not required for dbauthz since GetUsers is already
41534165
// authenticated.
41544166
func (q *querier) GetAuthorizedUsers(ctx context.Context, arg database.GetUsersParams, _ rbac.PreparedAuthorized) ([]database.GetUsersRow, error) {

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1484,6 +1484,14 @@ func (s *MethodTestSuite) TestWorkspace() {
14841484
// No asserts here because SQLFilter.
14851485
check.Args(database.GetWorkspacesParams{}, emptyPreparedAuthorized{}).Asserts()
14861486
}))
1487+
s.Run("GetWorkspacesAndAgents", s.Subtest(func(db database.Store, check *expects) {
1488+
// No asserts here because SQLFilter.
1489+
check.Args().Asserts()
1490+
}))
1491+
s.Run("GetAuthorizedWorkspacesAndAgents", s.Subtest(func(db database.Store, check *expects) {
1492+
// No asserts here because SQLFilter.
1493+
check.Args(emptyPreparedAuthorized{}).Asserts()
1494+
}))
14871495
s.Run("GetLatestWorkspaceBuildByWorkspaceID", s.Subtest(func(db database.Store, check *expects) {
14881496
ws := dbgen.Workspace(s.T(), db, database.Workspace{})
14891497
b := dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{WorkspaceID: ws.ID})

‎coderd/database/dbmem/dbmem.go

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6731,6 +6731,11 @@ func (q *FakeQuerier) GetWorkspaces(ctx context.Context, arg database.GetWorkspa
67316731
return workspaceRows, err
67326732
}
67336733

6734+
func (q *FakeQuerier) GetWorkspacesAndAgents(ctx context.Context) ([]database.GetWorkspacesAndAgentsRow, error) {
6735+
// No auth filter.
6736+
return q.GetAuthorizedWorkspacesAndAgents(ctx, nil)
6737+
}
6738+
67346739
func (q *FakeQuerier) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]database.Workspace, error) {
67356740
q.mutex.RLock()
67366741
defer q.mutex.RUnlock()
@@ -11097,6 +11102,63 @@ func (q *FakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.
1109711102
return q.convertToWorkspaceRowsNoLock(ctx, workspaces, int64(beforePageCount), arg.WithSummary), nil
1109811103
}
1109911104

11105+
func (q *FakeQuerier) GetAuthorizedWorkspacesAndAgents(ctx context.Context, prepared rbac.PreparedAuthorized) ([]database.GetWorkspacesAndAgentsRow, error) {
11106+
q.mutex.RLock()
11107+
defer q.mutex.RUnlock()
11108+
11109+
if prepared != nil {
11110+
// Call this to match the same function calls as the SQL implementation.
11111+
_, err := prepared.CompileToSQL(ctx, rbac.ConfigWithoutACL())
11112+
if err != nil {
11113+
return nil, err
11114+
}
11115+
}
11116+
workspaces := make([]database.Workspace, 0)
11117+
for _, workspace := range q.workspaces {
11118+
if prepared != nil && prepared.Authorize(ctx, workspace.RBACObject()) == nil {
11119+
workspaces = append(workspaces, workspace)
11120+
}
11121+
}
11122+
11123+
out := make([]database.GetWorkspacesAndAgentsRow, 0, len(workspaces))
11124+
for _, w := range workspaces {
11125+
// psql constraints ensure all these exist
11126+
build, err := q.getLatestWorkspaceBuildByWorkspaceIDNoLock(ctx, w.ID)
11127+
if err != nil {
11128+
return nil, xerrors.Errorf("get latest build: %w", err)
11129+
}
11130+
11131+
job, err := q.getProvisionerJobByIDNoLock(ctx, build.JobID)
11132+
if err != nil {
11133+
return nil, xerrors.Errorf("get provisioner job: %w", err)
11134+
}
11135+
11136+
resource, err := q.getWorkspaceResourcesByJobIDNoLock(ctx, job.ID)
11137+
if err != nil || len(resource) == 0 {
11138+
return nil, xerrors.Errorf("get workspace resources: %w", err)
11139+
}
11140+
11141+
agents, err := q.getWorkspaceAgentsByResourceIDsNoLock(ctx, []uuid.UUID{resource[0].ID})
11142+
if err != nil {
11143+
return nil, xerrors.Errorf("get workspace agents: %w", err)
11144+
}
11145+
agentIDs := make([]uuid.UUID, 0, len(agents))
11146+
for _, a := range agents {
11147+
agentIDs = append(agentIDs, a.ID)
11148+
}
11149+
11150+
out = append(out, database.GetWorkspacesAndAgentsRow{
11151+
WorkspaceID: w.ID,
11152+
WorkspaceName: w.Name,
11153+
JobStatus: job.JobStatus,
11154+
Transition: build.Transition,
11155+
AgentIds: agentIDs,
11156+
})
11157+
}
11158+
11159+
return out, nil
11160+
}
11161+
1110011162
func (q *FakeQuerier) GetAuthorizedUsers(ctx context.Context, arg database.GetUsersParams, prepared rbac.PreparedAuthorized) ([]database.GetUsersRow, error) {
1110111163
if err := validateDatabaseType(arg); err != nil {
1110211164
return nil, err

‎coderd/database/dbmetrics/dbmetrics.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: 30 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/modelqueries.go

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,7 @@ func (q *sqlQuerier) GetTemplateGroupRoles(ctx context.Context, id uuid.UUID) ([
221221

222222
type workspaceQuerier interface {
223223
GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspacesParams, prepared rbac.PreparedAuthorized) ([]GetWorkspacesRow, error)
224+
GetAuthorizedWorkspacesAndAgents(ctx context.Context, prepared rbac.PreparedAuthorized) ([]GetWorkspacesAndAgentsRow, error)
224225
}
225226

226227
// GetAuthorizedWorkspaces returns all workspaces that the user is authorized to access.
@@ -312,6 +313,45 @@ func (q *sqlQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspa
312313
return items, nil
313314
}
314315

316+
func (q *sqlQuerier) GetAuthorizedWorkspacesAndAgents(ctx context.Context, prepared rbac.PreparedAuthorized) ([]GetWorkspacesAndAgentsRow, error) {
317+
authorizedFilter, err := prepared.CompileToSQL(ctx, rbac.ConfigWorkspaces())
318+
if err != nil {
319+
return nil, xerrors.Errorf("compile authorized filter: %w", err)
320+
}
321+
filtered, err := insertAuthorizedFilter(getWorkspacesAndAgents, fmt.Sprintf(" WHERE %s", authorizedFilter))
322+
if err != nil {
323+
return nil, xerrors.Errorf("insert authorized filter: %w", err)
324+
}
325+
326+
query := fmt.Sprintf("-- name: GetAuthorizedWorkspaces :many\n%s", filtered)
327+
rows, err := q.db.QueryContext(ctx, query)
328+
if err != nil {
329+
return nil, err
330+
}
331+
defer rows.Close()
332+
var items []GetWorkspacesAndAgentsRow
333+
for rows.Next() {
334+
var i GetWorkspacesAndAgentsRow
335+
if err := rows.Scan(
336+
&i.WorkspaceID,
337+
&i.WorkspaceName,
338+
&i.JobStatus,
339+
&i.Transition,
340+
pq.Array(&i.AgentIds),
341+
); err != nil {
342+
return nil, err
343+
}
344+
items = append(items, i)
345+
}
346+
if err := rows.Close(); err != nil {
347+
return nil, err
348+
}
349+
if err := rows.Err(); err != nil {
350+
return nil, err
351+
}
352+
return items, nil
353+
}
354+
315355
type userQuerier interface {
316356
GetAuthorizedUsers(ctx context.Context, arg GetUsersParams, prepared rbac.PreparedAuthorized) ([]GetUsersRow, error)
317357
}

‎coderd/database/querier.go

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

‎coderd/database/querier_test.go

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -612,6 +612,70 @@ func TestGetWorkspaceAgentUsageStatsAndLabels(t *testing.T) {
612612
})
613613
}
614614

615+
func TestGetWorkspaceAndAgents(t *testing.T) {
616+
t.Parallel()
617+
if testing.Short() {
618+
t.SkipNow()
619+
}
620+
621+
sqlDB := testSQLDB(t)
622+
err := migrations.Up(sqlDB)
623+
require.NoError(t, err)
624+
db := database.New(sqlDB)
625+
626+
org := dbgen.Organization(t, db, database.Organization{})
627+
user := dbgen.User(t, db, database.User{})
628+
tpl := dbgen.Template(t, db, database.Template{
629+
OrganizationID: org.ID,
630+
CreatedBy: user.ID,
631+
})
632+
633+
pending := createTemplateVersion(t, db, tpl, tvArgs{
634+
Status: database.ProvisionerJobStatusPending,
635+
CreateWorkspace: true,
636+
CreateAgent: true,
637+
})
638+
failed := createTemplateVersion(t, db, tpl, tvArgs{
639+
Status: database.ProvisionerJobStatusFailed,
640+
CreateWorkspace: true,
641+
CreateAgent: true,
642+
})
643+
succeeded := createTemplateVersion(t, db, tpl, tvArgs{
644+
Status: database.ProvisionerJobStatusSucceeded,
645+
CreateWorkspace: true,
646+
CreateAgent: true,
647+
})
648+
deleted := createTemplateVersion(t, db, tpl, tvArgs{
649+
Status: database.ProvisionerJobStatusSucceeded,
650+
WorkspaceTransition: database.WorkspaceTransitionDelete,
651+
CreateWorkspace: true,
652+
CreateAgent: false,
653+
})
654+
655+
ctx := testutil.Context(t, testutil.WaitLong)
656+
rows, err := db.GetWorkspacesAndAgents(ctx)
657+
require.NoError(t, err)
658+
659+
require.Len(t, rows, 4)
660+
for _, row := range rows {
661+
switch row.WorkspaceID {
662+
case pending.ID:
663+
require.Len(t, row.AgentIds, 1)
664+
require.Equal(t, database.ProvisionerJobStatusPending, row.JobStatus)
665+
case failed.ID:
666+
require.Len(t, row.AgentIds, 1)
667+
require.Equal(t, database.ProvisionerJobStatusFailed, row.JobStatus)
668+
case succeeded.ID:
669+
require.Len(t, row.AgentIds, 1)
670+
require.Equal(t, database.ProvisionerJobStatusSucceeded, row.JobStatus)
671+
case deleted.ID:
672+
require.Len(t, row.AgentIds, 0)
673+
require.Equal(t, database.ProvisionerJobStatusSucceeded, row.JobStatus)
674+
require.Equal(t, database.WorkspaceTransitionDelete, row.Transition)
675+
}
676+
}
677+
}
678+
615679
func TestInsertWorkspaceAgentLogs(t *testing.T) {
616680
t.Parallel()
617681
if testing.Short() {
@@ -1537,6 +1601,7 @@ type tvArgs struct {
15371601
Status database.ProvisionerJobStatus
15381602
// CreateWorkspace is true if we should create a workspace for the template version
15391603
CreateWorkspace bool
1604+
CreateAgent bool
15401605
WorkspaceTransition database.WorkspaceTransition
15411606
}
15421607

@@ -1607,12 +1672,16 @@ func createTemplateVersion(t testing.TB, db database.Store, tpl database.Templat
16071672
if args.WorkspaceTransition != "" {
16081673
trans = args.WorkspaceTransition
16091674
}
1675+
16101676
buildJob := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
16111677
Type: database.ProvisionerJobTypeWorkspaceBuild,
16121678
CompletedAt: now,
16131679
InitiatorID: tpl.CreatedBy,
16141680
OrganizationID: tpl.OrganizationID,
16151681
})
1682+
resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
1683+
JobID: buildJob.ID,
1684+
})
16161685
dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
16171686
WorkspaceID: wrk.ID,
16181687
TemplateVersionID: version.ID,
@@ -1621,6 +1690,11 @@ func createTemplateVersion(t testing.TB, db database.Store, tpl database.Templat
16211690
InitiatorID: tpl.CreatedBy,
16221691
JobID: buildJob.ID,
16231692
})
1693+
if args.CreateAgent {
1694+
dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
1695+
ResourceID: resource.ID,
1696+
})
1697+
}
16241698
}
16251699
return version
16261700
}

0 commit comments

Comments
 (0)