This commit is contained in:
vet
2026-05-28 00:16:19 +08:00
commit 52446ccf3f
54 changed files with 4617 additions and 0 deletions

1031
internal/api/router.go Normal file

File diff suppressed because it is too large Load Diff

View 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)
}

View 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
View 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
View 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
}

View 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
}

View 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())
}
}

View 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, &params); 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
}

View 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
}

View 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, &params); 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
}

View 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
}

View 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)
}
}

View 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"`
}

View 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"`
}

View 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,
}
}

View 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)
}
}

View 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"`
}

View 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"`
}

View 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[:])
}

View 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
}

View 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
}

View File

@@ -0,0 +1,5 @@
package mongo
import "errors"
var ErrInvalidPassword = errors.New("invalid password")

View 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
}

View 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
}

View 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
}

View 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
}

View 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
}
}

View 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))
}
}

View 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
}