easyai-ai-gateway/apps/api/cmd/backfill-binary-results/main.go

package main

import (
	"context"
	"flag"
	"log/slog"
	"os"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

func main() {
	apply := flag.Bool("apply", false, "persist compacted results; default is dry-run")
	requireClean := flag.Bool("require-clean", false, "exit non-zero unless a complete dry-run finds no results requiring URL migration")
	batchSize := flag.Int("batch-size", 100, "rows per batch, maximum 100")
	maxBatches := flag.Int("max-batches", 10, "maximum batches for one invocation")
	afterID := flag.String("after-id", "", "resume after this task UUID")
	flag.Parse()

	logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
	cfg := config.Load()
	if err := cfg.Validate(); err != nil {
		logger.Error("invalid gateway configuration", "error", err)
		os.Exit(1)
	}
	if *batchSize < 1 || *batchSize > 100 || *maxBatches < 1 {
		logger.Error("invalid backfill bounds", "batchSize", *batchSize, "maxBatches", *maxBatches)
		os.Exit(1)
	}
	if *apply && *requireClean {
		logger.Error("--apply and --require-clean cannot be combined")
		os.Exit(1)
	}

	ctx := context.Background()
	db, err := store.Connect(ctx, cfg.DatabaseURL)
	if err != nil {
		logger.Error("connect postgres failed", "error", err)
		os.Exit(1)
	}
	defer db.Close()

	service := runner.New(cfg, db, logger)
	cursor := *afterID
	scanned := 0
	matched := 0
	blockingMatched := 0
	updated := 0
	expired := 0
	expiredLocalPlaceholders := 0
	complete := false
	for batch := 0; batch < *maxBatches; batch++ {
		items, err := db.ListTaskBinaryResultBackfillBatch(ctx, cursor, *batchSize)
		if err != nil {
			logger.Error("list binary result backfill batch failed", "afterId", cursor, "error", err)
			os.Exit(1)
		}
		if len(items) == 0 {
			complete = true
			break
		}
		for _, item := range items {
			cursor = item.ID
			scanned++
			if !runner.TaskResultNeedsURLMigration(item.Result) {
				continue
			}
			matched++
			isExpired := item.FinishedAt.Before(time.Now().Add(-time.Duration(localResultTTLHours(cfg)) * time.Hour))
			hasLocalPlaceholder := runner.TaskResultHasLocalPlaceholder(item.Result)
			if isExpired && hasLocalPlaceholder {
				expiredLocalPlaceholders++
				if *apply {
					logger.Warn("skip expired local result placeholder without overwriting stored result", "taskId", item.ID)
				}
				continue
			}
			blockingMatched++
			if !*apply {
				continue
			}
			persistent, changed, err := service.MigrateTaskResultToURLs(ctx, item.ID, item.Result)
			if err != nil {
				logger.Error("materialize historical binary result failed", "taskId", item.ID, "error", err)
				os.Exit(1)
			}
			if !changed {
				continue
			}
			ok, err := db.UpdateTaskBinaryResultBackfill(ctx, item.ID, persistent)
			if err != nil {
				logger.Error("update historical binary result failed", "taskId", item.ID, "error", err)
				os.Exit(1)
			}
			if !ok {
				continue
			}
			updated++
			if isExpired {
				expired++
			}
		}
		if len(items) < *batchSize {
			complete = true
			break
		}
	}
	logger.Info("binary result backfill completed",
		"apply", *apply,
		"scanned", scanned,
		"matched", matched,
		"blockingMatched", blockingMatched,
		"updated", updated,
		"expired", expired,
		"expiredLocalPlaceholders", expiredLocalPlaceholders,
		"complete", complete,
		"resumeAfterId", cursor,
	)
	if *requireClean && (!complete || blockingMatched > 0) {
		logger.Error("binary result URL migration gate failed", "complete", complete, "blockingMatched", blockingMatched, "expiredLocalPlaceholders", expiredLocalPlaceholders, "resumeAfterId", cursor)
		os.Exit(1)
	}
}

func localResultTTLHours(cfg config.Config) int {
	if cfg.LocalResultTTLHours <= 0 {
		return 24
	}
	return cfg.LocalResultTTLHours
}