fetch_ml/internal/queue/sqlite/queue.go

// Package sqlite provides a SQLite-based queue implementation
package sqlite

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/jfraeys/fetch_ml/internal/domain"
	_ "github.com/mattn/go-sqlite3"
)

// Queue implements a SQLite-based task queue
type Queue struct {
	db     *sql.DB
	ctx    context.Context
	cancel context.CancelFunc
}

// NewQueue creates a new SQLite queue instance
func NewQueue(path string) (*Queue, error) {
	if path == "" {
		return nil, fmt.Errorf("sqlite queue path is required")
	}

	db, err := sql.Open("sqlite3", fmt.Sprintf("file:%s?_busy_timeout=5000&_foreign_keys=on", path))
	if err != nil {
		return nil, err
	}

	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(1)

	ctx, cancel := context.WithCancel(context.Background())
	q := &Queue{db: db, ctx: ctx, cancel: cancel}

	if err := q.initSchema(); err != nil {
		_ = db.Close()
		cancel()
		return nil, err
	}

	go q.leaseReclaimer()
	go q.kvJanitor()
	return q, nil
}

// Close closes the queue
func (q *Queue) Close() error {
	q.cancel()
	return q.db.Close()
}

// initSchema initializes the database schema
func (q *Queue) initSchema() error {
	stmts := []string{
		"PRAGMA journal_mode=WAL;",
		"PRAGMA synchronous=NORMAL;",
		`CREATE TABLE IF NOT EXISTS tasks (
			id TEXT PRIMARY KEY,
			job_name TEXT,
			status TEXT,
			priority INTEGER,
			created_at INTEGER,
			updated_at INTEGER,
			payload BLOB
		);`,
		`CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status);`,
		`CREATE INDEX IF NOT EXISTS idx_tasks_priority ON tasks(priority);`,
		`CREATE TABLE IF NOT EXISTS kv (
			key TEXT PRIMARY KEY,
			value BLOB,
			expires_at INTEGER
		);`,
	}

	for _, stmt := range stmts {
		if _, err := q.db.Exec(stmt); err != nil {
			return fmt.Errorf("failed to execute schema statement: %w", err)
		}
	}
	return nil
}

// AddTask adds a task to the queue
func (q *Queue) AddTask(task *domain.Task) error {
	if task == nil {
		return errors.New("task is nil")
	}
	if task.ID == "" {
		return errors.New("task ID is required")
	}

	payload, err := json.Marshal(task)
	if err != nil {
		return fmt.Errorf("failed to marshal task: %w", err)
	}

	_, err = q.db.Exec(
		"INSERT INTO tasks (id, job_name, status, priority, created_at, updated_at, payload) VALUES (?, ?, ?, ?, ?, ?, ?)",
		task.ID,
		task.JobName,
		task.Status,
		task.Priority,
		time.Now().Unix(),
		time.Now().Unix(),
		payload,
	)
	if err != nil {
		return fmt.Errorf("failed to insert task: %w", err)
	}

	return nil
}

// GetTask retrieves a task by ID
func (q *Queue) GetTask(id string) (*domain.Task, error) {
	if id == "" {
		return nil, errors.New("task ID is required")
	}

	var payload []byte
	err := q.db.QueryRow("SELECT payload FROM tasks WHERE id = ?", id).Scan(&payload)
	if err == sql.ErrNoRows {
		return nil, fmt.Errorf("task not found: %s", id)
	}
	if err != nil {
		return nil, fmt.Errorf("failed to query task: %w", err)
	}

	var task domain.Task
	if err := json.Unmarshal(payload, &task); err != nil {
		return nil, fmt.Errorf("failed to unmarshal task: %w", err)
	}

	return &task, nil
}

// ListTasks lists all tasks in the queue
func (q *Queue) ListTasks() ([]*domain.Task, error) {
	rows, err := q.db.Query("SELECT payload FROM tasks ORDER BY priority DESC, created_at ASC")
	if err != nil {
		return nil, fmt.Errorf("failed to query tasks: %w", err)
	}
	defer rows.Close()

	var tasks []*domain.Task
	for rows.Next() {
		var payload []byte
		if err := rows.Scan(&payload); err != nil {
			continue
		}

		var task domain.Task
		if err := json.Unmarshal(payload, &task); err != nil {
			continue
		}
		tasks = append(tasks, &task)
	}

	return tasks, nil
}

// CancelTask cancels a task
func (q *Queue) CancelTask(id string) error {
	if id == "" {
		return errors.New("task ID is required")
	}

	_, err := q.db.Exec("DELETE FROM tasks WHERE id = ? AND status = 'pending'", id)
	if err != nil {
		return fmt.Errorf("failed to cancel task: %w", err)
	}

	return nil
}

// UpdateTask updates a task
func (q *Queue) UpdateTask(task *domain.Task) error {
	if task == nil || task.ID == "" {
		return errors.New("task is nil or missing ID")
	}

	payload, err := json.Marshal(task)
	if err != nil {
		return fmt.Errorf("failed to marshal task: %w", err)
	}

	_, err = q.db.Exec(
		"UPDATE tasks SET job_name = ?, status = ?, priority = ?, updated_at = ?, payload = ? WHERE id = ?",
		task.JobName,
		task.Status,
		task.Priority,
		time.Now().Unix(),
		payload,
		task.ID,
	)
	if err != nil {
		return fmt.Errorf("failed to update task: %w", err)
	}

	return nil
}

// leaseReclaimer periodically reclaims expired leases
func (q *Queue) leaseReclaimer() {
	ticker := time.NewTicker(1 * time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-q.ctx.Done():
			return
		case <-ticker.C:
			// Reclaim expired leases
		}
	}
}

// kvJanitor periodically cleans up expired KV entries
func (q *Queue) kvJanitor() {
	ticker := time.NewTicker(5 * time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-q.ctx.Done():
			return
		case <-ticker.C:
			// Clean up expired KV entries
		}
	}
}