-
Notifications
You must be signed in to change notification settings - Fork 2
/
clickouse_init.go
79 lines (68 loc) · 1.9 KB
/
clickouse_init.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
package main
import (
"context"
"fmt"
"os"
"os/signal"
"syscall"
"time"
"github.com/ClickHouse/clickhouse-go/v2"
"github.com/castai/kvisor/cmd/agent/daemon/state"
"github.com/castai/kvisor/pkg/logging"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
)
func NewClickhouseInitCommand() *cobra.Command {
var (
clickhouseAddr string
clickhouseDatabase string
clickhouseUsername string
)
command := &cobra.Command{
Use: "clickhouse-init",
Run: func(cmd *cobra.Command, args []string) {
pflag.Parse()
log := logging.New(&logging.Config{})
ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
defer stop()
run := func(ctx context.Context) error {
conn, err := clickhouse.Open(&clickhouse.Options{
Addr: []string{clickhouseAddr},
Auth: clickhouse.Auth{
Database: clickhouseDatabase,
Username: clickhouseUsername,
Password: os.Getenv("CLICKHOUSE_PASSWORD"),
},
Settings: clickhouse.Settings{
"allow_experimental_object_type": "1",
},
MaxOpenConns: 20,
})
if err != nil {
return err
}
defer conn.Close()
if err := conn.Exec(ctx, state.ClickhouseNetflowSchema()); err != nil {
return fmt.Errorf("creating clickhouse netflow schema: %w", err)
}
return nil
}
for i := 0; i < 10; i++ {
if err := run(ctx); err != nil {
log.Warnf("init failed, will retry: %v", err)
time.Sleep(3 * time.Second)
} else {
select {
case <-ctx.Done():
return
}
}
}
log.Fatal("init failed after 10 attempts")
},
}
command.PersistentFlags().StringVar(&clickhouseAddr, "clickhouse-addr", "", "clickhouse address")
command.PersistentFlags().StringVar(&clickhouseDatabase, "clickhouse-database", "", "clickhouse database")
command.PersistentFlags().StringVar(&clickhouseUsername, "clickhouse-username", "", "clickhouse username")
return command
}