68 lines
1.8 KiB
Go
68 lines
1.8 KiB
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)
|
||
|
|
}
|
||
|
|
}
|