diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index 33a4f03a5f..4fc510f780 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -10472,8 +10472,10 @@ function downloadCellBinaryValue(rowIndex: number, columnIndex: number, mode: Bi async function downloadDetailBinaryValue(detail: DataGridCellDetail | null, mode: BinaryCellDownloadMode) { if (!detail || !canDownloadDetailBinaryValue(detail)) return; try { + const sourceResult = props.result; if (!(await hydrateLargeValueCell(detail.rowId, detail.colIndex))) return; - const resolvedDetail = cellDetailFor(detail.rowNumber - 1, detail.colIndex); + if (props.result !== sourceResult) return; + const resolvedDetail = cellDetailFor(displayRowIndexById(detail.rowId), detail.colIndex); if (!resolvedDetail) return; const payload = binaryCellDownloadPayload(resolvedDetail.value, mode, resolvedDetail.type, resolvedDatabaseType.value); const fileName = binaryCellDownloadFileName({ diff --git a/apps/desktop/src/components/grid/__tests__/DataGridBinaryDownload.spec.ts b/apps/desktop/src/components/grid/__tests__/DataGridBinaryDownload.spec.ts new file mode 100644 index 0000000000..947ce00c29 --- /dev/null +++ b/apps/desktop/src/components/grid/__tests__/DataGridBinaryDownload.spec.ts @@ -0,0 +1,158 @@ +// @vitest-environment happy-dom + +import { createApp, defineComponent, h, nextTick, ref, type App, type ComponentPublicInstance } from "vue"; +import { createPinia, setActivePinia } from "pinia"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import i18n from "@/i18n"; +import { TooltipProvider } from "@/components/ui/tooltip"; +import type { QueryResult } from "@/types/database"; +import type { DataGridCellDetail } from "@/lib/dataGrid/dataGridDetail"; + +const mocks = vi.hoisted(() => ({ hydrate: vi.fn(), save: vi.fn() })); + +vi.mock("@/composables/useDataGridLargeValues", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + useDataGridLargeValues: (options: Parameters[0]) => ({ ...actual.useDataGridLargeValues(options), hydrateLargeValueCell: mocks.hydrate }), + }; +}); + +vi.mock("@/lib/dataGrid/binaryCellDownload", async (importOriginal) => ({ + ...(await importOriginal()), + downloadBinaryCellPayload: mocks.save, +})); + +vi.mock("@/composables/useDataGridColumnResize", async (importOriginal) => { + const actual = await importOriginal(); + const { ref } = await import("vue"); + return { + ...actual, + useDataGridColumnResize: () => ({ + initColumnWidths: vi.fn(), + onResizeStart: vi.fn(), + autoFitColumn: vi.fn(), + renderedColumnWidths: ref([120, 120]), + totalWidth: ref(240), + columnVars: ref({ "--total-w": "240px" }), + getIsResizing: () => false, + }), + }; +}); + +import DataGrid from "../DataGrid.vue"; +import { useSettingsStore } from "@/stores/settingsStore"; + +const mounted: Array<{ app: App; host: HTMLElement }> = []; + +afterEach(() => { + for (const { app, host } of mounted.splice(0)) { + app.unmount(); + host.remove(); + } + vi.resetAllMocks(); +}); + +async function mountGrid() { + const pinia = createPinia(); + setActivePinia(pinia); + useSettingsStore().updateEditorSettings({ dataGridRenderMode: "canvas" }); + const result = ref({ + columns: ["ID", "CONTENT"], + column_types: ["int4", "bytea"], + rows: [ + [3, "0xff"], + [2, "\\x01..."], + [1, "0xfe"], + ], + affected_rows: 0, + execution_time_ms: 1, + large_value_cells: [{ row_index: 1, column_index: 1, original_bytes: 4 }], + }); + const grid = ref(); + const app = createApp( + defineComponent({ + setup: () => () => h(TooltipProvider, {}, { default: () => h(DataGrid, { ref: grid, result: result.value, databaseType: "postgres", context: "table-data" }) }), + }), + ); + const host = document.createElement("div"); + document.body.append(host); + app.use(pinia); + app.use(i18n); + app.mount(host); + mounted.push({ app, host }); + await nextTick(); + // Exercise the actual component handlers with its reactive display-row mapping. + const state = ( + grid.value!.$ as unknown as { + setupState: { + cellDetailFor: (row: number, column: number) => DataGridCellDetail; + downloadDetailBinaryValue: (detail: DataGridCellDetail, mode: "binary") => Promise; + localColumnFilters: Record>; + }; + } + ).setupState; + let finish!: (success: boolean) => void; + mocks.hydrate.mockImplementation( + () => + new Promise((resolve) => { + finish = resolve; + }), + ); + mocks.save.mockResolvedValue({ kind: "cancelled" }); + return { result, state, finish: (success = true) => finish(success) }; +} + +describe("DataGrid binary download identity", () => { + it("downloads the selected record when local filtering changes the display position during hydration", async () => { + const { result, state, finish } = await mountGrid(); + const detail = state.cellDetailFor(1, 1); + const pending = state.downloadDetailBinaryValue(detail, "binary"); + state.localColumnFilters = { 0: new Set(["num:2", "num:1"]) }; + await nextTick(); + expect(state.cellDetailFor(1, 1).rowId).not.toBe(detail.rowId); + result.value.rows[1]![1] = "\\x01020304"; + result.value.large_value_cells = []; + finish(); + await pending; + expect(mocks.save).toHaveBeenCalledOnce(); + expect(mocks.save.mock.calls[0]![0].data).toEqual(new Uint8Array([1, 2, 3, 4])); + }); + + it("does not download from a replacement result that reuses the same row index", async () => { + const { result, state, finish } = await mountGrid(); + const pending = state.downloadDetailBinaryValue(state.cellDetailFor(1, 1), "binary"); + result.value = { + ...result.value, + rows: [ + [9, "0xff"], + [8, "0xfe"], + ], + large_value_cells: [], + }; + await nextTick(); + finish(); + await pending; + expect(mocks.save).not.toHaveBeenCalled(); + }); + + it("does not substitute another record when the selected row is filtered out", async () => { + const { result, state, finish } = await mountGrid(); + const pending = state.downloadDetailBinaryValue(state.cellDetailFor(1, 1), "binary"); + state.localColumnFilters = { 0: new Set(["num:3", "num:1"]) }; + result.value.rows[1]![1] = "\\x01020304"; + result.value.large_value_cells = []; + await nextTick(); + finish(); + await pending; + expect(mocks.save).not.toHaveBeenCalled(); + }); + + it("does not download a preview after hydration fails", async () => { + const { state, finish } = await mountGrid(); + const pending = state.downloadDetailBinaryValue(state.cellDetailFor(1, 1), "binary"); + finish(false); + await pending; + expect(mocks.save).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/desktop/src/composables/__tests__/useDataGridLargeValues.spec.ts b/apps/desktop/src/composables/__tests__/useDataGridLargeValues.spec.ts index f11ba3bad2..7d6e32a1af 100644 --- a/apps/desktop/src/composables/__tests__/useDataGridLargeValues.spec.ts +++ b/apps/desktop/src/composables/__tests__/useDataGridLargeValues.spec.ts @@ -1,9 +1,10 @@ // @vitest-environment happy-dom -import { computed, createApp, defineComponent, h, ref } from "vue"; +import { computed, createApp, defineComponent, h, ref, type Ref } from "vue"; import { afterEach, describe, expect, it, vi } from "vitest"; +import { binaryCellDisplayText, binaryCellDownloadPayload } from "@/lib/dataGrid/binaryCellDownload"; import { useDataGridLargeValues } from "@/composables/useDataGridLargeValues"; -import type { QueryResult } from "@/types/database"; +import type { DatabaseType, QueryResult } from "@/types/database"; const mocks = vi.hoisted(() => ({ buildDataGridContextFilterCondition: vi.fn(), @@ -24,10 +25,161 @@ const disposers: Array<() => void> = []; afterEach(() => { for (const dispose of disposers.splice(0)) dispose(); - vi.clearAllMocks(); + vi.resetAllMocks(); }); +function mountLargeValues(databaseType: DatabaseType, result: Ref) { + const toast = vi.fn(); + const formatCell = (value: unknown, columnIndex?: number, originalBytes?: number) => binaryCellDisplayText(value, result.value.column_types?.[columnIndex ?? 0], originalBytes, databaseType) ?? String(value ?? ""); + const item = (rowId = 1) => ({ id: rowId, sourceIndex: rowId - 1, data: result.value.rows[rowId - 1]!, isNew: false, isDeleted: false, isDirtyCol: [false, false] }); + let largeValues: ReturnType | undefined; + const app = createApp( + defineComponent({ + setup() { + largeValues = useDataGridLargeValues({ + result: computed(() => result.value), + tableMeta: computed(() => ({ + schema: "public", + tableName: "APP_DATA", + tableType: "TABLE", + columns: result.value.columns.map((name, index) => ({ name, data_type: result.value.column_types![index]!, is_nullable: index !== 0, column_default: null, is_primary_key: index === 0, extra: null })), + primaryKeys: ["ID"], + })), + databaseType: computed(() => databaseType), + connectionId: computed(() => `${databaseType}-1`), + executionDatabase: computed(() => "MAXIMO"), + resultSourceColumns: computed(() => ["ID", "APP"]), + allColumnTypes: computed(() => result.value.column_types ?? []), + renderedGridColumns: computed(() => []), + showTranspose: ref(false), + displayRowCount: computed(() => 1), + gridScrollerElement: () => null, + displayItemAt: () => undefined, + getRowItem: (rowId) => (rowId >= 1 && rowId <= result.value.rows.length ? item(rowId) : undefined), + tableColumnForGridColumn: () => undefined, + formatCell, + formatCellCached: formatCell, + scheduleCanvasDraw: vi.fn(), + clearCellFormatCache: vi.fn(), + invalidateResultEstimate: vi.fn(), + connectionIdentifierQuote: () => '"', + getConnectionConfig: () => undefined, + includeDatabaseName: computed(() => true), + globalQueryTimeoutSecs: computed(() => 30), + resultLifecycle: { beginOperation: () => 1, isCurrent: () => true }, + largeValueResolutionVersion: ref(0), + runtimeScope: { disposed: false, addCleanup: () => () => undefined, dispose: () => undefined }, + uuid: () => "execution-1", + translate: (key) => key, + translateBackendError: (error) => String(error), + appendDebugLog: vi.fn(), + toast, + cloneRow: vi.fn(), + cloneRows: vi.fn(), + }); + return () => h("div"); + }, + }), + ); + app.mount(document.createElement("div")); + disposers.push(() => app.unmount()); + if (!largeValues) throw new Error("useDataGridLargeValues was not mounted"); + return { largeValues, toast, item }; +} + +function previewResult(columnType = "bytea", value = "\\x00017f80ff...", originalBytes = 9645) { + return ref({ + columns: ["ID", "APP"], + column_types: ["int4", columnType], + rows: [[1, value]], + affected_rows: 0, + execution_time_ms: 1, + large_value_cells: [{ row_index: 0, column_index: 1, original_bytes: originalBytes }], + }); +} + describe("useDataGridLargeValues", () => { + it.each([ + ["postgres", 3], + ["mysql", 1], + ] as const)("preserves %s fetch batching when a small binary column shares a row with another large value", async (databaseType, expectedBatches) => { + const result = previewResult("bytea", "0x00...", 8193); + result.value.columns.push("DOCUMENT"); + result.value.column_types!.push("text"); + result.value.rows = [1, 2, 3].map((id) => [id, "0x00...", "large document..."]); + result.value.large_value_cells = [0, 1, 2].flatMap((row_index) => [ + { row_index, column_index: 1, original_bytes: 8193 }, + { row_index, column_index: 2, original_bytes: 64 * 1024 * 1024 }, + ]); + const { largeValues } = mountLargeValues(databaseType, result); + mocks.buildDataGridContextFilterCondition.mockImplementation(({ value }) => Promise.resolve(`"ID" = ${value}`)); + mocks.buildTableSelectSql.mockImplementation((options) => Promise.resolve(JSON.stringify(options))); + mocks.executeMulti.mockImplementation(async (_connection, _database, sql) => { + const options = JSON.parse(sql); + const ids = [...options.whereInput.matchAll(/"ID" = (\d+)/g)].map((match) => Number(match[1])); + return [{ columns: ["ID", "APP", "DOCUMENT"], rows: ids.map((id) => [id, "0x00", "full document"]), affected_rows: 0, execution_time_ms: 1 }]; + }); + + const resolved = await largeValues.resolveLargeValueCells([1, 2, 3], [1]); + + expect([...resolved.keys()]).toEqual([1, 2, 3]); + expect(mocks.executeMulti).toHaveBeenCalledTimes(expectedBatches); + expect(mocks.buildTableSelectSql.mock.calls.every(([options]) => options.limit === 3 / expectedBatches)).toBe(true); + }); + + it.each(["0x", "0X", "\\x"])("shows the full bytea size and downloads all bytes after hydrating a %s preview", async (prefix) => { + const bytes = Uint8Array.from({ length: 9645 }, (_, index) => index % 256); + const hex = Buffer.from(bytes).toString("hex"); + const result = previewResult("bytea", `${prefix}${hex.slice(0, 3276 * 2)}...`, bytes.length); + const { largeValues, toast, item } = mountLargeValues("postgres", result); + mocks.buildDataGridContextFilterCondition.mockResolvedValue('"ID" = 1'); + mocks.buildTableSelectSql.mockResolvedValue('SELECT "ID", "APP" FROM "public"."APP_DATA" WHERE "ID" = 1 LIMIT 1'); + mocks.executeMulti.mockResolvedValue([{ columns: ["ID", "APP"], rows: [[1, `${prefix}${hex}`]], affected_rows: 0, execution_time_ms: 1 }]); + + expect(largeValues.formatGridItemCell(item(), 1)).toBe("BYTEA [9.4 KB]"); + expect(() => binaryCellDownloadPayload(result.value.rows[0]![1], "binary", "bytea", "postgres")).toThrow(); + await expect(largeValues.hydrateLargeValueCell(1, 1)).resolves.toBe(true); + + const selectOptions = mocks.buildTableSelectSql.mock.calls[0]![0]; + expect(selectOptions).toMatchObject({ databaseType: "postgres", columns: ["ID", "APP"], primaryKeys: ["ID"], limit: 1 }); + expect(selectOptions).not.toHaveProperty("largeValuePreviewSize"); + expect(result.value.large_value_cells).toEqual([]); + expect(largeValues.formatGridItemCell(item(), 1)).toBe("BYTEA [9.4 KB]"); + expect(binaryCellDownloadPayload(result.value.rows[0]![1], "binary", "bytea", "postgres").data).toEqual(bytes); + await expect(largeValues.hydrateLargeValueCell(1, 1)).resolves.toBe(true); + expect(mocks.executeMulti).toHaveBeenCalledTimes(1); + expect(toast).not.toHaveBeenCalled(); + }); + + it("retains an incomplete bytea preview when the full value cannot be fetched", async () => { + const result = previewResult(); + const { largeValues, toast, item } = mountLargeValues("postgres", result); + mocks.buildDataGridContextFilterCondition.mockResolvedValue('"ID" = 1'); + mocks.buildTableSelectSql.mockResolvedValue('SELECT "ID", "APP" FROM "APP_DATA" WHERE "ID" = 1'); + mocks.executeMulti.mockRejectedValue(new Error("connection lost")); + + await expect(largeValues.hydrateLargeValueCell(1, 1)).resolves.toBe(false); + + expect(largeValues.isLargeValuePreview(item(), 1)).toBe(true); + expect(result.value.rows[0]![1]).toBe("\\x00017f80ff..."); + expect(() => binaryCellDownloadPayload(result.value.rows[0]![1], "binary", "bytea", "postgres")).toThrow(); + expect(toast).toHaveBeenCalledWith("grid.largeValueLoadFailedWithMessage", 5000); + }); + + it.each([ + ["postgres", " BYTEA ", 64 * 1024 * 1024], + ["postgres", "text", undefined], + ["postgres", "bytea[]", undefined], + ["db2", "BLOB", undefined], + ["mysql", "LONGBLOB", 64 * 1024 * 1024], + ] as const)("uses only known original sizes for %s %s", (databaseType, columnType, expected) => { + const result = previewResult(columnType, "0x00...", 64 * 1024 * 1024); + const { largeValues, item } = mountLargeValues(databaseType, result); + expect(largeValues.largeValueOriginalBytes(item(), 1)).toBe(expected); + expect(largeValues.largeValueOriginalBytes({ ...item(), isDirtyCol: [false, true] }, 1)).toBeUndefined(); + expect(largeValues.largeValueOriginalBytes({ ...item(), isNew: true }, 1)).toBeUndefined(); + }); + it("hydrates a deferred DB2 BLOB through its primary key", async () => { const result = ref({ columns: ["ID", "APP"], @@ -37,7 +189,6 @@ describe("useDataGridLargeValues", () => { execution_time_ms: 1, large_value_cells: [{ row_index: 0, column_index: 1, original_bytes: 64 * 1024 * 1024 }], }); - const toast = vi.fn(); mocks.buildDataGridContextFilterCondition.mockResolvedValue('"ID" = 1'); mocks.buildTableSelectSql.mockResolvedValue('SELECT "ID", "APP" FROM "APP_DATA" WHERE "ID" = 1 FETCH FIRST 1 ROWS ONLY'); mocks.executeMulti.mockResolvedValue([ @@ -49,71 +200,7 @@ describe("useDataGridLargeValues", () => { }, ]); - let largeValues: ReturnType | undefined; - const app = createApp( - defineComponent({ - setup() { - largeValues = useDataGridLargeValues({ - result: computed(() => result.value), - tableMeta: computed(() => ({ - schema: "APP", - tableName: "APP_DATA", - tableType: "TABLE", - columns: [ - { name: "ID", data_type: "BIGINT", is_nullable: false, column_default: null, is_primary_key: true, extra: null }, - { name: "APP", data_type: "BLOB", is_nullable: true, column_default: null, is_primary_key: false, extra: null }, - ], - primaryKeys: ["ID"], - })), - databaseType: computed(() => "db2"), - connectionId: computed(() => "db2-1"), - executionDatabase: computed(() => "MAXIMO"), - resultSourceColumns: computed(() => ["ID", "APP"]), - allColumnTypes: computed(() => ["BIGINT", "BLOB"]), - renderedGridColumns: computed(() => []), - showTranspose: ref(false), - displayRowCount: computed(() => 1), - gridScrollerElement: () => null, - displayItemAt: () => undefined, - getRowItem: (rowId) => - rowId === 1 - ? { - id: 1, - sourceIndex: 0, - data: result.value.rows[0]!, - isNew: false, - isDeleted: false, - isDirtyCol: [false, false], - } - : undefined, - tableColumnForGridColumn: () => undefined, - formatCell: (value) => String(value ?? ""), - formatCellCached: (value) => String(value ?? ""), - scheduleCanvasDraw: vi.fn(), - clearCellFormatCache: vi.fn(), - invalidateResultEstimate: vi.fn(), - connectionIdentifierQuote: () => '"', - getConnectionConfig: () => undefined, - includeDatabaseName: computed(() => true), - globalQueryTimeoutSecs: computed(() => 30), - resultLifecycle: { beginOperation: () => 1, isCurrent: () => true }, - largeValueResolutionVersion: ref(0), - runtimeScope: { disposed: false, addCleanup: () => () => undefined, dispose: () => undefined }, - uuid: () => "execution-1", - translate: (key) => key, - translateBackendError: (error) => String(error), - appendDebugLog: vi.fn(), - toast, - cloneRow: vi.fn(), - cloneRows: vi.fn(), - }); - return () => h("div"); - }, - }), - ); - app.mount(document.createElement("div")); - disposers.push(() => app.unmount()); - if (!largeValues) throw new Error("useDataGridLargeValues was not mounted"); + const { largeValues, toast } = mountLargeValues("db2", result); await expect(largeValues.hydrateLargeValueCell(1, 1)).resolves.toBe(true); diff --git a/apps/desktop/src/composables/useDataGridLargeValues.ts b/apps/desktop/src/composables/useDataGridLargeValues.ts index 818ae3f655..931dcbd680 100644 --- a/apps/desktop/src/composables/useDataGridLargeValues.ts +++ b/apps/desktop/src/composables/useDataGridLargeValues.ts @@ -111,7 +111,9 @@ export function useDataGridLargeValues(options: UseDataGridLargeValuesOptions) { } function largeValueOriginalBytes(item: Pick | undefined, columnIndex: number): number | undefined { - if (options.databaseType.value !== "mysql" || !item || item.isNew || item.isDraft || item.sourceIndex === undefined || item.isDirtyCol[columnIndex]) return undefined; + // Other PostgreSQL preview types still carry an internal estimate, not an exact size. + const hasKnownByteSize = options.databaseType.value === "mysql" || (options.databaseType.value === "postgres" && options.allColumnTypes.value[columnIndex]?.trim().toLowerCase() === "bytea"); + if (!hasKnownByteSize || !item || item.isNew || item.isDraft || item.sourceIndex === undefined || item.isDirtyCol[columnIndex]) return undefined; return largeValueCellsByKey.value.get(largeValueCellKey(item.sourceIndex, columnIndex))?.original_bytes; } @@ -383,11 +385,15 @@ export function useDataGridLargeValues(options: UseDataGridLargeValuesOptions) { options.runtimeScope.addCleanup(pauseVisibleLargeValuePreviewHydration); function chunkLargeValueRequests(requests: LargeValueCellRequest[]): LargeValueCellRequest[][] { + // PostgreSQL full-value SQL uses SELECT *, so this column's exact byte size + // cannot bound other large columns in the same row. Keep the one-row batches + // used before bytea markers started reporting exact sizes. + const maxRows = options.databaseType.value === "postgres" ? 1 : LARGE_VALUE_FETCH_MAX_ROWS; const chunks: LargeValueCellRequest[][] = []; let chunk: LargeValueCellRequest[] = []; let bytes = 0; for (const request of requests) { - if (chunk.length > 0 && (chunk.length >= LARGE_VALUE_FETCH_MAX_ROWS || bytes + request.originalBytes > LARGE_VALUE_FETCH_TARGET_BYTES)) { + if (chunk.length > 0 && (chunk.length >= maxRows || bytes + request.originalBytes > LARGE_VALUE_FETCH_TARGET_BYTES)) { chunks.push(chunk); chunk = []; bytes = 0; diff --git a/crates/dbx-core/src/query/mod.rs b/crates/dbx-core/src/query/mod.rs index 16f6e6a940..fe45816660 100644 --- a/crates/dbx-core/src/query/mod.rs +++ b/crates/dbx-core/src/query/mod.rs @@ -573,6 +573,10 @@ fn truncate_server_large_value_preview( ServerLargeValuePreviewKind::Text => text.char_indices().nth(preview_size).map(|(index, _)| index), ServerLargeValuePreviewKind::Binary => text .strip_prefix("0x") + .or_else(|| text.strip_prefix("0X")) + // PostgreSQL text-protocol queries (including rows with timestamps) + // return bytea as \x hex instead of the binary driver's 0x form. + .or_else(|| text.strip_prefix("\\x")) .filter(|hex| hex.len() > preview_size.saturating_mul(2)) .map(|_| 2usize.saturating_add(preview_size.saturating_mul(2))), ServerLargeValuePreviewKind::Vector | ServerLargeValuePreviewKind::Deferred => unreachable!(), @@ -12058,7 +12062,18 @@ for line in sys.stdin: column_sortables: vec![true, true], spatial_columns: Vec::new(), spatial_values: Vec::new(), - rows: vec![vec![serde_json::json!("0x0102030405"), serde_json::json!("B:4:5")]], + rows: ["0x", "0X", "\\x"] + .into_iter() + .flat_map(|prefix| { + [0, 1, 4, 5].into_iter().map(move |bytes| { + vec![ + serde_json::json!(format!("{prefix}{}", "ab".repeat(bytes))), + serde_json::json!(format!("B:4:{bytes}")), + ] + }) + }) + .chain([vec![serde_json::Value::Null, serde_json::Value::Null]]) + .collect(), affected_rows: 0, execution_time_ms: 0, server_execute_time_us: None, @@ -12072,8 +12087,129 @@ for line in sys.stdin: let cells = extract_server_large_value_markers(&mut result); - assert_eq!(result.rows, vec![vec![serde_json::json!("0x01020304...")]]); - assert_eq!(cells, vec![db::LargeValueCell { row_index: 0, column_index: 0, original_bytes: 5 }]); + assert_eq!(result.columns, vec!["raw_value"]); + assert_eq!(result.column_types, vec!["bytea"]); + for (prefix_index, prefix) in ["0x", "0X", "\\x"].into_iter().enumerate() { + for (size_index, bytes) in [0, 1, 4, 5].into_iter().enumerate() { + let suffix = if bytes > 4 { "..." } else { "" }; + assert_eq!( + result.rows[prefix_index * 4 + size_index], + vec![serde_json::json!(format!("{prefix}{}{suffix}", "ab".repeat(bytes.min(4))))] + ); + } + } + assert_eq!(result.rows[12], vec![serde_json::Value::Null]); + assert_eq!( + cells, + [3, 7, 11] + .into_iter() + .map(|row_index| db::LargeValueCell { row_index, column_index: 0, original_bytes: 5 }) + .collect::>() + ); + } + + #[tokio::test] + #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at PostgreSQL (read-only queries)"] + async fn live_postgres_bytea_preview_preserves_size_and_full_value() { + use crate::sql_dialect::{build_table_data_select_sql, TableDataSelectSqlOptions}; + use sha2::{Digest, Sha256}; + + let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); + let pool = db::postgres::connect(&url, Duration::from_secs(5)).await.expect("connect PostgreSQL"); + let preview_size = 3276; + let sizes = [None, Some(0), Some(1), Some(3275), Some(3276), Some(3277), Some(9645), Some(261007)]; + let pattern = db::hex_encode(&(0..=255).collect::>()); + let inputs = sizes + .iter() + .enumerate() + .map(|(id, size)| { + format!("({id}, {})", size.map_or_else(|| "NULL::integer".to_string(), |n| n.to_string())) + }) + .collect::>() + .join(", "); + // A CTE supplies repeatable bytea data without creating database objects. + let fixture = format!( + "WITH payloads AS (SELECT id, substring(decode(repeat('{pattern}', (size + 255) / 256), 'hex') \ + from 1 for size) AS content FROM (VALUES {inputs}) AS inputs(id, size)), \ + bytea_preview AS (SELECT id, content, now() AS created_at, octet_length(content) AS bytes, \ + encode(sha256(content), 'hex') AS checksum FROM payloads) " + ); + + for text_protocol in [false, true] { + let mut columns = vec!["id".to_string(), "content".to_string()]; + let mut column_types = vec!["int4".to_string(), "bytea".to_string()]; + if text_protocol { + // A temporal column makes the driver choose PostgreSQL's text protocol. + columns.push("created_at".to_string()); + column_types.push("timestamptz".to_string()); + } + let select = build_table_data_select_sql(TableDataSelectSqlOptions { + database_type: Some(DatabaseType::Postgres), + table_name: "bytea_preview".to_string(), + primary_keys: vec!["id".to_string()], + columns: columns.clone(), + column_types: column_types.clone(), + large_value_preview_size: Some(preview_size), + limit: Some(sizes.len()), + order_by: Some("id".to_string()), + ..Default::default() + }); + let raw = db::postgres::execute_query(&pool, &format!("{fixture}{select}")).await.unwrap(); + let result = ExecuteMultiResult::success_with_optional_server_large_values(raw, true); + assert_eq!(result.result.columns, columns); + assert_eq!(result.result.column_types, column_types); + assert_eq!(result.result.rows.len(), sizes.len()); + let prefix = if text_protocol { "\\x" } else { "0x" }; + + for (row_index, size) in sizes.into_iter().enumerate() { + let Some(size) = size else { + assert!(result.result.rows[row_index][1].is_null()); + continue; + }; + let bytes = (0..size).map(|index| (index % 256) as u8).collect::>(); + let suffix = if size > preview_size { "..." } else { "" }; + assert_eq!( + result.result.rows[row_index][1], + serde_json::json!(format!("{prefix}{}{suffix}", db::hex_encode(&bytes[..size.min(preview_size)]))) + ); + if size <= preview_size { + continue; + } + // The download path fetches by primary key without preview options. + let full_select = build_table_data_select_sql(TableDataSelectSqlOptions { + database_type: Some(DatabaseType::Postgres), + table_name: "bytea_preview".to_string(), + primary_keys: vec!["id".to_string()], + columns: ["id", "content", "bytes", "checksum"].map(str::to_string).to_vec(), + where_input: Some(format!("id = {row_index}")), + limit: Some(1), + ..Default::default() + }); + let full = db::postgres::execute_query(&pool, &format!("{fixture}{full_select}")).await.unwrap(); + assert_eq!(full.rows.len(), 1); + // The table-data builder can select *, so map columns by name like + // the grid does. A returned timestamp also makes this fetch use text. + let column = |name: &str| &full.rows[0][full.columns.iter().position(|c| c == name).unwrap()]; + let value = column("content").as_str().unwrap(); + let hex = value.strip_prefix("0x").or_else(|| value.strip_prefix("\\x")).unwrap(); + assert_eq!(hex.len(), size * 2); + assert!(hex == db::hex_encode(&bytes), "full bytea differs for row {row_index}"); + let actual_size = + column("bytes").as_u64().or_else(|| column("bytes").as_str().and_then(|value| value.parse().ok())); + assert_eq!(actual_size, Some(size as u64)); + assert_eq!(column("checksum"), &serde_json::json!(db::hex_encode(&Sha256::digest(&bytes)))); + } + assert_eq!( + result.large_value_cells, + sizes + .into_iter() + .enumerate() + .filter_map(|(row_index, size)| size + .filter(|size| *size > preview_size) + .map(|original_bytes| { db::LargeValueCell { row_index, column_index: 1, original_bytes } })) + .collect::>() + ); + } } #[test] diff --git a/crates/dbx-sql-dialect/src/sql_dialect/table_select.rs b/crates/dbx-sql-dialect/src/sql_dialect/table_select.rs index 977c35f655..9ee37500ac 100644 --- a/crates/dbx-sql-dialect/src/sql_dialect/table_select.rs +++ b/crates/dbx-sql-dialect/src/sql_dialect/table_select.rs @@ -176,6 +176,8 @@ fn build_large_value_preview_columns(options: &TableDataSelectSqlOptions) -> Opt }; let marker = if database_type == Some(DatabaseType::Mysql) { format!("CONCAT('{marker_kind}:{preview_size}:', LENGTH({quoted})) AS {marker_alias}") + } else if database_type == Some(DatabaseType::Postgres) && kind == LargeValuePreviewKind::Binary { + format!("'{marker_kind}:{preview_size}:' || octet_length({quoted})::text AS {marker_alias}") } else { format!("'{marker_kind}:{preview_size}' AS {marker_alias}") }; diff --git a/crates/dbx-sql-dialect/src/sql_dialect/tests.rs b/crates/dbx-sql-dialect/src/sql_dialect/tests.rs index 6609901b76..dcf3131966 100644 --- a/crates/dbx-sql-dialect/src/sql_dialect/tests.rs +++ b/crates/dbx-sql-dialect/src/sql_dialect/tests.rs @@ -1021,8 +1021,8 @@ fn builds_postgres_table_data_large_value_previews() { schema: Some("public".to_string()), table_name: "large_rows".to_string(), primary_keys: vec!["id".to_string()], - columns: vec!["id".to_string(), "payload".to_string(), "metadata".to_string()], - column_types: vec!["integer".to_string(), "text".to_string(), "jsonb".to_string()], + columns: vec!["id".to_string(), "payload".to_string(), "metadata".to_string(), "content".to_string()], + column_types: vec!["integer".to_string(), "text".to_string(), "jsonb".to_string(), "bytea".to_string()], large_value_preview_size: Some(8192), limit: Some(100), ..Default::default() @@ -1032,6 +1032,8 @@ fn builds_postgres_table_data_large_value_previews() { assert!(sql.contains("'T:8192' AS \"__DBX_LARGE_VALUE_BYTES_T_1\"")); assert!(sql.contains("left(\"metadata\"::text, 8193) AS \"metadata\"")); assert!(sql.contains("'T:8192' AS \"__DBX_LARGE_VALUE_BYTES_K_2\"")); + assert!(sql.contains("substring(\"content\" from 1 for 8193) AS \"content\"")); + assert!(sql.contains("'B:8192:' || octet_length(\"content\")::text AS \"__DBX_LARGE_VALUE_BYTES_B_3\"")); } #[test] @@ -1138,7 +1140,7 @@ fn preserves_postgres_array_types_in_large_value_previews() { assert!(sql.contains("'T:8' AS \"__DBX_LARGE_VALUE_BYTES_K_11\"")); assert!(sql.contains("'T:8' AS \"__DBX_LARGE_VALUE_BYTES_S_12\"")); assert!(sql.contains("'V:8' AS \"__DBX_LARGE_VALUE_BYTES_V_13\"")); - assert!(sql.contains("'B:8' AS \"__DBX_LARGE_VALUE_BYTES_B_14\"")); + assert!(sql.contains("'B:8:' || octet_length(\"bytea_value\")::text AS \"__DBX_LARGE_VALUE_BYTES_B_14\"")); } #[test]