init
This commit is contained in:
vet
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