diff --git a/apps/desktop/src/components/transfer/__tests__/ObjectSelectionTree.spec.ts b/apps/desktop/src/components/transfer/__tests__/ObjectSelectionTree.spec.ts index b681fcbaf0..1ce83562ae 100644 --- a/apps/desktop/src/components/transfer/__tests__/ObjectSelectionTree.spec.ts +++ b/apps/desktop/src/components/transfer/__tests__/ObjectSelectionTree.spec.ts @@ -28,7 +28,7 @@ vi.mock("@/components/ui/button", () => ({ Button: passthrough("button") })); import ObjectSelectionTree from "../ObjectSelectionTree.vue"; -type Kind = "TABLE" | "VIEW" | "FUNCTION"; +type Kind = "TABLE" | "VIEW" | "FUNCTION" | "SEQUENCE"; interface TreeGroup { kind: Kind; @@ -39,6 +39,12 @@ interface TreeGroup { const VIEWS: TreeGroup = { kind: "VIEW", label: "Views", items: ["v1", "v2", "v3"] }; const FUNCTIONS: TreeGroup = { kind: "FUNCTION", label: "Fns", items: ["f1", "f2"] }; const TABLES: TreeGroup = { kind: "TABLE", label: "Tables", items: ["t1", "t2"] }; +const POSTGRES_FUNCTIONS: TreeGroup = { + kind: "FUNCTION", + label: "Fns", + items: ["_st_beststride", "_st_coveredby", "box", "box2d"], +}; +const POSTGRES_SEQUENCES: TreeGroup = { kind: "SEQUENCE", label: "Sequences", items: ["biz_banner_id_seq"] }; function mountTree(init: { groups?: TreeGroup[]; disabledGroups?: Kind[]; disabledHints?: Record; selection?: Record; search?: string }) { const { groups = [VIEWS, FUNCTIONS], disabledGroups = [], disabledHints = {}, selection = {}, search = "" } = init; @@ -127,6 +133,18 @@ describe("ObjectSelectionTree interaction", () => { expect(state.selection.VIEW).toEqual(["v1", "v3"]); }); + it("updates expanded PostgreSQL object lists immediately while searching", async () => { + const { container, app } = mountTree({ groups: [POSTGRES_FUNCTIONS, POSTGRES_SEQUENCES] }); + cleanup = () => app.unmount(); + + expect(container.querySelector('label[data-test="item-FUNCTION-box"]')).not.toBeNull(); + await typeSearch(container, "biz_ban"); + + expect(container.querySelector('label[data-test="item-FUNCTION-box"]')).toBeNull(); + expect(container.querySelector('label[data-test="item-SEQUENCE-biz_banner_id_seq"]')).not.toBeNull(); + expect(container.querySelector('[data-test="group-FUNCTION"]')?.textContent).toContain("无匹配"); + }); + it("deselect all only clears visible enabled selections, keeping hidden ones", async () => { const { state, container, app } = mountTree({ selection: { VIEW: ["v1", "v2", "v3"], FUNCTION: ["f1", "f2"] }, diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index eb1fa1b55e..19ffafaeb3 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -1826,6 +1826,19 @@ struct PostgresSequenceSnapshot { owner_column: Option, } +#[derive(Debug, Clone, PartialEq, Eq)] +struct PostgresTransferSequence { + name: String, + data_type: String, + start_value: String, + min_value: String, + max_value: String, + increment: String, + cycle: bool, + cache_value: String, + last_value: Option, +} + fn postgres_sequence_qualified_name(schema: &str, sequence_name: &str) -> String { if schema.trim().is_empty() { quote_identifier(sequence_name, &DatabaseType::Postgres) @@ -1838,6 +1851,31 @@ fn postgres_sequence_qualified_name(schema: &str, sequence_name: &str) -> String } } +fn generate_postgres_transfer_sequence_create_ddl(sequence: &PostgresTransferSequence, schema: &str) -> String { + let qualified_name = postgres_sequence_qualified_name(schema, &sequence.name); + let cycle = if sequence.cycle { "CYCLE" } else { "NO CYCLE" }; + format!( + "CREATE SEQUENCE IF NOT EXISTS {qualified_name}\n AS {data_type}\n START WITH {start_value}\n INCREMENT BY {increment}\n MINVALUE {min_value}\n MAXVALUE {max_value}\n CACHE {cache_value}\n {cycle}", + data_type = sequence.data_type, + start_value = sequence.start_value, + increment = sequence.increment, + min_value = sequence.min_value, + max_value = sequence.max_value, + cache_value = sequence.cache_value, + ) +} + +fn generate_postgres_transfer_sequence_setval_sql(sequence: &PostgresTransferSequence, schema: &str) -> Option { + let last_value = sequence.last_value.as_deref()?.trim(); + if last_value.is_empty() { + return None; + } + Some(format!( + "SELECT setval({}, {last_value}, true)", + quote_postgres_string_literal(&postgres_sequence_qualified_name(schema, &sequence.name)) + )) +} + /// Reuse an existing target sequence only when it is already bound to the same /// target table column; otherwise the later `OWNED BY` rebind would silently /// change unrelated objects. @@ -4353,6 +4391,84 @@ async fn get_postgres_sequence_snapshots_for_transfer( .collect()) } +fn postgres_selected_sequences_sql(schema: &str, names: &[String]) -> Option { + if names.is_empty() { + return None; + } + let name_list = names.iter().map(|name| quote_string_literal(name)).collect::>().join(", "); + Some(format!( + "SELECT c.relname, \ + COALESCE(format_type(s.seqtypid, NULL), 'bigint'), \ + COALESCE(s.seqstart::text, '1'), \ + COALESCE(s.seqmin::text, '1'), \ + COALESCE(s.seqmax::text, '9223372036854775807'), \ + COALESCE(s.seqincrement::text, '1'), \ + CASE WHEN COALESCE(s.seqcycle, false) THEN 'true' ELSE 'false' END, \ + COALESCE(s.seqcache::text, '1'), \ + pg_sequence_last_value(c.oid)::text \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + LEFT JOIN pg_catalog.pg_sequence s ON s.seqrelid = c.oid \ + WHERE c.relkind = 'S' AND n.nspname = {} AND c.relname IN ({name_list}) \ + ORDER BY c.relname", + quote_string_literal(schema) + )) +} + +async fn get_postgres_selected_sequences_for_transfer( + state: &AppState, + pool_key: &str, + schema: &str, + names: &[String], +) -> Result, String> { + let Some(sql) = postgres_selected_sequences_sql(schema, names) else { + return Ok(Vec::new()); + }; + Ok(execute_on_pool(state, pool_key, &sql) + .await? + .rows + .into_iter() + .filter_map(|row| { + Some(PostgresTransferSequence { + name: json_string_cell(&row, 0)?, + data_type: json_string_cell(&row, 1)?, + start_value: json_string_cell(&row, 2)?, + min_value: json_string_cell(&row, 3)?, + max_value: json_string_cell(&row, 4)?, + increment: json_string_cell(&row, 5)?, + cycle: json_string_cell(&row, 6).as_deref() == Some("true"), + cache_value: json_string_cell(&row, 7)?, + last_value: json_string_cell(&row, 8), + }) + }) + .collect()) +} + +async fn get_existing_postgres_sequence_names_for_transfer( + state: &AppState, + pool_key: &str, + schema: &str, + names: &[String], +) -> Result, String> { + if names.is_empty() { + return Ok(HashSet::new()); + } + let name_list = names.iter().map(|name| quote_string_literal(name)).collect::>().join(", "); + let sql = format!( + "SELECT c.relname \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE c.relkind = 'S' AND n.nspname = {} AND c.relname IN ({name_list})", + quote_string_literal(schema) + ); + Ok(execute_on_pool(state, pool_key, &sql) + .await? + .rows + .into_iter() + .filter_map(|row| json_string_cell(&row, 0)) + .collect()) +} + /// Create owned PostgreSQL sequences before executing reused table DDL because /// serial defaults still reference `nextval('...')` in `CREATE TABLE`. async fn prepare_postgres_owned_sequences_for_transfer( @@ -4455,6 +4571,21 @@ pub fn selected_object_names(selections: &[TransferObjectSelection], kind: &Tran selections.iter().filter(|s| &s.object_type == kind).flat_map(|s| s.names.clone()).collect::>() } +fn selected_postgres_sequence_names(request: &TransferRequest) -> Vec { + let mut names = selected_object_names(&request.objects, &TransferObjectKind::Sequence); + names.sort(); + names.dedup(); + names +} + +fn postgres_transfer_relation_names(request: &TransferRequest) -> Vec { + let mut names = request.tables.clone(); + names.extend(selected_postgres_sequence_names(request)); + names.sort(); + names.dedup(); + names +} + /// Whether a kind participates in a transfer. An empty selection is the legacy /// PG→PG default: every kind participates (views, functions, triggers, /// materialized views are all transferred). Once the caller explicitly selects @@ -5288,12 +5419,13 @@ pub async fn preview_transfer_ownership( return Ok(TransferOwnershipPreview { missing_owners: Vec::new(), target_owner: String::new() }); } + let relation_names = postgres_transfer_relation_names(request); let statements = get_postgres_ownership_statements_for_transfer( state, source_pool_key, &request.source_schema, &request.target_schema, - &request.tables, + &relation_names, ) .await?; let roles = distinct_postgres_ownership_roles(&statements); @@ -6334,7 +6466,26 @@ where get_postgres_extension_sources_for_transfer(state, source_pool_key, &request.source_schema).await?; let enum_types = get_postgres_enum_sources_for_transfer(state, source_pool_key, &request.source_schema).await?; let domains = get_postgres_domain_sources_for_transfer(state, source_pool_key, &request.source_schema).await?; - let total_steps = extensions.len() + enum_types.len() + domains.len(); + let selected_sequence_names = selected_postgres_sequence_names(request); + let selected_sequences = get_postgres_selected_sequences_for_transfer( + state, + source_pool_key, + &request.source_schema, + &selected_sequence_names, + ) + .await?; + let existing_sequence_names = get_existing_postgres_sequence_names_for_transfer( + state, + target_pool_key, + &request.target_schema, + &selected_sequence_names, + ) + .await?; + let selected_sequences = selected_sequences + .into_iter() + .filter(|sequence| !existing_sequence_names.contains(&sequence.name)) + .collect::>(); + let total_steps = extensions.len() + enum_types.len() + domains.len() + selected_sequences.len(); let table_index = 0; let mut completed_steps = 0_u64; @@ -6401,6 +6552,36 @@ where .map_err(|e| format!("Failed to create PostgreSQL domain {}: {e}", domain.domain_name))?; } + for sequence in selected_sequences { + if is_cancelled(&request.transfer_id).await { + return Err("Cancelled".to_string()); + } + completed_steps += 1; + progress_callback(TransferProgress { + transfer_id: request.transfer_id.clone(), + table: format!("sequence: {}", sequence.name), + table_index, + total_tables: request.tables.len(), + rows_transferred: completed_steps, + total_rows: Some(total_steps as u64), + status: TransferStatus::Running, + error: None, + terminal: false, + }); + execute_on_pool( + state, + target_pool_key, + &generate_postgres_transfer_sequence_create_ddl(&sequence, &request.target_schema), + ) + .await + .map_err(|e| format!("Failed to create PostgreSQL sequence {}: {e}", sequence.name))?; + if let Some(setval_sql) = generate_postgres_transfer_sequence_setval_sql(&sequence, &request.target_schema) { + execute_on_pool(state, target_pool_key, &setval_sql) + .await + .map_err(|e| format!("Failed to restore PostgreSQL sequence {} value: {e}", sequence.name))?; + } + } + Ok(()) } @@ -6453,6 +6634,7 @@ where &request.tables, ) .await?; + let relation_names = postgres_transfer_relation_names(request); let ownership_statements = if matches!(request.ownership_policy, TransferOwnershipPolicy::Skip) { Vec::new() } else { @@ -6461,7 +6643,7 @@ where source_pool_key, &request.source_schema, &request.target_schema, - &request.tables, + &relation_names, ) .await? }; @@ -6483,7 +6665,7 @@ where source_pool_key, &request.source_schema, &request.target_schema, - &request.tables, + &relation_names, ) .await?; let materialized_view_step_count = materialized_views @@ -7622,6 +7804,53 @@ mod tests { assert_eq!(filtered.len(), 1); assert_eq!(filtered[0].name, "v1"); } + + #[test] + fn selected_postgres_sequences_are_prepared_without_changing_default_requests() { + let mut request = test_transfer_request(vec!["biz_banner"]); + assert!(selected_postgres_sequence_names(&request).is_empty()); + + request.objects = vec![ + TransferObjectSelection { object_type: TransferObjectKind::Table, names: vec!["biz_banner".into()] }, + TransferObjectSelection { + object_type: TransferObjectKind::Sequence, + names: vec!["biz_banner_id_seq".into(), "biz_banner_id_seq".into()], + }, + ]; + + assert_eq!(selected_postgres_sequence_names(&request), vec!["biz_banner_id_seq"]); + assert_eq!(postgres_transfer_relation_names(&request), vec!["biz_banner", "biz_banner_id_seq"]); + let sql = postgres_selected_sequences_sql("public", &selected_postgres_sequence_names(&request)).unwrap(); + assert!(sql.contains("c.relname IN ('biz_banner_id_seq')")); + assert!(sql.contains("pg_sequence_last_value(c.oid)::text")); + } + + #[test] + fn postgres_selected_sequence_ddl_preserves_definition_and_value() { + let sequence = PostgresTransferSequence { + name: "biz_banner_id_seq".into(), + data_type: "bigint".into(), + start_value: "5".into(), + min_value: "-10".into(), + max_value: "999".into(), + increment: "2".into(), + cycle: true, + cache_value: "7".into(), + last_value: Some("41".into()), + }; + + assert_eq!( + generate_postgres_transfer_sequence_create_ddl(&sequence, "archive"), + "CREATE SEQUENCE IF NOT EXISTS \"archive\".\"biz_banner_id_seq\"\n AS bigint\n START WITH 5\n INCREMENT BY 2\n MINVALUE -10\n MAXVALUE 999\n CACHE 7\n CYCLE" + ); + assert_eq!( + generate_postgres_transfer_sequence_setval_sql(&sequence, "archive"), + Some("SELECT setval('\"archive\".\"biz_banner_id_seq\"', 41, true)".into()) + ); + + let never_called = PostgresTransferSequence { last_value: None, ..sequence }; + assert_eq!(generate_postgres_transfer_sequence_setval_sql(&never_called, "archive"), None); + } } mod transfer_content_mode_tests { diff --git a/crates/dbx-core/tests/live_postgres_transfer.rs b/crates/dbx-core/tests/live_postgres_transfer.rs index 87411ef250..4467e5fa49 100644 --- a/crates/dbx-core/tests/live_postgres_transfer.rs +++ b/crates/dbx-core/tests/live_postgres_transfer.rs @@ -4,7 +4,7 @@ use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; use dbx_core::storage::Storage; use dbx_core::transfer::{ get_db_type, transfer_postgres_schema_dependencies, transfer_postgres_schema_objects, transfer_table, TransferMode, - TransferOwnershipPolicy, TransferRequest, TransferTableNameCase, + TransferObjectKind, TransferObjectSelection, TransferOwnershipPolicy, TransferRequest, TransferTableNameCase, }; use serde_json::json; @@ -540,3 +540,189 @@ async fn live_postgres_transfer_skips_create_ddl_for_existing_target_table() { let _ = postgres::execute_batch(&target_pool, &[cleanup_sql[1].clone()]).await; let _ = std::fs::remove_dir_all(dir); } + +#[tokio::test] +#[ignore = "requires PostgreSQL URLs via DBX_LIVE_PG_TRANSFER_SOURCE_URL and DBX_LIVE_PG_TRANSFER_TARGET_URL"] +async fn live_postgres_transfer_creates_selected_sequence_before_referencing_table() { + let source_url = std::env::var("DBX_LIVE_PG_TRANSFER_SOURCE_URL").expect("DBX_LIVE_PG_TRANSFER_SOURCE_URL"); + let target_url = std::env::var("DBX_LIVE_PG_TRANSFER_TARGET_URL").unwrap_or_else(|_| source_url.clone()); + + let source_pool = postgres::connect(&source_url, std::time::Duration::from_secs(5)).await.unwrap(); + let target_pool = postgres::connect(&target_url, std::time::Duration::from_secs(5)).await.unwrap(); + let source_database = query_scalar(&source_pool, "SELECT current_database()").await.as_str().unwrap().to_string(); + let target_database = query_scalar(&target_pool, "SELECT current_database()").await.as_str().unwrap().to_string(); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let source_schema = format!("dbx_src_sequence_{}", &suffix[..8]); + let target_schema = format!("dbx_dst_sequence_{}", &suffix[..8]); + let cleanup_sql = [ + format!("DROP SCHEMA IF EXISTS \"{}\" CASCADE", source_schema), + format!("DROP SCHEMA IF EXISTS \"{}\" CASCADE", target_schema), + ]; + let _ = postgres::execute_batch(&source_pool, &[cleanup_sql[0].clone()]).await; + let _ = postgres::execute_batch(&target_pool, &[cleanup_sql[1].clone()]).await; + + postgres::execute_batch( + &source_pool, + &[ + format!("CREATE SCHEMA \"{}\"", source_schema), + format!( + "CREATE SEQUENCE \"{}\".\"biz_banner_id_seq\" AS bigint START WITH 5 INCREMENT BY 2 MINVALUE 1 MAXVALUE 999 CACHE 1 CYCLE", + source_schema + ), + format!( + "CREATE TABLE \"{}\".\"biz_banner\" (\ + \"id\" bigint DEFAULT nextval('\"{}\".\"biz_banner_id_seq\"'::regclass) PRIMARY KEY,\ + \"name\" text NOT NULL\ + )", + source_schema, source_schema + ), + format!("INSERT INTO \"{}\".\"biz_banner\" (\"name\") VALUES ('first')", source_schema), + format!("SELECT setval('\"{}\".\"biz_banner_id_seq\"', 41, true)", source_schema), + format!("INSERT INTO \"{}\".\"biz_banner\" (\"name\") VALUES ('second')", source_schema), + format!("GRANT USAGE ON SEQUENCE \"{}\".\"biz_banner_id_seq\" TO PUBLIC", source_schema), + ], + ) + .await + .unwrap(); + + let dir = std::env::temp_dir().join(format!("dbx-live-sequence-transfer-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let source_connection_id = "live-sequence-source"; + let target_connection_id = "live-sequence-target"; + let source_pool_key = format!("{source_connection_id}:{source_database}"); + let target_pool_key = format!("{target_connection_id}:{target_database}"); + state.connections.write().await.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); + state.connections.write().await.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + state + .configs + .write() + .await + .insert(source_connection_id.to_string(), postgres_test_config(source_connection_id, &source_database)); + state + .configs + .write() + .await + .insert(target_connection_id.to_string(), postgres_test_config(target_connection_id, &target_database)); + + let request = TransferRequest { + transfer_id: format!("live-sequence-transfer-{suffix}"), + source_connection_id: source_connection_id.to_string(), + source_database: source_database.clone(), + source_schema: source_schema.clone(), + source_catalog: None, + target_connection_id: target_connection_id.to_string(), + target_database: target_database.clone(), + target_schema: target_schema.clone(), + target_catalog: None, + tables: vec!["biz_banner".to_string()], + create_table: true, + content: dbx_core::transfer::TransferContent::default(), + objects: vec![ + TransferObjectSelection { object_type: TransferObjectKind::Table, names: vec!["biz_banner".to_string()] }, + TransferObjectSelection { + object_type: TransferObjectKind::Sequence, + names: vec!["biz_banner_id_seq".to_string()], + }, + ], + mode: TransferMode::Append, + target_table_name_case: TransferTableNameCase::Preserve, + ownership_policy: TransferOwnershipPolicy::Preserve, + batch_size: 100, + }; + let source_db_type = get_db_type(&state, source_connection_id).await.unwrap(); + let target_db_type = get_db_type(&state, target_connection_id).await.unwrap(); + + transfer_postgres_schema_dependencies(&state, &request, &source_pool_key, &target_pool_key, |_| {}).await.unwrap(); + let transferred = transfer_table( + &state, + &request, + "biz_banner", + 0, + &source_db_type, + &target_db_type, + &source_pool_key, + &target_pool_key, + |_| {}, + ) + .await + .unwrap(); + let outcome = + transfer_postgres_schema_objects(&state, &request, &source_pool_key, &target_pool_key, |_| {}).await.unwrap(); + + assert_eq!(transferred, 2); + assert!(outcome.failed.is_empty()); + assert_eq!( + query_scalar(&target_pool, &format!("SELECT count(*) FROM \"{}\".\"biz_banner\"", target_schema)).await, + json!(2) + ); + assert_eq!( + query_scalar( + &target_pool, + &format!( + "SELECT format_type(s.seqtypid, NULL) || ':' || s.seqstart || ':' || s.seqincrement || ':' || s.seqmin || ':' || s.seqmax || ':' || s.seqcache || ':' || s.seqcycle \ + FROM pg_catalog.pg_sequence s \ + JOIN pg_catalog.pg_class c ON c.oid = s.seqrelid \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = '{}' AND c.relname = 'biz_banner_id_seq'", + target_schema + ) + ) + .await, + json!("bigint:5:2:1:999:1:true") + ); + assert_eq!( + query_scalar(&target_pool, &format!("SELECT nextval('\"{}\".\"biz_banner_id_seq\"'::regclass)", target_schema)) + .await, + json!(45) + ); + assert_eq!( + query_scalar( + &target_pool, + &format!( + "SELECT column_default LIKE '%{}%' AND column_default NOT LIKE '%{}%' \ + FROM information_schema.columns \ + WHERE table_schema = '{}' AND table_name = 'biz_banner' AND column_name = 'id'", + target_schema, source_schema, target_schema + ) + ) + .await, + json!(true) + ); + assert_eq!( + query_scalar( + &target_pool, + &format!( + "SELECT count(*) FROM pg_catalog.pg_depend d \ + JOIN pg_catalog.pg_class c ON c.oid = d.objid \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = '{}' AND c.relname = 'biz_banner_id_seq' AND d.deptype IN ('a', 'i')", + target_schema + ) + ) + .await, + json!(0) + ); + assert_eq!( + query_scalar( + &target_pool, + &format!( + "SELECT count(*) \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + JOIN LATERAL aclexplode(c.relacl) a ON true \ + WHERE n.nspname = '{}' AND c.relname = 'biz_banner_id_seq' \ + AND c.relkind = 'S' AND a.grantee = 0 AND a.privilege_type = 'USAGE'", + target_schema + ) + ) + .await, + json!(1) + ); + + let _ = postgres::execute_batch(&source_pool, &[cleanup_sql[0].clone()]).await; + let _ = postgres::execute_batch(&target_pool, &[cleanup_sql[1].clone()]).await; + let _ = std::fs::remove_dir_all(dir); +}