forked from wrenn/wrenn
v0.2.0 (#50)
Co-authored-by: Tasnim Kabir Sadik <tksadik@omukk.dev> Reviewed-on: wrenn/wrenn#50
This commit is contained in:
@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user