Skip to content

chore: optimize GetPrebuiltWorkspaces query #18717

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open
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
142 changes: 142 additions & 0 deletions coderd/database/querier_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5021,3 +5021,145 @@ func requireUsersMatch(t testing.TB, expected []database.User, found []database.
t.Helper()
require.ElementsMatch(t, expected, database.ConvertUserRows(found), msg)
}

// TestGetRunningPrebuiltWorkspaces ensures the correct behavior of the
// GetRunningPrebuiltWorkspaces query.
func TestGetRunningPrebuiltWorkspaces(t *testing.T) {
t.Parallel()

if !dbtestutil.WillUsePostgres() {
t.Skip("Test requires PostgreSQL for complex queries")
}

ctx := testutil.Context(t, testutil.WaitLong)
db, _ := dbtestutil.NewDB(t)

// Given: a prebuilt workspace with a successful start build and a stop build.
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
template := dbgen.Template(t, db, database.Template{
CreatedBy: user.ID,
OrganizationID: org.ID,
})
templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
TemplateID: uuid.NullUUID{UUID: template.ID, Valid: true},
OrganizationID: org.ID,
CreatedBy: user.ID,
})
preset := dbgen.Preset(t, db, database.InsertPresetParams{
TemplateVersionID: templateVersion.ID,
DesiredInstances: sql.NullInt32{Int32: 1, Valid: true},
})

// Create a prebuild workspace (owned by system user)
prebuildSystemUser := uuid.MustParse("c42fdf75-3097-471c-8c33-fb52454d81c0")
stoppedPrebuild := dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: prebuildSystemUser,
TemplateID: template.ID,
Name: "test-prebuild",
Deleted: false,
})

// Create a successful START build
stoppedPrebuildJob1 := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
OrganizationID: org.ID,
InitiatorID: database.PrebuildsSystemUserID,
Provisioner: database.ProvisionerTypeEcho,
Type: database.ProvisionerJobTypeWorkspaceBuild,
StartedAt: sql.NullTime{Time: dbtime.Now().Add(-time.Minute), Valid: true},
CompletedAt: sql.NullTime{Time: dbtime.Now(), Valid: true},
Error: sql.NullString{},
ErrorCode: sql.NullString{},
})
stoppedPrebuiltWorkspaceBuild1 := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
WorkspaceID: stoppedPrebuild.ID,
TemplateVersionID: templateVersion.ID,
TemplateVersionPresetID: uuid.NullUUID{UUID: preset.ID, Valid: true},
JobID: stoppedPrebuildJob1.ID,
BuildNumber: 1,
Transition: database.WorkspaceTransitionStart,
InitiatorID: database.PrebuildsSystemUserID,
Reason: database.BuildReasonInitiator,
})

// Create a STOP build (making this prebuild "not running")
stoppedPrebuildWorkspaceJob2 := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
OrganizationID: org.ID,
InitiatorID: database.PrebuildsSystemUserID,
Provisioner: database.ProvisionerTypeEcho,
Type: database.ProvisionerJobTypeWorkspaceBuild,
StartedAt: sql.NullTime{Time: dbtime.Now().Add(-time.Minute), Valid: true},
CompletedAt: sql.NullTime{Time: dbtime.Now(), Valid: true},
Error: sql.NullString{},
ErrorCode: sql.NullString{},
})
stoppedPrebuiltWorkspaceBuild2 := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
WorkspaceID: stoppedPrebuild.ID,
TemplateVersionID: templateVersion.ID,
TemplateVersionPresetID: uuid.NullUUID{UUID: preset.ID, Valid: true},
JobID: stoppedPrebuildWorkspaceJob2.ID,
BuildNumber: 2,
Transition: database.WorkspaceTransitionStop,
InitiatorID: database.PrebuildsSystemUserID,
Reason: database.BuildReasonInitiator,
})

// Create a second running prebuild workspace with a successful start build
// and no stop build.
runningPrebuild := dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: prebuildSystemUser,
TemplateID: template.ID,
Name: "test-running-prebuild",
Deleted: false,
})
// Create a successful START build for the running prebuild workspace
runningPrebuildJob1 := dbgen.ProvisionerJob(t, db, nil,
database.ProvisionerJob{
OrganizationID: org.ID,
InitiatorID: user.ID,
Provisioner: database.ProvisionerTypeEcho,
Type: database.ProvisionerJobTypeWorkspaceBuild,
StartedAt: sql.NullTime{Time: dbtime.Now().Add(-time.Minute), Valid: true},
CompletedAt: sql.NullTime{Time: dbtime.Now(), Valid: true},
Error: sql.NullString{},
ErrorCode: sql.NullString{},
})
runningPrebuiltWorkspaceBuild1 := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
WorkspaceID: runningPrebuild.ID,
TemplateVersionID: templateVersion.ID,
TemplateVersionPresetID: uuid.NullUUID{UUID: preset.ID, Valid: true},
JobID: runningPrebuildJob1.ID,
BuildNumber: 1,
Transition: database.WorkspaceTransitionStart,
InitiatorID: database.PrebuildsSystemUserID,
Reason: database.BuildReasonInitiator,
})

// Create a third running regular workspace (not a prebuild) with a successful
// start build.
_ = dbgen.Workspace(t, db, database.WorkspaceTable{
OwnerID: user.ID,
TemplateID: template.ID,
Name: "test-running-regular-workspace",
Deleted: false,
})

// Given: assert test invariants.
require.Equal(t, database.WorkspaceTransitionStart, stoppedPrebuiltWorkspaceBuild1.Transition)
require.Equal(t, int32(1), stoppedPrebuiltWorkspaceBuild1.BuildNumber)
require.Equal(t, database.ProvisionerJobStatusSucceeded, stoppedPrebuildJob1.JobStatus)
require.Equal(t, database.WorkspaceTransitionStop, stoppedPrebuiltWorkspaceBuild2.Transition)
require.Equal(t, int32(2), stoppedPrebuiltWorkspaceBuild2.BuildNumber)
require.Equal(t, database.ProvisionerJobStatusSucceeded, stoppedPrebuildWorkspaceJob2.JobStatus)
require.Equal(t, database.WorkspaceTransitionStart, runningPrebuiltWorkspaceBuild1.Transition)
require.Equal(t, int32(1), runningPrebuiltWorkspaceBuild1.BuildNumber)
require.Equal(t, database.ProvisionerJobStatusSucceeded, runningPrebuildJob1.JobStatus)

// When: we query for running prebuild workspaces
runningPrebuilds, err := db.GetRunningPrebuiltWorkspaces(ctx)
require.NoError(t, err)

// Then: the stopped prebuild workspace should not be returned.
require.Len(t, runningPrebuilds, 1, "expected only one running prebuilt workspace")
require.Equal(t, runningPrebuild.ID, runningPrebuilds[0].ID, "expected the running prebuilt workspace to be returned")
}
57 changes: 46 additions & 11 deletions coderd/database/queries.sql.go

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

58 changes: 47 additions & 11 deletions coderd/database/queries/prebuilds.sql
Original file line number Diff line number Diff line change
Expand Up @@ -49,18 +49,54 @@ WHERE tvp.desired_instances IS NOT NULL -- Consider only presets that have a pre
AND (t.id = sqlc.narg('template_id')::uuid OR sqlc.narg('template_id') IS NULL);

-- name: GetRunningPrebuiltWorkspaces :many
WITH latest_prebuilds AS (
SELECT
workspaces.id,
workspaces.name,
workspaces.template_id,
workspace_latest_builds.template_version_id,
workspace_latest_builds.job_id,
workspaces.created_at
FROM workspace_latest_builds
JOIN workspaces ON workspaces.id = workspace_latest_builds.workspace_id
WHERE workspace_latest_builds.transition = 'start'::workspace_transition
AND workspace_latest_builds.job_status = 'succeeded'::provisioner_job_status
AND workspaces.owner_id = 'c42fdf75-3097-471c-8c33-fb52454d81c0'::UUID
AND NOT workspaces.deleted
),
workspace_latest_presets AS (
SELECT DISTINCT ON (latest_prebuilds.id)
latest_prebuilds.id AS workspace_id,
workspace_builds.template_version_preset_id AS current_preset_id
FROM latest_prebuilds
JOIN workspace_builds ON workspace_builds.workspace_id = latest_prebuilds.id
WHERE workspace_builds.transition = 'start'::workspace_transition
AND workspace_builds.template_version_preset_id IS NOT NULL
ORDER BY latest_prebuilds.id, workspace_builds.build_number DESC
),
ready_agents AS (
SELECT
latest_prebuilds.job_id,
BOOL_AND(workspace_agents.lifecycle_state = 'ready'::workspace_agent_lifecycle_state)::boolean AS ready
FROM latest_prebuilds
JOIN workspace_resources ON workspace_resources.job_id = latest_prebuilds.job_id
JOIN workspace_agents ON workspace_agents.resource_id = workspace_resources.id
WHERE workspace_agents.deleted = false
AND workspace_agents.parent_id IS NULL
GROUP BY latest_prebuilds.job_id
)
SELECT
p.id,
p.name,
p.template_id,
b.template_version_id,
p.current_preset_id AS current_preset_id,
p.ready,
p.created_at
FROM workspace_prebuilds p
INNER JOIN workspace_latest_builds b ON b.workspace_id = p.id
WHERE (b.transition = 'start'::workspace_transition
AND b.job_status = 'succeeded'::provisioner_job_status);
latest_prebuilds.id,
latest_prebuilds.name,
latest_prebuilds.template_id,
latest_prebuilds.template_version_id,
workspace_latest_presets.current_preset_id,
COALESCE(ready_agents.ready, false)::boolean AS ready,
latest_prebuilds.created_at
FROM latest_prebuilds
LEFT JOIN ready_agents ON ready_agents.job_id = latest_prebuilds.job_id
LEFT JOIN workspace_latest_presets ON workspace_latest_presets.workspace_id = latest_prebuilds.id
;

-- name: CountInProgressPrebuilds :many
-- CountInProgressPrebuilds returns the number of in-progress prebuilds, grouped by preset ID and transition.
Expand Down
Loading
Loading