fix: restore reviewable migration evidence

2026-07-10 18:26:48 +08:00
commit d1bbb5370c
42 changed files with 6419 additions and 0 deletions
@@ -0,0 +1,221 @@
+package db
+
+import (
+	"context"
+	"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
+	mu                sync.RWMutex
+	conn              clickhouse.Conn
+	healthy           bool
+}
+
+// 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 {
+	cctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+	opts, err := ClickHouseOptions(p.dsn)
+	if err != nil {
+		return err
+	}
+	conn, err := clickhouse.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.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 && value[0] != "" {
+		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 释放底层连接。
+func (p *Pool) Close() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.conn != nil {
+		_ = p.conn.Close()
+		p.conn = nil
+	}
+	p.healthy = false
+}
@@ -0,0 +1,127 @@
+package db
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestClickHouseOptionsHostPort(t *testing.T) {
+	opts, err := ClickHouseOptions("localhost:9000")
+	if err != nil {
+		t.Fatalf("ClickHouseOptions: %v", err)
+	}
+	if len(opts.Addr) != 1 || opts.Addr[0] != "localhost:9000" {
+		t.Fatalf("Addr=%v want [localhost:9000]", opts.Addr)
+	}
+	if opts.Auth.Username != "" || opts.Auth.Password != "" || opts.Auth.Database != "" {
+		t.Fatalf("Auth=%+v want empty", opts.Auth)
+	}
+}
+
+func TestClickHouseOptionsURLWithAuthAndDatabase(t *testing.T) {
+	opts, err := ClickHouseOptions("clickhouse://user%40name:p%2Fass@localhost:9000/token%20thief?compress=zstd")
+	if err != nil {
+		t.Fatalf("ClickHouseOptions: %v", err)
+	}
+	if len(opts.Addr) != 1 || opts.Addr[0] != "localhost:9000" {
+		t.Fatalf("Addr=%v want [localhost:9000]", opts.Addr)
+	}
+	if opts.Auth.Username != "user@name" {
+		t.Fatalf("Username=%q want user@name", opts.Auth.Username)
+	}
+	if opts.Auth.Password != "p/ass" {
+		t.Fatalf("Password=%q want p/ass", opts.Auth.Password)
+	}
+	if opts.Auth.Database != "token thief" {
+		t.Fatalf("Database=%q want token thief", opts.Auth.Database)
+	}
+	if opts.Compression == nil {
+		t.Fatal("Compression=nil want enabled")
+	}
+}
+
+func TestClickHouseOptionsURLWithDatabaseOnly(t *testing.T) {
+	opts, err := ClickHouseOptions("clickhouse://localhost:9000/tokenthief")
+	if err != nil {
+		t.Fatalf("ClickHouseOptions: %v", err)
+	}
+	if len(opts.Addr) != 1 || opts.Addr[0] != "localhost:9000" {
+		t.Fatalf("Addr=%v want [localhost:9000]", opts.Addr)
+	}
+	if opts.Auth.Database != "tokenthief" {
+		t.Fatalf("Database=%q want tokenthief", opts.Auth.Database)
+	}
+}
+
+func TestClickHouseOptionsTLS(t *testing.T) {
+	opts, err := ClickHouseOptions("clickhouses://localhost:9440/tokenthief?skip_verify=true")
+	if err != nil {
+		t.Fatalf("ClickHouseOptions: %v", err)
+	}
+	if opts.TLS == nil || !opts.TLS.InsecureSkipVerify {
+		t.Fatalf("TLS=%+v want InsecureSkipVerify", opts.TLS)
+	}
+}
+
+func TestClickHouseOptionsRejectsInvalidDSN(t *testing.T) {
+	tests := []string{
+		"localhost",
+		"localhost:http",
+		"clickhouse://localhost/tokenthief",
+		"clickhouse:///tokenthief",
+		"clickhouse://localhost:9000/db#fragment",
+		"clickhouse://localhost:9000/db?unknown=true",
+		"clickhouse://localhost:9000/db?compress=snappy",
+		"clickhouse://localhost:9000/db?compress=lz4&compress=zstd",
+		"clickhouse://localhost:9000/db?compress=%zz",
+		"clickhouse://localhost:9000/db?secure=true",
+		"clickhouses://localhost:9440/db?secure=false",
+		"clickhouses://localhost:9440/db?skip_verify=maybe",
+		"clickhouse://localhost:9000/db?skip_verify=true",
+		"https://localhost:9440/db",
+	}
+	for _, dsn := range tests {
+		t.Run(dsn, func(t *testing.T) {
+			if _, err := ClickHouseOptions(dsn); err == nil {
+				t.Fatalf("ClickHouseOptions(%q) succeeded", dsn)
+			}
+		})
+	}
+}
+
+func TestValidateProxyLogsSchema(t *testing.T) {
+	columns := append([]schemaColumn(nil), proxyLogsColumns...)
+	table := schemaTable{
+		engine:       "MergeTree",
+		partitionKey: "toYYYYMM(started_at)",
+		sortingKey:   "started_at, request_id",
+	}
+	if err := validateProxyLogsSchema(columns, table); err != nil {
+		t.Fatalf("validateProxyLogsSchema: %v", err)
+	}
+
+	badColumns := append([]schemaColumn(nil), columns...)
+	badColumns[8].typ = "UInt16"
+	tests := []struct {
+		name    string
+		columns []schemaColumn
+		table   schemaTable
+	}{
+		{"missing column", columns[:16], table},
+		{"wrong type", badColumns, table},
+		{"wrong engine", columns, schemaTable{"ReplacingMergeTree", table.partitionKey, table.sortingKey}},
+		{"wrong partition", columns, schemaTable{table.engine, "toDate(started_at)", table.sortingKey}},
+		{"wrong sorting", columns, schemaTable{table.engine, table.partitionKey, "request_id"}},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if err := validateProxyLogsSchema(test.columns, test.table); err == nil {
+				t.Fatal("validateProxyLogsSchema succeeded")
+			}
+		})
+	}
+
+	if !reflect.DeepEqual(columns, proxyLogsColumns) {
+		t.Fatal("schema validation mutated columns")
+	}
+}
@@ -0,0 +1,134 @@
+package db
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/ClickHouse/clickhouse-go/v2"
+)
+
+const schemaSQL = `
+CREATE TABLE IF NOT EXISTS proxy_logs (
+    request_id String,
+    method String,
+    path String,
+    query String,
+    client_ip String,
+    request_headers String,
+    request_body String,
+    request_truncated Bool DEFAULT false,
+    status_code Int32,
+    response_headers String,
+    response_body String,
+    response_truncated Bool DEFAULT false,
+    is_stream Bool DEFAULT false,
+    latency_ms Int64,
+    started_at DateTime64(3),
+    finished_at DateTime64(3),
+    error String
+)
+ENGINE = MergeTree
+PARTITION BY toYYYYMM(started_at)
+ORDER BY (started_at, request_id)
+`
+
+func migrate(ctx context.Context, conn clickhouse.Conn) error {
+	if err := conn.Exec(ctx, schemaSQL); err != nil {
+		return fmt.Errorf("create proxy_logs: %w", err)
+	}
+
+	rows, err := conn.Query(ctx, `
+SELECT name, type
+FROM system.columns
+WHERE database = currentDatabase() AND table = 'proxy_logs'
+ORDER BY position`)
+	if err != nil {
+		return fmt.Errorf("query proxy_logs columns: %w", err)
+	}
+	var columns []schemaColumn
+	for rows.Next() {
+		var column schemaColumn
+		if err := rows.Scan(&column.name, &column.typ); err != nil {
+			rows.Close()
+			return fmt.Errorf("scan proxy_logs columns: %w", err)
+		}
+		columns = append(columns, column)
+	}
+	if err := rows.Err(); err != nil {
+		rows.Close()
+		return fmt.Errorf("read proxy_logs columns: %w", err)
+	}
+	rows.Close()
+
+	var table schemaTable
+	err = conn.QueryRow(ctx, `
+SELECT engine, partition_key, sorting_key
+FROM system.tables
+WHERE database = currentDatabase() AND name = 'proxy_logs'`).Scan(
+		&table.engine, &table.partitionKey, &table.sortingKey,
+	)
+	if err != nil {
+		return fmt.Errorf("query proxy_logs table: %w", err)
+	}
+	if err := validateProxyLogsSchema(columns, table); err != nil {
+		return fmt.Errorf("incompatible proxy_logs schema: %w", err)
+	}
+	return nil
+}
+
+type schemaColumn struct {
+	name string
+	typ  string
+}
+
+type schemaTable struct {
+	engine       string
+	partitionKey string
+	sortingKey   string
+}
+
+var proxyLogsColumns = []schemaColumn{
+	{"request_id", "String"},
+	{"method", "String"},
+	{"path", "String"},
+	{"query", "String"},
+	{"client_ip", "String"},
+	{"request_headers", "String"},
+	{"request_body", "String"},
+	{"request_truncated", "Bool"},
+	{"status_code", "Int32"},
+	{"response_headers", "String"},
+	{"response_body", "String"},
+	{"response_truncated", "Bool"},
+	{"is_stream", "Bool"},
+	{"latency_ms", "Int64"},
+	{"started_at", "DateTime64(3)"},
+	{"finished_at", "DateTime64(3)"},
+	{"error", "String"},
+}
+
+func validateProxyLogsSchema(columns []schemaColumn, table schemaTable) error {
+	if len(columns) != len(proxyLogsColumns) {
+		return fmt.Errorf("got %d columns, want %d", len(columns), len(proxyLogsColumns))
+	}
+	for i, want := range proxyLogsColumns {
+		if columns[i] != want {
+			return fmt.Errorf("column %d is %s %s, want %s %s", i+1, columns[i].name, columns[i].typ, want.name, want.typ)
+		}
+	}
+	if table.engine != "MergeTree" {
+		return fmt.Errorf("engine is %q, want MergeTree", table.engine)
+	}
+	if compactExpression(table.partitionKey) != "toYYYYMM(started_at)" {
+		return fmt.Errorf("partition key is %q, want toYYYYMM(started_at)", table.partitionKey)
+	}
+	if compactExpression(table.sortingKey) != "started_at,request_id" {
+		return fmt.Errorf("sorting key is %q, want started_at, request_id", table.sortingKey)
+	}
+	return nil
+}
+
+func compactExpression(value string) string {
+	return strings.Join(strings.Fields(value), "")
+}