From d2189e14425991e78a4526606e0dcf25fab2f1e9 Mon Sep 17 00:00:00 2001 From: Syngnat Date: Fri, 5 Jun 2026 21:35:03 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(duckdb):=20=E4=BF=AE?= =?UTF-8?q?=E5=A4=8D=20DuckDB=20=E6=9F=A5=E8=AF=A2=E8=AF=AF=E7=94=A8?= =?UTF-8?q?=E8=BF=9E=E6=8E=A5=E8=B6=85=E6=97=B6=E5=AF=BC=E8=87=B4=E4=B8=AD?= =?UTF-8?q?=E6=96=AD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 DuckDB 查询上下文策略，避免将连接超时直接作为查询执行超时 - 调整 DBQueryWithCancel、DBQueryMulti、DBQueryIsolated 统一走查询上下文工厂 - 补充 DuckDB 查询不继承连接超时与网络型数据库保留 deadline 的回归测试 --- frontend/package.json.md5 | 2 +- internal/app/methods_db.go | 29 ++++++++--------- internal/app/methods_db_cancel_test.go | 23 +++++++++++++ internal/app/methods_db_multi_test.go | 45 ++++++++++++++++++++++++++ 4 files changed, 83 insertions(+), 16 deletions(-) diff --git a/frontend/package.json.md5 b/frontend/package.json.md5 index bed8925..7396e24 100755 --- a/frontend/package.json.md5 +++ b/frontend/package.json.md5 @@ -1 +1 @@ -0295a42fd931778d85157816d79d29e5 \ No newline at end of file +d0464f9da25e9356e61652e638c99ffe \ No newline at end of file diff --git a/internal/app/methods_db.go b/internal/app/methods_db.go index 9a15d3a..f0a2b38 100644 --- a/internal/app/methods_db.go +++ b/internal/app/methods_db.go @@ -23,6 +23,17 @@ func normalizeTestConnectionConfig(config connection.ConnectionConfig) connectio return normalized } +func newQueryExecutionContext(config connection.ConnectionConfig) (context.Context, context.CancelFunc) { + if strings.EqualFold(strings.TrimSpace(config.Type), "duckdb") { + return context.WithCancel(context.Background()) + } + timeoutSeconds := config.Timeout + if timeoutSeconds <= 0 { + timeoutSeconds = 30 + } + return utils.ContextWithTimeout(time.Duration(timeoutSeconds) * time.Second) +} + func validateTestConnectionInput(config connection.ConnectionConfig) error { dbType := strings.ToLower(strings.TrimSpace(config.Type)) if dbType == "" { @@ -600,11 +611,7 @@ func (a *App) DBQueryWithCancel(config connection.ConnectionConfig, dbName strin } query = sanitizeSQLForPgLike(resolveDDLDBType(config), query) - timeoutSeconds := runConfig.Timeout - if timeoutSeconds <= 0 { - timeoutSeconds = 30 - } - ctx, cancel := utils.ContextWithTimeout(time.Duration(timeoutSeconds) * time.Second) + ctx, cancel := newQueryExecutionContext(runConfig) defer cancel() // Store cancel function for potential manual cancellation @@ -707,11 +714,7 @@ func (a *App) DBQueryMulti(config connection.ConnectionConfig, dbName string, qu } query = sanitizeSQLForPgLike(resolveDDLDBType(config), query) - timeoutSeconds := runConfig.Timeout - if timeoutSeconds <= 0 { - timeoutSeconds = 30 - } - ctx, cancel := utils.ContextWithTimeout(time.Duration(timeoutSeconds) * time.Second) + ctx, cancel := newQueryExecutionContext(runConfig) defer cancel() a.queryMu.Lock() @@ -1033,11 +1036,7 @@ func (a *App) DBQueryIsolated(config connection.ConnectionConfig, dbName string, }() query = sanitizeSQLForPgLike(resolveDDLDBType(config), query) - timeoutSeconds := runConfig.Timeout - if timeoutSeconds <= 0 { - timeoutSeconds = 30 - } - ctx, cancel := utils.ContextWithTimeout(time.Duration(timeoutSeconds) * time.Second) + ctx, cancel := newQueryExecutionContext(runConfig) defer cancel() isReadQuery := isReadOnlySQLQuery(runConfig.Type, query) diff --git a/internal/app/methods_db_cancel_test.go b/internal/app/methods_db_cancel_test.go index 36ffe39..c41fadf 100644 --- a/internal/app/methods_db_cancel_test.go +++ b/internal/app/methods_db_cancel_test.go @@ -147,3 +147,26 @@ func TestDBQueryWithCancel_QueryIDPropagation(t *testing.T) { t.Fatalf("Expected QueryID 'test-query-id' in result, got: %s", result.QueryID) } } + +func TestNewQueryExecutionContext_UsesTimeoutForNetworkDatabases(t *testing.T) { + ctx, cancel := newQueryExecutionContext(connection.ConnectionConfig{Type: "mysql", Timeout: 7}) + defer cancel() + + deadline, ok := ctx.Deadline() + if !ok { + t.Fatal("expected network database query context to carry a deadline") + } + remaining := time.Until(deadline) + if remaining <= 0 || remaining > 8*time.Second { + t.Fatalf("expected deadline around 7s, got remaining=%s", remaining) + } +} + +func TestNewQueryExecutionContext_DoesNotApplyConnectTimeoutToDuckDBQueries(t *testing.T) { + ctx, cancel := newQueryExecutionContext(connection.ConnectionConfig{Type: "duckdb", Timeout: 1}) + defer cancel() + + if _, ok := ctx.Deadline(); ok { + t.Fatal("expected DuckDB query context to avoid connection-timeout deadline") + } +} diff --git a/internal/app/methods_db_multi_test.go b/internal/app/methods_db_multi_test.go index 1d56bed..6b1ac6d 100644 --- a/internal/app/methods_db_multi_test.go +++ b/internal/app/methods_db_multi_test.go @@ -14,6 +14,7 @@ type fakeBatchWriteDB struct { execCalls int execQueries []string lastQuery string + lastCtx context.Context queryCalls int queryMap map[string][]map[string]interface{} fieldMap map[string][]string @@ -87,12 +88,14 @@ func (f *fakeBatchWriteDB) GetTriggers(dbName, tableName string) ([]connection.T } func (f *fakeBatchWriteDB) ExecContext(ctx context.Context, query string) (int64, error) { + f.lastCtx = ctx f.execCalls++ f.execQueries = append(f.execQueries, query) return 1, nil } func (f *fakeBatchWriteDB) QueryContext(ctx context.Context, query string) ([]map[string]interface{}, []string, error) { + f.lastCtx = ctx f.queryCalls++ if err := f.queryErr[query]; err != nil { return nil, nil, err @@ -395,6 +398,48 @@ func TestDBQueryWithCancelReturnsMessagesForSQLServerQuery(t *testing.T) { } } +func TestDBQueryWithCancel_DuckDBQueriesDoNotInheritConnectTimeout(t *testing.T) { + originalNewDatabaseFunc := newDatabaseFunc + originalVerifyDriverAgentRevisionFunc := verifyDriverAgentRevisionFunc + t.Cleanup(func() { + newDatabaseFunc = originalNewDatabaseFunc + verifyDriverAgentRevisionFunc = originalVerifyDriverAgentRevisionFunc + }) + + query := "SELECT 1" + fakeDB := &fakeBatchWriteDB{ + queryMap: map[string][]map[string]interface{}{ + query: { + {"value": 1}, + }, + }, + fieldMap: map[string][]string{ + query: {"value"}, + }, + queryErr: map[string]error{}, + } + newDatabaseFunc = func(dbType string) (db.Database, error) { + return fakeDB, nil + } + verifyDriverAgentRevisionFunc = func(config connection.ConnectionConfig) error { + return nil + } + + app := NewAppWithSecretStore(secretstore.NewUnavailableStore("test")) + config := connection.ConnectionConfig{Type: "duckdb", Host: ":memory:", Timeout: 1} + + result := app.DBQueryWithCancel(config, "main", query, "duckdb-no-deadline-test") + if !result.Success { + t.Fatalf("expected DuckDB DBQueryWithCancel success, got failure: %s", result.Message) + } + if fakeDB.lastCtx == nil { + t.Fatal("expected DuckDB query path to receive a context") + } + if _, ok := fakeDB.lastCtx.Deadline(); ok { + t.Fatal("expected DuckDB query context to avoid connection-timeout deadline") + } +} + func TestDBQueryMultiUsesBatchWriteExecerForAllWriteStatements(t *testing.T) { originalNewDatabaseFunc := newDatabaseFunc t.Cleanup(func() {