From 01ea8a12d8ccc9bb5ab533562b4e1c196dd9eb9e Mon Sep 17 00:00:00 2001 From: Spike Curtis Date: Wed, 19 Aug 2026 14:25:54 +0000 Subject: [PATCH] refactor(coderd): add related-data selection to workspace queries --- coderd/aitasks.go | 4 +- coderd/workspacebuilds.go | 169 +++++++++------ coderd/workspaces.go | 79 ++++--- coderd/wsrelateddata.go | 87 ++++++++ coderd/wsrelateddata_internal_test.go | 294 ++++++++++++++++++++++++++ 5 files changed, 535 insertions(+), 98 deletions(-) create mode 100644 coderd/wsrelateddata.go create mode 100644 coderd/wsrelateddata_internal_test.go diff --git a/coderd/aitasks.go b/coderd/aitasks.go index a5b5252638941..545db5d0896a3 100644 --- a/coderd/aitasks.go +++ b/coderd/aitasks.go @@ -487,7 +487,7 @@ func (api *API) convertTasks(ctx context.Context, requesterID uuid.UUID, dbTasks } // Gather associated data and convert to API workspaces. - data, err := api.workspaceData(ctx, workspaces) + data, err := api.workspaceData(ctx, workspaces, allWorkspaceRelated()) if err != nil { return nil, xerrors.Errorf("fetch workspace data: %w", err) } @@ -547,7 +547,7 @@ func (api *API) taskGet(rw http.ResponseWriter, r *http.Request) { return } - data, err := api.workspaceData(ctx, []database.Workspace{workspace}) + data, err := api.workspaceData(ctx, []database.Workspace{workspace}, allWorkspaceRelated()) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Internal error fetching workspace resources.", diff --git a/coderd/workspacebuilds.go b/coderd/workspacebuilds.go index 070135e6d1c7e..e42dc1119b522 100644 --- a/coderd/workspacebuilds.go +++ b/coderd/workspacebuilds.go @@ -51,7 +51,7 @@ func (api *API) workspaceBuild(rw http.ResponseWriter, r *http.Request) { workspaceBuild := httpmw.WorkspaceBuildParam(r) workspace := httpmw.WorkspaceParam(r) - data, err := api.workspaceBuildsData(ctx, []database.WorkspaceBuild{workspaceBuild}) + data, err := api.workspaceBuildsData(ctx, []database.WorkspaceBuild{workspaceBuild}, allLatestBuildRelated()) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Internal error getting workspace build data.", @@ -190,7 +190,7 @@ func (api *API) workspaceBuilds(rw http.ResponseWriter, r *http.Request) { return } - data, err := api.workspaceBuildsData(ctx, workspaceBuilds) + data, err := api.workspaceBuildsData(ctx, workspaceBuilds, allLatestBuildRelated()) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Internal error getting workspace build data.", @@ -281,7 +281,7 @@ func (api *API) workspaceBuildByBuildNumber(rw http.ResponseWriter, r *http.Requ return } - data, err := api.workspaceBuildsData(ctx, []database.WorkspaceBuild{workspaceBuild}) + data, err := api.workspaceBuildsData(ctx, []database.WorkspaceBuild{workspaceBuild}, allLatestBuildRelated()) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Internal error getting workspace build data.", @@ -1110,39 +1110,59 @@ type workspaceBuildsData struct { provisionerDaemons []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow } -func (api *API) workspaceBuildsData(ctx context.Context, workspaceBuilds []database.WorkspaceBuild) (workspaceBuildsData, error) { +func (api *API) workspaceBuildsData(ctx context.Context, workspaceBuilds []database.WorkspaceBuild, cfg latestBuildRelated) (workspaceBuildsData, error) { jobIDs := make([]uuid.UUID, 0, len(workspaceBuilds)) for _, build := range workspaceBuilds { jobIDs = append(jobIDs, build.JobID) } - jobs, err := api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{ - IDs: jobIDs, - StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(), - }) - if err != nil && !errors.Is(err, sql.ErrNoRows) { - return workspaceBuildsData{}, xerrors.Errorf("get provisioner jobs: %w", err) - } - pendingJobIDs := []uuid.UUID{} - for _, job := range jobs { - if job.ProvisionerJob.JobStatus == database.ProvisionerJobStatusPending { - pendingJobIDs = append(pendingJobIDs, job.ProvisionerJob.ID) + + var ( + jobs []database.GetProvisionerJobsByIDsWithQueuePositionRow + pendingJobProvisioners []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow + ) + if cfg.Job != nil { + var err error + jobs, err = api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{ + IDs: jobIDs, + StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(), + }) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return workspaceBuildsData{}, xerrors.Errorf("get provisioner jobs: %w", err) + } + pendingJobIDs := []uuid.UUID{} + for _, job := range jobs { + if job.ProvisionerJob.JobStatus == database.ProvisionerJobStatusPending { + pendingJobIDs = append(pendingJobIDs, job.ProvisionerJob.ID) + } } - } - pendingJobProvisioners, err := api.Database.GetEligibleProvisionerDaemonsByProvisionerJobIDs(ctx, pendingJobIDs) - if err != nil && !errors.Is(err, sql.ErrNoRows) { - return workspaceBuildsData{}, xerrors.Errorf("get provisioner daemons: %w", err) + pendingJobProvisioners, err = api.Database.GetEligibleProvisionerDaemonsByProvisionerJobIDs(ctx, pendingJobIDs) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return workspaceBuildsData{}, xerrors.Errorf("get provisioner daemons: %w", err) + } } - templateVersionIDs := make([]uuid.UUID, 0, len(workspaceBuilds)) - for _, build := range workspaceBuilds { - templateVersionIDs = append(templateVersionIDs, build.TemplateVersionID) + var templateVersions []database.TemplateVersion + if cfg.TemplateVersion { + templateVersionIDs := make([]uuid.UUID, 0, len(workspaceBuilds)) + for _, build := range workspaceBuilds { + templateVersionIDs = append(templateVersionIDs, build.TemplateVersionID) + } + + var err error + // nolint:gocritic // Getting template versions by ID is a system function. + templateVersions, err = api.Database.GetTemplateVersionsByIDs(dbauthz.AsSystemRestricted(ctx), templateVersionIDs) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return workspaceBuildsData{}, xerrors.Errorf("get template versions: %w", err) + } } - // nolint:gocritic // Getting template versions by ID is a system function. - templateVersions, err := api.Database.GetTemplateVersionsByIDs(dbauthz.AsSystemRestricted(ctx), templateVersionIDs) - if err != nil && !errors.Is(err, sql.ErrNoRows) { - return workspaceBuildsData{}, xerrors.Errorf("get template versions: %w", err) + if cfg.Resources == nil { + return workspaceBuildsData{ + jobs: jobs, + templateVersions: templateVersions, + provisionerDaemons: pendingJobProvisioners, + }, nil } // nolint:gocritic // Getting workspace resources by job ID is a system function. @@ -1164,19 +1184,36 @@ func (api *API) workspaceBuildsData(ctx context.Context, workspaceBuilds []datab resourceIDs = append(resourceIDs, resource.ID) } - // nolint:gocritic // Getting workspace resource metadata by resource ID is a system function. - metadata, err := api.Database.GetWorkspaceResourceMetadataByResourceIDs(dbauthz.AsSystemRestricted(ctx), resourceIDs) - if err != nil && !errors.Is(err, sql.ErrNoRows) { - return workspaceBuildsData{}, xerrors.Errorf("fetching resource metadata: %w", err) + var ( + metadata []database.WorkspaceResourceMetadatum + agents []database.WorkspaceAgent + egRes errgroup.Group + ) + if cfg.Resources.Metadata { + egRes.Go(func() (err error) { + // nolint:gocritic // Getting workspace resource metadata by resource ID is a system function. + metadata, err = api.Database.GetWorkspaceResourceMetadataByResourceIDs(dbauthz.AsSystemRestricted(ctx), resourceIDs) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return xerrors.Errorf("fetching resource metadata: %w", err) + } + return nil + }) } - - // nolint:gocritic // Getting workspace agents by resource IDs is a system function. - agents, err := api.Database.GetWorkspaceAgentsByResourceIDs(dbauthz.AsSystemRestricted(ctx), resourceIDs) - if err != nil && !errors.Is(err, sql.ErrNoRows) { - return workspaceBuildsData{}, xerrors.Errorf("get workspace agents: %w", err) + if cfg.Resources.Agents != nil { + egRes.Go(func() (err error) { + // nolint:gocritic // Getting workspace agents by resource IDs is a system function. + agents, err = api.Database.GetWorkspaceAgentsByResourceIDs(dbauthz.AsSystemRestricted(ctx), resourceIDs) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return xerrors.Errorf("get workspace agents: %w", err) + } + return nil + }) + } + if err := egRes.Wait(); err != nil { + return workspaceBuildsData{}, err } - if len(resources) == 0 { + if cfg.Resources.Agents == nil || len(agents) == 0 { return workspaceBuildsData{ jobs: jobs, templateVersions: templateVersions, @@ -1185,6 +1222,7 @@ func (api *API) workspaceBuildsData(ctx context.Context, workspaceBuilds []datab provisionerDaemons: pendingJobProvisioners, }, nil } + agentsCfg := cfg.Resources.Agents agentIDs := make([]uuid.UUID, 0) for _, agent := range agents { @@ -1198,35 +1236,44 @@ func (api *API) workspaceBuildsData(ctx context.Context, workspaceBuilds []datab ) var eg errgroup.Group - eg.Go(func() (err error) { - // nolint:gocritic // Getting workspace apps by agent IDs is a system function. - apps, err = api.Database.GetWorkspaceAppsByAgentIDs(dbauthz.AsSystemRestricted(ctx), agentIDs) - return err - }) - eg.Go(func() (err error) { - // nolint:gocritic // Getting workspace scripts by agent IDs is a system function. - scripts, err = api.Database.GetWorkspaceAgentScriptsByAgentIDs(dbauthz.AsSystemRestricted(ctx), agentIDs) - return err - }) - eg.Go(func() error { - // nolint:gocritic // Getting workspace agent log sources by agent IDs is a system function. - logSources, err = api.Database.GetWorkspaceAgentLogSourcesByAgentIDs(dbauthz.AsSystemRestricted(ctx), agentIDs) - return err - }) - err = eg.Wait() - if err != nil { + if agentsCfg.Apps != nil { + eg.Go(func() (err error) { + // nolint:gocritic // Getting workspace apps by agent IDs is a system function. + apps, err = api.Database.GetWorkspaceAppsByAgentIDs(dbauthz.AsSystemRestricted(ctx), agentIDs) + return err + }) + } + if agentsCfg.Scripts { + eg.Go(func() (err error) { + // nolint:gocritic // Getting workspace scripts by agent IDs is a system function. + scripts, err = api.Database.GetWorkspaceAgentScriptsByAgentIDs(dbauthz.AsSystemRestricted(ctx), agentIDs) + return err + }) + } + if agentsCfg.LogSources { + eg.Go(func() error { + var err error + // nolint:gocritic // Getting workspace agent log sources by agent IDs is a system function. + logSources, err = api.Database.GetWorkspaceAgentLogSourcesByAgentIDs(dbauthz.AsSystemRestricted(ctx), agentIDs) + return err + }) + } + if err := eg.Wait(); err != nil { return workspaceBuildsData{}, err } - appIDs := make([]uuid.UUID, 0) - for _, app := range apps { - appIDs = append(appIDs, app.ID) - } + var statuses []database.WorkspaceAppStatus + if cfg.appStatuses() { + appIDs := make([]uuid.UUID, 0) + for _, app := range apps { + appIDs = append(appIDs, app.ID) + } - // nolint:gocritic // Getting workspace app statuses by app IDs is a system function. - statuses, err := api.Database.GetWorkspaceAppStatusesByAppIDs(dbauthz.AsSystemRestricted(ctx), appIDs) - if err != nil && !errors.Is(err, sql.ErrNoRows) { - return workspaceBuildsData{}, xerrors.Errorf("get workspace app statuses: %w", err) + // nolint:gocritic // Getting workspace app statuses by app IDs is a system function. + statuses, err = api.Database.GetWorkspaceAppStatusesByAppIDs(dbauthz.AsSystemRestricted(ctx), appIDs) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return workspaceBuildsData{}, xerrors.Errorf("get workspace app statuses: %w", err) + } } return workspaceBuildsData{ diff --git a/coderd/workspaces.go b/coderd/workspaces.go index d48f6d5dd72e1..5f3455c17eeb3 100644 --- a/coderd/workspaces.go +++ b/coderd/workspaces.go @@ -96,7 +96,7 @@ func (api *API) workspace(rw http.ResponseWriter, r *http.Request) { return } - data, err := api.workspaceData(ctx, []database.Workspace{workspace}) + data, err := api.workspaceData(ctx, []database.Workspace{workspace}, allWorkspaceRelated()) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Internal error fetching workspace resources.", @@ -224,7 +224,7 @@ func (api *API) workspaces(rw http.ResponseWriter, r *http.Request) { return } - data, err := api.workspaceData(ctx, workspaces) + data, err := api.workspaceData(ctx, workspaces, allWorkspaceRelated()) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Internal error fetching workspace resources.", @@ -313,7 +313,7 @@ func (api *API) workspaceByOwnerAndName(rw http.ResponseWriter, r *http.Request) return } - data, err := api.workspaceData(ctx, []database.Workspace{workspace}) + data, err := api.workspaceData(ctx, []database.Workspace{workspace}, allWorkspaceRelated()) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Internal error fetching workspace resources.", @@ -1562,7 +1562,7 @@ func (api *API) putWorkspaceDormant(rw http.ResponseWriter, r *http.Request) { return } - data, err := api.workspaceData(ctx, []database.Workspace{workspace}) + data, err := api.workspaceData(ctx, []database.Workspace{workspace}, allWorkspaceRelated()) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Internal error fetching workspace resources.", @@ -2140,7 +2140,7 @@ func (api *API) watchWorkspace( return } - data, err := api.workspaceData(ctx, []database.Workspace{workspace}) + data, err := api.workspaceData(ctx, []database.Workspace{workspace}, allWorkspaceRelated()) if err != nil { _ = sendEvent(codersdk.ServerSentEvent{ Type: codersdk.ServerSentEventTypeError, @@ -2678,7 +2678,7 @@ func (api *API) allowWorkspaceSharing(ctx context.Context, rw http.ResponseWrite // does not have the correct perms to read a given template, the template will // not be returned. // So the caller must check the templates & users exist before using them. -func (api *API) workspaceData(ctx context.Context, workspaces []database.Workspace) (workspaceData, error) { +func (api *API) workspaceData(ctx context.Context, workspaces []database.Workspace, cfg workspaceRelated) (workspaceData, error) { workspaceIDs := make([]uuid.UUID, 0, len(workspaces)) templateIDs := make([]uuid.UUID, 0, len(workspaces)) for _, workspace := range workspaces { @@ -2692,41 +2692,50 @@ func (api *API) workspaceData(ctx context.Context, workspaces []database.Workspa appStatuses []database.WorkspaceAppStatus eg errgroup.Group ) - eg.Go(func() (err error) { - templates, err = api.Database.GetTemplatesWithFilter(ctx, database.GetTemplatesWithFilterParams{ - IDs: templateIDs, + if cfg.Template { + eg.Go(func() (err error) { + templates, err = api.Database.GetTemplatesWithFilter(ctx, database.GetTemplatesWithFilterParams{ + IDs: templateIDs, + }) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return xerrors.Errorf("get templates: %w", err) + } + return nil }) - if err != nil && !errors.Is(err, sql.ErrNoRows) { - return xerrors.Errorf("get templates: %w", err) - } - return nil - }) - eg.Go(func() (err error) { - // This query must be run as system restricted to be efficient. - // nolint:gocritic - builds, err = api.Database.GetLatestWorkspaceBuildsByWorkspaceIDs(dbauthz.AsSystemRestricted(ctx), workspaceIDs) - if err != nil && !errors.Is(err, sql.ErrNoRows) { - return xerrors.Errorf("get workspace builds: %w", err) - } - return nil - }) - eg.Go(func() (err error) { - // This query must be run as system restricted to be efficient. - // nolint:gocritic - appStatuses, err = api.Database.GetLatestWorkspaceAppStatusesByWorkspaceIDs(dbauthz.AsSystemRestricted(ctx), workspaceIDs) - if err != nil && !errors.Is(err, sql.ErrNoRows) { - return xerrors.Errorf("get workspace app statuses: %w", err) - } - return nil - }) + } + if cfg.LatestBuild != nil { + eg.Go(func() (err error) { + // This query must be run as system restricted to be efficient. + // nolint:gocritic + builds, err = api.Database.GetLatestWorkspaceBuildsByWorkspaceIDs(dbauthz.AsSystemRestricted(ctx), workspaceIDs) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return xerrors.Errorf("get workspace builds: %w", err) + } + return nil + }) + } + if cfg.LatestBuild.appStatuses() { + eg.Go(func() (err error) { + // This query must be run as system restricted to be efficient. + // nolint:gocritic + appStatuses, err = api.Database.GetLatestWorkspaceAppStatusesByWorkspaceIDs(dbauthz.AsSystemRestricted(ctx), workspaceIDs) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return xerrors.Errorf("get workspace app statuses: %w", err) + } + return nil + }) + } err := eg.Wait() if err != nil { return workspaceData{}, err } - data, err := api.workspaceBuildsData(ctx, builds) - if err != nil { - return workspaceData{}, xerrors.Errorf("get workspace builds data: %w", err) + var data workspaceBuildsData + if cfg.LatestBuild != nil { + data, err = api.workspaceBuildsData(ctx, builds, *cfg.LatestBuild) + if err != nil { + return workspaceData{}, xerrors.Errorf("get workspace builds data: %w", err) + } } apiBuilds, err := api.convertWorkspaceBuilds( diff --git a/coderd/wsrelateddata.go b/coderd/wsrelateddata.go new file mode 100644 index 0000000000000..7e1f7e7640098 --- /dev/null +++ b/coderd/wsrelateddata.go @@ -0,0 +1,87 @@ +package coderd + +// workspaceRelated selects which workspace-related database objects to load when +// building a codersdk.Workspace. Loading a fully populated workspace is +// expensive, so callers use this to avoid querying data they will not use. +// +// The type is a tree that mirrors the parent/child relationships between those +// objects: a build has a job, resources, and a template version; a resource has +// agents; an agent has apps; and so on. Branch nodes are pointers that are +// non-nil when selected; leaf nodes are bools. Modeling it as a tree makes +// selecting a child without its parent unrepresentable, which is exactly the +// constraint loading requires: a parent must be queried to learn the +// identifiers of its children. +// +// A zero value (nil branches) selects nothing but the workspace itself. +// allWorkspaceRelated selects everything. +// +// The trailing comment on each field is that node's dotted path from the root +// of the tree, e.g. latest_build.resources.agents. +type workspaceRelated struct { + Template bool // template + LatestBuild *latestBuildRelated // latest_build +} + +type latestBuildRelated struct { + Job *jobRelated // latest_build.job + Resources *resourcesRelated // latest_build.resources + TemplateVersion bool // latest_build.template_version +} + +type jobRelated struct { + QueuePosition bool // latest_build.job.queue_position +} + +type resourcesRelated struct { + Metadata bool // latest_build.resources.metadata + Agents *agentsRelated // latest_build.resources.agents +} + +type agentsRelated struct { + Apps *appsRelated // latest_build.resources.agents.apps + Scripts bool // latest_build.resources.agents.scripts + LogSources bool // latest_build.resources.agents.log_sources +} + +type appsRelated struct { + Statuses bool // latest_build.resources.agents.apps.statuses +} + +// allWorkspaceRelated returns a selection that loads every related object. It +// reproduces the behavior of callers that have not been narrowed to a specific +// subset. +func allWorkspaceRelated() workspaceRelated { + latestBuild := allLatestBuildRelated() + return workspaceRelated{ + Template: true, + LatestBuild: &latestBuild, + } +} + +// allLatestBuildRelated returns the latest-build subtree with every node +// selected. +func allLatestBuildRelated() latestBuildRelated { + return latestBuildRelated{ + Job: &jobRelated{QueuePosition: true}, + Resources: &resourcesRelated{ + Metadata: true, + Agents: &agentsRelated{ + Apps: &appsRelated{Statuses: true}, + Scripts: true, + LogSources: true, + }, + }, + TemplateVersion: true, + } +} + +// appStatuses reports whether app statuses +// (latest_build.resources.agents.apps.statuses) are selected. It is nil-safe so +// callers holding a possibly-nil subtree can descend without a chain of guards. +func (c *latestBuildRelated) appStatuses() bool { + return c != nil && + c.Resources != nil && + c.Resources.Agents != nil && + c.Resources.Agents.Apps != nil && + c.Resources.Agents.Apps.Statuses +} diff --git a/coderd/wsrelateddata_internal_test.go b/coderd/wsrelateddata_internal_test.go new file mode 100644 index 0000000000000..ff7c586975e7e --- /dev/null +++ b/coderd/wsrelateddata_internal_test.go @@ -0,0 +1,294 @@ +package coderd + +import ( + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbmock" + "github.com/coder/coder/v2/testutil" +) + +// TestAllWorkspaceRelated verifies that allWorkspaceRelated selects every node +// in the related-data hierarchy. +func TestAllWorkspaceRelated(t *testing.T) { + t.Parallel() + + all := allWorkspaceRelated() + require.True(t, all.Template) + require.NotNil(t, all.LatestBuild) + require.NotNil(t, all.LatestBuild.Job) + require.True(t, all.LatestBuild.Job.QueuePosition) + require.True(t, all.LatestBuild.TemplateVersion) + require.NotNil(t, all.LatestBuild.Resources) + require.True(t, all.LatestBuild.Resources.Metadata) + require.NotNil(t, all.LatestBuild.Resources.Agents) + require.NotNil(t, all.LatestBuild.Resources.Agents.Apps) + require.True(t, all.LatestBuild.Resources.Agents.Apps.Statuses) + require.True(t, all.LatestBuild.Resources.Agents.Scripts) + require.True(t, all.LatestBuild.Resources.Agents.LogSources) + require.True(t, all.LatestBuild.appStatuses()) +} + +// TestLatestBuildRelatedAppStatuses verifies the nil-safe appStatuses accessor +// only reports true when the full apps.statuses path is present. +func TestLatestBuildRelatedAppStatuses(t *testing.T) { + t.Parallel() + + require.False(t, (*latestBuildRelated)(nil).appStatuses()) + require.False(t, (&latestBuildRelated{}).appStatuses()) + require.False(t, (&latestBuildRelated{Resources: &resourcesRelated{}}).appStatuses()) + require.False(t, (&latestBuildRelated{Resources: &resourcesRelated{Agents: &agentsRelated{}}}).appStatuses()) + require.False(t, (&latestBuildRelated{Resources: &resourcesRelated{Agents: &agentsRelated{Apps: &appsRelated{}}}}).appStatuses()) + require.True(t, (&latestBuildRelated{Resources: &resourcesRelated{Agents: &agentsRelated{Apps: &appsRelated{Statuses: true}}}}).appStatuses()) +} + +// TestWorkspaceBuildsDataQueryGating asserts that workspaceBuildsData only +// issues the database queries implied by the selection. gomock is strict, so +// any query that is not set up here fails the test if invoked. +func TestWorkspaceBuildsDataQueryGating(t *testing.T) { + t.Parallel() + + build := database.WorkspaceBuild{ + ID: uuid.New(), + JobID: uuid.New(), + WorkspaceID: uuid.New(), + TemplateVersionID: uuid.New(), + } + resource := database.WorkspaceResource{ID: uuid.New(), JobID: build.JobID} + agent := database.WorkspaceAgent{ID: uuid.New(), ResourceID: resource.ID} + app := database.WorkspaceApp{ID: uuid.New(), AgentID: agent.ID} + + expectJob := func(db *dbmock.MockStore) { + db.EXPECT().GetProvisionerJobsByIDsWithQueuePosition(gomock.Any(), gomock.Any()). + Return([]database.GetProvisionerJobsByIDsWithQueuePositionRow{}, nil) + db.EXPECT().GetEligibleProvisionerDaemonsByProvisionerJobIDs(gomock.Any(), gomock.Any()). + Return([]database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow{}, nil) + } + expectTemplateVersion := func(db *dbmock.MockStore) { + db.EXPECT().GetTemplateVersionsByIDs(gomock.Any(), gomock.Any()). + Return([]database.TemplateVersion{}, nil) + } + expectResources := func(db *dbmock.MockStore) { + db.EXPECT().GetWorkspaceResourcesByJobIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceResource{resource}, nil) + } + expectAgents := func(db *dbmock.MockStore) { + db.EXPECT().GetWorkspaceAgentsByResourceIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceAgent{agent}, nil) + } + expectApps := func(db *dbmock.MockStore) { + db.EXPECT().GetWorkspaceAppsByAgentIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceApp{app}, nil) + } + + cases := []struct { + name string + cfg latestBuildRelated + setup func(*dbmock.MockStore) + }{ + { + name: "BuildOnly", + cfg: latestBuildRelated{}, + setup: func(*dbmock.MockStore) {}, + }, + { + name: "Job", + cfg: latestBuildRelated{Job: &jobRelated{}}, + setup: expectJob, + }, + { + name: "TemplateVersion", + cfg: latestBuildRelated{TemplateVersion: true}, + setup: expectTemplateVersion, + }, + { + name: "Resources", + cfg: latestBuildRelated{Resources: &resourcesRelated{}}, + setup: expectResources, + }, + { + name: "Metadata", + cfg: latestBuildRelated{Resources: &resourcesRelated{Metadata: true}}, + setup: func(db *dbmock.MockStore) { + expectResources(db) + db.EXPECT().GetWorkspaceResourceMetadataByResourceIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceResourceMetadatum{}, nil) + }, + }, + { + name: "Agents", + cfg: latestBuildRelated{Resources: &resourcesRelated{Agents: &agentsRelated{}}}, + setup: func(db *dbmock.MockStore) { + expectResources(db) + expectAgents(db) + }, + }, + { + name: "Apps", + cfg: latestBuildRelated{Resources: &resourcesRelated{Agents: &agentsRelated{Apps: &appsRelated{}}}}, + setup: func(db *dbmock.MockStore) { + expectResources(db) + expectAgents(db) + expectApps(db) + }, + }, + { + name: "AppStatuses", + cfg: latestBuildRelated{Resources: &resourcesRelated{Agents: &agentsRelated{Apps: &appsRelated{Statuses: true}}}}, + setup: func(db *dbmock.MockStore) { + expectResources(db) + expectAgents(db) + expectApps(db) + db.EXPECT().GetWorkspaceAppStatusesByAppIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceAppStatus{}, nil) + }, + }, + { + name: "Scripts", + cfg: latestBuildRelated{Resources: &resourcesRelated{Agents: &agentsRelated{Scripts: true}}}, + setup: func(db *dbmock.MockStore) { + expectResources(db) + expectAgents(db) + db.EXPECT().GetWorkspaceAgentScriptsByAgentIDs(gomock.Any(), gomock.Any()). + Return([]database.GetWorkspaceAgentScriptsByAgentIDsRow{}, nil) + }, + }, + { + name: "LogSources", + cfg: latestBuildRelated{Resources: &resourcesRelated{Agents: &agentsRelated{LogSources: true}}}, + setup: func(db *dbmock.MockStore) { + expectResources(db) + expectAgents(db) + db.EXPECT().GetWorkspaceAgentLogSourcesByAgentIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceAgentLogSource{}, nil) + }, + }, + { + name: "All", + cfg: allLatestBuildRelated(), + setup: func(db *dbmock.MockStore) { + expectJob(db) + expectTemplateVersion(db) + expectResources(db) + expectAgents(db) + expectApps(db) + db.EXPECT().GetWorkspaceResourceMetadataByResourceIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceResourceMetadatum{}, nil) + db.EXPECT().GetWorkspaceAppStatusesByAppIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceAppStatus{}, nil) + db.EXPECT().GetWorkspaceAgentScriptsByAgentIDs(gomock.Any(), gomock.Any()). + Return([]database.GetWorkspaceAgentScriptsByAgentIDsRow{}, nil) + db.EXPECT().GetWorkspaceAgentLogSourcesByAgentIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceAgentLogSource{}, nil) + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + tc.setup(db) + + api := &API{Options: &Options{Database: db}} + _, err := api.workspaceBuildsData(ctx, []database.WorkspaceBuild{build}, tc.cfg) + require.NoError(t, err) + }) + } +} + +// TestWorkspaceBuildsDataResourcesShortCircuit verifies that when resources are +// requested but none exist, no agent-level queries are issued even though the +// deeper nodes are selected. +func TestWorkspaceBuildsDataResourcesShortCircuit(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + build := database.WorkspaceBuild{ID: uuid.New(), JobID: uuid.New(), TemplateVersionID: uuid.New()} + // No resources returned: the agent, app, and status queries must be skipped. + db.EXPECT().GetWorkspaceResourcesByJobIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceResource{}, nil) + + cfg := latestBuildRelated{Resources: &resourcesRelated{Agents: &agentsRelated{Apps: &appsRelated{Statuses: true}}}} + api := &API{Options: &Options{Database: db}} + _, err := api.workspaceBuildsData(ctx, []database.WorkspaceBuild{build}, cfg) + require.NoError(t, err) +} + +// TestWorkspaceDataQueryGating asserts that workspaceData gates its +// workspace-level queries (templates, latest builds, and latest app statuses) +// on the selection. +func TestWorkspaceDataQueryGating(t *testing.T) { + t.Parallel() + + workspace := database.Workspace{ID: uuid.New(), TemplateID: uuid.New()} + + cases := []struct { + name string + cfg workspaceRelated + setup func(*dbmock.MockStore) + }{ + { + name: "None", + cfg: workspaceRelated{}, + setup: func(*dbmock.MockStore) {}, + }, + { + name: "Template", + cfg: workspaceRelated{Template: true}, + setup: func(db *dbmock.MockStore) { + db.EXPECT().GetTemplatesWithFilter(gomock.Any(), gomock.Any()). + Return([]database.Template{}, nil) + }, + }, + { + name: "LatestBuild", + cfg: workspaceRelated{LatestBuild: &latestBuildRelated{}}, + setup: func(db *dbmock.MockStore) { + // No builds returned, so no build-subtree queries run. + db.EXPECT().GetLatestWorkspaceBuildsByWorkspaceIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceBuild{}, nil) + }, + }, + { + name: "AppStatuses", + cfg: workspaceRelated{LatestBuild: &latestBuildRelated{ + Resources: &resourcesRelated{Agents: &agentsRelated{Apps: &appsRelated{Statuses: true}}}, + }}, + setup: func(db *dbmock.MockStore) { + // The workspace-level latest app status query is gated on the + // apps.statuses node. With no builds, only the workspace-level + // status, build, and (empty) resource queries run. + db.EXPECT().GetLatestWorkspaceAppStatusesByWorkspaceIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceAppStatus{}, nil) + db.EXPECT().GetLatestWorkspaceBuildsByWorkspaceIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceBuild{}, nil) + db.EXPECT().GetWorkspaceResourcesByJobIDs(gomock.Any(), gomock.Any()). + Return([]database.WorkspaceResource{}, nil) + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + tc.setup(db) + + api := &API{Options: &Options{Database: db}} + _, err := api.workspaceData(ctx, []database.Workspace{workspace}, tc.cfg) + require.NoError(t, err) + }) + } +}