Add AI job management endpoints
All checks were successful
CI / test (push) Successful in 13s
Build and Deploy / build-and-deploy (push) Successful in 24s

This commit is contained in:
Grendgi
2026-06-08 13:57:28 +03:00
parent 59e1073d96
commit 7452b1d5f2
4 changed files with 301 additions and 0 deletions

View File

@@ -128,6 +128,67 @@ WHERE id = $1
return job, err
}
func (s *Store) ListJobs(ctx context.Context, filter model.JobFilter) ([]*model.Job, error) {
normalizeFilter(&filter)
const q = `
SELECT ` + jobSelectColumns + `
FROM ai_jobs
WHERE ($1 = '' OR owner_service = $1)
AND ($2 = '' OR owner_ref = $2)
AND ($3 = '' OR task_type = $3)
AND ($4 = '' OR model_profile = $4)
AND (cardinality($5::text[]) = 0 OR status = ANY($5::text[]))
AND (cardinality($6::text[]) = 0 OR COALESCE(NULLIF(error_code, ''), 'unknown') = ANY($6::text[]))
ORDER BY created_at DESC
LIMIT $7 OFFSET $8
`
rows, err := s.pool.Query(ctx, q,
filter.OwnerService,
filter.OwnerRef,
filter.TaskType,
filter.ModelProfile,
filter.Statuses,
filter.ErrorCodes,
filter.Limit,
filter.Offset,
)
if err != nil {
return nil, err
}
defer rows.Close()
var out []*model.Job
for rows.Next() {
job, err := scanJob(rows)
if err != nil {
return nil, err
}
out = append(out, job)
}
return out, rows.Err()
}
func normalizeFilter(filter *model.JobFilter) {
filter.OwnerService = strings.TrimSpace(filter.OwnerService)
filter.OwnerRef = strings.TrimSpace(filter.OwnerRef)
filter.TaskType = strings.TrimSpace(filter.TaskType)
filter.ModelProfile = strings.TrimSpace(filter.ModelProfile)
if filter.Statuses == nil {
filter.Statuses = []string{}
}
if filter.ErrorCodes == nil {
filter.ErrorCodes = []string{}
}
if filter.Limit <= 0 {
filter.Limit = 100
}
if filter.Limit > 500 {
filter.Limit = 500
}
if filter.Offset < 0 {
filter.Offset = 0
}
}
func (s *Store) RetryJob(ctx context.Context, id uuid.UUID) (*model.Job, error) {
const q = `
UPDATE ai_jobs
@@ -155,6 +216,87 @@ RETURNING id, owner_service, owner_ref, task_type, model_profile, priority, stat
return job, err
}
func (s *Store) RetryJobs(ctx context.Context, filter model.JobFilter) (int, error) {
normalizeFilter(&filter)
const q = `
WITH picked AS (
SELECT id
FROM ai_jobs
WHERE status IN ('failed', 'running')
AND ($1 = '' OR owner_service = $1)
AND ($2 = '' OR owner_ref = $2)
AND ($3 = '' OR task_type = $3)
AND ($4 = '' OR model_profile = $4)
AND (cardinality($5::text[]) = 0 OR COALESCE(NULLIF(error_code, ''), 'unknown') = ANY($5::text[]))
ORDER BY updated_at ASC
LIMIT $6
)
UPDATE ai_jobs j
SET status = 'pending',
attempts = 0,
started_at = NULL,
completed_at = NULL,
error_code = NULL,
error_message = NULL,
worker_id = NULL,
heartbeat_at = NULL,
scheduled_at = NOW(),
updated_at = NOW()
FROM picked
WHERE j.id = picked.id
`
tag, err := s.pool.Exec(ctx, q,
filter.OwnerService,
filter.OwnerRef,
filter.TaskType,
filter.ModelProfile,
filter.ErrorCodes,
filter.Limit,
)
if err != nil {
return 0, err
}
return int(tag.RowsAffected()), nil
}
func (s *Store) CancelJobs(ctx context.Context, filter model.JobFilter) (int, error) {
normalizeFilter(&filter)
const q = `
WITH picked AS (
SELECT id
FROM ai_jobs
WHERE status IN ('pending', 'running')
AND ($1 = '' OR owner_service = $1)
AND ($2 = '' OR owner_ref = $2)
AND ($3 = '' OR task_type = $3)
AND ($4 = '' OR model_profile = $4)
AND (cardinality($5::text[]) = 0 OR status = ANY($5::text[]))
ORDER BY updated_at ASC
LIMIT $6
)
UPDATE ai_jobs j
SET status = 'cancelled',
completed_at = NOW(),
worker_id = NULL,
heartbeat_at = NULL,
updated_at = NOW()
FROM picked
WHERE j.id = picked.id
`
tag, err := s.pool.Exec(ctx, q,
filter.OwnerService,
filter.OwnerRef,
filter.TaskType,
filter.ModelProfile,
filter.Statuses,
filter.Limit,
)
if err != nil {
return 0, err
}
return int(tag.RowsAffected()), nil
}
func (s *Store) ClaimJobs(ctx context.Context, in model.ClaimJobs) ([]*model.Job, error) {
if in.Limit <= 0 {
in.Limit = 1
@@ -326,6 +468,27 @@ ORDER BY task_type, model_profile, status
return nil, err
}
ownerRows, err := s.pool.Query(ctx, `
SELECT owner_service, task_type, model_profile, status, count(*)
FROM ai_jobs
GROUP BY owner_service, task_type, model_profile, status
ORDER BY owner_service, task_type, model_profile, status
`)
if err != nil {
return nil, err
}
defer ownerRows.Close()
for ownerRows.Next() {
var stat model.OwnerStat
if err := ownerRows.Scan(&stat.OwnerService, &stat.TaskType, &stat.ModelProfile, &stat.Status, &stat.Total); err != nil {
return nil, err
}
out.Owners = append(out.Owners, stat)
}
if err := ownerRows.Err(); err != nil {
return nil, err
}
errorRows, err := s.pool.Query(ctx, `
SELECT task_type, model_profile, COALESCE(NULLIF(error_code, ''), 'unknown') AS error_code,
count(*) AS total,