From 5bafb68919aa68a1f148d00d4c04f65e15bd9263 Mon Sep 17 00:00:00 2001 From: zipg Date: Mon, 24 Aug 2026 18:37:30 +0800 Subject: [PATCH] fix(grid): continue Oracle manual transaction cursor for infinite scroll --- apps/desktop/src/components/grid/DataGrid.vue | 31 ++- .../DataGridReadonlyTextSelection.spec.ts | 3 +- .../dataGrid/dataGridInfiniteScroll.spec.ts | 35 ++- apps/desktop/src/lib/backend/http.ts | 2 +- apps/desktop/src/lib/backend/tauri.ts | 4 +- .../lib/dataGrid/dataGridInfiniteScroll.ts | 27 ++ .../queryStore.hiddenPrimaryKey.spec.ts | 59 ++++ apps/desktop/src/stores/queryStore.ts | 10 +- crates/dbx-core/src/query.rs | 259 ++++++++++++++++-- src-tauri/src/commands/query.rs | 10 +- 10 files changed, 410 insertions(+), 30 deletions(-) diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index 447122ca74..05ac623332 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -197,7 +197,16 @@ import { tableDataVisiblePreviewRowRange, type ResultScopedRowCache, } from "@/lib/dataGrid/dataGridLargeValues"; -import { dataGridBottomScrollTop, dataGridScrollPosition, isDataGridAtScrollBottom, isDataGridNearScrollBottom, isDataGridPrefixAppend, shouldCheckInfiniteScrollAfterScroll, type DataGridScrollPosition } from "@/lib/dataGrid/dataGridInfiniteScroll"; +import { + dataGridBottomScrollTop, + dataGridInfiniteScrollAppendCompletion, + dataGridScrollPosition, + didDataGridInfiniteScrollContextChange, + isDataGridAtScrollBottom, + isDataGridNearScrollBottom, + shouldCheckInfiniteScrollAfterScroll, + type DataGridScrollPosition, +} from "@/lib/dataGrid/dataGridInfiniteScroll"; import { resolveDataGridWheelScroll } from "@/lib/dataGrid/dataGridWheel"; import { CANVAS_DATA_GRID_ROW_HEIGHT, MAX_CANVAS_DATA_GRID_PIXEL_RATIO, canvasDataGridActionOverlayWidth, canvasDataGridActionReservedWidth, dataGridSearchMatchKey, drawCanvasDataGrid, type CanvasDevicePixelSize } from "@/lib/dataGrid/canvasDataGridRenderer"; import { DATA_GRID_DARK_STRIPED_ROW_BG, DATA_GRID_LIGHT_STRIPED_ROW_BG, dataGridActiveRowBackground } from "@/lib/dataGrid/dataGridPaintTheme"; @@ -3533,7 +3542,8 @@ function currentOrderBy(): string | undefined { watch( () => [props.countSql ?? "", props.tableMeta?.schema ?? "", props.tableMeta?.tableName ?? "", currentWhereInput() ?? "", props.database ?? "", props.connectionId ?? ""], - () => { + (values, previousValues) => { + if (!didDataGridInfiniteScrollContextChange(values, previousValues)) return; manualTotalRowCount.value = undefined; // Reset infinite-scroll allLoaded when query context changes infiniteScrollAllLoaded = false; @@ -9757,7 +9767,22 @@ watch( preservedDetailsOnNextResult = null; const shouldPreserveTranspose = preserveTransposeOnNextResult.value; preserveTransposeOnNextResult.value = false; - if (isDataGridPrefixAppend(previousResult, result)) return; + const appendCompletion = dataGridInfiniteScrollAppendCompletion(previousResult, result, { + pageSize: pageSize.value, + maxRows: infiniteScrollMaxRows.value, + }); + if (appendCompletion) { + if (infiniteScrollEnabled.value) { + currentPage.value = appendCompletion.loadedPage; + lastInfiniteScrollPage = Math.max(0, appendCompletion.loadedPage - 1); + infiniteScrollAllLoaded = appendCompletion.allLoaded; + infiniteScrollRequestedOffset = undefined; + infiniteScrollRequestedLimit = undefined; + infiniteScrollLoading.value = false; + isInfiniteScrollPaginating.value = false; + } + return; + } if (getResetScrollAfterResult()) { clearResetScrollAfterResult(); resetGridVerticalScroll(); diff --git a/apps/desktop/src/components/grid/__tests__/DataGridReadonlyTextSelection.spec.ts b/apps/desktop/src/components/grid/__tests__/DataGridReadonlyTextSelection.spec.ts index c4ea03441a..f2d4dc2011 100644 --- a/apps/desktop/src/components/grid/__tests__/DataGridReadonlyTextSelection.spec.ts +++ b/apps/desktop/src/components/grid/__tests__/DataGridReadonlyTextSelection.spec.ts @@ -33,7 +33,8 @@ describe("DataGridReadonlyTextSelection", () => { it("closes a stale read-only selection when a result is replaced", () => { const resultWatcher = dataGridSource.match(/watch\(\s*\(\) => props\.result,\s*\(result, previousResult\) => \{[\s\S]*?\n\);/)?.[0] ?? ""; - expect(resultWatcher).toContain("if (isDataGridPrefixAppend(previousResult, result)) return;"); + expect(resultWatcher).toContain("const appendCompletion = dataGridInfiniteScrollAppendCompletion(previousResult, result"); + expect(resultWatcher).toContain("if (appendCompletion) {"); expect(resultWatcher).toContain("closeReadonlyCellTextSelection();"); }); diff --git a/apps/desktop/src/lib/__tests__/dataGrid/dataGridInfiniteScroll.spec.ts b/apps/desktop/src/lib/__tests__/dataGrid/dataGridInfiniteScroll.spec.ts index 409e098f12..e4f3ee1fbe 100644 --- a/apps/desktop/src/lib/__tests__/dataGrid/dataGridInfiniteScroll.spec.ts +++ b/apps/desktop/src/lib/__tests__/dataGrid/dataGridInfiniteScroll.spec.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { dataGridBottomScrollTop, isDataGridAtScrollBottom, restoredDataGridScrollLeft } from "@/lib/dataGrid/dataGridInfiniteScroll"; +import { dataGridBottomScrollTop, dataGridInfiniteScrollAppendCompletion, didDataGridInfiniteScrollContextChange, isDataGridAtScrollBottom, restoredDataGridScrollLeft } from "@/lib/dataGrid/dataGridInfiniteScroll"; describe("data grid bottom anchoring", () => { it("keeps DOM rows anchored when scrollbar padding increases the scroll height", () => { @@ -44,3 +44,36 @@ describe("data grid horizontal restoration", () => { expect(restoredDataGridScrollLeft(320, 500, 600)).toBe(0); }); }); + +describe("data grid infinite-scroll append completion", () => { + const firstPageRows = Array.from({ length: 100 }, (_, index) => [index + 1]); + + it("stops after an incomplete cursor-backed final segment", () => { + const completion = dataGridInfiniteScrollAppendCompletion({ rows: firstPageRows, session_id: "oracle-go-4", has_more: true }, { rows: [...firstPageRows, [101], [102], [103], [104]], appended_from_row_count: 100, has_more: false }, { pageSize: 100, maxRows: 100_000 }); + + expect(completion).toEqual({ loadedPage: 2, allLoaded: true }); + }); + + it("honors cursor exhaustion when the total is an exact page multiple", () => { + const completion = dataGridInfiniteScrollAppendCompletion( + { rows: firstPageRows, session_id: "oracle-go-5", has_more: true }, + { rows: [...firstPageRows, ...Array.from({ length: 100 }, (_, index) => [index + 101])], appended_from_row_count: 100, has_more: false }, + { pageSize: 100, maxRows: 100_000 }, + ); + + expect(completion).toEqual({ loadedPage: 2, allLoaded: true }); + }); + + it("keeps offset pagination available after a full non-cursor segment", () => { + const completion = dataGridInfiniteScrollAppendCompletion({ rows: firstPageRows, has_more: false }, { rows: [...firstPageRows, ...Array.from({ length: 100 }, (_, index) => [index + 101])], appended_from_row_count: 100, has_more: false }, { pageSize: 100, maxRows: 100_000 }); + + expect(completion).toEqual({ loadedPage: 2, allLoaded: false }); + }); + + it("does not reset a completed append when equivalent metadata is replaced", () => { + const context = ["SELECT COUNT(*) FROM users", "public", "users", "", "app", "connection-1"]; + + expect(didDataGridInfiniteScrollContextChange([...context], context)).toBe(false); + expect(didDataGridInfiniteScrollContextChange([...context.slice(0, 2), "audit", ...context.slice(3)], context)).toBe(true); + }); +}); diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index f05ef8d269..193ac546ae 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -1248,7 +1248,7 @@ export async function beginManualTransaction(_connectionId: string, _database: s throw new Error("Manual transaction management is only available in the desktop app."); } -export async function executeInManualTransaction(_txnSessionId: string, _sql: string, _database: string, _schema?: string, _maxRows?: number, _tableDataPreview?: boolean): Promise { +export async function executeInManualTransaction(_txnSessionId: string, _sql: string, _database: string, _schema?: string, _maxRows?: number, _tableDataPreview?: boolean, _pageSize?: number, _resultSessionId?: string): Promise { throw new Error("Manual transaction management is only available in the desktop app."); } diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 764bf146e9..40367c33e1 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -1428,7 +1428,7 @@ export async function beginManualTransaction(connectionId: string, database: str return invoke("begin_manual_transaction", { connectionId, database, schema, catalog }); } -export async function executeInManualTransaction(txnSessionId: string, sql: string, database: string, schema?: string, maxRows?: number, tableDataPreview?: boolean): Promise { +export async function executeInManualTransaction(txnSessionId: string, sql: string, database: string, schema?: string, maxRows?: number, tableDataPreview?: boolean, pageSize?: number, resultSessionId?: string): Promise { return invoke("execute_in_manual_transaction", { txnSessionId, sql, @@ -1436,6 +1436,8 @@ export async function executeInManualTransaction(txnSessionId: string, sql: stri schema, maxRows, tableDataPreview, + pageSize, + resultSessionId, }); } diff --git a/apps/desktop/src/lib/dataGrid/dataGridInfiniteScroll.ts b/apps/desktop/src/lib/dataGrid/dataGridInfiniteScroll.ts index cb3170934c..484450c2c6 100644 --- a/apps/desktop/src/lib/dataGrid/dataGridInfiniteScroll.ts +++ b/apps/desktop/src/lib/dataGrid/dataGridInfiniteScroll.ts @@ -12,6 +12,13 @@ export interface DataGridScrollMetrics { export interface DataGridAppendResult { rows: readonly unknown[]; appended_from_row_count?: number; + session_id?: string | null; + has_more?: boolean; +} + +export interface DataGridInfiniteScrollAppendCompletion { + loadedPage: number; + allLoaded: boolean; } export function dataGridScrollPosition(top: number, left: number): DataGridScrollPosition { @@ -40,6 +47,26 @@ export function isDataGridPrefixAppend(previous: DataGridAppendResult | undefine return previous.rows.every((row, index) => row === next.rows[index]); } +export function dataGridInfiniteScrollAppendCompletion(previous: DataGridAppendResult | undefined, next: DataGridAppendResult, options: { pageSize: number; maxRows: number }): DataGridInfiniteScrollAppendCompletion | undefined { + if (!isDataGridPrefixAppend(previous, next)) return undefined; + + const pageSize = Math.max(1, Math.floor(options.pageSize)); + const maxRows = Math.max(1, Math.floor(options.maxRows)); + const appendedFromRowCount = next.appended_from_row_count!; + const appendedRowCount = next.rows.length - appendedFromRowCount; + const requestedRowCount = Math.min(pageSize, Math.max(0, maxRows - appendedFromRowCount)); + const cursorExhausted = !!previous?.session_id && previous.has_more === true && next.has_more === false; + + return { + loadedPage: Math.max(1, Math.ceil(next.rows.length / pageSize)), + allLoaded: next.rows.length >= maxRows || appendedRowCount < requestedRowCount || cursorExhausted, + }; +} + +export function didDataGridInfiniteScrollContextChange(current: readonly string[], previous: readonly string[] | undefined): boolean { + return !previous || current.length !== previous.length || current.some((value, index) => value !== previous[index]); +} + export function isDataGridAtScrollBottom(metrics: DataGridScrollMetrics, tolerance = 1): boolean { const maxScrollTop = Math.max(0, metrics.scrollHeight - metrics.clientHeight); return maxScrollTop - metrics.scrollTop <= tolerance; diff --git a/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts index 8e8f58424f..108ee312d9 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.hiddenPrimaryKey.spec.ts @@ -812,6 +812,65 @@ describe("queryStore hidden primary key editing", () => { expect(store.tabs.find((tab) => tab.id === tabId)?.result?.large_value_cells).toEqual([{ row_index: 0, column_index: 1, original_bytes: 81920 }]); }); + it("continues the Oracle cursor for page 2 in a manual transaction", async () => { + getConnectionConfig.mockReturnValue({ id: "oracle-1", name: "Oracle", db_type: "oracle", database: "ORCL", query_timeout_secs: 30 }); + analyzeEditableQueryEditability.mockResolvedValue({ editable: false, reason: "complex-query" }); + prepareQueryPaginationExecutionPlan.mockImplementation(async (options) => ({ + sqlToExecute: options.sql, + pageSql: options.sql, + pageLimit: options.pagination.limit, + pageOffset: options.pagination.offset, + countSql: undefined, + useAgentResultSession: true, + })); + executeInManualTransaction + .mockResolvedValueOnce([ + { + columns: ["ID"], + rows: Array.from({ length: 100 }, (_, index) => [index + 1]), + affected_rows: 0, + execution_time_ms: 1, + session_id: "oracle-go-1", + has_more: true, + }, + ]) + .mockResolvedValueOnce([ + { + columns: ["ID"], + rows: Array.from({ length: 100 }, (_, index) => [index + 101]), + affected_rows: 0, + execution_time_ms: 1, + has_more: false, + }, + ]); + + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("oracle-1", "ORCL", "Query"); + // Tabs default to auto-commit; exercise the explicit manual-transaction path. + store.setAutoCommit(tabId, false); + const sql = "SELECT ID FROM APP.EVENTS ORDER BY ID"; + + await store.executeTabSql(tabId, sql); + const tab = store.tabs.find((item) => item.id === tabId)!; + expect(tab.autoCommit).toBe(false); + expect(executeInManualTransaction).toHaveBeenNthCalledWith(1, "txn-1", sql, "ORCL", undefined, expect.any(Number), false, 100, undefined); + expect(tab.result?.rows).toHaveLength(100); + + await store.executeTabSql(tabId, sql, { + resultBaseSql: sql, + pagination: { limit: 100, offset: 100, sessionId: "oracle-go-1" }, + appendResult: { maxRows: 10_000 }, + preserveResultDuringExecution: true, + preserveTotalRowCountDuringExecution: true, + replaceActiveResultInGroup: true, + }); + + expect(executeInManualTransaction).toHaveBeenNthCalledWith(2, "txn-1", sql, "ORCL", undefined, expect.any(Number), false, 100, "oracle-go-1"); + expect(tab.result?.rows).toHaveLength(200); + expect(tab.result?.rows[100]).toEqual([101]); + }); + it("does not start an Oracle manual transaction after cancellation during metadata loading", async () => { const columnsGate = deferred>>(); getConnectionConfig.mockReturnValue({ id: "oracle-1", name: "Oracle", db_type: "oracle", database: "ORCL", query_timeout_secs: 30 }); diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index d6444f3415..ef438ce573 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -5139,10 +5139,14 @@ export const useQueryStore = defineStore("query", () => { executionPromise = (async () => { const txnSessionId = tab.txnSessionId; if (!txnSessionId) throw new Error("Manual transaction session was not initialized"); + const executeInTransaction = (sessionId: string) => + useAgentResultSession + ? api.executeInManualTransaction(sessionId, sqlToExecute, executionDatabase, executionSchema, agentProtocolQueryResultMaxRows(queryResultMaxRows), useOracleLobPreview, pageLimit, options?.pagination?.sessionId) + : api.executeInManualTransaction(sessionId, sqlToExecute, executionDatabase, executionSchema, pageLimit ?? agentProtocolQueryResultMaxRows(queryResultMaxRows), useOracleLobPreview); try { - return await api.executeInManualTransaction(txnSessionId, sqlToExecute, executionDatabase, executionSchema, pageLimit ?? agentProtocolQueryResultMaxRows(queryResultMaxRows), useOracleLobPreview); + return await executeInTransaction(txnSessionId); } catch (error) { - if (manualTransactionRecoveryAttempted || !isManualTransactionSessionExpired(error)) throw error; + if (options?.pagination?.sessionId || manualTransactionRecoveryAttempted || !isManualTransactionSessionExpired(error)) throw error; manualTransactionRecoveryAttempted = true; tab.txnSessionId = undefined; tab.txnAutoRolledBack = true; @@ -5150,7 +5154,7 @@ export const useQueryStore = defineStore("query", () => { const refreshedSessionId = await api.beginManualTransaction(executionConnectionId, executionDatabase, executionSchema, executionCatalog); tab.txnSessionId = refreshedSessionId; queryExecutionLog("info", "manual-txn:restarted", { traceId, txnSessionId: refreshedSessionId, elapsed: elapsed() }); - return api.executeInManualTransaction(refreshedSessionId, sqlToExecute, executionDatabase, executionSchema, pageLimit ?? agentProtocolQueryResultMaxRows(queryResultMaxRows), useOracleLobPreview); + return executeInTransaction(refreshedSessionId); } })(); } else { diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 147b5b8b9e..efe051972d 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -4577,9 +4577,24 @@ pub async fn execute_in_manual_transaction( schema: Option<&str>, max_rows: Option, ) -> Result, String> { - execute_in_manual_transaction_with_options(state, txn_session_id, sql, database, schema, max_rows, false) - .await - .map(|results| results.into_iter().map(ExecuteMultiResult::into_query_result).collect()) + execute_in_manual_transaction_with_options( + state, + txn_session_id, + sql, + database, + schema, + ManualTransactionExecutionOptions { max_rows, ..Default::default() }, + ) + .await + .map(|results| results.into_iter().map(ExecuteMultiResult::into_query_result).collect()) +} + +#[derive(Clone, Debug, Default)] +pub struct ManualTransactionExecutionOptions { + pub max_rows: Option, + pub table_data_preview: bool, + pub page_size: Option, + pub result_session_id: Option, } pub async fn execute_in_manual_transaction_with_options( @@ -4588,8 +4603,7 @@ pub async fn execute_in_manual_transaction_with_options( sql: &str, database: &str, schema: Option<&str>, - max_rows: Option, - table_data_preview: bool, + options: ManualTransactionExecutionOptions, ) -> Result, String> { const TXN_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(MANUAL_TRANSACTION_IDLE_TIMEOUT_SECS); @@ -4610,9 +4624,21 @@ pub async fn execute_in_manual_transaction_with_options( if statements.is_empty() { return Ok(vec![ExecuteMultiResult::success_with_optional_server_large_values( empty_query_result(0), - table_data_preview, + options.table_data_preview, )]); } + // A result-session cursor can only track one result set, so pagination is + // single-statement only. Multi-statement scripts predate pagination: keep + // the legacy sequential execution and ignore the pagination options rather + // than failing the whole script. + let options = if statements.len() != 1 && (options.page_size.is_some() || options.result_session_id.is_some()) { + ManualTransactionExecutionOptions { page_size: None, result_session_id: None, ..options } + } else { + options + }; + if options.result_session_id.is_some() && options.page_size.is_none() { + return Err("Manual transaction result pagination requires a page size".to_string()); + } // Read-only check while the session is still in the map. If this fails the // session remains intact. @@ -4649,7 +4675,7 @@ pub async fn execute_in_manual_transaction_with_options( .map(|session| Arc::clone(&session.connection)) .ok_or(MANUAL_TRANSACTION_SESSION_NOT_FOUND_ERROR)? }; - let row_limit = max_rows.unwrap_or(MAX_ROWS).max(1); + let row_limit = options.max_rows.unwrap_or(MAX_ROWS).max(1); let mut results = Vec::with_capacity(statements.len()); let mut conn = connection.lock().await; @@ -4667,7 +4693,9 @@ pub async fn execute_in_manual_transaction_with_options( database, schema, row_limit, - table_data_preview, + options.table_data_preview, + options.page_size, + options.result_session_id.as_deref(), ) .await } @@ -4680,7 +4708,7 @@ pub async fn execute_in_manual_transaction_with_options( Ok(query_result) => { results.push(ExecuteMultiResult::success_with_optional_server_large_values( query_result, - table_data_preview, + options.table_data_preview, )); } Err(e) => { @@ -4911,8 +4939,46 @@ async fn release_manual_txn_session_pool(state: &AppState, connection_id: &str, } } -fn manual_txn_agent_query_options(row_limit: usize, table_data_preview: bool) -> QueryExecutionOptions { - QueryExecutionOptions { max_rows: Some(row_limit.max(1)), table_data_preview, ..QueryExecutionOptions::default() } +fn manual_txn_agent_query_options( + row_limit: usize, + table_data_preview: bool, + page_size: Option, + result_session_id: Option<&str>, +) -> QueryExecutionOptions { + QueryExecutionOptions { + max_rows: Some(row_limit.max(1)), + table_data_preview, + page_size, + result_session_id: result_session_id.map(str::to_owned), + ..QueryExecutionOptions::default() + } +} + +#[derive(Debug)] +enum ManualTxnAgentQueryRequest { + Execute(serde_json::Value), + ExecutePage(serde_json::Value), + FetchPage(serde_json::Value), +} + +fn manual_txn_agent_query_request( + sql: &str, + database: Option<&str>, + schema: Option<&str>, + options: QueryExecutionOptions, +) -> ManualTxnAgentQueryRequest { + if let Some(session_id) = options.result_session_id.as_deref() { + return ManualTxnAgentQueryRequest::FetchPage(agent_fetch_query_page_params( + session_id, + options.page_size.unwrap_or(MAX_ROWS), + )); + } + if options.page_size.is_some() { + return ManualTxnAgentQueryRequest::ExecutePage(agent_execute_query_page_params( + sql, database, schema, options, + )); + } + ManualTxnAgentQueryRequest::Execute(agent_execute_query_params(sql, database, schema, options)) } async fn execute_manual_txn_agent_statement( @@ -4923,16 +4989,31 @@ async fn execute_manual_txn_agent_statement( schema: Option<&str>, row_limit: usize, table_data_preview: bool, + page_size: Option, + result_session_id: Option<&str>, ) -> Result { let sql = sql_for_execution_context(db_type, statement, schema); let execution_schema = schema_for_execution_context(db_type, schema); - let options = manual_txn_agent_query_options(row_limit, table_data_preview); - let params = - agent_execute_query_params(&sql, Some(database).filter(|value| !value.is_empty()), execution_schema, options); + let options = manual_txn_agent_query_options(row_limit, table_data_preview, page_size, result_session_id); + let request = manual_txn_agent_query_request( + &sql, + Some(database).filter(|value| !value.is_empty()), + execution_schema, + options, + ); let mut locked = client.lock().await; - locked - .execute_query_typed_with_timeout::(params, None) - .await + let result = match request { + ManualTxnAgentQueryRequest::Execute(params) => { + locked.execute_query_typed_with_timeout::(params, None).await + } + ManualTxnAgentQueryRequest::ExecutePage(params) => { + locked.execute_query_page_typed_with_timeout_and_cancel::(params, None, None).await + } + ManualTxnAgentQueryRequest::FetchPage(params) => { + locked.fetch_query_page_typed_with_timeout_and_cancel::(params, None, None).await + } + }; + result .map(|result| truncate_result_with_max_rows(result, Some(row_limit.max(1)))) .map_err(|error| error.into_legacy_string()) } @@ -7589,7 +7670,7 @@ for line in sys.stdin: "SELECT * FROM documents", Some("ORCL"), Some("APP"), - manual_txn_agent_query_options(250, true), + manual_txn_agent_query_options(250, true, None, None), ); assert_eq!(params["maxRows"], 250); @@ -7810,6 +7891,148 @@ for line in sys.stdin: assert_eq!(params["pageSize"], 500); } + #[test] + fn manual_transaction_agent_query_uses_cursor_for_later_pages() { + let first_page = manual_txn_agent_query_request( + "SELECT ID FROM APP.EVENTS ORDER BY ID", + Some("ORCL"), + Some("APP"), + manual_txn_agent_query_options(10_000, false, Some(100), None), + ); + let ManualTxnAgentQueryRequest::ExecutePage(first_params) = first_page else { + panic!("first page must start an Agent query cursor"); + }; + assert_eq!(first_params["pageSize"], 100); + assert_eq!(first_params["maxRows"], 10_000); + + let second_page = manual_txn_agent_query_request( + "SELECT ID FROM APP.EVENTS ORDER BY ID", + Some("ORCL"), + Some("APP"), + manual_txn_agent_query_options(10_000, false, Some(100), Some("oracle-go-1")), + ); + let ManualTxnAgentQueryRequest::FetchPage(second_params) = second_page else { + panic!("later pages must continue the existing Agent query cursor"); + }; + assert_eq!(second_params["sessionId"], "oracle-go-1"); + assert_eq!(second_params["pageSize"], 100); + assert!(second_params.get("sql").is_none()); + } + + /// Spawns a fake Python agent and registers a manual transaction session in + /// the app state so `execute_in_manual_transaction_with_options` can run + /// end to end without a live database. + #[cfg(unix)] + async fn manual_transaction_test_state(db_type: DatabaseType) -> (AppState, String, std::path::PathBuf) { + use std::io::Write; + + let mut script = tempfile::NamedTempFile::new().unwrap(); + write!( + script, + r#"import json +import sys + +print(json.dumps({{"ready": True}}), flush=True) +for line in sys.stdin: + request = json.loads(line) + method = request.get("method") + if method == "execute_query": + sql = request.get("params", {{}}).get("sql", "") + row = [sql] + else: + # commit_manual_transaction / rollback_manual_transaction / disconnect + row = [] + result = {{ + "columns": ["stmt"], + "column_types": [], + "column_sortables": [], + "rows": [row], + "affected_rows": 0, + "execution_time_ms": 1, + "truncated": False, + "session_id": None, + "has_more": False + }} + print(json.dumps({{"jsonrpc": "2.0", "id": request["id"], "result": result}}), flush=True) +"# + ) + .unwrap(); + script.flush().unwrap(); + + let client = AgentDriverClient::spawn( + AgentLaunchSpec::new("python3").with_args([script.path().to_string_lossy().to_string()]), + ) + .await + .unwrap(); + + let dir = std::env::temp_dir().join(format!("dbx-manual-txn-{}", uuid::Uuid::new_v4().simple())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let mut config = test_connection_config(db_type); + config.id = "agent-conn".to_string(); + state.configs.write().await.insert(config.id.clone(), config); + + let txn_session_id = uuid::Uuid::new_v4().to_string(); + let client_session_id = format!("manual-txn-{}", uuid::Uuid::new_v4()); + // The session pool is not registered in the state map; on commit or + // rollback the detach reports "not found" and the cleanup guard is + // disarmed, matching production flows where the pool was already gone. + let cleanup_guard = state + .workload_session_pool_cleanup_guard("agent-conn", None, &client_session_id) + .await + .expect("agent connections get a session-scoped cleanup guard"); + state.transaction_sessions.write().await.insert( + txn_session_id.clone(), + TransactionSession { + connection: Arc::new(tokio::sync::Mutex::new(TxnConnection::Agent { + client: Arc::new(crate::db::agent_driver::PooledAgentClient::new(client)), + client_session_id, + database: None, + cleanup_guard, + })), + pool_key: "agent-conn".to_string(), + last_activity: std::time::Instant::now(), + busy: false, + connection_id: "agent-conn".to_string(), + database: "ORCL".to_string(), + schema: None, + }, + ); + (state, txn_session_id, dir) + } + + /// Regression: multi-statement scripts under a manual transaction must keep + /// executing sequentially even when pagination options are set — they are + /// ignored instead of failing the script (single-statement pagination stays + /// cursor-based). + #[cfg(unix)] + #[tokio::test] + async fn manual_transaction_multi_statement_script_ignores_pagination_options() { + let (state, txn_session_id, dir) = manual_transaction_test_state(DatabaseType::Oracle).await; + + let results = execute_in_manual_transaction_with_options( + &state, + &txn_session_id, + "SELECT 1 FROM DUAL; SELECT 2 FROM DUAL", + "ORCL", + None, + ManualTransactionExecutionOptions { max_rows: Some(100), page_size: Some(100), ..Default::default() }, + ) + .await + .expect("multi-statement script executes sequentially"); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].result.rows[0][0], serde_json::json!("SELECT 1 FROM DUAL")); + assert_eq!(results[1].result.rows[0][0], serde_json::json!("SELECT 2 FROM DUAL")); + // Both statements ran on the plain execute path — no cursor session was + // opened, so nothing leaks a query cursor on the agent. + assert!(results.iter().all(|result| result.result.session_id.is_none())); + + assert!(rollback_manual_transaction(&state, &txn_session_id).await.is_ok()); + let _ = std::fs::remove_dir_all(dir); + } + #[test] fn agent_close_query_session_params_include_session() { let params = agent_close_query_session_params("session-1"); diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index ae0c03b8b5..cdac5e0bb6 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -463,6 +463,8 @@ pub async fn execute_in_manual_transaction( schema: Option, max_rows: Option, table_data_preview: Option, + page_size: Option, + result_session_id: Option, ) -> Result, ManualTransactionCommandError> { dbx_core::query::execute_in_manual_transaction_with_options( &state, @@ -470,8 +472,12 @@ pub async fn execute_in_manual_transaction( &sql, &database, schema.as_deref(), - max_rows, - table_data_preview.unwrap_or(false), + dbx_core::query::ManualTransactionExecutionOptions { + max_rows, + table_data_preview: table_data_preview.unwrap_or(false), + page_size, + result_session_id, + }, ) .await .map_err(|error| {