Files
NianAIGC/backend/internal/postgres/jobs.go

509 lines
19 KiB
Go

package postgres
import (
"context"
"database/sql"
"encoding/json"
"fmt"
"strings"
"time"
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs"
)
const jobColumns = `id, owner_id, external_client_id, capability, provider, req_key, status, prompt,
input_asset_ids, input_urls, output_asset_ids, provider_task_id, provider_dispatch_started_at, request_payload, response_payload,
error, retry_of, idempotency_key, idempotency_fingerprint, priority, attempts, max_attempts,
scheduled_at, locked_at, locked_by, started_at, completed_at, webhook_url, webhook_attempts,
webhook_last_status, usage_context, billing, dispatch_ready_at, finalized_at, created_at, updated_at`
const ClaimJobsSQL = `SELECT ` + jobColumns + ` FROM public.claim_generation_jobs($1::text, $2::integer, $3::integer)`
const WriteJobBillingSQL = `UPDATE public.generation_jobs SET billing = $2::jsonb, updated_at = now() WHERE id = $1::text RETURNING id`
const ActivateJobCreationSQL = `UPDATE public.generation_jobs
SET billing = $2::jsonb, dispatch_ready_at = COALESCE(dispatch_ready_at, now()), updated_at = now()
WHERE id = $1::text AND (dispatch_ready_at IS NULL OR billing = $2::jsonb) RETURNING id`
const WriteJobOutputAssetIDsSQL = `UPDATE public.generation_jobs SET output_asset_ids = $2::text[], updated_at = now() WHERE id = $1::text RETURNING id`
const WriteJobBillingFencedSQL = `UPDATE public.generation_jobs SET billing = $2::jsonb, updated_at = now() WHERE id = $1::text AND status = $3::text AND locked_by = $4::text RETURNING id`
const WriteJobOutputAssetIDsFencedSQL = `UPDATE public.generation_jobs SET output_asset_ids = $2::text[], updated_at = now() WHERE id = $1::text AND status = $3::text AND locked_by = $4::text RETURNING id`
const FailJobCreationSQL = `UPDATE public.generation_jobs
SET status = $2::text, error = $3::jsonb, billing = $4::jsonb,
completed_at = COALESCE(completed_at, now()), locked_at = NULL, locked_by = NULL, updated_at = now()
WHERE id = $1::text RETURNING id`
func (db *Database) ListJobs(ctx context.Context, filter jobs.ListFilter) ([]jobs.Job, error) {
if err := db.requirePostgres("list generation jobs"); err != nil {
return nil, err
}
clauses := []string{"owner_id = $1::text"}
args := []any{filter.Scope.OwnerID}
if filter.Scope.ExternalClientID != "" {
args = append(args, filter.Scope.ExternalClientID)
clauses = append(clauses, fmt.Sprintf("external_client_id = $%d::text", len(args)))
}
if filter.Status != "" {
args = append(args, string(filter.Status))
clauses = append(clauses, fmt.Sprintf("status = $%d::text", len(args)))
}
if filter.Capability != "" {
args = append(args, filter.Capability)
clauses = append(clauses, fmt.Sprintf("capability = $%d::text", len(args)))
}
if filter.Before != nil {
args = append(args, *filter.Before)
clauses = append(clauses, fmt.Sprintf("created_at < $%d::timestamptz", len(args)))
}
args = append(args, filter.Limit)
query := `SELECT ` + jobColumns + ` FROM public.generation_jobs WHERE ` + strings.Join(clauses, " AND ") +
fmt.Sprintf(" ORDER BY created_at DESC LIMIT $%d::integer", len(args))
return db.queryJobs(ctx, query, args...)
}
func (db *Database) FindJob(ctx context.Context, scope jobs.Scope, id string) (jobs.Job, bool, error) {
if err := db.requirePostgres("find generation job"); err != nil {
return jobs.Job{}, false, err
}
query := `SELECT ` + jobColumns + ` FROM public.generation_jobs WHERE id = $1::text AND owner_id = $2::text`
args := []any{id, scope.OwnerID}
if scope.ExternalClientID != "" {
query += ` AND external_client_id = $3::text`
args = append(args, scope.ExternalClientID)
}
query += ` LIMIT 1`
return db.queryOneJob(ctx, query, args...)
}
func (db *Database) FindIdempotentJob(ctx context.Context, scope jobs.Scope, key string) (jobs.Job, bool, error) {
if err := db.requirePostgres("find idempotent generation job"); err != nil {
return jobs.Job{}, false, err
}
query := `SELECT ` + jobColumns + ` FROM public.generation_jobs
WHERE owner_id = $1::text AND external_client_id = $2::text AND idempotency_key = $3::text LIMIT 1`
return db.queryOneJob(ctx, query, scope.OwnerID, scope.ExternalClientID, key)
}
func (db *Database) CreateJob(ctx context.Context, job jobs.Job) (jobs.Job, error) {
if err := db.requirePostgres("create generation job"); err != nil {
return jobs.Job{}, err
}
const columns = `id, owner_id, external_client_id, capability, provider, req_key, status, prompt,
input_asset_ids, input_urls, output_asset_ids, provider_task_id, provider_dispatch_started_at, request_payload, response_payload,
error, retry_of, idempotency_key, idempotency_fingerprint, priority, attempts, max_attempts,
scheduled_at, locked_at, locked_by, started_at, completed_at, webhook_url, webhook_attempts,
webhook_last_status, usage_context, billing, dispatch_ready_at, finalized_at, created_at, updated_at`
placeholders := make([]string, 36)
for index := range placeholders {
placeholders[index] = fmt.Sprintf("$%d", index+1)
}
query := `INSERT INTO public.generation_jobs (` + columns + `) VALUES (` + strings.Join(placeholders, ", ") + `) RETURNING ` + jobColumns
row, found, err := db.queryOneJob(ctx, query, jobArguments(job)...)
if err != nil {
if sqlState(err) == "23505" && job.ExternalClientID != "" && job.IdempotencyKey != "" {
return jobs.Job{}, jobs.ErrUniqueIdempotency
}
return jobs.Job{}, fmt.Errorf("create generation job: %w", err)
}
if !found {
return jobs.Job{}, fmt.Errorf("create generation job returned no row")
}
return row, nil
}
func (db *Database) UpdateJob(ctx context.Context, id string, patch jobs.Patch) (jobs.Job, error) {
if err := db.requirePostgres("update generation job"); err != nil {
return jobs.Job{}, err
}
sets := make([]string, 0, 12)
args := []any{id}
add := func(column, cast string, value any) {
args = append(args, value)
sets = append(sets, fmt.Sprintf("%s = $%d%s", column, len(args), cast))
}
if patch.Status != nil {
add("status", "::text", string(*patch.Status))
}
if patch.Error != nil {
raw, err := json.Marshal(patch.Error)
if err != nil {
return jobs.Job{}, fmt.Errorf("encode generation job error: %w", err)
}
add("error", "::jsonb", raw)
}
if patch.ClearError {
sets = append(sets, "error = NULL")
}
if patch.Attempts != nil {
add("attempts", "::integer", *patch.Attempts)
}
if patch.ScheduledAt != nil {
add("scheduled_at", "::timestamptz", *patch.ScheduledAt)
}
if patch.CompletedAt != nil {
add("completed_at", "::timestamptz", *patch.CompletedAt)
}
if patch.DispatchReadyAt != nil {
add("dispatch_ready_at", "::timestamptz", *patch.DispatchReadyAt)
}
if patch.FinalizedAt != nil {
add("finalized_at", "::timestamptz", *patch.FinalizedAt)
}
if patch.ProviderTaskID != nil {
add("provider_task_id", "::text", *patch.ProviderTaskID)
}
if patch.ProviderDispatchStartedAt != nil {
add("provider_dispatch_started_at", "::timestamptz", *patch.ProviderDispatchStartedAt)
}
if patch.SetResponsePayload {
add("response_payload", "::jsonb", optionalJSON(patch.ResponsePayload))
}
if patch.ClearProviderTaskID {
sets = append(sets, "provider_task_id = NULL")
}
if patch.ClearProviderDispatch {
sets = append(sets, "provider_dispatch_started_at = NULL")
}
if patch.ClearLease {
sets = append(sets, "locked_at = NULL", "locked_by = NULL")
}
if patch.WebhookAttempts != nil {
add("webhook_attempts", "::integer", *patch.WebhookAttempts)
}
if patch.SetWebhookStatus {
add("webhook_last_status", "::jsonb", patch.WebhookLastStatus)
}
if len(sets) == 0 {
return db.mustFindJobByID(ctx, id)
}
sets = append(sets, "updated_at = now()")
where := []string{"id = $1::text"}
if len(patch.ExpectedStatuses) != 0 {
values := make([]string, len(patch.ExpectedStatuses))
for index, status := range patch.ExpectedStatuses {
values[index] = string(status)
}
args = append(args, values)
where = append(where, fmt.Sprintf("status = ANY($%d::text[])", len(args)))
}
if patch.ExpectedLockedBy != nil {
args = append(args, *patch.ExpectedLockedBy)
where = append(where, fmt.Sprintf("locked_by = $%d::text", len(args)))
}
query := `UPDATE public.generation_jobs SET ` + strings.Join(sets, ", ") + ` WHERE ` + strings.Join(where, " AND ") + ` RETURNING ` + jobColumns
job, found, err := db.queryOneJob(ctx, query, args...)
if err != nil {
return jobs.Job{}, fmt.Errorf("update generation job: %w", err)
}
if !found {
if len(patch.ExpectedStatuses) != 0 || patch.ExpectedLockedBy != nil {
return jobs.Job{}, jobs.ErrTransitionConflict
}
return jobs.Job{}, fmt.Errorf("generation job not found: %s", id)
}
return job, nil
}
func (db *Database) DeleteJob(ctx context.Context, id string) error {
if err := db.requirePostgres("delete generation job"); err != nil {
return err
}
rows, err := db.querier.Query(ctx, `DELETE FROM public.generation_jobs WHERE id = $1::text RETURNING id`, id)
if err != nil {
return fmt.Errorf("delete generation job: %w", err)
}
defer rows.Close()
if !rows.Next() {
return fmt.Errorf("generation job not found: %s", id)
}
var deleted string
if err := rows.Scan(&deleted); err != nil {
return fmt.Errorf("scan deleted generation job: %w", err)
}
return nil
}
// WriteBilling and WriteOutputAssetIDs implement the deliberately narrow
// orchestration state seam without widening jobs.Patch into a generic update
// bag. Both fail closed when the target disappeared during processing.
func (db *Database) WriteBilling(ctx context.Context, id string, value json.RawMessage) error {
if len(value) == 0 || !json.Valid(value) {
return fmt.Errorf("write generation billing: invalid JSON")
}
return db.writeJobState(ctx, "write generation billing", WriteJobBillingSQL, id, value)
}
func (db *Database) ActivateCreation(ctx context.Context, id string, value json.RawMessage) error {
if len(value) == 0 || !json.Valid(value) {
return fmt.Errorf("activate generation creation: invalid JSON")
}
return db.writeJobState(ctx, "activate generation creation", ActivateJobCreationSQL, id, value)
}
func (db *Database) WriteOutputAssetIDs(ctx context.Context, id string, values []string) error {
if values == nil {
values = []string{}
}
return db.writeJobState(ctx, "write generation output assets", WriteJobOutputAssetIDsSQL, id, values)
}
func (db *Database) WriteBillingFenced(ctx context.Context, id string, value json.RawMessage, status jobs.Status, lockedBy string) error {
if len(value) == 0 || !json.Valid(value) || lockedBy == "" {
return fmt.Errorf("write generation billing: invalid fenced state")
}
return db.writeJobStateConflict(ctx, "write generation billing", WriteJobBillingFencedSQL, id, value, string(status), lockedBy)
}
func (db *Database) WriteOutputAssetIDsFenced(ctx context.Context, id string, values []string, status jobs.Status, lockedBy string) error {
if lockedBy == "" {
return fmt.Errorf("write generation output assets: invalid fenced state")
}
if values == nil {
values = []string{}
}
return db.writeJobStateConflict(ctx, "write generation output assets", WriteJobOutputAssetIDsFencedSQL, id, values, string(status), lockedBy)
}
// FailCreation commits the externally observable result of a failed charge in
// one statement. A job must never remain queued with a wallet-side failure.
func (db *Database) FailCreation(ctx context.Context, job jobs.Job) error {
if job.Status != jobs.StatusFailed || job.Error == nil || len(job.Billing) == 0 || !json.Valid(job.Billing) {
return fmt.Errorf("fail generation creation: invalid terminal state")
}
errorJSON, err := json.Marshal(job.Error)
if err != nil {
return fmt.Errorf("fail generation creation: encode error: %w", err)
}
return db.writeJobStateArgs(ctx, "fail generation creation", FailJobCreationSQL, job.ID, string(job.Status), json.RawMessage(errorJSON), job.Billing)
}
func (db *Database) writeJobState(ctx context.Context, operation, query, id string, value any) error {
return db.writeJobStateArgs(ctx, operation, query, id, value)
}
func (db *Database) writeJobStateArgs(ctx context.Context, operation, query, id string, values ...any) error {
if err := db.requirePostgres(operation); err != nil {
return err
}
args := make([]any, 0, len(values)+1)
args = append(args, id)
args = append(args, values...)
rows, err := db.querier.Query(ctx, query, args...)
if err != nil {
return fmt.Errorf("%s: %w", operation, err)
}
defer rows.Close()
if !rows.Next() {
if err := rows.Err(); err != nil {
return fmt.Errorf("%s: %w", operation, err)
}
return fmt.Errorf("%s: generation job not found", operation)
}
var returnedID string
if err := rows.Scan(&returnedID); err != nil {
return fmt.Errorf("%s: scan result: %w", operation, err)
}
if err := rows.Err(); err != nil {
return fmt.Errorf("%s: %w", operation, err)
}
return nil
}
func (db *Database) writeJobStateConflict(ctx context.Context, operation, query, id string, values ...any) error {
err := db.writeJobStateArgs(ctx, operation, query, id, values...)
if err != nil && strings.Contains(err.Error(), "generation job not found") {
return jobs.ErrTransitionConflict
}
return err
}
func (db *Database) ClaimJobs(ctx context.Context, workerID string, limit, lockTimeoutSeconds int) ([]jobs.Job, error) {
if err := db.requirePostgres("claim generation jobs"); err != nil {
return nil, err
}
limit = max(1, min(limit, 20))
if lockTimeoutSeconds <= 0 {
lockTimeoutSeconds = 300
}
return db.queryJobs(ctx, ClaimJobsSQL, workerID, limit, lockTimeoutSeconds)
}
func (db *Database) mustFindJobByID(ctx context.Context, id string) (jobs.Job, error) {
job, found, err := db.queryOneJob(ctx, `SELECT `+jobColumns+` FROM public.generation_jobs WHERE id = $1::text LIMIT 1`, id)
if err != nil {
return jobs.Job{}, err
}
if !found {
return jobs.Job{}, fmt.Errorf("generation job not found: %s", id)
}
return job, nil
}
func (db *Database) queryJobs(ctx context.Context, query string, args ...any) ([]jobs.Job, error) {
rows, err := db.querier.Query(ctx, query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]jobs.Job, 0)
for rows.Next() {
item, err := scanJob(rows)
if err != nil {
return nil, err
}
items = append(items, item)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
func (db *Database) queryOneJob(ctx context.Context, query string, args ...any) (jobs.Job, bool, error) {
rows, err := db.querier.Query(ctx, query, args...)
if err != nil {
return jobs.Job{}, false, err
}
defer rows.Close()
if !rows.Next() {
return jobs.Job{}, false, rows.Err()
}
job, err := scanJob(rows)
if err != nil {
return jobs.Job{}, false, err
}
if err := rows.Err(); err != nil {
return jobs.Job{}, false, err
}
return job, true, nil
}
func scanJob(rows Rows) (jobs.Job, error) {
var job jobs.Job
var externalClientID, prompt, providerTaskID, retryOf, idempotencyKey, fingerprint sql.NullString
var lockedBy, webhookURL sql.NullString
var lockedAt, startedAt, completedAt, dispatchReadyAt, finalizedAt, providerDispatchStartedAt sql.NullTime
var requestPayload, responsePayload, errorPayload, webhookStatus, usageContext, billing []byte
var status string
err := rows.Scan(
&job.ID, &job.OwnerID, &externalClientID, &job.Capability, &job.Provider, &job.ReqKey, &status, &prompt,
&job.InputAssetIDs, &job.InputURLs, &job.OutputAssetIDs, &providerTaskID, &providerDispatchStartedAt, &requestPayload, &responsePayload,
&errorPayload, &retryOf, &idempotencyKey, &fingerprint, &job.Priority, &job.Attempts, &job.MaxAttempts,
&job.ScheduledAt, &lockedAt, &lockedBy, &startedAt, &completedAt, &webhookURL, &job.WebhookAttempts,
&webhookStatus, &usageContext, &billing, &dispatchReadyAt, &finalizedAt, &job.CreatedAt, &job.UpdatedAt,
)
if err != nil {
return jobs.Job{}, fmt.Errorf("scan generation job: %w", err)
}
job.Status = jobs.Status(status)
job.ExternalClientID = nullString(externalClientID)
job.Prompt = nullString(prompt)
job.ProviderTaskID = nullString(providerTaskID)
job.ProviderDispatchStartedAt = nullTimePointer(providerDispatchStartedAt)
job.RetryOf = nullString(retryOf)
job.IdempotencyKey = nullString(idempotencyKey)
job.IdempotencyFingerprint = nullString(fingerprint)
job.LockedBy = nullString(lockedBy)
job.WebhookURL = nullString(webhookURL)
job.LockedAt = nullTimePointer(lockedAt)
job.StartedAt = nullTimePointer(startedAt)
job.CompletedAt = nullTimePointer(completedAt)
job.DispatchReadyAt = nullTimePointer(dispatchReadyAt)
job.FinalizedAt = nullTimePointer(finalizedAt)
job.RequestPayload = normalizedJSON(requestPayload, `{}`)
job.ResponsePayload = normalizedOptionalJSON(responsePayload)
job.WebhookLastStatus = normalizedOptionalJSON(webhookStatus)
job.UsageContext = normalizedOptionalJSON(usageContext)
job.Billing = normalizedOptionalJSON(billing)
if len(errorPayload) > 0 && string(errorPayload) != "null" {
if err := json.Unmarshal(errorPayload, &job.Error); err != nil {
return jobs.Job{}, fmt.Errorf("decode generation job error: %w", err)
}
}
if job.InputAssetIDs == nil {
job.InputAssetIDs = []string{}
}
if job.InputURLs == nil {
job.InputURLs = []string{}
}
if job.OutputAssetIDs == nil {
job.OutputAssetIDs = []string{}
}
return job, nil
}
func jobArguments(job jobs.Job) []any {
errorPayload, _ := json.Marshal(job.Error)
if job.Error == nil {
errorPayload = nil
}
return []any{
job.ID, job.OwnerID, optionalDatabaseText(job.ExternalClientID), job.Capability, job.Provider, job.ReqKey,
string(job.Status), optionalDatabaseText(job.Prompt), job.InputAssetIDs, job.InputURLs, job.OutputAssetIDs,
optionalDatabaseText(job.ProviderTaskID), job.ProviderDispatchStartedAt, job.RequestPayload, optionalJSON(job.ResponsePayload), optionalJSON(errorPayload),
optionalDatabaseText(job.RetryOf), optionalDatabaseText(job.IdempotencyKey), optionalDatabaseText(job.IdempotencyFingerprint),
job.Priority, job.Attempts, job.MaxAttempts, job.ScheduledAt, job.LockedAt, optionalDatabaseText(job.LockedBy),
job.StartedAt, job.CompletedAt, optionalDatabaseText(job.WebhookURL), job.WebhookAttempts,
optionalJSON(job.WebhookLastStatus), optionalJSON(job.UsageContext), optionalJSON(job.Billing), job.DispatchReadyAt, job.FinalizedAt, job.CreatedAt, job.UpdatedAt,
}
}
func (db *Database) requirePostgres(operation string) error {
if db == nil || db.config.Backend != BackendPostgres || db.querier == nil {
return fmt.Errorf("%s: PostgreSQL is unavailable", operation)
}
return nil
}
func nullString(value sql.NullString) string {
if value.Valid {
return value.String
}
return ""
}
func nullTimePointer(value sql.NullTime) *time.Time {
if !value.Valid {
return nil
}
copy := value.Time
return &copy
}
func optionalJSON(value []byte) any {
if len(value) == 0 || string(value) == "null" {
return nil
}
return value
}
func normalizedJSON(value []byte, fallback string) json.RawMessage {
if len(value) == 0 || string(value) == "null" {
return json.RawMessage(fallback)
}
return append(json.RawMessage(nil), value...)
}
func normalizedOptionalJSON(value []byte) json.RawMessage {
if len(value) == 0 || string(value) == "null" {
return nil
}
return append(json.RawMessage(nil), value...)
}
type sqlStateError interface{ SQLState() string }
func sqlState(err error) string {
for err != nil {
if state, ok := err.(sqlStateError); ok {
return state.SQLState()
}
type unwrapper interface{ Unwrap() error }
wrapped, ok := err.(unwrapper)
if !ok {
break
}
err = wrapped.Unwrap()
}
return ""
}
var _ jobs.Store = (*Database)(nil)