From 507da4fe966876dfb9f01fc0b1269d22d064820f Mon Sep 17 00:00:00 2001 From: MiMoCode Date: Fri, 10 Jul 2026 18:55:34 +0800 Subject: [PATCH] test: prove unhealthy schema startup boundary --- db/clickhouse.go | 7 +- db/clickhouse_test.go | 177 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 183 insertions(+), 1 deletion(-) diff --git a/db/clickhouse.go b/db/clickhouse.go index dac9589..69778f9 100644 --- a/db/clickhouse.go +++ b/db/clickhouse.go @@ -18,6 +18,7 @@ import ( type Pool struct { dsn string reconnectInterval time.Duration + open func(*clickhouse.Options) (clickhouse.Conn, error) mu sync.RWMutex conn clickhouse.Conn healthy bool @@ -40,7 +41,11 @@ func (p *Pool) connect(ctx context.Context) error { if err != nil { return err } - conn, err := clickhouse.Open(opts) + open := p.open + if open == nil { + open = clickhouse.Open + } + conn, err := open(opts) if err != nil { return err } diff --git a/db/clickhouse_test.go b/db/clickhouse_test.go index b5c0682..5f23d73 100644 --- a/db/clickhouse_test.go +++ b/db/clickhouse_test.go @@ -1,8 +1,14 @@ package db import ( + "context" + "errors" "reflect" + "strings" "testing" + + "github.com/ClickHouse/clickhouse-go/v2" + "github.com/ClickHouse/clickhouse-go/v2/lib/driver" ) func TestClickHouseOptionsHostPort(t *testing.T) { @@ -126,3 +132,174 @@ func TestValidateProxyLogsSchema(t *testing.T) { t.Fatal("schema validation mutated columns") } } + +func TestConnectRejectsUnhealthySchemaWithoutReplacingHealthyConnection(t *testing.T) { + tests := []struct { + name string + candidate *fakeClickHouseConn + }{ + {"migration exec", newFakeClickHouseConn().withExecError(errors.New("migration failed"))}, + {"metadata query", newFakeClickHouseConn().withQueryError(errors.New("metadata query failed"))}, + {"metadata scan", newFakeClickHouseConn().withRows(&fakeRows{values: schemaValues(), scanErrAt: 0})}, + {"metadata iteration", newFakeClickHouseConn().withRows(&fakeRows{values: schemaValues(), err: errors.New("metadata iteration failed")})}, + {"table metadata", newFakeClickHouseConn().withTableError(errors.New("table metadata failed"))}, + {"incompatible schema", newFakeClickHouseConn().withColumnType(8, "UInt16")}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + old := newFakeClickHouseConn() + pool := &Pool{dsn: "localhost:9000", conn: old, healthy: true} + pool.open = func(*clickhouse.Options) (clickhouse.Conn, error) { return test.candidate, nil } + + if err := pool.connect(context.Background()); err == nil { + t.Fatal("connect succeeded with unhealthy schema") + } + if test.candidate.closeCalls != 1 { + t.Fatalf("candidate close calls = %d, want 1", test.candidate.closeCalls) + } + if old.closeCalls != 0 || pool.Get() != old || !pool.Healthy() { + t.Fatalf("old connection was not preserved: closes=%d current=%T healthy=%v", old.closeCalls, pool.Get(), pool.Healthy()) + } + assertNonDestructiveMigration(t, test.candidate.statements) + }) + } +} + +func TestInitialConnectWithIncompatibleSchemaStaysUnhealthy(t *testing.T) { + candidate := newFakeClickHouseConn().withColumnType(8, "UInt16") + pool := &Pool{dsn: "localhost:9000"} + pool.open = func(*clickhouse.Options) (clickhouse.Conn, error) { return candidate, nil } + + if err := pool.connect(context.Background()); err == nil { + t.Fatal("connect succeeded with incompatible schema") + } + if pool.Get() != nil || pool.Healthy() { + t.Fatalf("pool connection = %T, healthy = %v; want nil, false", pool.Get(), pool.Healthy()) + } + if candidate.closeCalls != 1 { + t.Fatalf("candidate close calls = %d, want 1", candidate.closeCalls) + } + assertNonDestructiveMigration(t, candidate.statements) +} + +func assertNonDestructiveMigration(t *testing.T, statements []string) { + t.Helper() + if len(statements) == 0 || !strings.Contains(strings.ToUpper(statements[0]), "CREATE TABLE IF NOT EXISTS") { + t.Fatalf("migration statements = %q, want CREATE TABLE IF NOT EXISTS first", statements) + } + for _, statement := range statements { + upper := strings.ToUpper(statement) + if strings.Contains(upper, "DROP ") || strings.Contains(upper, "TRUNCATE ") || strings.Contains(upper, "ALTER ") { + t.Fatalf("destructive migration statement executed: %q", statement) + } + } +} + +type fakeClickHouseConn struct { + rows *fakeRows + row driver.Row + execErr error + queryErr error + statements []string + closeCalls int +} + +func newFakeClickHouseConn() *fakeClickHouseConn { + return &fakeClickHouseConn{ + rows: &fakeRows{values: schemaValues(), scanErrAt: -1}, + row: fakeRow{values: []string{"MergeTree", "toYYYYMM(started_at)", "started_at, request_id"}}, + } +} + +func (c *fakeClickHouseConn) withQueryError(err error) *fakeClickHouseConn { + c.queryErr = err + return c +} +func (c *fakeClickHouseConn) withExecError(err error) *fakeClickHouseConn { + c.execErr = err + return c +} +func (c *fakeClickHouseConn) withRows(rows *fakeRows) *fakeClickHouseConn { c.rows = rows; return c } +func (c *fakeClickHouseConn) withTableError(err error) *fakeClickHouseConn { + c.row = fakeRow{err: err} + return c +} +func (c *fakeClickHouseConn) withColumnType(index int, typ string) *fakeClickHouseConn { + c.rows.values[index][1] = typ + return c +} + +func (c *fakeClickHouseConn) Contributors() []string { return nil } +func (c *fakeClickHouseConn) ServerVersion() (*driver.ServerVersion, error) { return nil, nil } +func (c *fakeClickHouseConn) Select(context.Context, any, string, ...any) error { return nil } +func (c *fakeClickHouseConn) Query(_ context.Context, query string, _ ...any) (driver.Rows, error) { + c.statements = append(c.statements, query) + return c.rows, c.queryErr +} +func (c *fakeClickHouseConn) QueryRow(_ context.Context, query string, _ ...any) driver.Row { + c.statements = append(c.statements, query) + return c.row +} +func (c *fakeClickHouseConn) PrepareBatch(context.Context, string, ...driver.PrepareBatchOption) (driver.Batch, error) { + return nil, errors.New("unexpected PrepareBatch") +} +func (c *fakeClickHouseConn) Exec(_ context.Context, query string, _ ...any) error { + c.statements = append(c.statements, query) + return c.execErr +} +func (c *fakeClickHouseConn) AsyncInsert(context.Context, string, bool, ...any) error { return nil } +func (c *fakeClickHouseConn) Ping(context.Context) error { return nil } +func (c *fakeClickHouseConn) Stats() driver.Stats { return driver.Stats{} } +func (c *fakeClickHouseConn) Close() error { c.closeCalls++; return nil } + +type fakeRows struct { + values [][]string + index int + scanErrAt int + err error +} + +func (r *fakeRows) Next() bool { return r.index < len(r.values) } +func (r *fakeRows) Scan(dest ...any) error { + if r.index == r.scanErrAt { + return errors.New("metadata scan failed") + } + for i := range dest { + *(dest[i].(*string)) = r.values[r.index][i] + } + r.index++ + return nil +} +func (r *fakeRows) ScanStruct(any) error { return nil } +func (r *fakeRows) ColumnTypes() []driver.ColumnType { return nil } +func (r *fakeRows) Totals(...any) error { return nil } +func (r *fakeRows) Columns() []string { return []string{"name", "type"} } +func (r *fakeRows) Close() error { return nil } +func (r *fakeRows) Err() error { return r.err } +func (r *fakeRows) HasData() bool { return len(r.values) > 0 } + +type fakeRow struct { + values []string + err error +} + +func (r fakeRow) Err() error { return r.err } +func (r fakeRow) Scan(dest ...any) error { + if r.err != nil { + return r.err + } + for i := range dest { + *(dest[i].(*string)) = r.values[i] + } + return nil +} +func (r fakeRow) ScanStruct(any) error { return r.err } + +func schemaValues() [][]string { + values := make([][]string, len(proxyLogsColumns)) + for i, column := range proxyLogsColumns { + values[i] = []string{column.name, column.typ} + } + return values +}