Add AI job management endpoints
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user