Skip to content

Commit ddc81b7

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

File tree

14 files changed

+402
-38
lines changed

14 files changed

+402
-38
lines changed

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2719,6 +2719,10 @@ func (q *querier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesP
27192719
return q.db.GetAuthorizedWorkspaces(ctx, arg, prep)
27202720
}
27212721

2722+
func (q *querier) GetWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID) ([]database.GetWorkspacesAndAgentsByOwnerIDRow, error) {
2723+
return q.db.GetWorkspacesAndAgentsByOwnerID(ctx, ownerID)
2724+
}
2725+
27222726
func (q *querier) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]database.Workspace, error) {
27232727
return q.db.GetWorkspacesEligibleForTransition(ctx, now)
27242728
}

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1470,6 +1470,14 @@ 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.Workspace{})
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+
check.Args(ws.OwnerID).Asserts()
1480+
}))
14731481
s.Run("GetLatestWorkspaceBuildByWorkspaceID", s.Subtest(func(db database.Store, check *expects) {
14741482
ws := dbgen.Workspace(s.T(), db, database.Workspace{})
14751483
b := dbgen.WorkspaceBuild(s.T(), db, database.WorkspaceBuild{WorkspaceID: ws.ID})

‎coderd/database/dbmem/dbmem.go

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6794,6 +6794,60 @@ func (q *FakeQuerier) GetWorkspaces(ctx context.Context, arg database.GetWorkspa
67946794
return workspaceRows, err
67956795
}
67966796

6797+
func (q *FakeQuerier) GetWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID) ([]database.GetWorkspacesAndAgentsByOwnerIDRow, error) {
6798+
q.mutex.RLock()
6799+
defer q.mutex.RUnlock()
6800+
6801+
workspaces := make([]database.Workspace, 0)
6802+
for _, workspace := range q.workspaces {
6803+
if workspace.OwnerID == ownerID && !workspace.Deleted {
6804+
workspaces = append(workspaces, workspace)
6805+
}
6806+
}
6807+
6808+
out := make([]database.GetWorkspacesAndAgentsByOwnerIDRow, 0, len(workspaces))
6809+
for _, w := range workspaces {
6810+
// these always exist
6811+
build, err := q.getLatestWorkspaceBuildByWorkspaceIDNoLock(ctx, w.ID)
6812+
if err != nil {
6813+
return nil, xerrors.Errorf("get latest build: %w", err)
6814+
}
6815+
6816+
job, err := q.getProvisionerJobByIDNoLock(ctx, build.JobID)
6817+
if err != nil {
6818+
return nil, xerrors.Errorf("get provisioner job: %w", err)
6819+
}
6820+
6821+
outAgents := make([]database.AgentIDNamePair, 0)
6822+
resources, err := q.getWorkspaceResourcesByJobIDNoLock(ctx, job.ID)
6823+
if err != nil {
6824+
return nil, xerrors.Errorf("get workspace resources: %w", err)
6825+
}
6826+
if len(resources) > 0 {
6827+
agents, err := q.getWorkspaceAgentsByResourceIDsNoLock(ctx, []uuid.UUID{resources[0].ID})
6828+
if err != nil {
6829+
return nil, xerrors.Errorf("get workspace agents: %w", err)
6830+
}
6831+
for _, a := range agents {
6832+
outAgents = append(outAgents, database.AgentIDNamePair{
6833+
ID: a.ID,
6834+
Name: a.Name,
6835+
})
6836+
}
6837+
}
6838+
6839+
out = append(out, database.GetWorkspacesAndAgentsByOwnerIDRow{
6840+
ID: w.ID,
6841+
Name: w.Name,
6842+
JobStatus: job.JobStatus,
6843+
Transition: build.Transition,
6844+
Agents: outAgents,
6845+
})
6846+
}
6847+
6848+
return out, nil
6849+
}
6850+
67976851
func (q *FakeQuerier) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]database.Workspace, error) {
67986852
q.mutex.RLock()
67996853
defer q.mutex.RUnlock()

‎coderd/database/dbmetrics/dbmetrics.go

Lines changed: 7 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: 15 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/000266_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/000266_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/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: 159 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -612,6 +612,85 @@ func TestGetWorkspaceAgentUsageStatsAndLabels(t *testing.T) {
612612
})
613613
}
614614

615+
func TestGetWorkspacesAndAgentsByOwnerID(t *testing.T) {
616+
t.Parallel()
617+
if testing.Short() {
618+
t.SkipNow()
619+
}
620+
621+
ctx := testutil.Context(t, testutil.WaitLong)
622+
sqlDB := testSQLDB(t)
623+
err := migrations.Up(sqlDB)
624+
require.NoError(t, err)
625+
db := database.New(sqlDB)
626+
627+
org := dbgen.Organization(t, db, database.Organization{})
628+
owner := dbgen.User(t, db, database.User{
629+
RBACRoles: []string{rbac.RoleOwner().String()},
630+
})
631+
tpl := dbgen.Template(t, db, database.Template{
632+
OrganizationID: org.ID,
633+
CreatedBy: owner.ID,
634+
})
635+
636+
pendingID := uuid.New()
637+
createTemplateVersion(t, db, tpl, tvArgs{
638+
Status: database.ProvisionerJobStatusPending,
639+
CreateWorkspace: true,
640+
WorkspaceID: pendingID,
641+
CreateAgent: true,
642+
})
643+
failedID := uuid.New()
644+
createTemplateVersion(t, db, tpl, tvArgs{
645+
Status: database.ProvisionerJobStatusFailed,
646+
CreateWorkspace: true,
647+
CreateAgent: true,
648+
WorkspaceID: failedID,
649+
})
650+
succeededID := uuid.New()
651+
createTemplateVersion(t, db, tpl, tvArgs{
652+
Status: database.ProvisionerJobStatusSucceeded,
653+
WorkspaceTransition: database.WorkspaceTransitionStart,
654+
CreateWorkspace: true,
655+
WorkspaceID: succeededID,
656+
CreateAgent: true,
657+
ExtraAgents: 1,
658+
ExtraBuilds: 2,
659+
})
660+
deletedID := uuid.New()
661+
createTemplateVersion(t, db, tpl, tvArgs{
662+
Status: database.ProvisionerJobStatusSucceeded,
663+
WorkspaceTransition: database.WorkspaceTransitionDelete,
664+
CreateWorkspace: true,
665+
WorkspaceID: deletedID,
666+
CreateAgent: false,
667+
})
668+
669+
ownerRows, err := db.GetWorkspacesAndAgentsByOwnerID(ctx, owner.ID)
670+
require.NoError(t, err)
671+
require.Len(t, ownerRows, 4)
672+
for _, row := range ownerRows {
673+
switch row.ID {
674+
case pendingID:
675+
require.Len(t, row.Agents, 1)
676+
require.Equal(t, database.ProvisionerJobStatusPending, row.JobStatus)
677+
case failedID:
678+
require.Len(t, row.Agents, 1)
679+
require.Equal(t, database.ProvisionerJobStatusFailed, row.JobStatus)
680+
case succeededID:
681+
require.Len(t, row.Agents, 2)
682+
require.Equal(t, database.ProvisionerJobStatusSucceeded, row.JobStatus)
683+
require.Equal(t, database.WorkspaceTransitionStart, row.Transition)
684+
case deletedID:
685+
require.Len(t, row.Agents, 0)
686+
require.Equal(t, database.ProvisionerJobStatusSucceeded, row.JobStatus)
687+
require.Equal(t, database.WorkspaceTransitionDelete, row.Transition)
688+
default:
689+
t.Fatalf("unexpected workspace ID: %s", row.ID)
690+
}
691+
}
692+
}
693+
615694
func TestInsertWorkspaceAgentLogs(t *testing.T) {
616695
t.Parallel()
617696
if testing.Short() {
@@ -1537,7 +1616,11 @@ type tvArgs struct {
15371616
Status database.ProvisionerJobStatus
15381617
// CreateWorkspace is true if we should create a workspace for the template version
15391618
CreateWorkspace bool
1619+
WorkspaceID uuid.UUID
1620+
CreateAgent bool
15401621
WorkspaceTransition database.WorkspaceTransition
1622+
ExtraAgents int
1623+
ExtraBuilds int
15411624
}
15421625

15431626
// createTemplateVersion is a helper function to create a version with its dependencies.
@@ -1554,49 +1637,18 @@ func createTemplateVersion(t testing.TB, db database.Store, tpl database.Templat
15541637
CreatedBy: tpl.CreatedBy,
15551638
})
15561639

1557-
earlier := sql.NullTime{
1558-
Time: dbtime.Now().Add(time.Second * -30),
1559-
Valid: true,
1560-
}
1561-
now := sql.NullTime{
1562-
Time: dbtime.Now(),
1563-
Valid: true,
1564-
}
1565-
j := database.ProvisionerJob{
1640+
latestJob := database.ProvisionerJob{
15661641
ID: version.JobID,
1567-
CreatedAt: earlier.Time,
1568-
UpdatedAt: earlier.Time,
15691642
Error: sql.NullString{},
15701643
OrganizationID: tpl.OrganizationID,
15711644
InitiatorID: tpl.CreatedBy,
15721645
Type: database.ProvisionerJobTypeTemplateVersionImport,
15731646
}
1574-
1575-
switch args.Status {
1576-
case database.ProvisionerJobStatusRunning:
1577-
j.StartedAt = earlier
1578-
case database.ProvisionerJobStatusPending:
1579-
case database.ProvisionerJobStatusFailed:
1580-
j.StartedAt = earlier
1581-
j.CompletedAt = now
1582-
j.Error = sql.NullString{
1583-
String: "failed",
1584-
Valid: true,
1585-
}
1586-
j.ErrorCode = sql.NullString{
1587-
String: "failed",
1588-
Valid: true,
1589-
}
1590-
case database.ProvisionerJobStatusSucceeded:
1591-
j.StartedAt = earlier
1592-
j.CompletedAt = now
1593-
default:
1594-
t.Fatalf("invalid status: %s", args.Status)
1595-
}
1596-
1597-
dbgen.ProvisionerJob(t, db, nil, j)
1647+
setJobStatus(t, args.Status, &latestJob)
1648+
dbgen.ProvisionerJob(t, db, nil, latestJob)
15981649
if args.CreateWorkspace {
15991650
wrk := dbgen.Workspace(t, db, database.Workspace{
1651+
ID: args.WorkspaceID,
16001652
CreatedAt: time.Time{},
16011653
UpdatedAt: time.Time{},
16021654
OwnerID: tpl.CreatedBy,
@@ -1607,24 +1659,93 @@ func createTemplateVersion(t testing.TB, db database.Store, tpl database.Templat
16071659
if args.WorkspaceTransition != "" {
16081660
trans = args.WorkspaceTransition
16091661
}
1610-
buildJob := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
1662+
latestJob = database.ProvisionerJob{
16111663
Type: database.ProvisionerJobTypeWorkspaceBuild,
1612-
CompletedAt: now,
16131664
InitiatorID: tpl.CreatedBy,
16141665
OrganizationID: tpl.OrganizationID,
1666+
}
1667+
setJobStatus(t, args.Status, &latestJob)
1668+
latestJob = dbgen.ProvisionerJob(t, db, nil, latestJob)
1669+
latestResource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
1670+
JobID: latestJob.ID,
16151671
})
16161672
dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
16171673
WorkspaceID: wrk.ID,
16181674
TemplateVersionID: version.ID,
16191675
BuildNumber: 1,
16201676
Transition: trans,
16211677
InitiatorID: tpl.CreatedBy,
1622-
JobID: buildJob.ID,
1678+
JobID: latestJob.ID,
16231679
})
1680+
for i := 0; i < args.ExtraBuilds; i++ {
1681+
latestJob = database.ProvisionerJob{
1682+
Type: database.ProvisionerJobTypeWorkspaceBuild,
1683+
InitiatorID: tpl.CreatedBy,
1684+
OrganizationID: tpl.OrganizationID,
1685+
}
1686+
setJobStatus(t, args.Status, &latestJob)
1687+
latestJob = dbgen.ProvisionerJob(t, db, nil, latestJob)
1688+
latestResource = dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
1689+
JobID: latestJob.ID,
1690+
})
1691+
dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
1692+
WorkspaceID: wrk.ID,
1693+
TemplateVersionID: version.ID,
1694+
BuildNumber: int32(i) + 2,
1695+
Transition: trans,
1696+
InitiatorID: tpl.CreatedBy,
1697+
JobID: latestJob.ID,
1698+
})
1699+
}
1700+
1701+
if args.CreateAgent {
1702+
dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
1703+
ResourceID: latestResource.ID,
1704+
})
1705+
}
1706+
for i := 0; i < args.ExtraAgents; i++ {
1707+
dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
1708+
ResourceID: latestResource.ID,
1709+
})
1710+
}
16241711
}
16251712
return version
16261713
}
16271714

1715+
func setJobStatus(t testing.TB, status database.ProvisionerJobStatus, j *database.ProvisionerJob) {
1716+
t.Helper()
1717+
1718+
earlier := sql.NullTime{
1719+
Time: dbtime.Now().Add(time.Second * -30),
1720+
Valid: true,
1721+
}
1722+
now := sql.NullTime{
1723+
Time: dbtime.Now(),
1724+
Valid: true,
1725+
}
1726+
switch status {
1727+
case database.ProvisionerJobStatusRunning:
1728+
j.StartedAt = earlier
1729+
case database.ProvisionerJobStatusPending:
1730+
case database.ProvisionerJobStatusFailed:
1731+
j.StartedAt = earlier
1732+
j.CompletedAt = now
1733+
j.Error = sql.NullString{
1734+
String: "failed",
1735+
Valid: true,
1736+
}
1737+
j.ErrorCode = sql.NullString{
1738+
String: "failed",
1739+
Valid: true,
1740+
}
1741+
case database.ProvisionerJobStatusSucceeded:
1742+
j.StartedAt = earlier
1743+
j.CompletedAt = now
1744+
default:
1745+
t.Fatalf("invalid status: %s", status)
1746+
}
1747+
}
1748+
16281749
func TestArchiveVersions(t *testing.T) {
16291750
t.Parallel()
16301751
if testing.Short() {

0 commit comments

Comments
 (0)