From f84b7d591f75c8ffeb94ebdb2aad48e70e1d986c Mon Sep 17 00:00:00 2001 From: CaIon Date: Mon, 22 Jun 2026 19:32:32 +0800 Subject: [PATCH] refactor(log): simplify ClickHouse log deletion and add unit tests Replace the per-batch ClickHouse mutation loop with a single ALTER TABLE ... DELETE, since ClickHouse DELETE is a heavy mutation that rewrites data parts and per-batch mutations are pathologically slow. Add deterministic unit tests covering ClickHouse DSN handling, main-database rejection, TTL DDL generation, log ordering, request_id backfill, and display id assignment. --- model/clickhouse_log_test.go | 126 +++++++++++++++++++++++++++++++++++ model/log.go | 33 ++++----- 2 files changed, 138 insertions(+), 21 deletions(-) create mode 100644 model/clickhouse_log_test.go diff --git a/model/clickhouse_log_test.go b/model/clickhouse_log_test.go new file mode 100644 index 00000000..7d84fea8 --- /dev/null +++ b/model/clickhouse_log_test.go @@ -0,0 +1,126 @@ +package model + +import ( + "os" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestIsClickHouseDSN(t *testing.T) { + cases := []struct { + dsn string + want bool + }{ + {"clickhouse://default:pass@localhost:9000/logs", true}, + {"tcp://localhost:9000/logs", true}, + {"http://localhost:8123/logs", true}, + {"https://localhost:8443/logs", true}, + {"postgres://root:pass@localhost:5432/db", false}, + {"postgresql://root:pass@localhost:5432/db", false}, + {"root:pass@tcp(localhost:3306)/db", false}, + {"local", false}, + {"", false}, + } + for _, c := range cases { + assert.Equalf(t, c.want, isClickHouseDSN(c.dsn), "dsn=%q", c.dsn) + } +} + +func TestNormalizeClickHouseDSN(t *testing.T) { + // https without secure gets secure=true appended + normalized := normalizeClickHouseDSN("https://default:pass@localhost:8443/logs") + assert.Contains(t, normalized, "secure=true") + assert.True(t, strings.HasPrefix(normalized, "https://")) + + // https that already specifies secure is left untouched + assert.Equal(t, + "https://localhost:8443/logs?secure=false", + normalizeClickHouseDSN("https://localhost:8443/logs?secure=false"), + ) + + // non-https schemes are returned verbatim + assert.Equal(t, "clickhouse://localhost:9000/logs", normalizeClickHouseDSN("clickhouse://localhost:9000/logs")) + assert.Equal(t, "tcp://localhost:9000/logs", normalizeClickHouseDSN("tcp://localhost:9000/logs")) +} + +func TestChooseDBRejectsClickHouseForMainDatabase(t *testing.T) { + original, had := os.LookupEnv("SQL_DSN") + t.Cleanup(func() { + if had { + require.NoError(t, os.Setenv("SQL_DSN", original)) + } else { + require.NoError(t, os.Unsetenv("SQL_DSN")) + } + }) + require.NoError(t, os.Setenv("SQL_DSN", "clickhouse://default:pass@localhost:9000/logs")) + + db, dbType, err := chooseDB("SQL_DSN", false) + require.Error(t, err) + assert.Nil(t, db) + assert.Equal(t, common.DatabaseType(""), dbType) + assert.Contains(t, err.Error(), "does not support ClickHouse") +} + +func TestClickHouseLogTTLExpression(t *testing.T) { + assert.Equal(t, "", clickHouseLogTTLExpression(0)) + assert.Equal(t, "", clickHouseLogTTLExpression(-5)) + assert.Equal(t, "toDateTime(created_at) + INTERVAL 30 DAY DELETE", clickHouseLogTTLExpression(30)) +} + +func TestClickHouseLogTTLClause(t *testing.T) { + assert.Equal(t, "", clickHouseLogTTLClause(0)) + assert.Equal(t, "\nTTL toDateTime(created_at) + INTERVAL 7 DAY DELETE", clickHouseLogTTLClause(7)) +} + +func TestClickHouseLogCreateTableSQL(t *testing.T) { + withoutTTL := clickHouseLogCreateTableSQL(0) + assert.Contains(t, withoutTTL, "CREATE TABLE IF NOT EXISTS logs") + assert.Contains(t, withoutTTL, "ENGINE = MergeTree()") + assert.Contains(t, withoutTTL, "PARTITION BY toYYYYMM(toDateTime(created_at))") + assert.Contains(t, withoutTTL, "ORDER BY (created_at, request_id)") + assert.NotContains(t, withoutTTL, "TTL ") + + withTTL := clickHouseLogCreateTableSQL(30) + assert.Contains(t, withTTL, "ORDER BY (created_at, request_id)") + assert.Contains(t, withTTL, "TTL toDateTime(created_at) + INTERVAL 30 DAY DELETE") +} + +func TestClickHouseCreateTableHasTTL(t *testing.T) { + assert.True(t, clickHouseCreateTableHasTTL("CREATE TABLE logs (...)\nTTL toDateTime(created_at) + INTERVAL 30 DAY DELETE")) + assert.True(t, clickHouseCreateTableHasTTL("CREATE TABLE logs (...) TTL toDateTime(created_at)")) + assert.False(t, clickHouseCreateTableHasTTL("CREATE TABLE logs (...)\nORDER BY (created_at, request_id)")) +} + +func TestClickHouseLogOrder(t *testing.T) { + assert.Equal(t, "created_at desc, request_id desc", clickHouseLogOrder("")) + assert.Equal(t, "logs.created_at desc, logs.request_id desc", clickHouseLogOrder("logs.")) +} + +func TestEnsureLogRequestId(t *testing.T) { + empty := &Log{} + ensureLogRequestId(empty) + assert.NotEmpty(t, empty.RequestId, "empty request id should be backfilled") + + existing := &Log{RequestId: "fixed-request-id"} + ensureLogRequestId(existing) + assert.Equal(t, "fixed-request-id", existing.RequestId, "existing request id must be preserved") + + assert.NotPanics(t, func() { ensureLogRequestId(nil) }) +} + +func TestAssignDisplayLogIds(t *testing.T) { + logs := []*Log{{}, {}, {}} + + assignDisplayLogIds(logs, 0) + assert.Equal(t, []int{1, 2, 3}, []int{logs[0].Id, logs[1].Id, logs[2].Id}) + + assignDisplayLogIds(logs, 20) + assert.Equal(t, []int{21, 22, 23}, []int{logs[0].Id, logs[1].Id, logs[2].Id}) + + assert.NotPanics(t, func() { assignDisplayLogIds(nil, 0) }) +} diff --git a/model/log.go b/model/log.go index 720bb6c8..544638c8 100644 --- a/model/log.go +++ b/model/log.go @@ -683,33 +683,24 @@ func DeleteOldLogBatch(ctx context.Context, targetTimestamp int64, limit int) (i } if common.UsingLogDatabase(common.DatabaseTypeClickHouse) { - var batchCount int64 - if err := LOG_DB.WithContext(ctx).Raw(` -SELECT count() FROM ( - SELECT created_at, request_id - FROM logs - WHERE created_at < ? - ORDER BY created_at ASC, request_id ASC - LIMIT ? -)`, targetTimestamp, limit).Scan(&batchCount).Error; err != nil { + // ClickHouse DELETE is a heavy mutation that rewrites data parts, so + // per-batch mutations would be pathologically slow. Remove all matching + // rows in a single synchronous mutation regardless of limit; the reported + // count lets the caller's progress loop complete in one pass. + total, err := CountOldLog(ctx, targetTimestamp) + if err != nil { return 0, err } - if batchCount == 0 { + if total == 0 { return 0, nil } - - if err := LOG_DB.WithContext(ctx).Exec(` -ALTER TABLE logs DELETE WHERE (created_at, request_id) IN ( - SELECT created_at, request_id - FROM logs - WHERE created_at < ? - ORDER BY created_at ASC, request_id ASC - LIMIT ? -) SETTINGS mutations_sync = 1`, targetTimestamp, limit).Error; err != nil { + if err := LOG_DB.WithContext(ctx).Exec( + "ALTER TABLE logs DELETE WHERE created_at < ? SETTINGS mutations_sync = 1", + targetTimestamp, + ).Error; err != nil { return 0, err } - - return batchCount, nil + return total, nil } result := LOG_DB.WithContext(ctx).Where("created_at < ?", targetTimestamp).Limit(limit).Delete(&Log{})