192 lines
5.1 KiB
Go
192 lines
5.1 KiB
Go
package db
|
|
|
|
import (
|
|
"database/sql"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
)
|
|
|
|
// EnqueueTask adds a new task to the queue
|
|
func (db *DB) EnqueueTask(taskType string, payload interface{}) error {
|
|
payloadBytes, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to marshal task payload: %w", err)
|
|
}
|
|
|
|
query := `
|
|
INSERT INTO tasks (task_type, payload, status)
|
|
VALUES (?, ?, 'PENDING')
|
|
`
|
|
_, err = db.Exec(query, taskType, string(payloadBytes))
|
|
return err
|
|
}
|
|
|
|
// ClaimNextTask atomically finds and claims the next pending task
|
|
func (db *DB) ClaimNextTask() (*Task, error) {
|
|
query := `
|
|
UPDATE tasks
|
|
SET status = 'RUNNING', updated_at = CURRENT_TIMESTAMP
|
|
WHERE id = (
|
|
SELECT id FROM tasks
|
|
WHERE status = 'PENDING' AND next_retry_at <= CURRENT_TIMESTAMP
|
|
ORDER BY created_at ASC
|
|
LIMIT 1
|
|
)
|
|
RETURNING id, task_type, payload, status, error_message, retry_count, next_retry_at, created_at, updated_at
|
|
`
|
|
|
|
var task Task
|
|
var errMsg sql.NullString
|
|
err := db.QueryRow(query).Scan(
|
|
&task.ID, &task.TaskType, &task.Payload, &task.Status, &errMsg,
|
|
&task.RetryCount, &task.NextRetryAt, &task.CreatedAt, &task.UpdatedAt,
|
|
)
|
|
|
|
if err == sql.ErrNoRows {
|
|
return nil, nil // No tasks available
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if errMsg.Valid {
|
|
task.ErrorMessage = errMsg.String
|
|
}
|
|
|
|
return &task, nil
|
|
}
|
|
|
|
// CompleteTask marks a task as COMPLETED
|
|
func (db *DB) CompleteTask(id int64) error {
|
|
query := `UPDATE tasks SET status = 'COMPLETED', updated_at = CURRENT_TIMESTAMP WHERE id = ?`
|
|
_, err := db.Exec(query, id)
|
|
return err
|
|
}
|
|
|
|
// FailTask marks a task as FAILED or PENDING (for retry) using exponential backoff
|
|
func (db *DB) FailTask(id int64, errStr string) error {
|
|
// First get current retry count
|
|
var retryCount int
|
|
err := db.QueryRow(`SELECT retry_count FROM tasks WHERE id = ?`, id).Scan(&retryCount)
|
|
if err != nil {
|
|
log.Printf("[DB] FailTask: failed to get retry_count for task %d: %v", id, err)
|
|
return err
|
|
}
|
|
|
|
retryCount++
|
|
|
|
// Max retries = 5 (1s, 2s, 4s, 8s, 16s)
|
|
status := "PENDING"
|
|
if retryCount > 5 {
|
|
status = "FAILED"
|
|
}
|
|
|
|
// Exponential backoff
|
|
backoffSeconds := 1 << retryCount // 2^retryCount
|
|
|
|
query := `
|
|
UPDATE tasks
|
|
SET status = ?,
|
|
error_message = ?,
|
|
retry_count = ?,
|
|
next_retry_at = datetime('now', '+' || ? || ' seconds'),
|
|
updated_at = CURRENT_TIMESTAMP
|
|
WHERE id = ?
|
|
`
|
|
_, err = db.Exec(query, status, errStr, retryCount, fmt.Sprintf("%d", backoffSeconds), id)
|
|
return err
|
|
}
|
|
|
|
// HasPendingFetchTask checks if there is already a PENDING or RUNNING fetch task for a given canonical ID
|
|
func (db *DB) HasPendingFetchTask(canonicalID string) (bool, error) {
|
|
query := `
|
|
SELECT COUNT(*) FROM tasks
|
|
WHERE task_type = 'FETCH'
|
|
AND status IN ('PENDING', 'RUNNING')
|
|
AND payload LIKE ?
|
|
`
|
|
likePattern := fmt.Sprintf(`%%"canonical_id":"%s"%%`, canonicalID)
|
|
|
|
var count int
|
|
err := db.QueryRow(query, likePattern).Scan(&count)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
return count > 0, nil
|
|
}
|
|
|
|
// ClearTaskQueue deletes all tasks or only pending/failed ones
|
|
func (db *DB) ClearTaskQueue() error {
|
|
query := `DELETE FROM tasks WHERE status IN ('PENDING', 'FAILED', 'COMPLETED')`
|
|
_, err := db.Exec(query)
|
|
return err
|
|
}
|
|
|
|
// GetQueueStatus returns counts of tasks by status
|
|
func (db *DB) GetQueueStatus() (map[string]int, error) {
|
|
query := `SELECT status, COUNT(*) FROM tasks GROUP BY status`
|
|
rows, err := db.Query(query)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
|
|
counts := make(map[string]int)
|
|
for rows.Next() {
|
|
var status string
|
|
var count int
|
|
if err := rows.Scan(&status, &count); err != nil {
|
|
return nil, err
|
|
}
|
|
counts[status] = count
|
|
}
|
|
return counts, nil
|
|
}
|
|
|
|
// ReprocessAllJobs wipes all job analysis results and re-queues PROCESS tasks for all existing raw jobs.
|
|
func (db *DB) ReprocessAllJobs() error {
|
|
tx, err := db.Begin()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
|
|
// 1. Delete all rows from job_analyses
|
|
if _, err := tx.Exec(`DELETE FROM job_analyses`); err != nil {
|
|
return fmt.Errorf("failed to clear job_analyses: %w", err)
|
|
}
|
|
|
|
// 2. Delete all existing PROCESS and ANALYZE tasks
|
|
if _, err := tx.Exec(`DELETE FROM tasks WHERE task_type IN ('PROCESS', 'ANALYZE')`); err != nil {
|
|
return fmt.Errorf("failed to clear old tasks: %w", err)
|
|
}
|
|
|
|
// 3. Queue up new PROCESS tasks for all raw jobs
|
|
rows, err := tx.Query(`SELECT id FROM raw_jobs`)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get raw jobs: %w", err)
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
|
|
var rawJobIDs []int64
|
|
for rows.Next() {
|
|
var rawJobID int64
|
|
if err := rows.Scan(&rawJobID); err != nil {
|
|
return err
|
|
}
|
|
rawJobIDs = append(rawJobIDs, rawJobID)
|
|
}
|
|
rows.Close() // Explicitly close to release any read locks
|
|
|
|
// 4. Enqueue PROCESS task for each job
|
|
for _, rawJobID := range rawJobIDs {
|
|
payloadJSON := fmt.Sprintf(`{"raw_job_id":%d}`, rawJobID)
|
|
if _, err := tx.Exec(`INSERT INTO tasks (task_type, payload, status) VALUES ('PROCESS', ?, 'PENDING')`, payloadJSON); err != nil {
|
|
return fmt.Errorf("failed to enqueue process task for job %d: %w", rawJobID, err)
|
|
}
|
|
}
|
|
|
|
return tx.Commit()
|
|
}
|