修复日志持久化可靠性与数据库安全默认值

db/clickhouse_test.go

@@ -6,6 +6,7 @@ import (
 	"reflect"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/ClickHouse/clickhouse-go/v2"
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
@@ -133,6 +134,29 @@ func TestValidateProxyLogsSchema(t *testing.T) {
 	}
 }
 
+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
@@ -183,6 +207,24 @@ func TestInitialConnectWithIncompatibleSchemaStaysUnhealthy(t *testing.T) {
 	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{})
@@ -234,6 +276,85 @@ func TestOldGenerationCannotMarkReplacementUnhealthy(t *testing.T) {
 	}
 }
 
+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") {