forgebot/internal/database/tasks.go

package database

import (
	"context"
	"time"

	"github.com/jackc/pgx/v5"

	"git.unkin.net/unkin/forgebot/pkg/models"
)

func (db *DB) CreateTask(ctx context.Context, req models.CreateTaskRequest) (*models.Task, error) {
	task := &models.Task{
		Command:      req.Command,
		Skill:        req.Skill,
		Repository:   req.Repository,
		Ref:          req.Ref,
		IssueNumber:  req.IssueNumber,
		PRNumber:     req.PRNumber,
		CommentID:    req.CommentID,
		Body:         req.Body,
		Author:       req.Author,
		ExtraTools:   req.ExtraTools,
		ParentTaskID: req.ParentTaskID,
		PoolRef:      req.PoolRef,
		Status:       models.StatusPending,
	}
	if task.ExtraTools == nil {
		task.ExtraTools = []string{}
	}

	err := db.Pool.QueryRow(ctx, `
		INSERT INTO tasks (
			parent_task_id, command, skill, repository, ref,
			issue_number, pr_number, comment_id, body, author,
			extra_tools, pool_ref
		) VALUES (
			NULLIF($1, ''), $2, $3, $4, $5,
			$6, $7, $8, $9, $10,
			$11, $12
		) RETURNING id, created_at`,
		task.ParentTaskID, task.Command, task.Skill, task.Repository, task.Ref,
		task.IssueNumber, task.PRNumber, task.CommentID, task.Body, task.Author,
		task.ExtraTools, task.PoolRef,
	).Scan(&task.ID, &task.CreatedAt)
	if err != nil {
		return nil, err
	}
	return task, nil
}

func (db *DB) GetTask(ctx context.Context, id string) (*models.Task, error) {
	task := &models.Task{}
	var parentID *string
	var startedAt, completedAt *time.Time

	err := db.Pool.QueryRow(ctx, `
		SELECT id, parent_task_id, command, skill, repository, ref,
			issue_number, pr_number, comment_id, body, author,
			extra_tools, status, pool_ref, job_name, result, error_message,
			created_at, started_at, completed_at
		FROM tasks WHERE id = $1`, id,
	).Scan(
		&task.ID, &parentID, &task.Command, &task.Skill, &task.Repository, &task.Ref,
		&task.IssueNumber, &task.PRNumber, &task.CommentID, &task.Body, &task.Author,
		&task.ExtraTools, &task.Status, &task.PoolRef, &task.JobName, &task.Result, &task.ErrorMessage,
		&task.CreatedAt, &startedAt, &completedAt,
	)
	if err != nil {
		return nil, err
	}
	if parentID != nil {
		task.ParentTaskID = *parentID
	}
	task.StartedAt = startedAt
	task.CompletedAt = completedAt
	return task, nil
}

func (db *DB) ListPendingTasks(ctx context.Context) ([]models.Task, error) {
	return db.listTasksByStatus(ctx, string(models.StatusPending))
}

func (db *DB) listTasksByStatus(ctx context.Context, status string) ([]models.Task, error) {
	rows, err := db.Pool.Query(ctx, `
		SELECT id, parent_task_id, command, skill, repository, ref,
			issue_number, pr_number, comment_id, body, author,
			extra_tools, status, pool_ref, job_name, result, error_message,
			created_at, started_at, completed_at
		FROM tasks WHERE status = $1
		ORDER BY created_at ASC`, status,
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return scanTasks(rows)
}

func (db *DB) ListTasks(ctx context.Context, status string, repository string) ([]models.Task, error) {
	query := `SELECT id, parent_task_id, command, skill, repository, ref,
		issue_number, pr_number, comment_id, body, author,
		extra_tools, status, pool_ref, job_name, result, error_message,
		created_at, started_at, completed_at
		FROM tasks WHERE 1=1`
	args := []any{}
	argIdx := 1

	if status != "" {
		query += " AND status = $" + itoa(argIdx)
		args = append(args, status)
		argIdx++
	}
	if repository != "" {
		query += " AND repository = $" + itoa(argIdx)
		args = append(args, repository)
		argIdx++
	}
	query += " ORDER BY created_at DESC LIMIT 100"

	rows, err := db.Pool.Query(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return scanTasks(rows)
}

func (db *DB) UpdateTaskStatus(ctx context.Context, id string, req models.UpdateTaskRequest) error {
	if req.Status == models.StatusRunning {
		_, err := db.Pool.Exec(ctx, `
			UPDATE tasks SET status = $2, job_name = COALESCE(NULLIF($3, ''), job_name), started_at = NOW()
			WHERE id = $1`, id, req.Status, req.JobName)
		return err
	}
	if req.Status == models.StatusSucceeded || req.Status == models.StatusFailed {
		_, err := db.Pool.Exec(ctx, `
			UPDATE tasks SET status = $2, result = COALESCE(NULLIF($3, ''), result),
				error_message = COALESCE(NULLIF($4, ''), error_message), completed_at = NOW()
			WHERE id = $1`, id, req.Status, req.Message, req.ErrorMessage)
		return err
	}
	_, err := db.Pool.Exec(ctx, `UPDATE tasks SET status = $2 WHERE id = $1`, id, req.Status)
	return err
}

func scanTasks(rows pgx.Rows) ([]models.Task, error) {
	var tasks []models.Task
	for rows.Next() {
		var task models.Task
		var parentID *string
		var startedAt, completedAt *time.Time

		err := rows.Scan(
			&task.ID, &parentID, &task.Command, &task.Skill, &task.Repository, &task.Ref,
			&task.IssueNumber, &task.PRNumber, &task.CommentID, &task.Body, &task.Author,
			&task.ExtraTools, &task.Status, &task.PoolRef, &task.JobName, &task.Result, &task.ErrorMessage,
			&task.CreatedAt, &startedAt, &completedAt,
		)
		if err != nil {
			return nil, err
		}
		if parentID != nil {
			task.ParentTaskID = *parentID
		}
		task.StartedAt = startedAt
		task.CompletedAt = completedAt
		tasks = append(tasks, task)
	}
	return tasks, rows.Err()
}

func itoa(i int) string {
	return string(rune('0'+i)) + ""
}