xinmi/scripts/wiki_backfill_trigger.go

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strconv"
	"time"

	"github.com/Tencent/WeKnora/internal/common"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/google/uuid"
	"github.com/hibiken/asynq"
)

type scope struct {
	TenantID        uint64 `json:"tenant_id"`
	KnowledgeBaseID string `json:"knowledge_base_id"`
	Language        string `json:"language"`
}

func run() error {
	var scopes []scope
	if err := json.NewDecoder(os.Stdin).Decode(&scopes); err != nil {
		return err
	}
	if len(scopes) == 0 || len(scopes) > 100 {
		return errors.New("invalid scope count")
	}
	for _, s := range scopes {
		if s.TenantID == 0 {
			return errors.New("tenant is required")
		}
		if _, err := uuid.Parse(s.KnowledgeBaseID); err != nil {
			return errors.New("invalid knowledge base")
		}
	}
	db, err := strconv.Atoi(os.Getenv("REDIS_DB"))
	if err != nil || os.Getenv("REDIS_ADDR") == "" {
		return errors.New("missing runtime Redis configuration")
	}
	client := asynq.NewClient(asynq.RedisClientOpt{Addr: os.Getenv("REDIS_ADDR"), Username: os.Getenv("REDIS_USERNAME"), Password: os.Getenv("REDIS_PASSWORD"), DB: db, TLSConfig: common.RedisTLSConfig()})
	defer client.Close()
	if err := client.Ping(); err != nil {
		return errors.New("queue is unavailable")
	}
	for _, s := range scopes {
		body, err := json.Marshal(s)
		if err != nil {
			return err
		}
		_, err = client.Enqueue(asynq.NewTask(types.TypeWikiIngest, body), asynq.Queue(types.QueueWiki), asynq.MaxRetry(10), asynq.Timeout(time.Hour), asynq.Unique(time.Minute))
		if err != nil && !errors.Is(err, asynq.ErrDuplicateTask) {
			return errors.New("wiki trigger enqueue failed")
		}
	}
	fmt.Printf("{\"triggered_scopes\":%d}\n", len(scopes))
	return nil
}

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