NianAIGC/backend/internal/postgres/jobs.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, 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, created_at, updated_at`

const ClaimJobsSQL = `SELECT ` + jobColumns + ` FROM public.claim_generation_jobs($1::text, $2::integer, $3::integer)`

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, 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, created_at, updated_at`
	placeholders := make([]string, 33)
	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.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.ProviderTaskID != nil {
		add("provider_task_id", "::text", *patch.ProviderTaskID)
	}
	if patch.ClearProviderTaskID {
		sets = append(sets, "provider_task_id = 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()")
	query := `UPDATE public.generation_jobs SET ` + strings.Join(sets, ", ") + ` WHERE id = $1::text 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 {
		return jobs.Job{}, fmt.Errorf("generation job not found: %s", id)
	}
	return job, nil
}

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 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, &requestPayload, &responsePayload,
		&errorPayload, &retryOf, &idempotencyKey, &fingerprint, &job.Priority, &job.Attempts, &job.MaxAttempts,
		&job.ScheduledAt, &lockedAt, &lockedBy, &startedAt, &completedAt, &webhookURL, &job.WebhookAttempts,
		&webhookStatus, &usageContext, &billing, &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.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.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.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.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)