-
Notifications
You must be signed in to change notification settings - Fork 0
/
seed.go
95 lines (75 loc) · 1.69 KB
/
seed.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
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
package main
import (
"context"
"encoding/json"
"fmt"
"log"
"github.com/abelanger5/postgres-fair-queue/internal/cmdutils"
"github.com/abelanger5/postgres-fair-queue/internal/dbsqlc"
"github.com/spf13/cobra"
)
// seedCmd represents the seed command
var seedCmd = &cobra.Command{
Use: "seed",
Short: "seed the database with some tasks.",
Run: func(cmd *cobra.Command, args []string) {
ctx, cancel := cmdutils.NewInterruptContext()
defer cancel()
seed(ctx)
},
}
var seedCount int
var partitions int
var batchSize int
func init() {
rootCmd.AddCommand(seedCmd)
seedCmd.PersistentFlags().IntVarP(
&seedCount,
"count",
"c",
10,
"The number of tasks to seed.",
)
seedCmd.PersistentFlags().IntVarP(
&partitions,
"partitions",
"p",
10,
"The number of partitions to seed.",
)
seedCmd.PersistentFlags().IntVarP(
&batchSize,
"batch-size",
"b",
100,
"The batch size for seeding.",
)
}
type seedData struct {
Group int `json:"group"`
Index int `json:"index"`
}
func seed(ctx context.Context) {
batchCount := 0
for remaining := seedCount; remaining > 0; remaining -= min(batchSize, remaining) {
j := min(batchSize, remaining)
log.Printf("seeding batch of %d\n", j)
batch := make([]dbsqlc.CreateTaskParams, 0)
for i := 0; i < j; i++ {
data, _ := json.Marshal(seedData{
Group: batchCount % partitions,
Index: i,
})
batch = append(batch, dbsqlc.CreateTaskParams{
Args: data,
GroupKey: fmt.Sprintf("group-%d", batchCount%partitions),
})
}
err := queries.CreateTask(ctx, pool, batch).Close()
if err != nil {
log.Fatalf("could not seed tasks: %v", err)
}
log.Printf("seeded tasks %d\n", len(batch))
batchCount++
}
}