flatrender/services/render/internal/db/snapshotjobs.go

package db

import (
	"context"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

// SnapshotJob is an async "render one frame of a scene" job (status row).
type SnapshotJob struct {
	ID       uuid.UUID `json:"id"`
	SceneID  uuid.UUID `json:"scene_id"`
	SceneKey string    `json:"scene_key"`
	Status   string    `json:"status"` // queued | running | done | error
	ImageURL *string   `json:"image_url,omitempty"`
	Error    *string   `json:"error,omitempty"`
}

// SnapshotClaim is the minimal info a node needs to render a claimed snapshot.
type SnapshotClaim struct {
	ID        uuid.UUID
	ProjectID uuid.UUID
	SceneID   uuid.UUID
	SceneKey  string
	CompName  string
	Frame     int
}

// EnqueueSceneSnapshots clears any prior snapshot jobs for the project and queues
// one fresh job per active scene (comp_name defaults to the scene key — a comp of
// that name; the node falls back to the render comp when it does not exist).
// Returns the number of jobs queued.
func (s *Store) EnqueueSceneSnapshots(ctx context.Context, projectID uuid.UUID) (int, error) {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return 0, err
	}
	defer tx.Rollback(ctx)

	if _, err = tx.Exec(ctx, `DELETE FROM render.snapshot_jobs WHERE project_id = $1`, projectID); err != nil {
		return 0, err
	}
	tag, err := tx.Exec(ctx, `
		INSERT INTO render.snapshot_jobs (project_id, scene_id, scene_key, comp_name, frame, status)
		SELECT $1, sc.id, sc.key, sc.key, 0, 'queued'
		FROM content.scenes sc
		WHERE sc.project_id = $1 AND sc.deleted_at IS NULL AND sc.is_active = true`, projectID)
	if err != nil {
		return 0, err
	}
	if err = tx.Commit(ctx); err != nil {
		return 0, err
	}
	return int(tag.RowsAffected()), nil
}

// ListSnapshotJobs returns the snapshot jobs for a project (for polling).
func (s *Store) ListSnapshotJobs(ctx context.Context, projectID uuid.UUID) ([]SnapshotJob, error) {
	rows, err := s.pool.Query(ctx,
		`SELECT id, scene_id, scene_key, status, image_url, error
		 FROM render.snapshot_jobs WHERE project_id = $1 ORDER BY created_at`, projectID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []SnapshotJob
	for rows.Next() {
		var j SnapshotJob
		if err := rows.Scan(&j.ID, &j.SceneID, &j.SceneKey, &j.Status, &j.ImageURL, &j.Error); err != nil {
			return nil, err
		}
		out = append(out, j)
	}
	return out, rows.Err()
}

// ClaimSnapshotJob atomically grabs the oldest queued snapshot for a node.
// Returns nil when the queue is empty.
func (s *Store) ClaimSnapshotJob(ctx context.Context, nodeID uuid.UUID) (*SnapshotClaim, error) {
	var c SnapshotClaim
	err := s.pool.QueryRow(ctx, `
		UPDATE render.snapshot_jobs SET status = 'running', node_id = $1, updated_at = NOW()
		WHERE id = (
			SELECT id FROM render.snapshot_jobs
			WHERE status = 'queued'
			ORDER BY created_at
			LIMIT 1 FOR UPDATE SKIP LOCKED
		)
		RETURNING id, project_id, scene_id, scene_key, comp_name, frame`,
		nodeID).Scan(&c.ID, &c.ProjectID, &c.SceneID, &c.SceneKey, &c.CompName, &c.Frame)
	if err != nil {
		if err == pgx.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}
	return &c, nil
}

// SetSnapshotResult marks a running snapshot done with its image URL and writes
// that URL onto the content scene (same DB, cross-schema) so the studio + admin
// render a real thumbnail.
func (s *Store) SetSnapshotResult(ctx context.Context, id uuid.UUID, imageURL string) error {
	var sceneID uuid.UUID
	err := s.pool.QueryRow(ctx,
		`UPDATE render.snapshot_jobs SET status = 'done', image_url = $2, error = NULL, updated_at = NOW()
		 WHERE id = $1 AND status = 'running' RETURNING scene_id`, id, imageURL).Scan(&sceneID)
	if err != nil {
		if err == pgx.ErrNoRows {
			return nil // job no longer running (cancelled/regenerated) — ignore late result
		}
		return err
	}
	_, err = s.pool.Exec(ctx,
		`UPDATE content.scenes SET snapshot_url = $2, updated_at = NOW() WHERE id = $1`, sceneID, imageURL)
	return err
}

func (s *Store) SetSnapshotError(ctx context.Context, id uuid.UUID, msg string) error {
	_, err := s.pool.Exec(ctx,
		`UPDATE render.snapshot_jobs SET status = 'error', error = $2, updated_at = NOW() WHERE id = $1 AND status = 'running'`,
		id, msg)
	return err
}

// GetSnapshotJobMeta returns the project id + scene key for a job (to build the
// object key the node uploads its rendered still to).
func (s *Store) GetSnapshotJobMeta(ctx context.Context, id uuid.UUID) (uuid.UUID, string, error) {
	var pid uuid.UUID
	var key string
	err := s.pool.QueryRow(ctx,
		`SELECT project_id, scene_key FROM render.snapshot_jobs WHERE id = $1`, id).Scan(&pid, &key)
	return pid, key, err
}