491 lines
16 KiB
Go
491 lines
16 KiB
Go
package db
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/ClickHouse/clickhouse-go/v2"
|
|
"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
|
|
)
|
|
|
|
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",
|
|
"clickhouses://localhost:9440/db?skip_verify=",
|
|
"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")
|
|
}
|
|
}
|
|
|
|
func TestMissingProxyLogColumnsReturnsAdditiveDDL(t *testing.T) {
|
|
columns := append([]schemaColumn(nil), proxyLogsColumns...)
|
|
columns = append(columns[:7], columns[8:]...)
|
|
|
|
statements, err := missingProxyLogColumns(columns)
|
|
if err != nil {
|
|
t.Fatalf("missingProxyLogColumns: %v", err)
|
|
}
|
|
want := []string{"ALTER TABLE proxy_logs ADD COLUMN IF NOT EXISTS request_truncated Bool DEFAULT false"}
|
|
if !reflect.DeepEqual(statements, want) {
|
|
t.Fatalf("statements=%q want %q", statements, want)
|
|
}
|
|
}
|
|
|
|
func TestValidateProxyLogsSchemaAllowsExtraColumns(t *testing.T) {
|
|
columns := append([]schemaColumn(nil), proxyLogsColumns...)
|
|
columns = append(columns, schemaColumn{"deployment", "String"})
|
|
table := schemaTable{"MergeTree", "toYYYYMM(started_at)", "started_at, request_id"}
|
|
if err := validateProxyLogsSchema(columns, table); err != nil {
|
|
t.Fatalf("validateProxyLogsSchema: %v", err)
|
|
}
|
|
}
|
|
|
|
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 TestMigrationDoesNotAlterIncompatibleTable(t *testing.T) {
|
|
values := schemaValues()
|
|
values = append(values[:7], values[8:]...)
|
|
candidate := newFakeClickHouseConn().withRows(&fakeRows{values: values, scanErrAt: -1})
|
|
candidate.row = fakeRow{values: []string{"ReplacingMergeTree", "toYYYYMM(started_at)", "started_at, request_id"}}
|
|
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 table")
|
|
}
|
|
for _, statement := range candidate.statements {
|
|
if strings.Contains(strings.ToUpper(statement), "ALTER TABLE") {
|
|
t.Fatalf("migration altered incompatible table: %q", statement)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestClosePreventsConcurrentConnectFromPublishing(t *testing.T) {
|
|
old := newFakeClickHouseConn()
|
|
pingStarted := make(chan struct{})
|
|
pingRelease := make(chan struct{})
|
|
candidate := newFakeClickHouseConn().withBlockedPing(pingStarted, pingRelease)
|
|
pool := &Pool{dsn: "localhost:9000", conn: old, healthy: true}
|
|
pool.open = func(*clickhouse.Options) (clickhouse.Conn, error) { return candidate, nil }
|
|
|
|
connectDone := make(chan error, 1)
|
|
go func() { connectDone <- pool.connect(context.Background()) }()
|
|
<-pingStarted
|
|
|
|
if err := pool.Close(context.Background()); err != nil {
|
|
t.Fatalf("Close: %v", err)
|
|
}
|
|
close(pingRelease)
|
|
if err := <-connectDone; err == nil {
|
|
t.Fatal("connect succeeded after pool was closed")
|
|
}
|
|
|
|
if pool.Get() != nil || pool.Healthy() {
|
|
t.Fatalf("closed pool connection = %T, healthy = %v; want nil, false", pool.Get(), pool.Healthy())
|
|
}
|
|
if old.closeCalls != 1 {
|
|
t.Fatalf("old connection close calls = %d, want 1", old.closeCalls)
|
|
}
|
|
if candidate.closeCalls != 1 {
|
|
t.Fatalf("candidate close calls = %d, want 1", candidate.closeCalls)
|
|
}
|
|
}
|
|
|
|
func TestOldGenerationCannotMarkReplacementUnhealthy(t *testing.T) {
|
|
pool := &Pool{conn: newFakeClickHouseConn(), healthy: true, generation: 1}
|
|
oldGeneration := pool.Generation()
|
|
|
|
pool.mu.Lock()
|
|
pool.conn = newFakeClickHouseConn()
|
|
pool.generation++
|
|
pool.healthy = true
|
|
pool.mu.Unlock()
|
|
|
|
pool.MarkUnhealthyGeneration(oldGeneration)
|
|
if !pool.Healthy() {
|
|
t.Fatal("old connection failure marked replacement connection unhealthy")
|
|
}
|
|
pool.MarkUnhealthyGeneration(pool.Generation())
|
|
if pool.Healthy() {
|
|
t.Fatal("current connection failure did not mark pool unhealthy")
|
|
}
|
|
}
|
|
|
|
func TestConnectionLeaseDefersReplacementCloseUntilRelease(t *testing.T) {
|
|
old := newFakeClickHouseConn()
|
|
candidate := newFakeClickHouseConn()
|
|
pool := &Pool{dsn: "localhost:9000", conn: old, healthy: true, generation: 1}
|
|
pool.open = func(*clickhouse.Options) (clickhouse.Conn, error) { return candidate, nil }
|
|
|
|
conn, generation, release := pool.Acquire()
|
|
if conn != old || generation != 1 || release == nil {
|
|
t.Fatalf("Acquire=(%T,%d,%v), want old,1,release", conn, generation, release != nil)
|
|
}
|
|
if err := pool.connect(context.Background()); err != nil {
|
|
t.Fatalf("connect: %v", err)
|
|
}
|
|
if old.closeCalls != 0 {
|
|
t.Fatalf("leased old connection closed %d times", old.closeCalls)
|
|
}
|
|
release()
|
|
if old.closeCalls != 1 {
|
|
t.Fatalf("old connection close calls=%d want 1 after release", old.closeCalls)
|
|
}
|
|
release()
|
|
if old.closeCalls != 1 {
|
|
t.Fatalf("release was not idempotent: close calls=%d", old.closeCalls)
|
|
}
|
|
}
|
|
|
|
func TestCloseWaitsForConnectionLease(t *testing.T) {
|
|
conn := newFakeClickHouseConn()
|
|
pool := &Pool{conn: conn, healthy: true, generation: 1}
|
|
_, _, release := pool.Acquire()
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
if err := pool.Close(ctx); !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("Close error=%v want canceled", err)
|
|
}
|
|
if conn.closeCalls != 0 {
|
|
t.Fatalf("leased connection closed %d times", conn.closeCalls)
|
|
}
|
|
if acquired, _, _ := pool.Acquire(); acquired != nil {
|
|
t.Fatal("Acquire succeeded after Close")
|
|
}
|
|
release()
|
|
if conn.closeCalls != 1 {
|
|
t.Fatalf("connection close calls=%d want 1 after release", conn.closeCalls)
|
|
}
|
|
}
|
|
|
|
func TestCloseWaitsForRetiredConnectionLease(t *testing.T) {
|
|
old := newFakeClickHouseConn()
|
|
current := newFakeClickHouseConn()
|
|
pool := &Pool{conn: old, healthy: true, generation: 1}
|
|
_, _, release := pool.Acquire()
|
|
pool.mu.Lock()
|
|
pool.retireCurrentLocked()
|
|
pool.conn = current
|
|
pool.generation++
|
|
pool.ensureLeaseLocked(pool.generation, current)
|
|
pool.mu.Unlock()
|
|
|
|
closeDone := make(chan error, 1)
|
|
go func() { closeDone <- pool.Close(context.Background()) }()
|
|
select {
|
|
case err := <-closeDone:
|
|
t.Fatalf("Close returned before retired lease release: %v", err)
|
|
case <-time.After(25 * time.Millisecond):
|
|
}
|
|
if old.closeCalls != 0 {
|
|
t.Fatalf("retired leased connection closed %d times", old.closeCalls)
|
|
}
|
|
release()
|
|
if err := <-closeDone; err != nil {
|
|
t.Fatalf("Close: %v", err)
|
|
}
|
|
if old.closeCalls != 1 {
|
|
t.Fatalf("retired connection close calls=%d want 1", old.closeCalls)
|
|
}
|
|
}
|
|
|
|
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
|
|
pingStarted chan struct{}
|
|
pingRelease chan struct{}
|
|
}
|
|
|
|
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) withBlockedPing(started, release chan struct{}) *fakeClickHouseConn {
|
|
c.pingStarted = started
|
|
c.pingRelease = release
|
|
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 {
|
|
if c.pingStarted != nil {
|
|
close(c.pingStarted)
|
|
<-c.pingRelease
|
|
}
|
|
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
|
|
}
|