Files

375 lines
18 KiB
SQL

-- 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;