weichen
2de4eda16d
Track workspace size and age, expose an admin list and cleanup route, and refuse TTL or quota deletion until a frozen manifest matches verified receipts.
113 lines
3.1 KiB
Go
113 lines
3.1 KiB
Go
package workspace
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
// RegistryDDL is the table used by PGRegistry. Goose migration 00062 is the deploy path.
|
|
const RegistryDDL = `
|
|
CREATE TABLE IF NOT EXISTS agentx_workspace_registry (
|
|
id TEXT PRIMARY KEY,
|
|
run_id TEXT NOT NULL DEFAULT '',
|
|
bytes BIGINT NOT NULL DEFAULT 0,
|
|
created_at TIMESTAMPTZ NOT NULL,
|
|
last_seen TIMESTAMPTZ NOT NULL,
|
|
frozen JSONB,
|
|
receipts JSONB NOT NULL DEFAULT '[]'::jsonb
|
|
)`
|
|
|
|
// PGRegistry persists workspace records in Postgres.
|
|
type PGRegistry struct {
|
|
pool *pgxpool.Pool
|
|
}
|
|
|
|
func NewPGRegistry(pool *pgxpool.Pool) *PGRegistry {
|
|
return &PGRegistry{pool: pool}
|
|
}
|
|
|
|
func (p *PGRegistry) Ensure(ctx context.Context) error {
|
|
_, err := p.pool.Exec(ctx, RegistryDDL)
|
|
return err
|
|
}
|
|
|
|
func (p *PGRegistry) Upsert(ctx context.Context, rec Record) error {
|
|
if err := validateID(rec.ID); err != nil {
|
|
return err
|
|
}
|
|
if rec.CreatedAt.IsZero() {
|
|
rec.CreatedAt = time.Now().UTC()
|
|
}
|
|
if rec.LastSeen.IsZero() {
|
|
rec.LastSeen = rec.CreatedAt
|
|
}
|
|
frozen, err := json.Marshal(rec.Frozen)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if rec.Frozen == nil {
|
|
frozen = []byte("null")
|
|
}
|
|
receipts, err := json.Marshal(rec.Receipts)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if rec.Receipts == nil {
|
|
receipts = []byte("[]")
|
|
}
|
|
_, err = p.pool.Exec(ctx, `
|
|
INSERT INTO agentx_workspace_registry (id, run_id, bytes, created_at, last_seen, frozen, receipts)
|
|
VALUES ($1,$2,$3,$4,$5,$6::jsonb,$7::jsonb)
|
|
ON CONFLICT (id) DO UPDATE SET
|
|
run_id=EXCLUDED.run_id, bytes=EXCLUDED.bytes, last_seen=EXCLUDED.last_seen,
|
|
frozen=EXCLUDED.frozen, receipts=EXCLUDED.receipts`,
|
|
rec.ID, rec.RunID, rec.Bytes, rec.CreatedAt, rec.LastSeen, frozen, receipts)
|
|
return err
|
|
}
|
|
|
|
func (p *PGRegistry) Get(ctx context.Context, id string) (Record, bool, error) {
|
|
row := p.pool.QueryRow(ctx, `SELECT id, run_id, bytes, created_at, last_seen, frozen, receipts FROM agentx_workspace_registry WHERE id=$1`, id)
|
|
rec, err := scanRecord(row)
|
|
if err == pgx.ErrNoRows {
|
|
return Record{}, false, nil
|
|
}
|
|
if err != nil {
|
|
return Record{}, false, err
|
|
}
|
|
return rec, true, nil
|
|
}
|
|
|
|
func scanRecord(row pgx.Row) (Record, error) {
|
|
var rec Record
|
|
var frozen, receipts []byte
|
|
if err := row.Scan(&rec.ID, &rec.RunID, &rec.Bytes, &rec.CreatedAt, &rec.LastSeen, &frozen, &receipts); err != nil {
|
|
return Record{}, err
|
|
}
|
|
if len(frozen) > 0 && string(frozen) != "null" {
|
|
rec.Frozen = &FrozenManifest{}
|
|
if err := json.Unmarshal(frozen, rec.Frozen); err != nil {
|
|
return Record{}, err
|
|
}
|
|
}
|
|
if len(receipts) > 0 {
|
|
if err := json.Unmarshal(receipts, &rec.Receipts); err != nil {
|
|
return Record{}, err
|
|
}
|
|
}
|
|
return rec, nil
|
|
}
|
|
|
|
// TouchSchema creates an isolated schema for tests and points the pool at it.
|
|
func TouchSchema(ctx context.Context, pool *pgxpool.Pool, schema string) error {
|
|
if _, err := pool.Exec(ctx, fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s", schema)); err != nil {
|
|
return err
|
|
}
|
|
_, err := pool.Exec(ctx, fmt.Sprintf("SET search_path TO %s", schema))
|
|
return err
|
|
}
|