/
populate.go
65 lines (56 loc) · 1.48 KB
/
populate.go
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
package cloudspanner
import (
"bytes"
"context"
"fmt"
"os"
"time"
"cloud.google.com/go/spanner"
)
func (s AppSpanner) populate(start int, end int) error {
err := s.populateTable(start, end, "Users (Id, Gold)", "(%d,10000)")
if err != nil {
return err
}
err = s.populateTable(start, end, "UserItems (Id, Amount)", "(%d,0)")
if err != nil {
return err
}
return nil
}
func (s AppSpanner) PopulateMany(numUsers int, populationBatchSize int) error {
for startPos, endPos := 1, 0; endPos < numUsers; {
if endPos+populationBatchSize < numUsers {
endPos += populationBatchSize
} else {
endPos = numUsers
}
err := s.populate(startPos, endPos)
if err != nil {
return err
}
startPos = endPos + 1
}
return nil
}
func (s AppSpanner) populateTable(start int, end int, columns string, rowTmpl string) error {
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
defer cancel()
_, err := s.client.ReadWriteTransaction(ctx, func(ctx context.Context, txn *spanner.ReadWriteTransaction) error {
b := bytes.NewBufferString("INSERT " + columns + " VALUES ")
for i := start; i < end; i++ {
b.WriteString(fmt.Sprintf(rowTmpl+",", spannerKey(uint32(i))))
}
b.WriteString(fmt.Sprintf(rowTmpl, spannerKey(uint32(end))))
stmt := spanner.Statement{
SQL: b.String(),
}
rowCount, err := txn.Update(ctx, stmt)
if err != nil {
return err
}
fmt.Fprintf(os.Stdout, "%s: %d record(s) inserted.\n", columns, rowCount)
return err
})
return err
}