Skip to content

Commit 1638825

Browse files
ethanndicksonethanndickson
committed
chore: add db query to retrieve workspaces & their agents
1 parent c6e4428 commit 1638825

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
@@ -2705,6 +2705,14 @@ func (q *querier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesP
27052705
return q.db.GetAuthorizedWorkspaces(ctx, arg, prep)
27062706
}
27072707

2708+
func (q *querier) GetWorkspacesAndAgents(ctx context.Context) ([]database.GetWorkspacesAndAgentsRow, error) {
2709+
prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceWorkspace.Type)
2710+
if err != nil {
2711+
return nil, xerrors.Errorf("(dev error) prepare sql filter: %w", err)
2712+
}
2713+
return q.db.GetAuthorizedWorkspacesAndAgents(ctx, prep)
2714+
}
2715+
27082716
func (q *querier) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]database.Workspace, error) {
27092717
return q.db.GetWorkspacesEligibleForTransition(ctx, now)
27102718
}
@@ -4142,6 +4150,10 @@ func (q *querier) GetAuthorizedWorkspaces(ctx context.Context, arg database.GetW
41424150
return q.GetWorkspaces(ctx, arg)
41434151
}
41444152

4153+
func (q *querier) GetAuthorizedWorkspacesAndAgents(ctx context.Context, _ rbac.PreparedAuthorized) ([]database.GetWorkspacesAndAgentsRow, error) {
4154+
return q.GetWorkspacesAndAgents(ctx)
4155+
}
4156+
41454157
// GetAuthorizedUsers is not required for dbauthz since GetUsers is already
41464158
// authenticated.
41474159
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
@@ -6714,6 +6714,11 @@ func (q *FakeQuerier) GetWorkspaces(ctx context.Context, arg database.GetWorkspa
67146714
return workspaceRows, err
67156715
}
67166716

6717+
func (q *FakeQuerier) GetWorkspacesAndAgents(ctx context.Context) ([]database.GetWorkspacesAndAgentsRow, error) {
6718+
// No auth filter.
6719+
return q.GetAuthorizedWorkspacesAndAgents(ctx, nil)
6720+
}
6721+
67176722
func (q *FakeQuerier) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]database.Workspace, error) {
67186723
q.mutex.RLock()
67196724
defer q.mutex.RUnlock()
@@ -11080,6 +11085,63 @@ func (q *FakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.
1108011085
return q.convertToWorkspaceRowsNoLock(ctx, workspaces, int64(beforePageCount), arg.WithSummary), nil
1108111086
}
1108211087

11088+
func (q *FakeQuerier) GetAuthorizedWorkspacesAndAgents(ctx context.Context, prepared rbac.PreparedAuthorized) ([]database.GetWorkspacesAndAgentsRow, error) {
11089+
q.mutex.RLock()
11090+
defer q.mutex.RUnlock()
11091+
11092+
if prepared != nil {
11093+
// Call this to match the same function calls as the SQL implementation.
11094+
_, err := prepared.CompileToSQL(ctx, rbac.ConfigWithoutACL())
11095+
if err != nil {
11096+
return nil, err
11097+
}
11098+
}
11099+
workspaces := make([]database.Workspace, 0)
11100+
for _, workspace := range q.workspaces {
11101+
if prepared != nil && prepared.Authorize(ctx, workspace.RBACObject()) == nil {
11102+
workspaces = append(workspaces, workspace)
11103+
}
11104+
}
11105+
11106+
out := make([]database.GetWorkspacesAndAgentsRow, 0, len(workspaces))
11107+
for _, w := range workspaces {
11108+
// psql constraints ensure all these exist
11109+
build, err := q.getLatestWorkspaceBuildByWorkspaceIDNoLock(ctx, w.ID)
11110+
if err != nil {
11111+
return nil, xerrors.Errorf("get latest build: %w", err)
11112+
}
11113+
11114+
job, err := q.getProvisionerJobByIDNoLock(ctx, build.JobID)
11115+
if err != nil {
11116+
return nil, xerrors.Errorf("get provisioner job: %w", err)
11117+
}
11118+
11119+
resource, err := q.getWorkspaceResourcesByJobIDNoLock(ctx, job.ID)
11120+
if err != nil || len(resource) == 0 {
11121+
return nil, xerrors.Errorf("get workspace resources: %w", err)
11122+
}
11123+
11124+
agents, err := q.getWorkspaceAgentsByResourceIDsNoLock(ctx, []uuid.UUID{resource[0].ID})
11125+
if err != nil {
11126+
return nil, xerrors.Errorf("get workspace agents: %w", err)
11127+
}
11128+
agentIDs := make([]uuid.UUID, 0, len(agents))
11129+
for _, a := range agents {
11130+
agentIDs = append(agentIDs, a.ID)
11131+
}
11132+
11133+
out = append(out, database.GetWorkspacesAndAgentsRow{
11134+
WorkspaceID: w.ID,
11135+
WorkspaceName: w.Name,
11136+
JobStatus: job.JobStatus,
11137+
Transition: build.Transition,
11138+
AgentIds: agentIDs,
11139+
})
11140+
}
11141+
11142+
return out, nil
11143+
}
11144+
1108311145
func (q *FakeQuerier) GetAuthorizedUsers(ctx context.Context, arg database.GetUsersParams, prepared rbac.PreparedAuthorized) ([]database.GetUsersRow, error) {
1108411146
if err := validateDatabaseType(arg); err != nil {
1108511147
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)