// Copyright 2026 Arctel.net // SPDX-License-Identifier: Apache-2.0 // Package database 提供数据库连接与基础设施 package database import ( "Wavelet/pkg/config" "context" "fmt" "log" "net/url" "strconv" "strings" "time" "github.com/ClickHouse/clickhouse-go/v2" "github.com/ClickHouse/clickhouse-go/v2/lib/driver" "go.opentelemetry.io/otel/attribute" clickhouseDriver "gorm.io/driver/clickhouse" "gorm.io/gorm" "gorm.io/plugin/opentelemetry/tracing" ) const ( clickhouseMaxExecTime = 60 // ClickHouse 最大执行时间(秒) clickhouseReadTimeoutFactor = 2 // ReadTimeout 为 DialTimeout 的倍数 ) var ( // ChConn ClickHouse 原生连接实例,用于批量写入 ChConn driver.Conn chDB *gorm.DB ) func init() { if !config.Config.ClickHouse.Enabled { return } cfg := config.Config.ClickHouse if cfg.Database == "" { log.Fatalf("[ClickHouse] database name is required (expected: wavelet)\n") } opts := buildClickHouseOptions() var err error ChConn, err = clickhouse.Open(opts) if err != nil { log.Fatalf("[ClickHouse] init connection failed: %v\n", err) } if err = ChConn.Ping(context.Background()); err != nil { log.Fatalf("[ClickHouse] ping failed: %v\n", err) } chDB, err = gorm.Open(clickhouseDriver.New(clickhouseDriver.Config{ DSN: buildClickHouseDSN(), }), &gorm.Config{ SkipDefaultTransaction: true, }) if err != nil { log.Fatalf("[ClickHouse] init gorm connection failed: %v\n", err) } if err = chDB.Use( tracing.NewPlugin( tracing.WithoutMetrics(), tracing.WithAttributes( attribute.String("db.instance", cfg.Database), attribute.String("db.system", "ClickHouse"), ), ), ); err != nil { log.Fatalf("[ClickHouse] init trace failed: %v\n", err) } sqlDB, err := chDB.DB() if err != nil { log.Fatalf("[ClickHouse] load sql db failed: %v\n", err) } sqlDB.SetMaxIdleConns(cfg.MaxIdleConn) sqlDB.SetMaxOpenConns(cfg.MaxOpenConn) sqlDB.SetConnMaxLifetime(time.Duration(cfg.ConnMaxLifetime) * time.Second) log.Println("[ClickHouse] connection established successfully") } func buildClickHouseOptions() *clickhouse.Options { cfg := config.Config.ClickHouse return &clickhouse.Options{ Addr: cfg.Hosts, Auth: clickhouse.Auth{ Database: cfg.Database, Username: cfg.Username, Password: cfg.Password, }, Settings: clickhouse.Settings{ "max_execution_time": clickhouseMaxExecTime, }, Compression: &clickhouse.Compression{ Method: clickhouse.CompressionLZ4, }, DialTimeout: time.Duration(cfg.DialTimeout) * time.Second, MaxOpenConns: cfg.MaxOpenConn, MaxIdleConns: cfg.MaxIdleConn, ConnMaxLifetime: time.Duration(cfg.ConnMaxLifetime) * time.Second, ReadTimeout: time.Duration(cfg.DialTimeout*clickhouseReadTimeoutFactor) * time.Second, BlockBufferSize: cfg.BlockBufferSize, } } func buildClickHouseDSN() string { cfg := config.Config.ClickHouse chURL := &url.URL{ Scheme: "clickhouse", Host: strings.Join(cfg.Hosts, ","), Path: "/" + cfg.Database, } if cfg.Username != "" || cfg.Password != "" { chURL.User = url.UserPassword(cfg.Username, cfg.Password) } query := chURL.Query() query.Set("dial_timeout", fmt.Sprintf("%ds", cfg.DialTimeout)) query.Set("read_timeout", fmt.Sprintf("%ds", cfg.DialTimeout*clickhouseReadTimeoutFactor)) query.Set("max_execution_time", strconv.Itoa(clickhouseMaxExecTime)) chURL.RawQuery = query.Encode() return chURL.String() } // ChDB returns a context-aware GORM ClickHouse instance. func ChDB(ctx context.Context) *gorm.DB { if chDB == nil { return nil } return chDB.WithContext(ctx) } // SetChDBForTest sets the package-level ClickHouse GORM instance for testing. func SetChDBForTest(d *gorm.DB) { chDB = d } // SetChConnForTest sets the package-level native ClickHouse connection for testing. func SetChConnForTest(c driver.Conn) { ChConn = c }