Skip to content

Commit f941e78

Browse files
ethanndicksonethanndickson
authored
chore: add db query to retrieve workspaces & their agents (#14792)
Second PR for #14716. Adds a query that, given a user ID, returns all the workspaces they own, that can also be `ActionRead` by the requesting user. ``` type GetWorkspacesAndAgentsByOwnerIDRow struct { WorkspaceID uuid.UUID `db:"workspace_id" json:"workspace_id"` WorkspaceName string `db:"workspace_name" json:"workspace_name"` JobStatus ProvisionerJobStatus `db:"job_status" json:"job_status"` Transition WorkspaceTransition `db:"transition" json:"transition"` Agents []AgentIDNamePair `db:"agents" json:"agents"` } ``` `JobStatus` and `Transition` are set using the latest build/job of the workspace. Deleted workspaces are not included.
1 parent 31506e6 commit f941e78

File tree

15 files changed

+551
-38
lines changed

15 files changed

+551
-38
lines changed

‎coderd/database/dbauthz/dbauthz.go

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

2792+
func (q *querier) GetWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID) ([]database.GetWorkspacesAndAgentsByOwnerIDRow, error) {
2793+
prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceWorkspace.Type)
2794+
if err != nil {
2795+
return nil, xerrors.Errorf("(dev error) prepare sql filter: %w", err)
2796+
}
2797+
return q.db.GetAuthorizedWorkspacesAndAgentsByOwnerID(ctx, ownerID, prep)
2798+
}
2799+
27922800
func (q *querier) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]database.WorkspaceTable, error) {
27932801
return q.db.GetWorkspacesEligibleForTransition(ctx, now)
27942802
}
@@ -4242,6 +4250,10 @@ func (q *querier) GetAuthorizedWorkspaces(ctx context.Context, arg database.GetW
42424250
return q.GetWorkspaces(ctx, arg)
42434251
}
42444252

4253+
func (q *querier) GetAuthorizedWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID, _ rbac.PreparedAuthorized) ([]database.GetWorkspacesAndAgentsByOwnerIDRow, error) {
4254+
return q.GetWorkspacesAndAgentsByOwnerID(ctx, ownerID)
4255+
}
4256+
42454257
// GetAuthorizedUsers is not required for dbauthz since GetUsers is already
42464258
// authenticated.
42474259
func (q *querier) GetAuthorizedUsers(ctx context.Context, arg database.GetUsersParams, _ rbac.PreparedAuthorized) ([]database.GetUsersRow, error) {

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1470,6 +1470,24 @@ func (s *MethodTestSuite) TestWorkspace() {
14701470
// No asserts here because SQLFilter.
14711471
check.Args(database.GetWorkspacesParams{}, emptyPreparedAuthorized{}).Asserts()
14721472
}))
1473+
s.Run("GetWorkspacesAndAgentsByOwnerID", s.Subtest(func(db database.Store, check *expects) {
1474+
ws := dbgen.Workspace(s.T(), db, database.WorkspaceTable{})
1475+
build := dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{WorkspaceID: ws.ID, JobID: uuid.New()})
1476+
_ = dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{ID: build.JobID, Type: database.ProvisionerJobTypeWorkspaceBuild})
1477+
res := dbgen.WorkspaceResource(s.T(), db, database.WorkspaceResource{JobID: build.JobID})
1478+
_ = dbgen.WorkspaceAgent(s.T(), db, database.WorkspaceAgent{ResourceID: res.ID})
1479+
// No asserts here because SQLFilter.
1480+
check.Args(ws.OwnerID).Asserts()
1481+
}))
1482+
s.Run("GetAuthorizedWorkspacesAndAgentsByOwnerID", s.Subtest(func(db database.Store, check *expects) {
1483+
ws := dbgen.Workspace(s.T(), db, database.WorkspaceTable{})
1484+
build := dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{WorkspaceID: ws.ID, JobID: uuid.New()})
1485+
_ = dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{ID: build.JobID, Type: database.ProvisionerJobTypeWorkspaceBuild})
1486+
res := dbgen.WorkspaceResource(s.T(), db, database.WorkspaceResource{JobID: build.JobID})
1487+
_ = dbgen.WorkspaceAgent(s.T(), db, database.WorkspaceAgent{ResourceID: res.ID})
1488+
// No asserts here because SQLFilter.
1489+
check.Args(ws.OwnerID, emptyPreparedAuthorized{}).Asserts()
1490+
}))
14731491
s.Run("GetLatestWorkspaceBuildByWorkspaceID", s.Subtest(func(db database.Store, check *expects) {
14741492
ws := dbgen.Workspace(s.T(), db, database.WorkspaceTable{})
14751493
b := dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{WorkspaceID: ws.ID})

‎coderd/database/dbmem/dbmem.go

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6839,6 +6839,11 @@ func (q *FakeQuerier) GetWorkspaces(ctx context.Context, arg database.GetWorkspa
68396839
return workspaceRows, err
68406840
}
68416841

6842+
func (q *FakeQuerier) GetWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID) ([]database.GetWorkspacesAndAgentsByOwnerIDRow, error) {
6843+
// No auth filter.
6844+
return q.GetAuthorizedWorkspacesAndAgentsByOwnerID(ctx, ownerID, nil)
6845+
}
6846+
68426847
func (q *FakeQuerier) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]database.WorkspaceTable, error) {
68436848
q.mutex.RLock()
68446849
defer q.mutex.RUnlock()
@@ -11224,6 +11229,67 @@ func (q *FakeQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg database.
1122411229
return q.convertToWorkspaceRowsNoLock(ctx, workspaces, int64(beforePageCount), arg.WithSummary), nil
1122511230
}
1122611231

11232+
func (q *FakeQuerier) GetAuthorizedWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID, prepared rbac.PreparedAuthorized) ([]database.GetWorkspacesAndAgentsByOwnerIDRow, error) {
11233+
q.mutex.RLock()
11234+
defer q.mutex.RUnlock()
11235+
11236+
if prepared != nil {
11237+
// Call this to match the same function calls as the SQL implementation.
11238+
_, err := prepared.CompileToSQL(ctx, rbac.ConfigWithoutACL())
11239+
if err != nil {
11240+
return nil, err
11241+
}
11242+
}
11243+
workspaces := make([]database.WorkspaceTable, 0)
11244+
for _, workspace := range q.workspaces {
11245+
if workspace.OwnerID == ownerID && !workspace.Deleted {
11246+
workspaces = append(workspaces, workspace)
11247+
}
11248+
}
11249+
11250+
out := make([]database.GetWorkspacesAndAgentsByOwnerIDRow, 0, len(workspaces))
11251+
for _, w := range workspaces {
11252+
// these always exist
11253+
build, err := q.getLatestWorkspaceBuildByWorkspaceIDNoLock(ctx, w.ID)
11254+
if err != nil {
11255+
return nil, xerrors.Errorf("get latest build: %w", err)
11256+
}
11257+
11258+
job, err := q.getProvisionerJobByIDNoLock(ctx, build.JobID)
11259+
if err != nil {
11260+
return nil, xerrors.Errorf("get provisioner job: %w", err)
11261+
}
11262+
11263+
outAgents := make([]database.AgentIDNamePair, 0)
11264+
resources, err := q.getWorkspaceResourcesByJobIDNoLock(ctx, job.ID)
11265+
if err != nil {
11266+
return nil, xerrors.Errorf("get workspace resources: %w", err)
11267+
}
11268+
if len(resources) > 0 {
11269+
agents, err := q.getWorkspaceAgentsByResourceIDsNoLock(ctx, []uuid.UUID{resources[0].ID})
11270+
if err != nil {
11271+
return nil, xerrors.Errorf("get workspace agents: %w", err)
11272+
}
11273+
for _, a := range agents {
11274+
outAgents = append(outAgents, database.AgentIDNamePair{
11275+
ID: a.ID,
11276+
Name: a.Name,
11277+
})
11278+
}
11279+
}
11280+
11281+
out = append(out, database.GetWorkspacesAndAgentsByOwnerIDRow{
11282+
ID: w.ID,
11283+
Name: w.Name,
11284+
JobStatus: job.JobStatus,
11285+
Transition: build.Transition,
11286+
Agents: outAgents,
11287+
})
11288+
}
11289+
11290+
return out, nil
11291+
}
11292+
1122711293
func (q *FakeQuerier) GetAuthorizedUsers(ctx context.Context, arg database.GetUsersParams, prepared rbac.PreparedAuthorized) ([]database.GetUsersRow, error) {
1122811294
if err := validateDatabaseType(arg); err != nil {
1122911295
return nil, err

‎coderd/database/dbmetrics/querymetrics.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/dump.sql

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

‎coderd/database/migrations/000273_workspace_updates.down.sql

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
DROP TYPE agent_id_name_pair;

‎coderd/database/migrations/000273_workspace_updates.up.sql

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
CREATE TYPE agent_id_name_pair AS (
2+
id uuid,
3+
name text
4+
);

‎coderd/database/modelqueries.go

Lines changed: 44 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+
GetAuthorizedWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID, prepared rbac.PreparedAuthorized) ([]GetWorkspacesAndAgentsByOwnerIDRow, error)
224225
}
225226

226227
// GetAuthorizedWorkspaces returns all workspaces that the user is authorized to access.
@@ -320,6 +321,49 @@ func (q *sqlQuerier) GetAuthorizedWorkspaces(ctx context.Context, arg GetWorkspa
320321
return items, nil
321322
}
322323

324+
func (q *sqlQuerier) GetAuthorizedWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID, prepared rbac.PreparedAuthorized) ([]GetWorkspacesAndAgentsByOwnerIDRow, error) {
325+
authorizedFilter, err := prepared.CompileToSQL(ctx, rbac.ConfigWorkspaces())
326+
if err != nil {
327+
return nil, xerrors.Errorf("compile authorized filter: %w", err)
328+
}
329+
330+
// In order to properly use ORDER BY, OFFSET, and LIMIT, we need to inject the
331+
// authorizedFilter between the end of the where clause and those statements.
332+
filtered, err := insertAuthorizedFilter(getWorkspacesAndAgentsByOwnerID, fmt.Sprintf(" AND %s", authorizedFilter))
333+
if err != nil {
334+
return nil, xerrors.Errorf("insert authorized filter: %w", err)
335+
}
336+
337+
// The name comment is for metric tracking
338+
query := fmt.Sprintf("-- name: GetAuthorizedWorkspacesAndAgentsByOwnerID :many\n%s", filtered)
339+
rows, err := q.db.QueryContext(ctx, query, ownerID)
340+
if err != nil {
341+
return nil, err
342+
}
343+
defer rows.Close()
344+
var items []GetWorkspacesAndAgentsByOwnerIDRow
345+
for rows.Next() {
346+
var i GetWorkspacesAndAgentsByOwnerIDRow
347+
if err := rows.Scan(
348+
&i.ID,
349+
&i.Name,
350+
&i.JobStatus,
351+
&i.Transition,
352+
pq.Array(&i.Agents),
353+
); err != nil {
354+
return nil, err
355+
}
356+
items = append(items, i)
357+
}
358+
if err := rows.Close(); err != nil {
359+
return nil, err
360+
}
361+
if err := rows.Err(); err != nil {
362+
return nil, err
363+
}
364+
return items, nil
365+
}
366+
323367
type userQuerier interface {
324368
GetAuthorizedUsers(ctx context.Context, arg GetUsersParams, prepared rbac.PreparedAuthorized) ([]GetUsersRow, error)
325369
}

‎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.

0 commit comments

Comments
 (0)