fix(jdbc): bound oversized responses and stream LOB snapshots

This commit is contained in:
efeire authored and GitHub committed 2026-10-24 16:21:30 +08:00
1 parent b7f655f379
commit b7ab31bdf5
57 files changed
+1665 -91

No files matched your search

+3
View File
@@ -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'
@@ -4502,6 +4502,7 @@ async function refreshSavedRows(request: { dirtyRows: ReadonlyMap<number, Readon
const tableMeta = props.tableMeta;
const connectionId = props.connectionId;
if (!tableMeta || !connectionId) return false;
if (connectionStore.getConfig(connectionId)?.db_type === "jdbc" && resolvedDatabaseType.value === "oracle" && tableMeta.columns.some((column) => /^(?: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,
@@ -144,7 +144,7 @@ export function useDataGridActions(activeTab: ComputedRef<QueryTab | undefined>)
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,
@@ -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,
@@ -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,
@@ -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")];
+3
View File
@@ -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<QueryResult[]> {
throw new Error("Manual transaction management is only available in the desktop app.");
}
+4
View File
@@ -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<QueryResult[]> {
return invokeBackend("execute_in_manual_transaction", {
txnSessionId,
@@ -2296,6 +2299,7 @@ export async function executeInManualTransaction(
executionId,
timeoutSecs,
boundStatements,
deferLobs,
});
}
@@ -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<string, never> {
export function tableDataLargeValuePreviewOptions(databaseType: DatabaseType | undefined, columns: readonly ColumnInfo[], primaryKeys: readonly string[], pageSize?: number, jdbcDeferredLobs = false): { columnTypes: string[]; largeValuePreviewSize: number } | Record<string, never> {
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;
@@ -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]);
});
});
@@ -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<QueryResult[]>; closeSession: (sessionId: string) => Promise<unknown>; isCancelled: () => boolean }): Promise<QueryResult[]> {
export async function executeCursorOffsetPage(options: {
offset: number;
limit: number;
fetchPage: (pageSize: number, sessionId?: string) => Promise<QueryResult[]>;
closeSession: (sessionId: string) => Promise<unknown>;
isCancelled: () => boolean;
releaseResults?: (results: QueryResult[]) => Promise<void>;
}): Promise<QueryResult[]> {
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<string, number> = {};
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);
@@ -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<QueryResult[]>((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" };
@@ -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");
});
+59 -38
View File
@@ -2308,12 +2308,13 @@ export const useQueryStore = defineStore("query", () => {
async function releaseResultLargeValues(...results: Array<QueryResult | undefined>) {
const released = new Set<string>();
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<QueryResult[]> => {
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<QueryResult[]>;
@@ -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) {
+3
View File
@@ -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(),
+1
View File
@@ -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(),
+1
View File
@@ -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![],
+1
View File
@@ -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![],
+1
View File
@@ -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![],
+199 -21
View File
@@ -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<db::LargeValueCell>,
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<db::QueryResult> 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<usize>,
pub fetch_size: Option<usize>,
pub page_size: Option<usize>,
@@ -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::<serde_json::Value>(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::<serde_json::Value>(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<Vec<dbx_types::types::BlobBoundStatement>>,
pub max_rows: Option<usize>,
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<db::QueryResult, String> {
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::<db::QueryResult>(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::<db::QueryResult>(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(),
+1
View File
@@ -1450,6 +1450,7 @@ fn query_result(columns: Vec<String>, rows: Vec<Vec<serde_json::Value>>, 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(),
@@ -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![],
+1
View File
@@ -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![],
+15
View File
@@ -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![],
@@ -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![],
+96
View File
@@ -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<serde_json::Value> = 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;
}
@@ -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![],
+11
View File
@@ -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![],
@@ -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(),
@@ -887,6 +887,7 @@ fn server_messages_query_result(messages: Vec<String>, 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<QueryResult>, result: Option<SqlS
results.push(QueryResult {
columns: result.columns,
column_types: result.column_types,
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -2006,6 +2009,7 @@ fn push_sqlserver_ordered_events(
result: QueryResult {
columns: vec![],
column_types: vec![],
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -3895,6 +3899,7 @@ async fn execute_query_with_max_rows_inner(
QueryResult {
columns: vec![],
column_types: Vec::new(),
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -3962,6 +3967,7 @@ async fn execute_batch_with_max_rows_metadata_inner(
QueryResult {
columns: vec![],
column_types: Vec::new(),
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -4062,6 +4068,7 @@ pub async fn execute_simple_batch_with_max_rows_metadata(
result: QueryResult {
columns: vec![],
column_types: Vec::new(),
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -4686,6 +4693,7 @@ mod tests {
let empty = QueryResult {
columns: vec![],
column_types: vec![],
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -4712,6 +4720,7 @@ mod tests {
QueryResult {
columns: vec![],
column_types: vec![],
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -4733,6 +4742,7 @@ mod tests {
let select = QueryResult {
columns: vec!["id".to_string()],
column_types: vec!["int".to_string()],
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -5992,6 +6002,7 @@ mod tests {
let context = completion_context_from_query_result(QueryResult {
columns: vec!["default_schema".to_string(), "engine_edition".to_string()],
column_types: vec![],
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -6555,6 +6566,7 @@ mod tests {
let mut result = QueryResult {
columns: vec!["id".to_string(), "__dbx_row_num".to_string()],
column_types: vec!["int".to_string(), "bigint".to_string()],
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -602,6 +602,7 @@ fn limited_query_result(result: ChJsonResult, execution_time_ms: u128, max_rows:
QueryResult {
columns,
column_types,
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -865,6 +866,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![],
@@ -479,6 +479,7 @@ fn query_result(
QueryResult {
columns: result.columns,
column_types: Vec::new(),
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -605,6 +605,7 @@ fn couchdb_table_result(
QueryResult {
columns,
column_types,
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -625,6 +626,7 @@ fn couchdb_raw_json_response_result(status: u16, body_text: impl Into<String>, 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![],
@@ -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(),
@@ -358,6 +358,7 @@ fn build_query_result(rows: Vec<serde_json::Value>, 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![],
@@ -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<QueryResult, String> {
}
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![],
@@ -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(),
@@ -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![],
@@ -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(),
+2
View File
@@ -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<String>, 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![],
+2
View File
@@ -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![],
@@ -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![],
@@ -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![],
@@ -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<Value>, start: Instant) -> QueryResult {
QueryResult {
columns,
column_types: Vec::new(),
deferred_large_value_cells: Vec::new(),
column_sortables: vec![],
spatial_columns: vec![],
spatial_values: vec![],
@@ -349,6 +349,7 @@ fn series_result_to_query_result(series: Vec<SeriesResult>, 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<Vec<Value>>, 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,
+1
View File
@@ -2792,6 +2792,7 @@ fn query_result(columns: Vec<String>, rows: Vec<Vec<Value>>, 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![],
+1
View File
@@ -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(),
+6
View File
@@ -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![],
+1
View File
@@ -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(),
+17
View File
@@ -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<LargeValueCell>,
#[serde(default, deserialize_with = "deserialize_vec_or_default")]
pub columns: Vec<String>,
/// 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(),
+10
View File
@@ -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<ExecuteQueryRequest>,
) -> Result<Json<dbx_core::db::QueryResult>, 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<ExecuteQueryRequest>,
) -> Result<Json<dbx_core::db::QueryResult>, 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![],
@@ -204,11 +204,18 @@ public final class DbxJdbcPlugin {
String sharedConnectionKey = "";
Connection sharedConnection;
boolean manualTransactionActive;
boolean deferLobs;
final Map<String, QuerySession> querySessions = new HashMap<>();
final Map<String, CancellationRequest> 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<String> 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<String> 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<String> 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 {
@@ -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<String, Entry> 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<String> checkpoint() { return Set.copyOf(entries.keySet()); }
void releaseSince(Set<String> 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) { }
}
}
@@ -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<Map<Integer, Object>> 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<Integer, Object> 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);
});
});
}
}
@@ -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<Long> lengths = new ArrayList<>();
final List<String> 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() {}
}
}
@@ -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());
}
+4
View File
@@ -247,6 +247,7 @@ pub async fn execute_multi(
continue_on_error: Option<bool>,
execution_mode: Option<dbx_core::query::QueryExecutionMode>,
preserve_explicit_transaction: Option<bool>,
defer_lobs: Option<bool>,
) -> Result<Vec<dbx_core::query::ExecuteMultiResult>, 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<String>,
timeout_secs: Option<u64>,
bound_statements: Option<Vec<db::BlobBoundStatement>>,
defer_lobs: Option<bool>,
) -> Result<Vec<dbx_core::query::ExecuteMultiResult>, 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),