package db import ( "context" "errors" "fmt" "log" "net" "net/url" "strconv" "strings" "sync" "time" "github.com/ClickHouse/clickhouse-go/v2" ) // Pool 封装 ClickHouse 连接并维护健康状态,DB 故障时不阻塞调用方。 type Pool struct { dsn string reconnectInterval time.Duration open func(*clickhouse.Options) (clickhouse.Conn, error) mu sync.RWMutex conn clickhouse.Conn healthy bool closed bool } var errPoolClosed = errors.New("db pool is closed") // NewPool 创建 Pool。即使首次连接失败也返回非 nil 实例,后台会持续重试。 func NewPool(ctx context.Context, dsn string, reconnectInterval time.Duration) *Pool { p := &Pool{dsn: dsn, reconnectInterval: reconnectInterval} if err := p.connect(ctx); err != nil { log.Printf("[db] initial connect failed: %v (service continues without DB)", err) } go p.watch(ctx) return p } func (p *Pool) connect(ctx context.Context) error { p.mu.RLock() closed := p.closed p.mu.RUnlock() if closed { return errPoolClosed } cctx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() opts, err := ClickHouseOptions(p.dsn) if err != nil { return err } open := p.open if open == nil { open = clickhouse.Open } conn, err := open(opts) if err != nil { return err } if err := conn.Ping(cctx); err != nil { _ = conn.Close() return err } // 先 migrate,再 swap:避免新连接 migrate 失败时取代掉旧的可用连接。 mctx, cancel := context.WithTimeout(ctx, 10*time.Second) defer cancel() if err := migrate(mctx, conn); err != nil { _ = conn.Close() log.Printf("[db] migrate failed: %v", err) return err } p.mu.Lock() if p.closed { p.mu.Unlock() _ = conn.Close() return errPoolClosed } if p.conn != nil { _ = p.conn.Close() } p.conn = conn p.healthy = true p.mu.Unlock() log.Printf("[db] connected and migrated") return nil } // ClickHouseOptions converts the supported CLICKHOUSE_URL subset into driver options. func ClickHouseOptions(raw string) (*clickhouse.Options, error) { if !strings.Contains(raw, "://") { if err := validateHostPort(raw); err != nil { return nil, fmt.Errorf("invalid CLICKHOUSE_URL: %w", err) } return &clickhouse.Options{Addr: []string{raw}}, nil } u, err := url.Parse(raw) if err != nil { return nil, fmt.Errorf("invalid CLICKHOUSE_URL: %w", err) } if u.Scheme != "clickhouse" && u.Scheme != "clickhouses" { return nil, fmt.Errorf("invalid CLICKHOUSE_URL: unsupported scheme %q", u.Scheme) } if u.Fragment != "" { return nil, fmt.Errorf("invalid CLICKHOUSE_URL: fragment is not allowed") } if err := validateHostPort(u.Host); err != nil { return nil, fmt.Errorf("invalid CLICKHOUSE_URL: %w", err) } query, err := url.ParseQuery(u.RawQuery) if err != nil { return nil, fmt.Errorf("invalid CLICKHOUSE_URL: query: %w", err) } for key, values := range query { switch key { case "secure", "skip_verify", "compress": default: return nil, fmt.Errorf("invalid CLICKHOUSE_URL: unsupported parameter %q", key) } if len(values) != 1 { return nil, fmt.Errorf("invalid CLICKHOUSE_URL: parameter %q must occur once", key) } } wantSecure := u.Scheme == "clickhouses" if value, ok := query["secure"]; ok { secure, err := strconv.ParseBool(value[0]) if err != nil { return nil, fmt.Errorf("invalid CLICKHOUSE_URL: secure: %w", err) } if secure != wantSecure { return nil, fmt.Errorf("invalid CLICKHOUSE_URL: secure conflicts with %s", u.Scheme) } } if _, ok := query["skip_verify"]; ok && !wantSecure { return nil, fmt.Errorf("invalid CLICKHOUSE_URL: skip_verify requires clickhouses") } if value, ok := query["skip_verify"]; ok { if _, err := strconv.ParseBool(value[0]); err != nil { return nil, fmt.Errorf("invalid CLICKHOUSE_URL: skip_verify: %w", err) } } if value, ok := query["compress"]; ok { switch value[0] { case "true", "false", "none", "zstd", "lz4", "lz4hc", "gzip", "deflate", "br": default: return nil, fmt.Errorf("invalid CLICKHOUSE_URL: unsupported compression %q", value[0]) } } // Let the driver decode userinfo/path and interpret compression and TLS values. if wantSecure { query.Set("secure", "true") u.RawQuery = query.Encode() } opts, err := clickhouse.ParseDSN(u.String()) if err != nil { return nil, fmt.Errorf("invalid CLICKHOUSE_URL: %w", err) } return opts, nil } func validateHostPort(address string) error { if address == "" { return fmt.Errorf("missing host and port") } host, port, err := net.SplitHostPort(address) if err != nil { return fmt.Errorf("expected host:port: %w", err) } if host == "" { return fmt.Errorf("missing host") } n, err := strconv.ParseUint(port, 10, 16) if err != nil || n == 0 { return fmt.Errorf("invalid port %q", port) } return nil } // watch 定期探活;不健康时尝试重连。 func (p *Pool) watch(ctx context.Context) { t := time.NewTicker(p.reconnectInterval) defer t.Stop() for { select { case <-ctx.Done(): return case <-t.C: if p.Healthy() { if conn := p.Get(); conn != nil { pctx, cancel := context.WithTimeout(ctx, 3*time.Second) if err := conn.Ping(pctx); err != nil { log.Printf("[db] ping failed, marking unhealthy: %v", err) p.MarkUnhealthy() } cancel() } continue } if err := p.connect(ctx); err != nil { log.Printf("[db] reconnect failed: %v", err) } } } } // Healthy 报告连接是否可用。 func (p *Pool) Healthy() bool { p.mu.RLock() defer p.mu.RUnlock() return p.healthy } // MarkUnhealthy 由调用方在写入失败后调用。 func (p *Pool) MarkUnhealthy() { p.mu.Lock() p.healthy = false p.mu.Unlock() } // Get 返回当前连接,可能为 nil。 func (p *Pool) Get() clickhouse.Conn { p.mu.RLock() defer p.mu.RUnlock() return p.conn } // Close starts closing the underlying connection and waits within ctx. func (p *Pool) Close(ctx context.Context) error { p.mu.Lock() p.closed = true conn := p.conn p.conn = nil p.healthy = false p.mu.Unlock() if conn == nil { return nil } done := make(chan error, 1) go func() { done <- conn.Close() }() select { case err := <-done: return err case <-ctx.Done(): return ctx.Err() } }