restic-manager/internal/restic/restore.go

package restic

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os/exec"
	"strings"
)

// RestoreStatus mirrors the JSON `status` lines `restic restore --json`
// emits while restoring. Field names track restic's wire format; we
// project a subset (the rest are cosmetic).
type RestoreStatus struct {
	MessageType    string  `json:"message_type"`
	SecondsElapsed int64   `json:"seconds_elapsed"`
	PercentDone    float64 `json:"percent_done"`
	TotalFiles     int64   `json:"total_files"`
	FilesRestored  int64   `json:"files_restored"`
	FilesSkipped   int64   `json:"files_skipped"`
	TotalBytes     int64   `json:"total_bytes"`
	BytesRestored  int64   `json:"bytes_restored"`
	BytesSkipped   int64   `json:"bytes_skipped"`
}

// RestoreSummary is the final summary line emitted after a successful
// restore. Newer restic prints it; older clients leave us with no
// summary, in which case the agent skips the stats and the live UI
// just sees percent reach 100%.
type RestoreSummary struct {
	MessageType    string `json:"message_type"`
	SecondsElapsed int64  `json:"seconds_elapsed"`
	TotalFiles     int64  `json:"total_files"`
	FilesRestored  int64  `json:"files_restored"`
	FilesSkipped   int64  `json:"files_skipped"`
	TotalBytes     int64  `json:"total_bytes"`
	BytesRestored  int64  `json:"bytes_restored"`
	BytesSkipped   int64  `json:"bytes_skipped"`
}

// RunRestore executes `restic restore <snapshotID> --target <dir>
// [--include <p>...]` with --json and pumps progress events into
// handle. paths is the operator-selected list (each becomes an
// `--include` flag); preserveOwner controls --no-ownership.
//
// inPlace toggles target semantics:
//   - true  → target is "/" and ownership is preserved
//   - false → target is targetDir and --no-ownership is passed
//
// targetDir is created on demand by restic itself.
func (e Env) RunRestore(ctx context.Context, snapshotID string, paths []string, inPlace bool, targetDir string, handle LineHandler) (*RestoreSummary, error) {
	if snapshotID == "" {
		return nil, fmt.Errorf("restic restore: snapshot id required")
	}
	if !inPlace && targetDir == "" {
		return nil, fmt.Errorf("restic restore: target dir required for non-in-place restore")
	}

	args := []string{"restore", "--json", snapshotID}
	target := targetDir
	if inPlace {
		target = "/"
	}
	args = append(args, "--target", target)
	if !inPlace {
		args = append(args, "--no-ownership")
	}
	for _, p := range paths {
		args = append(args, "--include", p)
	}

	cmd := e.resticCmd(ctx, args...)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("restic restore: stdout pipe: %w", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("restic restore: stderr pipe: %w", err)
	}

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("restic restore: start: %w", err)
	}

	var summary *RestoreSummary
	done := make(chan error, 2)
	go func() { done <- pumpRestoreStdout(stdout, handle, &summary) }()
	go func() { done <- pumpStderr(stderr, handle) }()
	for i := 0; i < 2; i++ {
		if err := <-done; err != nil && handle != nil {
			handle("event", fmt.Sprintf("pump error: %v", err), nil)
		}
	}
	werr := cmd.Wait()
	if werr != nil {
		var ee *exec.ExitError
		if errors.As(werr, &ee) {
			return summary, fmt.Errorf("restic restore: exit %d", ee.ExitCode())
		}
		return summary, fmt.Errorf("restic restore: %w", werr)
	}
	return summary, nil
}

// pumpRestoreStdout is the restore variant of pumpStdout: it emits
// `event` lines for the parsed status/summary objects (so the runner
// can shape them into job.progress) and forwards everything else as
// stdout — but unlike backup we include the raw status JSON in
// log.stream too because restore is short and the live log audience
// genuinely benefits from the per-file traffic. Actually — we mirror
// backup's behavior and DROP raw status lines from log.stream
// (they'd drown the log on a fast restore); the progress envelope
// covers them.
func pumpRestoreStdout(r io.Reader, handle LineHandler, summary **RestoreSummary) error {
	scanner := bufio.NewScanner(r)
	scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
	for scanner.Scan() {
		line := scanner.Text()
		if handle == nil {
			continue
		}
		if !strings.HasPrefix(line, "{") {
			handle("stdout", line, nil)
			continue
		}
		var probe struct {
			MessageType string `json:"message_type"`
		}
		if err := json.Unmarshal([]byte(line), &probe); err != nil {
			handle("stdout", line, nil)
			continue
		}
		switch probe.MessageType {
		case "status":
			var ev RestoreStatus
			if json.Unmarshal([]byte(line), &ev) == nil {
				// Don't tee status lines to log.stream — too chatty.
				handle("event", line, ev)
				continue
			}
		case "summary":
			var ev RestoreSummary
			if json.Unmarshal([]byte(line), &ev) == nil {
				if summary != nil {
					s := ev
					*summary = &s
				}
				handle("event", line, ev)
				continue
			}
		case "verbose_status":
			handle("event", line, nil)
			continue
		}
		handle("stdout", line, nil)
	}
	return scanner.Err()
}

// RunDiff executes `restic diff --json <a> <b>` and forwards every
// line to handle as stdout. Restic emits per-line "change" objects
// plus a final "statistics" object; we don't parse them server-side —
// the operator reads the raw output on the live job log page.
func (e Env) RunDiff(ctx context.Context, snapshotA, snapshotB string, handle LineHandler) error {
	if snapshotA == "" || snapshotB == "" {
		return fmt.Errorf("restic diff: two snapshot ids required")
	}
	cmd := e.resticCmd(ctx, "diff", "--json", snapshotA, snapshotB)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("restic diff: stdout pipe: %w", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("restic diff: stderr pipe: %w", err)
	}
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("restic diff: start: %w", err)
	}
	done := make(chan error, 2)
	// diff output isn't huge; pumpStderr-ish line-by-line forwarding
	// is fine.
	go func() {
		s := bufio.NewScanner(stdout)
		s.Buffer(make([]byte, 0, 64*1024), 1024*1024)
		for s.Scan() {
			if handle != nil {
				handle("stdout", s.Text(), nil)
			}
		}
		done <- s.Err()
	}()
	go func() { done <- pumpStderr(stderr, handle) }()
	for i := 0; i < 2; i++ {
		if err := <-done; err != nil && handle != nil {
			handle("event", fmt.Sprintf("pump error: %v", err), nil)
		}
	}
	werr := cmd.Wait()
	if werr != nil {
		var ee *exec.ExitError
		if errors.As(werr, &ee) {
			return fmt.Errorf("restic diff: exit %d", ee.ExitCode())
		}
		return fmt.Errorf("restic diff: %w", werr)
	}
	return nil
}