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

package store

import (
	"context"
	"errors"
	"fmt"

	"github.com/jackc/pgx/v5"

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

// moveCandidatesSQL pairs a removed Compose service with a live one of the
// same project and service name on another agent. A pair counts only when it
// is the single one for both rows, and the live row is new (newer than the
// removed one, never updated, no jobs) while the removed one has something
// worth moving (a desired ref, auto mode or jobs).
const moveCandidatesSQL = `
WITH pairs AS (
	SELECT f.id AS from_id, t.id AS to_id
	FROM services f
	JOIN services t ON t.kind = 'compose' AND t.removed_at IS NULL AND NOT t.self
	 AND t.compose_project = f.compose_project AND t.compose_service = f.compose_service
	 AND t.agent_id IS DISTINCT FROM f.agent_id
	WHERE f.kind = 'compose' AND f.removed_at IS NOT NULL AND NOT f.self
), uniq AS (
	SELECT from_id, to_id FROM pairs p
	WHERE (SELECT count(*) FROM pairs q WHERE q.from_id = p.from_id) = 1
	  AND (SELECT count(*) FROM pairs q WHERE q.to_id   = p.to_id)   = 1
)
SELECT u.from_id, u.to_id, COALESCE(fa.name, ''), COALESCE(ta.name, '')
FROM uniq u
JOIN services f ON f.id = u.from_id
JOIN services t ON t.id = u.to_id
LEFT JOIN agents fa ON fa.id = f.agent_id
LEFT JOIN agents ta ON ta.id = t.agent_id
WHERE t.first_seen_at > f.first_seen_at
  AND t.desired_ref = ''
  AND NOT EXISTS (SELECT 1 FROM update_jobs j WHERE j.service_id = t.id)
  AND (f.desired_ref <> '' OR f.update_mode = 'auto'
       OR EXISTS (SELECT 1 FROM update_jobs j WHERE j.service_id = f.id))
ORDER BY u.to_id`

// MoveCandidates returns the removed Compose services that look like they
// moved to another agent; see model.Move. It only suggests: nothing is merged.
func (s *Store) MoveCandidates(ctx context.Context) ([]model.Move, error) {
	rows, err := s.pool.Query(ctx, moveCandidatesSQL)
	if err != nil {
		return nil, fmt.Errorf("move candidates: %w", err)
	}
	defer rows.Close()
	var out []model.Move
	for rows.Next() {
		var m model.Move
		if err := rows.Scan(&m.FromID, &m.ToID, &m.FromAgent, &m.ToAgent); err != nil {
			return nil, fmt.Errorf("scan move: %w", err)
		}
		out = append(out, m)
	}
	return out, rows.Err()
}

// MergeSources returns up to 50 removed services that could be merged into
// the service toID: the same kind (Compose or Swarm), not the updater itself,
// those with the same name first, then the most recently removed. Returns
// model.ErrNotFound if toID does not exist.
func (s *Store) MergeSources(ctx context.Context, toID int64) ([]model.Service, error) {
	to, err := s.GetService(ctx, toID)
	if err != nil {
		return nil, err
	}
	if to.Kind == protocol.KindContainer {
		return nil, nil
	}
	rows, err := s.pool.Query(ctx, `
		SELECT `+serviceColumns+`
		FROM services s LEFT JOIN agents a ON a.id = s.agent_id
		WHERE s.removed_at IS NOT NULL AND s.kind = $1 AND NOT s.self AND s.id <> $2
		ORDER BY (s.compose_project, s.compose_service, s.swarm_service) = ($3, $4, $5) DESC,
		         s.removed_at DESC, s.id DESC
		LIMIT 50`,
		string(to.Kind), toID, to.ComposeProject, to.ComposeService, to.SwarmService)
	if err != nil {
		return nil, fmt.Errorf("merge sources: %w", err)
	}
	defer rows.Close()
	var out []model.Service
	for rows.Next() {
		m, err := scanService(rows)
		if err != nil {
			return nil, fmt.Errorf("scan service: %w", err)
		}
		out = append(out, m)
	}
	return out, rows.Err()
}

// mergeRow is what MergeService needs to know about one of the two rows.
type mergeRow struct {
	kind    string
	agentID *int64
	removed bool
	self    bool
}

// MergeService moves the live service toID onto the removed service fromID:
// the old row keeps its id, first_seen_at, update mode and desired ref, takes
// the location and inventory of the live row (and its available suggestion),
// becomes live, and the live row is deleted. The jobs of the live row move to
// the old one. This is the only place a services row is deleted.
//
// It returns *model.MergeRefusedError when the pair cannot be merged,
// model.ErrJobActive when either row has a queued or running job,
// model.ErrRevoked when the live row's agent is revoked, model.ErrNotFound
// when a row does not exist, and model.ErrConflict when the live row moved to
// another agent while the merge started.
//
// Locks are taken in the order inventory writes use: the agent row first,
// then the service rows by id. The live row is deleted before the old one is
// updated because location_key is unique.
func (s *Store) MergeService(ctx context.Context, fromID, toID int64) error {
	if fromID == toID {
		return &model.MergeRefusedError{Reason: model.ReasonMergeSame}
	}
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("begin merge: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()

	var toAgent *int64
	err = tx.QueryRow(ctx, `SELECT agent_id FROM services WHERE id = $1`, toID).Scan(&toAgent)
	if errors.Is(err, pgx.ErrNoRows) {
		return model.ErrNotFound
	}
	if err != nil {
		return fmt.Errorf("merge: find new service: %w", err)
	}
	if toAgent != nil {
		if err := lockActiveAgent(ctx, tx, *toAgent); err != nil {
			return err
		}
	}

	rows, err := tx.Query(ctx, `
		SELECT id, kind, agent_id, removed_at IS NOT NULL, self FROM services
		WHERE id IN ($1, $2) ORDER BY id FOR UPDATE`, fromID, toID)
	if err != nil {
		return fmt.Errorf("merge: lock services: %w", err)
	}
	locked := map[int64]mergeRow{}
	for rows.Next() {
		var id int64
		var r mergeRow
		if err := rows.Scan(&id, &r.kind, &r.agentID, &r.removed, &r.self); err != nil {
			rows.Close()
			return fmt.Errorf("merge: scan service: %w", err)
		}
		locked[id] = r
	}
	rows.Close()
	if err := rows.Err(); err != nil {
		return fmt.Errorf("merge: lock services: %w", err)
	}
	from, okFrom := locked[fromID]
	to, okTo := locked[toID]
	if !okFrom || !okTo {
		return model.ErrNotFound
	}
	if (to.agentID == nil) != (toAgent == nil) || (toAgent != nil && *to.agentID != *toAgent) {
		return model.ErrConflict
	}
	if reason := mergeRefusal(from, to); reason != "" {
		return &model.MergeRefusedError{Reason: reason}
	}

	var active bool
	if err := tx.QueryRow(ctx, `
		SELECT EXISTS (SELECT 1 FROM update_jobs
		               WHERE service_id IN ($1, $2) AND status IN ('queued', 'running'))`,
		fromID, toID).Scan(&active); err != nil {
		return fmt.Errorf("merge: check active jobs: %w", err)
	}
	if active {
		return model.ErrJobActive
	}

	if _, err := tx.Exec(ctx, `UPDATE update_jobs SET service_id = $1 WHERE service_id = $2`, fromID, toID); err != nil {
		return fmt.Errorf("merge: move jobs: %w", err)
	}
	var t model.Service
	var lastTagDigest string
	err = tx.QueryRow(ctx, `
		SELECT agent_id, location_key, compose_project, compose_service, working_dir, config_files, built,
		       swarm_cluster_id, swarm_service, container_name, image_repo, actual_ref, actual_digest, status, self,
		       updated_at, available_ref, available_digest, available_source, available_at, last_poll_error,
		       last_tag_digest
		FROM services WHERE id = $1`, toID).Scan(
		&t.AgentID, &t.LocationKey, &t.ComposeProject, &t.ComposeService, &t.WorkingDir, &t.ConfigFiles, &t.Built,
		&t.SwarmClusterID, &t.SwarmService, &t.ContainerName, &t.ImageRepo, &t.ActualRef, &t.ActualDigest, &t.Status, &t.Self,
		&t.UpdatedAt, &t.AvailableRef, &t.AvailableDigest, &t.AvailableSource, &t.AvailableAt, &t.LastPollError,
		&lastTagDigest)
	if err != nil {
		return fmt.Errorf("merge: read new service: %w", err)
	}
	if _, err := tx.Exec(ctx, `DELETE FROM services WHERE id = $1`, toID); err != nil {
		return fmt.Errorf("merge: delete new service: %w", err)
	}
	if _, err := tx.Exec(ctx, `
		UPDATE services SET
			agent_id = $2, location_key = $3, compose_project = $4, compose_service = $5, working_dir = $6,
			config_files = $7, built = $8, swarm_cluster_id = $9, swarm_service = $10, container_name = $11,
			image_repo = $12, actual_ref = $13, actual_digest = $14, status = $15, self = $16, updated_at = $17,
			available_ref = $18, available_digest = $19, available_source = $20, available_at = $21,
			last_poll_error = $22,
			last_tag_digest = CASE WHEN last_tag_digest <> '' THEN last_tag_digest ELSE $23 END,
			removed_at = NULL, awaiting_inventory = false
		WHERE id = $1`,
		fromID, t.AgentID, t.LocationKey, t.ComposeProject, t.ComposeService, t.WorkingDir,
		t.ConfigFiles, t.Built, t.SwarmClusterID, t.SwarmService, t.ContainerName,
		t.ImageRepo, t.ActualRef, t.ActualDigest, t.Status, t.Self, t.UpdatedAt,
		t.AvailableRef, t.AvailableDigest, t.AvailableSource, t.AvailableAt,
		t.LastPollError, lastTagDigest); err != nil {
		return fmt.Errorf("merge: update old service: %w", err)
	}
	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("commit merge: %w", err)
	}
	return nil
}

// mergeRefusal returns why from (the old row) cannot be merged with to (the
// live one), or "".
func mergeRefusal(from, to mergeRow) string {
	switch {
	case !from.removed:
		return model.ReasonMergeLiveOld
	case to.removed:
		return model.ReasonMergeGoneNew
	case from.kind != to.kind:
		return model.ReasonMergeKind
	case from.kind == string(protocol.KindContainer):
		return model.ReasonMergeContainer
	case from.self || to.self:
		return model.ReasonSelf
	}
	return ""
}