1
0
forked from wrenn/wrenn
Co-authored-by: Tasnim Kabir Sadik <tksadik@omukk.dev>

Reviewed-on: wrenn/wrenn#50
This commit is contained in:
2026-05-24 21:10:37 +00:00
parent 4707f16c76
commit 05ddf62399
203 changed files with 15815 additions and 9344 deletions

View File

@ -11,17 +11,24 @@ import (
"github.com/jackc/pgx/v5/pgtype"
)
const bulkRestoreRunning = `-- name: BulkRestoreRunning :exec
const bulkRestoreMissingToStatus = `-- name: BulkRestoreMissingToStatus :exec
UPDATE sandboxes
SET status = 'running',
SET status = $2,
last_updated = NOW()
WHERE id = ANY($1::uuid[]) AND status = 'missing'
`
type BulkRestoreMissingToStatusParams struct {
Column1 []pgtype.UUID `json:"column_1"`
Status string `json:"status"`
}
// Called by the reconciler when a host comes back online and its sandboxes are
// confirmed alive. Restores only sandboxes that are in 'missing' state.
func (q *Queries) BulkRestoreRunning(ctx context.Context, dollar_1 []pgtype.UUID) error {
_, err := q.db.Exec(ctx, bulkRestoreRunning, dollar_1)
// confirmed alive. Restores only sandboxes currently in 'missing' state to the
// given target status (typically 'running' or 'paused' based on the live state
// reported by the host agent's ListSandboxes RPC).
func (q *Queries) BulkRestoreMissingToStatus(ctx context.Context, arg BulkRestoreMissingToStatusParams) error {
_, err := q.db.Exec(ctx, bulkRestoreMissingToStatus, arg.Column1, arg.Status)
return err
}
@ -375,7 +382,7 @@ const markSandboxesMissingByHost = `-- name: MarkSandboxesMissingByHost :exec
UPDATE sandboxes
SET status = 'missing',
last_updated = NOW()
WHERE host_id = $1 AND status IN ('running', 'starting', 'pending')
WHERE host_id = $1 AND status IN ('running', 'starting', 'pending', 'pausing', 'resuming', 'stopping')
`
// Called when the host monitor marks a host unreachable.
@ -403,6 +410,23 @@ func (q *Queries) UpdateLastActive(ctx context.Context, arg UpdateLastActivePara
return err
}
const updateSandboxDiskSize = `-- name: UpdateSandboxDiskSize :exec
UPDATE sandboxes
SET disk_size_mb = $2,
last_updated = NOW()
WHERE id = $1
`
type UpdateSandboxDiskSizeParams struct {
ID pgtype.UUID `json:"id"`
DiskSizeMb int32 `json:"disk_size_mb"`
}
func (q *Queries) UpdateSandboxDiskSize(ctx context.Context, arg UpdateSandboxDiskSizeParams) error {
_, err := q.db.Exec(ctx, updateSandboxDiskSize, arg.ID, arg.DiskSizeMb)
return err
}
const updateSandboxMetadata = `-- name: UpdateSandboxMetadata :exec
UPDATE sandboxes
SET metadata = $2,
@ -470,6 +494,61 @@ func (q *Queries) UpdateSandboxRunning(ctx context.Context, arg UpdateSandboxRun
return i, err
}
const updateSandboxRunningIf = `-- name: UpdateSandboxRunningIf :one
UPDATE sandboxes
SET status = 'running',
host_ip = $3,
guest_ip = $4,
started_at = $5,
last_active_at = $5,
last_updated = NOW()
WHERE id = $1 AND status = $2
RETURNING id, team_id, host_id, template, status, vcpus, memory_mb, timeout_sec, disk_size_mb, guest_ip, host_ip, created_at, started_at, last_active_at, last_updated, template_id, template_team_id, metadata
`
type UpdateSandboxRunningIfParams struct {
ID pgtype.UUID `json:"id"`
Status string `json:"status"`
HostIp string `json:"host_ip"`
GuestIp string `json:"guest_ip"`
StartedAt pgtype.Timestamptz `json:"started_at"`
}
// Conditionally transition a sandbox to running only if the current status
// matches the expected value. Prevents races where a user destroys a sandbox
// while the create/resume goroutine is still in-flight.
func (q *Queries) UpdateSandboxRunningIf(ctx context.Context, arg UpdateSandboxRunningIfParams) (Sandbox, error) {
row := q.db.QueryRow(ctx, updateSandboxRunningIf,
arg.ID,
arg.Status,
arg.HostIp,
arg.GuestIp,
arg.StartedAt,
)
var i Sandbox
err := row.Scan(
&i.ID,
&i.TeamID,
&i.HostID,
&i.Template,
&i.Status,
&i.Vcpus,
&i.MemoryMb,
&i.TimeoutSec,
&i.DiskSizeMb,
&i.GuestIp,
&i.HostIp,
&i.CreatedAt,
&i.StartedAt,
&i.LastActiveAt,
&i.LastUpdated,
&i.TemplateID,
&i.TemplateTeamID,
&i.Metadata,
)
return i, err
}
const updateSandboxStatus = `-- name: UpdateSandboxStatus :one
UPDATE sandboxes
SET status = $2,
@ -508,3 +587,46 @@ func (q *Queries) UpdateSandboxStatus(ctx context.Context, arg UpdateSandboxStat
)
return i, err
}
const updateSandboxStatusIf = `-- name: UpdateSandboxStatusIf :one
UPDATE sandboxes
SET status = $3,
last_updated = NOW()
WHERE id = $1 AND status = $2
RETURNING id, team_id, host_id, template, status, vcpus, memory_mb, timeout_sec, disk_size_mb, guest_ip, host_ip, created_at, started_at, last_active_at, last_updated, template_id, template_team_id, metadata
`
type UpdateSandboxStatusIfParams struct {
ID pgtype.UUID `json:"id"`
Status string `json:"status"`
Status_2 string `json:"status_2"`
}
// Atomically update status only when the current status matches the expected value.
// Prevents background goroutines from overwriting a status that has since changed
// (e.g. user destroyed a sandbox while Create was in-flight).
func (q *Queries) UpdateSandboxStatusIf(ctx context.Context, arg UpdateSandboxStatusIfParams) (Sandbox, error) {
row := q.db.QueryRow(ctx, updateSandboxStatusIf, arg.ID, arg.Status, arg.Status_2)
var i Sandbox
err := row.Scan(
&i.ID,
&i.TeamID,
&i.HostID,
&i.Template,
&i.Status,
&i.Vcpus,
&i.MemoryMb,
&i.TimeoutSec,
&i.DiskSizeMb,
&i.GuestIp,
&i.HostIp,
&i.CreatedAt,
&i.StartedAt,
&i.LastActiveAt,
&i.LastUpdated,
&i.TemplateID,
&i.TemplateTeamID,
&i.Metadata,
)
return i, err
}