From 58d902c07f3dd19d8a989b6829f49b0261834e37 Mon Sep 17 00:00:00 2001 From: zipg Date: Mon, 14 Sep 2026 13:49:59 +0800 Subject: [PATCH] fix(query): make oracle offset jump consume agent result session --- apps/desktop/src/stores/queryStore.ts | 39 +++++++++++++- crates/dbx-core/src/query_result_sql.rs | 62 +++++++++++++++++++++++ packages/app-tests/queryStore.test.ts | 67 ++++++++++++++++++++----- 3 files changed, 154 insertions(+), 14 deletions(-) diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 8c6dbf00d7..c2e958ff46 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -6532,6 +6532,10 @@ export const useQueryStore = defineStore("query", () => { } const executionSchema = connectionQueryExecutionSchema(conn, targetDatabase, targetSchema, tab.mode === "data"); + // Jumping to a non-zero offset without a live cursor session: the plan + // could not rewrite the SQL for these engines, so a plain execution + // would return the first page again (#8993). + const isOffsetJumpPage = typeof pageOffset === "number" && pageOffset > 0 && !options?.pagination?.sessionId; const frontendTimeoutSecs = frontendQueryTimeoutSecsForSql(sqlToExecute, effectiveDbType, queryTimeoutSecs); const sourceLabelDatabase = targetDatabase || conn?.database; const executionClientSessionId = options?.pagination?.clientSessionId ?? (tab.mode === "query" || tab.mode === "data" ? tabClientSessionId(tab) : undefined); @@ -6582,6 +6586,35 @@ export const useQueryStore = defineStore("query", () => { clientSession: Boolean(executionClientSessionId), }); executionDispatched = true; + if (useAgentResultSession && tab.mode === "query" && typeof pageOffset === "number" && pageOffset > 0 && !options?.pagination?.sessionId && !(tab.batchSqlExecution && tab.batchSqlExecution.total > 1)) { + return (async () => { + let sessionId: string | undefined; + let skipped = 0; + while (true) { + const pageResults = await api.executeMulti(executionConnectionId, executionDatabase, sqlToExecute, executionSchema, executionId, { + ...executionOptions, + resultSessionId: sessionId, + }); + const page = pageResults[0]; + if (!page) return pageResults; + if (skipped + page.rows.length > pageOffset) { + const start = pageOffset - skipped; + const limit = typeof pageLimit === "number" ? pageLimit : page.rows.length - start; + return [{ ...page, rows: page.rows.slice(start, start + limit) }]; + } + skipped += page.rows.length; + if (!page.has_more || !page.session_id) { + if (page.has_more) { + // The cursor session ended before the requested offset; + // returning the short page would show the wrong rows. + throw new Error("Result session ended before the requested page offset"); + } + return [{ ...page, rows: [] }]; + } + sessionId = page.session_id; + } + })(); + } return tab.batchSqlExecution && tab.batchSqlExecution.total > 1 ? api.executeMultiWithProgress( executionConnectionId, @@ -6633,8 +6666,12 @@ export const useQueryStore = defineStore("query", () => { executionPromise = (async () => { const txnSessionId = tab.txnSessionId; if (!txnSessionId) throw new Error("Manual transaction session was not initialized"); + // Offset jumps inside a manual transaction keep the legacy + // single-shot call: the session-consume loop only runs in the + // auto-commit path, and opening a cursor session here would + // strand it after returning the first page. const executeInTransaction = (sessionId: string) => - useAgentResultSession + useAgentResultSession && !isOffsetJumpPage ? api.executeInManualTransaction(sessionId, sqlToExecute, executionDatabase, executionSchema, agentProtocolQueryResultMaxRows(queryResultMaxRows), useOracleLobPreview, pageLimit, options?.pagination?.sessionId, classificationSql) : api.executeInManualTransaction(sessionId, sqlToExecute, executionDatabase, executionSchema, pageLimit ?? agentProtocolQueryResultMaxRows(queryResultMaxRows), useOracleLobPreview, undefined, undefined, classificationSql); try { diff --git a/crates/dbx-core/src/query_result_sql.rs b/crates/dbx-core/src/query_result_sql.rs index 3fbbc3a4a9..8ec918ecb1 100644 --- a/crates/dbx-core/src/query_result_sql.rs +++ b/crates/dbx-core/src/query_result_sql.rs @@ -209,6 +209,19 @@ pub fn build_query_pagination_execution_plan( plan.page_sql = paginated.sql; plan.page_limit = Some(options.pagination.limit); plan.page_offset = Some(options.pagination.offset); + if options.use_agent_cursor + && matches!( + pagination_strategy(options.database_type, PaginationContext::UserQuery), + TablePaginationStrategy::AgentMaxRows | TablePaginationStrategy::Unbounded + ) + { + // The dialect cannot rewrite the statement with server-side + // pagination (Oracle user queries, generic JDBC), so the page + // metadata alone cannot reach a non-zero offset. Keep the Agent + // result session enabled and let the client consume through to + // the requested offset (#8993). + plan.use_agent_result_session = true; + } } else if can_use_first_page_cursor && options.database_type != Some(DatabaseType::Highgo) { // Kingbase JDBC may buffer an entire result in auto-commit mode, so use // LIMIT/OFFSET whenever the statement can be rewritten safely. Keep the @@ -4606,6 +4619,55 @@ WHERE u.id = picked.id; assert!(plan.use_agent_result_session); } + #[test] + fn oracle_offset_jump_keeps_agent_result_session() { + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: "SELECT * FROM events".to_string(), + query_base_sql: "SELECT * FROM events".to_string(), + database_type: Some(DatabaseType::Oracle), + pagination: QueryPagination { limit: 100, offset: 200, session_id: None }, + use_agent_cursor: true, + first_page_uses_actual_sql: false, + }); + + // Oracle user queries are Unbounded: the SQL stays unchanged, so the + // client must consume the agent result session through to the offset. + assert_eq!(plan.sql_to_execute, "SELECT * FROM events;"); + assert_eq!(plan.page_limit, Some(100)); + assert_eq!(plan.page_offset, Some(200)); + assert!(plan.use_agent_result_session); + } + + #[test] + fn jdbc_offset_jump_keeps_agent_result_session() { + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: "SELECT * FROM events".to_string(), + query_base_sql: "SELECT * FROM events".to_string(), + database_type: Some(DatabaseType::Jdbc), + pagination: QueryPagination { limit: 100, offset: 200, session_id: None }, + use_agent_cursor: true, + first_page_uses_actual_sql: false, + }); + + assert!(plan.use_agent_result_session); + assert_eq!(plan.page_offset, Some(200)); + } + + #[test] + fn rewritten_offset_jump_does_not_use_agent_result_session() { + let plan = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { + sql: "SELECT * FROM events".to_string(), + query_base_sql: "SELECT * FROM events".to_string(), + database_type: Some(DatabaseType::Postgres), + pagination: QueryPagination { limit: 100, offset: 200, session_id: None }, + use_agent_cursor: true, + first_page_uses_actual_sql: false, + }); + + assert!(plan.sql_to_execute.contains("LIMIT")); + assert!(!plan.use_agent_result_session); + } + #[test] fn highgo_prefers_server_pagination_over_agent_cursor() { let first_page = build_query_pagination_execution_plan(QueryPaginationExecutionPlanOptions { diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index dadcf2b1e5..b3a674fba5 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -1686,7 +1686,10 @@ test("result run pin and close state persist across a restart", async () => { store = useQueryStore(); await store.initOpenTabs(); const restored = store.tabs.find((item) => item.id === tabId); - assert.deepEqual(restored?.resultRuns?.map((run) => run.id), ["run-1"]); + assert.deepEqual( + restored?.resultRuns?.map((run) => run.id), + ["run-1"], + ); assert.equal(restored?.resultRuns?.[0]?.pinned, true); } finally { restoreStorage(); @@ -1732,7 +1735,10 @@ test("bulk result-run close leaves all runs untouched when the selected run is u assert.equal(await store.closeOtherResultRuns(tabId, "run-2"), false); assert.equal(await store.closeResultRunsToLeft(tabId, "run-2"), false); assert.equal(await store.closeResultRunsToRight(tabId, "run-2"), false); - assert.deepEqual(tab.resultRuns?.map((run) => run.id), ["run-1", "run-2", "run-3"]); + assert.deepEqual( + tab.resultRuns?.map((run) => run.id), + ["run-1", "run-2", "run-3"], + ); assert.equal(tab.activeResultRunId, "run-1"); assert.deepEqual(tab.result?.rows, [[1]]); }); @@ -1787,10 +1793,7 @@ test("bulk result-run close does not rewrite a deleted session-backed snapshot", resultSessionId: "session-1", resultCacheKey: "tab:tab-1:run:run-1", }; - tab.resultRuns = [ - removedRun, - { id: "run-2", title: "Run 2", sequence: 2, sql: "select 2", createdAt: 2, result: { columns: ["two"], rows: [[2]], affected_rows: 0, execution_time_ms: 1 } }, - ]; + tab.resultRuns = [removedRun, { id: "run-2", title: "Run 2", sequence: 2, sql: "select 2", createdAt: 2, result: { columns: ["two"], rows: [[2]], affected_rows: 0, execution_time_ms: 1 } }]; tab.activeResultRunId = removedRun.id; tab.result = removedRun.result; tab.resultSessionId = removedRun.resultSessionId; @@ -1889,7 +1892,10 @@ test("closing a tab releases result payloads retained by deactivated grids", () store.closeTab(tabId, { force: true }); - assert.equal(store.tabs.some((item) => item.id === tabId), false); + assert.equal( + store.tabs.some((item) => item.id === tabId), + false, + ); assert.deepEqual(retainedResult.columns, []); assert.deepEqual(retainedResult.rows, []); assert.equal(retainedResult.mongo_documents, undefined); @@ -2181,10 +2187,12 @@ test("auto-saved result stays visible until the next run is ready", async () => assert.deepEqual(tab.result?.columns, ["run_1"]); assert.deepEqual(tab.resultRuns?.[0]?.result?.rows, [[1]]); - resolveSecondExecution?.(new Response(JSON.stringify([{ columns: ["run_2"], rows: [[2]], affected_rows: 0, execution_time_ms: 1 }]), { - status: 200, - headers: { "Content-Type": "application/json" }, - })); + resolveSecondExecution?.( + new Response(JSON.stringify([{ columns: ["run_2"], rows: [[2]], affected_rows: 0, execution_time_ms: 1 }]), { + status: 200, + headers: { "Content-Type": "application/json" }, + }), + ); await execution; assert.equal(tab.resultRuns?.length, 2); @@ -8578,7 +8586,16 @@ for (const paginationMode of [ const tab = store.tabs.find((item) => item.id === tabId); assert.ok(tab); - globalThis.fetch = withConnectionHealthMock(async (input) => { + let executeMultiCalls = 0; + const sessionPageBodies: Array = []; + const agentSessionPages = paginationMode.useAgentResultSession + ? [ + { rows: Array.from({ length: 100 }, (_, index) => [index + 1]), has_more: true, session_id: "sess-1" }, + { rows: Array.from({ length: 100 }, (_, index) => [index + 101]), has_more: true, session_id: "sess-1" }, + { rows: Array.from({ length: 37 }, (_, index) => [index + 201]), has_more: false, session_id: "sess-1" }, + ] + : []; + globalThis.fetch = withConnectionHealthMock(async (input, init) => { const url = String(input); if (url === "/api/query/prepare-pagination-plan") { return new Response( @@ -8593,14 +8610,26 @@ for (const paginationMode of [ ); } if (url === "/api/query/execute-multi") { + const pageIndex = executeMultiCalls; + let resultSessionId: string | undefined; + try { + const rawBody = input instanceof Request ? await input.clone().text() : init?.body; + resultSessionId = typeof rawBody === "string" ? JSON.parse(rawBody)?.resultSessionId : undefined; + } catch { + resultSessionId = undefined; + } + sessionPageBodies.push(resultSessionId); + executeMultiCalls += 1; + const page = paginationMode.useAgentResultSession ? agentSessionPages[pageIndex] : undefined; return new Response( JSON.stringify([ { columns: ["id"], - rows: Array.from({ length: 37 }, (_, index) => [index + 201]), + rows: page ? page.rows : Array.from({ length: 37 }, (_, index) => [index + 201]), affected_rows: 0, execution_time_ms: 1, ...paginationMode.result, + ...(page ? { has_more: page.has_more, session_id: page.session_id } : {}), }, ]), { status: 200, headers: { "Content-Type": "application/json" } }, @@ -8625,6 +8654,18 @@ for (const paginationMode of [ assert.equal(tab.resultTotalRowCount, 237); assert.equal(tab.resultTotalRowCountLoading, false); assert.equal(countRequests, 0); + assert.equal(tab.result?.rows.length, 37); + assert.equal(tab.result?.rows[0]?.[0], 201); + assert.equal(tab.result?.rows[36]?.[0], 237); + if (paginationMode.useAgentResultSession) { + // The offset jump consumed the cursor session page by page (#8993). + assert.equal(executeMultiCalls, 3); + assert.equal(sessionPageBodies[0], undefined); + assert.equal(sessionPageBodies[1], "sess-1"); + assert.equal(sessionPageBodies[2], "sess-1"); + } else { + assert.equal(executeMultiCalls, 1); + } } finally { globalThis.fetch = originalFetch; restoreStorage();