package-updater/internal/master/store/jobs.go

package store

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"

	"github.com/shcizo/package-updater/internal/master/model"
	"github.com/shcizo/package-updater/internal/protocol"
)

const lostJobMessage = "lost: the agent did not know this job when it reconnected"

// failActiveJobsSQL fails every queued or running job of agent $1 with
// message $2 at time $3.
const failActiveJobsSQL = `
UPDATE update_jobs SET status = 'failed', error = $2, finished_at = $3
WHERE agent_id = $1 AND status IN ('queued', 'running')`

const jobSelectSQL = `
SELECT j.id, j.service_id, s.kind, s.compose_project, s.compose_service, s.swarm_service, s.container_name,
       j.agent_id, COALESCE(a.name, ''), j.from_ref, j.to_ref, j.trigger, j.requested_by, j.status,
       %s, j.output_truncated, j.error, j.created_at, j.started_at, j.finished_at
FROM update_jobs j
JOIN services s ON s.id = j.service_id
LEFT JOIN agents a ON a.id = j.agent_id`

func scanJob(row pgx.Row) (model.Job, error) {
	var j model.Job
	var svc model.Service
	var kind, trigger, status string
	err := row.Scan(&j.ID, &j.ServiceID, &kind, &svc.ComposeProject, &svc.ComposeService, &svc.SwarmService, &svc.ContainerName,
		&j.AgentID, &j.AgentName, &j.FromRef, &j.ToRef, &trigger, &j.RequestedBy, &status,
		&j.Output, &j.OutputTruncated, &j.Error, &j.CreatedAt, &j.StartedAt, &j.FinishedAt)
	svc.Kind = protocol.Kind(kind)
	j.ServiceName = svc.DisplayName()
	j.Trigger = model.JobTrigger(trigger)
	j.Status = model.JobStatus(status)
	return j, err
}

func (s *Store) getJob(ctx context.Context, id int64) (model.Job, error) {
	j, err := scanJob(s.pool.QueryRow(ctx, fmt.Sprintf(jobSelectSQL, "j.output")+` WHERE j.id = $1`, id))
	if errors.Is(err, pgx.ErrNoRows) {
		return model.Job{}, model.ErrNotFound
	}
	if err != nil {
		return model.Job{}, fmt.Errorf("get job: %w", err)
	}
	return j, nil
}

// CreateJob inserts a queued job. If the service already has a queued or
// running job, it returns that job and model.ErrJobActive (detected via the
// unique index, so it holds under concurrent calls).
func (s *Store) CreateJob(ctx context.Context, j model.NewJob) (model.Job, error) {
	var id int64
	err := s.pool.QueryRow(ctx, `
		INSERT INTO update_jobs (service_id, agent_id, from_ref, to_ref, trigger, requested_by, status, created_at)
		VALUES ($1, $2, $3, $4, $5, $6, 'queued', $7)
		RETURNING id`,
		j.ServiceID, j.AgentID, j.FromRef, j.ToRef, string(j.Trigger), j.RequestedBy, j.CreatedAt).Scan(&id)
	var pgErr *pgconn.PgError
	if errors.As(err, &pgErr) && pgErr.Code == "23505" && pgErr.ConstraintName == "update_jobs_one_active" {
		active, err := scanJob(s.pool.QueryRow(ctx, fmt.Sprintf(jobSelectSQL, "j.output")+
			` WHERE j.service_id = $1 AND j.status IN ('queued', 'running')`, j.ServiceID))
		if err != nil {
			return model.Job{}, fmt.Errorf("find active job: %w", err)
		}
		return active, model.ErrJobActive
	}
	if err != nil {
		return model.Job{}, fmt.Errorf("create job: %w", err)
	}
	return s.getJob(ctx, id)
}

// FailJob sets a queued or running job to failed with msg and finished_at =
// at. Returns false (no error) if the job is no longer active.
func (s *Store) FailJob(ctx context.Context, id int64, msg string, at time.Time) (bool, error) {
	tag, err := s.pool.Exec(ctx, `
		UPDATE update_jobs SET status = 'failed', error = $2, finished_at = $3
		WHERE id = $1 AND status IN ('queued', 'running')`, id, msg, at)
	if err != nil {
		return false, fmt.Errorf("fail job: %w", err)
	}
	return tag.RowsAffected() == 1, nil
}

// GetJob returns one job with its output. Returns model.ErrNotFound if missing.
func (s *Store) GetJob(ctx context.Context, id int64) (model.Job, error) {
	return s.getJob(ctx, id)
}

// ListJobs returns jobs newest first, without Output.
func (s *Store) ListJobs(ctx context.Context, f model.JobFilter) ([]model.Job, error) {
	limit := f.Limit
	if limit <= 0 {
		limit = 100
	}
	rows, err := s.pool.Query(ctx, fmt.Sprintf(jobSelectSQL, "''")+`
		WHERE ($1 = 0 OR j.service_id = $1)
		ORDER BY j.id DESC
		LIMIT $2`, f.ServiceID, limit)
	if err != nil {
		return nil, fmt.Errorf("list jobs: %w", err)
	}
	defer rows.Close()
	var out []model.Job
	for rows.Next() {
		j, err := scanJob(rows)
		if err != nil {
			return nil, fmt.Errorf("scan job: %w", err)
		}
		out = append(out, j)
	}
	return out, rows.Err()
}

// FailLostJobs fails every queued or running job sent to agentID whose ID is
// not in known, with a "lost" error. Returns the IDs it failed.
func (s *Store) FailLostJobs(ctx context.Context, agentID int64, known []int64, at time.Time) ([]int64, error) {
	if known == nil {
		known = []int64{}
	}
	rows, err := s.pool.Query(ctx, `
		UPDATE update_jobs SET status = 'failed', error = $3, finished_at = $4
		WHERE agent_id = $1 AND status IN ('queued', 'running') AND NOT (id = ANY($2))
		RETURNING id`, agentID, known, lostJobMessage, at)
	if err != nil {
		return nil, fmt.Errorf("fail lost jobs: %w", err)
	}
	defer rows.Close()
	var ids []int64
	for rows.Next() {
		var id int64
		if err := rows.Scan(&id); err != nil {
			return nil, fmt.Errorf("scan lost job: %w", err)
		}
		ids = append(ids, id)
	}
	return ids, rows.Err()
}

// JobProgress appends chunk to the output of a queued or running job sent to
// agentID, keeping the last protocol.MaxJobOutput characters, and moves a
// queued job to running (started_at = at if unset). Returns false if no such
// active job belongs to agentID.
func (s *Store) JobProgress(ctx context.Context, agentID, jobID int64, chunk string, at time.Time) (bool, error) {
	tag, err := s.pool.Exec(ctx, `
		UPDATE update_jobs SET
			status           = 'running',
			started_at       = COALESCE(started_at, $4),
			output_truncated = output_truncated OR length(output) + length($3::text) > $5,
			output           = right(output || $3::text, $5)
		WHERE id = $1 AND agent_id = $2 AND status IN ('queued', 'running')`,
		jobID, agentID, chunk, at, protocol.MaxJobOutput)
	if err != nil {
		return false, fmt.Errorf("job progress: %w", err)
	}
	return tag.RowsAffected() == 1, nil
}

// FinishJob records the agent's result for a queued or running job sent to
// agentID: status, error, output (replacing what progress assembled),
// started_at (kept if already set), finished_at; to_ref = f.Ref when set. On
// success it also sets the service's desired_ref = f.Ref, in the same
// transaction. Returns false if no such active job belongs to agentID.
func (s *Store) FinishJob(ctx context.Context, agentID, jobID int64, f model.JobFinish) (bool, error) {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return false, fmt.Errorf("finish job: begin: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()

	var serviceID int64
	err = tx.QueryRow(ctx, `
		UPDATE update_jobs SET
			status = $3, error = $4, output = $5, output_truncated = $6,
			to_ref = CASE WHEN $7::text <> '' THEN $7::text ELSE to_ref END,
			started_at = COALESCE(started_at, $8), finished_at = $9
		WHERE id = $1 AND agent_id = $2 AND status IN ('queued', 'running')
		RETURNING service_id`,
		jobID, agentID, string(f.Status), f.Error, f.Output, f.OutputTruncated, f.Ref, f.StartedAt, f.FinishedAt).Scan(&serviceID)
	if errors.Is(err, pgx.ErrNoRows) {
		return false, nil
	}
	if err != nil {
		return false, fmt.Errorf("finish job: %w", err)
	}
	if f.Status == model.JobSucceeded {
		if _, err := tx.Exec(ctx, `UPDATE services SET desired_ref = $2 WHERE id = $1`, serviceID, f.Ref); err != nil {
			return false, fmt.Errorf("finish job: desired ref: %w", err)
		}
	}
	if err := tx.Commit(ctx); err != nil {
		return false, fmt.Errorf("finish job: commit: %w", err)
	}
	return true, nil
}