init
This commit is contained in:
1031
internal/api/router.go
Normal file
1031
internal/api/router.go
Normal file
File diff suppressed because it is too large
Load Diff
33
internal/executor/service.go
Normal file
33
internal/executor/service.go
Normal file
@@ -0,0 +1,33 @@
|
||||
package executor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"scheduler-backend/internal/jobdef"
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
registry *jobdef.Registry
|
||||
runtime jobdef.Runtime
|
||||
}
|
||||
|
||||
func NewService(registry *jobdef.Registry, runtime jobdef.Runtime) *Service {
|
||||
return &Service{
|
||||
registry: registry,
|
||||
runtime: runtime,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) Execute(ctx context.Context, handlerKey string, req jobdef.ExecuteRequest) error {
|
||||
handler, ok := s.registry.Get(handlerKey)
|
||||
if !ok {
|
||||
return fmt.Errorf("handler %s not found", handlerKey)
|
||||
}
|
||||
|
||||
if req.LogCollector == nil {
|
||||
req.LogCollector = jobdef.NewLogCollector()
|
||||
}
|
||||
|
||||
return handler.Run(ctx, s.runtime, req)
|
||||
}
|
||||
28
internal/executor/service_test.go
Normal file
28
internal/executor/service_test.go
Normal file
@@ -0,0 +1,28 @@
|
||||
package executor
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"scheduler-backend/internal/jobdef"
|
||||
"scheduler-backend/pkg/log"
|
||||
)
|
||||
|
||||
func TestExecuteRunsRegisteredHandler(t *testing.T) {
|
||||
registry := jobdef.NewRegistry(jobdef.SampleHandler{})
|
||||
service := NewService(
|
||||
registry,
|
||||
jobdef.Runtime{
|
||||
Logger: log.New(),
|
||||
},
|
||||
)
|
||||
|
||||
err := service.Execute(context.Background(), "sample-handler", jobdef.ExecuteRequest{
|
||||
ExecutionID: "exec-1",
|
||||
JobID: "job-1",
|
||||
TriggerType: "manual",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Execute() error = %v", err)
|
||||
}
|
||||
}
|
||||
103
internal/jobdef/handler.go
Normal file
103
internal/jobdef/handler.go
Normal file
@@ -0,0 +1,103 @@
|
||||
package jobdef
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
|
||||
"scheduler-backend/pkg/config"
|
||||
"scheduler-backend/pkg/log"
|
||||
)
|
||||
|
||||
const logTailLines = 2000
|
||||
|
||||
type LogCollector struct {
|
||||
mu sync.Mutex
|
||||
lines []string
|
||||
}
|
||||
|
||||
func NewLogCollector() *LogCollector {
|
||||
return &LogCollector{}
|
||||
}
|
||||
|
||||
func (c *LogCollector) Appendf(format string, args ...any) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
c.mu.Lock()
|
||||
c.lines = append(c.lines, fmt.Sprintf(format, args...))
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
func (c *LogCollector) String() string {
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return strings.Join(c.lines, "\n")
|
||||
}
|
||||
|
||||
type FlushResult struct {
|
||||
LogText string
|
||||
LogFile string
|
||||
}
|
||||
|
||||
func (c *LogCollector) Flush(executionID string) FlushResult {
|
||||
if c == nil {
|
||||
return FlushResult{}
|
||||
}
|
||||
c.mu.Lock()
|
||||
lines := c.lines
|
||||
c.mu.Unlock()
|
||||
|
||||
if len(lines) == 0 {
|
||||
return FlushResult{}
|
||||
}
|
||||
|
||||
if len(lines) <= logTailLines {
|
||||
return FlushResult{LogText: strings.Join(lines, "\n")}
|
||||
}
|
||||
|
||||
filename := executionID + ".log"
|
||||
dir := filepath.Join("logs", "joblog")
|
||||
_ = os.MkdirAll(dir, 0o755)
|
||||
fullPath := filepath.Join(dir, filename)
|
||||
_ = os.WriteFile(fullPath, []byte(strings.Join(lines, "\n")+"\n"), 0o644)
|
||||
|
||||
tail := lines[len(lines)-logTailLines:]
|
||||
header := fmt.Sprintf("[完整日志: %s (%d行)]", filename, len(lines))
|
||||
|
||||
return FlushResult{
|
||||
LogText: header + "\n" + strings.Join(tail, "\n"),
|
||||
LogFile: filename,
|
||||
}
|
||||
}
|
||||
|
||||
type ExecuteRequest struct {
|
||||
ExecutionID string `json:"executionID"`
|
||||
JobID string `json:"jobID"`
|
||||
TriggerType string `json:"triggerType"`
|
||||
Params json.RawMessage `json:"params"`
|
||||
LogCollector *LogCollector `json:"-"`
|
||||
}
|
||||
|
||||
type Runtime struct {
|
||||
Config config.Config
|
||||
Logger log.Logger
|
||||
MetaDB *mongo.Database
|
||||
BusinessDB *mongo.Database
|
||||
}
|
||||
|
||||
type Handler interface {
|
||||
Key() string
|
||||
Name() string
|
||||
Description() string
|
||||
Run(ctx context.Context, runtime Runtime, req ExecuteRequest) error
|
||||
}
|
||||
62
internal/jobdef/loader.go
Normal file
62
internal/jobdef/loader.go
Normal file
@@ -0,0 +1,62 @@
|
||||
package jobdef
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"scheduler-backend/internal/store/model"
|
||||
"scheduler-backend/pkg/log"
|
||||
)
|
||||
|
||||
type JobConfigUpserter interface {
|
||||
UpsertByHandlerKey(ctx context.Context, item *model.JobConfig) error
|
||||
}
|
||||
|
||||
type configFile struct {
|
||||
Name string `json:"name"`
|
||||
HandlerKey string `json:"handlerKey"`
|
||||
ScheduleType string `json:"scheduleType"`
|
||||
ScheduleValue string `json:"scheduleValue"`
|
||||
DefaultParams string `json:"defaultParams"`
|
||||
}
|
||||
|
||||
func SyncJobConfigs(ctx context.Context, dir string, store JobConfigUpserter, logger log.Logger) error {
|
||||
matches, err := filepath.Glob(filepath.Join(dir, "*.json"))
|
||||
if err != nil {
|
||||
return fmt.Errorf("glob job config files: %w", err)
|
||||
}
|
||||
if len(matches) == 0 {
|
||||
logger.Info("no job config files found", "dir", dir)
|
||||
return nil
|
||||
}
|
||||
|
||||
for _, path := range matches {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read %s: %w", path, err)
|
||||
}
|
||||
var cfg configFile
|
||||
if err := json.Unmarshal(data, &cfg); err != nil {
|
||||
return fmt.Errorf("parse %s: %w", path, err)
|
||||
}
|
||||
if cfg.HandlerKey == "" {
|
||||
return fmt.Errorf("%s: handlerKey is required", path)
|
||||
}
|
||||
|
||||
item := &model.JobConfig{
|
||||
Name: cfg.Name,
|
||||
HandlerKey: cfg.HandlerKey,
|
||||
ScheduleType: cfg.ScheduleType,
|
||||
ScheduleValue: cfg.ScheduleValue,
|
||||
DefaultParams: cfg.DefaultParams,
|
||||
}
|
||||
if err := store.UpsertByHandlerKey(ctx, item); err != nil {
|
||||
return fmt.Errorf("upsert %s: %w", cfg.HandlerKey, err)
|
||||
}
|
||||
logger.Info("synced job config", "handlerKey", cfg.HandlerKey, "name", cfg.Name, "file", filepath.Base(path))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
28
internal/jobdef/registry.go
Normal file
28
internal/jobdef/registry.go
Normal file
@@ -0,0 +1,28 @@
|
||||
package jobdef
|
||||
|
||||
type Registry struct {
|
||||
items map[string]Handler
|
||||
}
|
||||
|
||||
func NewRegistry(handlers ...Handler) *Registry {
|
||||
items := make(map[string]Handler, len(handlers))
|
||||
for _, handler := range handlers {
|
||||
items[handler.Key()] = handler
|
||||
}
|
||||
|
||||
return &Registry{items: items}
|
||||
}
|
||||
|
||||
func (r *Registry) Get(key string) (Handler, bool) {
|
||||
handler, ok := r.items[key]
|
||||
return handler, ok
|
||||
}
|
||||
|
||||
func (r *Registry) List() []Handler {
|
||||
result := make([]Handler, 0, len(r.items))
|
||||
for _, handler := range r.items {
|
||||
result = append(result, handler)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
16
internal/jobdef/registry_test.go
Normal file
16
internal/jobdef/registry_test.go
Normal file
@@ -0,0 +1,16 @@
|
||||
package jobdef
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestRegistryGetReturnsRegisteredHandler(t *testing.T) {
|
||||
registry := NewRegistry(SampleHandler{})
|
||||
|
||||
handler, ok := registry.Get("sample-handler")
|
||||
if !ok {
|
||||
t.Fatal("expected registered handler")
|
||||
}
|
||||
|
||||
if handler.Key() != "sample-handler" {
|
||||
t.Fatalf("handler key = %q, want sample-handler", handler.Key())
|
||||
}
|
||||
}
|
||||
44
internal/jobdef/s3_migrate_handler.go
Normal file
44
internal/jobdef/s3_migrate_handler.go
Normal file
@@ -0,0 +1,44 @@
|
||||
package jobdef
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"scheduler-backend/internal/s3migrate"
|
||||
)
|
||||
|
||||
type S3MigrateHandler struct{}
|
||||
|
||||
func (S3MigrateHandler) Key() string {
|
||||
return "s3-migrate"
|
||||
}
|
||||
|
||||
func (S3MigrateHandler) Name() string {
|
||||
return "S3 Migrate"
|
||||
}
|
||||
|
||||
func (S3MigrateHandler) Description() string {
|
||||
return "Migrate objects between S3-compatible storage engines (MinIO, AWS S3, etc.). Streams files via presigned URLs and updates MongoDB engine records."
|
||||
}
|
||||
|
||||
func (S3MigrateHandler) Run(ctx context.Context, runtime Runtime, req ExecuteRequest) error {
|
||||
var params s3migrate.Params
|
||||
if err := json.Unmarshal(req.Params, ¶ms); err != nil {
|
||||
return fmt.Errorf("parse params: %w", err)
|
||||
}
|
||||
|
||||
logf := func(msg string, args ...any) {
|
||||
line := fmt.Sprintf(msg, args...)
|
||||
req.LogCollector.Appendf("%s", line)
|
||||
runtime.Logger.Info(line, "jobID", req.JobID, "executionID", req.ExecutionID)
|
||||
}
|
||||
|
||||
if err := s3migrate.Run(ctx, runtime.BusinessDB, params, logf); err != nil {
|
||||
logf("s3-migrate failed: %v", err)
|
||||
return fmt.Errorf("s3-migrate failed: %w", err)
|
||||
}
|
||||
|
||||
logf("s3-migrate finished")
|
||||
return nil
|
||||
}
|
||||
31
internal/jobdef/sample_handler.go
Normal file
31
internal/jobdef/sample_handler.go
Normal file
@@ -0,0 +1,31 @@
|
||||
package jobdef
|
||||
|
||||
import "context"
|
||||
|
||||
type SampleHandler struct{}
|
||||
|
||||
func (SampleHandler) Key() string {
|
||||
return "sample-handler"
|
||||
}
|
||||
|
||||
func (SampleHandler) Name() string {
|
||||
return "Sample Handler"
|
||||
}
|
||||
|
||||
func (SampleHandler) Description() string {
|
||||
return "Writes a sample startup log for scheduler plumbing."
|
||||
}
|
||||
|
||||
func (SampleHandler) Run(ctx context.Context, runtime Runtime, req ExecuteRequest) error {
|
||||
req.LogCollector.Appendf("sample handler executed: jobID=%s executionID=%s triggerType=%s", req.JobID, req.ExecutionID, req.TriggerType)
|
||||
runtime.Logger.Info(
|
||||
"sample handler executed",
|
||||
"jobID",
|
||||
req.JobID,
|
||||
"executionID",
|
||||
req.ExecutionID,
|
||||
"triggerType",
|
||||
req.TriggerType,
|
||||
)
|
||||
return nil
|
||||
}
|
||||
49
internal/jobdef/url_rewrite_handler.go
Normal file
49
internal/jobdef/url_rewrite_handler.go
Normal file
@@ -0,0 +1,49 @@
|
||||
package jobdef
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"scheduler-backend/internal/urlrewrite"
|
||||
)
|
||||
|
||||
type URLRewriteHandler struct{}
|
||||
|
||||
func (URLRewriteHandler) Key() string {
|
||||
return "url-rewrite"
|
||||
}
|
||||
|
||||
func (URLRewriteHandler) Name() string {
|
||||
return "URL Rewrite"
|
||||
}
|
||||
|
||||
func (URLRewriteHandler) Description() string {
|
||||
return "Batch-rewrite URL prefixes in MongoDB (user avatars, message media, favorites). Supports dry-run, apply, rollback, and cache invalidation."
|
||||
}
|
||||
|
||||
func (URLRewriteHandler) Run(ctx context.Context, runtime Runtime, req ExecuteRequest) error {
|
||||
var params urlrewrite.Params
|
||||
if err := json.Unmarshal(req.Params, ¶ms); err != nil {
|
||||
return fmt.Errorf("parse params: %w", err)
|
||||
}
|
||||
|
||||
logf := func(msg string, args ...any) {
|
||||
line := fmt.Sprintf(msg, args...)
|
||||
req.LogCollector.Appendf("%s", line)
|
||||
runtime.Logger.Info(line, "jobID", req.JobID, "executionID", req.ExecutionID)
|
||||
}
|
||||
|
||||
redisCfg := urlrewrite.RedisConfig{
|
||||
Addr: runtime.Config.RedisAddr,
|
||||
Password: runtime.Config.RedisPassword,
|
||||
}
|
||||
batchID, err := urlrewrite.Run(ctx, runtime.BusinessDB, params, redisCfg, logf)
|
||||
if err != nil {
|
||||
logf("url-rewrite failed: %v", err)
|
||||
return fmt.Errorf("url-rewrite %s failed (batch_id=%s): %w", params.Mode, batchID, err)
|
||||
}
|
||||
|
||||
logf("url-rewrite finished, mode=%s batch_id=%s", params.Mode, batchID)
|
||||
return nil
|
||||
}
|
||||
220
internal/s3migrate/runner.go
Normal file
220
internal/s3migrate/runner.go
Normal file
@@ -0,0 +1,220 @@
|
||||
package s3migrate
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
)
|
||||
|
||||
const (
|
||||
collectionName = "s3"
|
||||
presignExpiry = time.Hour
|
||||
)
|
||||
|
||||
type Params struct {
|
||||
SourceEngine string `json:"sourceEngine"`
|
||||
SourceEndpoint string `json:"sourceEndpoint"`
|
||||
SourceBucket string `json:"sourceBucket"`
|
||||
SourceAccessKey string `json:"sourceAccessKey"`
|
||||
SourceSecretKey string `json:"sourceSecretKey"`
|
||||
SourceRegion string `json:"sourceRegion"`
|
||||
|
||||
DestEngine string `json:"destEngine"`
|
||||
DestEndpoint string `json:"destEndpoint"`
|
||||
DestBucket string `json:"destBucket"`
|
||||
DestAccessKey string `json:"destAccessKey"`
|
||||
DestSecretKey string `json:"destSecretKey"`
|
||||
DestRegion string `json:"destRegion"`
|
||||
}
|
||||
|
||||
type objectRecord struct {
|
||||
Name string `bson:"name"`
|
||||
Engine string `bson:"engine"`
|
||||
Key string `bson:"key"`
|
||||
Size int64 `bson:"size"`
|
||||
ContentType string `bson:"content_type"`
|
||||
}
|
||||
|
||||
type LogFunc func(msg string, args ...any)
|
||||
|
||||
func Run(ctx context.Context, db *mongo.Database, params Params, logf LogFunc) error {
|
||||
if err := validateParams(params); err != nil {
|
||||
return err
|
||||
}
|
||||
if params.SourceRegion == "" {
|
||||
params.SourceRegion = "us-east-1"
|
||||
}
|
||||
if params.DestRegion == "" {
|
||||
params.DestRegion = "us-east-1"
|
||||
}
|
||||
|
||||
srcPresign := newPresignClient(params.SourceEndpoint, params.SourceAccessKey, params.SourceSecretKey, params.SourceRegion)
|
||||
dstPresign := newPresignClient(params.DestEndpoint, params.DestAccessKey, params.DestSecretKey, params.DestRegion)
|
||||
|
||||
coll := db.Collection(collectionName)
|
||||
count, err := coll.CountDocuments(ctx, bson.M{"engine": params.SourceEngine})
|
||||
if err != nil {
|
||||
return fmt.Errorf("count source objects: %w", err)
|
||||
}
|
||||
logf("source engine=%s count=%d", params.SourceEngine, count)
|
||||
|
||||
if count == 0 {
|
||||
logf("no objects to migrate")
|
||||
return nil
|
||||
}
|
||||
|
||||
cursor, err := coll.Find(ctx, bson.M{"engine": params.SourceEngine}, options.Find().SetBatchSize(100))
|
||||
if err != nil {
|
||||
return fmt.Errorf("find source objects: %w", err)
|
||||
}
|
||||
defer cursor.Close(ctx)
|
||||
|
||||
var migrated, skipped, failed int64
|
||||
for cursor.Next(ctx) {
|
||||
var obj objectRecord
|
||||
if err := cursor.Decode(&obj); err != nil {
|
||||
logf("decode error: %v", err)
|
||||
failed++
|
||||
continue
|
||||
}
|
||||
|
||||
exists, err := objectExistsInDest(ctx, coll, params.DestEngine, obj.Name)
|
||||
if err != nil {
|
||||
logf("check dest error name=%s: %v", obj.Name, err)
|
||||
failed++
|
||||
continue
|
||||
}
|
||||
if exists {
|
||||
skipped++
|
||||
continue
|
||||
}
|
||||
|
||||
if err := copyObject(ctx, srcPresign, params.SourceBucket, dstPresign, params.DestBucket, obj); err != nil {
|
||||
logf("copy error name=%s: %v", obj.Name, err)
|
||||
failed++
|
||||
continue
|
||||
}
|
||||
|
||||
if _, err := coll.UpdateOne(ctx,
|
||||
bson.M{"engine": params.SourceEngine, "name": obj.Name},
|
||||
bson.M{"$set": bson.M{"engine": params.DestEngine}},
|
||||
); err != nil {
|
||||
logf("update engine error name=%s: %v", obj.Name, err)
|
||||
failed++
|
||||
continue
|
||||
}
|
||||
|
||||
migrated++
|
||||
if migrated%100 == 0 {
|
||||
logf("progress: migrated=%d skipped=%d failed=%d", migrated, skipped, failed)
|
||||
}
|
||||
}
|
||||
|
||||
logf("complete: migrated=%d skipped=%d failed=%d total=%d", migrated, skipped, failed, count)
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateParams(p Params) error {
|
||||
if p.SourceEngine == "" {
|
||||
return fmt.Errorf("sourceEngine is required")
|
||||
}
|
||||
if p.DestEngine == "" {
|
||||
return fmt.Errorf("destEngine is required")
|
||||
}
|
||||
if p.SourceEngine == p.DestEngine {
|
||||
return fmt.Errorf("sourceEngine and destEngine must be different")
|
||||
}
|
||||
if p.SourceEndpoint == "" || p.SourceBucket == "" || p.SourceAccessKey == "" || p.SourceSecretKey == "" {
|
||||
return fmt.Errorf("source S3 config (endpoint, bucket, accessKey, secretKey) is required")
|
||||
}
|
||||
if p.DestEndpoint == "" || p.DestBucket == "" || p.DestAccessKey == "" || p.DestSecretKey == "" {
|
||||
return fmt.Errorf("dest S3 config (endpoint, bucket, accessKey, secretKey) is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func newPresignClient(endpoint, accessKey, secretKey, region string) *s3.PresignClient {
|
||||
cfg := aws.Config{
|
||||
Region: region,
|
||||
Credentials: credentials.NewStaticCredentialsProvider(accessKey, secretKey, ""),
|
||||
BaseEndpoint: &endpoint,
|
||||
}
|
||||
client := s3.NewFromConfig(cfg, func(o *s3.Options) {
|
||||
o.UsePathStyle = true
|
||||
})
|
||||
return s3.NewPresignClient(client)
|
||||
}
|
||||
|
||||
func objectExistsInDest(ctx context.Context, coll *mongo.Collection, destEngine, name string) (bool, error) {
|
||||
count, err := coll.CountDocuments(ctx, bson.M{"engine": destEngine, "name": name}, options.Count().SetLimit(1))
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
func copyObject(ctx context.Context, srcPresign *s3.PresignClient, srcBucket string, dstPresign *s3.PresignClient, dstBucket string, obj objectRecord) error {
|
||||
getReq, err := srcPresign.PresignGetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: &srcBucket,
|
||||
Key: &obj.Key,
|
||||
}, s3.WithPresignExpires(presignExpiry))
|
||||
if err != nil {
|
||||
return fmt.Errorf("presign get: %w", err)
|
||||
}
|
||||
|
||||
contentType := obj.ContentType
|
||||
putReq, err := dstPresign.PresignPutObject(ctx, &s3.PutObjectInput{
|
||||
Bucket: &dstBucket,
|
||||
Key: &obj.Key,
|
||||
ContentType: &contentType,
|
||||
}, s3.WithPresignExpires(presignExpiry))
|
||||
if err != nil {
|
||||
return fmt.Errorf("presign put: %w", err)
|
||||
}
|
||||
|
||||
downloadResp, err := http.Get(getReq.URL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("download: %w", err)
|
||||
}
|
||||
defer downloadResp.Body.Close()
|
||||
|
||||
if downloadResp.StatusCode == http.StatusNotFound {
|
||||
return fmt.Errorf("source object not found: %s", obj.Key)
|
||||
}
|
||||
if downloadResp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("download status: %s", downloadResp.Status)
|
||||
}
|
||||
|
||||
uploadReq, err := http.NewRequestWithContext(ctx, http.MethodPut, putReq.URL, downloadResp.Body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create upload request: %w", err)
|
||||
}
|
||||
if downloadResp.ContentLength > 0 {
|
||||
uploadReq.ContentLength = downloadResp.ContentLength
|
||||
} else {
|
||||
uploadReq.ContentLength = obj.Size
|
||||
}
|
||||
if contentType != "" {
|
||||
uploadReq.Header.Set("Content-Type", contentType)
|
||||
}
|
||||
|
||||
uploadResp, err := http.DefaultClient.Do(uploadReq)
|
||||
if err != nil {
|
||||
return fmt.Errorf("upload: %w", err)
|
||||
}
|
||||
defer func() { _, _ = io.Copy(io.Discard, uploadResp.Body); uploadResp.Body.Close() }()
|
||||
|
||||
if uploadResp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("upload status: %s", uploadResp.Status)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
171
internal/scheduler/service.go
Normal file
171
internal/scheduler/service.go
Normal file
@@ -0,0 +1,171 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/go-co-op/gocron/v2"
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
|
||||
"scheduler-backend/internal/executor"
|
||||
"scheduler-backend/internal/jobdef"
|
||||
"scheduler-backend/internal/store/model"
|
||||
)
|
||||
|
||||
type JobConfigStore interface {
|
||||
ListEnabled(ctx context.Context) ([]model.JobConfig, error)
|
||||
UpdateRunState(ctx context.Context, id primitive.ObjectID, status string, lastRunAt *time.Time) error
|
||||
UpdateNextRunAt(ctx context.Context, id primitive.ObjectID, nextRunAt *time.Time) error
|
||||
}
|
||||
|
||||
type JobExecutionStore interface {
|
||||
Create(ctx context.Context, item *model.JobExecution) error
|
||||
UpdateResult(ctx context.Context, item *model.JobExecution) error
|
||||
}
|
||||
|
||||
type Service struct {
|
||||
scheduler gocron.Scheduler
|
||||
jobConfigStore JobConfigStore
|
||||
executionStore JobExecutionStore
|
||||
executor *executor.Service
|
||||
runtime jobdef.Runtime
|
||||
}
|
||||
|
||||
func NewService(
|
||||
scheduler gocron.Scheduler,
|
||||
jobConfigStore JobConfigStore,
|
||||
executionStore JobExecutionStore,
|
||||
executor *executor.Service,
|
||||
runtime jobdef.Runtime,
|
||||
) *Service {
|
||||
return &Service{
|
||||
scheduler: scheduler,
|
||||
jobConfigStore: jobConfigStore,
|
||||
executionStore: executionStore,
|
||||
executor: executor,
|
||||
runtime: runtime,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) RegisterEnabledJobs(ctx context.Context) error {
|
||||
items, err := s.jobConfigStore.ListEnabled(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, item := range items {
|
||||
if err := s.registerJob(ctx, item); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
s.scheduler.Start()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) Shutdown() error {
|
||||
return s.scheduler.Shutdown()
|
||||
}
|
||||
|
||||
func (s *Service) registerJob(ctx context.Context, item model.JobConfig) error {
|
||||
switch item.ScheduleType {
|
||||
case "manual":
|
||||
return s.jobConfigStore.UpdateNextRunAt(ctx, item.ID, nil)
|
||||
case "cron":
|
||||
job, err := s.scheduler.NewJob(
|
||||
gocron.CronJob(item.ScheduleValue, false),
|
||||
gocron.NewTask(func() {
|
||||
s.runScheduledJob(context.Background(), item)
|
||||
}),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("register cron job %s: %w", item.Name, err)
|
||||
}
|
||||
next, err := job.NextRun()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.jobConfigStore.UpdateNextRunAt(ctx, item.ID, &next)
|
||||
case "interval":
|
||||
duration, err := time.ParseDuration(item.ScheduleValue)
|
||||
if err != nil {
|
||||
return fmt.Errorf("parse interval for %s: %w", item.Name, err)
|
||||
}
|
||||
job, err := s.scheduler.NewJob(
|
||||
gocron.DurationJob(duration),
|
||||
gocron.NewTask(func() {
|
||||
s.runScheduledJob(context.Background(), item)
|
||||
}),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("register interval job %s: %w", item.Name, err)
|
||||
}
|
||||
next, err := job.NextRun()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.jobConfigStore.UpdateNextRunAt(ctx, item.ID, &next)
|
||||
default:
|
||||
return fmt.Errorf("unsupported schedule type %s", item.ScheduleType)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) runScheduledJob(ctx context.Context, job model.JobConfig) {
|
||||
now := time.Now()
|
||||
params := json.RawMessage(job.DefaultParams)
|
||||
record := model.JobExecution{
|
||||
JobConfigID: job.ID,
|
||||
TriggerType: "schedule",
|
||||
Status: "running",
|
||||
ParamsSnapshot: string(params),
|
||||
StartedAt: &now,
|
||||
CreatedAt: now,
|
||||
}
|
||||
|
||||
if err := s.executionStore.Create(ctx, &record); err != nil {
|
||||
s.runtime.Logger.Error("create scheduled execution failed", "jobID", job.ID.Hex(), "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
execReq := jobdef.ExecuteRequest{
|
||||
ExecutionID: record.ID.Hex(),
|
||||
JobID: job.ID.Hex(),
|
||||
TriggerType: "schedule",
|
||||
Params: params,
|
||||
LogCollector: jobdef.NewLogCollector(),
|
||||
}
|
||||
runErr := s.executor.Execute(ctx, job.HandlerKey, execReq)
|
||||
|
||||
finishedAt := time.Now()
|
||||
record.FinishedAt = &finishedAt
|
||||
record.DurationMs = finishedAt.Sub(now).Milliseconds()
|
||||
flushed := execReq.LogCollector.Flush(record.ID.Hex())
|
||||
if runErr != nil {
|
||||
record.Status = "failed"
|
||||
record.ErrorMessage = runErr.Error()
|
||||
record.ResultSummary = "scheduled execution failed"
|
||||
if flushed.LogText != "" {
|
||||
record.LogText = flushed.LogText
|
||||
} else {
|
||||
record.LogText = runErr.Error()
|
||||
}
|
||||
record.LogFile = flushed.LogFile
|
||||
_ = s.jobConfigStore.UpdateRunState(ctx, job.ID, "failed", &finishedAt)
|
||||
} else {
|
||||
record.Status = "success"
|
||||
record.ResultSummary = "scheduled execution succeeded"
|
||||
if flushed.LogText != "" {
|
||||
record.LogText = flushed.LogText
|
||||
} else {
|
||||
record.LogText = "handler executed successfully"
|
||||
}
|
||||
record.LogFile = flushed.LogFile
|
||||
_ = s.jobConfigStore.UpdateRunState(ctx, job.ID, "success", &finishedAt)
|
||||
}
|
||||
|
||||
if err := s.executionStore.UpdateResult(ctx, &record); err != nil {
|
||||
s.runtime.Logger.Error("update scheduled execution failed", "executionID", record.ID.Hex(), "error", err)
|
||||
}
|
||||
}
|
||||
18
internal/store/model/admin_profile.go
Normal file
18
internal/store/model/admin_profile.go
Normal file
@@ -0,0 +1,18 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
)
|
||||
|
||||
type AdminProfile struct {
|
||||
ID primitive.ObjectID `bson:"_id,omitempty" json:"id"`
|
||||
Account string `bson:"account" json:"account"`
|
||||
Nickname string `bson:"nickname" json:"nickname"`
|
||||
Avatar string `bson:"avatar" json:"avatar"`
|
||||
Remark string `bson:"remark" json:"remark"`
|
||||
PasswordHash string `bson:"passwordHash" json:"-"`
|
||||
CreatedAt time.Time `bson:"createdAt" json:"createdAt"`
|
||||
UpdatedAt time.Time `bson:"updatedAt" json:"updatedAt"`
|
||||
}
|
||||
19
internal/store/model/admin_user.go
Normal file
19
internal/store/model/admin_user.go
Normal file
@@ -0,0 +1,19 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
)
|
||||
|
||||
type AdminUser struct {
|
||||
ID primitive.ObjectID `bson:"_id,omitempty" json:"id"`
|
||||
Account string `bson:"account" json:"account"`
|
||||
Nickname string `bson:"nickname" json:"nickname"`
|
||||
Status string `bson:"status" json:"status"`
|
||||
Role string `bson:"role" json:"role"`
|
||||
Remark string `bson:"remark" json:"remark"`
|
||||
PasswordHash string `bson:"passwordHash" json:"-"`
|
||||
CreatedAt time.Time `bson:"createdAt" json:"createdAt"`
|
||||
UpdatedAt time.Time `bson:"updatedAt" json:"updatedAt"`
|
||||
}
|
||||
33
internal/store/model/job_config.go
Normal file
33
internal/store/model/job_config.go
Normal file
@@ -0,0 +1,33 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
)
|
||||
|
||||
type JobConfig struct {
|
||||
ID primitive.ObjectID `bson:"_id,omitempty" json:"id"`
|
||||
Name string `bson:"name" json:"name"`
|
||||
HandlerKey string `bson:"handlerKey" json:"handlerKey"`
|
||||
Enabled bool `bson:"enabled" json:"enabled"`
|
||||
ScheduleType string `bson:"scheduleType" json:"scheduleType"`
|
||||
ScheduleValue string `bson:"scheduleValue" json:"scheduleValue"`
|
||||
DefaultParams string `bson:"defaultParams" json:"defaultParams"`
|
||||
LastStatus string `bson:"lastStatus" json:"lastStatus"`
|
||||
LastRunAt *time.Time `bson:"lastRunAt,omitempty" json:"lastRunAt,omitempty"`
|
||||
NextRunAt *time.Time `bson:"nextRunAt,omitempty" json:"nextRunAt,omitempty"`
|
||||
CreatedAt time.Time `bson:"createdAt" json:"createdAt"`
|
||||
UpdatedAt time.Time `bson:"updatedAt" json:"updatedAt"`
|
||||
}
|
||||
|
||||
func NewJobConfig(name, handlerKey string) JobConfig {
|
||||
now := time.Now()
|
||||
return JobConfig{
|
||||
Name: name,
|
||||
HandlerKey: handlerKey,
|
||||
LastStatus: "idle",
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
}
|
||||
10
internal/store/model/job_config_test.go
Normal file
10
internal/store/model/job_config_test.go
Normal file
@@ -0,0 +1,10 @@
|
||||
package model
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestJobConfigDefaultsToIdleStatus(t *testing.T) {
|
||||
job := NewJobConfig("daily-check", "daily-check")
|
||||
if job.LastStatus != "idle" {
|
||||
t.Fatalf("LastStatus = %q, want idle", job.LastStatus)
|
||||
}
|
||||
}
|
||||
23
internal/store/model/job_execution.go
Normal file
23
internal/store/model/job_execution.go
Normal file
@@ -0,0 +1,23 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
)
|
||||
|
||||
type JobExecution struct {
|
||||
ID primitive.ObjectID `bson:"_id,omitempty" json:"id"`
|
||||
JobConfigID primitive.ObjectID `bson:"jobConfigId" json:"jobConfigId"`
|
||||
TriggerType string `bson:"triggerType" json:"triggerType"`
|
||||
Status string `bson:"status" json:"status"`
|
||||
ParamsSnapshot string `bson:"paramsSnapshot" json:"paramsSnapshot"`
|
||||
StartedAt *time.Time `bson:"startedAt,omitempty" json:"startedAt,omitempty"`
|
||||
FinishedAt *time.Time `bson:"finishedAt,omitempty" json:"finishedAt,omitempty"`
|
||||
DurationMs int64 `bson:"durationMs" json:"durationMs"`
|
||||
ResultSummary string `bson:"resultSummary" json:"resultSummary"`
|
||||
ErrorMessage string `bson:"errorMessage" json:"errorMessage"`
|
||||
LogText string `bson:"logText" json:"logText"`
|
||||
LogFile string `bson:"logFile,omitempty" json:"logFile,omitempty"`
|
||||
CreatedAt time.Time `bson:"createdAt" json:"createdAt"`
|
||||
}
|
||||
19
internal/store/model/system_config.go
Normal file
19
internal/store/model/system_config.go
Normal file
@@ -0,0 +1,19 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
)
|
||||
|
||||
type SystemConfig struct {
|
||||
ID primitive.ObjectID `bson:"_id,omitempty" json:"id"`
|
||||
Key string `bson:"key" json:"key"`
|
||||
Title string `bson:"title" json:"title"`
|
||||
Value string `bson:"value" json:"value"`
|
||||
ValueType string `bson:"valueType" json:"valueType"`
|
||||
Description string `bson:"description" json:"description"`
|
||||
Enabled bool `bson:"enabled" json:"enabled"`
|
||||
CreatedAt time.Time `bson:"createdAt" json:"createdAt"`
|
||||
UpdatedAt time.Time `bson:"updatedAt" json:"updatedAt"`
|
||||
}
|
||||
125
internal/store/mongo/admin_profile_store.go
Normal file
125
internal/store/mongo/admin_profile_store.go
Normal file
@@ -0,0 +1,125 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
mongodriver "go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
|
||||
"scheduler-backend/internal/store/model"
|
||||
)
|
||||
|
||||
const defaultAdminPassword = "admin123"
|
||||
|
||||
type AdminProfileStore struct {
|
||||
collection *mongodriver.Collection
|
||||
}
|
||||
|
||||
func NewAdminProfileStore(db *mongodriver.Database) *AdminProfileStore {
|
||||
return &AdminProfileStore{
|
||||
collection: db.Collection("admin_profiles"),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AdminProfileStore) Get(ctx context.Context) (*model.AdminProfile, error) {
|
||||
findCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
var item model.AdminProfile
|
||||
err := s.collection.FindOne(findCtx, bson.D{}, options.FindOne().SetSort(bson.D{{Key: "updatedAt", Value: -1}})).Decode(&item)
|
||||
if err == nil {
|
||||
return &item, nil
|
||||
}
|
||||
if !errors.Is(err, mongodriver.ErrNoDocuments) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
item = model.AdminProfile{
|
||||
ID: primitive.NewObjectID(),
|
||||
Account: "admin",
|
||||
Nickname: "调度管理员",
|
||||
Avatar: "",
|
||||
Remark: "default admin profile",
|
||||
PasswordHash: hashPassword(defaultAdminPassword),
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
if _, err := s.collection.InsertOne(findCtx, item); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
func (s *AdminProfileStore) Upsert(ctx context.Context, profile *model.AdminProfile) error {
|
||||
updateCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
now := time.Now()
|
||||
profile.UpdatedAt = now
|
||||
if profile.CreatedAt.IsZero() {
|
||||
profile.CreatedAt = now
|
||||
}
|
||||
|
||||
if profile.ID.IsZero() {
|
||||
current, err := s.Get(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
profile.ID = current.ID
|
||||
if profile.PasswordHash == "" {
|
||||
profile.PasswordHash = current.PasswordHash
|
||||
}
|
||||
if profile.CreatedAt.IsZero() {
|
||||
profile.CreatedAt = current.CreatedAt
|
||||
}
|
||||
}
|
||||
|
||||
_, err := s.collection.UpdateByID(updateCtx, profile.ID, bson.M{
|
||||
"$set": bson.M{
|
||||
"account": profile.Account,
|
||||
"nickname": profile.Nickname,
|
||||
"avatar": profile.Avatar,
|
||||
"remark": profile.Remark,
|
||||
"passwordHash": profile.PasswordHash,
|
||||
"createdAt": profile.CreatedAt,
|
||||
"updatedAt": profile.UpdatedAt,
|
||||
},
|
||||
}, options.Update().SetUpsert(true))
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *AdminProfileStore) ChangePassword(ctx context.Context, oldPassword, newPassword string) error {
|
||||
current, err := s.Get(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if current.PasswordHash != hashPassword(oldPassword) {
|
||||
return ErrInvalidPassword
|
||||
}
|
||||
|
||||
current.PasswordHash = hashPassword(newPassword)
|
||||
return s.Upsert(ctx, current)
|
||||
}
|
||||
|
||||
func (s *AdminProfileStore) Authenticate(ctx context.Context, account, password string) (*model.AdminProfile, error) {
|
||||
current, err := s.Get(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if current.Account != account || current.PasswordHash != hashPassword(password) {
|
||||
return nil, ErrInvalidPassword
|
||||
}
|
||||
return current, nil
|
||||
}
|
||||
|
||||
func hashPassword(password string) string {
|
||||
sum := sha256.Sum256([]byte(password))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
152
internal/store/mongo/admin_user_store.go
Normal file
152
internal/store/mongo/admin_user_store.go
Normal file
@@ -0,0 +1,152 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
mongodriver "go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
|
||||
"scheduler-backend/internal/store/model"
|
||||
)
|
||||
|
||||
type AdminUserStore struct {
|
||||
collection *mongodriver.Collection
|
||||
}
|
||||
|
||||
func NewAdminUserStore(db *mongodriver.Database) *AdminUserStore {
|
||||
return &AdminUserStore{
|
||||
collection: db.Collection("admin_users"),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AdminUserStore) List(ctx context.Context) ([]model.AdminUser, error) {
|
||||
findCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
cursor, err := s.collection.Find(findCtx, bson.D{}, options.Find().SetSort(bson.D{{Key: "updatedAt", Value: -1}}))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer cursor.Close(findCtx)
|
||||
|
||||
items := make([]model.AdminUser, 0)
|
||||
for cursor.Next(findCtx) {
|
||||
var item model.AdminUser
|
||||
if err := cursor.Decode(&item); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
if err := cursor.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if len(items) == 0 {
|
||||
defaultAdmin := model.AdminUser{
|
||||
Account: "admin",
|
||||
Nickname: "调度管理员",
|
||||
Status: "active",
|
||||
Role: "super_admin",
|
||||
Remark: "default admin user",
|
||||
PasswordHash: hashPassword(defaultAdminPassword),
|
||||
}
|
||||
if err := s.Create(ctx, &defaultAdmin); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, defaultAdmin)
|
||||
}
|
||||
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (s *AdminUserStore) GetByID(ctx context.Context, id primitive.ObjectID) (*model.AdminUser, error) {
|
||||
findCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
var item model.AdminUser
|
||||
err := s.collection.FindOne(findCtx, bson.M{"_id": id}).Decode(&item)
|
||||
if err != nil {
|
||||
if errors.Is(err, mongodriver.ErrNoDocuments) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
func (s *AdminUserStore) Create(ctx context.Context, item *model.AdminUser) error {
|
||||
insertCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
now := time.Now()
|
||||
item.CreatedAt = now
|
||||
item.UpdatedAt = now
|
||||
if item.Status == "" {
|
||||
item.Status = "active"
|
||||
}
|
||||
if item.Role == "" {
|
||||
item.Role = "super_admin"
|
||||
}
|
||||
|
||||
result, err := s.collection.InsertOne(insertCtx, item)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if oid, ok := result.InsertedID.(primitive.ObjectID); ok {
|
||||
item.ID = oid
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *AdminUserStore) Update(ctx context.Context, item *model.AdminUser) error {
|
||||
updateCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
item.UpdatedAt = time.Now()
|
||||
_, err := s.collection.UpdateByID(updateCtx, item.ID, bson.M{
|
||||
"$set": bson.M{
|
||||
"account": item.Account,
|
||||
"nickname": item.Nickname,
|
||||
"status": item.Status,
|
||||
"role": item.Role,
|
||||
"remark": item.Remark,
|
||||
"updatedAt": item.UpdatedAt,
|
||||
},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *AdminUserStore) ChangePassword(ctx context.Context, id primitive.ObjectID, newPassword string) error {
|
||||
updateCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
_, err := s.collection.UpdateByID(updateCtx, id, bson.M{
|
||||
"$set": bson.M{
|
||||
"passwordHash": hashPassword(newPassword),
|
||||
"updatedAt": time.Now(),
|
||||
},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *AdminUserStore) Authenticate(ctx context.Context, account, password string) (*model.AdminUser, error) {
|
||||
findCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
var item model.AdminUser
|
||||
err := s.collection.FindOne(findCtx, bson.M{"account": account}).Decode(&item)
|
||||
if err != nil {
|
||||
if errors.Is(err, mongodriver.ErrNoDocuments) {
|
||||
return nil, ErrInvalidPassword
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
if item.Status == "disabled" || item.PasswordHash != hashPassword(password) {
|
||||
return nil, ErrInvalidPassword
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
37
internal/store/mongo/client.go
Normal file
37
internal/store/mongo/client.go
Normal file
@@ -0,0 +1,37 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
mongodriver "go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
|
||||
"scheduler-backend/pkg/config"
|
||||
)
|
||||
|
||||
type Databases struct {
|
||||
Client *mongodriver.Client
|
||||
MetaDB *mongodriver.Database
|
||||
BusinessDB *mongodriver.Database
|
||||
}
|
||||
|
||||
func Connect(ctx context.Context, cfg config.Config) (*Databases, error) {
|
||||
client, err := mongodriver.Connect(ctx, options.Client().ApplyURI(cfg.MongoURI()))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
pingCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := client.Ping(pingCtx, nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &Databases{
|
||||
Client: client,
|
||||
MetaDB: client.Database(cfg.SchedulerMongoDatabase),
|
||||
BusinessDB: client.Database(cfg.BusinessMongoDatabase),
|
||||
}, nil
|
||||
}
|
||||
5
internal/store/mongo/errors.go
Normal file
5
internal/store/mongo/errors.go
Normal file
@@ -0,0 +1,5 @@
|
||||
package mongo
|
||||
|
||||
import "errors"
|
||||
|
||||
var ErrInvalidPassword = errors.New("invalid password")
|
||||
222
internal/store/mongo/job_config_store.go
Normal file
222
internal/store/mongo/job_config_store.go
Normal file
@@ -0,0 +1,222 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
mongodriver "go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
|
||||
"scheduler-backend/internal/store/model"
|
||||
)
|
||||
|
||||
type JobConfigStore struct {
|
||||
collection *mongodriver.Collection
|
||||
}
|
||||
|
||||
func NewJobConfigStore(db *mongodriver.Database) *JobConfigStore {
|
||||
return &JobConfigStore{
|
||||
collection: db.Collection("job_configs"),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *JobConfigStore) List(ctx context.Context) ([]model.JobConfig, error) {
|
||||
findCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
cursor, err := s.collection.Find(
|
||||
findCtx,
|
||||
bson.D{},
|
||||
options.Find().SetSort(bson.D{{Key: "updatedAt", Value: -1}}),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer cursor.Close(findCtx)
|
||||
|
||||
items := make([]model.JobConfig, 0)
|
||||
for cursor.Next(findCtx) {
|
||||
var item model.JobConfig
|
||||
if err := cursor.Decode(&item); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
|
||||
if err := cursor.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (s *JobConfigStore) ListEnabled(ctx context.Context) ([]model.JobConfig, error) {
|
||||
findCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
cursor, err := s.collection.Find(
|
||||
findCtx,
|
||||
bson.M{"enabled": true},
|
||||
options.Find().SetSort(bson.D{{Key: "updatedAt", Value: -1}}),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer cursor.Close(findCtx)
|
||||
|
||||
items := make([]model.JobConfig, 0)
|
||||
for cursor.Next(findCtx) {
|
||||
var item model.JobConfig
|
||||
if err := cursor.Decode(&item); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
|
||||
if err := cursor.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (s *JobConfigStore) GetByID(ctx context.Context, id primitive.ObjectID) (*model.JobConfig, error) {
|
||||
findCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
var item model.JobConfig
|
||||
err := s.collection.FindOne(findCtx, bson.M{"_id": id}).Decode(&item)
|
||||
if err != nil {
|
||||
if errors.Is(err, mongodriver.ErrNoDocuments) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
func (s *JobConfigStore) Create(ctx context.Context, item *model.JobConfig) error {
|
||||
insertCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
now := time.Now()
|
||||
item.CreatedAt = now
|
||||
item.UpdatedAt = now
|
||||
if item.LastStatus == "" {
|
||||
item.LastStatus = "idle"
|
||||
}
|
||||
|
||||
result, err := s.collection.InsertOne(insertCtx, item)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if oid, ok := result.InsertedID.(primitive.ObjectID); ok {
|
||||
item.ID = oid
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *JobConfigStore) Update(ctx context.Context, item *model.JobConfig) error {
|
||||
updateCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
item.UpdatedAt = time.Now()
|
||||
update := bson.M{
|
||||
"$set": bson.M{
|
||||
"name": item.Name,
|
||||
"handlerKey": item.HandlerKey,
|
||||
"enabled": item.Enabled,
|
||||
"scheduleType": item.ScheduleType,
|
||||
"scheduleValue": item.ScheduleValue,
|
||||
"defaultParams": item.DefaultParams,
|
||||
"updatedAt": item.UpdatedAt,
|
||||
},
|
||||
}
|
||||
|
||||
_, err := s.collection.UpdateByID(updateCtx, item.ID, update)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *JobConfigStore) ToggleEnabled(ctx context.Context, id primitive.ObjectID, enabled bool) error {
|
||||
updateCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
_, err := s.collection.UpdateByID(updateCtx, id, bson.M{
|
||||
"$set": bson.M{
|
||||
"enabled": enabled,
|
||||
"updatedAt": time.Now(),
|
||||
},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *JobConfigStore) UpdateRunState(ctx context.Context, id primitive.ObjectID, status string, lastRunAt *time.Time) error {
|
||||
updateCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
set := bson.M{
|
||||
"lastStatus": status,
|
||||
"updatedAt": time.Now(),
|
||||
}
|
||||
if lastRunAt != nil {
|
||||
set["lastRunAt"] = *lastRunAt
|
||||
}
|
||||
|
||||
_, err := s.collection.UpdateByID(updateCtx, id, bson.M{
|
||||
"$set": set,
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *JobConfigStore) UpsertByHandlerKey(ctx context.Context, item *model.JobConfig) error {
|
||||
upsertCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
now := time.Now()
|
||||
_, err := s.collection.UpdateOne(
|
||||
upsertCtx,
|
||||
bson.M{"handlerKey": item.HandlerKey},
|
||||
bson.M{
|
||||
"$set": bson.M{
|
||||
"name": item.Name,
|
||||
"handlerKey": item.HandlerKey,
|
||||
"scheduleType": item.ScheduleType,
|
||||
"scheduleValue": item.ScheduleValue,
|
||||
"defaultParams": item.DefaultParams,
|
||||
"updatedAt": now,
|
||||
},
|
||||
"$setOnInsert": bson.M{
|
||||
"enabled": false,
|
||||
"lastStatus": "idle",
|
||||
"createdAt": now,
|
||||
},
|
||||
},
|
||||
options.Update().SetUpsert(true),
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *JobConfigStore) UpdateNextRunAt(ctx context.Context, id primitive.ObjectID, nextRunAt *time.Time) error {
|
||||
updateCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
set := bson.M{
|
||||
"updatedAt": time.Now(),
|
||||
}
|
||||
if nextRunAt != nil {
|
||||
set["nextRunAt"] = *nextRunAt
|
||||
} else {
|
||||
set["nextRunAt"] = nil
|
||||
}
|
||||
|
||||
_, err := s.collection.UpdateByID(updateCtx, id, bson.M{
|
||||
"$set": set,
|
||||
})
|
||||
return err
|
||||
}
|
||||
109
internal/store/mongo/job_execution_store.go
Normal file
109
internal/store/mongo/job_execution_store.go
Normal file
@@ -0,0 +1,109 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
mongodriver "go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
|
||||
"scheduler-backend/internal/store/model"
|
||||
)
|
||||
|
||||
type JobExecutionStore struct {
|
||||
collection *mongodriver.Collection
|
||||
}
|
||||
|
||||
func NewJobExecutionStore(db *mongodriver.Database) *JobExecutionStore {
|
||||
return &JobExecutionStore{
|
||||
collection: db.Collection("job_executions"),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *JobExecutionStore) List(ctx context.Context) ([]model.JobExecution, error) {
|
||||
findCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
cursor, err := s.collection.Find(
|
||||
findCtx,
|
||||
bson.D{},
|
||||
options.Find().SetSort(bson.D{{Key: "createdAt", Value: -1}}),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer cursor.Close(findCtx)
|
||||
|
||||
items := make([]model.JobExecution, 0)
|
||||
for cursor.Next(findCtx) {
|
||||
var item model.JobExecution
|
||||
if err := cursor.Decode(&item); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
|
||||
if err := cursor.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (s *JobExecutionStore) GetByID(ctx context.Context, id primitive.ObjectID) (*model.JobExecution, error) {
|
||||
findCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
var item model.JobExecution
|
||||
err := s.collection.FindOne(findCtx, bson.M{"_id": id}).Decode(&item)
|
||||
if err != nil {
|
||||
if errors.Is(err, mongodriver.ErrNoDocuments) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
func (s *JobExecutionStore) Create(ctx context.Context, item *model.JobExecution) error {
|
||||
insertCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
if item.CreatedAt.IsZero() {
|
||||
item.CreatedAt = time.Now()
|
||||
}
|
||||
|
||||
result, err := s.collection.InsertOne(insertCtx, item)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if oid, ok := result.InsertedID.(primitive.ObjectID); ok {
|
||||
item.ID = oid
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *JobExecutionStore) UpdateResult(ctx context.Context, item *model.JobExecution) error {
|
||||
updateCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
_, err := s.collection.UpdateByID(updateCtx, item.ID, bson.M{
|
||||
"$set": bson.M{
|
||||
"status": item.Status,
|
||||
"startedAt": item.StartedAt,
|
||||
"finishedAt": item.FinishedAt,
|
||||
"durationMs": item.DurationMs,
|
||||
"resultSummary": item.ResultSummary,
|
||||
"errorMessage": item.ErrorMessage,
|
||||
"logText": item.LogText,
|
||||
"logFile": item.LogFile,
|
||||
},
|
||||
})
|
||||
return err
|
||||
}
|
||||
113
internal/store/mongo/system_config_store.go
Normal file
113
internal/store/mongo/system_config_store.go
Normal file
@@ -0,0 +1,113 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
mongodriver "go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
|
||||
"scheduler-backend/internal/store/model"
|
||||
)
|
||||
|
||||
type SystemConfigStore struct {
|
||||
collection *mongodriver.Collection
|
||||
}
|
||||
|
||||
func NewSystemConfigStore(db *mongodriver.Database) *SystemConfigStore {
|
||||
return &SystemConfigStore{
|
||||
collection: db.Collection("system_configs"),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SystemConfigStore) List(ctx context.Context) ([]model.SystemConfig, error) {
|
||||
findCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
cursor, err := s.collection.Find(findCtx, bson.D{}, options.Find().SetSort(bson.D{{Key: "updatedAt", Value: -1}}))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer cursor.Close(findCtx)
|
||||
|
||||
items := make([]model.SystemConfig, 0)
|
||||
for cursor.Next(findCtx) {
|
||||
var item model.SystemConfig
|
||||
if err := cursor.Decode(&item); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
if err := cursor.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (s *SystemConfigStore) GetByID(ctx context.Context, id primitive.ObjectID) (*model.SystemConfig, error) {
|
||||
findCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
var item model.SystemConfig
|
||||
err := s.collection.FindOne(findCtx, bson.M{"_id": id}).Decode(&item)
|
||||
if err != nil {
|
||||
if errors.Is(err, mongodriver.ErrNoDocuments) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &item, nil
|
||||
}
|
||||
|
||||
func (s *SystemConfigStore) Create(ctx context.Context, item *model.SystemConfig) error {
|
||||
insertCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
now := time.Now()
|
||||
item.CreatedAt = now
|
||||
item.UpdatedAt = now
|
||||
|
||||
result, err := s.collection.InsertOne(insertCtx, item)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if oid, ok := result.InsertedID.(primitive.ObjectID); ok {
|
||||
item.ID = oid
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SystemConfigStore) Update(ctx context.Context, item *model.SystemConfig) error {
|
||||
updateCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
item.UpdatedAt = time.Now()
|
||||
_, err := s.collection.UpdateByID(updateCtx, item.ID, bson.M{
|
||||
"$set": bson.M{
|
||||
"key": item.Key,
|
||||
"title": item.Title,
|
||||
"value": item.Value,
|
||||
"valueType": item.ValueType,
|
||||
"description": item.Description,
|
||||
"enabled": item.Enabled,
|
||||
"updatedAt": item.UpdatedAt,
|
||||
},
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *SystemConfigStore) ToggleEnabled(ctx context.Context, id primitive.ObjectID, enabled bool) error {
|
||||
updateCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
_, err := s.collection.UpdateByID(updateCtx, id, bson.M{
|
||||
"$set": bson.M{
|
||||
"enabled": enabled,
|
||||
"updatedAt": time.Now(),
|
||||
},
|
||||
})
|
||||
return err
|
||||
}
|
||||
18
internal/urlrewrite/redis.go
Normal file
18
internal/urlrewrite/redis.go
Normal file
@@ -0,0 +1,18 @@
|
||||
package urlrewrite
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
func newRedisClient(ctx context.Context, addr, password string) (*redis.Client, error) {
|
||||
rdb := redis.NewClient(&redis.Options{
|
||||
Addr: addr,
|
||||
Password: password,
|
||||
})
|
||||
if err := rdb.Ping(ctx).Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return rdb, nil
|
||||
}
|
||||
65
internal/urlrewrite/rewrite.go
Normal file
65
internal/urlrewrite/rewrite.go
Normal file
@@ -0,0 +1,65 @@
|
||||
package urlrewrite
|
||||
|
||||
import "encoding/json"
|
||||
|
||||
func rewriteStringValue(input, oldPrefix, newPrefix string) (string, bool) {
|
||||
if oldPrefix == "" || newPrefix == "" {
|
||||
return input, false
|
||||
}
|
||||
if len(input) >= len(oldPrefix) && input[:len(oldPrefix)] == oldPrefix {
|
||||
return newPrefix + input[len(oldPrefix):], true
|
||||
}
|
||||
return input, false
|
||||
}
|
||||
|
||||
func rewriteJSONContent(input, oldPrefix, newPrefix string) (string, bool, error) {
|
||||
var data any
|
||||
if err := json.Unmarshal([]byte(input), &data); err != nil {
|
||||
rewritten, changed := rewriteStringValue(input, oldPrefix, newPrefix)
|
||||
if changed {
|
||||
return rewritten, true, nil
|
||||
}
|
||||
return "", false, err
|
||||
}
|
||||
rewritten, changed := rewriteAny(data, oldPrefix, newPrefix)
|
||||
if !changed {
|
||||
return input, false, nil
|
||||
}
|
||||
out, err := json.Marshal(rewritten)
|
||||
if err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
return string(out), true, nil
|
||||
}
|
||||
|
||||
func rewriteAny(value any, oldPrefix, newPrefix string) (any, bool) {
|
||||
switch v := value.(type) {
|
||||
case map[string]any:
|
||||
changed := false
|
||||
for key, child := range v {
|
||||
newChild, childChanged := rewriteAny(child, oldPrefix, newPrefix)
|
||||
if childChanged {
|
||||
v[key] = newChild
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
return v, changed
|
||||
case []any:
|
||||
changed := false
|
||||
for i, child := range v {
|
||||
newChild, childChanged := rewriteAny(child, oldPrefix, newPrefix)
|
||||
if childChanged {
|
||||
v[i] = newChild
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
return v, changed
|
||||
case string:
|
||||
if rewritten, changed, err := rewriteJSONContent(v, oldPrefix, newPrefix); err == nil && changed {
|
||||
return rewritten, true
|
||||
}
|
||||
return rewriteStringValue(v, oldPrefix, newPrefix)
|
||||
default:
|
||||
return value, false
|
||||
}
|
||||
}
|
||||
170
internal/urlrewrite/rewrite_test.go
Normal file
170
internal/urlrewrite/rewrite_test.go
Normal file
@@ -0,0 +1,170 @@
|
||||
package urlrewrite
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRewriteStringValue(t *testing.T) {
|
||||
oldPrefix := "https://s3.jizhying.com/images/"
|
||||
newPrefix := "https://dp9pkdckmd09t.cloudfront.net/"
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "replace matching prefix",
|
||||
input: "https://s3.jizhying.com/images/openim/data/hash/a.jpg",
|
||||
want: "https://dp9pkdckmd09t.cloudfront.net/openim/data/hash/a.jpg",
|
||||
},
|
||||
{
|
||||
name: "keep non matching url",
|
||||
input: "https://example.com/a.jpg",
|
||||
want: "https://example.com/a.jpg",
|
||||
},
|
||||
{
|
||||
name: "keep plain text",
|
||||
input: "hello",
|
||||
want: "hello",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, changed := rewriteStringValue(tt.input, oldPrefix, newPrefix)
|
||||
if got != tt.want {
|
||||
t.Fatalf("rewriteStringValue() = %q, want %q", got, tt.want)
|
||||
}
|
||||
if (tt.input != tt.want) != changed {
|
||||
t.Fatalf("rewriteStringValue() changed = %v, want %v", changed, tt.input != tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRewriteJSONContent(t *testing.T) {
|
||||
oldPrefix := "https://s3.jizhying.com/images/"
|
||||
newPrefix := "https://dp9pkdckmd09t.cloudfront.net/"
|
||||
|
||||
input := `{"sourcePicture":{"url":"https://s3.jizhying.com/images/openim/data/hash/src.jpg"},"snapshotPicture":{"url":"https://s3.jizhying.com/images/openim/data/hash/snap.jpg"},"nested":{"fileElem":{"sourceUrl":"https://s3.jizhying.com/images/openim/data/hash/file.zip"}}}`
|
||||
want := `{"nested":{"fileElem":{"sourceUrl":"https://dp9pkdckmd09t.cloudfront.net/openim/data/hash/file.zip"}},"snapshotPicture":{"url":"https://dp9pkdckmd09t.cloudfront.net/openim/data/hash/snap.jpg"},"sourcePicture":{"url":"https://dp9pkdckmd09t.cloudfront.net/openim/data/hash/src.jpg"}}`
|
||||
|
||||
got, changed, err := rewriteJSONContent(input, oldPrefix, newPrefix)
|
||||
if err != nil {
|
||||
t.Fatalf("rewriteJSONContent() error = %v", err)
|
||||
}
|
||||
if !changed {
|
||||
t.Fatalf("rewriteJSONContent() changed = false, want true")
|
||||
}
|
||||
if got != want {
|
||||
t.Fatalf("rewriteJSONContent() = %s, want %s", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRewriteJSONContentNestedJSONString(t *testing.T) {
|
||||
oldPrefix := "https://s3.jizhying.com/images/"
|
||||
newPrefix := "https://dp9pkdckmd09t.cloudfront.net/"
|
||||
|
||||
input := `{"detail":"{\"group\":{\"faceURL\":\"https://s3.jizhying.com/images/openim/data/hash/group.jpg\"},\"entrantUser\":{\"faceURL\":\"https://s3.jizhying.com/images/openim/data/hash/user.jpg\"}}"}`
|
||||
want := `{"detail":"{\"group\":{\"faceURL\":\"https://dp9pkdckmd09t.cloudfront.net/openim/data/hash/group.jpg\"},\"entrantUser\":{\"faceURL\":\"https://dp9pkdckmd09t.cloudfront.net/openim/data/hash/user.jpg\"}}"}`
|
||||
|
||||
got, changed, err := rewriteJSONContent(input, oldPrefix, newPrefix)
|
||||
if err != nil {
|
||||
t.Fatalf("rewriteJSONContent() error = %v", err)
|
||||
}
|
||||
if !changed {
|
||||
t.Fatalf("rewriteJSONContent() changed = false, want true")
|
||||
}
|
||||
var gotValue map[string]any
|
||||
var wantValue map[string]any
|
||||
if err := json.Unmarshal([]byte(got), &gotValue); err != nil {
|
||||
t.Fatalf("unmarshal got error = %v", err)
|
||||
}
|
||||
if err := json.Unmarshal([]byte(want), &wantValue); err != nil {
|
||||
t.Fatalf("unmarshal want error = %v", err)
|
||||
}
|
||||
var gotDetail any
|
||||
var wantDetail any
|
||||
if err := json.Unmarshal([]byte(gotValue["detail"].(string)), &gotDetail); err != nil {
|
||||
t.Fatalf("unmarshal got detail error = %v", err)
|
||||
}
|
||||
if err := json.Unmarshal([]byte(wantValue["detail"].(string)), &wantDetail); err != nil {
|
||||
t.Fatalf("unmarshal want detail error = %v", err)
|
||||
}
|
||||
gotValue["detail"] = gotDetail
|
||||
wantValue["detail"] = wantDetail
|
||||
if !reflect.DeepEqual(gotValue, wantValue) {
|
||||
t.Fatalf("rewriteJSONContent() = %#v, want %#v", gotValue, wantValue)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRewriteJSONContentInvalidJSON(t *testing.T) {
|
||||
_, changed, err := rewriteJSONContent("not-json", "https://old/", "https://new/")
|
||||
if err == nil {
|
||||
t.Fatalf("rewriteJSONContent() error = nil, want non-nil")
|
||||
}
|
||||
if changed {
|
||||
t.Fatalf("rewriteJSONContent() changed = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCollectMsgFieldUpdates(t *testing.T) {
|
||||
oldPrefix := "https://s3.jizhying.com/images/"
|
||||
newPrefix := "https://dp9pkdckmd09t.cloudfront.net/"
|
||||
|
||||
doc := msgDoc{
|
||||
Doc: "si_a_b:0",
|
||||
Msgs: []msgEntryModel{
|
||||
{
|
||||
Msg: &msgPayload{
|
||||
Content: `{"sourcePicture":{"url":"https://s3.jizhying.com/images/openim/data/hash/1.jpg"}}`,
|
||||
SenderFaceURL: "https://s3.jizhying.com/images/openim/data/hash/avatar.jpg",
|
||||
},
|
||||
},
|
||||
{
|
||||
Msg: &msgPayload{
|
||||
Content: `{"text":"hello"}`,
|
||||
SenderFaceURL: "",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
updates, changed, samples := collectMsgFieldUpdates(doc, oldPrefix, newPrefix, 5)
|
||||
if !changed {
|
||||
t.Fatalf("collectMsgFieldUpdates() changed = false, want true")
|
||||
}
|
||||
if len(samples) != 2 {
|
||||
t.Fatalf("collectMsgFieldUpdates() samples len = %d, want 2", len(samples))
|
||||
}
|
||||
if len(updates) != 2 {
|
||||
t.Fatalf("collectMsgFieldUpdates() updates len = %d, want 2", len(updates))
|
||||
}
|
||||
if updates[0].Path != "msgs.0.msg.content" {
|
||||
t.Fatalf("first update path = %s", updates[0].Path)
|
||||
}
|
||||
if updates[0].NewValue != `{"sourcePicture":{"url":"https://dp9pkdckmd09t.cloudfront.net/openim/data/hash/1.jpg"}}` {
|
||||
t.Fatalf("content update = %s", updates[0].NewValue)
|
||||
}
|
||||
if updates[1].Path != "msgs.0.msg.sender_face_url" {
|
||||
t.Fatalf("second update path = %s", updates[1].Path)
|
||||
}
|
||||
if updates[1].NewValue != "https://dp9pkdckmd09t.cloudfront.net/openim/data/hash/avatar.jpg" {
|
||||
t.Fatalf("sender_face_url update = %s", updates[1].NewValue)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildBackupDocs(t *testing.T) {
|
||||
updates := []fieldUpdate{
|
||||
{Path: "face_url", OldValue: "https://old/a.jpg", NewValue: "https://new/a.jpg"},
|
||||
{Path: "thumbnail", OldValue: "https://old/b.jpg", NewValue: "https://new/b.jpg"},
|
||||
}
|
||||
|
||||
docs := buildBackupDocs("batch-1", "attributes", "doc-1", updates)
|
||||
if len(docs) != 2 {
|
||||
t.Fatalf("buildBackupDocs() len = %d, want 2", len(docs))
|
||||
}
|
||||
}
|
||||
500
internal/urlrewrite/runner.go
Normal file
500
internal/urlrewrite/runner.go
Normal file
@@ -0,0 +1,500 @@
|
||||
package urlrewrite
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
)
|
||||
|
||||
const (
|
||||
backupCollectionName = "url_rewrite_backups"
|
||||
msgCachePrefix = "MSG_CACHE:"
|
||||
|
||||
ModeDryRun = "dry-run"
|
||||
ModeApply = "apply"
|
||||
ModeRollback = "rollback"
|
||||
ModeInvalidateCache = "invalidate-cache"
|
||||
)
|
||||
|
||||
type Params struct {
|
||||
OldPrefix string `json:"oldPrefix"`
|
||||
NewPrefix string `json:"newPrefix"`
|
||||
Mode string `json:"mode"`
|
||||
BatchID string `json:"batchId"`
|
||||
SampleSize int `json:"sampleSize"`
|
||||
}
|
||||
|
||||
type RedisConfig struct {
|
||||
Addr string
|
||||
Password string
|
||||
}
|
||||
|
||||
type Summary struct {
|
||||
Collection string
|
||||
Scanned int64
|
||||
Updated int64
|
||||
Samples []string
|
||||
}
|
||||
|
||||
type fieldUpdate struct {
|
||||
Path string
|
||||
OldValue string
|
||||
NewValue string
|
||||
}
|
||||
|
||||
type msgDoc struct {
|
||||
ID any `bson:"_id"`
|
||||
Doc string `bson:"doc_id"`
|
||||
Msgs []msgEntryModel `bson:"msgs"`
|
||||
}
|
||||
|
||||
type msgEntryModel struct {
|
||||
Msg *msgPayload `bson:"msg"`
|
||||
}
|
||||
|
||||
type msgPayload struct {
|
||||
Content string `bson:"content"`
|
||||
SenderFaceURL string `bson:"sender_face_url"`
|
||||
Seq int64 `bson:"seq"`
|
||||
}
|
||||
|
||||
type LogFunc func(msg string, args ...any)
|
||||
|
||||
func Run(ctx context.Context, db *mongo.Database, params Params, redisCfg RedisConfig, logf LogFunc) (string, error) {
|
||||
if err := validateParams(params); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if params.SampleSize <= 0 {
|
||||
params.SampleSize = 5
|
||||
}
|
||||
if params.Mode == "" {
|
||||
params.Mode = ModeDryRun
|
||||
}
|
||||
if params.Mode == ModeApply && params.BatchID == "" {
|
||||
params.BatchID = "urlrewrite-" + time.Now().Format("20060102150405")
|
||||
}
|
||||
|
||||
backupColl := db.Collection(backupCollectionName)
|
||||
|
||||
if params.Mode == ModeRollback {
|
||||
return params.BatchID, rollbackBatch(ctx, backupColl, db, params, logf)
|
||||
}
|
||||
if params.Mode == ModeInvalidateCache {
|
||||
return params.BatchID, invalidateBatchCache(ctx, db, redisCfg, params, logf)
|
||||
}
|
||||
|
||||
summaries := []Summary{
|
||||
rewriteSimpleFieldCollection(ctx, db.Collection("wallets"), backupColl, "real_name_auth.id_card_photo_front", params),
|
||||
rewriteSimpleFieldCollection(ctx, db.Collection("wallets"), backupColl, "real_name_auth.id_card_photo_back", params),
|
||||
rewriteSimpleFieldCollection(ctx, db.Collection("favorites"), backupColl, "content", params),
|
||||
rewriteSimpleFieldCollection(ctx, db.Collection("favorites"), backupColl, "thumbnail", params),
|
||||
rewriteSimpleFieldCollection(ctx, db.Collection("favorites"), backupColl, "link_url", params),
|
||||
rewriteSimpleFieldCollection(ctx, db.Collection("attributes"), backupColl, "face_url", params),
|
||||
rewriteSimpleFieldCollection(ctx, db.Collection("attribute"), backupColl, "face_url", params),
|
||||
rewriteSimpleFieldCollection(ctx, db.Collection("user"), backupColl, "face_url", params),
|
||||
rewriteSimpleFieldCollection(ctx, db.Collection("group"), backupColl, "face_url", params),
|
||||
rewriteSimpleFieldCollection(ctx, db.Collection("group_member"), backupColl, "face_url", params),
|
||||
}
|
||||
|
||||
msgSummary, err := rewriteMsgCollection(ctx, db.Collection("msg"), backupColl, params)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("rewrite msg collection: %w", err)
|
||||
}
|
||||
summaries = append(summaries, msgSummary)
|
||||
|
||||
for _, s := range summaries {
|
||||
logf("[%s] scanned=%d updated=%d", s.Collection, s.Scanned, s.Updated)
|
||||
for _, sample := range s.Samples {
|
||||
logf("[%s] sample: %s", s.Collection, sample)
|
||||
}
|
||||
}
|
||||
|
||||
switch params.Mode {
|
||||
case ModeDryRun:
|
||||
logf("dry-run complete, rerun with mode=apply to write changes")
|
||||
case ModeApply:
|
||||
logf("apply complete, batch_id=%s", params.BatchID)
|
||||
}
|
||||
return params.BatchID, nil
|
||||
}
|
||||
|
||||
func validateParams(p Params) error {
|
||||
switch p.Mode {
|
||||
case ModeDryRun, ModeApply, ModeRollback, ModeInvalidateCache:
|
||||
case "":
|
||||
default:
|
||||
return fmt.Errorf("invalid mode: %s", p.Mode)
|
||||
}
|
||||
if p.Mode == ModeRollback || p.Mode == ModeInvalidateCache {
|
||||
if p.BatchID == "" {
|
||||
return fmt.Errorf("batchId is required for %s", p.Mode)
|
||||
}
|
||||
} else {
|
||||
if p.OldPrefix == "" || p.NewPrefix == "" {
|
||||
return fmt.Errorf("oldPrefix and newPrefix are required")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func invalidateBatchCache(ctx context.Context, db *mongo.Database, redisCfg RedisConfig, params Params, logf LogFunc) error {
|
||||
if redisCfg.Addr == "" {
|
||||
return fmt.Errorf("REDIS_ADDR is not configured")
|
||||
}
|
||||
backupColl := db.Collection(backupCollectionName)
|
||||
cursor, err := backupColl.Find(ctx, bson.M{
|
||||
"batch_id": params.BatchID,
|
||||
"collection": "msg",
|
||||
"field": bson.M{"$regex": primitive.Regex{Pattern: `^msgs\.\d+\.msg\.`}},
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("find backup docs: %w", err)
|
||||
}
|
||||
defer cursor.Close(ctx)
|
||||
|
||||
type backupDoc struct {
|
||||
DocumentID any `bson:"document_id"`
|
||||
Field string `bson:"field"`
|
||||
}
|
||||
|
||||
docIndexes := make(map[string]map[int]struct{})
|
||||
for cursor.Next(ctx) {
|
||||
var item backupDoc
|
||||
if err := cursor.Decode(&item); err != nil {
|
||||
return fmt.Errorf("decode backup doc: %w", err)
|
||||
}
|
||||
docID, ok := item.DocumentID.(primitive.ObjectID)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
msgIndex, ok := parseMsgFieldIndex(item.Field)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
id := docID.Hex()
|
||||
if _, exists := docIndexes[id]; !exists {
|
||||
docIndexes[id] = make(map[int]struct{})
|
||||
}
|
||||
docIndexes[id][msgIndex] = struct{}{}
|
||||
}
|
||||
|
||||
msgColl := db.Collection("msg")
|
||||
var keys []string
|
||||
for hexID, indexes := range docIndexes {
|
||||
docObjectID, err := primitive.ObjectIDFromHex(hexID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("parse object id %s: %w", hexID, err)
|
||||
}
|
||||
var doc msgDoc
|
||||
if err := msgColl.FindOne(ctx, bson.M{"_id": docObjectID}).Decode(&doc); err != nil {
|
||||
return fmt.Errorf("find msg doc %s: %w", hexID, err)
|
||||
}
|
||||
conversationID := trimDocIDSuffix(doc.Doc)
|
||||
if conversationID == "" {
|
||||
continue
|
||||
}
|
||||
for idx := range indexes {
|
||||
if idx < 0 || idx >= len(doc.Msgs) || doc.Msgs[idx].Msg == nil || doc.Msgs[idx].Msg.Seq <= 0 {
|
||||
continue
|
||||
}
|
||||
keys = append(keys, msgCachePrefix+conversationID+":"+strconv.Itoa(int(doc.Msgs[idx].Msg.Seq)))
|
||||
}
|
||||
}
|
||||
|
||||
keys = deduplicateStrings(keys)
|
||||
if len(keys) == 0 {
|
||||
logf("invalidate-cache complete, batch_id=%s keys=0", params.BatchID)
|
||||
return nil
|
||||
}
|
||||
|
||||
rdb, err := newRedisClient(ctx, redisCfg.Addr, redisCfg.Password)
|
||||
if err != nil {
|
||||
return fmt.Errorf("connect redis: %w", err)
|
||||
}
|
||||
defer rdb.Close()
|
||||
|
||||
if err := rdb.Del(ctx, keys...).Err(); err != nil {
|
||||
return fmt.Errorf("redis del: %w", err)
|
||||
}
|
||||
logf("invalidate-cache complete, batch_id=%s keys=%d", params.BatchID, len(keys))
|
||||
return nil
|
||||
}
|
||||
|
||||
func rewriteSimpleFieldCollection(ctx context.Context, coll, backupColl *mongo.Collection, field string, params Params) Summary {
|
||||
filter := bson.M{field: bson.M{"$regex": primitive.Regex{Pattern: "^" + regexp.QuoteMeta(params.OldPrefix)}}}
|
||||
cursor, err := coll.Find(ctx, filter)
|
||||
if err != nil {
|
||||
return Summary{Collection: coll.Name() + "." + field, Samples: []string{fmt.Sprintf("find error: %v", err)}}
|
||||
}
|
||||
defer cursor.Close(ctx)
|
||||
|
||||
summary := Summary{Collection: coll.Name() + "." + field}
|
||||
for cursor.Next(ctx) {
|
||||
summary.Scanned++
|
||||
var doc bson.M
|
||||
if err := cursor.Decode(&doc); err != nil {
|
||||
if len(summary.Samples) < params.SampleSize {
|
||||
summary.Samples = append(summary.Samples, fmt.Sprintf("decode error: %v", err))
|
||||
}
|
||||
continue
|
||||
}
|
||||
original, ok := nestedString(doc, field)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
rewritten, changed := rewriteStringValue(original, params.OldPrefix, params.NewPrefix)
|
||||
if !changed {
|
||||
continue
|
||||
}
|
||||
updates := []fieldUpdate{{
|
||||
Path: field,
|
||||
OldValue: original,
|
||||
NewValue: rewritten,
|
||||
}}
|
||||
summary.Updated++
|
||||
if len(summary.Samples) < params.SampleSize {
|
||||
summary.Samples = append(summary.Samples, fmt.Sprintf("%s -> %s", original, rewritten))
|
||||
}
|
||||
if params.Mode != ModeApply {
|
||||
continue
|
||||
}
|
||||
if err := insertBackupDocs(ctx, backupColl, params.BatchID, coll.Name(), doc["_id"], updates); err != nil {
|
||||
if len(summary.Samples) < params.SampleSize {
|
||||
summary.Samples = append(summary.Samples, fmt.Sprintf("backup error: %v", err))
|
||||
}
|
||||
continue
|
||||
}
|
||||
if _, err := coll.UpdateByID(ctx, doc["_id"], bson.M{"$set": bson.M{field: rewritten}}); err != nil {
|
||||
if len(summary.Samples) < params.SampleSize {
|
||||
summary.Samples = append(summary.Samples, fmt.Sprintf("update error: %v", err))
|
||||
}
|
||||
}
|
||||
}
|
||||
return summary
|
||||
}
|
||||
|
||||
func rewriteMsgCollection(ctx context.Context, coll, backupColl *mongo.Collection, params Params) (Summary, error) {
|
||||
cursor, err := coll.Find(ctx, bson.M{
|
||||
"$or": bson.A{
|
||||
bson.M{"msgs.msg.content": bson.M{"$regex": primitive.Regex{Pattern: regexp.QuoteMeta(params.OldPrefix)}}},
|
||||
bson.M{"msgs.msg.sender_face_url": bson.M{"$regex": primitive.Regex{Pattern: "^" + regexp.QuoteMeta(params.OldPrefix)}}},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return Summary{}, fmt.Errorf("find msg docs: %w", err)
|
||||
}
|
||||
defer cursor.Close(ctx)
|
||||
|
||||
summary := Summary{Collection: coll.Name()}
|
||||
for cursor.Next(ctx) {
|
||||
summary.Scanned++
|
||||
var doc msgDoc
|
||||
if err := cursor.Decode(&doc); err != nil {
|
||||
if len(summary.Samples) < params.SampleSize {
|
||||
summary.Samples = append(summary.Samples, fmt.Sprintf("decode error: %v", err))
|
||||
}
|
||||
continue
|
||||
}
|
||||
updates, changed, samples := collectMsgFieldUpdates(doc, params.OldPrefix, params.NewPrefix, params.SampleSize-len(summary.Samples))
|
||||
if !changed {
|
||||
continue
|
||||
}
|
||||
summary.Updated++
|
||||
summary.Samples = append(summary.Samples, samples...)
|
||||
if params.Mode != ModeApply {
|
||||
continue
|
||||
}
|
||||
if err := insertBackupDocs(ctx, backupColl, params.BatchID, coll.Name(), doc.ID, updates); err != nil {
|
||||
if len(summary.Samples) < params.SampleSize {
|
||||
summary.Samples = append(summary.Samples, fmt.Sprintf("backup error: %v", err))
|
||||
}
|
||||
continue
|
||||
}
|
||||
if _, err := coll.UpdateByID(ctx, doc.ID, bson.M{"$set": toUpdateMap(updates)}); err != nil {
|
||||
return summary, fmt.Errorf("update msg doc: %w", err)
|
||||
}
|
||||
}
|
||||
return summary, nil
|
||||
}
|
||||
|
||||
func collectMsgFieldUpdates(doc msgDoc, oldPrefix, newPrefix string, sampleBudget int) ([]fieldUpdate, bool, []string) {
|
||||
var updates []fieldUpdate
|
||||
var samples []string
|
||||
for i := range doc.Msgs {
|
||||
if doc.Msgs[i].Msg == nil {
|
||||
continue
|
||||
}
|
||||
if doc.Msgs[i].Msg.Content != "" {
|
||||
rewritten, contentChanged, err := rewriteJSONContent(doc.Msgs[i].Msg.Content, oldPrefix, newPrefix)
|
||||
if err == nil && contentChanged {
|
||||
updates = append(updates, fieldUpdate{
|
||||
Path: fmt.Sprintf("msgs.%d.msg.content", i),
|
||||
OldValue: doc.Msgs[i].Msg.Content,
|
||||
NewValue: rewritten,
|
||||
})
|
||||
if len(samples) < sampleBudget {
|
||||
samples = append(samples, fmt.Sprintf("doc=%s msgs.%d.msg.content updated", doc.Doc, i))
|
||||
}
|
||||
}
|
||||
}
|
||||
if doc.Msgs[i].Msg.SenderFaceURL != "" {
|
||||
rewritten, faceChanged := rewriteStringValue(doc.Msgs[i].Msg.SenderFaceURL, oldPrefix, newPrefix)
|
||||
if faceChanged {
|
||||
updates = append(updates, fieldUpdate{
|
||||
Path: fmt.Sprintf("msgs.%d.msg.sender_face_url", i),
|
||||
OldValue: doc.Msgs[i].Msg.SenderFaceURL,
|
||||
NewValue: rewritten,
|
||||
})
|
||||
if len(samples) < sampleBudget {
|
||||
samples = append(samples, fmt.Sprintf("doc=%s msgs.%d.msg.sender_face_url updated", doc.Doc, i))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return updates, len(updates) > 0, samples
|
||||
}
|
||||
|
||||
func toUpdateMap(updates []fieldUpdate) bson.M {
|
||||
sets := bson.M{}
|
||||
for _, u := range updates {
|
||||
sets[u.Path] = u.NewValue
|
||||
}
|
||||
return sets
|
||||
}
|
||||
|
||||
func buildBackupDocs(batchID, collection string, documentID any, updates []fieldUpdate) []any {
|
||||
now := time.Now()
|
||||
docs := make([]any, 0, len(updates))
|
||||
for _, u := range updates {
|
||||
docs = append(docs, bson.M{
|
||||
"batch_id": batchID,
|
||||
"collection": collection,
|
||||
"document_id": documentID,
|
||||
"field": u.Path,
|
||||
"old_value": u.OldValue,
|
||||
"new_value": u.NewValue,
|
||||
"created_at": now,
|
||||
})
|
||||
}
|
||||
return docs
|
||||
}
|
||||
|
||||
func insertBackupDocs(ctx context.Context, backupColl *mongo.Collection, batchID, collection string, documentID any, updates []fieldUpdate) error {
|
||||
if backupColl == nil || len(updates) == 0 {
|
||||
return nil
|
||||
}
|
||||
_, err := backupColl.InsertMany(ctx, buildBackupDocs(batchID, collection, documentID, updates))
|
||||
return err
|
||||
}
|
||||
|
||||
func rollbackBatch(ctx context.Context, backupColl *mongo.Collection, db *mongo.Database, params Params, logf LogFunc) error {
|
||||
cursor, err := backupColl.Find(ctx, bson.M{"batch_id": params.BatchID}, options.Find().SetSort(bson.D{
|
||||
{Key: "collection", Value: 1},
|
||||
{Key: "document_id", Value: 1},
|
||||
{Key: "field", Value: 1},
|
||||
}))
|
||||
if err != nil {
|
||||
return fmt.Errorf("find backup docs: %w", err)
|
||||
}
|
||||
defer cursor.Close(ctx)
|
||||
|
||||
type rollbackDoc struct {
|
||||
Collection string `bson:"collection"`
|
||||
Field string `bson:"field"`
|
||||
OldValue string `bson:"old_value"`
|
||||
DocumentID any `bson:"document_id"`
|
||||
}
|
||||
|
||||
grouped := map[string]map[any][]fieldUpdate{}
|
||||
for cursor.Next(ctx) {
|
||||
var doc rollbackDoc
|
||||
if err := cursor.Decode(&doc); err != nil {
|
||||
return fmt.Errorf("decode backup doc: %w", err)
|
||||
}
|
||||
if _, ok := grouped[doc.Collection]; !ok {
|
||||
grouped[doc.Collection] = map[any][]fieldUpdate{}
|
||||
}
|
||||
grouped[doc.Collection][doc.DocumentID] = append(grouped[doc.Collection][doc.DocumentID], fieldUpdate{
|
||||
Path: doc.Field,
|
||||
NewValue: doc.OldValue,
|
||||
})
|
||||
}
|
||||
if len(grouped) == 0 {
|
||||
return fmt.Errorf("no backup records found for batch_id %s", params.BatchID)
|
||||
}
|
||||
|
||||
var collectionNames []string
|
||||
for name := range grouped {
|
||||
collectionNames = append(collectionNames, name)
|
||||
}
|
||||
sort.Strings(collectionNames)
|
||||
|
||||
for _, name := range collectionNames {
|
||||
coll := db.Collection(name)
|
||||
for documentID, updates := range grouped[name] {
|
||||
if _, err := coll.UpdateByID(ctx, documentID, bson.M{"$set": toUpdateMap(updates)}); err != nil {
|
||||
return fmt.Errorf("rollback %s: %w", name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
logf("rollback complete, batch_id=%s", params.BatchID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func parseMsgFieldIndex(path string) (int, bool) {
|
||||
parts := strings.Split(path, ".")
|
||||
if len(parts) < 4 || parts[0] != "msgs" {
|
||||
return 0, false
|
||||
}
|
||||
idx, err := strconv.Atoi(parts[1])
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return idx, true
|
||||
}
|
||||
|
||||
func trimDocIDSuffix(docID string) string {
|
||||
pos := strings.LastIndex(docID, ":")
|
||||
if pos <= 0 {
|
||||
return docID
|
||||
}
|
||||
return docID[:pos]
|
||||
}
|
||||
|
||||
func deduplicateStrings(values []string) []string {
|
||||
seen := make(map[string]struct{}, len(values))
|
||||
result := make([]string, 0, len(values))
|
||||
for _, v := range values {
|
||||
if _, ok := seen[v]; ok {
|
||||
continue
|
||||
}
|
||||
seen[v] = struct{}{}
|
||||
result = append(result, v)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func nestedString(doc bson.M, path string) (string, bool) {
|
||||
current := any(doc)
|
||||
for _, segment := range strings.Split(path, ".") {
|
||||
switch m := current.(type) {
|
||||
case bson.M:
|
||||
current = m[segment]
|
||||
case map[string]any:
|
||||
current = m[segment]
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
val, ok := current.(string)
|
||||
return val, ok
|
||||
}
|
||||
Reference in New Issue
Block a user