-- name: ClaimPrebuiltWorkspace :one

UPDATE workspaces w

SET owner_id = @new_user_id::uuid,

name = @new_name::text,

updated_at = @now::timestamptz,

-- Update autostart_schedule, next_start_at and ttl according to template and workspace-level

-- configurations, allowing the workspace to be managed by the lifecycle executor as expected.

autostart_schedule = @autostart_schedule,

next_start_at = @next_start_at,

ttl = @workspace_ttl,

-- Update last_used_at during claim to ensure the claimed workspace is treated as recently used.

-- This avoids unintended dormancy caused by prebuilds having stale usage timestamps.

last_used_at = @now::timestamptz,

-- Clear dormant and deletion timestamps as a safeguard to ensure a clean lifecycle state after claim.

-- These fields should not be set on prebuilds, but we defensively reset them here to prevent

-- accidental dormancy or deletion by the lifecycle executor.

dormant_at = NULL,

deleting_at = NULL

WHERE w.id IN (

SELECT p.id

FROM workspace_prebuilds p

INNER JOIN workspace_latest_builds b ON b.workspace_id = p.id

INNER JOIN templates t ON p.template_id = t.id

WHERE (b.transition = 'start'::workspace_transition

AND b.job_status IN ('succeeded'::provisioner_job_status))

-- The prebuilds system should never try to claim a prebuild for an inactive template version.

-- Nevertheless, this filter is here as a defensive measure:

AND b.template_version_id = t.active_version_id

AND p.current_preset_id = @preset_id::uuid

AND p.ready

AND NOT t.deleted

LIMIT 1 FOR UPDATE OF p SKIP LOCKED -- Ensure that a concurrent request will not select the same prebuild.

)

RETURNING w.id, w.name;

-- name: GetTemplatePresetsWithPrebuilds :many

-- GetTemplatePresetsWithPrebuilds retrieves template versions with configured presets and prebuilds.

-- It also returns the number of desired instances for each preset.

-- If template_id is specified, only template versions associated with that template will be returned.

SELECT

t.id AS template_id,

t.name AS template_name,

o.id AS organization_id,

o.name AS organization_name,

tv.id AS template_version_id,

tv.name AS template_version_name,

tv.id = t.active_version_id AS using_active_version,

tvp.id,

tvp.name,

tvp.desired_instances AS desired_instances,

tvp.scheduling_timezone,

tvp.invalidate_after_secs AS ttl,

tvp.prebuild_status,

tvp.last_invalidated_at,

t.deleted,

t.deprecated != '' AS deprecated

FROM templates t

INNER JOIN template_versions tv ON tv.template_id = t.id

INNER JOIN template_version_presets tvp ON tvp.template_version_id = tv.id

INNER JOIN organizations o ON o.id = t.organization_id

WHERE tvp.desired_instances IS NOT NULL -- Consider only presets that have a prebuild configuration.

-- AND NOT t.deleted -- We don't exclude deleted templates because there's no constraint in the DB preventing a soft deletion on a template while workspaces are running.

AND (t.id = sqlc.narg('template_id')::uuid OR sqlc.narg('template_id') IS NULL);

-- name: GetRunningPrebuiltWorkspaces :many

WITH latest_prebuilds AS (

-- All workspaces that match the following criteria:

-- 1. Owned by prebuilds user

-- 2. Not deleted

-- 3. Latest build is a 'start' transition

-- 4. Latest build was successful

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 (

-- For each of the above workspaces, the preset_id of the most recent

-- successful start transition.

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 (

-- For each of the above workspaces, check if all agents are ready.

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

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

ORDER BY latest_prebuilds.id;

-- name: CountInProgressPrebuilds :many

-- CountInProgressPrebuilds returns the number of in-progress prebuilds, grouped by preset ID and transition.

-- Prebuild considered in-progress if it's in the "pending", "starting", "stopping", or "deleting" state.

SELECT t.id AS template_id, wpb.template_version_id, wpb.transition, COUNT(wpb.transition)::int AS count, wlb.template_version_preset_id as preset_id

FROM workspace_latest_builds wlb

INNER JOIN workspace_prebuild_builds wpb ON wpb.id = wlb.id

-- We only need these counts for active template versions.

-- It doesn't influence whether we create or delete prebuilds

-- for inactive template versions. This is because we never create

-- prebuilds for inactive template versions, we always delete

-- running prebuilds for inactive template versions, and we ignore

-- prebuilds that are still building.

INNER JOIN templates t ON t.active_version_id = wlb.template_version_id

WHERE wlb.job_status IN ('pending'::provisioner_job_status, 'running'::provisioner_job_status)

-- AND NOT t.deleted -- We don't exclude deleted templates because there's no constraint in the DB preventing a soft deletion on a template while workspaces are running.

GROUP BY t.id, wpb.template_version_id, wpb.transition, wlb.template_version_preset_id;

-- GetPresetsBackoff groups workspace builds by preset ID.

-- Each preset is associated with exactly one template version ID.

-- For each group, the query checks up to N of the most recent jobs that occurred within the

-- lookback period, where N equals the number of desired instances for the corresponding preset.

-- If at least one of the job within a group has failed, we should backoff on the corresponding preset ID.

-- Query returns a list of preset IDs for which we should backoff.

-- Only active template versions with configured presets are considered.

-- We also return the number of failed workspace builds that occurred during the lookback period.

--

-- NOTE:

-- - To **decide whether to back off**, we look at up to the N most recent builds (within the defined lookback period).

-- - To **calculate the number of failed builds**, we consider all builds within the defined lookback period.

--

-- The number of failed builds is used downstream to determine the backoff duration.

-- name: GetPresetsBackoff :many

WITH filtered_builds AS (

-- Only select builds which are for prebuild creations

SELECT wlb.template_version_id, wlb.created_at, tvp.id AS preset_id, wlb.job_status, tvp.desired_instances

FROM template_version_presets tvp

INNER JOIN workspace_latest_builds wlb ON wlb.template_version_preset_id = tvp.id

INNER JOIN workspaces w ON wlb.workspace_id = w.id

INNER JOIN template_versions tv ON wlb.template_version_id = tv.id

INNER JOIN templates t ON tv.template_id = t.id AND t.active_version_id = tv.id

WHERE tvp.desired_instances IS NOT NULL -- Consider only presets that have a prebuild configuration.

AND wlb.transition = 'start'::workspace_transition

AND w.owner_id = 'c42fdf75-3097-471c-8c33-fb52454d81c0'

AND NOT t.deleted

),

time_sorted_builds AS (

-- Group builds by preset, then sort each group by created_at.

SELECT fb.template_version_id, fb.created_at, fb.preset_id, fb.job_status, fb.desired_instances,

ROW_NUMBER() OVER (PARTITION BY fb.preset_id ORDER BY fb.created_at DESC) as rn

FROM filtered_builds fb

),

failed_count AS (

-- Count failed builds per preset in the given period

SELECT preset_id, COUNT(*) AS num_failed

FROM filtered_builds

WHERE job_status = 'failed'::provisioner_job_status

AND created_at >= @lookback::timestamptz

GROUP BY preset_id

)

SELECT

tsb.template_version_id,

tsb.preset_id,

COALESCE(fc.num_failed, 0)::int AS num_failed,

MAX(tsb.created_at)::timestamptz AS last_build_at

FROM time_sorted_builds tsb

LEFT JOIN failed_count fc ON fc.preset_id = tsb.preset_id

WHERE tsb.rn <= tsb.desired_instances -- Fetch the last N builds, where N is the number of desired instances; if any fail, we backoff

AND tsb.job_status = 'failed'::provisioner_job_status

AND created_at >= @lookback::timestamptz

GROUP BY tsb.template_version_id, tsb.preset_id, fc.num_failed;

-- GetPresetsAtFailureLimit groups workspace builds by preset ID.

-- Each preset is associated with exactly one template version ID.

-- For each preset, the query checks the last hard_limit builds.

-- If all of them failed, the preset is considered to have hit the hard failure limit.

-- The query returns a list of preset IDs that have reached this failure threshold.

-- Only active template versions with configured presets are considered.

-- name: GetPresetsAtFailureLimit :many

WITH filtered_builds AS (

-- Only select builds which are for prebuild creations

SELECT wlb.template_version_id, wlb.created_at, tvp.id AS preset_id, wlb.job_status, tvp.desired_instances

FROM template_version_presets tvp

INNER JOIN workspace_latest_builds wlb ON wlb.template_version_preset_id = tvp.id

INNER JOIN workspaces w ON wlb.workspace_id = w.id

INNER JOIN template_versions tv ON wlb.template_version_id = tv.id

INNER JOIN templates t ON tv.template_id = t.id AND t.active_version_id = tv.id

WHERE tvp.desired_instances IS NOT NULL -- Consider only presets that have a prebuild configuration.

AND wlb.transition = 'start'::workspace_transition

AND w.owner_id = 'c42fdf75-3097-471c-8c33-fb52454d81c0'

),

time_sorted_builds AS (

-- Group builds by preset, then sort each group by created_at.

SELECT fb.template_version_id, fb.created_at, fb.preset_id, fb.job_status, fb.desired_instances,

ROW_NUMBER() OVER (PARTITION BY fb.preset_id ORDER BY fb.created_at DESC) as rn

FROM filtered_builds fb

)

SELECT

tsb.template_version_id,

tsb.preset_id

FROM time_sorted_builds tsb

-- For each preset, check the last hard_limit builds.

-- If all of them failed, the preset is considered to have hit the hard failure limit.

WHERE tsb.rn <= @hard_limit::bigint

AND tsb.job_status = 'failed'::provisioner_job_status

GROUP BY tsb.template_version_id, tsb.preset_id

HAVING COUNT(*) = @hard_limit::bigint;

-- name: GetPrebuildMetrics :many

SELECT

t.name as template_name,

tvp.name as preset_name,

o.name as organization_name,

COUNT(*) as created_count,

COUNT(*) FILTER (WHERE pj.job_status = 'failed'::provisioner_job_status) as failed_count,

COUNT(*) FILTER (

WHERE w.owner_id != 'c42fdf75-3097-471c-8c33-fb52454d81c0'::uuid -- The system user responsible for prebuilds.

) as claimed_count

FROM workspaces w

INNER JOIN workspace_prebuild_builds wpb ON wpb.workspace_id = w.id

INNER JOIN templates t ON t.id = w.template_id

INNER JOIN template_version_presets tvp ON tvp.id = wpb.template_version_preset_id

INNER JOIN provisioner_jobs pj ON pj.id = wpb.job_id

INNER JOIN organizations o ON o.id = w.organization_id

WHERE NOT t.deleted AND wpb.build_number = 1

GROUP BY t.name, tvp.name, o.name

ORDER BY t.name, tvp.name, o.name;

-- name: FindMatchingPresetID :one

-- FindMatchingPresetID finds a preset ID that is the largest exact subset of the provided parameters.

-- It returns the preset ID if a match is found, or NULL if no match is found.

-- The query finds presets where all preset parameters are present in the provided parameters,

-- and returns the preset with the most parameters (largest subset).

WITH provided_params AS (

SELECT

unnest(@parameter_names::text[]) AS name,

unnest(@parameter_values::text[]) AS value

),

preset_matches AS (

SELECT

tvp.id AS template_version_preset_id,

COALESCE(COUNT(tvpp.name), 0) AS total_preset_params,

COALESCE(COUNT(pp.name), 0) AS matching_params

FROM template_version_presets tvp

LEFT JOIN template_version_preset_parameters tvpp ON tvpp.template_version_preset_id = tvp.id

LEFT JOIN provided_params pp ON pp.name = tvpp.name AND pp.value = tvpp.value

WHERE tvp.template_version_id = @template_version_id

GROUP BY tvp.id

)

SELECT pm.template_version_preset_id

FROM preset_matches pm

WHERE pm.total_preset_params = pm.matching_params -- All preset parameters must match

ORDER BY pm.total_preset_params DESC -- Return the preset with the most parameters

LIMIT 1;

-- name: CountPendingNonActivePrebuilds :many

-- CountPendingNonActivePrebuilds returns the number of pending prebuilds for non-active template versions

SELECT

wpb.template_version_preset_id AS preset_id,

COUNT(*)::int AS count

FROM workspace_prebuild_builds wpb

INNER JOIN provisioner_jobs pj ON pj.id = wpb.job_id

INNER JOIN workspaces w ON w.id = wpb.workspace_id

INNER JOIN templates t ON t.id = w.template_id

WHERE

wpb.template_version_id != t.active_version_id

-- Only considers initial builds, i.e. created by the reconciliation loop

AND wpb.build_number = 1

-- Only consider 'start' transitions (provisioning), not 'stop'/'delete' (deprovisioning)

-- Deprovisioning jobs should complete naturally as they're already cleaning up resources

AND wpb.transition = 'start'::workspace_transition

-- Pending jobs that have not yet been picked up by a provisioner

AND pj.job_status = 'pending'::provisioner_job_status

AND pj.worker_id IS NULL

AND pj.canceled_at IS NULL

AND pj.completed_at IS NULL

GROUP BY wpb.template_version_preset_id;

-- name: UpdatePrebuildProvisionerJobWithCancel :many

-- Cancels all pending provisioner jobs for prebuilt workspaces on a specific preset from an

-- inactive template version.

-- This is an optimization to clean up stale pending jobs.

WITH jobs_to_cancel AS (

SELECT pj.id, w.id AS workspace_id, w.template_id, wpb.template_version_preset_id

FROM provisioner_jobs pj

INNER JOIN workspace_prebuild_builds wpb ON wpb.job_id = pj.id

INNER JOIN workspaces w ON w.id = wpb.workspace_id

INNER JOIN templates t ON t.id = w.template_id

WHERE

wpb.template_version_id != t.active_version_id

AND wpb.template_version_preset_id = @preset_id

-- Only considers initial builds, i.e. created by the reconciliation loop

AND wpb.build_number = 1

-- Only consider 'start' transitions (provisioning), not 'stop'/'delete' (deprovisioning)

-- Deprovisioning jobs should complete naturally as they're already cleaning up resources

AND wpb.transition = 'start'::workspace_transition

-- Pending jobs that have not yet been picked up by a provisioner

AND pj.job_status = 'pending'::provisioner_job_status

AND pj.worker_id IS NULL

AND pj.canceled_at IS NULL

AND pj.completed_at IS NULL

)

UPDATE provisioner_jobs

SET

canceled_at = @now::timestamptz,

completed_at = @now::timestamptz

FROM jobs_to_cancel

WHERE provisioner_jobs.id = jobs_to_cancel.id

RETURNING jobs_to_cancel.id, jobs_to_cancel.workspace_id, jobs_to_cancel.template_id, jobs_to_cancel.template_version_preset_id;

-- name: GetOrganizationsWithPrebuildStatus :many

-- GetOrganizationsWithPrebuildStatus returns organizations with prebuilds configured and their

-- membership status for the prebuilds system user (org membership, group existence, group membership).

WITH orgs_with_prebuilds AS (

-- Get unique organizations that have presets with prebuilds configured

SELECT DISTINCT o.id, o.name

FROM organizations o

INNER JOIN templates t ON t.organization_id = o.id

INNER JOIN template_versions tv ON tv.template_id = t.id

INNER JOIN template_version_presets tvp ON tvp.template_version_id = tv.id

WHERE tvp.desired_instances IS NOT NULL

),

prebuild_user_membership AS (

-- Check if the user is a member of the organizations

SELECT om.organization_id

FROM organization_members om

INNER JOIN orgs_with_prebuilds owp ON owp.id = om.organization_id

WHERE om.user_id = @user_id::uuid

),

prebuild_groups AS (

-- Check if the organizations have the prebuilds group

SELECT g.organization_id, g.id as group_id

FROM groups g

INNER JOIN orgs_with_prebuilds owp ON owp.id = g.organization_id

WHERE g.name = @group_name::text

),

prebuild_group_membership AS (

-- Check if the user is in the prebuilds group

SELECT pg.organization_id

FROM prebuild_groups pg

INNER JOIN group_members gm ON gm.group_id = pg.group_id

WHERE gm.user_id = @user_id::uuid

)

SELECT

owp.id AS organization_id,

owp.name AS organization_name,

(pum.organization_id IS NOT NULL)::boolean AS has_prebuild_user,

pg.group_id AS prebuilds_group_id,

(pgm.organization_id IS NOT NULL)::boolean AS has_prebuild_user_in_group

FROM orgs_with_prebuilds owp

LEFT JOIN prebuild_groups pg ON pg.organization_id = owp.id

LEFT JOIN prebuild_user_membership pum ON pum.organization_id = owp.id

LEFT JOIN prebuild_group_membership pgm ON pgm.organization_id = owp.id;