package ck import ( "context" "fmt" "reflect" "sig-pub/pkg/config" "strings" "time" "github.com/ClickHouse/clickhouse-go/v2/lib/driver" clickhousev2 "github.com/ClickHouse/clickhouse-go/v2" ) type ClickhouseBatchWriter struct { cfg config.ClickhouseConfig conn driver.Conn } func NewClickhouseBatchWriter(cfg config.ClickhouseConfig) *ClickhouseBatchWriter { return &ClickhouseBatchWriter{ cfg: cfg, } } func (c *ClickhouseBatchWriter) Init() (err error) { options, err := c.cfg.Clickhousev2Options() if err != nil { return } // initial clickhouse conn c.conn, err = clickhousev2.Open(options) if err != nil { return } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if err = c.conn.Ping(ctx); err != nil { if exception, ok := err.(*clickhousev2.Exception); ok { err = fmt.Errorf("exception [%d] %s %s", exception.Code, exception.Message, exception.StackTrace) } return } return } func reflectGormData(data any) (table string, columns []string, values map[string]any, err error) { // value := reflect.ValueOf(data) // tableNameM := value.MethodByName("TableName") // _ = tableNameM.Call(nil) defer func() { if r := recover(); r != nil { err = fmt.Errorf("%T parse gorm data error %v", data, r) } }() values = make(map[string]any) refValue := reflect.ValueOf(data) refType := reflect.TypeOf(data) tableNameM := refValue.MethodByName("TableName") if !tableNameM.IsValid() { err = fmt.Errorf("type %T not has method TableName", data) return } rsp := tableNameM.Call(nil) table = rsp[0].Interface().(string) if refValue.Kind() == reflect.Ptr { refValue = refValue.Elem() } if refType.Kind() == reflect.Ptr { refType = refType.Elem() } for i := 0; i < refType.NumField(); i++ { field := refType.Field(i) value := refValue.Field(i).Interface() // value tag := field.Tag.Get("gorm") column := getGormTagColumnName(tag) if column == "" || strings.HasPrefix(column, "-") { // 忽略字段 continue } columns = append(columns, column) values[column] = value } return } func getGormTagColumnName(tag string) (column string) { i1 := strings.Index(tag, "column:") if i1 < 0 { return } i2 := strings.Index(tag[i1+7:], ";") if i2 < 0 { return tag[i1+7:] } return tag[i1+7 : i2+7] } // InsertBatch 批量插入gorm标签的结构体, 结构体用指针! func InsertBatch[T any](c *ClickhouseBatchWriter, ctx context.Context, datas []T) (err error) { if len(datas) == 0 { return } var tColumns = make(map[string][]string, 3) var tBatchs = make(map[string]driver.Batch, 3) for _, data := range datas { // 优化原生批量插入: https://clickhouse.com/docs/en/integrations/go#batch-insert table, columns, values, e := reflectGormData(data) if e != nil { err = e return } // 初始化该表 prepare if _, ok := tColumns[table]; !ok { tColumns[table] = columns prepareSql := fmt.Sprintf("INSERT INTO %s(", table) for i, column := range columns { if i == 0 { prepareSql += column } else { prepareSql += ("," + column) } } prepareSql += ") SETTINGS async_insert=1, wait_for_async_insert=0" batch, e := c.conn.PrepareBatch(ctx, prepareSql) if e != nil { err = e return } tBatchs[table] = batch } columns = tColumns[table] batch := tBatchs[table] var args = make([]any, 0, len(columns)) for _, column := range columns { args = append(args, values[column]) } err = batch.Append(args...) if err != nil { return } } // 批量插入 for _, batch := range tBatchs { err = batch.Send() if err != nil { return } } return }