Skip to content

Commit ac9bc72

Browse files
ethanndicksonethanndickson
committed
chore: add db query to retrieve workspaces & their agents
1 parent 85d80ad commit ac9bc72

File tree

13 files changed

+395
-38
lines changed

13 files changed

+395
-38
lines changed

‎coderd/database/dbauthz/dbauthz.go

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

2768+
func (q *querier) GetWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID) ([]database.GetWorkspacesAndAgentsByOwnerIDRow, error) {
2769+
return q.db.GetWorkspacesAndAgentsByOwnerID(ctx, ownerID)
2770+
}
2771+
27682772
func (q *querier) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]database.WorkspaceTable, error) {
27692773
return q.db.GetWorkspacesEligibleForTransition(ctx, now)
27702774
}

‎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.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+
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.WorkspaceTable{})
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
@@ -6839,6 +6839,60 @@ 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+
q.mutex.RLock()
6844+
defer q.mutex.RUnlock()
6845+
6846+
workspaces := make([]database.WorkspaceTable, 0)
6847+
for _, workspace := range q.workspaces {
6848+
if workspace.OwnerID == ownerID && !workspace.Deleted {
6849+
workspaces = append(workspaces, workspace)
6850+
}
6851+
}
6852+
6853+
out := make([]database.GetWorkspacesAndAgentsByOwnerIDRow, 0, len(workspaces))
6854+
for _, w := range workspaces {
6855+
// these always exist
6856+
build, err := q.getLatestWorkspaceBuildByWorkspaceIDNoLock(ctx, w.ID)
6857+
if err != nil {
6858+
return nil, xerrors.Errorf("get latest build: %w", err)
6859+
}
6860+
6861+
job, err := q.getProvisionerJobByIDNoLock(ctx, build.JobID)
6862+
if err != nil {
6863+
return nil, xerrors.Errorf("get provisioner job: %w", err)
6864+
}
6865+
6866+
outAgents := make([]database.AgentIDNamePair, 0)
6867+
resources, err := q.getWorkspaceResourcesByJobIDNoLock(ctx, job.ID)
6868+
if err != nil {
6869+
return nil, xerrors.Errorf("get workspace resources: %w", err)
6870+
}
6871+
if len(resources) > 0 {
6872+
agents, err := q.getWorkspaceAgentsByResourceIDsNoLock(ctx, []uuid.UUID{resources[0].ID})
6873+
if err != nil {
6874+
return nil, xerrors.Errorf("get workspace agents: %w", err)
6875+
}
6876+
for _, a := range agents {
6877+
outAgents = append(outAgents, database.AgentIDNamePair{
6878+
ID: a.ID,
6879+
Name: a.Name,
6880+
})
6881+
}
6882+
}
6883+
6884+
out = append(out, database.GetWorkspacesAndAgentsByOwnerIDRow{
6885+
ID: w.ID,
6886+
Name: w.Name,
6887+
JobStatus: job.JobStatus,
6888+
Transition: build.Transition,
6889+
Agents: outAgents,
6890+
})
6891+
}
6892+
6893+
return out, nil
6894+
}
6895+
68426896
func (q *FakeQuerier) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]database.WorkspaceTable, error) {
68436897
q.mutex.RLock()
68446898
defer q.mutex.RUnlock()

‎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.WorkspaceTable{
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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy