[AI] chore: 纳入 dbseed 一次性建库播种工具
hgmj 从 91dj_video 建库时写的工具,user 表按 uid < 30000 过滤, 排除交易/分析/IM/AI/UGC 等大表。之前一直没入库。 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,165 @@
|
||||
// dbseed 把源库的一批集合播种到目标库,用于开新站。
|
||||
//
|
||||
// 与 scripts/contentsync 的区别:contentsync 是每天跑的增量同步,只认 _id 水位线;
|
||||
// dbseed 是一次性播种,可带过滤条件,按 _id 升序分批 upsert,中断后重跑幂等。
|
||||
//
|
||||
// 集合清单与过滤条件写在 collections.go,改范围改那里。
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
"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"
|
||||
)
|
||||
|
||||
func main() {
|
||||
var (
|
||||
srcURI = flag.String("src", "", "源 URI")
|
||||
srcDB = flag.String("src-db", "", "源库名")
|
||||
dstURI = flag.String("dst", "", "目标 URI")
|
||||
dstDB = flag.String("dst-db", "", "目标库名")
|
||||
only = flag.String("only", "", "只跑指定集合,逗号分隔")
|
||||
dryRun = flag.Bool("dry-run", false, "只统计不写入")
|
||||
batchSize = flag.Int("batch", 500, "每批条数")
|
||||
timeout = flag.Duration("timeout", 6*time.Hour, "整体超时")
|
||||
)
|
||||
flag.Parse()
|
||||
if *srcURI == "" || *srcDB == "" || *dstURI == "" || *dstDB == "" {
|
||||
fmt.Fprintln(os.Stderr, "必须指定 --src --src-db --dst --dst-db")
|
||||
os.Exit(2)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), *timeout)
|
||||
defer cancel()
|
||||
|
||||
srcCli, err := mongo.Connect(ctx, options.Client().ApplyURI(*srcURI))
|
||||
if err != nil {
|
||||
fatal("连接源库失败: %v", err)
|
||||
}
|
||||
defer srcCli.Disconnect(context.Background())
|
||||
dstCli, err := mongo.Connect(ctx, options.Client().ApplyURI(*dstURI))
|
||||
if err != nil {
|
||||
fatal("连接目标库失败: %v", err)
|
||||
}
|
||||
defer dstCli.Disconnect(context.Background())
|
||||
|
||||
src := srcCli.Database(*srcDB)
|
||||
dst := dstCli.Database(*dstDB)
|
||||
|
||||
wanted := map[string]bool{}
|
||||
for _, n := range strings.Split(*only, ",") {
|
||||
if n = strings.TrimSpace(n); n != "" {
|
||||
wanted[n] = true
|
||||
}
|
||||
}
|
||||
|
||||
started := time.Now()
|
||||
var total int
|
||||
failed := false
|
||||
for _, spec := range seedSpecs {
|
||||
if len(wanted) > 0 && !wanted[spec.Name] {
|
||||
continue
|
||||
}
|
||||
n, err := seedOne(ctx, src, dst, spec, *dryRun, *batchSize)
|
||||
if err != nil {
|
||||
fmt.Printf("[FAIL] %-26s %v\n", spec.Name, err)
|
||||
failed = true
|
||||
continue
|
||||
}
|
||||
total += n
|
||||
note := ""
|
||||
if spec.Filter != nil {
|
||||
note = " (带过滤)"
|
||||
}
|
||||
fmt.Printf("[ok] %-26s 写入 %7d%s\n", spec.Name, n, note)
|
||||
}
|
||||
fmt.Printf("\n合计写入 %d 条,耗时 %s\n", total, time.Since(started).Round(time.Second))
|
||||
if failed {
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func seedOne(ctx context.Context, src, dst *mongo.Database, spec SeedSpec,
|
||||
dryRun bool, batchSize int) (written int, err error) {
|
||||
|
||||
filter := bson.M{}
|
||||
if spec.Filter != nil {
|
||||
filter = spec.Filter
|
||||
}
|
||||
cur, err := src.Collection(spec.Name).Find(ctx, filter,
|
||||
options.Find().SetSort(bson.D{{Key: "_id", Value: 1}}).
|
||||
SetBatchSize(int32(batchSize)).SetNoCursorTimeout(true))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer cur.Close(ctx)
|
||||
|
||||
target := dst.Collection(spec.Name)
|
||||
buf := make([]mongo.WriteModel, 0, batchSize)
|
||||
scanned := 0
|
||||
lastLog := time.Now()
|
||||
|
||||
flush := func() error {
|
||||
if len(buf) == 0 {
|
||||
return nil
|
||||
}
|
||||
models := buf
|
||||
buf = make([]mongo.WriteModel, 0, batchSize)
|
||||
if dryRun {
|
||||
written += len(models)
|
||||
return nil
|
||||
}
|
||||
res, e := target.BulkWrite(ctx, models, options.BulkWrite().SetOrdered(false))
|
||||
if res != nil {
|
||||
written += int(res.UpsertedCount) + int(res.ModifiedCount)
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
for cur.Next(ctx) {
|
||||
var doc bson.M
|
||||
if e := cur.Decode(&doc); e != nil {
|
||||
return written, e
|
||||
}
|
||||
scanned++
|
||||
if time.Since(lastLog) >= 30*time.Second {
|
||||
fmt.Printf(" %-26s 已扫描 %d 条…\n", spec.Name, scanned)
|
||||
lastLog = time.Now()
|
||||
}
|
||||
id := doc["_id"]
|
||||
delete(doc, "_id")
|
||||
if len(doc) == 0 {
|
||||
continue
|
||||
}
|
||||
// 整条替换:目标是新库,没有需要保留的本地字段。
|
||||
// upsert 保证中断重跑幂等。
|
||||
buf = append(buf, mongo.NewUpdateOneModel().
|
||||
SetFilter(bson.M{"_id": id}).
|
||||
SetUpdate(bson.M{"$set": doc}).
|
||||
SetUpsert(true))
|
||||
if len(buf) >= batchSize {
|
||||
if e := flush(); e != nil {
|
||||
return written, e
|
||||
}
|
||||
}
|
||||
}
|
||||
if e := cur.Err(); e != nil {
|
||||
return written, e
|
||||
}
|
||||
return written, flush()
|
||||
}
|
||||
|
||||
var _ = primitive.NilObjectID
|
||||
|
||||
func fatal(f string, a ...any) {
|
||||
fmt.Fprintf(os.Stderr, f+"\n", a...)
|
||||
os.Exit(1)
|
||||
}
|
||||
Reference in New Issue
Block a user