feat: complete remaining Go backend modules
This commit is contained in:
zn-admin committed
1 parent
cea2751dc5
commit
aef5a97165
145 files changed
+18376
-199
No files matched your search
@@ -12,12 +12,23 @@ import (
|
||||
)
|
||||
|
||||
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,
|
||||
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, created_at, updated_at`
|
||||
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 {
|
||||
@@ -75,11 +86,11 @@ func (db *Database) CreateJob(ctx context.Context, job jobs.Job) (jobs.Job, erro
|
||||
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,
|
||||
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, created_at, updated_at`
|
||||
placeholders := make([]string, 33)
|
||||
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)
|
||||
}
|
||||
@@ -117,6 +128,9 @@ func (db *Database) UpdateJob(ctx context.Context, id string, patch jobs.Patch)
|
||||
}
|
||||
add("error", "::jsonb", raw)
|
||||
}
|
||||
if patch.ClearError {
|
||||
sets = append(sets, "error = NULL")
|
||||
}
|
||||
if patch.Attempts != nil {
|
||||
add("attempts", "::integer", *patch.Attempts)
|
||||
}
|
||||
@@ -126,12 +140,27 @@ func (db *Database) UpdateJob(ctx context.Context, id string, patch jobs.Patch)
|
||||
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")
|
||||
}
|
||||
@@ -145,17 +174,146 @@ func (db *Database) UpdateJob(ctx context.Context, id string, patch jobs.Patch)
|
||||
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
|
||||
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
|
||||
@@ -221,15 +379,15 @@ 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 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, &requestPayload, &responsePayload,
|
||||
&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, &job.CreatedAt, &job.UpdatedAt,
|
||||
&webhookStatus, &usageContext, &billing, &dispatchReadyAt, &finalizedAt, &job.CreatedAt, &job.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return jobs.Job{}, fmt.Errorf("scan generation job: %w", err)
|
||||
@@ -238,6 +396,7 @@ func scanJob(rows Rows) (jobs.Job, error) {
|
||||
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)
|
||||
@@ -246,6 +405,8 @@ func scanJob(rows Rows) (jobs.Job, error) {
|
||||
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)
|
||||
@@ -276,11 +437,11 @@ func jobArguments(job jobs.Job) []any {
|
||||
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.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.CreatedAt, job.UpdatedAt,
|
||||
optionalJSON(job.WebhookLastStatus), optionalJSON(job.UsageContext), optionalJSON(job.Billing), job.DispatchReadyAt, job.FinalizedAt, job.CreatedAt, job.UpdatedAt,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user