83 lines
2.5 KiB
Go
83 lines
2.5 KiB
Go
|
package db
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
|
||
|
|
"github.com/google/uuid"
|
||
|
|
"github.com/jackc/pgx/v5"
|
||
|
|
)
|
||
|
|
|
||
|
|
// ScanJob is an async "scan this project's AE template" job.
|
||
|
|
type ScanJob struct {
|
||
|
|
ID uuid.UUID `json:"id"`
|
||
|
|
ProjectID uuid.UUID `json:"project_id"`
|
||
|
|
Status string `json:"status"` // queued | running | done | error
|
||
|
|
Engine string `json:"engine"`
|
||
|
|
Result json.RawMessage `json:"result,omitempty"` // ScanResult JSON, present when done
|
||
|
|
Error *string `json:"error,omitempty"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// ScanClaim is the minimal info a node needs to run a claimed scan.
|
||
|
|
type ScanClaim struct {
|
||
|
|
ID uuid.UUID
|
||
|
|
ProjectID uuid.UUID
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Store) CreateScanJob(ctx context.Context, projectID uuid.UUID, engine string) (uuid.UUID, error) {
|
||
|
|
var id uuid.UUID
|
||
|
|
err := s.pool.QueryRow(ctx,
|
||
|
|
`INSERT INTO render.scan_jobs (project_id, engine, status) VALUES ($1, $2, 'queued') RETURNING id`,
|
||
|
|
projectID, engine).Scan(&id)
|
||
|
|
return id, err
|
||
|
|
}
|
||
|
|
|
||
|
|
// ClaimScanJob atomically grabs the oldest queued ae-jsx scan for a node.
|
||
|
|
// Returns nil when the queue is empty.
|
||
|
|
func (s *Store) ClaimScanJob(ctx context.Context, nodeID uuid.UUID) (*ScanClaim, error) {
|
||
|
|
var c ScanClaim
|
||
|
|
err := s.pool.QueryRow(ctx, `
|
||
|
|
UPDATE render.scan_jobs SET status = 'running', node_id = $1, updated_at = NOW()
|
||
|
|
WHERE id = (
|
||
|
|
SELECT id FROM render.scan_jobs
|
||
|
|
WHERE status = 'queued' AND engine = 'ae-jsx'
|
||
|
|
ORDER BY created_at
|
||
|
|
LIMIT 1 FOR UPDATE SKIP LOCKED
|
||
|
|
)
|
||
|
|
RETURNING id, project_id`, nodeID).Scan(&c.ID, &c.ProjectID)
|
||
|
|
if err != nil {
|
||
|
|
if err == pgx.ErrNoRows {
|
||
|
|
return nil, nil
|
||
|
|
}
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return &c, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Store) SetScanResult(ctx context.Context, id uuid.UUID, resultJSON string) error {
|
||
|
|
_, err := s.pool.Exec(ctx,
|
||
|
|
`UPDATE render.scan_jobs SET status = 'done', result = $2::jsonb, error = NULL, updated_at = NOW() WHERE id = $1`,
|
||
|
|
id, resultJSON)
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Store) SetScanError(ctx context.Context, id uuid.UUID, msg string) error {
|
||
|
|
_, err := s.pool.Exec(ctx,
|
||
|
|
`UPDATE render.scan_jobs SET status = 'error', error = $2, updated_at = NOW() WHERE id = $1`, id, msg)
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Store) GetScanJob(ctx context.Context, id uuid.UUID) (*ScanJob, error) {
|
||
|
|
var j ScanJob
|
||
|
|
err := s.pool.QueryRow(ctx,
|
||
|
|
`SELECT id, project_id, status, engine, result, error FROM render.scan_jobs WHERE id = $1`,
|
||
|
|
id).Scan(&j.ID, &j.ProjectID, &j.Status, &j.Engine, &j.Result, &j.Error)
|
||
|
|
if err != nil {
|
||
|
|
if err == pgx.ErrNoRows {
|
||
|
|
return nil, nil
|
||
|
|
}
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return &j, nil
|
||
|
|
}
|