Thanks to visit codestin.com
Credit goes to github.com

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions coderd/database/dbauthz/dbauthz.go
Original file line number Diff line number Diff line change
Expand Up @@ -4528,6 +4528,12 @@ func (q *querier) GetProvisionerJobTimingsByJobID(ctx context.Context, jobID uui
return q.db.GetProvisionerJobTimingsByJobID(ctx, jobID)
}

func (q *querier) GetProvisionerJobsByIDs(ctx context.Context, ids []uuid.UUID) ([]database.ProvisionerJob, error) {
// TODO: Remove this once we have a proper rbac check for provisioner jobs.
// Details in https://github.com/coder/coder/issues/16160
return q.db.GetProvisionerJobsByIDs(ctx, ids)
}

func (q *querier) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, ids database.GetProvisionerJobsByIDsWithQueuePositionParams) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
// TODO: Remove this once we have a proper rbac check for provisioner jobs.
// Details in https://github.com/coder/coder/issues/16160
Expand Down
5 changes: 5 additions & 0 deletions coderd/database/dbauthz/dbauthz_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5427,6 +5427,11 @@ func (s *MethodTestSuite) TestSystemFunctions() {
dbm.EXPECT().GetWorkspaceAgentLogSourcesByAgentIDs(gomock.Any(), ids).Return([]database.WorkspaceAgentLogSource{}, nil).AnyTimes()
check.Args(ids).Asserts(rbac.ResourceSystem, policy.ActionRead)
}))
s.Run("GetProvisionerJobsByIDs", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
ids := []uuid.UUID{}
dbm.EXPECT().GetProvisionerJobsByIDs(gomock.Any(), ids).Return([]database.ProvisionerJob{}, nil).AnyTimes()
check.Args(ids).Asserts()
}))
s.Run("GetProvisionerJobsByIDsWithQueuePosition", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
arg := database.GetProvisionerJobsByIDsWithQueuePositionParams{}
dbm.EXPECT().GetProvisionerJobsByIDsWithQueuePosition(gomock.Any(), arg).Return([]database.GetProvisionerJobsByIDsWithQueuePositionRow{}, nil).AnyTimes()
Expand Down
8 changes: 8 additions & 0 deletions coderd/database/dbmetrics/querymetrics.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

15 changes: 15 additions & 0 deletions coderd/database/dbmock/dbmock.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 5 additions & 0 deletions coderd/database/querier.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

60 changes: 60 additions & 0 deletions coderd/database/queries.sql.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

14 changes: 14 additions & 0 deletions coderd/database/queries/provisionerjobs.sql
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,20 @@ WHERE
id = $1
FOR UPDATE;

-- name: GetProvisionerJobsByIDs :many
-- Fetches provisioner jobs by their IDs without computing queue position or
-- queue size. Callers that do not need the queue position should prefer this
-- over GetProvisionerJobsByIDsWithQueuePosition, whose window functions over
-- pending jobs and provisioner daemons are comparatively expensive.
SELECT
*
FROM
provisioner_jobs
WHERE
id = ANY(@ids :: uuid [ ])
ORDER BY
created_at;

-- name: GetProvisionerJobsByIDsWithQueuePosition :many
WITH filtered_provisioner_jobs AS (
-- Step 1: Filter provisioner_jobs
Expand Down
35 changes: 31 additions & 4 deletions coderd/workspacebuilds.go
Original file line number Diff line number Diff line change
Expand Up @@ -1110,6 +1110,36 @@ type workspaceBuildsData struct {
provisionerDaemons []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow
}

// provisionerJobsByIDs fetches provisioner jobs by ID, shaped as
// GetProvisionerJobsByIDsWithQueuePositionRow so callers can treat the result
// uniformly. When the selection requests the queue position it uses
// GetProvisionerJobsByIDsWithQueuePosition, whose queue position and size are
// computed with window functions over pending jobs and provisioner daemons and
// are comparatively expensive. Otherwise, it uses the cheaper
// GetProvisionerJobsByIDs and leaves QueuePosition and QueueSize zero.
func (api *API) provisionerJobsByIDs(ctx context.Context, jobIDs []uuid.UUID, cfg jobRelated) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
if cfg.QueuePosition {
return api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
IDs: jobIDs,
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
})
}

provisionerJobs, err := api.Database.GetProvisionerJobsByIDs(ctx, jobIDs)
if err != nil {
return nil, err
}
jobs := make([]database.GetProvisionerJobsByIDsWithQueuePositionRow, 0, len(provisionerJobs))
for _, job := range provisionerJobs {
jobs = append(jobs, database.GetProvisionerJobsByIDsWithQueuePositionRow{
ID: job.ID,
CreatedAt: job.CreatedAt,
ProvisionerJob: job,
})
}
Comment on lines +1132 to +1139

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I imagine there's not a nice way to do it, and we'll probably clean things up at their call sites in a follow up PR anyways, but it would be nice if we had a way to have Database.GetProvisionerJobsByIDs to return database.GetProvisionerJobsByIDsWithQueuePositionRow directly so we didn't have to do additional allocations inline here.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, the methods and their return types are autogenerated by sqlc, so I think it would be hard.

Fortunately these are not pointer types, so we get one allocation for the slice and we're good to go.

return jobs, nil
}

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 {
Expand All @@ -1122,10 +1152,7 @@ func (api *API) workspaceBuildsData(ctx context.Context, workspaceBuilds []datab
)
if cfg.Job != nil {
var err error
jobs, err = api.Database.GetProvisionerJobsByIDsWithQueuePosition(ctx, database.GetProvisionerJobsByIDsWithQueuePositionParams{
IDs: jobIDs,
StaleIntervalMS: provisionerdserver.StaleInterval.Milliseconds(),
})
jobs, err = api.provisionerJobsByIDs(ctx, jobIDs, *cfg.Job)
if err != nil && !errors.Is(err, sql.ErrNoRows) {
return workspaceBuildsData{}, xerrors.Errorf("get provisioner jobs: %w", err)
}
Expand Down
13 changes: 12 additions & 1 deletion coderd/wsrelateddata_internal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,12 @@ func TestWorkspaceBuildsDataQueryGating(t *testing.T) {
app := database.WorkspaceApp{ID: uuid.New(), AgentID: agent.ID}

expectJob := func(db *dbmock.MockStore) {
db.EXPECT().GetProvisionerJobsByIDs(gomock.Any(), gomock.Any()).
Return([]database.ProvisionerJob{}, nil)
db.EXPECT().GetEligibleProvisionerDaemonsByProvisionerJobIDs(gomock.Any(), gomock.Any()).
Return([]database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow{}, nil)
}
expectJobWithQueuePosition := func(db *dbmock.MockStore) {
db.EXPECT().GetProvisionerJobsByIDsWithQueuePosition(gomock.Any(), gomock.Any()).
Return([]database.GetProvisionerJobsByIDsWithQueuePositionRow{}, nil)
db.EXPECT().GetEligibleProvisionerDaemonsByProvisionerJobIDs(gomock.Any(), gomock.Any()).
Expand Down Expand Up @@ -100,6 +106,11 @@ func TestWorkspaceBuildsDataQueryGating(t *testing.T) {
cfg: latestBuildRelated{Job: &jobRelated{}},
setup: expectJob,
},
{
name: "JobWithQueuePosition",
cfg: latestBuildRelated{Job: &jobRelated{QueuePosition: true}},
setup: expectJobWithQueuePosition,
},
{
name: "TemplateVersion",
cfg: latestBuildRelated{TemplateVersion: true},
Expand Down Expand Up @@ -171,7 +182,7 @@ func TestWorkspaceBuildsDataQueryGating(t *testing.T) {
name: "All",
cfg: allLatestBuildRelated(),
setup: func(db *dbmock.MockStore) {
expectJob(db)
expectJobWithQueuePosition(db)
expectTemplateVersion(db)
expectResources(db)
expectAgents(db)
Expand Down
Loading