test: prove unhealthy schema startup boundary

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
 	}

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
+}