orbit/examples/file-conversion/main.go

// file-conversion 演示一个通用的文件格式转换任务：CSV 内容由 Worker 转为 JSON。
// 示例把小文件内容放进任务快照；生产大文件应传对象引用，并在业务侧保存稳定操作键。
package main

import (
	"context"
	"encoding/csv"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	orbit "git.yuntuo.hk/k2safe/orbit/sdk/go"
)

type conversionInput struct {
	// FileName 是输入文件名，仅用于结果追踪，不参与路径访问。
	FileName string `json:"file_name"`
	// Content 是示例中的小型文件正文；生产场景应替换为受控对象引用。
	Content string `json:"content"`
	// TargetFormat 固定为 json，演示处理器版本化能力。
	TargetFormat string `json:"target_format"`
}

type conversionOutput struct {
	// FileName 是转换后文件名。
	FileName string `json:"file_name"`
	// Format 是输出格式。
	Format string `json:"format"`
	// Content 是转换后的 JSON 文件正文。
	Content string `json:"content"`
	// Rows 是转换的记录数。
	Rows int `json:"rows"`
}

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func run() error {
	token, queue := os.Getenv("ORBIT_TOKEN"), os.Getenv("ORBIT_QUEUE_ID")
	if path := os.Getenv("ORBIT_CREDENTIALS_FILE"); path != "" {
		raw, err := os.ReadFile(path)
		if err != nil {
			return fmt.Errorf("read private credentials: %w", err)
		}
		var credentials struct {
			Token   string `json:"token"`
			QueueID string `json:"queue_id"`
		}
		if err := json.Unmarshal(raw, &credentials); err != nil {
			return errors.New("invalid credentials file")
		}
		if token == "" {
			token = credentials.Token
		}
		if queue == "" {
			queue = credentials.QueueID
		}
	}
	client, err := orbit.NewClient(os.Getenv("ORBIT_URL"), orbit.Options{Token: token})
	if err != nil {
		return err
	}
	runID, err := orbit.NewWorkerMessageID()
	if err != nil {
		return err
	}
	handler := "file.convert.csv-json." + runID
	worker, err := orbit.NewWorker(client, orbit.WorkerOptions{
		QueueID:         queue,
		Capacity:        1,
		ShutdownTimeout: 5 * time.Second,
		Handlers: []orbit.WorkerHandler{{
			Capability: orbit.WorkerCapability{Handler: handler, HandlerVersion: "v1", PayloadVersion: "v1"},
			Handle:     convertCSV,
		}},
	})
	if err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	payload, err := json.Marshal(conversionInput{
		FileName:     "records.csv",
		Content:      "name,amount\nalpha,2\nbeta,3\n",
		TargetFormat: "json",
	})
	if err != nil {
		return err
	}
	key := "file-conversion-" + runID
	accepted, err := client.SubmitTask(ctx, key, orbit.SubmitTaskRequest{
		QueueID:        queue,
		Handler:        handler,
		HandlerVersion: "v1",
		PayloadVersion: "v1",
		Payload:        payload,
		TaskTimeoutMS:  60000,
		Retry:          orbit.RetryPolicy{MaxAttempts: 1, AttemptTimeoutMS: 10000},
		RecoverySafety: orbit.RecoveryReconcileOnly,
	})
	if err != nil {
		return err
	}
	workerCtx, stop := context.WithCancel(ctx)
	done := make(chan error, 1)
	go func() { done <- worker.Run(workerCtx) }()
	waited := false
	defer func() {
		stop()
		if !waited {
			<-done
		}
	}()
	for {
		current, err := client.GetTask(ctx, accepted.Task.ID)
		if err != nil {
			return err
		}
		switch current.State {
		case orbit.TaskSucceeded:
			var output conversionOutput
			if err := json.Unmarshal(current.Result, &output); err != nil {
				return fmt.Errorf("decode conversion result: %w", err)
			}
			if output.FileName != "records.json" || output.Format != "json" || output.Rows != 2 || output.Content != `[{"amount":"2","name":"alpha"},{"amount":"3","name":"beta"}]` {
				return fmt.Errorf("unexpected conversion result: %+v", output)
			}
			stop()
			err = <-done
			waited = true
			if err != nil {
				return err
			}
			stats := worker.Stats()
			if stats.Started != 1 || stats.ReportsConfirmed != 1 || stats.Running != 0 {
				return fmt.Errorf("unexpected Worker completion counts: %+v", stats)
			}
			fmt.Printf("input=%s output=%s rows=%d state=%s worker_drained=true\n", "records.csv", output.FileName, output.Rows, current.State)
			return nil
		case orbit.TaskFailed, orbit.TaskNeedsReview:
			return fmt.Errorf("conversion task requires attention: %s", current.State)
		}
		select {
		case err := <-done:
			waited = true
			return fmt.Errorf("Worker stopped before conversion completed: %v", err)
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(100 * time.Millisecond):
		}
	}
}

func convertCSV(ctx context.Context, execution orbit.WorkerExecution) (json.RawMessage, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	var in conversionInput
	if err := json.Unmarshal(execution.Task.Payload, &in); err != nil {
		return nil, &orbit.WorkerFailure{Code: "INVALID_FILE"}
	}
	out, err := convert(in)
	if err != nil {
		return nil, &orbit.WorkerFailure{Code: "INVALID_FILE"}
	}
	return json.Marshal(out)
}

func convert(in conversionInput) (conversionOutput, error) {
	var out conversionOutput
	if strings.TrimSpace(in.FileName) == "" || in.TargetFormat != "json" || in.Content == "" {
		return out, errors.New("invalid conversion input")
	}
	reader := csv.NewReader(strings.NewReader(in.Content))
	header, err := reader.Read()
	if err != nil || len(header) == 0 {
		return out, errors.New("CSV header is required")
	}
	seen := make(map[string]struct{}, len(header))
	for i := range header {
		header[i] = strings.TrimSpace(header[i])
		if header[i] == "" {
			return out, errors.New("CSV header is empty")
		}
		if _, ok := seen[header[i]]; ok {
			return out, errors.New("CSV header is duplicated")
		}
		seen[header[i]] = struct{}{}
	}
	rows := make([]map[string]string, 0, 16)
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil || len(record) != len(header) {
			return out, errors.New("CSV row does not match header")
		}
		row := make(map[string]string, len(header))
		for i, value := range record {
			row[header[i]] = value
		}
		rows = append(rows, row)
	}
	encoded, err := json.Marshal(rows)
	if err != nil {
		return out, err
	}
	return conversionOutput{FileName: strings.TrimSuffix(in.FileName, ".csv") + ".json", Format: "json", Content: string(encoded), Rows: len(rows)}, nil
}