package db import ( "database/sql" "encoding/json" "fmt" "github.com/hireme/hireme/internal/logger" ) // QueueStats holds task queue counters type QueueStats struct { Pending int `json:"pending"` Running int `json:"running"` Failed int `json:"failed"` Completed int `json:"completed"` } // GetQueueStats returns counts for each task status func (db *DB) GetQueueStats() (*QueueStats, error) { query := ` SELECT COALESCE(SUM(CASE WHEN status = 'PENDING' THEN 1 ELSE 0 END), 0), COALESCE(SUM(CASE WHEN status = 'RUNNING' THEN 1 ELSE 0 END), 0), COALESCE(SUM(CASE WHEN status = 'FAILED' THEN 1 ELSE 0 END), 0), COALESCE(SUM(CASE WHEN status = 'COMPLETED' THEN 1 ELSE 0 END), 0) FROM tasks ` var stats QueueStats err := db.QueryRow(query).Scan(&stats.Pending, &stats.Running, &stats.Failed, &stats.Completed) if err != nil { return nil, err } return &stats, nil } // 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 { logger.Log.Errorw("FailTask: failed to get retry_count", "component", "DB", "task_id", id, "error", 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() }