knowledge-governance/internal/workflows/catalogrefresh/handler.go

package catalogrefresh

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
	"unicode"

	"knowledge-governance/internal/jobs"
)

type Request struct {
	ConnectionID string `json:"connection_id"`
	SourceType   string `json:"source_type"`
	SourceName   string `json:"source_name"`
	Stage        string `json:"stage"`
	Progress     int    `json:"progress"`
	Message      string `json:"message"`
	Processed    int    `json:"processed"`
	Total        int    `json:"total"`
}
type Progress struct {
	Stage                 string `json:"stage"`
	Progress              int    `json:"progress"`
	Message               string `json:"message"`
	Processed             int    `json:"processed"`
	Total                 int    `json:"total"`
	SourceName            string `json:"source_name,omitempty"`
	CurrentDirectory      string `json:"current_directory,omitempty"`
	ScannedDirectories    int    `json:"scanned_directories,omitempty"`
	DiscoveredDirectories int    `json:"discovered_directories,omitempty"`
	ScannedFiles          int    `json:"scanned_files,omitempty"`
	EstimatedTotalFiles   int    `json:"estimated_total_files,omitempty"`
	TotalIsEstimate       bool   `json:"total_is_estimate,omitempty"`
}
type Store interface {
	UpdateCatalogRefreshProgress(context.Context, string, string, Progress, time.Time) error
}

var progressForbiddenSubstrings = []string{"http://", "https://", "token", "secret", "password", "authorization", "bearer", "cookie", "/run/source-secrets", "api-key", "apikey"}

type Runner interface {
	Run(context.Context, Request, func(Progress) error) error
}
type Config struct {
	Store  Store
	Runner Runner
	Now    func() time.Time
}
type Handler struct {
	store  Store
	runner Runner
	now    func() time.Time
}

func NewHandler(config Config) (*Handler, error) {
	if config.Store == nil || config.Runner == nil {
		return nil, errors.New("catalog refresh dependencies are required")
	}
	if config.Now == nil {
		config.Now = time.Now
	}
	return &Handler{store: config.Store, runner: config.Runner, now: config.Now}, nil
}
func (handler *Handler) Handle(ctx context.Context, job jobs.Job) error {
	var request Request
	// 进度回写会把 stage/progress/scanned_* 等字段写回同一份 payload；重试时严格解码
	// 会把这些合法字段当成未知字段，把真实失败原因盖成 catalog_refresh_invalid_payload。
	decoder := json.NewDecoder(bytes.NewReader(job.Payload))
	if err := decoder.Decode(&request); err != nil || request.ConnectionID == "" || request.SourceName == "" || (request.SourceType != "pp" && request.SourceType != "dingtalk" && request.SourceType != "alist") || job.LeaseOwner == "" {
		return errors.New("catalog_refresh_invalid_payload")
	}
	var last Progress
	return handler.runner.Run(ctx, request, func(progress Progress) error {
		progress = clampProgress(last, progress)
		if !validProgress(progress) {
			return errors.New("catalog_refresh_invalid_progress")
		}
		last = progress
		return handler.store.UpdateCatalogRefreshProgress(ctx, job.ID, job.LeaseOwner, progress, handler.now().UTC())
	})
}
func validProgress(progress Progress) bool {
	if progress.Progress < 0 || progress.Progress > 100 || progress.Processed < 0 || progress.Total < 0 ||
		progress.ScannedDirectories < 0 || progress.DiscoveredDirectories < 0 || progress.ScannedFiles < 0 || progress.EstimatedTotalFiles < 0 ||
		strings.TrimSpace(progress.Stage) == "" || len(progress.Stage) > 32 ||
		strings.TrimSpace(progress.Message) == "" || len(progress.Message) > 300 ||
		len(progress.CurrentDirectory) > 1024 || len(progress.SourceName) > 255 {
		return false
	}
	for _, text := range []string{progress.Stage, progress.Message, progress.CurrentDirectory, progress.SourceName} {
		for _, r := range text {
			if unicode.IsControl(r) && r != '	' {
				return false
			}
		}
		lowered := strings.ToLower(text)
		for _, forbidden := range progressForbiddenSubstrings {
			if strings.Contains(lowered, forbidden) {
				return false
			}
		}
	}
	return true
}

// clampProgress enforces that display progress never regresses when discovered
// totals shift mid-scan (BFS estimates grow), and never clamps an explicit
// failure snapshot.
func clampProgress(previous, next Progress) Progress {
	if next.Stage == "failed" {
		return next
	}
	if next.Total < previous.Total && previous.Total >= next.Processed {
		next.Total = previous.Total
	}
	if next.Total > 0 {
		estimated := next.Processed * 100 / next.Total
		if next.Progress < previous.Progress && previous.Progress <= estimated {
			next.Progress = previous.Progress
		}
	}
	if next.ScannedDirectories < previous.ScannedDirectories {
		next.ScannedDirectories = previous.ScannedDirectories
	}
	if next.ScannedFiles < previous.ScannedFiles {
		next.ScannedFiles = previous.ScannedFiles
	}
	if next.DiscoveredDirectories < previous.DiscoveredDirectories {
		next.DiscoveredDirectories = previous.DiscoveredDirectories
	}
	if next.EstimatedTotalFiles < previous.EstimatedTotalFiles {
		next.EstimatedTotalFiles = previous.EstimatedTotalFiles
	}
	return next
}

type CommandRunner struct {
	PythonPath string
	ScriptPath string
}

func NewCommandRunner(pythonPath, scriptPath string) (*CommandRunner, error) {
	if !filepath.IsAbs(pythonPath) || !filepath.IsAbs(scriptPath) {
		return nil, errors.New("catalog refresh command paths must be absolute")
	}
	return &CommandRunner{PythonPath: pythonPath, ScriptPath: scriptPath}, nil
}
func (runner *CommandRunner) Run(ctx context.Context, request Request, report func(Progress) error) error {
	command := exec.CommandContext(ctx, runner.PythonPath, runner.ScriptPath, "--connection-id", request.ConnectionID, "--source-type", request.SourceType, "--source-name", request.SourceName)
	stdout, err := command.StdoutPipe()
	if err != nil {
		return errors.New("catalog_refresh_start_failed")
	}
	command.Stderr = io.Discard
	if err = command.Start(); err != nil {
		return errors.New("catalog_refresh_start_failed")
	}
	scanner := bufio.NewScanner(stdout)
	scanner.Buffer(make([]byte, 4096), 1<<20)
	for scanner.Scan() {
		decoder := json.NewDecoder(bytes.NewReader(scanner.Bytes()))
		decoder.DisallowUnknownFields()
		var progress Progress
		if decoder.Decode(&progress) != nil || !validProgress(progress) {
			_ = command.Process.Kill()
			_ = command.Wait()
			return errors.New("catalog_refresh_invalid_output")
		}
		if err = report(progress); err != nil {
			_ = command.Process.Kill()
			_ = command.Wait()
			return err
		}
	}
	if scanner.Err() != nil {
		_ = command.Process.Kill()
		_ = command.Wait()
		return errors.New("catalog_refresh_output_failed")
	}
	if err = command.Wait(); err != nil {
		return errors.New("catalog_refresh_failed")
	}
	return nil
}