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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
|
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, nodesCh chan Node, waysCh chan Way, relationsCh chan Relation, initial bool, blockSize int, worker int) error {
nodes := db.Collection("nodes")
ways := db.Collection("ways")
relations := db.Collection("relations")
opts := (new(options.BulkWriteOptions)).SetOrdered(false)
nodesBuffer := make([]mongo.WriteModel, 0, blockSize)
waysBuffer := make([]mongo.WriteModel, 0, blockSize)
relationsBuffer := make([]mongo.WriteModel, 0, blockSize)
nc := 0
wc := 0
rc := 0
for {
select {
case w := <-waysCh:
if initial {
um := mongo.NewInsertOneModel()
um.SetDocument(w)
waysBuffer = append(waysBuffer, um)
} else {
um := mongo.NewUpdateOneModel()
um.SetUpsert(true)
um.SetUpdate(w)
um.SetFilter(bson.M{"osm_id": w.OsmID})
waysBuffer = append(waysBuffer, um)
}
case n := <-nodesCh:
if initial {
um := mongo.NewInsertOneModel()
um.SetDocument(n)
nodesBuffer = append(nodesBuffer, um)
} else {
um := mongo.NewUpdateOneModel()
um.SetUpsert(true)
um.SetUpdate(n)
um.SetFilter(bson.M{"osm_id": n.OsmID})
nodesBuffer = append(nodesBuffer, um)
}
case r := <-relationsCh:
if initial {
um := mongo.NewInsertOneModel()
um.SetDocument(r)
relationsBuffer = append(relationsBuffer, um)
} else {
um := mongo.NewUpdateOneModel()
um.SetUpsert(true)
um.SetUpdate(r)
um.SetFilter(bson.M{"osm_id": r.OsmID})
relationsBuffer = append(relationsBuffer, um)
}
case <-ctx.Done():
log.Printf("[%d] saving last info in buffers...", worker)
if _, err := nodes.BulkWrite(context.Background(), nodesBuffer, opts); err != nil {
return err
}
if _, err := ways.BulkWrite(context.Background(), waysBuffer, opts); err != nil {
return err
}
if _, err := relations.BulkWrite(context.Background(), relationsBuffer, opts); err != nil {
return err
}
log.Printf("[%d] Done", worker)
return nil
}
if len(nodesBuffer) == blockSize {
nc++
log.Printf("[%d] nodes %d ways %d relations %d", worker, nc, wc, rc)
if _, err := nodes.BulkWrite(context.Background(), nodesBuffer, opts); err != nil {
return err
}
nodesBuffer = make([]mongo.WriteModel, 0)
}
if len(waysBuffer) == blockSize {
wc++
log.Printf("[%d] nodes %d ways %d relations %d", worker, nc, wc, rc)
if _, err := ways.BulkWrite(context.Background(), waysBuffer, opts); err != nil {
return err
}
waysBuffer = make([]mongo.WriteModel, 0)
}
if len(relationsBuffer) == blockSize {
rc++
log.Printf("[%d] nodes %d ways %d relations %d", worker, nc, wc, rc)
if _, err := relations.BulkWrite(context.Background(), relationsBuffer, opts); err != nil {
return err
}
relationsBuffer = make([]mongo.WriteModel, 0)
}
}
}
|
