From eb76253d2211601627d595c96bcc1f9d7195f781 Mon Sep 17 00:00:00 2001 From: t8y2 Date: Wed, 30 Sep 2026 15:53:02 +0800 Subject: [PATCH] feat(db2): enable data transfers and field lineage --- .../transfer/DataTransferDialog.vue | 6 +- .../__tests__/transferStrategy.spec.ts | 11 +- .../components/transfer/transferStrategy.ts | 4 + .../database/databaseFeatureSupport.spec.ts | 16 ++ .../assets/database-drivers.manifest.json | 4 +- crates/dbx-core/src/data/transfer.rs | 69 ++++-- crates/dbx-core/src/data/transfer/db2.rs | 133 +++++++++++ .../dbx-core/src/data/transfer/db2_tests.rs | 210 ++++++++++++++++++ plugins/connection-types/db2.yaml | 4 +- 9 files changed, 436 insertions(+), 21 deletions(-) create mode 100644 crates/dbx-core/src/data/transfer/db2.rs create mode 100644 crates/dbx-core/src/data/transfer/db2_tests.rs diff --git a/apps/desktop/src/components/transfer/DataTransferDialog.vue b/apps/desktop/src/components/transfer/DataTransferDialog.vue index 7a4f5a1a45..a17f26923e 100644 --- a/apps/desktop/src/components/transfer/DataTransferDialog.vue +++ b/apps/desktop/src/components/transfer/DataTransferDialog.vue @@ -6,7 +6,7 @@ import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle, Di import { Button } from "@/components/ui/button"; import { buildTransferObjectSelections, countTransferObjects } from "./transferSelections"; import { createTaskLoadTracker } from "./taskLoadTracker"; -import { confirmTransferWithProductionSafety, createTransferSubmission, rebuildUnavailableReason, resolveTransferStrategy, transferStrategyOptions, type TransferStrategy } from "./transferStrategy"; +import { confirmTransferWithProductionSafety, createTransferSubmission, rebuildUnavailableReason, resolveTransferStrategy, supportsTransferUpsert, transferStrategyOptions, type TransferStrategy } from "./transferStrategy"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; @@ -220,6 +220,7 @@ function isMongoConnection(id: string): boolean { const showTargetColumnQuoteOption = computed(() => ["gaussdb", "opengauss"].includes(connectionType(targetConnectionId.value) ?? "")); const rebuildDisabledReason = computed(() => rebuildUnavailableReason(transferContent.value, connectionType(targetConnectionId.value))); +const upsertSupported = computed(() => supportsTransferUpsert(connectionType(targetConnectionId.value))); const rebuildDisabledHint = computed(() => { if (rebuildDisabledReason.value === "dataOnly") return t("transfer.rebuildDataOnlyDisabled"); if (rebuildDisabledReason.value === "unsupported") return t("transfer.rebuildUnsupportedDisabled"); @@ -266,6 +267,7 @@ const canStart = computed(() => { (targetCatalogs.value.length <= 1 || !!targetCatalog.value) && (selectedTables.value.size > 0 || Object.values(selectedObjects.value).some((names) => names.size > 0)) && (targetTableStrategy.value !== "rebuild" || !rebuildDisabledReason.value) && + (transferContent.value === "structureOnly" || targetTableStrategy.value !== "upsert" || upsertSupported.value) && !sameSourceAndTarget ); }); @@ -1371,7 +1373,7 @@ async function saveConfigTask() { {{ t(transferContent === "structureOnly" ? "transfer.modeKeepExisting" : "transfer.modeAppend") }} {{ t("transfer.modeOverwrite") }} - {{ t("transfer.modeUpsert") }} + {{ t("transfer.modeUpsert") }} {{ t(transferContent === "structureOnly" ? "transfer.modeRebuildStructure" : "transfer.modeRebuild") }} diff --git a/apps/desktop/src/components/transfer/__tests__/transferStrategy.spec.ts b/apps/desktop/src/components/transfer/__tests__/transferStrategy.spec.ts index 742d3c5183..484515521f 100644 --- a/apps/desktop/src/components/transfer/__tests__/transferStrategy.spec.ts +++ b/apps/desktop/src/components/transfer/__tests__/transferStrategy.spec.ts @@ -3,7 +3,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import type { TransferOwnershipPreview, TransferRequest } from "@/lib/backend/api"; import type { ConnectionConfig } from "@/types/database"; import { useProductionSafetyStore } from "@/stores/productionSafetyStore"; -import { confirmTransferWithProductionSafety, createTransferSubmission, rebuildUnavailableReason, resolveTransferStrategy, transferStrategyOptions } from "../transferStrategy"; +import { confirmTransferWithProductionSafety, createTransferSubmission, rebuildUnavailableReason, resolveTransferStrategy, supportsTransferUpsert, transferStrategyOptions } from "../transferStrategy"; function deferred() { let resolve!: (value: T) => void; @@ -70,6 +70,15 @@ describe("transfer strategies", () => { expect(resolveTransferStrategy({})).toBe("append"); }); + it("does not expose unimplemented DB2 upsert or rebuild strategies", () => { + expect(supportsTransferUpsert("db2")).toBe(false); + expect(rebuildUnavailableReason("structureAndData", "db2")).toBe("unsupported"); + expect(rebuildUnavailableReason("structureOnly", "db2")).toBe("unsupported"); + expect(supportsTransferUpsert("postgres")).toBe(true); + expect(supportsTransferUpsert("mysql")).toBe(true); + expect(supportsTransferUpsert(undefined)).toBe(true); + }); + it.each([ ["dataOnly", "postgres", "dataOnly"], ["structureAndData", "mongodb", "unsupported"], diff --git a/apps/desktop/src/components/transfer/transferStrategy.ts b/apps/desktop/src/components/transfer/transferStrategy.ts index dd2f25febf..dacfaae6da 100644 --- a/apps/desktop/src/components/transfer/transferStrategy.ts +++ b/apps/desktop/src/components/transfer/transferStrategy.ts @@ -15,6 +15,10 @@ export function transferStrategyOptions(strategy: TransferStrategy): Pick(["mysql", "postgres", "sqlserver", "kingbase", "gaussdb", "opengauss", "kwdb", "goldendb", "sqlite", "duckdb", "cloudflare-d1"]); +export function supportsTransferUpsert(targetType: DatabaseType | undefined): boolean { + return targetType !== "db2"; +} + export function rebuildUnavailableReason(content: TransferContent, targetType: DatabaseType | undefined): "dataOnly" | "unsupported" | undefined { if (content === "dataOnly") return "dataOnly"; return targetType && REBUILD_TARGET_TYPES.has(targetType) ? undefined : "unsupported"; diff --git a/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts b/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts index 521098d3ff..b07e1dc4e8 100644 --- a/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts +++ b/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts @@ -11,6 +11,7 @@ import { supportsDatabaseNameCompletion, supportsDatabaseSchemaQualifier, supportsDatabaseSearch, + supportsFieldLineage, supportsObjectBrowser, supportsObjectBrowserTreeNode, supportsQueryExecution, @@ -19,6 +20,7 @@ import { supportsSqlInListPaste, supportsTableImport, supportsTableVacuum, + supportsTransfer, supportsTransaction, usesOracleStickyTransactionState, usesProvenReadOnlyStickyTransactionState, @@ -336,6 +338,20 @@ describe("supportsTableImport", () => { }); }); +describe("supportsFieldLineage", () => { + it("enables DB2 field lineage", () => { + expect(supportsFieldLineage("db2")).toBe(true); + }); +}); + +describe("DB2 transfer and namespace capabilities", () => { + it("enables table transfers without enabling database lifecycle operations", () => { + expect(supportsTransfer("db2")).toBe(true); + expect(connectionNamespaceCreationTarget({ db_type: "db2" })).toBeNull(); + expect(databaseNodeNamespaceCreationTarget({ db_type: "db2" }, { type: "database", database: "SAMPLE" })).toBe("schema"); + }); +}); + describe("database property editing", () => { it("allows MySQL-compatible charset and collation edits on database nodes", () => { expect(editableDatabasePropertyGroups({ db_type: "mysql" }, { type: "database", database: "app" })).toEqual(["charsetCollation"]); diff --git a/crates/dbx-core/assets/database-drivers.manifest.json b/crates/dbx-core/assets/database-drivers.manifest.json index 4b2626f3c9..4fbbdfd7d0 100644 --- a/crates/dbx-core/assets/database-drivers.manifest.json +++ b/crates/dbx-core/assets/database-drivers.manifest.json @@ -2045,10 +2045,10 @@ "tableDataEdit": true, "tableStructureEdit": false, "tableImport": true, - "dataTransfer": false, + "dataTransfer": true, "sqlFileExecution": true, "databaseCreate": false, - "fieldLineage": false, + "fieldLineage": true, "sqlExplain": false, "userAdmin": false, "driverManagement": true diff --git a/crates/dbx-core/src/data/transfer.rs b/crates/dbx-core/src/data/transfer.rs index cd246050af..8aae6f38df 100644 --- a/crates/dbx-core/src/data/transfer.rs +++ b/crates/dbx-core/src/data/transfer.rs @@ -17,6 +17,7 @@ use futures::{SinkExt, StreamExt}; #[path = "transfer/rebuild_tests.rs"] mod rebuild_tests; +mod db2; mod ddl_plan; use crate::connection::{config_for_pool_key, AppState, PoolKind}; @@ -2926,6 +2927,9 @@ pub fn escape_value_typed(val: &serde_json::Value, db_type: &DatabaseType, colum } }, serde_json::Value::Number(n) => { + if *db_type == DatabaseType::Db2 { + return db2::numeric_literal(&n.to_string(), column_type); + } if let Some(integer_literal) = normalize_integer_literal(&n.to_string(), db_type, column_type) { return integer_literal; } @@ -2941,6 +2945,9 @@ pub fn escape_value_typed(val: &serde_json::Value, db_type: &DatabaseType, colum } } serde_json::Value::String(s) => { + if *db_type == DatabaseType::Db2 { + return db2::string_literal(s, column_type); + } if let Some(json_array_literal) = format_starrocks_json_array_sql_literal(s, db_type, column_type) { return json_array_literal; } @@ -3000,11 +3007,15 @@ pub fn escape_value_typed(val: &serde_json::Value, db_type: &DatabaseType, colum } match db_type { DatabaseType::ClickHouse | DatabaseType::Databend => format_ch_array_sql_literal(arr), + DatabaseType::Db2 => db2::string_literal(&val.to_string(), None), _ => format_pg_array_sql_literal(arr), } } _ => { let s = val.to_string(); + if *db_type == DatabaseType::Db2 { + return db2::string_literal(&s, None); + } if *db_type == DatabaseType::H2 { return quote_string_literal(&s); } @@ -3418,7 +3429,7 @@ fn oracle_char_length_params(params: &str) -> String { /// targets to `CLOB` for file imports; the transfer path did not (#9886). fn target_text_type(target_db: &DatabaseType) -> &'static str { match target_db { - DatabaseType::Oracle | DatabaseType::OceanbaseOracle | DatabaseType::Dameng => "CLOB", + DatabaseType::Oracle | DatabaseType::OceanbaseOracle | DatabaseType::Dameng | DatabaseType::Db2 => "CLOB", _ => "TEXT", } } @@ -3427,6 +3438,11 @@ pub fn map_column_type(source_type: &str, source_db: &DatabaseType, target_db: & if source_db == target_db { return source_type.to_string(); } + if *target_db == DatabaseType::Db2 { + if let Some(mapped) = db2::map_column_type(source_type, source_db) { + return mapped; + } + } let t = source_type.to_lowercase(); let mut base = t.split('(').next().unwrap_or(&t).trim(); if matches!(source_db, DatabaseType::H2) { @@ -3771,9 +3787,11 @@ fn generate_create_table_ddl_with_column_quoting( }; let create_prefix = match target_db { - DatabaseType::Oracle | DatabaseType::OceanbaseOracle | DatabaseType::SqlServer | DatabaseType::Dameng => { - "CREATE TABLE" - } + DatabaseType::Oracle + | DatabaseType::OceanbaseOracle + | DatabaseType::SqlServer + | DatabaseType::Dameng + | DatabaseType::Db2 => "CREATE TABLE", _ => "CREATE TABLE IF NOT EXISTS", }; @@ -4630,6 +4648,8 @@ fn rewrite_transfer_source_table_ddl( let source_schema = if source_schema.trim().is_empty() { "PUBLIC" } else { source_schema }; let target_schema = if target_schema.trim().is_empty() { "PUBLIC" } else { target_schema }; Some(rewrite_h2_schema_qualifier(sql, source_schema, target_schema)) + } else if matches!((source_db_type, target_db_type), (DatabaseType::Db2, DatabaseType::Db2)) { + Some(rewrite_h2_schema_qualifier(sql, source_schema, target_schema)) } else if is_oracle_family_transfer_target(source_db_type) && is_oracle_family_transfer_target(target_db_type) { // Oracle-family sources return `DBMS_METADATA`-style DDL whose CREATE TABLE head // carries the owner schema (`"SALES"."T"`). Unlike the Oracle schema-object path @@ -9196,6 +9216,19 @@ fn transfer_cursor_sql( format!("SELECT {col_list} FROM {full_table}") } +fn uses_agent_transfer_cursor(db_type: &DatabaseType) -> bool { + matches!(db_type, DatabaseType::Kyuubi | DatabaseType::Impala | DatabaseType::Transwarp | DatabaseType::Db2) +} + +fn transfer_clear_table_sql(table: &str, schema: &str, db_type: &DatabaseType, catalog: Option<&str>) -> String { + let full_table = qualified_table(table, schema, db_type, catalog); + match db_type { + DatabaseType::Sqlite | DatabaseType::CloudflareD1 | DatabaseType::DuckDb => format!("DELETE FROM {full_table}"), + DatabaseType::Db2 => format!("TRUNCATE TABLE {full_table} IMMEDIATE"), + _ => format!("TRUNCATE TABLE {full_table}"), + } +} + async fn fetch_hive_server_transfer_batch( state: &AppState, pool_key: &str, @@ -9800,6 +9833,9 @@ async fn transfer_table_inner( where F: FnMut(TransferProgress), { + if *target_db_type == DatabaseType::Db2 { + db2::validate_request(request)?; + } if is_mongodb_transfer_type(source_db_type) || is_mongodb_transfer_type(target_db_type) { return transfer_mongodb_table( state, @@ -10179,6 +10215,14 @@ where col_names.clone() }; + if *target_db_type == DatabaseType::Db2 { + let sql = db2::generated_columns_sql(&request.target_schema, &target_table, &write_col_names); + let generated_columns = execute_on_pool(state, target_pool_key, &sql) + .await + .map_err(|error| format!("Failed to inspect DB2 generated columns before transfer: {error}"))?; + db2::validate_generated_columns(&generated_columns.rows)?; + } + if server_side_complex_copy { for (source_column, target_name) in writable_columns.iter().zip(&write_col_names) { let Some(target_column) = @@ -10204,14 +10248,12 @@ where // When drop_target_before_create is true, the target table was just created // and is already empty, so TRUNCATE is unnecessary. if request.mode == TransferMode::Overwrite && !request.drop_target_before_create { - let full_table = - qualified_table(&target_table, &request.target_schema, target_db_type, request.target_catalog.as_deref()); - let truncate_sql = match target_db_type { - DatabaseType::Sqlite | DatabaseType::CloudflareD1 | DatabaseType::DuckDb => { - format!("DELETE FROM {full_table}") - } - _ => format!("TRUNCATE TABLE {full_table}"), - }; + let truncate_sql = transfer_clear_table_sql( + &target_table, + &request.target_schema, + target_db_type, + request.target_catalog.as_deref(), + ); execute_on_pool(state, target_pool_key, &truncate_sql).await.map_err(|e| format!("Failed to truncate: {e}"))?; } @@ -10369,8 +10411,7 @@ where // A single Agent cursor keeps Hive-family rows in one query execution. Inceptor // rejects the generic LIMIT/OFFSET form, just like the other Agent cursor paths. // Re-running LIMIT/OFFSET pages is unstable for tables without a unique key. - let use_hive_server_cursor = - matches!(source_db_type, DatabaseType::Kyuubi | DatabaseType::Impala | DatabaseType::Transwarp); + let use_hive_server_cursor = uses_agent_transfer_cursor(source_db_type); let hive_server_transfer_sql = use_hive_server_cursor.then(|| { transfer_cursor_sql( &col_names, diff --git a/crates/dbx-core/src/data/transfer/db2.rs b/crates/dbx-core/src/data/transfer/db2.rs new file mode 100644 index 0000000000..6d1fed15eb --- /dev/null +++ b/crates/dbx-core/src/data/transfer/db2.rs @@ -0,0 +1,133 @@ +use super::{quote_string_literal, should_copy_data, DatabaseType, TransferMode, TransferRequest}; + +const LITERAL_CHUNK_BYTES: usize = 16_000; + +pub(super) fn map_column_type(source_type: &str, source_db: &DatabaseType) -> Option { + let normalized = source_type.trim().to_ascii_lowercase(); + let type_name = normalized.split('(').next().unwrap_or("").trim(); + let base = match type_name { + "character varying" => "varchar", + "character large object" => "clob", + "binary varying" => "varbinary", + "binary large object" => "blob", + _ => type_name.split(' ').next().unwrap_or(""), + }; + let parameters = normalized.split_once('(').and_then(|(_, rest)| rest.split_once(')')).map(|(value, _)| value); + let unsigned = normalized.split_whitespace().any(|part| part == "unsigned"); + if super::is_sqlite_transfer_dialect(source_db) && normalized.contains("int") { + return Some("BIGINT".to_string()); + } + let mapped = match base { + "tinyint" => "SMALLINT".to_string(), + "smallint" if unsigned => "INTEGER".to_string(), + "int" | "integer" if unsigned => "BIGINT".to_string(), + "bigint" if unsigned => "DECIMAL(20,0)".to_string(), + "bigserial" => "BIGINT".to_string(), + "smallserial" => "SMALLINT".to_string(), + "datetime" | "datetime2" | "smalldatetime" => match parameters { + Some(precision) => format!("TIMESTAMP({precision})"), + None => "TIMESTAMP".to_string(), + }, + "bit" if *source_db == DatabaseType::SqlServer => "BOOLEAN".to_string(), + "bit" | "varbit" => "CLOB".to_string(), + "varchar" | "nvarchar" | "varchar2" => match parameters.and_then(|length| length.trim().parse::().ok()) { + Some(length) if (1..=32672).contains(&length) => format!("VARCHAR({length})"), + _ => "CLOB".to_string(), + }, + "clob" => "CLOB".to_string(), + "varbinary" | "blob" => "BLOB".to_string(), + "decimal" | "numeric" | "number" => match parameters { + Some(parameters) => { + let precision = parameters.split(',').next()?.trim().parse::().ok()?; + if precision > 31 { + "CLOB".to_string() + } else { + format!("DECIMAL({parameters})") + } + } + None => "CLOB".to_string(), + }, + _ => return None, + }; + Some(mapped) +} + +pub(super) fn numeric_literal(value: &str, column_type: Option<&str>) -> String { + let normalized = column_type.unwrap_or("").trim().to_ascii_lowercase(); + let base = normalized.split('(').next().unwrap_or("").trim(); + if matches!(base, "decimal" | "numeric" | "number") { + let precision = normalized + .split_once('(') + .and_then(|(_, parameters)| parameters.split([',', ')']).next()) + .and_then(|precision| precision.trim().parse::().ok()); + if precision.is_none_or(|precision| precision > 31) { + return quote_string_literal(value); + } + } + value.to_string() +} + +pub(super) fn string_literal(value: &str, column_type: Option<&str>) -> String { + if column_type.is_some_and(super::is_binary_transfer_column_type) { + if let Some(hex) = value.strip_prefix("0x").or_else(|| value.strip_prefix("0X")) { + if hex.len() % 2 == 0 && hex.bytes().all(|byte| byte.is_ascii_hexdigit()) { + if hex.len() <= LITERAL_CHUNK_BYTES { + return format!("BX'{hex}'"); + } + return hex + .as_bytes() + .chunks(LITERAL_CHUNK_BYTES) + .map(|chunk| format!("BLOB(BX'{}')", std::str::from_utf8(chunk).expect("validated hexadecimal"))) + .collect::>() + .join(" || "); + } + } + } + if value.len() <= LITERAL_CHUNK_BYTES { + return quote_string_literal(value); + } + let mut chunks = Vec::new(); + let mut remaining = value; + while !remaining.is_empty() { + let mut end = remaining.len().min(LITERAL_CHUNK_BYTES); + while !remaining.is_char_boundary(end) { + end -= 1; + } + chunks.push(format!("CLOB({})", quote_string_literal(&remaining[..end]))); + remaining = &remaining[end..]; + } + chunks.join(" || ") +} + +pub(super) fn validate_request(request: &TransferRequest) -> Result<(), String> { + if request.drop_target_before_create { + return Err("DB2 transfer does not support rebuilding target tables".to_string()); + } + if should_copy_data(&request.content) && request.mode == TransferMode::Upsert { + return Err("DB2 transfer does not support upsert; choose append or overwrite".to_string()); + } + Ok(()) +} + +pub(super) fn generated_columns_sql(schema: &str, table: &str, columns: &[String]) -> String { + let schema = if schema.trim().is_empty() { "CURRENT SCHEMA".to_string() } else { quote_string_literal(schema) }; + let columns = columns.iter().map(|column| quote_string_literal(column)).collect::>().join(", "); + format!( + "SELECT COLNAME FROM SYSCAT.COLUMNS WHERE TABSCHEMA = {schema} AND TABNAME = {} AND GENERATED = 'A' AND COLNAME IN ({columns})", + quote_string_literal(table) + ) +} + +pub(super) fn validate_generated_columns(rows: &[Vec]) -> Result<(), String> { + if rows.is_empty() { + return Ok(()); + } + let columns = rows.iter().filter_map(|row| row.first()?.as_str()).collect::>().join(", "); + Err(format!( + "DB2 transfer cannot write GENERATED ALWAYS target columns ({columns}); target data has not been cleared" + )) +} + +#[cfg(test)] +#[path = "db2_tests.rs"] +mod tests; diff --git a/crates/dbx-core/src/data/transfer/db2_tests.rs b/crates/dbx-core/src/data/transfer/db2_tests.rs new file mode 100644 index 0000000000..183e5797af --- /dev/null +++ b/crates/dbx-core/src/data/transfer/db2_tests.rs @@ -0,0 +1,210 @@ +use super::*; +use crate::data::transfer; +use serde_json::json; + +fn column(name: &str, data_type: &str) -> crate::db::ColumnInfo { + crate::db::ColumnInfo { + name: name.to_string(), + data_type: data_type.to_string(), + is_nullable: true, + ..Default::default() + } +} + +fn request() -> TransferRequest { + serde_json::from_value(json!({ + "transferId": "db2-transfer", "sourceConnectionId": "source", + "sourceDatabase": "source", "sourceSchema": "public", + "targetConnectionId": "target", "targetDatabase": "SAMPLE", "targetSchema": "APP", + "tables": ["EVENTS"], "createTable": true, "content": "structureAndData", + "mode": "append", "batchSize": 1000 + })) + .unwrap() +} + +#[test] +fn db2_transfer_maps_common_cross_database_types() { + for (source, source_type, expected) in [ + (DatabaseType::Mysql, "tinyint(1)", "SMALLINT"), + (DatabaseType::Mysql, "smallint unsigned", "INTEGER"), + (DatabaseType::Mysql, "int unsigned", "BIGINT"), + (DatabaseType::Mysql, "bigint unsigned", "DECIMAL(20,0)"), + (DatabaseType::Mysql, "decimal(20,4) unsigned", "DECIMAL(20,4)"), + (DatabaseType::Mysql, "datetime", "TIMESTAMP"), + (DatabaseType::SqlServer, "datetime2(7)", "TIMESTAMP(7)"), + (DatabaseType::SqlServer, "nvarchar(max)", "CLOB"), + (DatabaseType::SqlServer, "bit", "BOOLEAN"), + (DatabaseType::Postgres, "bigserial", "BIGINT"), + (DatabaseType::Postgres, "numeric", "CLOB"), + (DatabaseType::Postgres, "numeric(38,8)", "CLOB"), + (DatabaseType::Postgres, "varchar(65535)", "CLOB"), + (DatabaseType::Postgres, "varchar(64)", "VARCHAR(64)"), + (DatabaseType::Postgres, "varchar", "CLOB"), + (DatabaseType::Postgres, "character varying(64)", "VARCHAR(64)"), + (DatabaseType::H2, "character large object", "CLOB"), + (DatabaseType::H2, "binary varying(128)", "BLOB"), + (DatabaseType::Postgres, "text", "CLOB"), + (DatabaseType::Postgres, "jsonb", "CLOB"), + (DatabaseType::Postgres, "bytea", "BLOB"), + (DatabaseType::Postgres, "uuid", "VARCHAR(36)"), + (DatabaseType::Postgres, "varbit(2048)", "CLOB"), + (DatabaseType::Sqlite, "tinyint", "BIGINT"), + ] { + assert_eq!(transfer::map_column_type(source_type, &source, &DatabaseType::Db2), expected, "{source_type}"); + } + assert_eq!(transfer::map_column_type("DECIMAL(20,4)", &DatabaseType::Db2, &DatabaseType::Db2), "DECIMAL(20,4)"); + assert_eq!(transfer::map_column_type("tinyint", &DatabaseType::Mysql, &DatabaseType::Postgres), "SMALLINT"); +} + +#[test] +fn db2_transfer_creates_schema_qualified_tables_with_primary_keys() { + let columns = [ + crate::db::ColumnInfo { is_nullable: false, is_primary_key: true, ..column("ID", "bigint") }, + column("BODY", "text"), + column("CREATED_AT", "datetime"), + ]; + let sql = transfer::generate_create_table_ddl( + &columns, + "EVENTS", + "source", + "APP", + &DatabaseType::Db2, + &DatabaseType::Mysql, + None, + None, + ); + assert_eq!( + sql, + "CREATE TABLE \"APP\".\"EVENTS\" (\n \"ID\" BIGINT NOT NULL,\n \"BODY\" CLOB,\n \"CREATED_AT\" TIMESTAMP,\n PRIMARY KEY (\"ID\")\n)" + ); +} + +#[test] +fn db2_transfer_uses_a_single_agent_cursor_instead_of_limit_offset() { + assert!(transfer::uses_agent_transfer_cursor(&DatabaseType::Db2)); + assert!(transfer::uses_agent_transfer_cursor(&DatabaseType::Impala)); + assert!(!transfer::uses_agent_transfer_cursor(&DatabaseType::Postgres)); + let sql = transfer::transfer_cursor_sql( + &["ID".to_string(), "BODY".to_string()], + "EVENTS", + "APP", + &DatabaseType::Db2, + None, + ); + assert_eq!(sql, "SELECT \"ID\", \"BODY\" FROM \"APP\".\"EVENTS\""); +} + +#[test] +fn db2_transfer_rewrites_reused_ddl_without_changing_literals_or_comments() { + let sql = r#"CREATE TABLE "SOURCE"."EVENTS" ( + "ID" INTEGER NOT NULL PRIMARY KEY, + "NOTE" VARCHAR(80) DEFAULT '"SOURCE"."EVENTS"', + "part""SOURCE"".value" INTEGER +); +COMMENT ON TABLE "SOURCE"."EVENTS" IS 'source table'; +-- keep "SOURCE"."EVENTS" in comments"#; + let rewritten = transfer::rewrite_transfer_source_table_ddl( + sql, + "SOURCE", + "TARGET", + &DatabaseType::Db2, + &DatabaseType::Db2, + "EVENTS", + "EVENTS", + ) + .unwrap(); + assert!(rewritten.starts_with("CREATE TABLE \"TARGET\".\"EVENTS\"")); + assert!(rewritten.contains("COMMENT ON TABLE \"TARGET\".\"EVENTS\"")); + assert!(rewritten.contains("DEFAULT '\"SOURCE\".\"EVENTS\"'")); + assert!(rewritten.contains("\"part\"\"SOURCE\"\".value\"")); + assert!(rewritten.contains("-- keep \"SOURCE\".\"EVENTS\" in comments")); +} + +#[test] +fn db2_transfer_writes_multiple_rows_without_changing_backslashes_or_binary_values() { + let sql = transfer::generate_insert_typed( + &["PATH".to_string(), "PAYLOAD".to_string()], + &[Some("CLOB".to_string()), Some("BLOB".to_string())], + &[vec![json!(r"C:\data\O'Brien"), json!("0x00ff")], vec![json!("中文"), json!("0x")]], + "EVENTS", + "APP", + &DatabaseType::Db2, + None, + ); + assert_eq!(sql, "INSERT INTO \"APP\".\"EVENTS\" (\"PATH\", \"PAYLOAD\") VALUES\n('C:\\data\\O''Brien', BX'00ff'),\n('中文', BX'')"); + assert_eq!(transfer::escape_value_typed(&json!("0xFF"), &DatabaseType::Db2, Some("VARCHAR(20)")), "'0xFF'"); + assert_eq!(transfer::escape_value_typed(&json!(true), &DatabaseType::Db2, Some("BOOLEAN")), "TRUE"); + assert_eq!(transfer::escape_value(&json!(null), &DatabaseType::Db2), "NULL"); +} + +#[test] +fn db2_transfer_preserves_json_and_splits_large_lob_literals() { + let numeric_text = "12345678901234567890123456789012.123456"; + let numeric: serde_json::Value = serde_json::from_str(numeric_text).unwrap(); + assert_eq!( + transfer::escape_value_typed(&numeric, &DatabaseType::Db2, Some("NUMERIC(38,6)")), + quote_string_literal(numeric_text) + ); + assert_eq!(transfer::escape_value_typed(&json!(123), &DatabaseType::Db2, Some("INTEGER")), "123"); + let value = json!({ "path": r"C:\data", "values": [1, 2] }); + assert_eq!(transfer::escape_value(&value, &DatabaseType::Db2), quote_string_literal(&value.to_string())); + assert_eq!(transfer::escape_value(&json!([1, 2]), &DatabaseType::Db2), "'[1,2]'"); + let text = "中'\\".repeat(8000); + let literal = string_literal(&text, Some("CLOB")); + assert!(literal.starts_with("CLOB('")); + let restored = literal + .split(" || ") + .map(|part| { + let quoted = part.strip_prefix("CLOB('").unwrap().strip_suffix("')").unwrap(); + assert!(quoted.len() < 32672); + quoted.replace("''", "'") + }) + .collect::(); + assert_eq!(restored, text); + let binary = format!("0x{}", "ff".repeat(16000)); + assert_eq!( + string_literal(&binary, Some("BLOB")), + format!("BLOB(BX'{}') || BLOB(BX'{}')", "ff".repeat(8000), "ff".repeat(8000)) + ); +} + +#[test] +fn db2_transfer_overwrite_uses_native_truncate_without_changing_other_engines() { + assert_eq!( + transfer::transfer_clear_table_sql("EVENTS", "APP", &DatabaseType::Db2, None), + "TRUNCATE TABLE \"APP\".\"EVENTS\" IMMEDIATE" + ); + assert_eq!( + transfer::transfer_clear_table_sql("events", "public", &DatabaseType::Postgres, None), + "TRUNCATE TABLE \"public\".\"events\"" + ); + assert_eq!( + transfer::transfer_clear_table_sql("events", "main", &DatabaseType::Sqlite, None), + "DELETE FROM \"main\".\"events\"" + ); +} + +#[test] +fn db2_transfer_rejects_unimplemented_modes_before_writing() { + let mut request = request(); + assert!(validate_request(&request).is_ok()); + request.mode = TransferMode::Overwrite; + assert!(validate_request(&request).is_ok()); + request.mode = TransferMode::Upsert; + assert!(validate_request(&request).unwrap_err().contains("does not support upsert")); + request.content = transfer::TransferContent::StructureOnly; + assert!(validate_request(&request).is_ok()); + request.drop_target_before_create = true; + assert!(validate_request(&request).unwrap_err().contains("does not support rebuilding")); +} + +#[test] +fn db2_transfer_checks_generated_always_columns_before_overwrite() { + assert!(validate_generated_columns(&[]).is_ok()); + assert!(validate_generated_columns(&[vec![json!("ID")]]).unwrap_err().contains("target data has not been cleared")); + assert_eq!( + generated_columns_sql("APP'S", "EVENTS", &["ID".to_string(), "O'Brien".to_string()]), + "SELECT COLNAME FROM SYSCAT.COLUMNS WHERE TABSCHEMA = 'APP''S' AND TABNAME = 'EVENTS' AND GENERATED = 'A' AND COLNAME IN ('ID', 'O''Brien')" + ); + assert!(generated_columns_sql("", "EVENTS", &["ID".to_string()]).contains("TABSCHEMA = CURRENT SCHEMA")); +} diff --git a/plugins/connection-types/db2.yaml b/plugins/connection-types/db2.yaml index 684784e355..c115878016 100644 --- a/plugins/connection-types/db2.yaml +++ b/plugins/connection-types/db2.yaml @@ -26,10 +26,10 @@ capabilities: tableDataEdit: true tableStructureEdit: false tableImport: true - dataTransfer: false + dataTransfer: true sqlFileExecution: true databaseCreate: false - fieldLineage: false + fieldLineage: true sqlExplain: false userAdmin: false driverManagement: true