From 3f228a04b049c31270326d89db99d1ad33906ab8 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Thu, 23 Jul 2026 21:52:42 +0800 Subject: [PATCH] fix(query): honor agent cursor row limits --- apps/desktop/src/stores/queryStore.ts | 22 ++- packages/app-tests/queryStore.test.ts | 214 +++++++++++++++++++++++++- 2 files changed, 231 insertions(+), 5 deletions(-) diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 60acecfbb5..1055afead4 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -42,7 +42,7 @@ import { connectionQueryExecutionSchema, connectionUsesDatabaseObjectTreeMode, e import { frontendQueryTimeoutSecsForSql, queryTimeoutSecsForConnection } from "@/lib/sql/queryTimeout"; import { queryResultNameFromPreamble, queryResultSourceLabel } from "@/lib/sql/queryResultSource"; import { sortDataGridRowIndexes, type DataGridSortDirection } from "@/lib/dataGrid/dataGridSort"; -import { normalizeResultPageSize } from "@/lib/dataGrid/paginationPageSize"; +import { MAX_RESULT_PAGE_SIZE, normalizeResultPageSize } from "@/lib/dataGrid/paginationPageSize"; import { executableStatementRanges, splitSqlStatementRanges } from "@/lib/sql/sqlStatementRanges"; import { externalSqlFileDisplayTitles, normalizeExternalSqlPath } from "@/lib/sql/sqlFileOpen"; import { clearDataGridPendingSnapshotsForTab } from "@/composables/useDataGridEditor"; @@ -131,6 +131,14 @@ function markQueryResultsRowsRaw(results: QueryResult[]): QueryResult[] { return results; } +function exactTotalFromIncompletePage(result: QueryResult, pageLimit: number | undefined, pageOffset: number | undefined, useAgentResultSession: boolean | undefined): number | undefined { + if (typeof pageLimit !== "number" || result.rows.length >= pageLimit || result.truncated === true) return undefined; + // Cursor-backed drivers must explicitly confirm exhaustion; an omitted + // has_more value may come from an older driver or an exhausted row cap. + if (useAgentResultSession && result.has_more !== false) return undefined; + return (pageOffset ?? 0) + result.rows.length; +} + function appendQueryResultSegment(previous: QueryResult, segment: QueryResult, maxRows: number): QueryResult { if (segment.execution_error) throw new Error(String(segment.rows[0]?.[0] ?? "Failed to load the next result segment")); if (previous.columns.length !== segment.columns.length || previous.columns.some((column, index) => column !== segment.columns[index])) { @@ -2814,6 +2822,7 @@ export const useQueryStore = defineStore("query", () => { result: QueryResult; pageLimit?: number; pageOffset?: number; + useAgentResultSession?: boolean; executionId: string; traceId: string; elapsed: () => string; @@ -2824,8 +2833,9 @@ export const useQueryStore = defineStore("query", () => { setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, undefined); return; } - if (typeof options.pageLimit === "number" && resultRowCount < options.pageLimit) { - setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, (options.pageOffset ?? 0) + resultRowCount); + const exactIncompletePageTotal = exactTotalFromIncompletePage(options.result, options.pageLimit, options.pageOffset, options.useAgentResultSession); + if (typeof exactIncompletePageTotal === "number") { + setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, exactIncompletePageTotal); return; } @@ -3461,6 +3471,9 @@ export const useQueryStore = defineStore("query", () => { ...(typeof pageLimit === "number" ? useAgentResultSession ? { + // Agent cursors apply maxRows cumulatively across fetched pages. + // Keep multi-page navigation available without allowing unbounded reads. + maxRows: MAX_RESULT_PAGE_SIZE, fetchSize: pageLimit, pageSize: pageLimit, resultSessionId: options?.pagination?.sessionId, @@ -3547,7 +3560,7 @@ export const useQueryStore = defineStore("query", () => { current.resultTotalRowCount = undefined; } const resultRowCount = current.result?.rows.length ?? 0; - const totalKnownFromIncompletePage = !!current.result && typeof pageLimit === "number" && resultRowCount < pageLimit; + const totalKnownFromIncompletePage = !!current.result && typeof exactTotalFromIncompletePage(current.result, pageLimit, pageOffset, useAgentResultSession) === "number"; const dataCountTarget = current.mode === "data" ? (() => { @@ -3594,6 +3607,7 @@ export const useQueryStore = defineStore("query", () => { result: current.result, pageLimit, pageOffset, + useAgentResultSession, executionId, traceId, elapsed, diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index 7978e4c2ec..923b5c09ed 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -3039,7 +3039,7 @@ test("jdbc query pagination uses result sessions without capping max rows to one assert.equal(prepareBody.options.useAgentCursor, true); assert.equal(executeBody.pageSize, 100); assert.equal(executeBody.fetchSize, 100); - assert.equal(executeBody.maxRows, undefined); + assert.equal(executeBody.maxRows, 100000); assert.equal(executeBody.clientSessionId, tabId); assert.equal(tab.resultSessionId, "session-1"); assert.equal(tab.result?.has_more, true); @@ -3049,6 +3049,71 @@ test("jdbc query pagination uses result sessions without capping max rows to one } }); +for (const pageSize of [1_000, 25_000, 100_000]) { + test(`oracle agent pagination keeps ${pageSize} rows within the bounded cursor limit`, async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + let executeBody: any; + + settingsStore.updateEditorSettings({ pageSize }); + connectionStore.addEphemeralConnection(oracleConn("oracle-1")); + const tabId = store.createTab("oracle-1", "ORCL", "Query", "query", "APP"); + + 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 LARGE_TABLE", + pageLimit: pageSize, + pageOffset: 0, + useAgentResultSession: true, + }), + { 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: ["ID"], + rows: [], + affected_rows: 0, + execution_time_ms: 1, + truncated: false, + has_more: false, + }, + ]), + { 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 LARGE_TABLE"); + + assert.equal(executeBody.pageSize, pageSize); + assert.equal(executeBody.fetchSize, pageSize); + assert.equal(executeBody.maxRows, 100000); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } + }); +} + test("mongo aggregate execution uses editor page size when pagination plan has no limit", async () => { const restoreStorage = installMemoryStorage(); setActivePinia(createPinia()); @@ -4689,6 +4754,153 @@ test("query execution keeps automatically counting total rows in the background" } }); +for (const resultState of [ + { label: "truncated", result: { truncated: true, has_more: false } }, + { label: "ambiguous exhaustion", result: {} }, +]) { + test(`oracle agent ${resultState.label} short page uses COUNT instead of inferring 10000`, async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + let countRequests = 0; + + settingsStore.updateEditorSettings({ pageSize: 100_000, autoCalculateTotalRows: true }); + connectionStore.addEphemeralConnection(oracleConn("oracle-1")); + const tabId = store.createTab("oracle-1", "ORCL", "Query", "query", "APP"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + + globalThis.fetch = withConnectionHealthMock(async (input) => { + const url = String(input); + if (url === "/api/query/prepare-pagination-plan") { + return new Response( + JSON.stringify({ + sqlToExecute: "SELECT ID FROM LARGE_TABLE", + pageLimit: 100_000, + pageOffset: 0, + countSql: "SELECT COUNT(*) FROM LARGE_TABLE", + useAgentResultSession: true, + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + ); + } + if (url === "/api/query/execute-multi") { + return new Response( + JSON.stringify([ + { + columns: ["ID"], + rows: Array.from({ length: 10_000 }, (_, index) => [index + 1]), + affected_rows: 0, + execution_time_ms: 1, + ...resultState.result, + }, + ]), + { status: 200, headers: { "Content-Type": "application/json" } }, + ); + } + if (url === "/api/query/execute") { + countRequests += 1; + return new Response(JSON.stringify({ columns: ["COUNT(*)"], rows: [[3_357_833]], 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 ID FROM LARGE_TABLE"); + await waitFor(() => tab.resultTotalRowCount === 3_357_833); + + assert.equal(countRequests, 1); + assert.equal(tab.result?.rows.length, 10_000); + assert.notEqual(tab.resultTotalRowCount, 10_000); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } + }); +} + +for (const paginationMode of [ + { label: "agent", useAgentResultSession: true, result: { truncated: false, has_more: false } }, + { label: "SQL", useAgentResultSession: false, result: {} }, +]) { + test(`${paginationMode.label} natural short page still infers the exact total`, async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const settingsStore = useSettingsStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + let countRequests = 0; + + settingsStore.updateEditorSettings({ autoCalculateTotalRows: true }); + connectionStore.addEphemeralConnection(paginationMode.useAgentResultSession ? oracleConn("conn-1") : conn("conn-1")); + const tabId = store.createTab("conn-1", "db", "Query", "query", "public"); + const tab = store.tabs.find((item) => item.id === tabId); + assert.ok(tab); + + globalThis.fetch = withConnectionHealthMock(async (input) => { + const url = String(input); + if (url === "/api/query/prepare-pagination-plan") { + return new Response( + JSON.stringify({ + sqlToExecute: "select id from users", + pageLimit: 100, + pageOffset: 200, + countSql: "select count(*) from users", + useAgentResultSession: paginationMode.useAgentResultSession, + }), + { status: 200, headers: { "Content-Type": "application/json" } }, + ); + } + if (url === "/api/query/execute-multi") { + return new Response( + JSON.stringify([ + { + columns: ["id"], + rows: Array.from({ length: 37 }, (_, index) => [index + 201]), + affected_rows: 0, + execution_time_ms: 1, + ...paginationMode.result, + }, + ]), + { status: 200, headers: { "Content-Type": "application/json" } }, + ); + } + if (url === "/api/query/execute") { + countRequests += 1; + return new Response("unexpected count", { status: 500 }); + } + 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 id from users", { pagination: { limit: 100, offset: 200 } }); + + assert.equal(tab.resultTotalRowCount, 237); + assert.equal(tab.resultTotalRowCountLoading, false); + assert.equal(countRequests, 0); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } + }); +} + test("paginated query execution keeps the previous total while refreshing it in the background", async () => { const restoreStorage = installMemoryStorage(); setActivePinia(createPinia());