Decompose GetTemplatePrebuildState into separate queries, reimplement logic in Go

This is in service of testability

Signed-off-by: Danny Kopping <danny@coder.com>
This commit is contained in:
Danny Kopping
2025-02-20 14:57:01 +00:00
parent 64d476545b
commit 4d97580666
10 changed files with 485 additions and 384 deletions

View File

@ -1977,6 +1977,13 @@ func (q *querier) GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUI
return q.db.GetParameterSchemasByJobID(ctx, jobID)
}
func (q *querier) GetPrebuildsInProgress(ctx context.Context) ([]database.GetPrebuildsInProgressRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceTemplate); err != nil {
return nil, err
}
return q.db.GetPrebuildsInProgress(ctx)
}
func (q *querier) GetPresetByWorkspaceBuildID(ctx context.Context, workspaceID uuid.UUID) (database.TemplateVersionPreset, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceTemplate); err != nil {
return database.TemplateVersionPreset{}, err
@ -2144,6 +2151,13 @@ func (q *querier) GetReplicasUpdatedAfter(ctx context.Context, updatedAt time.Ti
return q.db.GetReplicasUpdatedAfter(ctx, updatedAt)
}
func (q *querier) GetRunningPrebuilds(ctx context.Context) ([]database.GetRunningPrebuildsRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceTemplate); err != nil {
return nil, err
}
return q.db.GetRunningPrebuilds(ctx)
}
func (q *querier) GetRuntimeConfig(ctx context.Context, key string) (string, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
return "", err
@ -2268,12 +2282,11 @@ func (q *querier) GetTemplateParameterInsights(ctx context.Context, arg database
return q.db.GetTemplateParameterInsights(ctx, arg)
}
func (q *querier) GetTemplatePrebuildState(ctx context.Context, templateID uuid.UUID) ([]database.GetTemplatePrebuildStateRow, error) {
// TODO: authz
func (q *querier) GetTemplatePresetsWithPrebuilds(ctx context.Context, templateID uuid.NullUUID) ([]database.GetTemplatePresetsWithPrebuildsRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceTemplate); err != nil {
return nil, err
}
return q.db.GetTemplatePrebuildState(ctx, templateID)
return q.db.GetTemplatePresetsWithPrebuilds(ctx, templateID)
}
func (q *querier) GetTemplateUsageStats(ctx context.Context, arg database.GetTemplateUsageStatsParams) ([]database.TemplateUsageStat, error) {

View File

@ -3784,6 +3784,10 @@ func (q *FakeQuerier) GetParameterSchemasByJobID(_ context.Context, jobID uuid.U
return parameters, nil
}
func (q *FakeQuerier) GetPrebuildsInProgress(ctx context.Context) ([]database.GetPrebuildsInProgressRow, error) {
panic("not implemented")
}
func (q *FakeQuerier) GetPresetByWorkspaceBuildID(_ context.Context, workspaceBuildID uuid.UUID) (database.TemplateVersionPreset, error) {
q.mutex.RLock()
defer q.mutex.RUnlock()
@ -4489,6 +4493,10 @@ func (q *FakeQuerier) GetReplicasUpdatedAfter(_ context.Context, updatedAt time.
return replicas, nil
}
func (q *FakeQuerier) GetRunningPrebuilds(ctx context.Context) ([]database.GetRunningPrebuildsRow, error) {
panic("not implemented")
}
func (q *FakeQuerier) GetRuntimeConfig(_ context.Context, key string) (string, error) {
q.mutex.Lock()
defer q.mutex.Unlock()
@ -5528,7 +5536,7 @@ func (q *FakeQuerier) GetTemplateParameterInsights(ctx context.Context, arg data
return rows, nil
}
func (q *FakeQuerier) GetTemplatePrebuildState(ctx context.Context, templateID uuid.UUID) ([]database.GetTemplatePrebuildStateRow, error) {
func (q *FakeQuerier) GetTemplatePresetsWithPrebuilds(ctx context.Context, templateID uuid.NullUUID) ([]database.GetTemplatePresetsWithPrebuildsRow, error) {
panic("not implemented")
}

View File

@ -987,6 +987,13 @@ func (m queryMetricsStore) GetParameterSchemasByJobID(ctx context.Context, jobID
return schemas, err
}
func (m queryMetricsStore) GetPrebuildsInProgress(ctx context.Context) ([]database.GetPrebuildsInProgressRow, error) {
start := time.Now()
r0, r1 := m.s.GetPrebuildsInProgress(ctx)
m.queryLatencies.WithLabelValues("GetPrebuildsInProgress").Observe(time.Since(start).Seconds())
return r0, r1
}
func (m queryMetricsStore) GetPresetByWorkspaceBuildID(ctx context.Context, workspaceBuildID uuid.UUID) (database.TemplateVersionPreset, error) {
start := time.Now()
r0, r1 := m.s.GetPresetByWorkspaceBuildID(ctx, workspaceBuildID)
@ -1134,6 +1141,13 @@ func (m queryMetricsStore) GetReplicasUpdatedAfter(ctx context.Context, updatedA
return replicas, err
}
func (m queryMetricsStore) GetRunningPrebuilds(ctx context.Context) ([]database.GetRunningPrebuildsRow, error) {
start := time.Now()
r0, r1 := m.s.GetRunningPrebuilds(ctx)
m.queryLatencies.WithLabelValues("GetRunningPrebuilds").Observe(time.Since(start).Seconds())
return r0, r1
}
func (m queryMetricsStore) GetRuntimeConfig(ctx context.Context, key string) (string, error) {
start := time.Now()
r0, r1 := m.s.GetRuntimeConfig(ctx, key)
@ -1260,10 +1274,10 @@ func (m queryMetricsStore) GetTemplateParameterInsights(ctx context.Context, arg
return r0, r1
}
func (m queryMetricsStore) GetTemplatePrebuildState(ctx context.Context, templateID uuid.UUID) ([]database.GetTemplatePrebuildStateRow, error) {
func (m queryMetricsStore) GetTemplatePresetsWithPrebuilds(ctx context.Context, templateID uuid.NullUUID) ([]database.GetTemplatePresetsWithPrebuildsRow, error) {
start := time.Now()
r0, r1 := m.s.GetTemplatePrebuildState(ctx, templateID)
m.queryLatencies.WithLabelValues("GetTemplatePrebuildState").Observe(time.Since(start).Seconds())
r0, r1 := m.s.GetTemplatePresetsWithPrebuilds(ctx, templateID)
m.queryLatencies.WithLabelValues("GetTemplatePresetsWithPrebuilds").Observe(time.Since(start).Seconds())
return r0, r1
}

View File

@ -2031,6 +2031,21 @@ func (mr *MockStoreMockRecorder) GetParameterSchemasByJobID(ctx, jobID any) *gom
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetParameterSchemasByJobID", reflect.TypeOf((*MockStore)(nil).GetParameterSchemasByJobID), ctx, jobID)
}
// GetPrebuildsInProgress mocks base method.
func (m *MockStore) GetPrebuildsInProgress(ctx context.Context) ([]database.GetPrebuildsInProgressRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetPrebuildsInProgress", ctx)
ret0, _ := ret[0].([]database.GetPrebuildsInProgressRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetPrebuildsInProgress indicates an expected call of GetPrebuildsInProgress.
func (mr *MockStoreMockRecorder) GetPrebuildsInProgress(ctx any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPrebuildsInProgress", reflect.TypeOf((*MockStore)(nil).GetPrebuildsInProgress), ctx)
}
// GetPresetByWorkspaceBuildID mocks base method.
func (m *MockStore) GetPresetByWorkspaceBuildID(ctx context.Context, workspaceBuildID uuid.UUID) (database.TemplateVersionPreset, error) {
m.ctrl.T.Helper()
@ -2346,6 +2361,21 @@ func (mr *MockStoreMockRecorder) GetReplicasUpdatedAfter(ctx, updatedAt any) *go
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetReplicasUpdatedAfter", reflect.TypeOf((*MockStore)(nil).GetReplicasUpdatedAfter), ctx, updatedAt)
}
// GetRunningPrebuilds mocks base method.
func (m *MockStore) GetRunningPrebuilds(ctx context.Context) ([]database.GetRunningPrebuildsRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetRunningPrebuilds", ctx)
ret0, _ := ret[0].([]database.GetRunningPrebuildsRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetRunningPrebuilds indicates an expected call of GetRunningPrebuilds.
func (mr *MockStoreMockRecorder) GetRunningPrebuilds(ctx any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetRunningPrebuilds", reflect.TypeOf((*MockStore)(nil).GetRunningPrebuilds), ctx)
}
// GetRuntimeConfig mocks base method.
func (m *MockStore) GetRuntimeConfig(ctx context.Context, key string) (string, error) {
m.ctrl.T.Helper()
@ -2631,19 +2661,19 @@ func (mr *MockStoreMockRecorder) GetTemplateParameterInsights(ctx, arg any) *gom
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTemplateParameterInsights", reflect.TypeOf((*MockStore)(nil).GetTemplateParameterInsights), ctx, arg)
}
// GetTemplatePrebuildState mocks base method.
func (m *MockStore) GetTemplatePrebuildState(ctx context.Context, templateID uuid.UUID) ([]database.GetTemplatePrebuildStateRow, error) {
// GetTemplatePresetsWithPrebuilds mocks base method.
func (m *MockStore) GetTemplatePresetsWithPrebuilds(ctx context.Context, templateID uuid.NullUUID) ([]database.GetTemplatePresetsWithPrebuildsRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetTemplatePrebuildState", ctx, templateID)
ret0, _ := ret[0].([]database.GetTemplatePrebuildStateRow)
ret := m.ctrl.Call(m, "GetTemplatePresetsWithPrebuilds", ctx, templateID)
ret0, _ := ret[0].([]database.GetTemplatePresetsWithPrebuildsRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetTemplatePrebuildState indicates an expected call of GetTemplatePrebuildState.
func (mr *MockStoreMockRecorder) GetTemplatePrebuildState(ctx, templateID any) *gomock.Call {
// GetTemplatePresetsWithPrebuilds indicates an expected call of GetTemplatePresetsWithPrebuilds.
func (mr *MockStoreMockRecorder) GetTemplatePresetsWithPrebuilds(ctx, templateID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTemplatePrebuildState", reflect.TypeOf((*MockStore)(nil).GetTemplatePrebuildState), ctx, templateID)
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTemplatePresetsWithPrebuilds", reflect.TypeOf((*MockStore)(nil).GetTemplatePresetsWithPrebuilds), ctx, templateID)
}
// GetTemplateUsageStats mocks base method.

View File

@ -204,6 +204,7 @@ type sqlcQuerier interface {
GetOrganizations(ctx context.Context, arg GetOrganizationsParams) ([]Organization, error)
GetOrganizationsByUserID(ctx context.Context, userID uuid.UUID) ([]Organization, error)
GetParameterSchemasByJobID(ctx context.Context, jobID uuid.UUID) ([]ParameterSchema, error)
GetPrebuildsInProgress(ctx context.Context) ([]GetPrebuildsInProgressRow, error)
GetPresetByWorkspaceBuildID(ctx context.Context, workspaceBuildID uuid.UUID) (TemplateVersionPreset, error)
GetPresetParametersByTemplateVersionID(ctx context.Context, templateVersionID uuid.UUID) ([]TemplateVersionPresetParameter, error)
GetPresetsByTemplateVersionID(ctx context.Context, templateVersionID uuid.UUID) ([]TemplateVersionPreset, error)
@ -227,6 +228,7 @@ type sqlcQuerier interface {
GetQuotaConsumedForUser(ctx context.Context, arg GetQuotaConsumedForUserParams) (int64, error)
GetReplicaByID(ctx context.Context, id uuid.UUID) (Replica, error)
GetReplicasUpdatedAfter(ctx context.Context, updatedAt time.Time) ([]Replica, error)
GetRunningPrebuilds(ctx context.Context) ([]GetRunningPrebuildsRow, error)
GetRuntimeConfig(ctx context.Context, key string) (string, error)
GetTailnetAgents(ctx context.Context, id uuid.UUID) ([]TailnetAgent, error)
GetTailnetClientsForAgent(ctx context.Context, agentID uuid.UUID) ([]TailnetClient, error)
@ -269,7 +271,7 @@ type sqlcQuerier interface {
// created in the timeframe and return the aggregate usage counts of parameter
// values.
GetTemplateParameterInsights(ctx context.Context, arg GetTemplateParameterInsightsParams) ([]GetTemplateParameterInsightsRow, error)
GetTemplatePrebuildState(ctx context.Context, templateID uuid.UUID) ([]GetTemplatePrebuildStateRow, error)
GetTemplatePresetsWithPrebuilds(ctx context.Context, templateID uuid.NullUUID) ([]GetTemplatePresetsWithPrebuildsRow, error)
GetTemplateUsageStats(ctx context.Context, arg GetTemplateUsageStatsParams) ([]TemplateUsageStat, error)
GetTemplateVersionByID(ctx context.Context, id uuid.UUID) (TemplateVersion, error)
GetTemplateVersionByJobID(ctx context.Context, jobID uuid.UUID) (TemplateVersion, error)

View File

@ -5442,148 +5442,155 @@ func (q *sqlQuerier) ClaimPrebuild(ctx context.Context, arg ClaimPrebuildParams)
return i, err
}
const getTemplatePrebuildState = `-- name: GetTemplatePrebuildState :many
WITH
-- All prebuilds currently running
running_prebuilds AS (SELECT p.template_id,
b.template_version_id,
tvp_curr.id AS current_preset_id,
tvp_desired.id AS desired_preset_id,
COUNT(*) AS count,
SUM(CASE
WHEN p.lifecycle_state = 'ready'::workspace_agent_lifecycle_state THEN 1
ELSE 0 END) AS eligible,
STRING_AGG(p.id::text, ',') AS ids
FROM workspace_prebuilds p
INNER JOIN workspace_latest_build b ON b.workspace_id = p.id
INNER JOIN provisioner_jobs pj ON b.job_id = pj.id
INNER JOIN templates t ON p.template_id = t.id
LEFT JOIN template_version_presets tvp_curr
ON tvp_curr.id = b.template_version_preset_id
LEFT JOIN template_version_presets tvp_desired
ON tvp_desired.template_version_id = t.active_version_id
WHERE (b.transition = 'start'::workspace_transition
-- if a deletion job fails, the workspace will still be running
OR pj.job_status IN ('failed'::provisioner_job_status, 'canceled'::provisioner_job_status,
'unknown'::provisioner_job_status))
AND (tvp_curr.name = tvp_desired.name
OR tvp_desired.id IS NULL)
GROUP BY p.template_id, b.template_version_id, tvp_curr.id,
tvp_desired.id),
-- All templates which have been configured for prebuilds (any version)
templates_with_prebuilds AS (SELECT t.id AS template_id,
tv.id AS template_version_id,
tv.id = t.active_version_id AS using_active_version,
tvpp.preset_id,
tvp.name,
MAX(tvpp.desired_instances) AS desired_instances,
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 template_version_preset_prebuilds tvpp ON tvpp.preset_id = tvp.id
WHERE t.id = $1::uuid
GROUP BY t.id, tv.id, tvpp.preset_id, tvp.name),
-- Jobs relating to prebuilds current in-flight
prebuilds_in_progress AS (SELECT wpb.template_version_id, wpb.transition, COUNT(wpb.transition) AS count
FROM workspace_latest_build wlb
INNER JOIN provisioner_jobs pj ON wlb.job_id = pj.id
INNER JOIN workspace_prebuild_builds wpb ON wpb.id = wlb.id
WHERE pj.job_status NOT IN
('succeeded'::provisioner_job_status, 'canceled'::provisioner_job_status,
'failed'::provisioner_job_status)
GROUP BY wpb.template_version_id, wpb.transition)
SELECT t.template_id,
t.template_version_id,
t.preset_id,
t.using_active_version AS is_active,
MAX(CASE
WHEN p.template_version_id = t.template_version_id THEN p.ids
ELSE '' END)::text AS running_prebuild_ids,
COALESCE(MAX(CASE WHEN t.using_active_version THEN p.count ELSE 0 END),
0)::int AS actual, -- running prebuilds for active version
COALESCE(MAX(CASE WHEN t.using_active_version THEN p.eligible ELSE 0 END),
0)::int AS eligible, -- prebuilds which can be claimed
MAX(CASE WHEN t.using_active_version THEN t.desired_instances ELSE 0 END)::int AS desired, -- we only care about the active version's desired instances
COALESCE(MAX(CASE
WHEN p.template_version_id = t.template_version_id AND
t.using_active_version = false
THEN p.count
ELSE 0 END),
0)::int AS outdated, -- running prebuilds for inactive version
COALESCE(GREATEST(
(MAX(CASE WHEN t.using_active_version THEN p.count ELSE 0 END)::int
-
MAX(CASE WHEN t.using_active_version THEN t.desired_instances ELSE 0 END)),
0),
0) ::int AS extraneous, -- extra running prebuilds for active version
COALESCE(MAX(CASE
WHEN pip.transition = 'start'::workspace_transition THEN pip.count
ELSE 0 END),
0)::int AS starting,
COALESCE(MAX(CASE
WHEN pip.transition = 'stop'::workspace_transition THEN pip.count
ELSE 0 END),
0)::int AS stopping, -- not strictly needed, since prebuilds should never be left if a "stopped" state, but useful to know
COALESCE(MAX(CASE
WHEN pip.transition = 'delete'::workspace_transition THEN pip.count
ELSE 0 END),
0)::int AS deleting,
t.deleted::bool AS template_deleted,
t.deprecated::bool AS template_deprecated
FROM templates_with_prebuilds t
LEFT JOIN running_prebuilds p
ON (p.template_version_id = t.template_version_id AND p.current_preset_id = t.preset_id)
LEFT JOIN prebuilds_in_progress pip ON pip.template_version_id = t.template_version_id
WHERE (t.using_active_version = TRUE
OR p.count > 0)
GROUP BY t.template_id, t.template_version_id, t.preset_id, t.using_active_version, t.deleted, t.deprecated
const getPrebuildsInProgress = `-- name: GetPrebuildsInProgress :many
SELECT wpb.template_version_id, wpb.transition, COUNT(wpb.transition) AS count
FROM workspace_latest_build wlb
INNER JOIN provisioner_jobs pj ON wlb.job_id = pj.id
INNER JOIN workspace_prebuild_builds wpb ON wpb.id = wlb.id
WHERE pj.job_status NOT IN
('succeeded'::provisioner_job_status, 'canceled'::provisioner_job_status,
'failed'::provisioner_job_status)
GROUP BY wpb.template_version_id, wpb.transition
`
type GetTemplatePrebuildStateRow struct {
TemplateID uuid.UUID `db:"template_id" json:"template_id"`
TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"`
PresetID uuid.UUID `db:"preset_id" json:"preset_id"`
IsActive bool `db:"is_active" json:"is_active"`
RunningPrebuildIds string `db:"running_prebuild_ids" json:"running_prebuild_ids"`
Actual int32 `db:"actual" json:"actual"`
Eligible int32 `db:"eligible" json:"eligible"`
Desired int32 `db:"desired" json:"desired"`
Outdated int32 `db:"outdated" json:"outdated"`
Extraneous int32 `db:"extraneous" json:"extraneous"`
Starting int32 `db:"starting" json:"starting"`
Stopping int32 `db:"stopping" json:"stopping"`
Deleting int32 `db:"deleting" json:"deleting"`
TemplateDeleted bool `db:"template_deleted" json:"template_deleted"`
TemplateDeprecated bool `db:"template_deprecated" json:"template_deprecated"`
type GetPrebuildsInProgressRow struct {
TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"`
Transition WorkspaceTransition `db:"transition" json:"transition"`
Count int64 `db:"count" json:"count"`
}
func (q *sqlQuerier) GetTemplatePrebuildState(ctx context.Context, templateID uuid.UUID) ([]GetTemplatePrebuildStateRow, error) {
rows, err := q.db.QueryContext(ctx, getTemplatePrebuildState, templateID)
func (q *sqlQuerier) GetPrebuildsInProgress(ctx context.Context) ([]GetPrebuildsInProgressRow, error) {
rows, err := q.db.QueryContext(ctx, getPrebuildsInProgress)
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetTemplatePrebuildStateRow
var items []GetPrebuildsInProgressRow
for rows.Next() {
var i GetTemplatePrebuildStateRow
var i GetPrebuildsInProgressRow
if err := rows.Scan(&i.TemplateVersionID, &i.Transition, &i.Count); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const getRunningPrebuilds = `-- name: GetRunningPrebuilds :many
SELECT p.id AS workspace_id,
p.template_id,
b.template_version_id,
tvp_curr.id AS current_preset_id,
tvp_desired.id AS desired_preset_id,
CASE
WHEN p.lifecycle_state = 'ready'::workspace_agent_lifecycle_state THEN TRUE
ELSE FALSE END AS eligible
FROM workspace_prebuilds p
INNER JOIN workspace_latest_build b ON b.workspace_id = p.id
INNER JOIN provisioner_jobs pj ON b.job_id = pj.id
INNER JOIN templates t ON p.template_id = t.id
LEFT JOIN template_version_presets tvp_curr
ON tvp_curr.id = b.template_version_preset_id
LEFT JOIN template_version_presets tvp_desired
ON tvp_desired.template_version_id = t.active_version_id
WHERE (b.transition = 'start'::workspace_transition
-- if a deletion job fails, the workspace will still be running
OR pj.job_status IN ('failed'::provisioner_job_status, 'canceled'::provisioner_job_status,
'unknown'::provisioner_job_status))
AND (tvp_curr.name = tvp_desired.name
OR tvp_desired.id IS NULL)
`
type GetRunningPrebuildsRow struct {
WorkspaceID uuid.UUID `db:"workspace_id" json:"workspace_id"`
TemplateID uuid.UUID `db:"template_id" json:"template_id"`
TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"`
CurrentPresetID uuid.NullUUID `db:"current_preset_id" json:"current_preset_id"`
DesiredPresetID uuid.NullUUID `db:"desired_preset_id" json:"desired_preset_id"`
Eligible bool `db:"eligible" json:"eligible"`
}
func (q *sqlQuerier) GetRunningPrebuilds(ctx context.Context) ([]GetRunningPrebuildsRow, error) {
rows, err := q.db.QueryContext(ctx, getRunningPrebuilds)
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetRunningPrebuildsRow
for rows.Next() {
var i GetRunningPrebuildsRow
if err := rows.Scan(
&i.WorkspaceID,
&i.TemplateID,
&i.TemplateVersionID,
&i.CurrentPresetID,
&i.DesiredPresetID,
&i.Eligible,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const getTemplatePresetsWithPrebuilds = `-- name: GetTemplatePresetsWithPrebuilds :many
SELECT t.id AS template_id,
tv.id AS template_version_id,
tv.id = t.active_version_id AS using_active_version,
tvpp.preset_id,
tvp.name,
tvpp.desired_instances AS desired_instances,
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 template_version_preset_prebuilds tvpp ON tvpp.preset_id = tvp.id
WHERE (t.id = $1::uuid OR $1 IS NULL)
`
type GetTemplatePresetsWithPrebuildsRow struct {
TemplateID uuid.UUID `db:"template_id" json:"template_id"`
TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"`
UsingActiveVersion bool `db:"using_active_version" json:"using_active_version"`
PresetID uuid.UUID `db:"preset_id" json:"preset_id"`
Name string `db:"name" json:"name"`
DesiredInstances int32 `db:"desired_instances" json:"desired_instances"`
Deleted bool `db:"deleted" json:"deleted"`
Deprecated bool `db:"deprecated" json:"deprecated"`
}
func (q *sqlQuerier) GetTemplatePresetsWithPrebuilds(ctx context.Context, templateID uuid.NullUUID) ([]GetTemplatePresetsWithPrebuildsRow, error) {
rows, err := q.db.QueryContext(ctx, getTemplatePresetsWithPrebuilds, templateID)
if err != nil {
return nil, err
}
defer rows.Close()
var items []GetTemplatePresetsWithPrebuildsRow
for rows.Next() {
var i GetTemplatePresetsWithPrebuildsRow
if err := rows.Scan(
&i.TemplateID,
&i.TemplateVersionID,
&i.UsingActiveVersion,
&i.PresetID,
&i.IsActive,
&i.RunningPrebuildIds,
&i.Actual,
&i.Eligible,
&i.Desired,
&i.Outdated,
&i.Extraneous,
&i.Starting,
&i.Stopping,
&i.Deleting,
&i.TemplateDeleted,
&i.TemplateDeprecated,
&i.Name,
&i.DesiredInstances,
&i.Deleted,
&i.Deprecated,
); err != nil {
return nil, err
}

View File

@ -1,100 +1,51 @@
-- name: GetTemplatePrebuildState :many
WITH
-- All prebuilds currently running
running_prebuilds AS (SELECT p.template_id,
b.template_version_id,
tvp_curr.id AS current_preset_id,
tvp_desired.id AS desired_preset_id,
COUNT(*) AS count,
SUM(CASE
WHEN p.lifecycle_state = 'ready'::workspace_agent_lifecycle_state THEN 1
ELSE 0 END) AS eligible,
STRING_AGG(p.id::text, ',') AS ids
FROM workspace_prebuilds p
INNER JOIN workspace_latest_build b ON b.workspace_id = p.id
INNER JOIN provisioner_jobs pj ON b.job_id = pj.id
INNER JOIN templates t ON p.template_id = t.id
LEFT JOIN template_version_presets tvp_curr
ON tvp_curr.id = b.template_version_preset_id
LEFT JOIN template_version_presets tvp_desired
ON tvp_desired.template_version_id = t.active_version_id
WHERE (b.transition = 'start'::workspace_transition
-- if a deletion job fails, the workspace will still be running
OR pj.job_status IN ('failed'::provisioner_job_status, 'canceled'::provisioner_job_status,
'unknown'::provisioner_job_status))
AND (tvp_curr.name = tvp_desired.name
OR tvp_desired.id IS NULL)
GROUP BY p.template_id, b.template_version_id, tvp_curr.id,
tvp_desired.id),
-- All templates which have been configured for prebuilds (any version)
templates_with_prebuilds AS (SELECT t.id AS template_id,
tv.id AS template_version_id,
tv.id = t.active_version_id AS using_active_version,
tvpp.preset_id,
tvp.name,
MAX(tvpp.desired_instances) AS desired_instances,
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 template_version_preset_prebuilds tvpp ON tvpp.preset_id = tvp.id
WHERE t.id = @template_id::uuid
GROUP BY t.id, tv.id, tvpp.preset_id, tvp.name),
-- Jobs relating to prebuilds current in-flight
prebuilds_in_progress AS (SELECT wpb.template_version_id, wpb.transition, COUNT(wpb.transition) AS count
FROM workspace_latest_build wlb
INNER JOIN provisioner_jobs pj ON wlb.job_id = pj.id
INNER JOIN workspace_prebuild_builds wpb ON wpb.id = wlb.id
WHERE pj.job_status NOT IN
('succeeded'::provisioner_job_status, 'canceled'::provisioner_job_status,
'failed'::provisioner_job_status)
GROUP BY wpb.template_version_id, wpb.transition)
SELECT t.template_id,
t.template_version_id,
t.preset_id,
t.using_active_version AS is_active,
MAX(CASE
WHEN p.template_version_id = t.template_version_id THEN p.ids
ELSE '' END)::text AS running_prebuild_ids,
COALESCE(MAX(CASE WHEN t.using_active_version THEN p.count ELSE 0 END),
0)::int AS actual, -- running prebuilds for active version
COALESCE(MAX(CASE WHEN t.using_active_version THEN p.eligible ELSE 0 END),
0)::int AS eligible, -- prebuilds which can be claimed
MAX(CASE WHEN t.using_active_version THEN t.desired_instances ELSE 0 END)::int AS desired, -- we only care about the active version's desired instances
COALESCE(MAX(CASE
WHEN p.template_version_id = t.template_version_id AND
t.using_active_version = false
THEN p.count
ELSE 0 END),
0)::int AS outdated, -- running prebuilds for inactive version
COALESCE(GREATEST(
(MAX(CASE WHEN t.using_active_version THEN p.count ELSE 0 END)::int
-
MAX(CASE WHEN t.using_active_version THEN t.desired_instances ELSE 0 END)),
0),
0) ::int AS extraneous, -- extra running prebuilds for active version
COALESCE(MAX(CASE
WHEN pip.transition = 'start'::workspace_transition THEN pip.count
ELSE 0 END),
0)::int AS starting,
COALESCE(MAX(CASE
WHEN pip.transition = 'stop'::workspace_transition THEN pip.count
ELSE 0 END),
0)::int AS stopping, -- not strictly needed, since prebuilds should never be left if a "stopped" state, but useful to know
COALESCE(MAX(CASE
WHEN pip.transition = 'delete'::workspace_transition THEN pip.count
ELSE 0 END),
0)::int AS deleting,
t.deleted::bool AS template_deleted,
t.deprecated::bool AS template_deprecated
FROM templates_with_prebuilds t
LEFT JOIN running_prebuilds p
ON (p.template_version_id = t.template_version_id AND p.current_preset_id = t.preset_id)
LEFT JOIN prebuilds_in_progress pip ON pip.template_version_id = t.template_version_id
WHERE (t.using_active_version = TRUE
OR p.count > 0)
GROUP BY t.template_id, t.template_version_id, t.preset_id, t.using_active_version, t.deleted, t.deprecated;
-- name: GetRunningPrebuilds :many
SELECT p.id AS workspace_id,
p.template_id,
b.template_version_id,
tvp_curr.id AS current_preset_id,
tvp_desired.id AS desired_preset_id,
CASE
WHEN p.lifecycle_state = 'ready'::workspace_agent_lifecycle_state THEN TRUE
ELSE FALSE END AS eligible
FROM workspace_prebuilds p
INNER JOIN workspace_latest_build b ON b.workspace_id = p.id
INNER JOIN provisioner_jobs pj ON b.job_id = pj.id
INNER JOIN templates t ON p.template_id = t.id
LEFT JOIN template_version_presets tvp_curr
ON tvp_curr.id = b.template_version_preset_id
LEFT JOIN template_version_presets tvp_desired
ON tvp_desired.template_version_id = t.active_version_id
WHERE (b.transition = 'start'::workspace_transition
-- if a deletion job fails, the workspace will still be running
OR pj.job_status IN ('failed'::provisioner_job_status, 'canceled'::provisioner_job_status,
'unknown'::provisioner_job_status))
AND (tvp_curr.name = tvp_desired.name
OR tvp_desired.id IS NULL);
-- name: GetTemplatePresetsWithPrebuilds :many
SELECT t.id AS template_id,
tv.id AS template_version_id,
tv.id = t.active_version_id AS using_active_version,
tvpp.preset_id,
tvp.name,
tvpp.desired_instances AS desired_instances,
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 template_version_preset_prebuilds tvpp ON tvpp.preset_id = tvp.id
WHERE (t.id = sqlc.narg('template_id')::uuid OR sqlc.narg('template_id') IS NULL);
-- name: GetPrebuildsInProgress :many
SELECT wpb.template_version_id, wpb.transition, COUNT(wpb.transition) AS count
FROM workspace_latest_build wlb
INNER JOIN provisioner_jobs pj ON wlb.job_id = pj.id
INNER JOIN workspace_prebuild_builds wpb ON wpb.id = wlb.id
WHERE pj.job_status NOT IN
('succeeded'::provisioner_job_status, 'canceled'::provisioner_job_status,
'failed'::provisioner_job_status)
GROUP BY wpb.template_version_id, wpb.transition;
-- name: ClaimPrebuild :one
-- TODO: rewrite to use named CTE instead?