tangguo_server/scripts/dbseed/main.go

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