From c75ce783dff2052dae88c788ca2fb96da2c61f77 Mon Sep 17 00:00:00 2001 From: Guoyu Su Date: Mon, 28 Sep 2026 01:01:09 +0800 Subject: [PATCH] fix(doris): preserve aggregate-state bytes in query results --- apps/desktop/src/components/grid/DataGrid.vue | 7 +- .../grid/__tests__/DataGridSurfaces.spec.ts | 8 + .../__tests__/useDataGridEditor.spec.ts | 62 +- .../src/composables/useDataGridEditor.ts | 8 +- .../__tests__/binaryCellFileImport.spec.ts | 41 +- .../src/lib/dataGrid/binaryCellDownload.ts | 21 +- crates/dbx-core/src/data/database_export.rs | 37 +- crates/dbx-core/src/data/table_export.rs | 25 +- crates/dbx-core/src/data/transfer.rs | 98 +++ crates/dbx-driver-mysql/src/doris.rs | 17 + crates/dbx-driver-mysql/src/mysql.rs | 584 +++++++++++++++++- .../src/data_grid_extractors/sql.rs | 13 + .../src/data_grid_extractors/tests.rs | 33 + crates/dbx-sql-data/src/data_grid_sql.rs | 115 ++++ crates/dbx-types/src/types.rs | 31 +- crates/dbx-web/src/routes/transfer.rs | 6 + src-tauri/src/commands/transfer.rs | 2 + 17 files changed, 1066 insertions(+), 42 deletions(-) diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue index 069931fe35..48603d7d22 100644 --- a/apps/desktop/src/components/grid/DataGrid.vue +++ b/apps/desktop/src/components/grid/DataGrid.vue @@ -166,8 +166,11 @@ import { downloadBinaryCellPayload, formatBinaryCellByteSize, binaryCellUtf8Text, + hasUnsafeOpaqueAggregateStatePredicate, isBlobCellColumnType, isBinaryCellColumnType, + isOpaqueAggregateStateColumnType, + mergeOpaqueReadonlyColumnIndexes, openBinaryCellFile, parseBinaryCellBytes, retainBinaryCellDownloadMenuForHover, @@ -3878,7 +3881,7 @@ const editor = useDataGridEditor({ tableMeta: computed(() => props.tableMeta), sourceColumns: computed(() => props.sourceColumns), joinedWriteTargets: computed(() => props.joinedWriteTargets), - readonlyColumnIndexes: computed(() => (props.readonlyColumnIndexes ? new Set(props.readonlyColumnIndexes) : undefined)), + readonlyColumnIndexes: computed(() => mergeOpaqueReadonlyColumnIndexes(props.readonlyColumnIndexes, allColumnTypes.value)), canEditExistingRows, onExecuteSql: computed(() => props.onExecuteSql), customSaveHandler: computed(() => props.customSaveHandler), @@ -4088,6 +4091,7 @@ function canEditRowItem(item: RowItem | undefined): boolean { function canEditCellItem(item: RowItem | undefined, columnIndex: number): boolean { if (!canEditRowItem(item) || !item || !canEditColumn(columnIndex)) return false; + if (isOpaqueAggregateStateColumnType(allColumnTypes.value[columnIndex])) return false; if (!item.isNew && !item.isDraft && !canUpdateExistingRows.value) return false; if (isSavingNewRow(item)) return false; const column = props.result.columns[columnIndex] ?? ""; @@ -4351,6 +4355,7 @@ function isDecimalColumnType(dataType: string): boolean { function canDeleteRowItem(item: RowItem | undefined): boolean { if (!item) return false; + if (!item.isNew && canUseKeylessRowPredicate(props.databaseType, props.tableMeta?.primaryKeys ?? []) && hasUnsafeOpaqueAggregateStatePredicate(allColumnTypes.value, item.data)) return false; const canDelete = canDeleteGridRowItem({ editable: !!props.editable && canDeleteRows.value, isDraft: !!item.isDraft, diff --git a/apps/desktop/src/components/grid/__tests__/DataGridSurfaces.spec.ts b/apps/desktop/src/components/grid/__tests__/DataGridSurfaces.spec.ts index 5d22c6a2c4..08273551fe 100644 --- a/apps/desktop/src/components/grid/__tests__/DataGridSurfaces.spec.ts +++ b/apps/desktop/src/components/grid/__tests__/DataGridSurfaces.spec.ts @@ -113,6 +113,14 @@ import DataGridSearchBar from "@/components/grid/DataGridSearchBar.vue"; const cellDetailPanelSource = readFileSync("apps/desktop/src/components/grid/DataGridCellDetailPanel.vue", "utf8"); const cellDetailDialogSource = readFileSync("apps/desktop/src/components/grid/DataGridCellDetailDialog.vue", "utf8"); const binaryTextPreviewSource = readFileSync("apps/desktop/src/components/grid/DataGridCellDetailTextPreview.vue", "utf8"); +const dataGridSource = readFileSync("apps/desktop/src/components/grid/DataGrid.vue", "utf8"); + +describe("DataGrid opaque aggregate-state safety", () => { + it("blocks cell edits and keyless row deletion through the centralized gates", () => { + expect(dataGridSource).toMatch(/function canEditCellItem[\s\S]*isOpaqueAggregateStateColumnType\(allColumnTypes\.value\[columnIndex\]\)/); + expect(dataGridSource).toMatch(/function canDeleteRowItem[\s\S]*canUseKeylessRowPredicate[\s\S]*hasUnsafeOpaqueAggregateStatePredicate\(allColumnTypes\.value, item\.data\)/); + }); +}); function detail(patch: Partial = {}): DataGridCellDetail { return { diff --git a/apps/desktop/src/composables/__tests__/useDataGridEditor.spec.ts b/apps/desktop/src/composables/__tests__/useDataGridEditor.spec.ts index b772d48e0b..9ed06bce50 100644 --- a/apps/desktop/src/composables/__tests__/useDataGridEditor.spec.ts +++ b/apps/desktop/src/composables/__tests__/useDataGridEditor.spec.ts @@ -51,11 +51,13 @@ function createEditorWithResult( mongoCollectionGrid = false, quickEntry = false, editable = ref(true), + resultColumnTypes?: string[], ) { let editor: ReturnType; - const result = ref<{ columns: string[]; rows: CellValue[][] }>({ + const result = ref<{ columns: string[]; rows: CellValue[][]; column_types?: string[] }>({ columns: ["first", "hidden", "last"], rows: existingRows, + column_types: resultColumnTypes, }); editor = useDataGridEditor({ @@ -734,12 +736,66 @@ describe("useDataGridEditor appendPastedRowsToNewRow", () => { ]); }); - it("keeps explicitly read-only mapped columns out of editing and paste", () => { - const editor = createEditor(["first", "hidden", "last"], true, undefined, [0]); + it("preserves ordinary read-only values when cloning while still blocking edits and paste", () => { + const editor = createEditor(["first", "hidden", "last"], true, undefined, [0], [["Ada", "hidden", "Lovelace"]]); expect(editor.canEditColumn(0)).toBe(false); expect(editor.canEditColumn(2)).toBe(true); expect(editor.appendPastedRowsToNewRow(-1, [["Ada"]], [0])).toEqual({ ok: false, reason: "readonly-column" }); + editor.applyCellValue(0, 0, null); + expect(editor.dirtyRows.value.size).toBe(0); + editor.cloneRow(0); + expect(editor.newRows.value.at(-1)).toEqual(["Ada", "hidden", "Lovelace"]); + }); + + it("clears opaque aggregate states when cloning and prefers the result type over metadata", () => { + const opaque = createEditor( + ["first", "hidden", "last"], + true, + undefined, + [1], + [["Ada", "0x00ff", "Lovelace"]], + undefined, + [ + { name: "first", data_type: "varchar" }, + { name: "hidden", data_type: "varchar" }, + { name: "last", data_type: "varchar" }, + ], + false, + false, + ref(true), + ["varchar", "agg_state", "varchar"], + ); + opaque.cloneRow(0); + expect(opaque.newRows.value.at(-1)).toEqual(["Ada", null, "Lovelace"]); + + const metadataFallback = createEditor(["first", "hidden", "last"], true, undefined, [1], [["Ada", "0x00ff", "Lovelace"]], undefined, [ + { name: "first", data_type: "varchar" }, + { name: "hidden", data_type: "agg_state" }, + { name: "last", data_type: "varchar" }, + ]); + metadataFallback.cloneRow(0); + expect(metadataFallback.newRows.value.at(-1)).toEqual(["Ada", null, "Lovelace"]); + + const resultTypeWins = createEditor( + ["first", "hidden", "last"], + true, + undefined, + [1], + [["Ada", "ordinary", "Lovelace"]], + undefined, + [ + { name: "first", data_type: "varchar" }, + { name: "hidden", data_type: "agg_state" }, + { name: "last", data_type: "varchar" }, + ], + false, + false, + ref(true), + ["varchar", "varchar", "varchar"], + ); + resultTypeWins.cloneRow(0); + expect(resultTypeWins.newRows.value.at(-1)).toEqual(["Ada", "ordinary", "Lovelace"]); }); it("fills following blank new rows before adding more rows", () => { diff --git a/apps/desktop/src/composables/useDataGridEditor.ts b/apps/desktop/src/composables/useDataGridEditor.ts index c64dba8306..1753399406 100644 --- a/apps/desktop/src/composables/useDataGridEditor.ts +++ b/apps/desktop/src/composables/useDataGridEditor.ts @@ -6,6 +6,7 @@ import type { CellValue } from "@/lib/dataGrid/cellValue"; import { coerceDataGridCellValue, dataGridCellEditorText } from "@/lib/dataGrid/dataGridCellCoercion"; import { focusDataGridEditorWithoutScrolling, preserveDataGridScrollPosition } from "@/lib/dataGrid/dataGridEditorFocus"; import { normalizeDataGridSaveError } from "@/lib/dataGrid/dataGridSql"; +import { isOpaqueAggregateStateColumnType } from "@/lib/dataGrid/binaryCellDownload"; import { rowStatusFilterAfterAddingRow, type RowStatusFilter } from "@/lib/dataGrid/gridRowStatus"; import type { GridNewRowMeta, GridNewRowPlacement } from "@/lib/dataGrid/gridNewRowPlacement"; import { supportsDataGridTransaction } from "@/lib/table/tableEditing"; @@ -107,7 +108,7 @@ export interface DataGridSaveConfirmationRequest { } export interface UseDataGridEditorOptions { - result: ComputedRef<{ columns: string[]; rows: CellValue[][] }>; + result: ComputedRef<{ columns: string[]; rows: CellValue[][]; column_types?: string[] }>; editable: ComputedRef; databaseType: ComputedRef; connectionId: ComputedRef; @@ -1331,7 +1332,10 @@ export function useDataGridEditor(options: UseDataGridEditorOptions) { return (columnIndex: number) => { const columnName = sourceColumns.value?.[columnIndex] ?? result.value.columns[columnIndex]; if (columnName === undefined) return false; - return shouldClearClonedColumn(columnName, columnInfoByName.get(columnName.toLowerCase())); + const columnInfo = columnInfoByName.get(columnName.toLowerCase()); + const resultType = result.value.column_types?.[columnIndex]?.trim(); + if (isOpaqueAggregateStateColumnType(resultType || columnInfo?.data_type)) return true; + return shouldClearClonedColumn(columnName, columnInfo); }; } diff --git a/apps/desktop/src/lib/dataGrid/__tests__/binaryCellFileImport.spec.ts b/apps/desktop/src/lib/dataGrid/__tests__/binaryCellFileImport.spec.ts index f76c98684b..85001b5af4 100644 --- a/apps/desktop/src/lib/dataGrid/__tests__/binaryCellFileImport.spec.ts +++ b/apps/desktop/src/lib/dataGrid/__tests__/binaryCellFileImport.spec.ts @@ -1,5 +1,15 @@ import { describe, expect, it } from "vitest"; -import { binaryCellBytesToHexValue, canImportBinaryCellFile } from "@/lib/dataGrid/binaryCellDownload"; +import { + binaryCellBytesToHexValue, + binaryCellClipboardText, + binaryCellDisplayText, + binaryCellDownloadPayload, + canImportBinaryCellFile, + hasUnsafeOpaqueAggregateStatePredicate, + isBinaryCellColumnType, + isOpaqueAggregateStateColumnType, + mergeOpaqueReadonlyColumnIndexes, +} from "@/lib/dataGrid/binaryCellDownload"; describe("binary cell file import", () => { it("encodes arbitrary and empty files as prefixed hex cell values", () => { @@ -14,5 +24,34 @@ describe("binary cell file import", () => { expect(canImportBinaryCellFile("postgres", "text")).toBe(false); expect(canImportBinaryCellFile("sqlserver", "varbinary(max)")).toBe(false); expect(canImportBinaryCellFile("sqlite", "blob")).toBe(false); + expect(canImportBinaryCellFile("mysql", "agg_state")).toBe(false); + }); + + it("presents aggregate states as opaque byte-preserving binary values", () => { + const type = "agg_state"; + const value = "0x0001abff"; + expect(isOpaqueAggregateStateColumnType(type)).toBe(true); + expect(isOpaqueAggregateStateColumnType(" AGG_STATE < sum(int) > ")).toBe(true); + expect(isOpaqueAggregateStateColumnType("agg_state<")).toBe(false); + expect(isOpaqueAggregateStateColumnType("agg_state<>")).toBe(false); + expect(isOpaqueAggregateStateColumnType("😺agg_state")).toBe(false); + expect(isBinaryCellColumnType(type)).toBe(true); + expect(binaryCellDisplayText(value, type, undefined, "mysql")).toBe("AGG_STATE [4 bytes]"); + expect(binaryCellClipboardText(value, type, "mysql")).toBeNull(); + expect(binaryCellDownloadPayload(value, "binary", type, "mysql").data).toEqual(new Uint8Array([0x00, 0x01, 0xab, 0xff])); + }); + + it("merges opaque columns into the centralized readonly indexes", () => { + const readonly = mergeOpaqueReadonlyColumnIndexes([0, 3], ["varchar", "agg_state", "varbinary", "text"]); + expect([...readonly!]).toEqual([0, 3, 1]); + expect(mergeOpaqueReadonlyColumnIndexes(undefined, ["varchar", "varbinary"])).toBeUndefined(); + }); + + it("only blocks keyless predicates for non-null or missing opaque values", () => { + const types = ["varchar", "agg_state"]; + expect(hasUnsafeOpaqueAggregateStatePredicate(types, ["Ada", null])).toBe(false); + expect(hasUnsafeOpaqueAggregateStatePredicate(types, ["Ada", "0x00ff"])).toBe(true); + expect(hasUnsafeOpaqueAggregateStatePredicate(types, ["Ada"])).toBe(true); + expect(hasUnsafeOpaqueAggregateStatePredicate(["varchar", "text"], ["Ada", "ordinary"])).toBe(false); }); }); diff --git a/apps/desktop/src/lib/dataGrid/binaryCellDownload.ts b/apps/desktop/src/lib/dataGrid/binaryCellDownload.ts index f747463cca..cfa160bbba 100644 --- a/apps/desktop/src/lib/dataGrid/binaryCellDownload.ts +++ b/apps/desktop/src/lib/dataGrid/binaryCellDownload.ts @@ -117,6 +117,7 @@ const BINARY_STRING_TYPE_RE = /^(?:binary|varbinary)(?:\b|\()/i; const VARBINARY_TYPE_RE = /^varbinary(?:\b|\()/i; const BLOB_TYPE_RE = /^(?:blob|tinyblob|mediumblob|longblob)(?:\b|\()/i; const MYSQL_FILE_IMPORT_TYPE_RE = /^(?:blob|tinyblob|mediumblob|longblob|binary|varbinary)(?:\b|\()/i; +const OPAQUE_AGG_STATE_TYPE_RE = /^agg_state\s*<\s*[^>\s][^>]*>$/i; function copyBytesForBlob(bytes: Uint8Array): Uint8Array { return new Uint8Array(bytes); @@ -181,7 +182,23 @@ export function parseBinaryCellBytes(value: unknown, columnType?: string, databa export function isBinaryCellColumnType(columnType?: string): boolean { const type = (columnType ?? "").trim(); - return !!type && BINARY_TYPE_RE.test(type); + return !!type && (BINARY_TYPE_RE.test(type) || OPAQUE_AGG_STATE_TYPE_RE.test(type)); +} + +export function isOpaqueAggregateStateColumnType(columnType?: string): boolean { + return OPAQUE_AGG_STATE_TYPE_RE.test((columnType ?? "").trim()); +} + +export function hasUnsafeOpaqueAggregateStatePredicate(columnTypes: readonly (string | undefined)[], row: readonly CellValue[]): boolean { + return columnTypes.some((columnType, index) => isOpaqueAggregateStateColumnType(columnType) && row[index] !== null); +} + +export function mergeOpaqueReadonlyColumnIndexes(readonlyColumnIndexes: readonly number[] | undefined, columnTypes: readonly (string | undefined)[]): ReadonlySet | undefined { + const merged = new Set(readonlyColumnIndexes ?? []); + columnTypes.forEach((columnType, index) => { + if (isOpaqueAggregateStateColumnType(columnType)) merged.add(index); + }); + return merged.size ? merged : undefined; } export function isBlobCellColumnType(columnType?: string): boolean { @@ -294,6 +311,7 @@ export function binaryCellClipboardText(value: unknown, columnType?: string, dat // “显示是文本、编辑器却是十六进制”的不一致，故 blob 文本预览仅在 mysql 连接开启。 // binary/varbinary 的文本预览早于该特性存在（如 TDengine BINARY 文本），保持全库通用。 function isBinaryCellTextPreviewColumn(columnType: string | undefined, databaseType: DatabaseType | undefined): boolean { + if (isOpaqueAggregateStateColumnType(columnType)) return false; if (BINARY_STRING_TYPE_RE.test((columnType ?? "").trim())) return true; return isBlobCellColumnType(columnType) && databaseType === "mysql"; } @@ -367,6 +385,7 @@ function printableText(text: string): boolean { } function binaryCellDisplayLabel(columnType?: string): string { + if (isOpaqueAggregateStateColumnType(columnType)) return "AGG_STATE"; const base = (columnType ?? "") .trim() .split(/[(:\s]/)[0] diff --git a/crates/dbx-core/src/data/database_export.rs b/crates/dbx-core/src/data/database_export.rs index 69591d6ba1..cae389dc22 100644 --- a/crates/dbx-core/src/data/database_export.rs +++ b/crates/dbx-core/src/data/database_export.rs @@ -1241,7 +1241,6 @@ pub(crate) fn build_export_insert_statements_excluding_with_dialect( if options.columns.is_empty() || options.rows.is_empty() { return Ok(Vec::new()); } - let excluded_names: HashSet = exclude_columns.iter().map(|column| column.trim().to_ascii_uppercase()).collect(); let table = export_qualified_table_name( @@ -1274,6 +1273,18 @@ pub(crate) fn build_export_insert_statements_excluding_with_dialect( }) }) .collect::>(); + if insert_columns.iter().any(|(index, _, _)| { + options + .column_types + .get(*index) + .and_then(|column_type| column_type.as_deref()) + .is_some_and(crate::types::is_opaque_aggregate_state_type) + }) { + return Err( + "SQL INSERT export does not support Doris aggregate-state columns; use a representation export for canonical hex bytes" + .to_string(), + ); + } if insert_columns.is_empty() { // Only fail when the exclusion itself removed the last insertable column; // emptiness caused by other omission rules (e.g. generated columns) keeps @@ -6125,6 +6136,30 @@ mod tests { ); } + #[test] + fn sql_insert_export_allows_explicitly_excluded_opaque_column() { + let statements = build_export_insert_statements_excluding( + BuildExportInsertStatementsOptions { + database_type: Some(DatabaseType::Doris), + identifier_quote: None, + schema: None, + table_name: Some("states".to_string()), + qualified_table_name: None, + columns: vec!["id".to_string(), "v2".to_string()], + column_types: vec![Some("integer".to_string()), Some("agg_state".to_string())], + column_extras: Vec::new(), + spatial_columns: Vec::new(), + spatial_values: Vec::new(), + rows: vec![vec![json!(1), json!("0x00ff")]], + batch_size: Some(10), + }, + &["v2".to_string()], + ) + .unwrap(); + + assert_eq!(statements, vec!["INSERT INTO `states` (`id`) VALUES (1);"]); + } + #[test] fn sql_insert_export_reports_error_when_excluding_leaves_no_columns() { let result = build_export_insert_statements_excluding( diff --git a/crates/dbx-core/src/data/table_export.rs b/crates/dbx-core/src/data/table_export.rs index d50c3ec2e0..9675b520e8 100644 --- a/crates/dbx-core/src/data/table_export.rs +++ b/crates/dbx-core/src/data/table_export.rs @@ -33,7 +33,7 @@ use crate::transfer::{ count_sql_with_where_and_identifier_quote, execute_read_on_pool, execute_read_on_pool_with_max_rows, keyset_pagination_sql_with_identifier_quote, pagination_sql_with_filter_order_and_identifier_quote, }; -use crate::types::QueryResult; +use crate::types::{is_opaque_aggregate_state_type, QueryResult}; use crate::xlsx_export::{finish_streaming_xlsx_workbook, start_streaming_xlsx_workbook_with_options}; const DEFAULT_BATCH_SIZE: usize = 10_000; @@ -237,6 +237,18 @@ fn requested_mysql_sql_export_needs_column_metadata(database_type: DatabaseType, database_type == DatabaseType::Mysql && format.eq_ignore_ascii_case("sql") } +fn ensure_sql_insert_export_types_supported(format: &str, column_types: &[Option]) -> Result<(), String> { + if format.eq_ignore_ascii_case("sql") + && column_types.iter().flatten().any(|column_type| is_opaque_aggregate_state_type(column_type)) + { + return Err( + "SQL INSERT export does not support Doris aggregate-state columns; use CSV, JSON, TXT, or XLSX to export their canonical hex representation" + .to_string(), + ); + } + Ok(()) +} + /// Column EXTRA values supplied by the data grid, filtered and aligned exactly /// like [`resolve_requested_export_columns`] keeps name/type metadata aligned. fn resolve_requested_column_extras_by_position( @@ -1597,6 +1609,7 @@ async fn export_table_data_core_inner( if col_names.is_empty() { return Err("No columns found for table".to_string()); } + ensure_sql_insert_export_types_supported(&request.format, &column_types)?; let query_col_names = table_export_query_columns(request, &sql_context, &col_names)?; // Use keyset pagination when all PKs are in the selected (filtered) columns. @@ -2334,6 +2347,16 @@ mod tests { assert_eq!(serde_json::to_value(request).unwrap()["insertDialect"], json!("standard")); } + #[test] + fn sql_insert_export_rejects_opaque_aggregate_states_only() { + let types = vec![Some("int".to_string()), Some("agg_state".to_string())]; + assert!(ensure_sql_insert_export_types_supported("sql", &types) + .unwrap_err() + .contains("does not support Doris aggregate-state")); + assert!(ensure_sql_insert_export_types_supported("csv", &types).is_ok()); + assert!(ensure_sql_insert_export_types_supported("sql", &[Some("varbinary(32)".to_string())]).is_ok()); + } + #[cfg(unix)] struct ExternalDriverExportFixture { state: AppState, diff --git a/crates/dbx-core/src/data/transfer.rs b/crates/dbx-core/src/data/transfer.rs index c22fdb54fd..2bc021a196 100644 --- a/crates/dbx-core/src/data/transfer.rs +++ b/crates/dbx-core/src/data/transfer.rs @@ -43,6 +43,58 @@ static OCEANBASE_MYSQL_TABLE_OPTION_RE: std::sync::LazyLock = std::sync:: static MYSQL_COLLATE_CLAUSE_RE: std::sync::LazyLock = std::sync::LazyLock::new(|| { Regex::new(r"(?i)\bCOLLATE\s*=?\s*([A-Za-z0-9_]+)\b").expect("valid MySQL COLLATE clause regex") }); + +pub async fn ensure_transfer_source_types_supported( + state: &AppState, + request: &TransferRequest, + source_pool_key: &str, +) -> Result<(), String> { + if matches!(request.content, TransferContent::StructureOnly) { + return Ok(()); + } + let is_doris = { + let configs = state.configs.read().await; + configs + .get(&request.source_connection_id) + .is_some_and(|config| db::doris::is_native_profile(&config.db_type, config.driver_profile.as_deref())) + }; + if !is_doris { + return Ok(()); + } + for table in &request.tables { + let columns = get_columns_for_transfer( + state, + source_pool_key, + &request.source_connection_id, + &request.source_database, + &request.source_schema, + table, + request.source_catalog.as_deref(), + ) + .await?; + ensure_transfer_columns_supported(request, true, table, &columns)?; + } + Ok(()) +} + +fn ensure_transfer_columns_supported( + request: &TransferRequest, + is_doris_source: bool, + table: &str, + columns: &[db::ColumnInfo], +) -> Result<(), String> { + if matches!(request.content, TransferContent::StructureOnly) || !is_doris_source { + return Ok(()); + } + if let Some(column) = columns.iter().find(|column| crate::types::is_opaque_aggregate_state_type(&column.data_type)) + { + return Err(format!( + "Data transfer does not support Doris aggregate-state column `{}` in table `{table}`; export a representation format or use explicit Doris state functions instead", + column.name + )); + } + Ok(()) +} // An inline FK constraint *definition line*: an optional `CONSTRAINT ` // prefix followed by `FOREIGN KEY (`. Anchored to the line start so column // definitions whose COMMENT/DEFAULT strings mention "foreign key" never match. @@ -8840,6 +8892,14 @@ pub async fn rename_tables_to_backup( where F: FnMut(TransferProgress), { + let source_pool_key = ensure_transfer_pool( + state, + &request.source_connection_id, + &request.source_database, + request.source_catalog.as_deref(), + ) + .await?; + ensure_transfer_source_types_supported(state, request, &source_pool_key).await?; let total_tables = tables.len(); // Resolve target names first so the fail-fast check below sees the names that will @@ -9419,6 +9479,14 @@ where return Err(format!("No columns found for table {table}")); } + let is_doris_source = { + let configs = state.configs.read().await; + configs + .get(&request.source_connection_id) + .is_some_and(|config| db::doris::is_native_profile(&config.db_type, config.driver_profile.as_deref())) + }; + ensure_transfer_columns_supported(request, is_doris_source, table, &columns)?; + let writable_columns = writable_transfer_columns(&columns, source_db_type, target_db_type); let default_rows_only = mysql_generated_only_transfer(&columns, source_db_type, target_db_type); if writable_columns.is_empty() && !default_rows_only { @@ -12621,6 +12689,36 @@ CREATE TABLE "Other"."prefix""Source"."NAME" ("ID" INT);"#; } } + #[test] + fn transfer_column_validation_rejects_opaque_state_for_actual_doris_table() { + let request = test_transfer_request(vec!["first_state", "actual_state"]); + let columns = vec![test_column("id", "int"), test_column("v2", "agg_state")]; + + let error = ensure_transfer_columns_supported(&request, true, "actual_state", &columns).unwrap_err(); + + assert!(error.contains("`v2`"), "{error}"); + assert!(error.contains("`actual_state`"), "{error}"); + assert!(!error.contains("first_state"), "{error}"); + } + + #[test] + fn transfer_column_validation_is_doris_data_only_and_structure_aware() { + let columns = vec![test_column("v2", "AGG_STATE")]; + let data_request = test_transfer_request(vec!["states"]); + assert!(ensure_transfer_columns_supported(&data_request, false, "states", &columns).is_ok()); + + let structure_request = TransferRequest { content: TransferContent::StructureOnly, ..data_request }; + assert!(ensure_transfer_columns_supported(&structure_request, true, "states", &columns).is_ok()); + } + + #[test] + fn transfer_column_validation_allows_regular_doris_columns() { + let request = test_transfer_request(vec!["items"]); + let columns = vec![test_column("id", "bigint"), test_column("name", "varchar(64)")]; + + assert!(ensure_transfer_columns_supported(&request, true, "items", &columns).is_ok()); + } + #[test] fn transfer_request_defaults_preserve_table_name_case() { let request: TransferRequest = serde_json::from_value(json!({ diff --git a/crates/dbx-driver-mysql/src/doris.rs b/crates/dbx-driver-mysql/src/doris.rs index 0be6dc1ba1..be83b15f61 100644 --- a/crates/dbx-driver-mysql/src/doris.rs +++ b/crates/dbx-driver-mysql/src/doris.rs @@ -17,6 +17,14 @@ pub fn is_profile(db_type: &DatabaseType, driver_profile: Option<&str>) -> bool .is_some_and(|profile| matches!(profile.to_ascii_lowercase().as_str(), "doris" | "selectdb"))) } +/// Native Doris semantics only. SelectDB shares metadata/catalog handling but +/// must not inherit Doris-specific query preflight or opaque-state safeguards. +pub fn is_native_profile(db_type: &DatabaseType, driver_profile: Option<&str>) -> bool { + *db_type == DatabaseType::Doris + || (*db_type == DatabaseType::Mysql + && driver_profile.is_some_and(|profile| profile.eq_ignore_ascii_case("doris"))) +} + #[cfg(test)] mod tests { use super::*; @@ -29,4 +37,13 @@ mod tests { assert!(!is_profile(&DatabaseType::Mysql, Some("starrocks"))); assert!(!is_profile(&DatabaseType::Postgres, Some("doris"))); } + + #[test] + fn native_profile_excludes_selectdb_and_other_mysql_profiles() { + assert!(is_native_profile(&DatabaseType::Doris, None)); + assert!(is_native_profile(&DatabaseType::Mysql, Some("DORIS"))); + assert!(!is_native_profile(&DatabaseType::Mysql, Some("selectdb"))); + assert!(!is_native_profile(&DatabaseType::Mysql, Some("starrocks"))); + assert!(!is_native_profile(&DatabaseType::Mysql, None)); + } } diff --git a/crates/dbx-driver-mysql/src/mysql.rs b/crates/dbx-driver-mysql/src/mysql.rs index 5949fe4226..183d17acff 100644 --- a/crates/dbx-driver-mysql/src/mysql.rs +++ b/crates/dbx-driver-mysql/src/mysql.rs @@ -4,13 +4,14 @@ use mysql_async::consts::{ColumnFlags, ColumnType, StatusFlags}; use mysql_async::prelude::*; use percent_encoding::percent_decode_str; use rust_decimal::Decimal; -use sqlparser::ast::{AlterTableOperation, ObjectNamePart, Statement}; +use sqlparser::ast::{AlterTableOperation, Expr, ObjectNamePart, Statement, TableFactor, Visit, Visitor}; use sqlparser::dialect::MySqlDialect; use sqlparser::parser::Parser; use std::borrow::Cow; use std::collections::{HashMap, HashSet}; #[cfg(test)] use std::future::Future; +use std::ops::ControlFlow; use std::ops::Deref; use std::path::PathBuf; use std::sync::atomic::AtomicBool; @@ -24,10 +25,10 @@ use crate::models::connection::{ }; use crate::sql::{starts_with_executable_sql_keyword, starts_with_executable_sql_keyword_for_database}; use crate::types::{ - ColumnInfo, ColumnMetadataCapabilities, CompletionAssistantCandidate, CompletionAssistantCandidateKind, - CompletionAssistantMatchMode, CompletionAssistantObjectKind, CompletionAssistantRequest, - CompletionAssistantResponse, DatabaseInfo, ForeignKeyInfo, IndexInfo, LargeValueCell, ObjectInfo, ObjectStatistics, - QueryMessage, QueryResult, SpatialColumnBuilder, TableInfo, TriggerInfo, + is_opaque_aggregate_state_type, ColumnInfo, ColumnMetadataCapabilities, CompletionAssistantCandidate, + CompletionAssistantCandidateKind, CompletionAssistantMatchMode, CompletionAssistantObjectKind, + CompletionAssistantRequest, CompletionAssistantResponse, DatabaseInfo, ForeignKeyInfo, IndexInfo, LargeValueCell, + ObjectInfo, ObjectStatistics, QueryMessage, QueryResult, SpatialColumnBuilder, TableInfo, TriggerInfo, }; use dbx_types::metadata_filter::{table_name_filter_matches, TableNameFilter}; @@ -278,6 +279,7 @@ pub fn mysql_catalog_dialect(db_type: DatabaseType, driver_profile: Option<&str> #[derive(Clone, Copy, Debug, Default)] pub struct MySqlQueryDialect { supports_admin_show_results: bool, + is_doris: bool, } #[derive(Debug)] @@ -352,10 +354,231 @@ impl MySqlQueryDialect { || super::starrocks::is_profile(&db_type, driver_profile) || super::manticoresearch::is_profile(&db_type, driver_profile) || super::tidb::is_profile(&db_type, driver_profile), + is_doris: super::doris::is_native_profile(&db_type, driver_profile), } } } +type DorisOpaqueColumnTypes = HashMap<(Vec, Vec, Vec), String>; +type DorisOpaqueColumnKey = (Vec, Vec, Vec); + +fn insert_doris_opaque_column_type( + resolved: &mut DorisOpaqueColumnTypes, + ambiguous: &mut HashSet, + key: DorisOpaqueColumnKey, + data_type: String, +) { + if ambiguous.contains(&key) || resolved.insert(key.clone(), data_type).is_some() { + resolved.remove(&key); + ambiguous.insert(key); + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct DorisSourceTable { + database: Option, + table: String, +} + +#[derive(Default)] +struct DorisPreflightVisitor { + cte_names: HashSet, + sources: HashSet, + unsupported: bool, +} + +impl Visitor for DorisPreflightVisitor { + type Break = (); + + fn pre_visit_query(&mut self, query: &sqlparser::ast::Query) -> ControlFlow<()> { + if let Some(with) = &query.with { + for cte in &with.cte_tables { + if !self.cte_names.insert(cte.alias.name.value.to_ascii_lowercase()) { + self.unsupported = true; + return ControlFlow::Break(()); + } + } + } + ControlFlow::Continue(()) + } + + fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<()> { + if let Expr::Function(function) = expr { + if function + .name + .0 + .last() + .and_then(ObjectNamePart::as_ident) + .is_some_and(|name| name.value.eq_ignore_ascii_case("LAST_QUERY_ID")) + { + self.unsupported = true; + return ControlFlow::Break(()); + } + } + ControlFlow::Continue(()) + } + + fn pre_visit_table_factor(&mut self, factor: &TableFactor) -> ControlFlow<()> { + let TableFactor::Table { name, args: None, .. } = factor else { + if !matches!(factor, TableFactor::Derived { .. } | TableFactor::NestedJoin { .. }) { + self.unsupported = true; + return ControlFlow::Break(()); + } + return ControlFlow::Continue(()); + }; + let parts = name.0.iter().map(ObjectNamePart::as_ident).collect::>>(); + let Some(parts) = parts else { + self.unsupported = true; + return ControlFlow::Break(()); + }; + match parts.as_slice() { + [table] if !self.cte_names.contains(&table.value.to_ascii_lowercase()) => { + self.sources.insert(DorisSourceTable { database: None, table: table.value.clone() }); + } + [database, table] => { + self.sources + .insert(DorisSourceTable { database: Some(database.value.clone()), table: table.value.clone() }); + } + [_cte] => {} + _ => { + self.unsupported = true; + return ControlFlow::Break(()); + } + } + ControlFlow::Continue(()) + } +} + +fn doris_preflight_sources(sql: &str, dialect: MySqlQueryDialect) -> Option> { + if !dialect.is_doris { + return None; + } + let statements = Parser::parse_sql(&MySqlDialect {}, sql).ok()?; + let [Statement::Query(query)] = statements.as_slice() else { + return None; + }; + let mut visitor = DorisPreflightVisitor::default(); + let _ = query.visit(&mut visitor); + if visitor.unsupported || visitor.sources.is_empty() { + return None; + } + Some(visitor.sources.into_iter().collect()) +} + +async fn resolve_doris_opaque_column_types( + conn: &mut mysql_async::Conn, + sql: &str, + dialect: MySqlQueryDialect, +) -> DorisOpaqueColumnTypes { + let Some(mut sources) = doris_preflight_sources(sql, dialect) else { + return HashMap::new(); + }; + if sources.iter().any(|source| source.database.is_none()) { + let current_database = conn + .query_first::("SELECT DATABASE()") + .await + .ok() + .flatten() + .filter(|database| !database.is_empty()); + let Some(current_database) = current_database else { + return HashMap::new(); + }; + for source in &mut sources { + if source.database.is_none() { + source.database = Some(current_database.clone()); + } + } + } + sources.sort_by(|left, right| (&left.database, &left.table).cmp(&(&right.database, &right.table))); + sources.dedup(); + + let mut resolved = HashMap::new(); + let mut ambiguous = HashSet::new(); + for source in sources { + let Some(database) = source.database else { continue }; + match fetch_columns_show(conn, &database, &source.table).await { + Ok(columns) => { + for column in columns { + if is_opaque_aggregate_state_type(&column.data_type) { + let key = ( + database.as_bytes().to_vec(), + source.table.as_bytes().to_vec(), + column.name.as_bytes().to_vec(), + ); + insert_doris_opaque_column_type(&mut resolved, &mut ambiguous, key, column.data_type); + } + } + } + Err(error) => log::debug!( + "Failed to resolve Doris aggregate-state columns for `{database}`.`{}`: {error}", + source.table + ), + } + } + resolved +} + +fn effective_mysql_column_types(columns: &[mysql_async::Column], opaque_types: &DorisOpaqueColumnTypes) -> Vec { + columns + .iter() + .map(|column| { + opaque_types + .get(&(column.schema_ref().to_vec(), column.org_table_ref().to_vec(), column.org_name_ref().to_vec())) + .filter(|_| { + !column.schema_ref().is_empty() + && !column.org_table_ref().is_empty() + && !column.org_name_ref().is_empty() + && matches!( + column.column_type(), + ColumnType::MYSQL_TYPE_STRING + | ColumnType::MYSQL_TYPE_VAR_STRING + | ColumnType::MYSQL_TYPE_VARCHAR + | ColumnType::MYSQL_TYPE_BLOB + | ColumnType::MYSQL_TYPE_LONG_BLOB + | ColumnType::MYSQL_TYPE_MEDIUM_BLOB + | ColumnType::MYSQL_TYPE_TINY_BLOB + ) + }) + .cloned() + .unwrap_or_else(|| mysql_column_type_name(column)) + }) + .collect() +} + +fn mysql_value_to_json_with_effective_type( + row: &mysql_async::Row, + index: usize, + effective_type: Option<&str>, +) -> serde_json::Value { + if effective_type.is_some_and(is_opaque_aggregate_state_type) { + return match row.as_ref(index) { + None | Some(mysql_async::Value::NULL) => serde_json::Value::Null, + Some(mysql_async::Value::Bytes(bytes)) => super::binary_value_to_json(bytes), + _ => mysql_value_to_json(row, index), + }; + } + mysql_value_to_json(row, index) +} + +fn ensure_doris_opaque_cells_fit_budget( + row: &mysql_async::Row, + effective_types: &[String], + max_result_bytes: Option, +) -> Result<(), String> { + let Some(max_result_bytes) = max_result_bytes else { return Ok(()) }; + for (index, effective_type) in effective_types.iter().enumerate() { + if is_opaque_aggregate_state_type(effective_type) + && matches!(row.as_ref(index), Some(mysql_async::Value::Bytes(bytes)) if bytes.len().saturating_mul(2).saturating_add(2) > max_result_bytes) + { + return Err(format!( + "Doris aggregate-state value in column {} exceeds the interactive byte budget; use a streaming representation export to preserve the complete bytes", + index + 1 + )); + } + } + Ok(()) +} + pub enum MySqlQueryStreamItem { Columns { columns: Vec, column_types: Vec }, Row(Vec), @@ -989,6 +1212,7 @@ fn mysql_result_protected_column_indexes( fn mysql_row_to_json_with_srids( row: &mysql_async::Row, spatial_columns: &mut SpatialColumnBuilder, + effective_types: &[String], ) -> (Vec, Vec>) { let mut srids = vec![None; row.len()]; let values = (0..row.len()) @@ -998,7 +1222,7 @@ fn mysql_row_to_json_with_srids( .get(idx) .is_some_and(|column| column.column_type() == ColumnType::MYSQL_TYPE_GEOMETRY); if !is_geometry { - return mysql_value_to_json(row, idx); + return mysql_value_to_json_with_effective_type(row, idx, effective_types.get(idx).map(String::as_str)); } let Some(bytes) = row_get::, _>(row, idx) else { spatial_columns.observe(idx, None); @@ -1026,6 +1250,7 @@ fn mysql_row_to_json_with_srids_and_previews( row_index: usize, preview_bytes: Option, protected_indexes: &HashSet, + effective_types: &[String], ) -> (Vec, Vec>, Vec) { let mut srids = vec![None; row.len()]; let mut large_value_cells = Vec::new(); @@ -1053,13 +1278,20 @@ fn mysql_row_to_json_with_srids_and_previews( }; } + if effective_types.get(index).is_some_and(|column_type| is_opaque_aggregate_state_type(column_type)) { + return mysql_value_to_json_with_effective_type( + row, + index, + effective_types.get(index).map(String::as_str), + ); + } if let Some((preview, original_bytes)) = preview_bytes.and_then(|limit| mysql_bounded_value_preview(row, index, limit, protected_indexes)) { large_value_cells.push(LargeValueCell { row_index, column_index: index, original_bytes }); return preview; } - mysql_value_to_json(row, index) + mysql_value_to_json_with_effective_type(row, index, effective_types.get(index).map(String::as_str)) }) .collect(); (values, srids, large_value_cells) @@ -4533,6 +4765,25 @@ fn mysql_column_name(raw: String) -> Option { } } +fn needs_mysql_column_type_enrichment(columns: &[ColumnInfo]) -> bool { + columns.iter().any(|column| column.data_type.eq_ignore_ascii_case("unknown")) +} + +fn enrich_unknown_mysql_column_types(columns: &mut [ColumnInfo], show_columns: &[ColumnInfo]) { + for column in columns { + if !column.data_type.eq_ignore_ascii_case("unknown") { + continue; + } + let Some(show_column) = show_columns.iter().find(|show_column| show_column.name == column.name) else { + continue; + }; + let show_type = show_column.data_type.trim(); + if !show_type.is_empty() && !show_type.eq_ignore_ascii_case("unknown") { + column.data_type.clone_from(&show_column.data_type); + } + } +} + fn parse_mysql_enum_values(column_type: &str) -> Option> { let trimmed = column_type.trim(); if !trimmed.get(..5)?.eq_ignore_ascii_case("enum(") || !trimmed.ends_with(')') { @@ -4653,6 +4904,17 @@ where return get_columns_show(pool, database, table).await; } + if needs_mysql_column_type_enrichment(&columns) { + match fetch_columns_show(&mut conn, database, table).await { + Ok(show_columns) => enrich_unknown_mysql_column_types(&mut columns, &show_columns), + Err(error) => { + log::debug!( + "Failed to enrich unknown MySQL column types for `{database}`.`{table}` from SHOW COLUMNS: {error}" + ); + } + } + } + enrich_mysql_generated_column_expressions(&mut conn, database, table, &mut columns).await; Ok(columns) } @@ -4661,8 +4923,18 @@ pub async fn get_columns_show

(pool: &P, database: &str, table: &str) -> Resul where P: MySqlPoolAccess + ?Sized, { - let sql = show_columns_sql(database, table, true); let mut conn = get_conn_with_health_check(pool).await?; + let mut columns = fetch_columns_show(&mut conn, database, table).await?; + enrich_mysql_generated_column_expressions(&mut conn, database, table, &mut columns).await; + Ok(columns) +} + +async fn fetch_columns_show( + conn: &mut mysql_async::Conn, + database: &str, + table: &str, +) -> Result, String> { + let sql = show_columns_sql(database, table, true); let rows: Vec = match conn.query_iter(&sql).await { Ok(result) => result.collect_and_drop().await.map_err(|e| e.to_string())?, Err(_) => { @@ -4671,7 +4943,7 @@ where result.collect_and_drop().await.map_err(|e| e.to_string())? } }; - let mut columns: Vec = rows + let columns: Vec = rows .iter() .filter_map(|row| { let name = mysql_column_name(get_str_by_name(row, "Field"))?; @@ -4702,7 +4974,6 @@ where }) }) .collect(); - enrich_mysql_generated_column_expressions(&mut conn, database, table, &mut columns).await; Ok(columns) } @@ -5100,6 +5371,7 @@ async fn execute_result_set_with_text_protocol_on_conn( diagnostic_trace_id: Option<&str>, start: Instant, progress_clock: Option<&crate::execution::StreamProgressClock>, + opaque_types: &DorisOpaqueColumnTypes, ) -> Result { let diagnostics_enabled = diagnostic_trace_id.is_some() && log::log_enabled!(log::Level::Info); let dispatch_started_at = diagnostics_enabled.then(Instant::now); @@ -5130,7 +5402,7 @@ async fn execute_result_set_with_text_protocol_on_conn( })); } let columns: Vec = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect(); - let column_types: Vec = result.columns_ref().iter().map(mysql_column_type_name).collect(); + let column_types = effective_mysql_column_types(result.columns_ref(), opaque_types); let mut spatial_columns = mysql_spatial_column_builder(result.columns_ref()); let preview_bytes = max_result_bytes.map(|max_bytes| mysql_result_cell_preview_bytes(max_bytes, row_limit, result.columns_ref())); @@ -5149,7 +5421,7 @@ async fn execute_result_set_with_text_protocol_on_conn( .iter() .take(row_limit) .map(|row| { - let (values, srids) = mysql_row_to_json_with_srids(row, &mut spatial_columns); + let (values, srids) = mysql_row_to_json_with_srids(row, &mut spatial_columns, &column_types); spatial_values.push(srids); values }) @@ -5195,6 +5467,7 @@ async fn execute_result_set_with_text_protocol_on_conn( } let Some(row) = next_row else { break }; let row = row.map_err(|e| e.to_string())?; + ensure_doris_opaque_cells_fit_budget(&row, &column_types, max_result_bytes)?; if let Some(progress_clock) = progress_clock { progress_clock.mark(); } @@ -5210,6 +5483,7 @@ async fn execute_result_set_with_text_protocol_on_conn( row_index, preview_bytes, &protected_indexes, + &column_types, ); result_rows.push(values); spatial_values.push(srids); @@ -5281,6 +5555,7 @@ async fn execute_result_sets_with_text_protocol_on_conn( result_key_columns: &[String], diagnostic_trace_id: Option<&str>, start: Instant, + opaque_types: &DorisOpaqueColumnTypes, ) -> Result, String> { // Per-set warnings/info read after each `collect()` are only accurate when // the connection negotiated CLIENT_DEPRECATE_EOF: with legacy EOF packets @@ -5304,7 +5579,7 @@ async fn execute_result_sets_with_text_protocol_on_conn( while advance_to_result_set_with_columns(&mut result).await? { let columns: Vec = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect(); - let column_types: Vec = result.columns_ref().iter().map(mysql_column_type_name).collect(); + let column_types = effective_mysql_column_types(result.columns_ref(), opaque_types); let mut spatial_columns = mysql_spatial_column_builder(result.columns_ref()); let mut spatial_values = Vec::new(); let mut large_value_cells = Vec::new(); @@ -5325,7 +5600,7 @@ async fn execute_result_sets_with_text_protocol_on_conn( .iter() .take(row_limit) .map(|row| { - let (values, srids) = mysql_row_to_json_with_srids(row, &mut spatial_columns); + let (values, srids) = mysql_row_to_json_with_srids(row, &mut spatial_columns, &column_types); spatial_values.push(srids); values }) @@ -5351,6 +5626,7 @@ async fn execute_result_sets_with_text_protocol_on_conn( let Some(row) = next_row else { break }; let row = row.map_err(|e| e.to_string())?; if rows.len() < row_limit { + ensure_doris_opaque_cells_fit_budget(&row, &column_types, max_result_bytes)?; let row_index = rows.len(); let convert_started_at = diagnostics_enabled.then(Instant::now); let (values, srids, mut row_large_values) = mysql_row_to_json_with_srids_and_previews( @@ -5359,6 +5635,7 @@ async fn execute_result_sets_with_text_protocol_on_conn( row_index, preview_bytes, &protected_indexes, + &column_types, ); rows.push(values); spatial_values.push(srids); @@ -5506,13 +5783,14 @@ async fn execute_result_set_with_prepared_protocol_on_conn( diagnostic_trace_id: Option<&str>, start: Instant, progress_clock: Option<&crate::execution::StreamProgressClock>, + opaque_types: &DorisOpaqueColumnTypes, ) -> Result { let diagnostics_enabled = diagnostic_trace_id.is_some() && log::log_enabled!(log::Level::Info); let dispatch_started_at = diagnostics_enabled.then(Instant::now); let mut result = conn.exec_iter(sql, ()).await.map_err(|e| e.to_string())?; let dispatch_ms = dispatch_started_at.map_or(0, |started_at| started_at.elapsed().as_millis()); let columns: Vec = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect(); - let column_types: Vec = result.columns_ref().iter().map(mysql_column_type_name).collect(); + let column_types = effective_mysql_column_types(result.columns_ref(), opaque_types); let mut spatial_columns = mysql_spatial_column_builder(result.columns_ref()); let preview_bytes = max_result_bytes.map(|max_bytes| mysql_result_cell_preview_bytes(max_bytes, row_limit, result.columns_ref())); @@ -5538,6 +5816,7 @@ async fn execute_result_set_with_prepared_protocol_on_conn( } let Some(row) = next_row else { break }; let row = row.map_err(|e| e.to_string())?; + ensure_doris_opaque_cells_fit_budget(&row, &column_types, max_result_bytes)?; if let Some(progress_clock) = progress_clock { progress_clock.mark(); } @@ -5553,6 +5832,7 @@ async fn execute_result_set_with_prepared_protocol_on_conn( row_index, preview_bytes, &protected_indexes, + &column_types, ); result_rows.push(values); spatial_values.push(srids); @@ -5732,14 +6012,18 @@ pub async fn stream_query_result_on_conn( mut on_item: impl FnMut(MySqlQueryStreamItem) -> Result<(), String>, ) -> Result { let row_limit = max_rows.unwrap_or(usize::MAX); + let opaque_types = resolve_doris_opaque_column_types(conn, sql, dialect).await; if bare || prefers_text_protocol_query(sql, dialect) { - stream_query_result_text(conn, sql, row_limit, cancelled, spatial_as_wkb, &mut on_item).await + stream_query_result_text(conn, sql, row_limit, cancelled, spatial_as_wkb, &opaque_types, &mut on_item).await } else { - match stream_query_result_prepared(conn, sql, row_limit, cancelled, spatial_as_wkb, &mut on_item).await { + match stream_query_result_prepared(conn, sql, row_limit, cancelled, spatial_as_wkb, &opaque_types, &mut on_item) + .await + { Ok(rows) => Ok(rows), Err(err) if mysql_error_should_retry_with_text_protocol(&err) => { - stream_query_result_text(conn, sql, row_limit, cancelled, spatial_as_wkb, &mut on_item).await + stream_query_result_text(conn, sql, row_limit, cancelled, spatial_as_wkb, &opaque_types, &mut on_item) + .await } Err(err) => Err(err), } @@ -5752,6 +6036,7 @@ async fn stream_query_result_text( row_limit: usize, cancelled: &AtomicBool, spatial_as_wkb: bool, + opaque_types: &DorisOpaqueColumnTypes, on_item: &mut impl FnMut(MySqlQueryStreamItem) -> Result<(), String>, ) -> Result { let mut result = conn.query_iter(sql).await.map_err(|e| e.to_string())?; @@ -5759,8 +6044,8 @@ async fn stream_query_result_text( return Ok(0); } let columns: Vec = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect(); - let column_types: Vec = result.columns_ref().iter().map(mysql_column_type_name).collect(); - on_item(MySqlQueryStreamItem::Columns { columns, column_types })?; + let column_types = effective_mysql_column_types(result.columns_ref(), opaque_types); + on_item(MySqlQueryStreamItem::Columns { columns, column_types: column_types.clone() })?; let mut stream = result .stream::() @@ -5778,7 +6063,13 @@ async fn stream_query_result_text( } let row = row.map_err(|e| e.to_string())?; let values: Vec = (0..row.len()) - .map(|i| mysql_value_to_json_for_export(&row, i, spatial_as_wkb)) + .map(|i| { + if column_types.get(i).is_some_and(|column_type| is_opaque_aggregate_state_type(column_type)) { + Ok(mysql_value_to_json_with_effective_type(&row, i, column_types.get(i).map(String::as_str))) + } else { + mysql_value_to_json_for_export(&row, i, spatial_as_wkb) + } + }) .collect::>()?; on_item(MySqlQueryStreamItem::Row(values))?; rows_exported += 1; @@ -5793,6 +6084,7 @@ async fn stream_query_result_prepared( row_limit: usize, cancelled: &AtomicBool, spatial_as_wkb: bool, + opaque_types: &DorisOpaqueColumnTypes, on_item: &mut impl FnMut(MySqlQueryStreamItem) -> Result<(), String>, ) -> Result { let mut result = conn.exec_iter(sql, ()).await.map_err(|e| e.to_string())?; @@ -5800,8 +6092,8 @@ async fn stream_query_result_prepared( if columns.is_empty() { return Ok(0); } - let column_types: Vec = result.columns_ref().iter().map(mysql_column_type_name).collect(); - on_item(MySqlQueryStreamItem::Columns { columns, column_types })?; + let column_types = effective_mysql_column_types(result.columns_ref(), opaque_types); + on_item(MySqlQueryStreamItem::Columns { columns, column_types: column_types.clone() })?; let mut stream = result .stream::() @@ -5819,7 +6111,13 @@ async fn stream_query_result_prepared( } let row = row.map_err(|e| e.to_string())?; let values: Vec = (0..row.len()) - .map(|i| mysql_value_to_json_for_export(&row, i, spatial_as_wkb)) + .map(|i| { + if column_types.get(i).is_some_and(|column_type| is_opaque_aggregate_state_type(column_type)) { + Ok(mysql_value_to_json_with_effective_type(&row, i, column_types.get(i).map(String::as_str))) + } else { + mysql_value_to_json_for_export(&row, i, spatial_as_wkb) + } + }) .collect::>()?; on_item(MySqlQueryStreamItem::Row(values))?; rows_exported += 1; @@ -5942,7 +6240,7 @@ pub async fn execute_transaction_statement_on_conn( let mut truncated = false; while let Some(row) = query.next().await.map_err(transaction_error_from_mysql_error)? { if rows.len() < row_limit { - let (values, srids) = mysql_row_to_json_with_srids(&row, &mut spatial_columns); + let (values, srids) = mysql_row_to_json_with_srids(&row, &mut spatial_columns, &column_types); spatial_values.push(srids); rows.push(values); } else { @@ -6016,6 +6314,34 @@ pub async fn execute_query_on_conn_with_limits_progress( dialect: MySqlQueryDialect, diagnostic_trace_id: Option<&str>, progress_clock: Option<&crate::execution::StreamProgressClock>, +) -> Result { + let opaque_types = resolve_doris_opaque_column_types(conn, sql, dialect).await; + execute_query_on_conn_with_limits_progress_resolved( + conn, + sql, + bare, + max_rows, + max_result_bytes, + result_key_columns, + dialect, + diagnostic_trace_id, + progress_clock, + &opaque_types, + ) + .await +} + +async fn execute_query_on_conn_with_limits_progress_resolved( + conn: &mut mysql_async::Conn, + sql: &str, + bare: bool, + max_rows: Option, + max_result_bytes: Option, + result_key_columns: &[String], + dialect: MySqlQueryDialect, + diagnostic_trace_id: Option<&str>, + progress_clock: Option<&crate::execution::StreamProgressClock>, + opaque_types: &DorisOpaqueColumnTypes, ) -> Result { let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); @@ -6032,6 +6358,7 @@ pub async fn execute_query_on_conn_with_limits_progress( diagnostic_trace_id, start, progress_clock, + opaque_types, ) .await } else { @@ -6044,6 +6371,7 @@ pub async fn execute_query_on_conn_with_limits_progress( diagnostic_trace_id, start, progress_clock, + opaque_types, ) .await { @@ -6059,6 +6387,7 @@ pub async fn execute_query_on_conn_with_limits_progress( diagnostic_trace_id, start, progress_clock, + opaque_types, ) .await } @@ -6157,6 +6486,7 @@ pub async fn execute_query_results_on_conn_with_limits( dialect: MySqlQueryDialect, diagnostic_trace_id: Option<&str>, ) -> Result, String> { + let opaque_types = resolve_doris_opaque_column_types(conn, sql, dialect).await; if is_result_set_query(sql, dialect) && (bare || prefers_text_protocol_query(sql, dialect)) { let start = Instant::now(); execute_result_sets_with_text_protocol_on_conn( @@ -6168,10 +6498,11 @@ pub async fn execute_query_results_on_conn_with_limits( result_key_columns, diagnostic_trace_id, start, + &opaque_types, ) .await } else { - execute_query_on_conn_with_limits( + execute_query_on_conn_with_limits_progress_resolved( conn, sql, bare, @@ -6180,6 +6511,8 @@ pub async fn execute_query_results_on_conn_with_limits( result_key_columns, dialect, diagnostic_trace_id, + None, + &opaque_types, ) .await .map(|result| vec![result]) @@ -8305,6 +8638,197 @@ mod tests { assert_eq!(mysql_column_name(String::new()), None); } + #[test] + fn mysql_unknown_column_type_enrichment_preserves_information_schema_metadata() { + let mut columns = vec![ColumnInfo { + name: " state value".to_string(), + data_type: "unknown".to_string(), + resolved_schema: Some("analytics".to_string()), + is_nullable: true, + column_default: Some("seed".to_string()), + is_primary_key: true, + is_unique: true, + extra: Some("EXTRA".to_string()), + comment: Some("aggregate state".to_string()), + numeric_precision: Some(38), + numeric_scale: Some(4), + character_maximum_length: Some(512), + enum_values: Some(vec!["one".to_string()]), + character_set: Some("utf8mb4".to_string()), + collation: Some("utf8mb4_bin".to_string()), + metadata_capabilities: Some(ColumnMetadataCapabilities::all_supported()), + }]; + let original = columns[0].clone(); + let show_columns = vec![ColumnInfo { + name: " state value".to_string(), + data_type: "agg_state".to_string(), + ..Default::default() + }]; + + enrich_unknown_mysql_column_types(&mut columns, &show_columns); + + assert_eq!(columns[0].data_type, "agg_state"); + let mut expected = original; + expected.data_type = "agg_state".to_string(); + assert_eq!(serde_json::to_value(&columns[0]).unwrap(), serde_json::to_value(&expected).unwrap()); + assert_eq!(columns[0].metadata_capabilities, expected.metadata_capabilities); + } + + #[test] + fn mysql_unknown_column_type_enrichment_is_exact_and_conservative() { + let mut columns = vec![ + ColumnInfo { name: "known".to_string(), data_type: "varchar(20)".to_string(), ..Default::default() }, + ColumnInfo { name: " missing".to_string(), data_type: "unknown".to_string(), ..Default::default() }, + ColumnInfo { name: "blank".to_string(), data_type: "unknown".to_string(), ..Default::default() }, + ColumnInfo { name: "still_unknown".to_string(), data_type: "unknown".to_string(), ..Default::default() }, + ColumnInfo { name: " spaced".to_string(), data_type: "unknown".to_string(), ..Default::default() }, + ]; + let show_columns = vec![ + ColumnInfo { name: "known".to_string(), data_type: "bigint".to_string(), ..Default::default() }, + ColumnInfo { name: "missing".to_string(), data_type: "decimal(10,2)".to_string(), ..Default::default() }, + ColumnInfo { name: "blank".to_string(), data_type: " ".to_string(), ..Default::default() }, + ColumnInfo { name: "still_unknown".to_string(), data_type: "UNKNOWN".to_string(), ..Default::default() }, + ColumnInfo { + name: " spaced".to_string(), + data_type: "agg_state".to_string(), + ..Default::default() + }, + ]; + + assert!(needs_mysql_column_type_enrichment(&columns)); + enrich_unknown_mysql_column_types(&mut columns, &show_columns); + + assert_eq!(columns[0].data_type, "varchar(20)"); + assert_eq!(columns[1].data_type, "unknown"); + assert_eq!(columns[2].data_type, "unknown"); + assert_eq!(columns[3].data_type, "unknown"); + assert_eq!(columns[4].data_type, "agg_state"); + assert!(!needs_mysql_column_type_enrichment(&[ColumnInfo { + name: "known".to_string(), + data_type: "varchar(20)".to_string(), + ..Default::default() + }])); + } + + #[test] + fn doris_preflight_collects_static_sources_and_filters_ctes() { + let dialect = MySqlQueryDialect { is_doris: true, ..Default::default() }; + let mut sources = doris_preflight_sources( + "WITH recent AS (SELECT * FROM analytics.events) SELECT * FROM recent JOIN `raw`.`states` s ON 1=1", + dialect, + ) + .unwrap(); + sources.sort_by(|left, right| (&left.database, &left.table).cmp(&(&right.database, &right.table))); + assert_eq!( + sources, + vec![ + DorisSourceTable { database: Some("analytics".to_string()), table: "events".to_string() }, + DorisSourceTable { database: Some("raw".to_string()), table: "states".to_string() }, + ] + ); + } + + #[test] + fn doris_preflight_rejects_unsafe_or_observer_queries() { + let dialect = MySqlQueryDialect { is_doris: true, ..Default::default() }; + for sql in [ + "SELECT * FROM catalog.db.states", + "SELECT * FROM TABLE(generate_series(1, 2))", + "SELECT LAST_QUERY_ID() FROM states", + "SELECT sys.LAST_QUERY_ID() FROM states", + "WITH nested AS (SELECT LAST_QUERY_ID() FROM states) SELECT * FROM nested", + "WITH same_name AS (SELECT * FROM states) SELECT * FROM (WITH same_name AS (SELECT * FROM other_states) SELECT * FROM same_name) nested", + "SELECT * FROM states; SELECT * FROM other_states", + ] { + assert_eq!(doris_preflight_sources(sql, dialect), None, "{sql}"); + } + assert_eq!(doris_preflight_sources("SELECT * FROM states", MySqlQueryDialect::default()), None); + } + + #[test] + fn doris_query_enrichment_profile_is_narrow() { + assert!(MySqlQueryDialect::for_connection(DatabaseType::Doris, None).is_doris); + assert!(MySqlQueryDialect::for_connection(DatabaseType::Mysql, Some("doris")).is_doris); + assert!(!MySqlQueryDialect::for_connection(DatabaseType::Mysql, Some("selectdb")).is_doris); + assert!(!MySqlQueryDialect::for_connection(DatabaseType::Mysql, Some("starrocks")).is_doris); + assert!(!MySqlQueryDialect::for_connection(DatabaseType::Mysql, None).is_doris); + } + + #[test] + fn doris_effective_type_requires_complete_exact_origin() { + let direct = Column::new(ColumnType::MYSQL_TYPE_STRING) + .with_schema(b"analytics") + .with_org_table(b"states") + .with_name(b"renamed") + .with_org_name(b"v2") + .with_character_set(33); + let expression = Column::new(ColumnType::MYSQL_TYPE_STRING).with_name(b"v2").with_character_set(33); + let wrong_origin = Column::new(ColumnType::MYSQL_TYPE_STRING) + .with_schema(b"analytics") + .with_org_table(b"states") + .with_name(b"v2") + .with_org_name(b"ordinary_text") + .with_character_set(33); + let incompatible_wire = Column::new(ColumnType::MYSQL_TYPE_LONG) + .with_schema(b"analytics") + .with_org_table(b"states") + .with_name(b"v2") + .with_org_name(b"v2"); + let opaque_types = HashMap::from([( + (b"analytics".to_vec(), b"states".to_vec(), b"v2".to_vec()), + "agg_state".to_string(), + )]); + + assert_eq!( + effective_mysql_column_types(&[direct, expression, wrong_origin, incompatible_wire], &opaque_types), + vec!["agg_state", "char", "char", "int"] + ); + } + + #[test] + fn doris_duplicate_show_column_names_are_ambiguous() { + let key = (b"analytics".to_vec(), b"states".to_vec(), b"v2".to_vec()); + let mut resolved = HashMap::new(); + let mut ambiguous = HashSet::new(); + insert_doris_opaque_column_type(&mut resolved, &mut ambiguous, key.clone(), "agg_state".to_string()); + insert_doris_opaque_column_type(&mut resolved, &mut ambiguous, key.clone(), "agg_state".to_string()); + insert_doris_opaque_column_type(&mut resolved, &mut ambiguous, key.clone(), "agg_state".to_string()); + assert!(!resolved.contains_key(&key)); + assert!(ambiguous.contains(&key)); + } + + #[test] + fn doris_opaque_values_preserve_invalid_utf8_and_null() { + let column = Column::new(ColumnType::MYSQL_TYPE_STRING).with_name(b"v2").with_character_set(33); + let row = mysql_test_row_with_columns(vec![Value::Bytes(vec![0x00, 0xff, 0x01])], vec![column.clone()]); + assert_eq!( + mysql_value_to_json_with_effective_type(&row, 0, Some("agg_state")), + serde_json::json!("0x00ff01") + ); + assert!(ensure_doris_opaque_cells_fit_budget(&row, &["agg_state".to_string()], Some(7)) + .unwrap_err() + .contains("interactive byte budget")); + + let modest = mysql_test_row_with_columns(vec![Value::Bytes(vec![0xff; 9 * 1024])], vec![column.clone()]); + assert!(ensure_doris_opaque_cells_fit_budget(&modest, &["agg_state".to_string()], Some(20 * 1024),) + .is_ok()); + assert!(ensure_doris_opaque_cells_fit_budget(&modest, &["agg_state".to_string()], Some(18 * 1024),) + .is_err()); + + let null_row = mysql_test_row_with_columns(vec![Value::NULL], vec![column]); + assert_eq!( + mysql_value_to_json_with_effective_type(&null_row, 0, Some("agg_state")), + serde_json::Value::Null + ); + + let unexpected = + mysql_test_row_with_columns(vec![Value::Int(42)], vec![Column::new(ColumnType::MYSQL_TYPE_LONG)]); + assert_eq!( + mysql_value_to_json_with_effective_type(&unexpected, 0, Some("agg_state")), + serde_json::json!(42) + ); + } + #[test] fn parse_mysql_enum_values_preserves_mysql_literal_edges() { assert_eq!( @@ -8416,7 +8940,7 @@ mod tests { let mut spatial_columns = mysql_spatial_column_builder(row.columns_ref()); let (values, _srids, cells) = - mysql_row_to_json_with_srids_and_previews(&row, &mut spatial_columns, 7, Some(1024), &protected); + mysql_row_to_json_with_srids_and_previews(&row, &mut spatial_columns, 7, Some(1024), &protected, &[]); assert_eq!(values[0], serde_json::json!("k".repeat(2048))); assert_eq!(values[1].as_str().map(str::len), Some(1027)); @@ -8443,7 +8967,7 @@ mod tests { let mut spatial_columns = mysql_spatial_column_builder(row.columns_ref()); let (values, _srids, cells) = - mysql_row_to_json_with_srids_and_previews(&row, &mut spatial_columns, 0, Some(512), &protected); + mysql_row_to_json_with_srids_and_previews(&row, &mut spatial_columns, 0, Some(512), &protected, &[]); assert_eq!(values[0], serde_json::json!("k".repeat(2048))); assert_eq!(cells, vec![LargeValueCell { row_index: 0, column_index: 1, original_bytes: 4096 }]); @@ -8480,7 +9004,7 @@ mod tests { let mut spatial_columns = mysql_spatial_column_builder(row.columns_ref()); let (values, _srids, cells) = - mysql_row_to_json_with_srids_and_previews(&row, &mut spatial_columns, 0, Some(256), &protected); + mysql_row_to_json_with_srids_and_previews(&row, &mut spatial_columns, 0, Some(256), &protected, &[]); assert_eq!(values[0].as_str().map(str::len), Some(500)); assert!(values[1].as_str().is_some_and(|value| value.ends_with("..."))); @@ -8500,7 +9024,7 @@ mod tests { let mut spatial_columns = mysql_spatial_column_builder(row.columns_ref()); let (values, _srids, cells) = - mysql_row_to_json_with_srids_and_previews(&row, &mut spatial_columns, 0, Some(512), &HashSet::new()); + mysql_row_to_json_with_srids_and_previews(&row, &mut spatial_columns, 0, Some(512), &HashSet::new(), &[]); let preview = values[0].as_str().unwrap(); assert!(preview.starts_with("0xabab")); diff --git a/crates/dbx-sql-data/src/data_grid_extractors/sql.rs b/crates/dbx-sql-data/src/data_grid_extractors/sql.rs index 2218b5bc48..6a4dc5b5ed 100644 --- a/crates/dbx-sql-data/src/data_grid_extractors/sql.rs +++ b/crates/dbx-sql-data/src/data_grid_extractors/sql.rs @@ -9,6 +9,7 @@ use crate::data_grid_sql::{ is_grid_insert_omitted_column, is_non_identity_generated_column, supports_relational_copy_predicates, DataGridCopyInsertStatementOptions, DataGridCopyUpdateStatementOptions, DataGridTableMeta, }; +use dbx_types::types::is_opaque_aggregate_state_type; use serde_json::Value; use std::collections::{hash_map::Entry, HashMap, HashSet}; use std::io::Write; @@ -61,6 +62,7 @@ pub(super) fn write_sql_inserts( output: &mut dyn Write, ) -> Result { ensure_sql_builder_budget(context)?; + reject_selected_opaque_aggregate_states(context)?; let data = sql_selected_data(context, false)?; let statement = build_data_grid_copy_insert_statement(DataGridCopyInsertStatementOptions { database_type: context.request.database_type, @@ -90,6 +92,7 @@ pub(super) fn write_sql_updates( output: &mut dyn Write, ) -> Result { ensure_sql_builder_budget(context)?; + reject_selected_opaque_aggregate_states(context)?; let table_meta = context.request.table_meta.as_ref().ok_or_else(|| { DataGridExtractError::new( DataGridExtractErrorCode::MissingTableMetadata, @@ -133,6 +136,16 @@ pub(super) fn write_sql_updates( Ok(sql_metadata(data.omitted_columns)) } +fn reject_selected_opaque_aggregate_states(context: &ExtractContext<'_>) -> Result<(), DataGridExtractError> { + if context.selected_column_info.iter().flatten().any(|info| is_opaque_aggregate_state_type(&info.data_type)) { + return Err(DataGridExtractError::new( + DataGridExtractErrorCode::NoWritableColumns, + "Doris aggregate-state columns are opaque and cannot be copied as automatic SQL; use an explicit Doris state function instead.", + )); + } + Ok(()) +} + fn ensure_sql_builder_budget(context: &ExtractContext<'_>) -> Result<(), DataGridExtractError> { let value_bytes = context .request diff --git a/crates/dbx-sql-data/src/data_grid_extractors/tests.rs b/crates/dbx-sql-data/src/data_grid_extractors/tests.rs index 702d585c52..79c061ed01 100644 --- a/crates/dbx-sql-data/src/data_grid_extractors/tests.rs +++ b/crates/dbx-sql-data/src/data_grid_extractors/tests.rs @@ -1363,6 +1363,39 @@ fn sql_updates_require_table_primary_key_and_writable_columns() { ); } +#[test] +fn sql_copy_reports_selected_opaque_aggregate_states_before_generic_builder_errors() { + let opaque_column = DataGridColumnInfo { + name: "v2".to_string(), + data_type: "agg_state".to_string(), + is_nullable: true, + is_primary_key: false, + column_default: None, + extra: None, + }; + for extractor in [DataGridExtractorId::SqlInserts, DataGridExtractorId::SqlUpdates] { + let mut request = request(extractor); + request.columns = vec![column("v2", 0)]; + request.selected_column_indexes = vec![0]; + request.rows = vec![vec![json!("0x00ff")]]; + request.table_meta = Some(DataGridTableMeta { + catalog: None, + database: None, + schema: None, + table_name: "states".to_string(), + primary_keys: Vec::new(), + columns: Some(vec![opaque_column.clone()]), + }); + + let error = extract_data_grid_selection(request).expect_err("opaque SQL copy must fail explicitly"); + assert_eq!(error.code, DataGridExtractErrorCode::NoWritableColumns); + assert_eq!( + error.message, + "Doris aggregate-state columns are opaque and cannot be copied as automatic SQL; use an explicit Doris state function instead." + ); + } +} + #[test] fn rejects_requests_that_exceed_the_column_budget() { let mut request = request(DataGridExtractorId::Csv); diff --git a/crates/dbx-sql-data/src/data_grid_sql.rs b/crates/dbx-sql-data/src/data_grid_sql.rs index 5fee22ef61..c44e1ba1b2 100644 --- a/crates/dbx-sql-data/src/data_grid_sql.rs +++ b/crates/dbx-sql-data/src/data_grid_sql.rs @@ -34,6 +34,7 @@ use crate::sql_dialect::{ uses_single_row_insert_statements, uses_synthetic_row_id, uses_xugu_row_id, TablePaginationStrategy, }; use crate::value_literals::{format_ch_array_sql_literal, format_pg_array_sql_literal}; +use dbx_types::types::is_opaque_aggregate_state_type; const DBX_ROWID_COLUMN: &str = "__DBX_ROWID"; pub const DBX_NEO4J_ELEMENT_ID_COLUMN: &str = "__DBX_ELEMENT_ID"; @@ -441,6 +442,14 @@ pub fn build_data_grid_copy_update_statements(options: DataGridCopyUpdateStateme let primary_key_info = primary_keys.iter().map(|primary_key| column_info_for(column_info, primary_key)).collect::>(); + if writable_indexes + .iter() + .any(|(_, _, info)| info.is_some_and(|info| is_opaque_aggregate_state_type(&info.data_type))) + || primary_key_info.iter().any(|info| info.is_some_and(|info| is_opaque_aggregate_state_type(&info.data_type))) + { + return Vec::new(); + } + if writable_indexes.is_empty() { return Vec::new(); } @@ -541,6 +550,13 @@ pub fn build_data_grid_copy_insert_statement(options: DataGridCopyInsertStatemen .cloned() .collect(); + if insert_columns + .iter() + .any(|(_, _, info)| info.as_ref().is_some_and(|info| is_opaque_aggregate_state_type(&info.data_type))) + { + return None; + } + if insert_columns.is_empty() || options.rows.is_empty() { return None; } @@ -1199,6 +1215,9 @@ fn build_neo4j_data_grid_column_distinct_values_sql(options: &DataGridColumnDist } fn validate_data_grid_save(options: &DataGridSaveStatementOptions) -> Option { + if let Some(error) = validate_opaque_aggregate_state_write(options) { + return Some(error); + } if let Some(error) = validate_salesforce_id_column(options) { return Some(error); } @@ -1281,6 +1300,45 @@ fn validate_data_grid_save(options: &DataGridSaveStatementOptions) -> Option Option { + let save_columns = effective_columns(options); + let column_info = options.table_meta.columns.as_deref().unwrap_or(&[]); + let opaque_indexes = save_columns + .iter() + .enumerate() + .filter_map(|(index, column)| { + let column = column.as_deref()?; + column_info_for(column_info, column) + .is_some_and(|info| is_opaque_aggregate_state_type(&info.data_type)) + .then_some(index) + }) + .collect::>(); + if opaque_indexes.is_empty() { + return None; + } + if options.dirty_rows.iter().any(|(_, changes)| changes.iter().any(|(index, _)| opaque_indexes.contains(index))) + || options + .new_rows + .iter() + .any(|row| opaque_indexes.iter().any(|index| row.get(*index).is_some_and(|value| !value.is_null()))) + { + return Some("Doris aggregate-state columns are opaque and cannot be written automatically; use an explicit Doris state function instead.".to_string()); + } + if options.table_meta.primary_keys.is_empty() + && (!options.dirty_rows.is_empty() || !options.deleted_rows.is_empty()) + && options.dirty_rows.iter().map(|(index, _)| *index).chain(options.deleted_rows.iter().copied()).any( + |row_index| { + options.rows.get(row_index).is_some_and(|row| { + opaque_indexes.iter().any(|index| row.get(*index).is_some_and(|value| !value.is_null())) + }) + }, + ) + { + return Some("Cannot safely update or delete a keyless row whose predicate would contain an opaque Doris aggregate-state value.".to_string()); + } + None +} + fn validate_existing_row_primary_keys(options: &DataGridSaveStatementOptions) -> Option { let primary_keys = &options.table_meta.primary_keys; if primary_keys.is_empty() || (options.dirty_rows.is_empty() && options.deleted_rows.is_empty()) { @@ -10155,4 +10213,61 @@ mod tests { assert_eq!(result.validation_error, Some(r#"Column "LogTime" does not allow NULL."#.to_string())); assert!(result.statements.is_empty()); } + + #[test] + fn opaque_aggregate_state_blocks_keyless_predicates_and_copy_sql() { + let table_meta = DataGridTableMeta { + catalog: None, + database: Some("analytics".to_string()), + schema: None, + table_name: "states".to_string(), + primary_keys: vec![], + columns: Some(vec![ + column("name", "varchar", false, None), + column("v2", "agg_state", false, None), + ]), + }; + let save = prepare_data_grid_save(DataGridSaveStatementOptions { + database_type: Some(DatabaseType::Doris), + identifier_quote: None, + table_meta: table_meta.clone(), + columns: vec!["name".to_string(), "v2".to_string()], + source_columns: None, + rows: vec![vec![json!("before"), json!("0x00ff")]], + dirty_rows: vec![(0, vec![(0, json!("after"))])], + deleted_rows: vec![], + new_rows: vec![], + include_database_name: false, + }); + assert!(save.validation_error.as_deref().is_some_and(|error| error.contains("keyless row"))); + assert!(save.statements.is_empty()); + + assert!(build_data_grid_copy_insert_statement(DataGridCopyInsertStatementOptions { + database_type: Some(DatabaseType::Doris), + identifier_quote: None, + table_meta: Some(table_meta.clone()), + columns: vec!["name".to_string(), "v2".to_string()], + column_types: None, + source_columns: None, + rows: vec![vec![json!("before"), json!("0x00ff")]], + exclude_primary_keys: false, + include_computed_columns: false, + include_database_name: true, + insert_mode: DataGridCopyInsertMode::Merged, + }) + .is_none()); + + let mut keyed = table_meta; + keyed.primary_keys = vec!["name".to_string()]; + assert!(build_data_grid_copy_update_statements(DataGridCopyUpdateStatementOptions { + database_type: Some(DatabaseType::Doris), + identifier_quote: None, + table_meta: keyed, + columns: vec!["name".to_string(), "v2".to_string()], + source_columns: None, + rows: vec![vec![json!("before"), json!("0x00ff")]], + include_database_name: true, + }) + .is_empty()); + } } diff --git a/crates/dbx-types/src/types.rs b/crates/dbx-types/src/types.rs index 72b20ddaed..b803e84f7b 100644 --- a/crates/dbx-types/src/types.rs +++ b/crates/dbx-types/src/types.rs @@ -278,6 +278,22 @@ pub struct ColumnInfo { pub metadata_capabilities: Option, } +/// Doris aggregate-state columns contain opaque engine serialization, not a +/// value that DBX can safely edit or emit as an INSERT literal. +pub fn is_opaque_aggregate_state_type(data_type: &str) -> bool { + let data_type = data_type.trim(); + let Some(prefix) = data_type.get(.."agg_state".len()) else { return false }; + if !prefix.eq_ignore_ascii_case("agg_state") { + return false; + } + let Some(arguments) = + data_type.get("agg_state".len()..).map(str::trim_start).and_then(|value| value.strip_prefix('<')) + else { + return false; + }; + arguments.strip_suffix('>').is_some_and(|inner| !inner.trim().is_empty()) +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TableColumnsResult { pub table_name: String, @@ -1047,10 +1063,21 @@ pub struct CustomTypeDetails { #[cfg(test)] mod tests { use super::{ - CompletionAssistantCandidate, CompletionAssistantCandidateKind, ObjectInfo, ObjectSourceKind, QueryMessage, - SpatialColumn, SpatialColumnBuilder, TableInfo, + is_opaque_aggregate_state_type, CompletionAssistantCandidate, CompletionAssistantCandidateKind, ObjectInfo, + ObjectSourceKind, QueryMessage, SpatialColumn, SpatialColumnBuilder, TableInfo, }; + #[test] + fn opaque_aggregate_state_type_is_narrow() { + assert!(is_opaque_aggregate_state_type("agg_state")); + assert!(is_opaque_aggregate_state_type(" AGG_STATE ")); + assert!(!is_opaque_aggregate_state_type("agg_state")); + assert!(!is_opaque_aggregate_state_type("agg_state<")); + assert!(!is_opaque_aggregate_state_type("agg_state<>")); + assert!(!is_opaque_aggregate_state_type("😺agg_state")); + assert!(!is_opaque_aggregate_state_type("varchar")); + } + #[test] fn query_message_format_line_uppercases_severity() { let message = QueryMessage { diff --git a/crates/dbx-web/src/routes/transfer.rs b/crates/dbx-web/src/routes/transfer.rs index ecf22d174c..c27b800193 100644 --- a/crates/dbx-web/src/routes/transfer.rs +++ b/crates/dbx-web/src/routes/transfer.rs @@ -182,6 +182,12 @@ pub async fn start_transfer( } }; + if let Err(e) = transfer::ensure_transfer_source_types_supported(&app, &req, &source_pool_key).await { + send_transfer_progress(&progress_channel, &terminal_transfer_error(&req, e)); + finish_transfer_channel(&state_clone, &req.transfer_id, &progress_channel).await; + return; + } + let tables = req.tables.clone(); // Sort by FK dependency so referenced tables are transferred first, and // keep the foreign key metadata fetched along the way — MySQL-family diff --git a/src-tauri/src/commands/transfer.rs b/src-tauri/src/commands/transfer.rs index 469396140b..4bb090f05b 100644 --- a/src-tauri/src/commands/transfer.rs +++ b/src-tauri/src/commands/transfer.rs @@ -68,6 +68,8 @@ pub async fn start_transfer( ) .await?; + dbx_core::transfer::ensure_transfer_source_types_supported(&state, &request, &source_pool_key).await?; + tokio::spawn(async move { // Sort tables by FK dependency so referenced tables are transferred first, // and keep the foreign key metadata fetched along the way — MySQL-family