From b7ab31bdf597ed6e45836149b008b43cfddccf55 Mon Sep 17 00:00:00 2001 From: efeire Date: Sat, 24 Oct 2026 16:21:30 +0800 Subject: [PATCH] fix(jdbc): bound oversized responses and stream LOB snapshots --- .github/workflows/ci.yml | 3 + apps/desktop/src/components/grid/DataGrid.vue | 5 +- .../src/composables/useDataGridActions.ts | 2 +- .../src/composables/useNavigationTargets.ts | 4 +- .../composables/useSidebarDataOpenRuntime.ts | 2 +- .../dataGrid/dataGridLargeValues.spec.ts | 11 + apps/desktop/src/lib/backend/http.ts | 3 + apps/desktop/src/lib/backend/tauri.ts | 4 + .../src/lib/dataGrid/dataGridLargeValues.ts | 3 +- .../queryCursorOffsetPagination.spec.ts | 74 +++++ .../lib/query/queryCursorOffsetPagination.ts | 27 +- .../queryStore.lobSessionRelease.spec.ts | 76 +++++ ...eryStore.provenReadOnlyStickyState.spec.ts | 3 +- apps/desktop/src/stores/queryStore.ts | 97 +++--- crates/dbx-cli/src/main.rs | 3 + crates/dbx-core/src/ai/agent_tools.rs | 1 + crates/dbx-core/src/connection/mod.rs | 1 + crates/dbx-core/src/data/table_export.rs | 1 + crates/dbx-core/src/data/transfer.rs | 1 + crates/dbx-core/src/query/mod.rs | 220 +++++++++++-- crates/dbx-core/src/query/mongo_ops.rs | 1 + .../dbx-core/src/schema/agent_pg_sequences.rs | 1 + crates/dbx-core/src/schema/kingbase.rs | 1 + crates/dbx-core/src/schema/mod.rs | 15 + .../table_structure_sql/sqlite_rebuild.rs | 1 + crates/dbx-core/tests/jdbc_lob_context.rs | 96 ++++++ crates/dbx-driver-elasticsearch/src/lib.rs | 5 + crates/dbx-driver-mysql/src/mysql.rs | 11 + crates/dbx-driver-postgres/src/postgres.rs | 5 + crates/dbx-driver-sqlserver/src/sqlserver.rs | 12 + .../dbx-drivers/src/db/clickhouse_driver.rs | 2 + .../dbx-drivers/src/db/cloudflare_d1/mod.rs | 1 + crates/dbx-drivers/src/db/couchdb_driver.rs | 2 + crates/dbx-drivers/src/db/dynamodb_driver.rs | 2 + crates/dbx-drivers/src/db/influxdb3_driver.rs | 1 + crates/dbx-drivers/src/db/influxdb_driver.rs | 3 + .../dbx-drivers/src/db/meilisearch_driver.rs | 1 + crates/dbx-drivers/src/db/rqlite_driver.rs | 2 + .../dbx-drivers/src/db/salesforce_driver.rs | 3 + crates/dbx-drivers/src/db/solr_driver.rs | 2 + crates/dbx-drivers/src/db/sqlite.rs | 2 + crates/dbx-drivers/src/db/sqlite_worker.rs | 1 + crates/dbx-drivers/src/db/turso_driver.rs | 4 + crates/dbx-drivers/src/db/vector_driver.rs | 2 + .../src/db/victoriametrics_driver.rs | 2 + crates/dbx-mcp/src/backend.rs | 1 + crates/dbx-mcp/src/http.rs | 1 + crates/dbx-mcp/src/server.rs | 6 + crates/dbx-mcp/src/transaction.rs | 1 + crates/dbx-types/src/types.rs | 17 + crates/dbx-web/src/routes/query.rs | 10 + .../main/java/app/dbx/jdbc/DbxJdbcPlugin.java | 186 +++++++++-- .../java/app/dbx/jdbc/DeferredLobValues.java | 281 ++++++++++++++++ .../app/dbx/jdbc/DeferredLobValuesTest.java | 307 ++++++++++++++++++ .../dbx/jdbc/LargeResponseRecoveryTest.java | 177 ++++++++++ .../app/dbx/jdbc/StatementCancelTest.java | 46 +++ src-tauri/src/commands/query.rs | 4 + 57 files changed, 1665 insertions(+), 91 deletions(-) create mode 100644 apps/desktop/src/lib/query/__tests__/queryCursorOffsetPagination.spec.ts create mode 100644 crates/dbx-core/tests/jdbc_lob_context.rs create mode 100644 plugins/jdbc/src/main/java/app/dbx/jdbc/DeferredLobValues.java create mode 100644 plugins/jdbc/src/test/java/app/dbx/jdbc/DeferredLobValuesTest.java create mode 100644 plugins/jdbc/src/test/java/app/dbx/jdbc/LargeResponseRecoveryTest.java diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ec7e605f09..745f863109 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1176,6 +1176,9 @@ jobs: - name: JDBC plugin package check run: ./plugins/jdbc/package.sh + - name: JDBC response recovery tests + run: bash ./plugins/jdbc/gradlew -p plugins/jdbc --no-daemon test --tests app.dbx.jdbc.LargeResponseRecoveryTest --tests app.dbx.jdbc.StatementCancelTest --tests app.dbx.jdbc.DeferredLobValuesTest + offline-jdbc-release: needs: changes if: needs.changes.outputs.offline_jdbc == 'true' diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index adca21d73a..ce17311b86 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -4502,6 +4502,7 @@ async function refreshSavedRows(request: { dirtyRows: ReadonlyMap /^(?:CLOB|BLOB)$/i.test(column.data_type.trim()))) return false; const planResult = buildDataGridSavedRowRefreshPlan({ context: props.context, @@ -7736,7 +7737,7 @@ async function applyOrderBySearch(_value?: string, structuredOrderByOverride?: s columns: tableMeta.columns.map((column) => column.name), includeDatabaseName: settingsStore.editorSettings.generateSqlIncludeDatabaseName, primaryKeys: tableMeta.primaryKeys, - ...tableDataLargeValuePreviewOptions(resolvedDatabaseType.value, tableMeta.columns, tableMeta.primaryKeys, pageSize.value), + ...tableDataLargeValuePreviewOptions(resolvedDatabaseType.value, tableMeta.columns, tableMeta.primaryKeys, pageSize.value, !!props.connectionId && connectionStore.getConfig(props.connectionId)?.db_type === "jdbc"), orderBy: orderByClause, injectDefaultTimeSeriesWhere: true, limit: pageSize.value, @@ -7779,7 +7780,7 @@ async function applyWhereFilter() { columns: tableMeta.columns.map((column) => column.name), includeDatabaseName: settingsStore.editorSettings.generateSqlIncludeDatabaseName, primaryKeys: tableMeta.primaryKeys, - ...tableDataLargeValuePreviewOptions(resolvedDatabaseType.value, tableMeta.columns, tableMeta.primaryKeys, pageSize.value), + ...tableDataLargeValuePreviewOptions(resolvedDatabaseType.value, tableMeta.columns, tableMeta.primaryKeys, pageSize.value, !!props.connectionId && connectionStore.getConfig(props.connectionId)?.db_type === "jdbc"), orderBy: effectiveOrderBy(), limit: pageSize.value, injectDefaultTimeSeriesWhere: true, diff --git a/apps/desktop/src/composables/useDataGridActions.ts b/apps/desktop/src/composables/useDataGridActions.ts index 102652b484..7684c9c6c4 100644 --- a/apps/desktop/src/composables/useDataGridActions.ts +++ b/apps/desktop/src/composables/useDataGridActions.ts @@ -144,7 +144,7 @@ export function useDataGridActions(activeTab: ComputedRef) catalog: tableMeta?.catalog, columns: realColumns?.map((column) => column.name), primaryKeys, - ...tableDataLargeValuePreviewOptions(effectiveDbType, realColumns ?? [], primaryKeys, previewPageSize ?? limit), + ...tableDataLargeValuePreviewOptions(effectiveDbType, realColumns ?? [], primaryKeys, previewPageSize ?? limit, config?.db_type === "jdbc"), includeDatabaseName: settingsStore.editorSettings.generateSqlIncludeDatabaseName, includeRowId: useRowId, limit, diff --git a/apps/desktop/src/composables/useNavigationTargets.ts b/apps/desktop/src/composables/useNavigationTargets.ts index e57386f7a5..b4fc17a595 100644 --- a/apps/desktop/src/composables/useNavigationTargets.ts +++ b/apps/desktop/src/composables/useNavigationTargets.ts @@ -232,7 +232,7 @@ async function openTableTarget(target: NavigationTarget, options: NavigationOpen columns: eagerColumns.map((column) => column.name), primaryKeys: eagerPrimaryKeys, orderBy: defaultSort.orderBy, - ...tableDataLargeValuePreviewOptions(effectiveDbType, eagerColumns, eagerPrimaryKeys, pageLimit), + ...tableDataLargeValuePreviewOptions(effectiveDbType, eagerColumns, eagerPrimaryKeys, pageLimit, config.db_type === "jdbc"), whereInput: target.whereInput, limit: pageLimit, }); @@ -333,7 +333,7 @@ async function openTableTarget(target: NavigationTarget, options: NavigationOpen whereInput: target.whereInput, primaryKeys, columns: columns.map((column) => column.name), - ...tableDataLargeValuePreviewOptions(effectiveDbType, columns, primaryKeys, pageLimit), + ...tableDataLargeValuePreviewOptions(effectiveDbType, columns, primaryKeys, pageLimit, config.db_type === "jdbc"), includeRowId: true, orderBy: defaultSort.orderBy, limit: pageLimit, diff --git a/apps/desktop/src/composables/useSidebarDataOpenRuntime.ts b/apps/desktop/src/composables/useSidebarDataOpenRuntime.ts index 1a039629cd..5aed4dc2ef 100644 --- a/apps/desktop/src/composables/useSidebarDataOpenRuntime.ts +++ b/apps/desktop/src/composables/useSidebarDataOpenRuntime.ts @@ -433,7 +433,7 @@ export function useSidebarDataOpenRuntime() { catalog: node.catalog, columns: columns.map((column) => column.name), primaryKeys, - ...tableDataLargeValuePreviewOptions(effectiveDbType, columns, primaryKeys, limit), + ...tableDataLargeValuePreviewOptions(effectiveDbType, columns, primaryKeys, limit, config?.db_type === "jdbc"), includeDatabaseName: settingsStore.editorSettings.generateSqlIncludeDatabaseName, limit, includeRowId: withRowId, diff --git a/apps/desktop/src/lib/__tests__/dataGrid/dataGridLargeValues.spec.ts b/apps/desktop/src/lib/__tests__/dataGrid/dataGridLargeValues.spec.ts index 6b811b4114..8421447320 100644 --- a/apps/desktop/src/lib/__tests__/dataGrid/dataGridLargeValues.spec.ts +++ b/apps/desktop/src/lib/__tests__/dataGrid/dataGridLargeValues.spec.ts @@ -81,6 +81,17 @@ describe("data grid large-value metadata", () => { expect(tableDataLargeValuePreviewOptions("oracle", [column("ID", "NUMBER", true), column("DOC", "XMLTYPE")], ["ID"], 100)).toEqual({}); }); + it("keeps JDBC CLOB/BLOB locators intact while retaining the NCLOB SQL preview", () => { + const columns = [column("ID", "NUMBER", true), column("BODY", "CLOB"), column("DATA", "BLOB")]; + expect(tableDataLargeValuePreviewOptions("oracle", columns, ["ID"], 1000, true)).toEqual({}); + expect(tableDataLargeValuePreviewOptions("oracle", [...columns, column("NATIONAL", "NCLOB")], ["ID"], 1000, true)).toEqual({ + columnTypes: ["NUMBER", "", "", "NCLOB"], + largeValuePreviewSize: 4194, + }); + expect(columns.map((item) => item.data_type)).toEqual(["NUMBER", "CLOB", "BLOB"]); + expect(tableDataLargeValuePreviewOptions("oracle", columns, ["ID"], 1000)).toEqual({ columnTypes: ["NUMBER", "CLOB", "BLOB"], largeValuePreviewSize: 2097 }); + }); + it("disables marker parsing when a real column uses the reserved preview prefix", () => { const columns = [column("id", "bigint", true), column("__dbx_large_value_bytes_t_0", "text")]; diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index 97d5a34b0d..735e96da2e 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -1612,6 +1612,7 @@ export async function executeMulti( maxResultBytes?: number; resultKeyColumns?: string[]; tableDataPreview?: boolean; + deferLobs?: boolean; resultSessionId?: string; clientSessionId?: string; timeoutSecs?: number; @@ -1668,6 +1669,7 @@ export async function executeMultiWithProgress( maxResultBytes?: number; resultKeyColumns?: string[]; tableDataPreview?: boolean; + deferLobs?: boolean; resultSessionId?: string; clientSessionId?: string; timeoutSecs?: number; @@ -1839,6 +1841,7 @@ export async function executeInManualTransaction( _executionId?: string, _timeoutSecs?: number, _boundStatements?: BlobBoundStatement[], + _deferLobs?: boolean, ): 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 c587cb3d03..98498f0fa7 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -2065,6 +2065,7 @@ export async function executeMulti( maxResultBytes?: number; resultKeyColumns?: string[]; tableDataPreview?: boolean; + deferLobs?: boolean; resultSessionId?: string; clientSessionId?: string; timeoutSecs?: number; @@ -2141,6 +2142,7 @@ export async function executeMultiWithProgress( maxResultBytes?: number; resultKeyColumns?: string[]; tableDataPreview?: boolean; + deferLobs?: boolean; resultSessionId?: string; clientSessionId?: string; timeoutSecs?: number; @@ -2282,6 +2284,7 @@ export async function executeInManualTransaction( executionId?: string, timeoutSecs?: number, boundStatements?: BlobBoundStatement[], + deferLobs?: boolean, ): Promise { return invokeBackend("execute_in_manual_transaction", { txnSessionId, @@ -2296,6 +2299,7 @@ export async function executeInManualTransaction( executionId, timeoutSecs, boundStatements, + deferLobs, }); } diff --git a/apps/desktop/src/lib/dataGrid/dataGridLargeValues.ts b/apps/desktop/src/lib/dataGrid/dataGridLargeValues.ts index 5d75386f43..7d7db48bd6 100644 --- a/apps/desktop/src/lib/dataGrid/dataGridLargeValues.ts +++ b/apps/desktop/src/lib/dataGrid/dataGridLargeValues.ts @@ -220,7 +220,8 @@ export function tableDataPreviewRowBudget(options: { limit: number; offset: numb return remainingRows; } -export function tableDataLargeValuePreviewOptions(databaseType: DatabaseType | undefined, columns: readonly ColumnInfo[], primaryKeys: readonly string[], pageSize?: number): { columnTypes: string[]; largeValuePreviewSize: number } | Record { +export function tableDataLargeValuePreviewOptions(databaseType: DatabaseType | undefined, columns: readonly ColumnInfo[], primaryKeys: readonly string[], pageSize?: number, jdbcDeferredLobs = false): { columnTypes: string[]; largeValuePreviewSize: number } | Record { + if (databaseType === "oracle" && jdbcDeferredLobs) columns = columns.map((column) => (/^(?:CLOB|BLOB)$/i.test(column.data_type.trim()) ? { ...column, data_type: "" } : column)); if (!canUseTableDataLargeValuePreview(databaseType, columns, primaryKeys)) return {}; const keyColumns = new Set(primaryKeys.map((column) => column.toLocaleLowerCase())); const mysqlBudget = databaseType === "mysql" ? mysqlPreviewBudget(columns, keyColumns, pageSize) : null; diff --git a/apps/desktop/src/lib/query/__tests__/queryCursorOffsetPagination.spec.ts b/apps/desktop/src/lib/query/__tests__/queryCursorOffsetPagination.spec.ts new file mode 100644 index 0000000000..64204470fe --- /dev/null +++ b/apps/desktop/src/lib/query/__tests__/queryCursorOffsetPagination.spec.ts @@ -0,0 +1,74 @@ +import { describe, expect, it, vi } from "vitest"; +import type { QueryResult } from "@/types/database"; +import { executeCursorOffsetPage } from "../queryCursorOffsetPagination"; + +const page = (ref: string, more = true): QueryResult => ({ + columns: ["payload"], + rows: [[ref]], + affected_rows: 0, + execution_time_ms: 1, + session_id: more ? "cursor" : undefined, + has_more: more, + large_value_cells: [{ row_index: 0, column_index: 0, original_bytes: 67108864, value_ref: ref }], + large_value_refs: [ref], +}); + +describe("cursor offset original-value ownership", () => { + it("releases skipped pages but retains the requested page", async () => { + const skipped = page("skipped"); + const requested = page("requested", false); + const releaseResults = vi.fn(async () => {}); + const result = await executeCursorOffsetPage({ + offset: 1, + limit: 1, + fetchPage: vi.fn().mockResolvedValueOnce([skipped]).mockResolvedValueOnce([requested]), + closeSession: vi.fn(), + isCancelled: () => false, + releaseResults, + }); + expect(releaseResults.mock.calls).toEqual([[[skipped]]]); + expect(result[0]?.large_value_refs).toEqual(["requested"]); + }); + + it("releases a late page before closing its cancelled cursor", async () => { + const late = page("late"); + let cancelled = false; + const operations: string[] = []; + await expect( + executeCursorOffsetPage({ + offset: 1, + limit: 1, + fetchPage: async () => { + cancelled = true; + return [late]; + }, + isCancelled: () => cancelled, + releaseResults: async (results) => { + expect(results).toEqual([late]); + operations.push("release"); + }, + closeSession: async () => { + operations.push("close"); + }, + }), + ).rejects.toThrow("Query canceled"); + expect(operations).toEqual(["release", "close"]); + }); + + it("does not attach discarded references to an exhausted empty result", async () => { + const exhausted = page("discarded", false); + const releaseResults = vi.fn(async () => {}); + const result = await executeCursorOffsetPage({ + offset: 2, + limit: 1, + fetchPage: async () => [exhausted], + closeSession: vi.fn(), + isCancelled: () => false, + releaseResults, + }); + expect(result[0]?.rows).toEqual([]); + expect(result[0]?.large_value_cells).toEqual([]); + expect(result[0]?.large_value_refs).toEqual([]); + expect(releaseResults).toHaveBeenCalledWith([exhausted]); + }); +}); diff --git a/apps/desktop/src/lib/query/queryCursorOffsetPagination.ts b/apps/desktop/src/lib/query/queryCursorOffsetPagination.ts index 5a07c930c3..6251310020 100644 --- a/apps/desktop/src/lib/query/queryCursorOffsetPagination.ts +++ b/apps/desktop/src/lib/query/queryCursorOffsetPagination.ts @@ -1,7 +1,14 @@ import type { QueryResult } from "@/types/database"; /** Reach an offset using a forward-only cursor, keeping only the requested page. */ -export async function executeCursorOffsetPage(options: { offset: number; limit: number; fetchPage: (pageSize: number, sessionId?: string) => Promise; closeSession: (sessionId: string) => Promise; isCancelled: () => boolean }): Promise { +export async function executeCursorOffsetPage(options: { + offset: number; + limit: number; + fetchPage: (pageSize: number, sessionId?: string) => Promise; + closeSession: (sessionId: string) => Promise; + isCancelled: () => boolean; + releaseResults?: (results: QueryResult[]) => Promise; +}): Promise { let sessionId: string | undefined; let skipped = 0; let pageCount = 0; @@ -9,6 +16,12 @@ export async function executeCursorOffsetPage(options: { offset: number; limit: let completeTimings = true; let provenReadOnly: true | undefined; const timings: Record = {}; + let pendingResults: QueryResult[] = []; + const discardPending = async () => { + const discarded = pendingResults; + pendingResults = []; + await options.releaseResults?.(discarded); + }; try { while (true) { if (options.isCancelled()) throw new Error("Query canceled"); @@ -16,6 +29,7 @@ export async function executeCursorOffsetPage(options: { offset: number; limit: // offset instead of overshooting it with that shortened page size. const pageSize = skipped < options.offset ? Math.min(options.limit, options.offset - skipped) : options.limit; const results = await options.fetchPage(pageSize, sessionId); + pendingResults = results; const page = results[0]; if (!page) { if (sessionId) throw new Error("Result session ended before the requested page offset"); @@ -42,11 +56,20 @@ export async function executeCursorOffsetPage(options: { offset: number; limit: if (skipped === options.offset && (page.rows.length > 0 || !page.has_more)) return [result]; if (page.rows.length > pageSize) throw new Error("Result page exceeded the requested cursor page size"); skipped += page.rows.length; - if (!page.has_more) return [{ ...result, rows: [] }]; + if (!page.has_more) { + await discardPending(); + return [{ ...result, rows: [], large_value_cells: [], large_value_refs: [] }]; + } + await discardPending(); // Cassandra can return empty intermediate pages with a valid next cursor. // Continue advancing that cursor until rows arrive or it is exhausted. } } catch (error) { + try { + await discardPending(); + } catch (releaseError) { + console.warn("[DBX][offset-replay-large-value-release:error]", releaseError); + } if (sessionId) { try { await options.closeSession(sessionId); diff --git a/apps/desktop/src/stores/__tests__/queryStore.lobSessionRelease.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.lobSessionRelease.spec.ts index b8bbeb6a82..8ba8928388 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.lobSessionRelease.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.lobSessionRelease.spec.ts @@ -9,6 +9,7 @@ const api = vi.hoisted(() => ({ cancelQuery: vi.fn(), disconnectDb: vi.fn(), releaseLargeValue: vi.fn(), + executeInManualTransaction: vi.fn(), })); vi.mock("@/lib/backend/api", () => ({ @@ -41,6 +42,81 @@ describe("query Store LOB session release through public close and disconnect", api.releaseLargeValue.mockResolvedValue(true); }); + it("opts JDBC grid queries into deferred values and binds their original client session", async () => { + const { useQueryStore } = await import("@/stores/queryStore"); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const { useSettingsStore } = await import("@/stores/settingsStore"); + useSettingsStore().editorSettings.autoCalculateTotalRows = false; + const connections = useConnectionStore(); + connections.connections = [{ id: "jdbc-lob", name: "JDBC", db_type: "jdbc", database: "APP", read_only: false, connection_string: "jdbc:oracle:thin:@fixture", jdbc_driver_class: "oracle.jdbc.OracleDriver" } as ConnectionConfig]; + connections.connectedIds.add("jdbc-lob"); + const store = useQueryStore(); + const id = store.createTab("jdbc-lob", "APP", "LOB", "query", "APP"); + store.tabs.find((tab) => tab.id === id)!.autoCommit = true; + api.executeMulti.mockResolvedValueOnce([{ ...result("jdbc-page"), large_value_cells: [{ row_index: 0, column_index: 0, original_bytes: 67108864, value_ref: "jdbc-original" }] }]); + await store.executeTabSql(id, "SELECT PAYLOAD FROM FIXTURE"); + const options = api.executeMulti.mock.calls[0][5]; + expect(options.deferLobs).toBe(true); + expect(options.tableDataPreview).not.toBe(true); + expect(store.tabs.find((tab) => tab.id === id)?.result?.large_value_context).toEqual(expect.objectContaining({ connectionId: "jdbc-lob", database: "APP", clientSessionId: options.clientSessionId })); + await store.flushPendingPersist(); + }); + + it("keeps SQL marker parsing enabled for JDBC manual table data with mixed LOB types", async () => { + const { useQueryStore } = await import("@/stores/queryStore"); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const { useSettingsStore } = await import("@/stores/settingsStore"); + useSettingsStore().editorSettings.autoCalculateTotalRows = false; + const connections = useConnectionStore(); + connections.connections = [{ id: "jdbc-mixed", name: "JDBC", db_type: "jdbc", database: "APP", read_only: false, connection_string: "jdbc:oracle:thin:@fixture", jdbc_driver_class: "oracle.jdbc.OracleDriver" } as ConnectionConfig]; + connections.connectedIds.add("jdbc-mixed"); + const store = useQueryStore(); + const id = store.createTab("jdbc-mixed", "APP", "FIXTURE", "data", "APP"); + const tab = store.tabs.find((entry) => entry.id === id)!; + store.setTableMeta(id, { tableName: "FIXTURE", schema: "APP", primaryKeys: ["ID"], columns: ["NUMBER", "NCLOB", "CLOB"].map((data_type, index) => ({ name: ["ID", "N", "C"][index]!, data_type, is_nullable: false, column_default: null, is_primary_key: index === 0, extra: null })) }); + tab.autoCommit = false; + tab.txnSessionId = "mixed-transaction"; + api.executeInManualTransaction.mockResolvedValueOnce([{ columns: ["ID", "N", "C"], rows: [[1, "national preview", "clob preview"]], affected_rows: 0, execution_time_ms: 1 }]); + await store.executeTabSql(id, "SELECT ID,N,marker,C FROM FIXTURE"); + const args = api.executeInManualTransaction.mock.calls[0]; + expect(args[0]).toBe("mixed-transaction"); + expect(args[5]).toBe(true); + expect(args[12]).toBe(true); + await store.flushPendingPersist(); + }); + + it("releases a late manual result using its executed transaction, not the tab's replacement", async () => { + const { useQueryStore } = await import("@/stores/queryStore"); + const { useConnectionStore } = await import("@/stores/connectionStore"); + const { useSettingsStore } = await import("@/stores/settingsStore"); + useSettingsStore().editorSettings.autoCalculateTotalRows = false; + const connections = useConnectionStore(); + connections.connections = [{ id: "txn-lob", name: "OB", db_type: "oceanbase-oracle", database: "SYS", read_only: false } as ConnectionConfig]; + connections.connectedIds.add("txn-lob"); + const store = useQueryStore(); + const id = store.createTab("txn-lob", "SYS", "LOB", "query", "APP"); + const tab = store.tabs.find((tab) => tab.id === id)!; + tab.autoCommit = false; + tab.txnSessionId = "executed-transaction"; + let finish!: (results: QueryResult[]) => void; + api.executeInManualTransaction.mockImplementationOnce( + () => + new Promise((resolve) => { + finish = resolve; + }), + ); + const executing = store.executeTabSql(id, "SELECT PAYLOAD FROM FIXTURE"); + await vi.waitFor(() => expect(api.executeInManualTransaction).toHaveBeenCalledTimes(1)); + expect(api.executeInManualTransaction.mock.calls[0][0]).toBe("executed-transaction"); + tab.txnSessionId = "replacement-transaction"; + tab.executionId = "replacement-execution"; + finish([{ ...result("late-cursor"), large_value_cells: [{ row_index: 0, column_index: 0, original_bytes: 67108864, value_ref: "old-reference" }] }]); + await executing; + await vi.waitFor(() => expect(api.releaseLargeValue).toHaveBeenCalledWith(expect.objectContaining({ valueRef: "old-reference", txnSessionId: "executed-transaction" }))); + expect(api.releaseLargeValue.mock.calls.every(([request]) => request.txnSessionId !== "replacement-transaction")).toBe(true); + await store.flushPendingPersist(); + }); + it("preserves locator ownership from hydrated earlier pages when appending", async () => { const { appendQueryResultSegment } = await import("@/stores/queryStore"); const context = { connectionId: "ob-page", database: "SYS", clientSessionId: "original-session" }; diff --git a/apps/desktop/src/stores/__tests__/queryStore.provenReadOnlyStickyState.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.provenReadOnlyStickyState.spec.ts index 1831b409c0..270698af51 100644 --- a/apps/desktop/src/stores/__tests__/queryStore.provenReadOnlyStickyState.spec.ts +++ b/apps/desktop/src/stores/__tests__/queryStore.provenReadOnlyStickyState.spec.ts @@ -543,7 +543,8 @@ describe("queryStore Oracle/OceanBase manual-transaction sticky state (behavior // No classificationSql is sent and no sticky state exists for non-sticky // dialects: commit/rollback follow the legacy always-visible rule. const args = mocks.executeInManualTransaction.mock.calls[0]!; - expect(args.at(-1)).toBeUndefined(); + expect(args[8]).toBeUndefined(); + expect(args[12]).toBe(true); expect(tab.txnPossiblyDirty).toBeUndefined(); expect(tab.txnSessionId).toBe("txn-jdbc"); }); diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index e0468c5f61..b08476b21e 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -2308,12 +2308,13 @@ export const useQueryStore = defineStore("query", () => { async function releaseResultLargeValues(...results: Array) { const released = new Set(); for (const result of results) { - const context = result?.large_value_context; - if (!context) continue; + if (!result?.large_value_context) continue; + const context = result.large_value_context; const refs = new Set([...(result.large_value_refs ?? []), ...(result.large_value_cells ?? []).flatMap((cell) => (cell.value_ref ? [cell.value_ref] : []))]); for (const valueRef of refs) { - if (released.has(valueRef)) continue; - released.add(valueRef); + const key = resultLargeValueKey(result, valueRef); + if (released.has(key)) continue; + released.add(key); try { await api.releaseLargeValue({ ...context, valueRef }); } catch (error) { @@ -5620,7 +5621,7 @@ export const useQueryStore = defineStore("query", () => { catalog: tableMeta.catalog, columns: tableMeta.columns.map((column) => column.name), primaryKeys, - ...tableDataLargeValuePreviewOptions(effectiveDbType, tableMeta.columns, primaryKeys, limit), + ...tableDataLargeValuePreviewOptions(effectiveDbType, tableMeta.columns, primaryKeys, limit, conn?.db_type === "jdbc"), includeDatabaseName: settingsStore.editorSettings.generateSqlIncludeDatabaseName, includeRowId, whereInput: tab.whereInput, @@ -8958,6 +8959,21 @@ export const useQueryStore = defineStore("query", () => { } const sourceLabelDatabase = targetDatabase || conn?.database; const executionClientSessionId = options?.pagination?.clientSessionId ?? (tab.mode === "query" || tab.mode === "data" ? tabClientSessionId(tab) : undefined); + const deferLobs = conn?.db_type === "jdbc" && (tab.mode === "query" || tab.mode === "data"); + const bindLargeValueContext = (results: QueryResult[], txnSessionId?: string): QueryResult[] => { + for (const result of results) { + if (!result.large_value_cells?.some((cell) => cell.value_ref)) continue; + result.large_value_refs = result.large_value_cells.flatMap((cell) => (cell.value_ref ? [cell.value_ref] : [])); + result.large_value_context = { + connectionId: executionConnectionId, + database: executionDatabase, + clientSessionId: executionClientSessionId, + txnSessionId, + catalog: executionCatalog, + }; + } + return results; + }; if (tab.mode === "query" && tab.objectSource && !usesExternalExecutionTarget && effectiveDbType && tabs.value.includes(tab) && executionConnectionId === tab.connectionId && executionDatabase === tab.database) { await prepareObjectSourceRoutineRun(tab, effectiveDbType, executionDatabase, sqlToExecute, sqlStatementParameterOptions, traceId); } @@ -8967,14 +8983,15 @@ export const useQueryStore = defineStore("query", () => { return false; } + const dataTabMeta = tab.mode === "data" ? tableMetaForDataTab(tab) : undefined; + const useTableDataPreview = canUseTableDataLargeValuePreview(effectiveDbType, dataTabMeta?.columns ?? [], dataTabMeta?.primaryKeys ?? []); const executeWithoutManualTransaction = (): Promise => { queryExecutionLog("info", "execute-multi:start", { traceId, elapsed: elapsed() }); // Query and data tabs use a tab-scoped pool so repeated executions keep // connection-local state and avoid MySQL pool resets on every refresh. - const dataTabMeta = tab.mode === "data" ? tableMetaForDataTab(tab) : undefined; - const useTableDataPreview = canUseTableDataLargeValuePreview(effectiveDbType, dataTabMeta?.columns ?? [], dataTabMeta?.primaryKeys ?? []); const useJdbcDriverRowOffset = tab.mode === "data" && jdbcConnectionUsesDriverRowOffset(conn, effectiveDbType); const executionOptions = { + ...(deferLobs ? { deferLobs: true } : {}), ...(typeof pageLimit === "number" ? useAgentResultSession ? { @@ -9022,26 +9039,29 @@ export const useQueryStore = defineStore("query", () => { return executeCursorOffsetPage({ offset: pageOffset, limit: pageLimit, - fetchPage: (pageSize, sessionId) => api.executeMulti(executionConnectionId, executionDatabase, sqlToExecute, executionSchema, executionId, { ...executionOptions, pageSize, resultSessionId: sessionId }), + fetchPage: (pageSize, sessionId) => api.executeMulti(executionConnectionId, executionDatabase, sqlToExecute, executionSchema, executionId, { ...executionOptions, pageSize, resultSessionId: sessionId }).then((results) => bindLargeValueContext(results)), closeSession: (sessionId) => api.closeQuerySession(executionConnectionId, executionDatabase, sessionId, executionClientSessionId, executionCatalog), + releaseResults: (results) => releaseResultLargeValues(...results), isCancelled: isReplayCancelled, }); } - return tab.batchSqlExecution && tab.batchSqlExecution.total > 1 - ? api.executeMultiWithProgress( - executionConnectionId, - executionDatabase, - sqlToExecute, - (progress) => { - const current = findExecutionTab(id); - if (current?.executionId === executionId) { - applyBatchSqlProgress(current, progress, continueOnBatchError, batchResume?.startStatementIndex ?? 0); - } - }, - executionSchema, - { ...executionOptions, executionId }, - ) - : api.executeMulti(executionConnectionId, executionDatabase, sqlToExecute, executionSchema, executionId, executionOptions); + const request = + tab.batchSqlExecution && tab.batchSqlExecution.total > 1 + ? api.executeMultiWithProgress( + executionConnectionId, + executionDatabase, + sqlToExecute, + (progress) => { + const current = findExecutionTab(id); + if (current?.executionId === executionId) { + applyBatchSqlProgress(current, progress, continueOnBatchError, batchResume?.startStatementIndex ?? 0); + } + }, + executionSchema, + { ...executionOptions, executionId }, + ) + : api.executeMulti(executionConnectionId, executionDatabase, sqlToExecute, executionSchema, executionId, executionOptions); + return request.then((results) => bindLargeValueContext(results)); }; let executionPromise: Promise; @@ -9088,12 +9108,13 @@ export const useQueryStore = defineStore("query", () => { executionDatabase, executionSchema, useAgentResultSession || batchPagination ? agentProtocolQueryResultMaxRows(queryResultMaxRows) : (pageProbeLimit ?? pageLimit ?? agentProtocolQueryResultMaxRows(queryResultMaxRows)), - useLargeValuePreview, + deferLobs ? useTableDataPreview : useLargeValuePreview, useAgentResultSession ? pageSize : undefined, useAgentResultSession ? resultSessionId : undefined, resultSessionId ? undefined : classificationSql, ] as const; - const results = await (effectiveDbType === "sqlserver" ? api.executeInManualTransaction(...args, executionId, frontendTimeoutSecs) : api.executeInManualTransaction(...args)); + const results = await (effectiveDbType === "sqlserver" ? api.executeInManualTransaction(...args, executionId, frontendTimeoutSecs) : deferLobs ? api.executeInManualTransaction(...args, executionId, frontendTimeoutSecs, undefined, true) : api.executeInManualTransaction(...args)); + bindLargeValueContext(results, sessionId); for (const result of results) { if (result.session_id) result.manual_transaction_session_id = sessionId; } @@ -9110,6 +9131,7 @@ export const useQueryStore = defineStore("query", () => { limit: pageLimit, fetchPage, closeSession: (cursorId) => api.closeQuerySession(executionConnectionId, executionDatabase, cursorId, undefined, executionCatalog, sessionId), + releaseResults: (results) => releaseResultLargeValues(...results), isCancelled: () => isReplayCancelled() || tab.txnSessionId !== sessionId, }); } @@ -9153,11 +9175,17 @@ export const useQueryStore = defineStore("query", () => { } else { executionPromise = executeWithoutManualTransaction(); } - const responseResults = await withFrontendQueryTimeout(executionPromise, frontendTimeoutSecs, t("editor.queryTimeoutError", { seconds: frontendTimeoutSecs }), () => { + let responseWasAbandoned = false; + const observedExecution = executionPromise.then((results) => { + if (responseWasAbandoned) void releaseResultLargeValues(...results); + return results; + }); + const responseResults = await withFrontendQueryTimeout(observedExecution, frontendTimeoutSecs, t("editor.queryTimeoutError", { seconds: frontendTimeoutSecs }), () => { + responseWasAbandoned = true; void api.cancelQuery(executionId).catch((error) => queryExecutionLog("warn", "frontend-timeout:cancel-failed", { traceId, error })); }); for (const result of responseResults) { - if (effectiveDbType === "oceanbase-oracle" && !useLargeValuePreview && result.column_types?.some((type) => /^(N?CLOB|BLOB)$/i.test(type.trim()))) { + if (effectiveDbType === "oceanbase-oracle" && !useLargeValuePreview && !result.large_value_cells?.some((cell) => cell.value_ref) && result.column_types?.some((type) => /^(N?CLOB|BLOB)$/i.test(type.trim()))) { result.messages = [ ...(result.messages ?? []), { @@ -9167,18 +9195,11 @@ export const useQueryStore = defineStore("query", () => { }, ]; } - if (result.large_value_cells?.some((cell) => cell.value_ref)) { - result.large_value_refs = result.large_value_cells.flatMap((cell) => (cell.value_ref ? [cell.value_ref] : [])); - result.large_value_context = { - connectionId: executionConnectionId, - database: executionDatabase, - clientSessionId: executionClientSessionId, - txnSessionId: tab.autoCommit === false ? tab.txnSessionId : undefined, - catalog: executionCatalog, - }; - } } - if (findExecutionTab(id) !== tab || tab.executionId !== executionId || manualTransactionTargetEpoch(tab) !== executionTargetEpoch) return false; + if (findExecutionTab(id) !== tab || tab.executionId !== executionId || manualTransactionTargetEpoch(tab) !== executionTargetEpoch) { + void releaseResultLargeValues(...responseResults); + return false; + } // A single result has an unambiguous request boundary. This includes fetch and // transport, but excludes SQL preparation and the grid's later render work. if (clientRequestStartedAt !== undefined && responseResults.length === 1 && !responseResults[0]?.execution_error) { diff --git a/crates/dbx-cli/src/main.rs b/crates/dbx-cli/src/main.rs index 0922b58137..1b67b51e64 100644 --- a/crates/dbx-cli/src/main.rs +++ b/crates/dbx-cli/src/main.rs @@ -1116,6 +1116,7 @@ mod tests { Ok(QueryResult { columns: Vec::new(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -1218,6 +1219,7 @@ mod tests { let mut result = QueryResult { columns: Vec::new(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -1278,6 +1280,7 @@ mod tests { let result = QueryResult { columns: vec!["amount".to_string()], column_types: vec!["NUMBER(38,12)".to_string()], + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), diff --git a/crates/dbx-core/src/ai/agent_tools.rs b/crates/dbx-core/src/ai/agent_tools.rs index a6b4067969..2ec1799040 100644 --- a/crates/dbx-core/src/ai/agent_tools.rs +++ b/crates/dbx-core/src/ai/agent_tools.rs @@ -2576,6 +2576,7 @@ for line in sys.stdin: QueryResult { columns: columns.into_iter().map(str::to_string).collect(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), diff --git a/crates/dbx-core/src/connection/mod.rs b/crates/dbx-core/src/connection/mod.rs index 45e982c19d..558a618695 100644 --- a/crates/dbx-core/src/connection/mod.rs +++ b/crates/dbx-core/src/connection/mod.rs @@ -8510,6 +8510,7 @@ mod tests { let result = crate::types::QueryResult { columns: vec!["datcompatibility".to_string()], column_types: vec!["text".to_string()], + deferred_large_value_cells: Vec::new(), column_sortables: vec![true], spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-core/src/data/table_export.rs b/crates/dbx-core/src/data/table_export.rs index 802a502ada..14973f74db 100644 --- a/crates/dbx-core/src/data/table_export.rs +++ b/crates/dbx-core/src/data/table_export.rs @@ -731,6 +731,7 @@ async fn fetch_table_export_batch( return Ok(QueryResult { columns: col_names.to_vec(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-core/src/data/transfer.rs b/crates/dbx-core/src/data/transfer.rs index 7493a72a92..c43e27b6a1 100644 --- a/crates/dbx-core/src/data/transfer.rs +++ b/crates/dbx-core/src/data/transfer.rs @@ -15210,6 +15210,7 @@ CREATE TABLE "Other"."prefix""Source"."NAME" ("ID" INT);"#; db::QueryResult { columns: Vec::new(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-core/src/query/mod.rs b/crates/dbx-core/src/query/mod.rs index bdae602a25..87d910f67a 100644 --- a/crates/dbx-core/src/query/mod.rs +++ b/crates/dbx-core/src/query/mod.rs @@ -375,10 +375,11 @@ impl ExecuteMultiResult { } } - fn success_with_index(result: db::QueryResult, statement_index: usize) -> Self { + fn success_with_index(mut result: db::QueryResult, statement_index: usize) -> Self { + let large_value_cells = std::mem::take(&mut result.deferred_large_value_cells); Self { result, - large_value_cells: Vec::new(), + large_value_cells, execution_error: false, statement_index: Some(statement_index), error: None, @@ -397,6 +398,8 @@ impl ExecuteMultiResult { mut large_value_cells: Vec, table_data_preview: bool, ) -> Self { + large_value_cells = + merge_large_value_cells(large_value_cells, std::mem::take(&mut result.deferred_large_value_cells)); let server_cells = if table_data_preview { remap_large_value_cells_around_server_markers(&result, &mut large_value_cells); extract_server_large_value_markers(&mut result) @@ -427,8 +430,14 @@ impl ExecuteMultiResult { } fn success_with_optional_server_large_values(mut result: db::QueryResult, table_data_preview: bool) -> Self { - let large_value_cells = - if table_data_preview { extract_server_large_value_markers(&mut result) } else { Vec::new() }; + let mut deferred_cells = std::mem::take(&mut result.deferred_large_value_cells); + let large_value_cells = if table_data_preview { + remap_large_value_cells_around_server_markers(&result, &mut deferred_cells); + extract_server_large_value_markers(&mut result) + } else { + Vec::new() + }; + let large_value_cells = merge_large_value_cells(large_value_cells, deferred_cells); Self { result, large_value_cells, @@ -751,10 +760,11 @@ fn merge_large_value_cells( } impl From for ExecuteMultiResult { - fn from(result: db::QueryResult) -> Self { + fn from(mut result: db::QueryResult) -> Self { + let large_value_cells = std::mem::take(&mut result.deferred_large_value_cells); Self { result, - large_value_cells: Vec::new(), + large_value_cells, execution_error: false, statement_index: None, error: None, @@ -1209,6 +1219,7 @@ pub enum QueryExecutionMode { #[derive(Clone, Debug, Default)] pub struct QueryExecutionOptions { + pub defer_lobs: bool, pub max_rows: Option, pub fetch_size: Option, pub page_size: Option, @@ -2774,6 +2785,9 @@ pub(crate) fn external_driver_query_params( "schema": schema, "maxRows": options.max_rows.unwrap_or(MAX_ROWS), }); + if options.defer_lobs { + params["deferLobs"] = serde_json::json!(true); + } if let Some(fetch_size) = options.fetch_size { params["fetchSize"] = serde_json::json!(fetch_size); } @@ -3273,13 +3287,25 @@ async fn request_large_value_with_cancel( Arc::clone(&session.connection) }; let connection = wait_for_value_opt(cancel_token.clone(), timeout, connection.lock()).await?; - let TxnConnection::Agent { client, .. } = &*connection else { - return Err("LOB snapshot chunks require an agent transaction".to_string()); + let result = match &*connection { + TxnConnection::Agent { client, .. } => { + wait_for_value_opt(cancel_token.clone(), timeout, client.lock()) + .await? + .call_with_timeout_and_cancel(method, params, timeout, cancel_token.clone()) + .await + } + TxnConnection::ExternalDriver { session, config, .. } => { + let mut params = params; + params["connection"] = serde_json::json!(config.as_ref()); + wait_for_value_opt( + cancel_token.clone(), + timeout, + session.invoke_with_timeout::(method, params, timeout), + ) + .await? + } + _ => return Err("LOB snapshot chunks require a supported transaction".to_string()), }; - let result = wait_for_value_opt(cancel_token.clone(), timeout, client.lock()) - .await? - .call_with_timeout_and_cancel(method, params, timeout, cancel_token.clone()) - .await; if !wait_for_value_opt(cancel_token, timeout, state.transaction_sessions.read()).await?.contains_key(txn_id) { return Err("LOB transaction context expired".to_string()); } @@ -3298,11 +3324,23 @@ async fn request_large_value_with_cancel( .await? .ok_or_else(|| "LOB context has expired; rerun the query".to_string())?; let _activity = state.pool_activity_touch(pool_key.as_str()); - let PoolKind::Agent(client) = handle else { - return Err("LOB snapshot chunks require an agent connection".to_string()); - }; - let mut client = wait_for_value_opt(cancel_token.clone(), timeout, client.lock()).await?; - client.call_with_timeout_and_cancel(method, params, timeout, cancel_token).await + match handle { + PoolKind::Agent(client) => { + let mut client = wait_for_value_opt(cancel_token.clone(), timeout, client.lock()).await?; + client.call_with_timeout_and_cancel(method, params, timeout, cancel_token).await + } + PoolKind::ExternalDriver { session, config, .. } => { + let mut params = params; + params["connection"] = serde_json::json!(config.as_ref()); + wait_for_value_opt( + cancel_token, + timeout, + session.invoke_with_timeout::(method, params, timeout), + ) + .await? + } + _ => Err("LOB snapshot chunks require a supported connection".to_string()), + } } pub async fn write_large_value_snapshot( @@ -4607,6 +4645,7 @@ fn error_query_result(message: String) -> db::QueryResult { db::QueryResult { columns: vec!["Error".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -4627,6 +4666,7 @@ fn empty_query_result(execution_time_ms: u128) -> db::QueryResult { db::QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -5041,6 +5081,7 @@ async fn execute_statements_inner( Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -5646,6 +5687,7 @@ async fn exec_tx_pg_inner( (Ok(total_affected), Ok(_)) => Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -5743,6 +5785,7 @@ async fn exec_tx_mysql_inner( Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -5940,6 +5983,7 @@ async fn exec_tx_sqlite_inner( conn.execute_batch("COMMIT").map_err(|e| format!("COMMIT failed: {}", e)).map(|_| db::QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -6074,6 +6118,7 @@ async fn exec_tx_explicit_inner( Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -6502,6 +6547,7 @@ pub async fn execute_in_manual_transaction( #[derive(Clone, Debug, Default)] pub struct ManualTransactionExecutionOptions { + pub defer_lobs: bool, pub bound_statements: Option>, pub max_rows: Option, pub table_data_preview: bool, @@ -6723,8 +6769,16 @@ pub async fn execute_in_manual_transaction_with_options( } } TxnConnection::ExternalDriver { session, config, .. } => { - execute_manual_txn_external_driver_statement(session, config, statement, database, schema, row_limit) - .await + execute_manual_txn_external_driver_statement( + session, + config, + statement, + database, + schema, + row_limit, + options.defer_lobs, + ) + .await } }; match result { @@ -7244,9 +7298,11 @@ async fn execute_manual_txn_external_driver_statement( database: &str, schema: Option<&str>, row_limit: usize, + defer_lobs: bool, ) -> Result { let timeout_secs = config.effective_query_timeout_secs(); let options = QueryExecutionOptions { + defer_lobs, max_rows: Some(row_limit.max(1)), timeout_secs: Some(timeout_secs), ..QueryExecutionOptions::default() @@ -7336,6 +7392,7 @@ async fn execute_manual_txn_mysql_statement( return Ok(db::QueryResult { columns: Vec::new(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -7369,6 +7426,7 @@ async fn execute_manual_txn_mysql_statement( Ok(db::QueryResult { columns, column_types, + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -7390,6 +7448,7 @@ async fn execute_manual_txn_mysql_statement( Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -7452,6 +7511,7 @@ pub async fn commit_manual_transaction(state: &AppState, txn_session_id: &str) - Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -7486,6 +7546,7 @@ pub async fn rollback_manual_transaction(state: &AppState, txn_session_id: &str) Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -7509,6 +7570,54 @@ mod tests { use std::sync::atomic::Ordering; use tokio::time::timeout; + #[test] + fn deferred_cells_reach_all_multi_result_success_constructors_once() { + let sample = || { + serde_json::from_value::(serde_json::json!({ + "columns": ["__DBX_LARGE_VALUE_BYTES_user_column"], "rows": [["preview"]], + "affected_rows": 0, "execution_time_ms": 0, + "large_value_cells": [{ "row_index": 0, "column_index": 0, + "original_bytes": 67108864, "value_ref": "original-reference" }] + })) + .unwrap() + }; + for result in [ + ExecuteMultiResult::from(sample()), + ExecuteMultiResult::success_with_index(sample(), 0), + ExecuteMultiResult::success_with_index_and_large_values(sample(), 0, Vec::new(), false), + ExecuteMultiResult::success_with_optional_server_large_values(sample(), false), + ] { + assert!(result.result.deferred_large_value_cells.is_empty()); + assert_eq!(result.large_value_cells.len(), 1); + assert_eq!(result.result.columns, vec!["__DBX_LARGE_VALUE_BYTES_user_column"]); + let wire = serde_json::to_string(&result).unwrap(); + assert_eq!(wire.matches("\"large_value_cells\":").count(), 1); + assert!(wire.contains("original-reference")); + } + } + + #[test] + fn deferred_cells_follow_columns_when_nclob_markers_are_removed() { + let sample = || { + serde_json::from_value::(serde_json::json!({ + "columns": ["ID", "N", "__DBX_LARGE_VALUE_BYTES_T_1", "C"], + "rows": [[1, "national preview", "T:3", "clob preview"]], + "affected_rows": 0, "execution_time_ms": 0, + "large_value_cells": [{"row_index":0,"column_index":3,"original_bytes":67108864,"value_ref":"clob-original"}] + })).unwrap() + }; + for result in [ + ExecuteMultiResult::success_with_optional_server_large_values(sample(), true), + ExecuteMultiResult::success_with_index_and_large_values(sample(), 0, Vec::new(), true), + ] { + assert_eq!(result.result.columns, vec!["ID", "N", "C"]); + let deferred = result.large_value_cells.iter().find(|cell| cell.value_ref.is_some()).unwrap(); + assert_eq!(deferred.column_index, 2); + assert_eq!(deferred.value_ref.as_deref(), Some("clob-original")); + assert!(result.large_value_cells.iter().any(|cell| cell.column_index == 1 && cell.value_ref.is_none())); + } + } + mod manual_transaction_snapshot_tests { use super::*; use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -9964,6 +10073,7 @@ for line in sys.stdin: mysql_batch_result(db::QueryResult { columns: vec!["COUNT(*)".to_string()], column_types: vec!["BIGINT".to_string()], + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -10240,6 +10350,7 @@ for line in sys.stdin: let result_set = |value| db::QueryResult { columns: vec!["value".to_string()], column_types: vec!["INT".to_string()], + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -11133,7 +11244,30 @@ for line in sys.stdin: std::fs::write( &executable, format!( - "#!/bin/sh\nwhile IFS= read -r line; do\n id=$(printf '%s' \"$line\" | sed -E 's/.*\"id\":([0-9]+).*/\\1/')\n case \"$line\" in\n *'\"method\":\"executeInManualTransaction\"'*)\n echo executeInManualTransaction >> '{}'\n printf '{{\"id\":%s,\"result\":{{\"columns\":[\"value\"],\"column_types\":[],\"column_sortables\":[],\"rows\":[[42]],\"affected_rows\":0,\"execution_time_ms\":1,\"truncated\":false,\"session_id\":null,\"has_more\":false}}}}\\n' \"$id\"\n ;;\n *'\"method\":\"beginManualTransaction\"'*)\n echo beginManualTransaction >> '{}'\n printf '{{\"id\":%s,\"result\":{{\"ok\":true}}}}\\n' \"$id\"\n ;;\n *'\"method\":\"commitManualTransaction\"'*)\n echo commitManualTransaction >> '{}'\n printf '{{\"id\":%s,\"result\":{{\"ok\":true}}}}\\n' \"$id\"\n ;;\n esac\ndone\n", + r#"#!/bin/sh +while IFS= read -r line; do + id=$(printf '%s' "$line" | sed -E 's/.*"id":([0-9]+).*/\1/') + case "$line" in + *'"method":"executeInManualTransaction"'*) + echo executeInManualTransaction >> '{}' + printf '{{"id":%s,"result":{{"columns":["value"],"column_types":[],"column_sortables":[],"rows":[[42]],"affected_rows":0,"execution_time_ms":1,"truncated":false,"session_id":null,"has_more":false}}}}\n' "$id" + ;; + *'"method":"beginManualTransaction"'*) + echo beginManualTransaction >> '{}' + printf '{{"id":%s,"result":{{"ok":true}}}}\n' "$id" + ;; + *'"method":"read_large_value_chunk"'*|*'"method":"release_large_value"'*) + echo lob >> '{}' + printf '{{"id":%s,"result":{{"status":"ok","data":"original","next_offset":8,"eof":true,"value_kind":"text"}}}}\n' "$id" + ;; + *'"method":"commitManualTransaction"'*) + echo commitManualTransaction >> '{}' + printf '{{"id":%s,"result":{{"ok":true}}}}\n' "$id" + ;; + esac +done +"#, + calls.display(), calls.display(), calls.display(), calls.display() @@ -11231,11 +11365,45 @@ for line in sys.stdin: let results = execute_in_manual_transaction(&state, "txn-test", "SELECT 42", "dbx_test", None, Some(10)).await.unwrap(); assert_eq!(results[0].rows, vec![vec![serde_json::json!(42)]]); + let request: LargeValueRequest = serde_json::from_value(serde_json::json!({ + "connectionId":"jdbc-conn", "database":"dbx_test", "txnSessionId":"txn-test", + "valueRef":"original-ref", "offset":0, "limit":4096 + })) + .unwrap(); + for release in [false, true] { + let mut other_connection = request.clone(); + other_connection.connection_id = "other-connection".into(); + assert_eq!( + request_large_value(&state, other_connection, release).await.unwrap_err(), + "LOB transaction context expired or busy" + ); + state.transaction_sessions.write().await.get_mut("txn-test").unwrap().busy = true; + assert_eq!( + request_large_value(&state, request.clone(), release).await.unwrap_err(), + "LOB transaction context expired or busy" + ); + state.transaction_sessions.write().await.get_mut("txn-test").unwrap().busy = false; + state.transaction_sessions.write().await.get_mut("txn-test").unwrap().last_activity = + std::time::Instant::now() - Duration::from_secs(MANUAL_TRANSACTION_IDLE_TIMEOUT_SECS + 1); + assert_eq!( + request_large_value(&state, request.clone(), release).await.unwrap_err(), + "LOB transaction context expired or busy" + ); + state.transaction_sessions.write().await.get_mut("txn-test").unwrap().last_activity = + std::time::Instant::now(); + assert_eq!(request_large_value(&state, request.clone(), release).await.unwrap()["data"], "original"); + } commit_manual_transaction(&state, "txn-test").await.unwrap(); assert!(state.pool_handle(pool_key).await.is_none()); + for release in [false, true] { + assert_eq!( + request_large_value(&state, request.clone(), release).await.unwrap_err(), + MANUAL_TRANSACTION_SESSION_NOT_FOUND_ERROR + ); + } assert_eq!( std::fs::read_to_string(&calls).unwrap(), - "beginManualTransaction\nexecuteInManualTransaction\ncommitManualTransaction\n" + "beginManualTransaction\nexecuteInManualTransaction\nlob\nlob\ncommitManualTransaction\n" ); state.shutdown(Duration::from_secs(1)).await; @@ -11413,6 +11581,7 @@ for line in sys.stdin: Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -11440,6 +11609,7 @@ for line in sys.stdin: Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -12925,6 +13095,7 @@ for line in sys.stdin: let result = db::QueryResult { columns: vec!["id".to_string(), "nested".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -12959,6 +13130,7 @@ for line in sys.stdin: "note".to_string(), ], column_types: vec!["integer".to_string(), "text".to_string(), "bigint".to_string(), "text".to_string()], + deferred_large_value_cells: Vec::new(), column_sortables: vec![true, true, true, true], spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -13019,6 +13191,7 @@ for line in sys.stdin: "varchar2".to_string(), "varchar2".to_string(), ], + deferred_large_value_cells: Vec::new(), column_sortables: vec![true, true, true, true], spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -13079,6 +13252,7 @@ for line in sys.stdin: "text".to_string(), "text".to_string(), ], + deferred_large_value_cells: Vec::new(), column_sortables: vec![true; 5], spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -13145,6 +13319,7 @@ for line in sys.stdin: format!("{}B_0", crate::sql_dialect::DBX_LARGE_VALUE_BYTES_COLUMN_PREFIX), ], column_types: vec!["bytea".to_string(), "text".to_string()], + deferred_large_value_cells: Vec::new(), column_sortables: vec![true, true], spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -13307,6 +13482,7 @@ for line in sys.stdin: format!("{}C_1", crate::sql_dialect::DBX_LARGE_VALUE_BYTES_COLUMN_PREFIX), ], column_types: vec!["number".to_string(), "varchar2".to_string(), "varchar2".to_string()], + deferred_large_value_cells: Vec::new(), column_sortables: vec![true; 3], spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -13348,6 +13524,7 @@ for line in sys.stdin: let mut result = db::QueryResult { columns: vec!["embedding".to_string(), marker], column_types: vec!["text".to_string(), "text".to_string()], + deferred_large_value_cells: Vec::new(), column_sortables: vec![true, true], spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -13385,6 +13562,7 @@ for line in sys.stdin: format!("{}K_1", crate::sql_dialect::DBX_LARGE_VALUE_BYTES_COLUMN_PREFIX), ], column_types: vec!["text".to_string(), "text".to_string(), "text".to_string(), "text".to_string()], + deferred_large_value_cells: Vec::new(), column_sortables: vec![true; 4], spatial_columns: Vec::new(), spatial_values: Vec::new(), diff --git a/crates/dbx-core/src/query/mongo_ops.rs b/crates/dbx-core/src/query/mongo_ops.rs index b6d8e2291f..b47e60d3d2 100644 --- a/crates/dbx-core/src/query/mongo_ops.rs +++ b/crates/dbx-core/src/query/mongo_ops.rs @@ -1450,6 +1450,7 @@ fn query_result(columns: Vec, rows: Vec>, affecte QueryResult { columns, column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), diff --git a/crates/dbx-core/src/schema/agent_pg_sequences.rs b/crates/dbx-core/src/schema/agent_pg_sequences.rs index 338568a380..fdd91f6319 100644 --- a/crates/dbx-core/src/schema/agent_pg_sequences.rs +++ b/crates/dbx-core/src/schema/agent_pg_sequences.rs @@ -370,6 +370,7 @@ mod tests { "cycle".to_string(), ], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-core/src/schema/kingbase.rs b/crates/dbx-core/src/schema/kingbase.rs index b943244514..31257e9b72 100644 --- a/crates/dbx-core/src/schema/kingbase.rs +++ b/crates/dbx-core/src/schema/kingbase.rs @@ -275,6 +275,7 @@ mod tests { "nspname".to_string(), ], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-core/src/schema/mod.rs b/crates/dbx-core/src/schema/mod.rs index f3a637466f..21c8d0cffc 100644 --- a/crates/dbx-core/src/schema/mod.rs +++ b/crates/dbx-core/src/schema/mod.rs @@ -3314,6 +3314,7 @@ mod tests { db::QueryResult { columns: columns.iter().map(|column| (*column).to_string()).collect(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -4433,6 +4434,7 @@ done let result = db::QueryResult { columns: vec!["Table".to_string(), "Extra".to_string(), "CREATE TABLE".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -4463,6 +4465,7 @@ done let result = db::QueryResult { columns: vec!["Table".to_string(), "Create Table".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -4492,6 +4495,7 @@ done let result = db::QueryResult { columns: vec!["name".to_string(), "definition".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -4518,6 +4522,7 @@ done let result = db::QueryResult { columns: vec!["View".to_string(), "Extra".to_string(), "CREATE VIEW".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -4548,6 +4553,7 @@ done let result = db::QueryResult { columns: vec!["name".to_string(), "definition".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -5639,6 +5645,7 @@ for line in sys.stdin: let result = super::db::QueryResult { columns: vec!["table_name".to_string(), "table_type".to_string()], column_types: vec![], + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -5677,6 +5684,7 @@ for line in sys.stdin: "comment".to_string(), ], column_types: vec![], + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -5931,6 +5939,7 @@ for line in sys.stdin: let result = db::QueryResult { columns: vec!["COMMENTS".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -5951,6 +5960,7 @@ for line in sys.stdin: let empty = db::QueryResult { columns: vec!["COMMENTS".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -5979,6 +5989,7 @@ for line in sys.stdin: let result = db::QueryResult { columns: vec!["AUTO_INCREMENT".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -6024,6 +6035,7 @@ for line in sys.stdin: let result = db::QueryResult { columns: vec!["TABLE_NAME".to_string(), "COMMENTS".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -6198,6 +6210,7 @@ for line in sys.stdin: "COMMENTS".to_string(), ], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -6517,6 +6530,7 @@ done "TOTAL_BYTES".to_string(), ], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -6695,6 +6709,7 @@ done "TOTAL_BYTES".to_string(), ], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-core/src/schema/table_structure_sql/sqlite_rebuild.rs b/crates/dbx-core/src/schema/table_structure_sql/sqlite_rebuild.rs index 64662e97ad..f01490659b 100644 --- a/crates/dbx-core/src/schema/table_structure_sql/sqlite_rebuild.rs +++ b/crates/dbx-core/src/schema/table_structure_sql/sqlite_rebuild.rs @@ -678,6 +678,7 @@ fn execute_change_transaction( Ok(()) => Ok(db::QueryResult { columns: Vec::new(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-core/tests/jdbc_lob_context.rs b/crates/dbx-core/tests/jdbc_lob_context.rs new file mode 100644 index 0000000000..4ee8e36031 --- /dev/null +++ b/crates/dbx-core/tests/jdbc_lob_context.rs @@ -0,0 +1,96 @@ +#![cfg(unix)] + +use dbx_core::connection::{AppState, PoolKind}; +use dbx_core::models::connection::ConnectionConfig; +use dbx_core::plugins::{InstalledPlugin, PluginCompatibility, PluginDriverSession, PluginManifest, PluginRuntimeEnv}; +use dbx_core::query::{request_large_value, LargeValueRequest}; +use serde_json::json; +use std::{os::unix::fs::PermissionsExt, sync::Arc, time::Duration}; + +#[tokio::test] +async fn jdbc_lob_reads_and_releases_use_the_existing_workload_session_only() { + let directory = tempfile::tempdir().unwrap(); + let executable = directory.path().join("plugin.py"); + std::fs::write( + &executable, + r#"#!/usr/bin/env python3 +import json, pathlib, sys +for line in sys.stdin: + request = json.loads(line) + with pathlib.Path(__file__).with_name('calls.jsonl').open('a') as log: + log.write(json.dumps(request) + '\n') + result = {'status':'ok', 'data':'original', 'next_offset':8, 'eof':True, 'value_kind':'text'} + print(json.dumps({'id':request['id'], 'result':result}), flush=True) +"#, + ) + .unwrap(); + std::fs::set_permissions(&executable, std::fs::Permissions::from_mode(0o755)).unwrap(); + let session = Arc::new( + PluginDriverSession::start_for_test( + InstalledPlugin { + manifest: PluginManifest { + id: "jdbc".into(), + name: "JDBC".into(), + version: "test".into(), + protocol_version: 1, + executable: Some("plugin.py".into()), + ..Default::default() + }, + path: directory.path().into(), + compatibility: PluginCompatibility { + compatible: true, + backend_executable: Some(executable), + ..Default::default() + }, + provenance: None, + }, + "jdbc".into(), + PluginRuntimeEnv::default(), + ) + .await + .unwrap(), + ); + let state = + AppState::new(dbx_core::persistence::test_storage::open(&directory.path().join("state.db")).await.unwrap()); + let config: ConnectionConfig = serde_json::from_value(json!({ + "id":"jdbc-lob", "name":"fixture", "db_type":"jdbc", "host":"127.0.0.1", "port":1, + "username":"fixture", "password":"", "connection_string":"jdbc:fixture" + })) + .unwrap(); + state.configs.write().await.insert(config.id.clone(), config.clone()); + state + .update_connection_pools(|pools| { + pools.insert( + config.id.clone(), + PoolKind::ExternalDriver { + driver_id: "jdbc".into(), + config: Arc::new(config), + session: session.clone(), + }, + ); + }) + .await; + for release in [false, true] { + let mut request: LargeValueRequest = serde_json::from_value(json!({ + "connectionId":"jdbc-lob", "database":"", "valueRef":"original-ref", "offset":0, "limit":4096 + })) + .unwrap(); + request.client_session_id = Some("replacement-tab".into()); + assert_eq!( + request_large_value(&state, request.clone(), release).await.unwrap_err(), + "LOB context has expired; rerun the query" + ); + request.client_session_id = None; + assert_eq!(request_large_value(&state, request, release).await.unwrap()["data"], "original"); + } + let calls = std::fs::read_to_string(directory.path().join("calls.jsonl")).unwrap(); + let calls: Vec = calls.lines().map(|line| serde_json::from_str(line).unwrap()).collect(); + assert_eq!(calls.len(), 2, "Must not reconnect, replay SQL or access the replacement scope"); + assert_eq!(calls[0]["method"], "read_large_value_chunk"); + assert_eq!(calls[1]["method"], "release_large_value"); + for call in calls { + assert_eq!(call["params"]["valueRef"], "original-ref"); + assert_eq!(call["params"]["connection"]["id"], "jdbc-lob"); + } + state.shutdown(Duration::from_secs(2)).await; +} diff --git a/crates/dbx-driver-elasticsearch/src/lib.rs b/crates/dbx-driver-elasticsearch/src/lib.rs index 3dc9a37ac0..908c63de72 100644 --- a/crates/dbx-driver-elasticsearch/src/lib.rs +++ b/crates/dbx-driver-elasticsearch/src/lib.rs @@ -2375,6 +2375,7 @@ fn parse_elasticsearch_response_with_sql_parser( Ok(crate::types::QueryResult { columns, column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -2403,6 +2404,7 @@ fn parse_elasticsearch_response_with_sql_parser( Ok(crate::types::QueryResult { columns, column_types, + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -2577,6 +2579,7 @@ fn raw_json_response_result( crate::types::QueryResult { columns: vec!["status".to_string(), "response".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -2646,6 +2649,7 @@ fn parse_elasticsearch_rest_response_with_sql_parser( Ok(crate::types::QueryResult { columns: vec!["response".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -3154,6 +3158,7 @@ pub fn parse_tabular_sql_response( Some(QueryResult { columns: column_names, column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-driver-mysql/src/mysql.rs b/crates/dbx-driver-mysql/src/mysql.rs index cce7dbc168..35254dbecf 100644 --- a/crates/dbx-driver-mysql/src/mysql.rs +++ b/crates/dbx-driver-mysql/src/mysql.rs @@ -5566,6 +5566,7 @@ async fn execute_result_set_with_text_protocol_on_conn( return Ok(MySqlQueryResult::exact(QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -5611,6 +5612,7 @@ async fn execute_result_set_with_text_protocol_on_conn( return Ok(MySqlQueryResult::exact(QueryResult { columns, column_types, + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns, spatial_values, @@ -5708,6 +5710,7 @@ async fn execute_result_set_with_text_protocol_on_conn( result: QueryResult { columns, column_types, + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns, spatial_values, @@ -5842,6 +5845,7 @@ async fn execute_result_sets_with_text_protocol_on_conn( result: QueryResult { columns, column_types, + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns, spatial_values, @@ -5870,6 +5874,7 @@ async fn execute_result_sets_with_text_protocol_on_conn( results.push(MySqlQueryResult::exact(QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -6057,6 +6062,7 @@ async fn execute_result_set_with_prepared_protocol_on_conn( result: QueryResult { columns, column_types, + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns, spatial_values, @@ -6408,6 +6414,7 @@ pub async fn execute_transaction_statement_on_conn( QueryResult { columns, column_types, + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -6444,6 +6451,7 @@ pub async fn execute_transaction_statement_on_conn( QueryResult { columns, column_types, + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns, spatial_values, @@ -6597,6 +6605,7 @@ async fn execute_query_on_conn_with_limits_progress_resolved( return Ok(MySqlQueryResult::exact(QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -6636,6 +6645,7 @@ async fn execute_query_on_conn_with_limits_progress_resolved( Ok(MySqlQueryResult::exact(QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -6770,6 +6780,7 @@ pub async fn execute_non_result_batch_on_conn( let current_result = QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-driver-postgres/src/postgres.rs b/crates/dbx-driver-postgres/src/postgres.rs index 77976032ab..365bf33697 100644 --- a/crates/dbx-driver-postgres/src/postgres.rs +++ b/crates/dbx-driver-postgres/src/postgres.rs @@ -2179,6 +2179,7 @@ async fn execute_select_prepared( Ok(PreparedSelectOutcome::Complete(Box::new(QueryResult { columns, column_types, + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns, spatial_values, @@ -2288,6 +2289,7 @@ async fn execute_select_text( Ok(QueryResult { column_types: matching_pg_text_column_types(&columns, prepared_column_types), columns, + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns, spatial_values, @@ -8955,6 +8957,7 @@ pub async fn execute_query_with_max_rows( QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -9795,6 +9798,7 @@ async fn execute_query_with_max_rows_inner( affected.map_err(pg_error_to_string).map(|affected| QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], @@ -16444,6 +16448,7 @@ mod tests { "character_maximum_length".to_string(), ], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), diff --git a/crates/dbx-driver-sqlserver/src/sqlserver.rs b/crates/dbx-driver-sqlserver/src/sqlserver.rs index 17656037e9..5f56c15f6e 100644 --- a/crates/dbx-driver-sqlserver/src/sqlserver.rs +++ b/crates/dbx-driver-sqlserver/src/sqlserver.rs @@ -887,6 +887,7 @@ fn server_messages_query_result(messages: Vec, start: Instant) -> Option QueryResult { columns: vec![], column_types: vec![], + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -1010,6 +1011,7 @@ fn collect_first_result_finish( QueryResult { columns, column_types, + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns, spatial_values, @@ -1946,6 +1948,7 @@ fn push_sqlserver_result_set(results: &mut Vec, result: Option, s QueryResult { columns: vec!["status".to_string(), "response".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-drivers/src/db/dynamodb_driver.rs b/crates/dbx-drivers/src/db/dynamodb_driver.rs index ce8f911aed..b42a16e250 100644 --- a/crates/dbx-drivers/src/db/dynamodb_driver.rs +++ b/crates/dbx-drivers/src/db/dynamodb_driver.rs @@ -565,6 +565,7 @@ fn document_query_result(result: DocumentQueryResult, started: Instant) -> Query QueryResult { columns, column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -585,6 +586,7 @@ fn affected_query_result(affected_rows: u64, started: Instant) -> QueryResult { QueryResult { columns: Vec::new(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), diff --git a/crates/dbx-drivers/src/db/influxdb3_driver.rs b/crates/dbx-drivers/src/db/influxdb3_driver.rs index 259927682b..27f7030f27 100644 --- a/crates/dbx-drivers/src/db/influxdb3_driver.rs +++ b/crates/dbx-drivers/src/db/influxdb3_driver.rs @@ -358,6 +358,7 @@ fn build_query_result(rows: Vec, start: Instant) -> QueryResu } let affected = out_rows.len() as u64; QueryResult { + deferred_large_value_cells: Vec::new(), column_sortables: columns.iter().map(|_| false).collect(), spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-drivers/src/db/influxdb_driver.rs b/crates/dbx-drivers/src/db/influxdb_driver.rs index 15da4dd0fc..c94b79738f 100644 --- a/crates/dbx-drivers/src/db/influxdb_driver.rs +++ b/crates/dbx-drivers/src/db/influxdb_driver.rs @@ -644,6 +644,7 @@ pub async fn execute_query(client: &InfluxdbClient, database: &str, sql: &str) - if !series.is_empty() { let (columns, rows) = merge_v1_series(&series); return Ok(QueryResult { + deferred_large_value_cells: Vec::new(), column_sortables: columns.iter().map(|_| false).collect(), columns, column_types: vec![], @@ -664,6 +665,7 @@ pub async fn execute_query(client: &InfluxdbClient, database: &str, sql: &str) - Ok(QueryResult { columns: vec![], column_types: vec![], + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -863,6 +865,7 @@ fn parse_flux_csv(text: &str, start: Instant) -> Result { } let columns = headers.unwrap_or_default(); Ok(QueryResult { + deferred_large_value_cells: Vec::new(), column_sortables: columns.iter().map(|_| false).collect(), spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-drivers/src/db/meilisearch_driver.rs b/crates/dbx-drivers/src/db/meilisearch_driver.rs index c1da8907e5..a8ca9d4238 100644 --- a/crates/dbx-drivers/src/db/meilisearch_driver.rs +++ b/crates/dbx-drivers/src/db/meilisearch_driver.rs @@ -1905,6 +1905,7 @@ fn raw_response_result(status: u16, body: String, start: Instant, truncated: boo QueryResult { columns: vec!["status".to_string(), "response".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), diff --git a/crates/dbx-drivers/src/db/rqlite_driver.rs b/crates/dbx-drivers/src/db/rqlite_driver.rs index 9cd8678094..fe70c70fb5 100644 --- a/crates/dbx-drivers/src/db/rqlite_driver.rs +++ b/crates/dbx-drivers/src/db/rqlite_driver.rs @@ -319,6 +319,7 @@ pub async fn execute_query_with_max_rows( Ok(QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -374,6 +375,7 @@ fn query_result_from_rqlite_result( QueryResult { columns: result.columns, column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-drivers/src/db/salesforce_driver.rs b/crates/dbx-drivers/src/db/salesforce_driver.rs index e7965f5e17..301ee86aba 100644 --- a/crates/dbx-drivers/src/db/salesforce_driver.rs +++ b/crates/dbx-drivers/src/db/salesforce_driver.rs @@ -676,6 +676,7 @@ fn parse_soql_response(value: Value, elapsed_ms: u128, max_rows: usize) -> Query QueryResult { columns, column_types, + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -1176,6 +1177,7 @@ impl SfClient { Ok(QueryResult { columns: vec!["Id".to_string()], column_types: vec!["String".to_string()], + deferred_large_value_cells: Vec::new(), column_sortables: vec![false], spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -1212,6 +1214,7 @@ fn salesforce_affected_query_result(affected_rows: u64, started: Instant) -> Que QueryResult { columns: Vec::new(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), diff --git a/crates/dbx-drivers/src/db/solr_driver.rs b/crates/dbx-drivers/src/db/solr_driver.rs index 877af2425d..ac21342371 100644 --- a/crates/dbx-drivers/src/db/solr_driver.rs +++ b/crates/dbx-drivers/src/db/solr_driver.rs @@ -982,6 +982,7 @@ fn solr_table_result( QueryResult { columns, column_types, + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -1002,6 +1003,7 @@ fn solr_raw_json_response_result(status: u16, body_text: impl Into, star QueryResult { columns: vec!["status".to_string(), "response".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-drivers/src/db/sqlite.rs b/crates/dbx-drivers/src/db/sqlite.rs index 0da4bf7858..0ca7441040 100644 --- a/crates/dbx-drivers/src/db/sqlite.rs +++ b/crates/dbx-drivers/src/db/sqlite.rs @@ -2954,6 +2954,7 @@ fn execute_query_blocking( Ok(QueryResult { columns, column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -2973,6 +2974,7 @@ fn execute_query_blocking( Ok(QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-drivers/src/db/sqlite_worker.rs b/crates/dbx-drivers/src/db/sqlite_worker.rs index 8da6955e96..cb24332d65 100644 --- a/crates/dbx-drivers/src/db/sqlite_worker.rs +++ b/crates/dbx-drivers/src/db/sqlite_worker.rs @@ -123,6 +123,7 @@ impl SqliteWorkerSession<'_> { WorkerBody::Ok { columns, column_types, rows, affected_rows, truncated, .. } => Ok(QueryResult { columns: columns.unwrap_or_default(), column_types: column_types.unwrap_or_default(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-drivers/src/db/turso_driver.rs b/crates/dbx-drivers/src/db/turso_driver.rs index 9ee7f8903d..339470c3b4 100644 --- a/crates/dbx-drivers/src/db/turso_driver.rs +++ b/crates/dbx-drivers/src/db/turso_driver.rs @@ -387,6 +387,7 @@ pub async fn execute_query_with_max_rows( return Ok(QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -411,6 +412,7 @@ pub async fn execute_query_with_max_rows( Ok(QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -432,6 +434,7 @@ pub async fn execute_query_with_max_rows( Ok(QueryResult { columns: vec![], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -615,6 +618,7 @@ fn query_result_from_turso_result( QueryResult { columns: result.columns, column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-drivers/src/db/vector_driver.rs b/crates/dbx-drivers/src/db/vector_driver.rs index 11a38fef35..b735ebb403 100644 --- a/crates/dbx-drivers/src/db/vector_driver.rs +++ b/crates/dbx-drivers/src/db/vector_driver.rs @@ -978,6 +978,7 @@ fn json_to_query_result( QueryResult { columns: vec!["status".to_string(), "response".to_string()], column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -1017,6 +1018,7 @@ fn values_to_query_result(items: Vec, start: Instant) -> QueryResult { QueryResult { columns, column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-drivers/src/db/victoriametrics_driver.rs b/crates/dbx-drivers/src/db/victoriametrics_driver.rs index aa479fa105..e90aa4f5da 100644 --- a/crates/dbx-drivers/src/db/victoriametrics_driver.rs +++ b/crates/dbx-drivers/src/db/victoriametrics_driver.rs @@ -349,6 +349,7 @@ fn series_result_to_query_result(series: Vec, start: Instant) -> Q .chain(std::iter::once("double".to_string())) .chain(std::iter::repeat_n("string".to_string(), columns.len().saturating_sub(2))) .collect(), + deferred_large_value_cells: Vec::new(), column_sortables: vec![true; columns.len()], columns, affected_rows: rows.len() as u64, @@ -389,6 +390,7 @@ fn simple_result(rows: Vec>, value_type: &str, start: Instant) -> Que QueryResult { columns: vec!["timestamp".to_string(), "value".to_string()], column_types: vec!["timestamp".to_string(), value_type.to_string()], + deferred_large_value_cells: Vec::new(), column_sortables: vec![true, true], affected_rows: rows.len() as u64, rows, diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index 7c8ec93c87..7043e38163 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -2792,6 +2792,7 @@ fn query_result(columns: Vec, rows: Vec>, affected_rows: u64) dbx_core::db::QueryResult { columns, column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-mcp/src/http.rs b/crates/dbx-mcp/src/http.rs index ab516f1250..6965ba2141 100644 --- a/crates/dbx-mcp/src/http.rs +++ b/crates/dbx-mcp/src/http.rs @@ -242,6 +242,7 @@ mod tests { result: dbx_core::db::QueryResult { columns: Vec::new(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index c77d074857..4aae8662c5 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -4422,6 +4422,7 @@ fn empty_query_result() -> dbx_core::db::QueryResult { dbx_core::db::QueryResult { columns: Vec::new(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -5556,6 +5557,7 @@ mod tests { result: dbx_core::db::QueryResult { columns: Vec::new(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), @@ -6026,6 +6028,7 @@ mod tests { result: dbx_core::db::QueryResult { columns: vec![], column_types: vec![], + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -9103,6 +9106,7 @@ mod tests { result: dbx_core::db::QueryResult { columns: vec![], column_types: vec![], + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -9136,6 +9140,7 @@ mod tests { result: dbx_core::db::QueryResult { columns: vec![], column_types: vec![], + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], @@ -9170,6 +9175,7 @@ mod tests { result: dbx_core::db::QueryResult { columns: vec!["body".to_string()], column_types: vec![], + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], diff --git a/crates/dbx-mcp/src/transaction.rs b/crates/dbx-mcp/src/transaction.rs index d726953253..758727239e 100644 --- a/crates/dbx-mcp/src/transaction.rs +++ b/crates/dbx-mcp/src/transaction.rs @@ -799,6 +799,7 @@ mod tests { QueryResult { columns: Vec::new(), column_types: Vec::new(), + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), diff --git a/crates/dbx-types/src/types.rs b/crates/dbx-types/src/types.rs index 6c140b9231..4b8e43f579 100644 --- a/crates/dbx-types/src/types.rs +++ b/crates/dbx-types/src/types.rs @@ -597,6 +597,8 @@ pub struct LargeValueCell { #[derive(Debug, Clone, Deserialize)] pub struct QueryResult { + #[serde(default, rename = "large_value_cells", deserialize_with = "deserialize_vec_or_default")] + pub deferred_large_value_cells: Vec, #[serde(default, deserialize_with = "deserialize_vec_or_default")] pub columns: Vec, /// Database type name for each column, parallel to `columns`. May be empty @@ -1999,10 +2001,25 @@ mod tests { } } + #[test] + fn query_result_receives_deferred_cells_without_serializing_a_second_metadata_field() { + let mut wire = serde_json::to_value(bigint_result_sample()).unwrap(); + wire["large_value_cells"] = serde_json::json!([{ + "row_index": 0, "column_index": 0, "original_bytes": 67108864, "value_ref": "original-lob" + }]); + let result: super::QueryResult = serde_json::from_value(wire).unwrap(); + assert_eq!(result.deferred_large_value_cells.len(), 1); + assert_eq!(result.deferred_large_value_cells[0].value_ref.as_deref(), Some("original-lob")); + let serialized = serde_json::to_value(&result).unwrap(); + assert!(serialized.get("large_value_cells").is_none()); + assert!(serialized.get("deferred_large_value_cells").is_none()); + } + fn bigint_result_sample() -> super::QueryResult { super::QueryResult { columns: vec!["id".to_string()], column_types: vec!["bigint".to_string()], + deferred_large_value_cells: Vec::new(), column_sortables: Vec::new(), spatial_columns: Vec::new(), spatial_values: Vec::new(), diff --git a/crates/dbx-web/src/routes/query.rs b/crates/dbx-web/src/routes/query.rs index 56bafaaae5..3494d72a26 100644 --- a/crates/dbx-web/src/routes/query.rs +++ b/crates/dbx-web/src/routes/query.rs @@ -15,6 +15,8 @@ use dbx_core::query_cancel::RunningTaskMetadata; #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct ExecuteQueryRequest { + #[serde(default)] + pub defer_lobs: bool, pub connection_id: String, pub database: String, pub sql: String, @@ -480,6 +482,9 @@ pub async fn execute_query( headers: HeaderMap, Json(req): Json, ) -> Result, AppError> { + if req.defer_lobs { + return Err(AppError::bad_request("Deferred LOB results require the multi-result query endpoint")); + } let allow_database_switch = req.client_session_id.as_deref().is_some_and(|id| !id.trim().is_empty()); let database = super::mcp_policy::ensure_sql( &state, @@ -540,6 +545,9 @@ pub async fn execute_conditional_update( headers: HeaderMap, Json(req): Json, ) -> Result, AppError> { + if req.defer_lobs { + return Err(AppError::bad_request("Deferred LOB results require the multi-result query endpoint")); + } let allow_database_switch = req.client_session_id.as_deref().is_some_and(|id| !id.trim().is_empty()); let database = super::mcp_policy::ensure_sql( &state, @@ -673,6 +681,7 @@ pub async fn execute_multi( req.schema.as_deref(), Some(cancel_token), dbx_core::query::QueryExecutionOptions { + defer_lobs: req.defer_lobs, max_rows: req.max_rows, fetch_size: req.fetch_size, page_size: req.page_size, @@ -1893,6 +1902,7 @@ mod tests { result: dbx_core::db::QueryResult { columns: vec!["Error".to_string()], column_types: vec![], + deferred_large_value_cells: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], diff --git a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java index e66bf99e97..9d063808fb 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -204,11 +204,18 @@ public final class DbxJdbcPlugin { String sharedConnectionKey = ""; Connection sharedConnection; boolean manualTransactionActive; + boolean deferLobs; final Map querySessions = new HashMap<>(); final Map cancellationRequests = new HashMap<>(); volatile Statement activeStatement; volatile boolean closing; final java.util.concurrent.locks.ReentrantLock operation = new java.util.concurrent.locks.ReentrantLock(); + final DeferredLobValues largeValues = new DeferredLobValues(call -> { + Statement previous = activeStatement; + trackActiveStatement(call); + try { call.execute(); } + finally { trackActiveStatement(previous); } + }, () -> closing || isCancellationRequested(CANCELLATION_REQUEST.get())); } static final class CancellationRequest { @@ -530,13 +537,60 @@ public final class DbxJdbcPlugin { } private static void writeResponse(BufferedWriter writer, ObjectNode response) throws IOException { + if (response.has("error") && !responseFitsBudget(response, 64 * 1024 * 1024)) { + ObjectNode bounded = MAPPER.createObjectNode(); + bounded.set("id", response.path("id")); + bounded.set("error", MAPPER.createObjectNode().put("message", RESPONSE_LIMIT_MESSAGE)); + response = bounded; + } synchronized (writer) { writer.write(MAPPER.writeValueAsString(response)); - writer.newLine(); + writer.write('\n'); writer.flush(); } } + private static final String RESPONSE_LIMIT_MESSAGE = "JDBC response exceeds the 64 MiB transport limit. The statement was not retried; changes may already have taken effect. Reduce the selected data and verify any writes before retrying."; + + static void checkResponseBudget(JsonNode id, JsonNode result, JsonNode params, int maximum) throws SQLException { + ObjectNode response = MAPPER.createObjectNode(); + response.set("id", id); + response.set("result", result); + try { + if (responseFitsBudget(response, maximum)) return; + } catch (IOException error) { + throw new SQLException("Failed to encode JDBC response", error); + } + String sessionId = optionalText(result, "session_id"); + for (JsonNode cell : result.path("large_value_cells")) connectionState().largeValues.release(cell.path("value_ref").asText()); + if (sessionId == null) sessionId = optionalText(params, "sessionId"); + if (sessionId != null) closeQuerySession(sessionId); + throw new SQLException(RESPONSE_LIMIT_MESSAGE); + } + + private static boolean responseFitsBudget(JsonNode response, int maximum) throws IOException { + try { + MAPPER.writeValue(new java.io.OutputStream() { + private int remaining = maximum - 1; + @Override public void write(int value) throws IOException { + if (remaining-- <= 0) throw new ResponseLimitException(); + } + @Override public void write(byte[] bytes, int offset, int count) throws IOException { + if (count > remaining) throw new ResponseLimitException(); + remaining -= count; + } + }, response); + return true; + } catch (IOException error) { + Throwable cause = error; + while (cause != null && !(cause instanceof ResponseLimitException)) cause = cause.getCause(); + if (cause == null) throw error; + return false; + } + } + + private static final class ResponseLimitException extends IOException {} + private static void closeLogicalSession(JdbcConnectionState state) { state.closing = true; boolean locked = false; @@ -582,7 +636,7 @@ public final class DbxJdbcPlugin { return LOGICAL_SESSIONS.get(sessionId); } - private static JsonNode handleLogicalSession(String method, JsonNode params, JsonNode connection, + private static JsonNode handleLogicalSession(JsonNode requestId, String method, JsonNode params, JsonNode connection, JdbcConnectionState state) throws Exception { String id = requireText(params, "jdbcSessionId"); if ("closeJdbcSession".equals(method)) { @@ -595,10 +649,14 @@ public final class DbxJdbcPlugin { if (shuttingDown || state.closing) throw new SQLException("JDBC logical session is closed"); if ("openJdbcSession".equals(method)) return okResult(); CONNECTION_STATE.set(state); + Set checkpoint = state.largeValues.checkpoint(); try { registerDrivers(connection); - return handle(method, params, connection); + JsonNode result = handle(method, params, connection); + checkResponseBudget(requestId, result, params, 64 * 1024 * 1024); + return result; } catch (Exception | LinkageError error) { + state.largeValues.releaseSince(checkpoint); if ("connect".equals(method) && (state.sharedConnection == null || !isPostgresJdbcUrl(optionalText(connection, "connection_string")))) { closeSharedConnection(); @@ -626,6 +684,7 @@ public final class DbxJdbcPlugin { response.set("id", id.isMissingNode() ? MAPPER.getNodeFactory().numberNode(1) : id); JsonNode connection = MAPPER.createObjectNode(); + Set checkpoint = request.path("params").hasNonNull("jdbcSessionId") ? null : connectionState().largeValues.checkpoint(); try { String method = requireText(request, "method"); JsonNode params = request.path("params"); @@ -637,7 +696,7 @@ public final class DbxJdbcPlugin { return response; } if (params.hasNonNull("jdbcSessionId")) { - response.set("result", handleLogicalSession(method, params, connection, retainedState)); + response.set("result", handleLogicalSession(response.path("id"), method, params, connection, retainedState)); return response; } if ("close".equals(method)) { @@ -653,8 +712,11 @@ public final class DbxJdbcPlugin { return response; } registerDrivers(connection); - response.set("result", handle(method, params, connection)); + JsonNode result = handle(method, params, connection); + checkResponseBudget(response.path("id"), result, params, 64 * 1024 * 1024); + response.set("result", result); } catch (Throwable error) { + if (checkpoint != null) connectionState().largeValues.releaseSince(checkpoint); // The plugin protocol boundary must report linkage errors from vendor drivers instead of exiting silently. ObjectNode errorNode = MAPPER.createObjectNode(); errorNode.put("message", enrichDriverHint(connection, throwableMessage(error))); @@ -736,6 +798,8 @@ public final class DbxJdbcPlugin { } private static JsonNode handle(String method, JsonNode params, JsonNode connection) throws Exception { + connectionState().deferLobs = params.path("deferLobs").asBoolean(false) + && Set.of("executeQuery", "executeQueryPage", "execute_query_page", "executeInManualTransaction", "execute_in_manual_transaction").contains(method); String database = optionalText(params, "database"); boolean exactIdentity = params.path("exact_identity").asBoolean(false); if (exactIdentity && !driverQuirks(connection).useOracleMetadata()) { @@ -750,6 +814,11 @@ public final class DbxJdbcPlugin { } } return switch (method) { + case "read_large_value_chunk" -> MAPPER.valueToTree(connectionState().largeValues.fetch( + connectionState().sharedConnection, requireText(params, "valueRef"), + params.path("offset").asLong(-1), params.path("limit").asInt(0))); + case "release_large_value" -> MAPPER.valueToTree(Map.of("released", + connectionState().largeValues.release(requireText(params, "valueRef")))); case "testConnection" -> connectionTestResult(openConnection(connection)); case "connect" -> { openConnection(connection); @@ -1305,6 +1374,7 @@ public final class DbxJdbcPlugin { } private static void applyOracleProperties(JsonNode connection, Properties properties) { + properties.putIfAbsent("oracle.net.disableOob", "true"); properties.putIfAbsent("remarksReporting", "false"); properties.putIfAbsent("restrictGetTables", "true"); properties.putIfAbsent("includeSynonyms", "false"); @@ -1431,6 +1501,7 @@ public final class DbxJdbcPlugin { ZoneId timestampZone = tdengineTimestampZone(connection, conn); boolean bitStringColumns = usesBitStringColumns(connection); boolean fullPrecisionIntegers = usesFullPrecisionIntegers(connection); + Boolean deferredOceanbase = deferredLobMode(connection, conn); try (Statement statement = conn.createStatement()) { trackActiveStatement(statement); if (connectionState().closing) throw new SQLException("JDBC logical session is closed"); @@ -1441,6 +1512,7 @@ public final class DbxJdbcPlugin { ObjectNode result = MAPPER.createObjectNode(); ArrayNode columns = MAPPER.createArrayNode(); ArrayNode rows = MAPPER.createArrayNode(); + ArrayNode largeCells = MAPPER.createArrayNode(); boolean truncated = false; try (ResultSet rs = executed.resultSet()) { @@ -1460,17 +1532,16 @@ public final class DbxJdbcPlugin { truncated = true; break; } - ArrayNode row = MAPPER.createArrayNode(); - for (int i = 1; i <= columnCount; i++) { - row.add(MAPPER.valueToTree(readValue(rs, meta, i, preserveOracleDateTime, timestampZone, bitStringColumns, fullPrecisionIntegers))); - } - rows.add(row); + ResultRow row = readResultRow(rs, meta, preserveOracleDateTime, timestampZone, bitStringColumns, fullPrecisionIntegers, deferredOceanbase); + appendLargeCells(largeCells, row, rows.size()); + rows.add(row.values()); } } } result.set("columns", columns); result.set("rows", rows); + result.set("large_value_cells", largeCells); result.put("affected_rows", columns.isEmpty() ? Math.max(executed.updateCount(), 0) : 0); result.put("execution_time_ms", (System.nanoTime() - start) / 1_000_000); result.put("truncated", truncated); @@ -1559,8 +1630,9 @@ public final class DbxJdbcPlugin { private final ZoneId timestampZone; private final boolean bitStringColumns; private final boolean fullPrecisionIntegers; + private final Boolean deferredOceanbase; private int rowsReturned; - private ArrayNode pendingRow; + private ResultRow pendingRow; private QuerySession( String id, @@ -1575,7 +1647,8 @@ public final class DbxJdbcPlugin { boolean preserveOracleDateTime, ZoneId timestampZone, boolean bitStringColumns, - boolean fullPrecisionIntegers + boolean fullPrecisionIntegers, + Boolean deferredOceanbase ) { this.id = id; this.statement = statement; @@ -1590,6 +1663,7 @@ public final class DbxJdbcPlugin { this.timestampZone = timestampZone; this.bitStringColumns = bitStringColumns; this.fullPrecisionIntegers = fullPrecisionIntegers; + this.deferredOceanbase = deferredOceanbase; } } @@ -1661,7 +1735,8 @@ public final class DbxJdbcPlugin { preserveOracleDateTime, timestampZone, bitStringColumns, - usesFullPrecisionIntegers(connection) + usesFullPrecisionIntegers(connection), + deferredLobMode(connection, conn) ); connectionState().querySessions.put(sessionId, session); try { @@ -1724,44 +1799,53 @@ public final class DbxJdbcPlugin { if (connectionState().closing) throw new SQLException("JDBC logical session is closed"); int effectivePageSize = Math.max(1, pageSize); ArrayNode rows = MAPPER.createArrayNode(); + ArrayNode largeCells = MAPPER.createArrayNode(); boolean truncated = false; + try { while (rows.size() < effectivePageSize && session.rowsReturned < session.maxRows) { - ArrayNode row; + ResultRow row; if (session.pendingRow != null) { row = session.pendingRow; session.pendingRow = null; } else { if (!session.resultSet.next()) { closeQuerySession(session.id); - return queryPageResult(session, rows, false, false); + return queryPageResult(session, rows, largeCells, false, false); } - row = readRow(session.resultSet, session.meta, session.preserveOracleDateTime, session.timestampZone, session.bitStringColumns, session.fullPrecisionIntegers); + row = readResultRow(session.resultSet, session.meta, session.preserveOracleDateTime, session.timestampZone, session.bitStringColumns, session.fullPrecisionIntegers, session.deferredOceanbase); } - rows.add(row); + appendLargeCells(largeCells, row, rows.size()); + rows.add(row.values()); session.rowsReturned++; } if (session.rowsReturned >= session.maxRows) { truncated = session.pendingRow != null || session.resultSet.next(); closeQuerySession(session.id); - return queryPageResult(session, rows, truncated, false); + return queryPageResult(session, rows, largeCells, truncated, false); } boolean hasMore = session.resultSet.next(); if (!hasMore) { closeQuerySession(session.id); - return queryPageResult(session, rows, false, false); + return queryPageResult(session, rows, largeCells, false, false); } - session.pendingRow = readRow(session.resultSet, session.meta, session.preserveOracleDateTime, session.timestampZone, session.bitStringColumns, session.fullPrecisionIntegers); - return queryPageResult(session, rows, false, true); + session.pendingRow = readResultRow(session.resultSet, session.meta, session.preserveOracleDateTime, session.timestampZone, session.bitStringColumns, session.fullPrecisionIntegers, session.deferredOceanbase); + return queryPageResult(session, rows, largeCells, false, true); + } catch (SQLException | RuntimeException | LinkageError error) { + for (JsonNode cell : largeCells) connectionState().largeValues.release(cell.path("value_ref").asText()); + closeQuerySession(session.id); + throw error; + } } - private static ObjectNode queryPageResult(QuerySession session, ArrayNode rows, boolean truncated, boolean hasMore) { + private static ObjectNode queryPageResult(QuerySession session, ArrayNode rows, ArrayNode largeCells, boolean truncated, boolean hasMore) { ObjectNode result = MAPPER.createObjectNode(); result.set("columns", session.columns.deepCopy()); result.set("rows", rows); + result.set("large_value_cells", largeCells); result.put("affected_rows", 0); result.put("execution_time_ms", (System.nanoTime() - session.startNanos) / 1_000_000); result.put("truncated", truncated); @@ -1785,6 +1869,7 @@ public final class DbxJdbcPlugin { if (session == null) { return false; } + if (session.pendingRow != null) releaseRowValues(session.pendingRow); try { session.resultSet.close(); } catch (Exception ignored) { @@ -1804,6 +1889,62 @@ public final class DbxJdbcPlugin { } } + private static Boolean deferredLobMode(JsonNode config, Connection connection) throws SQLException { + if (!connectionState().deferLobs) return null; + String url = jdbcUrl(config); + if (isOracleUrl(url)) return false; + if (!urlMatchesPrefix(url, "jdbc:oceanbase:")) return null; + try { + var protocolMethod = connection.getClass().getMethod("getProtocol"); + Object protocol = protocolMethod.invoke(connection); + boolean oracle = (boolean) protocolMethod.getReturnType().getMethod("isOracleMode").invoke(protocol); + return oracle ? Boolean.TRUE : null; + } catch (ReflectiveOperationException error) { + throw new SQLException("Cannot verify OceanBase LOB mode", error); + } + } + + private record ResultRow(ArrayNode values, ArrayNode largeCells) {} + + private static ResultRow readResultRow(ResultSet rs, ResultSetMetaData meta, boolean preserveOracleDateTime, + ZoneId timestampZone, boolean bitStringColumns, boolean fullPrecisionIntegers, Boolean oceanbase) throws SQLException { + if (oceanbase == null) return new ResultRow(readRow(rs, meta, preserveOracleDateTime, timestampZone, bitStringColumns, fullPrecisionIntegers), MAPPER.createArrayNode()); + ArrayNode row = MAPPER.createArrayNode(); + ArrayNode cells = MAPPER.createArrayNode(); + DeferredLobValues values = connectionState().largeValues; + Set checkpoint = values.checkpoint(); + try { + for (int i = 1; i <= meta.getColumnCount(); i++) { + int type = meta.getColumnType(i); + if (type == Types.CLOB || type == Types.BLOB) { + Object locator = type == Types.BLOB ? rs.getBlob(i) : rs.getClob(i); + DeferredLobValues.Preview preview = values.capture(connectionState().sharedConnection, locator, oceanbase); + row.add(preview.text()); + if (preview.ref() != null) cells.add(MAPPER.createObjectNode() + .put("column_index", i - 1).put("original_bytes", preview.originalBytes()).put("value_ref", preview.ref())); + } else { + row.add(MAPPER.valueToTree(readValue(rs, meta, i, preserveOracleDateTime, timestampZone, bitStringColumns, fullPrecisionIntegers))); + } + } + return new ResultRow(row, cells); + } catch (SQLException error) { + values.releaseSince(checkpoint); + throw error; + } + } + + private static void appendLargeCells(ArrayNode cells, ResultRow row, int rowIndex) { + for (JsonNode cell : row.largeCells()) { + ObjectNode target = cell.deepCopy(); + target.put("row_index", rowIndex); + cells.add(target); + } + } + + private static void releaseRowValues(ResultRow row) { + for (JsonNode cell : row.largeCells()) connectionState().largeValues.release(cell.path("value_ref").asText()); + } + private static ArrayNode readRow( ResultSet rs, ResultSetMetaData meta, @@ -4504,6 +4645,7 @@ public final class DbxJdbcPlugin { private static void closeSharedConnection() { JdbcConnectionState state = connectionState(); closeAllQuerySessions(); + state.largeValues.clear(); if (state.sharedConnection != null) { if (state.manualTransactionActive) { try { diff --git a/plugins/jdbc/src/main/java/app/dbx/jdbc/DeferredLobValues.java b/plugins/jdbc/src/main/java/app/dbx/jdbc/DeferredLobValues.java new file mode 100644 index 0000000000..32ff357f6c --- /dev/null +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DeferredLobValues.java @@ -0,0 +1,281 @@ +package app.dbx.jdbc; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.sql.*; +import java.util.*; +import java.util.function.BooleanSupplier; + +final class DeferredLobValues { + static final int PREVIEW = 256; + static final int MAX_CHUNK = 4096; + private final Map entries = new LinkedHashMap<>(); + @FunctionalInterface + interface Executor { void execute(CallableStatement call) throws SQLException; } + private final Executor executor; + private final BooleanSupplier cancelled; + private static final int READ_AHEAD_SLOTS = 64; + private Window window; + + DeferredLobValues(Executor executor, BooleanSupplier cancelled) { + this.executor = executor; + this.cancelled = cancelled; + } + + record Preview(String text, String ref, long originalBytes) {} + record Chunk(String status, String data, long next_offset, boolean eof, String value_kind) {} + private record Block(String data, int count, long nativeCount, boolean eof) {} + private record Window(Entry owner, String data, long start, int count, long nativeLength) {} + private static final class Entry { + final Connection connection; + final Object locator; + final boolean oceanbase; + final boolean binary; + final long created = System.currentTimeMillis(); + long accessed = created; + long offset; + long nativeOffset; + Entry(Connection connection, Object locator, boolean oceanbase) { + this.connection = connection; + this.locator = locator; + this.oceanbase = oceanbase; + this.binary = locator instanceof Blob; + } + } + + Preview capture(Connection connection, Object locator, boolean oceanbase) throws SQLException { + if (locator == null) return new Preview(null, null, 0); + prune(); + Entry entry = new Entry(connection, locator, oceanbase); + boolean retained = false; + try { + if (entries.size() >= 4096) throw new SQLException("LOB reference limit reached; close older results"); + long length = entry.binary ? ((Blob) locator).length() : ((Clob) locator).length(); + if (length == 0) return new Preview(entry.binary ? "0x" : "", null, 0); + StringBuilder preview = new StringBuilder(); + boolean eof = false; + while (entry.offset < PREVIEW + 1 && !eof) { + Block block = read(entry, (int) (PREVIEW + 1 - entry.offset)); + preview.append(block.data); + advance(entry, block); + eof = block.eof; + } + String text = preview.toString(); + if (eof && entry.offset <= PREVIEW) return new Preview(entry.binary ? "0x" + text : text, null, length); + String ref = UUID.randomUUID().toString(); + entries.put(ref, entry); + retained = true; + String prefix = entry.binary ? "0x" + text.substring(0, PREVIEW * 2) + : text.substring(0, text.offsetByCodePoints(0, PREVIEW)); + return new Preview(prefix, ref, entry.binary ? length : 64L * 1024 * 1024); + } finally { + if (!retained) free(entry); + } + } + + Chunk fetch(Connection connection, String ref, long offset, int limit) throws SQLException { + if (offset < 0 || offset == Long.MAX_VALUE || limit < 1 || limit > MAX_CHUNK) + throw new IllegalArgumentException("Invalid LOB chunk offset or limit"); + prune(); + Entry entry = entries.get(ref); + if (entry == null || entry.connection != connection || connection == null || connection.isClosed()) { + release(ref); + return new Chunk("expired", "", offset, false, "text"); + } + entry.accessed = System.currentTimeMillis(); + try { + if (offset < entry.offset) { entry.offset = 0; entry.nativeOffset = 0; } + while (entry.offset < offset) { + Block skipped = readWindow(entry, (int) Math.min(MAX_CHUNK, offset - entry.offset)); + advance(entry, skipped); + if (skipped.eof && entry.offset < offset) throw new SQLException("LOB chunk offset exceeds original value"); + } + Block block = readWindow(entry, limit); + checkCancelled(); + advance(entry, block); + return new Chunk("ok", block.data, entry.offset, block.eof, entry.binary ? "binary" : "text"); + } catch (SQLException error) { + release(ref); + throw error; + } + } + + boolean release(String ref) { + Entry entry = entries.remove(ref); + if (entry == null) return false; + if (window != null && window.owner == entry) window = null; + free(entry); + return true; + } + + Set checkpoint() { return Set.copyOf(entries.keySet()); } + void releaseSince(Set checkpoint) { + for (String ref : List.copyOf(entries.keySet())) if (!checkpoint.contains(ref)) release(ref); + } + void clear() { window = null; entries.values().forEach(DeferredLobValues::free); entries.clear(); } + + private void prune() { + long now = System.currentTimeMillis(); + for (String ref : List.copyOf(entries.keySet())) { + Entry entry = entries.get(ref); + if (now - entry.accessed >= 5 * 60_000L || now - entry.created >= 30 * 60_000L) release(ref); + } + } + + private static void advance(Entry entry, Block block) { + entry.offset += block.count; + entry.nativeOffset += block.nativeCount; + } + + private void checkCancelled() throws SQLException { + if (cancelled.getAsBoolean()) throw new SQLException("JDBC request cancelled"); + } + + private Block readWindow(Entry entry, int limit) throws SQLException { + checkCancelled(); + if (window == null || window.owner != entry || entry.offset < window.start + || entry.offset >= window.start + window.count && entry.nativeOffset < window.nativeLength) { + window = null; + window = readAhead(entry); + } + Window current = window; + int start = Math.toIntExact(entry.offset - current.start); + int count = Math.min(limit, current.count - start); + String text = entry.binary ? current.data.substring(start * 2, (start + count) * 2) + : current.data.substring(current.data.offsetByCodePoints(0, start), current.data.offsetByCodePoints(0, start + count)); + long nativeCount = entry.binary || entry.oceanbase ? count : text.length(); + checkCancelled(); + return new Block(text, count, nativeCount, entry.nativeOffset + nativeCount >= current.nativeLength); + } + + private Window readAhead(Entry entry) throws SQLException { + int nativeAmount = entry.oceanbase ? (entry.binary ? 16384 : 8191) : MAX_CHUNK; + StringBuilder sql = new StringBuilder("DECLARE l " + (entry.binary ? "BLOB" : "CLOB") + + " := ?; o INTEGER := ?; z INTEGER; n INTEGER; b " + + (entry.binary ? "RAW(32767)" : "VARCHAR2(32767)") + "; BEGIN z:=DBMS_LOB.GETLENGTH(l); "); + for (int i = 0; i < READ_AHEAD_SLOTS; i++) { + sql.append("IF o<=z THEN n:=").append(nativeAmount).append("; DBMS_LOB.READ(l,n,o,b); o:=o+n; ELSE n:=0; b:=NULL; END IF; ?:=n; ?:=b; "); + } + sql.append("?:=z; END;"); + try (CallableStatement call = entry.connection.prepareCall(sql.toString())) { + if (entry.oceanbase) { + try { call.getClass().getMethod("setInternal").invoke(call); } + catch (ReflectiveOperationException error) { throw new SQLException("Unsupported OceanBase LOB statement", error); } + } + if (entry.binary) call.setBlob(1, (Blob) entry.locator); else call.setClob(1, (Clob) entry.locator); + call.setLong(2, entry.nativeOffset + 1); + for (int i = 0; i < READ_AHEAD_SLOTS; i++) { + call.registerOutParameter(3 + 2 * i, Types.INTEGER); + call.registerOutParameter(4 + 2 * i, entry.binary ? Types.VARBINARY : Types.VARCHAR); + } + int lengthIndex = 3 + 2 * READ_AHEAD_SLOTS; + call.registerOutParameter(lengthIndex, Types.BIGINT); + call.setQueryTimeout(20); + checkCancelled(); + executor.execute(call); + checkCancelled(); + long length = call.getLong(lengthIndex), nativeOffset = entry.nativeOffset; + if (length < 0 || nativeOffset > length) throw new SQLException("Invalid LOB chunk response"); + StringBuilder data = new StringBuilder(); + int count = 0; + for (int i = 0; i < READ_AHEAD_SLOTS; i++) { + int amount = call.getInt(3 + 2 * i); + if (amount < 0 || amount > nativeAmount + || nativeOffset < length && (amount == 0 || nativeOffset + amount > length) + || nativeOffset >= length && amount != 0) throw new SQLException("Invalid LOB chunk response"); + if (entry.binary) { + byte[] bytes = call.getBytes(4 + 2 * i); + if (bytes == null) bytes = new byte[0]; + if (bytes.length != amount) throw new SQLException("LOB byte count mismatch"); + data.append(HexFormat.of().withUpperCase().formatHex(bytes)); + count += amount; + } else { + String text; + if (entry.oceanbase) { + byte[] bytes = call.getBytes(4 + 2 * i); + try { text = bytes == null ? "" : StandardCharsets.UTF_8.newDecoder().decode(ByteBuffer.wrap(bytes)).toString(); } + catch (java.nio.charset.CharacterCodingException error) { throw new SQLException("Invalid UTF-8 in LOB chunk", error); } + } else { + text = call.getString(4 + 2 * i); + if (text == null) text = ""; + } + for (int j = 0; j < text.length(); j++) { + char ch = text.charAt(j); + if (Character.isHighSurrogate(ch)) { + if (++j >= text.length() || !Character.isLowSurrogate(text.charAt(j))) throw new SQLException("Invalid LOB character boundary"); + } else if (Character.isLowSurrogate(ch)) throw new SQLException("Invalid LOB character boundary"); + } + int points = text.codePointCount(0, text.length()); + if (amount != (entry.oceanbase ? points : text.length())) throw new SQLException("LOB character count mismatch"); + data.append(text); + count += points; + } + nativeOffset += amount; + checkCancelled(); + } + checkCancelled(); + return new Window(entry, data.toString(), entry.offset, count, length); + } + } + + private Block read(Entry entry, int limit) throws SQLException { + if (cancelled.getAsBoolean()) throw new SQLException("JDBC request cancelled"); + int nativeLimit = entry.binary || entry.oceanbase ? limit : Math.max(2, limit); + try (CallableStatement call = entry.connection.prepareCall( + "DECLARE l " + (entry.binary ? "BLOB" : "CLOB") + " := ?; n INTEGER := ?; o INTEGER := ?; " + + "b " + (entry.binary ? "RAW(32767)" : "VARCHAR2(32767)") + "; z INTEGER; " + + "BEGIN z:=DBMS_LOB.GETLENGTH(l); IF o>z THEN n:=0; ELSE DBMS_LOB.READ(l,n,o,b); END IF; " + + "? := n; ? := b; ? := z; END;")) { + if (entry.oceanbase) { + try { call.getClass().getMethod("setInternal").invoke(call); } + catch (ReflectiveOperationException error) { throw new SQLException("Unsupported OceanBase LOB statement", error); } + } + if (entry.binary) call.setBlob(1, (Blob) entry.locator); else call.setClob(1, (Clob) entry.locator); + call.setInt(2, nativeLimit); + call.setLong(3, entry.nativeOffset + 1); + call.registerOutParameter(4, Types.INTEGER); + call.registerOutParameter(5, entry.binary ? Types.VARBINARY : Types.VARCHAR); + call.registerOutParameter(6, Types.BIGINT); + call.setQueryTimeout(20); + if (cancelled.getAsBoolean()) throw new SQLException("JDBC request cancelled"); + executor.execute(call); + if (cancelled.getAsBoolean()) throw new SQLException("JDBC request cancelled"); + int amount = call.getInt(4); + long length = call.getLong(6); + if (amount < 0 || amount > nativeLimit || length < 0 + || entry.nativeOffset < length && (amount == 0 || entry.nativeOffset + amount > length) + || entry.nativeOffset >= length && amount != 0) throw new SQLException("Invalid LOB chunk response"); + if (entry.binary) { + byte[] bytes = call.getBytes(5); + if (bytes == null) bytes = new byte[0]; + if (bytes.length != amount) throw new SQLException("LOB byte count mismatch"); + return new Block(HexFormat.of().withUpperCase().formatHex(bytes), amount, amount, entry.nativeOffset + amount >= length); + } + String text; + if (entry.oceanbase) { + byte[] bytes = call.getBytes(5); + try { text = bytes == null ? "" : StandardCharsets.UTF_8.newDecoder().decode(ByteBuffer.wrap(bytes)).toString(); } + catch (java.nio.charset.CharacterCodingException error) { throw new SQLException("Invalid UTF-8 in LOB chunk", error); } + } else { + text = call.getString(5); + if (text == null) text = ""; + } + for (int i = 0; i < text.length(); i++) { + char ch = text.charAt(i); + if (Character.isHighSurrogate(ch)) { + if (++i >= text.length() || !Character.isLowSurrogate(text.charAt(i))) throw new SQLException("Invalid LOB character boundary"); + } else if (Character.isLowSurrogate(ch)) throw new SQLException("Invalid LOB character boundary"); + } + int count = text.codePointCount(0, text.length()); + if (amount != (entry.oceanbase ? count : text.length())) throw new SQLException("LOB character count mismatch"); + if (count > limit) { text = text.substring(0, text.offsetByCodePoints(0, limit)); count = limit; } + long nativeCount = entry.oceanbase ? count : text.length(); + return new Block(text, count, nativeCount, entry.nativeOffset + nativeCount >= length); + } + } + + private static void free(Entry entry) { + try { if (entry.binary) ((Blob) entry.locator).free(); else ((Clob) entry.locator).free(); } + catch (SQLException ignored) { } + } +} diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/DeferredLobValuesTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/DeferredLobValuesTest.java new file mode 100644 index 0000000000..a191ee3fd5 --- /dev/null +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/DeferredLobValuesTest.java @@ -0,0 +1,307 @@ +package app.dbx.jdbc; + +import org.junit.jupiter.api.Test; +import java.lang.reflect.Proxy; +import java.sql.*; +import java.util.*; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import javax.sql.rowset.serial.SerialBlob; +import javax.sql.rowset.serial.SerialClob; +import static org.junit.jupiter.api.Assertions.*; + +final class DeferredLobValuesTest { + public interface InternalStatement { void setInternal(); } + @Test + void nullAndEmptyValuesStayInline() throws Exception { + var registry = new DeferredLobValues(CallableStatement::execute, () -> false); + var connection = connection(); + assertNull(registry.capture(connection, null, false).text()); + assertEquals("", registry.capture(connection, new SerialClob(new char[0]), false).text()); + assertEquals("0x", registry.capture(connection, new SerialBlob(new byte[0]), false).text()); + assertTrue(registry.checkpoint().isEmpty()); + } + + @Test + void binaryChunksKeepOriginalBytesAndReleaseIsIdempotent() throws Exception { + var registry = new DeferredLobValues(CallableStatement::execute, () -> false); + var connection = connection(); + byte[] bytes = new byte[1025]; + for (int i = 0; i < bytes.length; i++) bytes[i] = (byte) i; + var preview = registry.capture(connection, new SerialBlob(bytes), false); + assertEquals("0x" + HexFormat.of().withUpperCase().formatHex(Arrays.copyOf(bytes, 256)), preview.text()); + assertEquals(bytes.length, preview.originalBytes()); + StringBuilder result = new StringBuilder(); + long offset = 0; + boolean eof; + do { + var chunk = registry.fetch(connection, preview.ref(), offset, 113); + assertEquals("binary", chunk.value_kind()); + result.append(chunk.data()); offset = chunk.next_offset(); eof = chunk.eof(); + } while (!eof); + assertEquals(HexFormat.of().withUpperCase().formatHex(bytes), result.toString()); + assertTrue(registry.release(preview.ref())); + assertFalse(registry.release(preview.ref())); + assertEquals("expired", registry.fetch(connection, preview.ref(), 0, 1).status()); + } + + @Test + void unicodeSingleCodePointChunksAndRewindPreserveTheValue() throws Exception { + var registry = new DeferredLobValues(CallableStatement::execute, () -> false); + var connection = connection(); + String value = "δΈ­πŸ™‚a".repeat(100); + var preview = registry.capture(connection, new SerialClob(value.toCharArray()), false); + assertEquals(256, preview.text().codePointCount(0, preview.text().length())); + StringBuilder result = new StringBuilder(); + long offset = 0; + do { + var chunk = registry.fetch(connection, preview.ref(), offset, 1); + result.append(chunk.data()); offset = chunk.next_offset(); + if (chunk.eof()) break; + } while (true); + assertEquals(value, result.toString()); + assertEquals("πŸ™‚", registry.fetch(connection, preview.ref(), 1, 1).data()); + registry.clear(); + assertTrue(registry.checkpoint().isEmpty()); + } + + @Test + void cancellationAndFailedChunksReleaseOnlyTheirReference() throws Exception { + AtomicBoolean cancelled = new AtomicBoolean(); + var registry = new DeferredLobValues(CallableStatement::execute, cancelled::get); + var connection = connection(); + var first = registry.capture(connection, new SerialBlob(new byte[300]), false); + var checkpoint = registry.checkpoint(); + var second = registry.capture(connection, new SerialBlob(new byte[300]), false); + registry.releaseSince(checkpoint); + assertEquals("expired", registry.fetch(connection, second.ref(), 0, 10).status()); + assertThrows(IllegalArgumentException.class, () -> registry.fetch(connection, first.ref(), -1, 10)); + assertEquals(checkpoint, registry.checkpoint()); + cancelled.set(true); + assertThrows(SQLException.class, () -> registry.fetch(connection, first.ref(), 0, 10)); + assertTrue(registry.checkpoint().isEmpty()); + } + + @Test + void cancellationAfterCallableExecutionDiscardsOnlyTheInterruptedLocator() throws Exception { + AtomicBoolean cancelled = new AtomicBoolean(); + AtomicBoolean cancelAfterRead = new AtomicBoolean(); + var registry = new DeferredLobValues(call -> { + call.execute(); + if (cancelAfterRead.get()) cancelled.set(true); + }, cancelled::get); + var connection = connection(); + var interrupted = registry.capture(connection, new SerialBlob(new byte[300]), false); + var survivor = registry.capture(connection, new SerialBlob(new byte[300]), false); + cancelAfterRead.set(true); + assertThrows(SQLException.class, () -> registry.fetch(connection, interrupted.ref(), 0, 10)); + cancelled.set(false); + cancelAfterRead.set(false); + assertEquals("expired", registry.fetch(connection, interrupted.ref(), 0, 10).status()); + assertEquals("ok", registry.fetch(connection, survivor.ref(), 0, 10).status()); + registry.clear(); + } + + @Test + void readAheadIsSharedButDeliveryRewindAndCancellationKeepTheirOwnBoundaries() throws Exception { + var calls = new AtomicInteger(); + var cancelled = new AtomicBoolean(); + var registry = new DeferredLobValues(CallableStatement::execute, cancelled::get); + var connection = connection(calls::incrementAndGet, () -> {}); + String value = "x".repeat(4095) + "πŸ™‚δΈ­" + "y".repeat(300000); + var first = registry.capture(connection, new SerialClob(value.toCharArray()), false); + var second = registry.capture(connection, new SerialBlob(new byte[400000]), false); + calls.set(0); + assertEquals("x", registry.fetch(connection, first.ref(), 0, 1).data()); + assertEquals(1, calls.get()); + assertEquals("πŸ™‚", registry.fetch(connection, first.ref(), 4095, 1).data()); + assertEquals("δΈ­", registry.fetch(connection, first.ref(), 4096, 1).data()); + assertEquals("x", registry.fetch(connection, first.ref(), 0, 1).data()); + assertEquals(1, calls.get()); + assertEquals("00", registry.fetch(connection, second.ref(), 0, 1).data()); + assertEquals(2, calls.get()); + assertEquals("x", registry.fetch(connection, first.ref(), 1, 1).data()); + assertEquals(3, calls.get()); + cancelled.set(true); + assertThrows(SQLException.class, () -> registry.fetch(connection, first.ref(), 2, 1)); + cancelled.set(false); + assertEquals("expired", registry.fetch(connection, first.ref(), 2, 1).status()); + assertEquals("ok", registry.fetch(connection, second.ref(), 1, 1).status()); + registry.clear(); + } + + @Test + void unicodeReadAheadCrossesItsWindowBoundaryAndEndsAtTheOriginalLength() throws Exception { + var registry = new DeferredLobValues(CallableStatement::execute, () -> false); + var connection = connection(); + String original = "a".repeat(262143) + "πŸ™‚δΈ­" + "z".repeat(4200) + "πŸ™‚"; + var preview = registry.capture(connection, new SerialClob(original.toCharArray()), false); + var actual = new StringBuilder(); + long offset = 0; + while (true) { + var chunk = registry.fetch(connection, preview.ref(), offset, 4096); + assertEquals(chunk.data().codePointCount(0, chunk.data().length()), chunk.next_offset() - offset); + actual.append(chunk.data()); + offset = chunk.next_offset(); + if (chunk.eof()) break; + } + assertEquals(original, actual.toString()); + var end = registry.fetch(connection, preview.ref(), offset, 1); + assertTrue(end.eof()); + assertEquals("", end.data()); + assertEquals("πŸ™‚", registry.fetch(connection, preview.ref(), 262143, 1).data()); + registry.clear(); + } + + @Test + void invalidLateOutSlotDiscardsTheWholeWindowAndOnlyItsReference() throws Exception { + var corrupt = new AtomicBoolean(); + var connection = connection(() -> {}, () -> {}, args -> { + if (corrupt.get() && args.containsKey(131)) args.put(129, 4097); + }); + var registry = new DeferredLobValues(CallableStatement::execute, () -> false); + var first = registry.capture(connection, new SerialBlob(new byte[400000]), false); + var survivor = registry.capture(connection, new SerialBlob(new byte[400000]), false); + corrupt.set(true); + assertThrows(SQLException.class, () -> registry.fetch(connection, first.ref(), 0, 1)); + assertEquals(Set.of(survivor.ref()), registry.checkpoint()); + corrupt.set(false); + assertEquals("expired", registry.fetch(connection, first.ref(), 0, 1).status()); + assertEquals("00", registry.fetch(connection, survivor.ref(), 0, 1).data()); + registry.clear(); + } + + @Test + void pruneAndClearDropTheOwnedWindowWithoutExpiringAnotherReference() throws Exception { + var registry = new DeferredLobValues(CallableStatement::execute, () -> false); + var connection = connection(); + var first = registry.capture(connection, new SerialBlob(new byte[400000]), false); + var survivor = registry.capture(connection, new SerialBlob(new byte[400000]), false); + registry.fetch(connection, first.ref(), 0, 1); + var windowField = DeferredLobValues.class.getDeclaredField("window"); + windowField.setAccessible(true); + assertNotNull(windowField.get(registry)); + var entriesField = DeferredLobValues.class.getDeclaredField("entries"); + entriesField.setAccessible(true); + Object entry = ((Map) entriesField.get(registry)).get(first.ref()); + var accessed = entry.getClass().getDeclaredField("accessed"); + accessed.setAccessible(true); + accessed.setLong(entry, System.currentTimeMillis() - 5 * 60_000L - 1); + assertEquals("expired", registry.fetch(connection, first.ref(), 1, 1).status()); + assertNull(windowField.get(registry)); + assertEquals("ok", registry.fetch(connection, survivor.ref(), 0, 1).status()); + assertNotNull(windowField.get(registry)); + registry.clear(); + assertNull(windowField.get(registry)); + assertTrue(registry.checkpoint().isEmpty()); + assertEquals("expired", registry.fetch(connection, survivor.ref(), 1, 1).status()); + } + + @Test + void oceanbaseWindowEndReturnsAShortPublicChunkBeforeTheNextWindow() throws Exception { + var calls = new AtomicInteger(); + var connection = connection(calls::incrementAndGet, () -> {}); + var registry = new DeferredLobValues(CallableStatement::execute, () -> false); + String original = "πŸ™‚".repeat(8192) + "δΈ­".repeat(8192) + "a".repeat(524224 - 16384 - 1) + "πŸ™‚ε°Ύ"; + var preview = registry.capture(connection, new SerialClob(original.toCharArray()), true); + calls.set(0); + assertThrows(IllegalArgumentException.class, () -> registry.fetch(connection, preview.ref(), 0, 4097)); + assertEquals(4096, registry.fetch(connection, preview.ref(), 0, 4096).next_offset()); + var last = registry.fetch(connection, preview.ref(), 520192, 4096); + assertEquals(4032, last.data().codePointCount(0, last.data().length())); + assertTrue(last.data().endsWith("πŸ™‚")); + assertEquals(524224, last.next_offset()); + assertFalse(last.eof()); + assertEquals(1, calls.get()); + var tail = registry.fetch(connection, preview.ref(), last.next_offset(), 4096); + assertEquals("ε°Ύ", tail.data()); + assertTrue(tail.eof()); + assertEquals(2, calls.get()); + assertEquals("πŸ™‚", registry.fetch(connection, preview.ref(), 524223, 1).data()); + assertEquals("πŸ™‚", registry.fetch(connection, preview.ref(), 0, 1).data()); + assertEquals("πŸ™‚", registry.fetch(connection, preview.ref(), 0, 1).data()); + assertEquals(3, calls.get()); + registry.clear(); + } + + @Test + void oceanbaseBinaryReadAheadKeepsAllBytesAcrossItsLargerWindow() throws Exception { + var calls = new AtomicInteger(); + var connection = connection(calls::incrementAndGet, () -> {}); + var registry = new DeferredLobValues(CallableStatement::execute, () -> false); + byte[] original = new byte[1048581]; + for (int i = 0; i < original.length; i++) original[i] = (byte) i; + var preview = registry.capture(connection, new SerialBlob(original), true); + calls.set(0); + var actual = new StringBuilder(); + long offset = 0; + while (true) { + var chunk = registry.fetch(connection, preview.ref(), offset, 4096); + actual.append(chunk.data()); + offset = chunk.next_offset(); + if (chunk.eof()) break; + } + assertEquals(HexFormat.of().withUpperCase().formatHex(original), actual.toString()); + assertEquals(2, calls.get()); + registry.clear(); + } + + // Simulate the JDBC boundary, including Oracle UTF-16 and OceanBase codepoint units. + private static Connection connection() { return connection(() -> {}, () -> {}); } + + static Connection connection(Runnable beforeExecute, Runnable onCancel) { + return connection(beforeExecute, onCancel, args -> {}); + } + + private static Connection connection(Runnable beforeExecute, Runnable onCancel, java.util.function.Consumer> afterExecute) { + return (Connection) Proxy.newProxyInstance(DeferredLobValuesTest.class.getClassLoader(), new Class[]{Connection.class}, (p, m, a) -> { + if (m.getName().equals("isClosed")) return false; + if (m.getName().equals("close")) return null; + if (!m.getName().equals("prepareCall")) throw new UnsupportedOperationException(m.getName()); + boolean batch = ((String) a[0]).contains("IF o<=z"); + var amount = java.util.regex.Pattern.compile("THEN n:=(\\d+);").matcher((String) a[0]); + int batchLimit = batch && amount.find() ? Integer.parseInt(amount.group(1)) : 0; + var oceanbase = new AtomicBoolean(); + Map args = new HashMap<>(); + return Proxy.newProxyInstance(DeferredLobValuesTest.class.getClassLoader(), new Class[]{CallableStatement.class, InternalStatement.class}, (q, method, values) -> { + String name = method.getName(); + if (name.equals("setInternal")) { oceanbase.set(true); return null; } + if (Set.of("setBlob", "setClob", "setInt", "setLong").contains(name)) { args.put((Integer) values[0], values[1]); return null; } + if (Set.of("registerOutParameter", "setQueryTimeout", "close").contains(name)) return null; + if (name.equals("cancel")) { onCancel.run(); return null; } + if (name.equals("execute")) { + beforeExecute.run(); + int limit = batch ? batchLimit : (Integer) args.get(2); + long start = (Long) args.get(batch ? 2 : 3); + Object lob = args.get(1); long length; + String clobText = oceanbase.get() && lob instanceof Clob clob ? clob.getSubString(1, (int) clob.length()) : null; + for (int slot = 0; slot < (batch ? 64 : 1); slot++) { + int amountIndex = batch ? 3 + slot * 2 : 4; + if (lob instanceof Blob blob) { + length = blob.length(); byte[] data = start > length ? new byte[0] : blob.getBytes(start, (int) Math.min(limit, length - start + 1)); + args.put(amountIndex, data.length); args.put(amountIndex + 1, data); + start += data.length; + } else if (oceanbase.get()) { + length = clobText.codePointCount(0, clobText.length()); + String text = start > length ? "" : clobText.substring(clobText.offsetByCodePoints(0, (int) start - 1), clobText.offsetByCodePoints(0, (int) Math.min(length, start - 1 + limit))); + int points = text.codePointCount(0, text.length()); + args.put(amountIndex, points); args.put(amountIndex + 1, text.getBytes(java.nio.charset.StandardCharsets.UTF_8)); + start += points; + } else { + Clob clob = (Clob) lob; length = clob.length(); + String text = start > length ? "" : clob.getSubString(start, (int) Math.min(limit, length - start + 1)); + if (!text.isEmpty() && Character.isHighSurrogate(text.charAt(text.length() - 1))) text = text.substring(0, text.length() - 1); + args.put(amountIndex, text.length()); args.put(amountIndex + 1, text); + start += text.length(); + } + args.put(batch ? 131 : 6, length); + } + afterExecute.accept(args); + return false; + } + if (Set.of("getInt", "getLong", "getString", "getBytes").contains(name)) return args.get((Integer) values[0]); + throw new UnsupportedOperationException(name); + }); + }); + } +} diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/LargeResponseRecoveryTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/LargeResponseRecoveryTest.java new file mode 100644 index 0000000000..ed65efa37b --- /dev/null +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/LargeResponseRecoveryTest.java @@ -0,0 +1,177 @@ +package app.dbx.jdbc; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; + +import java.io.BufferedReader; +import java.io.BufferedWriter; +import java.io.StringReader; +import java.io.Writer; +import java.util.ArrayList; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.*; + +final class LargeResponseRecoveryTest { + private static final int HOST_LINE_LIMIT = 64 * 1024 * 1024; + + @Test + void oracleCancellationDefaultsToInBandWithoutOverridingAnExplicitProperty() throws Exception { + var apply = DbxJdbcPlugin.class.getDeclaredMethod("applyOracleProperties", + com.fasterxml.jackson.databind.JsonNode.class, java.util.Properties.class); + apply.setAccessible(true); + var connection = new ObjectMapper().createObjectNode(); + var defaults = new java.util.Properties(); + apply.invoke(null, connection, defaults); + assertEquals("true", defaults.getProperty("oracle.net.disableOob")); + var explicit = new java.util.Properties(); + explicit.setProperty("oracle.net.disableOob", "false"); + apply.invoke(null, connection, explicit); + assertEquals("false", explicit.getProperty("oracle.net.disableOob")); + assertNull(defaults.getProperty("oracle.jdbc.javaNetNio")); + } + + @Test + void oversizedDriverErrorIsAlsoBoundedAndKeepsItsRequestId() throws Exception { + var mapper = new ObjectMapper(); + var response = mapper.createObjectNode().put("id", 17); + response.set("error", mapper.createObjectNode().put("message", "x".repeat(HOST_LINE_LIMIT))); + var write = DbxJdbcPlugin.class.getDeclaredMethod("writeResponse", BufferedWriter.class, + com.fasterxml.jackson.databind.node.ObjectNode.class); + write.setAccessible(true); + LineCollector output = new LineCollector(); + write.invoke(null, new BufferedWriter(output), response); + assertEquals(1, output.lines.size()); + assertTrue(output.lengths.get(0) < HOST_LINE_LIMIT); + var error = mapper.readTree(output.lines.get(0)); + assertEquals(17, error.path("id").asInt()); + assertTrue(error.path("error").path("message").asText().contains("not retried")); + assertFalse(error.has("result")); + } + + @Test + void oversizedLogicalPagePreservesOtherCursorsAndTheManualTransaction() throws Exception { + try { + call("openJdbcSession", "large"); + call("openJdbcSession", "survivor"); + call("connect", "large"); + call("connect", "survivor"); + call("executeQuery", "large", "CREATE TABLE IF NOT EXISTS PRESERVED(ID INT)"); + call("beginManualTransaction", "large"); + call("executeQuery", "large", "INSERT INTO PRESERVED VALUES(1)"); + var page = call("executeQueryPage", "survivor", "SELECT X FROM SYSTEM_RANGE(1,3)"); + String cursor = page.path("result").path("session_id").asText(); + assertFalse(cursor.isBlank(), page.toString()); + var oversized = call("executeQueryPage", "large", "SELECT CAST(REPEAT('x',67108864) AS CLOB) FROM SYSTEM_RANGE(1,2)"); + assertTrue(oversized.has("error")); + var next = call("executeQuery", "large", "SELECT 42"); + assertEquals(42, next.path("result").path("rows").path(0).path(0).asInt()); + var params = params("survivor").put("sessionId", cursor); + assertEquals(2, invoke("fetchQueryPage", params).path("result").path("rows").path(0).path(0).asInt()); + call("rollbackManualTransaction", "large"); + assertEquals(0, call("executeQuery", "large", "SELECT COUNT(*) FROM PRESERVED") + .path("result").path("rows").path(0).path(0).asInt()); + } finally { + call("closeJdbcSession", "large"); + call("closeJdbcSession", "survivor"); + invoke("close", new ObjectMapper().createObjectNode()); + var driverKey = DbxJdbcPlugin.class.getDeclaredField("logicalDriverKey"); + driverKey.setAccessible(true); + driverKey.set(null, null); + } + } + + private static com.fasterxml.jackson.databind.node.ObjectNode params(String session) { + var mapper = new ObjectMapper(); + var params = mapper.createObjectNode().put("jdbcSessionId", session).put("pageSize", 1).put("maxRows", 10); + params.set("connection", mapper.createObjectNode().put("connection_string", "jdbc:h2:mem:large-logical;DB_CLOSE_DELAY=-1").put("jdbc_driver_class", "org.h2.Driver")); + return params; + } + + private static com.fasterxml.jackson.databind.JsonNode call(String method, String session) throws Exception { + return invoke(method, params(session)); + } + + private static com.fasterxml.jackson.databind.JsonNode call(String method, String session, String sql) throws Exception { + return invoke(method, params(session).put("sql", sql)); + } + + private static com.fasterxml.jackson.databind.JsonNode invoke(String method, com.fasterxml.jackson.databind.node.ObjectNode params) throws Exception { + var request = new ObjectMapper().createObjectNode().put("id", 1).put("method", method); + request.set("params", params); + var handle = DbxJdbcPlugin.class.getDeclaredMethod("handleLine", String.class); + handle.setAccessible(true); + return (com.fasterxml.jackson.databind.JsonNode) handle.invoke(null, request.toString()); + } + + @Test + void responseBudgetCountsUtf8EscapingAndTheLineTerminator() throws Exception { + ObjectMapper mapper = new ObjectMapper(); + for (String value : List.of("plain", "δΈ­ζ–‡", "πŸ™‚", "\"\\\n\r\t")) { + var id = mapper.getNodeFactory().numberNode(1); + var result = mapper.createObjectNode().put("value", value); + var params = mapper.createObjectNode(); + var response = mapper.createObjectNode(); + response.set("id", id); + response.set("result", result); + int actual = mapper.writeValueAsBytes(response).length + 1; + assertDoesNotThrow(() -> DbxJdbcPlugin.checkResponseBudget(id, result, params, actual)); + assertThrows(java.sql.SQLException.class, + () -> DbxJdbcPlugin.checkResponseBudget(id, result, params, actual - 1)); + } + } + + @Test + void oversizedQueryReturnsAnErrorWithoutBreakingTheNextQuery() throws Exception { + String connection = "\"connection\":{\"connection_string\":\"jdbc:h2:mem:large-response\",\"jdbc_driver_class\":\"org.h2.Driver\"}"; + String input = "{\"id\":1,\"method\":\"connect\",\"params\":{" + connection + "}}\n" + + "{\"id\":2,\"method\":\"executeQuery\",\"params\":{" + connection + + ",\"sql\":\"SELECT CAST(REPEAT('x',67108864) AS CLOB)\"}}\n" + + "{\"id\":3,\"method\":\"executeQuery\",\"params\":{" + connection + ",\"sql\":\"SELECT 42\"}}\n" + + "{\"id\":4,\"method\":\"close\",\"params\":{}}\n"; + LineCollector output = new LineCollector(); + try { + DbxJdbcPlugin.serve(new BufferedReader(new StringReader(input)), new BufferedWriter(output)); + assertTrue(output.lengths.stream().allMatch(length -> length <= HOST_LINE_LIMIT), + "No response may exceed the production host JSONL limit: " + output.lengths); + ObjectMapper mapper = new ObjectMapper(); + var large = output.lines.stream().map(line -> { + try { return mapper.readTree(line); } + catch (Exception error) { throw new AssertionError(error); } + }).filter(line -> line.path("id").asInt() == 2).findFirst().orElseThrow(); + assertTrue(large.has("error"), "The oversized result must fail explicitly"); + var next = output.lines.stream().filter(line -> line.contains("\"id\":3")).findFirst().orElseThrow(); + assertEquals(42, mapper.readTree(next).path("result").path("rows").path(0).path(0).asInt()); + } finally { + var shuttingDown = DbxJdbcPlugin.class.getDeclaredField("shuttingDown"); + shuttingDown.setAccessible(true); + shuttingDown.set(null, false); + } + } + + private static final class LineCollector extends Writer { + final List lengths = new ArrayList<>(); + final List lines = new ArrayList<>(); + private final StringBuilder prefix = new StringBuilder(); + private long length; + + @Override + public void write(char[] chars, int offset, int count) { + for (int i = offset; i < offset + count; i++) { + char c = chars[i]; + length++; + if (c == '\n') { + lengths.add(length); + lines.add(prefix.toString()); + prefix.setLength(0); + length = 0; + } else if (prefix.length() < 4096) { + prefix.append(c); + } + } + } + + @Override public void flush() {} + @Override public void close() {} + } +} diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/StatementCancelTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/StatementCancelTest.java index 377abff9b3..3ce6742a38 100644 --- a/plugins/jdbc/src/test/java/app/dbx/jdbc/StatementCancelTest.java +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/StatementCancelTest.java @@ -154,6 +154,52 @@ final class StatementCancelTest { } } + @Test + void cancellingAnActiveLobChunkPreservesOtherRefsAndIgnoresStaleTokens() throws Exception { + var block = new java.util.concurrent.atomic.AtomicBoolean(); + var entered = new CountDownLatch(1); + var cancelled = new CountDownLatch(1); + var connection = DeferredLobValuesTest.connection(() -> { + if (!block.get()) return; + entered.countDown(); + try { assertTrue(cancelled.await(5, TimeUnit.SECONDS)); } + catch (InterruptedException error) { throw new AssertionError(error); } + }, cancelled::countDown); + var stateField = DbxJdbcPlugin.class.getDeclaredField("DEFAULT_CONNECTION_STATE"); + stateField.setAccessible(true); + Object state = stateField.get(null); + var connectionField = state.getClass().getDeclaredField("sharedConnection"); + connectionField.setAccessible(true); + var valuesField = state.getClass().getDeclaredField("largeValues"); + valuesField.setAccessible(true); + var values = (DeferredLobValues) valuesField.get(state); + values.clear(); + connectionField.set(state, connection); + var first = values.capture(connection, new javax.sql.rowset.serial.SerialBlob(new byte[300]), false); + var survivor = values.capture(connection, new javax.sql.rowset.serial.SerialBlob(new byte[300]), false); + try { + startServer(); + block.set(true); + write(1, "read_large_value_chunk", params().put("valueRef", first.ref()).put("offset", 0).put("limit", 10).put(DbxJdbcPlugin.CANCEL_TOKEN_PARAM, "lob-token")); + assertTrue(entered.await(5, TimeUnit.SECONDS)); + assertFalse(send(2, DbxJdbcPlugin.CANCEL_STATEMENT_METHOD, MAPPER.createObjectNode().put(DbxJdbcPlugin.CANCEL_TOKEN_PARAM, "stale-token")).path("result").path("cancelled").asBoolean()); + assertEquals(1, cancelled.getCount()); + assertTrue(send(3, DbxJdbcPlugin.CANCEL_STATEMENT_METHOD, MAPPER.createObjectNode().put(DbxJdbcPlugin.CANCEL_TOKEN_PARAM, "lob-token")).path("result").path("cancelled").asBoolean()); + assertTrue(await(1).has("error")); + block.set(false); + assertEquals("expired", send(4, "read_large_value_chunk", params().put("valueRef", first.ref()).put("offset", 0).put("limit", 1)).path("result").path("status").asText()); + assertEquals("ok", send(5, "read_large_value_chunk", params().put("valueRef", survivor.ref()).put("offset", 0).put("limit", 1)).path("result").path("status").asText()); + assertFalse(send(6, DbxJdbcPlugin.CANCEL_STATEMENT_METHOD, MAPPER.createObjectNode().put(DbxJdbcPlugin.CANCEL_TOKEN_PARAM, "lob-token")).path("result").path("cancelled").asBoolean()); + } finally { + block.set(false); + cancelled.countDown(); + stopServer(); + server = null; + values.clear(); + connectionField.set(state, null); + } + } + private void startServer() throws Exception { startServer(Executors.newSingleThreadExecutor()); } diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index aabee37650..8c5fc3369f 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -247,6 +247,7 @@ pub async fn execute_multi( continue_on_error: Option, execution_mode: Option, preserve_explicit_transaction: Option, + defer_lobs: Option, ) -> Result, BackendError> { let execution_id = execution_id.filter(|id| !id.trim().is_empty()); let registered_query = execution_id.as_ref().map(|id| { @@ -315,6 +316,7 @@ pub async fn execute_multi( continue_on_error: continue_on_error.unwrap_or(false), execution_mode: execution_mode.unwrap_or_default(), preserve_explicit_transaction: preserve_explicit_transaction.unwrap_or(false), + defer_lobs: defer_lobs.unwrap_or(false), }, progress, ) @@ -642,6 +644,7 @@ pub async fn execute_in_manual_transaction( execution_id: Option, timeout_secs: Option, bound_statements: Option>, + defer_lobs: Option, ) -> Result, ManualTransactionCommandError> { let messages = query_messages_callback(&app, execution_id.as_deref()); dbx_core::query::query_messages::with_query_messages( @@ -653,6 +656,7 @@ pub async fn execute_in_manual_transaction( &database, schema.as_deref(), dbx_core::query::ManualTransactionExecutionOptions { + defer_lobs: defer_lobs.unwrap_or(false), bound_statements, max_rows, table_data_preview: table_data_preview.unwrap_or(false),