From 53f2a4f215a101fefee4072c6e1d8ba4dbcb1b23 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 10 Jul 2026 21:45:13 +0800 Subject: [PATCH] fix(spark): apply selected database to queries --- .../__tests__/database/jdbcDialect.spec.ts | 24 ++++++++- apps/desktop/src/lib/database/jdbcDialect.ts | 12 +++++ apps/desktop/src/stores/queryStore.ts | 4 +- packages/app-tests/queryStore.test.ts | 54 +++++++++++++++++++ 4 files changed, 91 insertions(+), 3 deletions(-) diff --git a/apps/desktop/src/lib/__tests__/database/jdbcDialect.spec.ts b/apps/desktop/src/lib/__tests__/database/jdbcDialect.spec.ts index 2240ffc287..8c4ccdf154 100644 --- a/apps/desktop/src/lib/__tests__/database/jdbcDialect.spec.ts +++ b/apps/desktop/src/lib/__tests__/database/jdbcDialect.spec.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { effectiveDatabaseTypeForConnection, inferJdbcDialect } from "@/lib/database/jdbcDialect"; +import { connectionQueryExecutionSchema, effectiveDatabaseTypeForConnection, inferJdbcDialect } from "@/lib/database/jdbcDialect"; describe("jdbc dialect inference", () => { it("detects InterSystems IRIS and Caché JDBC connections", () => { @@ -47,3 +47,25 @@ describe("jdbc dialect inference", () => { ).toBe("sqlserver"); }); }); + +describe("query execution schema", () => { + it.each(["spark", "hive"] as const)("uses the selected database as the %s execution schema", (dbType) => { + expect(connectionQueryExecutionSchema({ db_type: dbType }, "ai_test", undefined, false)).toBe("ai_test"); + }); + + it("prefers an explicit schema for PostgreSQL", () => { + expect(connectionQueryExecutionSchema({ db_type: "postgres" }, "app", "reporting", false)).toBe("reporting"); + }); + + it("does not send a schema for MySQL database context", () => { + expect(connectionQueryExecutionSchema({ db_type: "mysql" }, "app", undefined, false)).toBeUndefined(); + }); + + it("does not change data-tab execution context", () => { + expect(connectionQueryExecutionSchema({ db_type: "spark" }, "ai_test", undefined, true)).toBeUndefined(); + }); + + it("keeps generic JDBC Databend schema fallback", () => { + expect(connectionQueryExecutionSchema({ db_type: "jdbc", connection_string: "jdbc:databend://localhost:8000/default" }, "analytics", undefined, false)).toBe("analytics"); + }); +}); diff --git a/apps/desktop/src/lib/database/jdbcDialect.ts b/apps/desktop/src/lib/database/jdbcDialect.ts index 59aee53b64..486c479701 100644 --- a/apps/desktop/src/lib/database/jdbcDialect.ts +++ b/apps/desktop/src/lib/database/jdbcDialect.ts @@ -3,6 +3,8 @@ import { isSchemaAware, usesDatabaseObjectTreeMode, usesTreeSchemaMode } from "@ type JdbcDialectConnection = Pick & Partial>; +const DATABASE_AS_EXECUTION_SCHEMA_TYPES = new Set(["hive", "spark"]); + const JDBC_DIALECT_MATCHERS: Array<{ type: DatabaseType; patterns: RegExp[] }> = [ { type: "databend", patterns: [/jdbc:databend:/i, /com\.databend\.jdbc\.DatabendDriver/i, /databend-jdbc/i] }, { type: "starrocks", patterns: [/starrocks/i] }, @@ -56,6 +58,16 @@ export function connectionUsesSchemaExecutionContext(connection?: Pick { pageOffset = options?.pagination?.offset ?? 0; } - const executionSchema = connectionUsesSchemaExecutionContext(conn) ? tab.schema || tab.database : tab.mode === "data" || connectionUsesDatabaseObjectTreeMode(conn) ? undefined : tab.schema; + const executionSchema = connectionQueryExecutionSchema(conn, tab.database, tab.schema, tab.mode === "data"); const frontendTimeoutSecs = frontendQueryTimeoutSecsForSql(sqlToExecute, effectiveDbType, queryTimeoutSecs); const sourceLabelDatabase = tab.database || conn?.database; diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index b0f9fc831c..0713cf8dee 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -57,6 +57,14 @@ function sqlServerConn(id: string): ConnectionConfig { }; } +function sparkConn(id: string): ConnectionConfig { + return { + ...conn(id), + db_type: "spark", + port: 10000, + }; +} + function withConnectionHealthMock(handler: typeof fetch): typeof fetch { return async (input, init) => { if (String(input) === "/api/connection/check-health") { @@ -3693,6 +3701,52 @@ test("query execution is scoped to the tab client session", async () => { } }); +test("Spark query execution applies the selected database as schema context", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + + connectionStore.addEphemeralConnection(sparkConn("spark-1")); + const tabId = store.createTab("spark-1", "ai_test", "Query"); + let executeBody: any; + + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + const url = String(input); + if (url === "/api/query/prepare-pagination-plan") { + return new Response(JSON.stringify({ sqlToExecute: "select * from user_orders_v2", useAgentResultSession: false }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + if (url === "/api/query/execute-multi") { + executeBody = JSON.parse(String(init?.body ?? "{}")); + return new Response(JSON.stringify([{ columns: ["order_id"], rows: [["ORD001"]], affected_rows: 0, execution_time_ms: 1 }]), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + if (url === "/api/query/analyze-editability") { + return new Response(JSON.stringify({ editable: false, reason: "complex-source" }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + await store.executeTabSql(tabId, "select * from user_orders_v2"); + + assert.equal(executeBody.database, "ai_test"); + assert.equal(executeBody.schema, "ai_test"); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + test("data tab execution uses a tab-scoped client session", async () => { const restoreStorage = installMemoryStorage(); setActivePinia(createPinia());