github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/cmd/agent/daemon/clickouse_init.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"os/signal"
     8  	"syscall"
     9  	"time"
    10  
    11  	"github.com/ClickHouse/clickhouse-go/v2"
    12  	"github.com/castai/kvisor/cmd/agent/daemon/state"
    13  	"github.com/castai/kvisor/pkg/logging"
    14  	"github.com/spf13/cobra"
    15  	"github.com/spf13/pflag"
    16  )
    17  
    18  func NewClickhouseInitCommand() *cobra.Command {
    19  	var (
    20  		clickhouseAddr     string
    21  		clickhouseDatabase string
    22  		clickhouseUsername string
    23  	)
    24  
    25  	command := &cobra.Command{
    26  		Use: "clickhouse-init",
    27  		Run: func(cmd *cobra.Command, args []string) {
    28  			pflag.Parse()
    29  			log := logging.New(&logging.Config{})
    30  
    31  			ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
    32  			defer stop()
    33  
    34  			run := func(ctx context.Context) error {
    35  				conn, err := clickhouse.Open(&clickhouse.Options{
    36  					Addr: []string{clickhouseAddr},
    37  					Auth: clickhouse.Auth{
    38  						Database: clickhouseDatabase,
    39  						Username: clickhouseUsername,
    40  						Password: os.Getenv("CLICKHOUSE_PASSWORD"),
    41  					},
    42  					Settings: clickhouse.Settings{
    43  						"allow_experimental_object_type": "1",
    44  					},
    45  					MaxOpenConns: 20,
    46  				})
    47  				if err != nil {
    48  					return err
    49  				}
    50  				defer conn.Close()
    51  
    52  				if err := conn.Exec(ctx, state.ClickhouseNetflowSchema()); err != nil {
    53  					return fmt.Errorf("creating clickhouse netflow schema: %w", err)
    54  				}
    55  
    56  				return nil
    57  			}
    58  
    59  			for i := 0; i < 10; i++ {
    60  				if err := run(ctx); err != nil {
    61  					log.Warnf("init failed, will retry: %v", err)
    62  					time.Sleep(3 * time.Second)
    63  				} else {
    64  					select {
    65  					case <-ctx.Done():
    66  						return
    67  					}
    68  				}
    69  			}
    70  
    71  			log.Fatal("init failed after 10 attempts")
    72  		},
    73  	}
    74  
    75  	command.PersistentFlags().StringVar(&clickhouseAddr, "clickhouse-addr", "", "clickhouse address")
    76  	command.PersistentFlags().StringVar(&clickhouseDatabase, "clickhouse-database", "", "clickhouse database")
    77  	command.PersistentFlags().StringVar(&clickhouseUsername, "clickhouse-username", "", "clickhouse username")
    78  	return command
    79  }