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

package ck

import (
	"context"
	"fmt"
	"sig-pub/pkg/config"
	"sig-pub/pkg/storage/persist"
	"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
}

// 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 := persist.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
}