token_thief/db/clickhouse.go

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()
	}
}