observe-server/common/olap/initialize.go

package olap

import (
	"context"
	"fmt"
	"log"
	"net"
	"time"

	extConfig "go-admin/config"

	"github.com/ClickHouse/clickhouse-go/v2"
	"github.com/go-admin-team/go-admin-core/sdk/config"
	"github.com/go-admin-team/go-admin-core/sdk/pkg"
	"github.com/uptrace/opentelemetry-go-extra/otelgorm"
	gormclickhouse "gorm.io/driver/clickhouse"
	"gorm.io/gorm"
	"gorm.io/gorm/schema"
)

func Setup() {
	if extConfig.ExtConfig.OLAPDB.Driver == "clickhouse" {
		setClickhouse()
		setClickhouseOrm()
	} else {
		log.Fatal("Have no other olap driver but clickhouse only.")
	}
}

func setClickhouse() {
	dialCount := 0
	debug := false
	if config.ApplicationConfig.Mode == "dev" {
		debug = true
	}
	OlapConn, err := clickhouse.Open(&clickhouse.Options{

		Addr: []string{extConfig.ExtConfig.OLAPDB.Addr},
		Auth: clickhouse.Auth{
			Database: extConfig.ExtConfig.OLAPDB.Database,
			Username: extConfig.ExtConfig.OLAPDB.Username,
			Password: extConfig.ExtConfig.OLAPDB.Password,
		},
		DialContext: func(ctx context.Context, addr string) (net.Conn, error) {
			dialCount++
			var d net.Dialer
			return d.DialContext(ctx, "tcp", addr)
		},
		Debug: debug,
		Debugf: func(format string, v ...interface{}) {
			fmt.Println(format, v)
		},
		Settings: clickhouse.Settings{
			"max_execution_time": 60,
		},
		Compression: &clickhouse.Compression{
			Method: clickhouse.CompressionLZ4,
		},
		DialTimeout:      time.Duration(10) * time.Second,
		MaxOpenConns:     extConfig.ExtConfig.OLAPDB.MaxOpenConns,
		MaxIdleConns:     extConfig.ExtConfig.OLAPDB.MaxIdleConns,
		ConnMaxLifetime:  time.Duration(10) * time.Minute,
		ConnOpenStrategy: clickhouse.ConnOpenInOrder,
		BlockBufferSize:  10,
	})
	if err != nil {
		log.Fatal(pkg.Red(extConfig.ExtConfig.OLAPDB.Driver+" connect error (conn):"), err)
	}
	defaultOlapConn = OlapConn
	log.Println(pkg.Green("clickhouse connect success (conn) !"))
}

func setClickhouseOrm() {
	dsn := fmt.Sprintf("clickhouse://%s:%s@%s/%s?dial_timeout=%ss&max_execution_time=%s",
		extConfig.ExtConfig.OLAPDB.Username,
		extConfig.ExtConfig.OLAPDB.Password,
		extConfig.ExtConfig.OLAPDB.Addr,
		extConfig.ExtConfig.OLAPDB.Database,
		"10",
		"60",
	)
	db, err := gorm.Open(gormclickhouse.Open(dsn), &gorm.Config{
		NamingStrategy: ChNamingStrategy{
			schema.NamingStrategy{TablePrefix: "otel_", SingularTable: false},
		},
	})
	if err != nil {
		log.Fatal(pkg.Red(extConfig.ExtConfig.OLAPDB.Driver+" connect error (orm) :"), err)
	}
	chdb, err := db.DB()
	if err != nil {
		log.Fatal(pkg.Red(extConfig.ExtConfig.OLAPDB.Driver+" get *sql.DB error :"), err)
	}
	chdb.SetMaxIdleConns(extConfig.ExtConfig.OLAPDB.MaxIdleConns)
	chdb.SetMaxOpenConns(extConfig.ExtConfig.OLAPDB.MaxOpenConns)
	chdb.SetConnMaxLifetime(time.Duration(10) * time.Minute)
	db.Use(otelgorm.NewPlugin())
	log.Println(pkg.Green(extConfig.ExtConfig.OLAPDB.Driver + " connect success (orm) !"))
	gormClickhouse = db
}