1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
|
package main
import (
"context"
"log"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
)
func write(ctx context.Context, db *mongo.Database, insertCh chan Object, initial bool, blockSize int, worker int) error {
nodes := db.Collection("items")
opts := (new(options.BulkWriteOptions)).SetOrdered(false)
buf := make([]mongo.WriteModel, 0, blockSize)
ic := 0
for {
select {
case w := <-insertCh:
if initial {
um := mongo.NewInsertOneModel()
um.SetDocument(w)
buf = append(buf, um)
} else {
um := mongo.NewUpdateOneModel()
um.SetUpsert(true)
um.SetUpdate(w)
um.SetFilter(bson.M{"osm_id": w.ID})
buf = append(buf, um)
}
case <-ctx.Done():
if len(buf) > 0 {
log.Printf("Worker: %d\tSaving last info in buffers...", worker)
if _, err := nodes.BulkWrite(context.Background(), buf, opts); err != nil {
return err
}
}
log.Printf("Worker: %d\tDone", worker)
return nil
}
if len(buf) == blockSize {
ic++
log.Printf("Worker: %d\tWriting block %d (%d objects)", worker, ic, ic*blockSize)
if _, err := nodes.BulkWrite(context.Background(), buf, opts); err != nil {
return err
}
buf = make([]mongo.WriteModel, 0)
}
}
}
|
