sig-pub/pkg/storage/ck/clickhouse_batch_writer.go

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
}