diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index 69c877c52..5a5eeb11c 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -607,6 +607,61 @@ fn generate_postgres_sequence_sync_sql(columns: &[db::ColumnInfo], table: &str, .collect() } +#[derive(Debug, Clone)] +struct PostgresOwnedSequence { + name: String, + owner_table: String, + owner_column: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct PostgresSequenceSnapshot { + name: String, + owner_table: Option, + owner_column: Option, +} + +fn postgres_sequence_qualified_name(schema: &str, sequence_name: &str) -> String { + if schema.trim().is_empty() { + quote_identifier(sequence_name, &DatabaseType::Postgres) + } else { + format!( + "{}.{}", + quote_identifier(schema, &DatabaseType::Postgres), + quote_identifier(sequence_name, &DatabaseType::Postgres) + ) + } +} + +/// 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. +fn validate_existing_postgres_sequence( + sequence: &PostgresOwnedSequence, + existing: Option<&PostgresSequenceSnapshot>, + schema: &str, +) -> Result { + let Some(existing) = existing else { + return Ok(true); + }; + + let owner_matches = match (existing.owner_table.as_deref(), existing.owner_column.as_deref()) { + (None, None) => true, + (Some(owner_table), Some(owner_column)) => { + owner_table == sequence.owner_table && owner_column == sequence.owner_column + } + _ => false, + }; + + if owner_matches { + return Ok(false); + } + + Err(format!( + "PostgreSQL sequence {} already exists with incompatible ownership", + postgres_sequence_qualified_name(schema, &sequence.name) + )) +} #[derive(Debug, Clone)] struct PostgresTriggerSource { table_name: String, @@ -2417,6 +2472,181 @@ async fn get_postgres_foreign_keys_for_transfer( db::postgres::list_foreign_keys(&pool, schema, table).await } +async fn get_postgres_owned_sequences_for_transfer( + state: &AppState, + pool_key: &str, + schema: &str, + tables: &[String], +) -> Result, String> { + if tables.is_empty() { + return Ok(Vec::new()); + } + + let pool = { + let connections = state.connections.read().await; + match connections.get(pool_key) { + Some(PoolKind::Postgres(pool)) => pool.clone(), + _ => return Ok(Vec::new()), + } + }; + let client = pool.get().await.map_err(|e| e.to_string())?; + let rows = client + .query( + "SELECT c.relname, \ + t.relname, \ + a.attname \ + FROM pg_class c \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + LEFT JOIN pg_sequence s ON s.seqrelid = c.oid \ + JOIN pg_depend d ON d.classid = 'pg_class'::regclass \ + AND d.objid = c.oid \ + AND d.refclassid = 'pg_class'::regclass \ + AND d.deptype IN ('a', 'i') \ + JOIN pg_class t ON t.oid = d.refobjid \ + JOIN pg_namespace tn ON tn.oid = t.relnamespace AND tn.nspname = n.nspname \ + JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = d.refobjsubid \ + WHERE c.relkind = 'S' AND n.nspname = $1 \ + ORDER BY t.relname, c.relname", + &[&schema], + ) + .await + .map_err(|e| e.to_string())?; + + let selected: HashSet<&str> = tables.iter().map(String::as_str).collect(); + Ok(rows + .iter() + .filter_map(|row| { + let owner_table = row.get::<_, String>(1); + if !selected.contains(owner_table.as_str()) { + return None; + } + Some(PostgresOwnedSequence { + name: row.get::<_, String>(0), + owner_table, + owner_column: row.get::<_, String>(2), + }) + }) + .collect()) +} + +async fn get_postgres_sequence_snapshots_for_transfer( + state: &AppState, + pool_key: &str, + schema: &str, +) -> Result, String> { + let pool = { + let connections = state.connections.read().await; + match connections.get(pool_key) { + Some(PoolKind::Postgres(pool)) => pool.clone(), + _ => return Ok(Vec::new()), + } + }; + let client = pool.get().await.map_err(|e| e.to_string())?; + let rows = client + .query( + "SELECT c.relname, \ + t.relname, \ + a.attname \ + FROM pg_class c \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + LEFT JOIN pg_sequence s ON s.seqrelid = c.oid \ + LEFT JOIN pg_depend d ON d.classid = 'pg_class'::regclass \ + AND d.objid = c.oid \ + AND d.refclassid = 'pg_class'::regclass \ + AND d.deptype IN ('a', 'i') \ + LEFT JOIN pg_class t ON t.oid = d.refobjid \ + LEFT JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = d.refobjsubid \ + WHERE c.relkind = 'S' AND n.nspname = $1 \ + ORDER BY c.relname", + &[&schema], + ) + .await + .map_err(|e| e.to_string())?; + + Ok(rows + .iter() + .map(|row| PostgresSequenceSnapshot { + name: row.get::<_, String>(0), + owner_table: row.get::<_, Option>(1), + owner_column: row.get::<_, Option>(2), + }) + .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( + state: &AppState, + request: &TransferRequest, + table: &str, + target_table: &str, + source_pool_key: &str, + target_pool_key: &str, + pg_compat_transfer: bool, + preserves_target_table_name: bool, + target_table_preexisting: bool, +) -> Result, String> { + if !(request.create_table && pg_compat_transfer && preserves_target_table_name && !target_table_preexisting) { + return Ok(Vec::new()); + } + + let owned_sequences = + get_postgres_owned_sequences_for_transfer(state, source_pool_key, &request.source_schema, &[table.to_string()]) + .await?; + if owned_sequences.is_empty() { + return Ok(Vec::new()); + } + + let existing_sequences = + get_postgres_sequence_snapshots_for_transfer(state, target_pool_key, &request.target_schema) + .await? + .into_iter() + .map(|sequence| (sequence.name.clone(), sequence)) + .collect::>(); + + for sequence in &owned_sequences { + let should_create = validate_existing_postgres_sequence( + sequence, + existing_sequences.get(&sequence.name), + &request.target_schema, + )?; + if should_create { + let create_sql = format!( + "CREATE SEQUENCE IF NOT EXISTS {}", + postgres_sequence_qualified_name(&request.target_schema, &sequence.name) + ); + execute_on_pool(state, target_pool_key, &create_sql) + .await + .map_err(|e| format!("Failed to create PostgreSQL sequence for {target_table}: {e}"))?; + } + } + + Ok(owned_sequences) +} + +/// Bind created or reused sequences after the table exists so +/// `pg_get_serial_sequence(...)` can find them during later sequence sync. +async fn bind_postgres_owned_sequences_for_transfer( + state: &AppState, + request: &TransferRequest, + target_table: &str, + target_pool_key: &str, + owned_sequences: &[PostgresOwnedSequence], +) -> Result<(), String> { + for sequence in owned_sequences { + let owner_sql = format!( + "ALTER SEQUENCE {} OWNED BY {}.{}", + postgres_sequence_qualified_name(&request.target_schema, &sequence.name), + qualified_table(&sequence.owner_table, &request.target_schema, &DatabaseType::Postgres), + quote_identifier(&sequence.owner_column, &DatabaseType::Postgres) + ); + execute_on_pool(state, target_pool_key, &owner_sql) + .await + .map_err(|e| format!("Failed to bind PostgreSQL sequence for {target_table}: {e}"))?; + } + Ok(()) +} + async fn get_postgres_schema_object_sources_for_transfer( state: &AppState, pool_key: &str, @@ -3263,6 +3493,18 @@ where .await .map_err(|e| format!("Failed to ensure schema exists: {e}"))?; } + let owned_sequences = prepare_postgres_owned_sequences_for_transfer( + state, + request, + table, + &target_table, + source_pool_key, + target_pool_key, + pg_compat_transfer, + preserves_target_table_name, + target_table_preexisting, + ) + .await?; let can_reuse_source_ddl = can_reuse_source_table_ddl(source_db_type, target_db_type, preserves_target_table_name); let ddl = if can_reuse_source_ddl { @@ -3329,6 +3571,14 @@ where log::warn!("[transfer] failed to set column comment for {}: {}", target_table, e); } } + bind_postgres_owned_sequences_for_transfer( + state, + request, + &target_table, + target_pool_key, + &owned_sequences, + ) + .await?; } } @@ -4694,6 +4944,119 @@ mod tests { ); } + #[test] + fn postgres_transfer_owned_sequence_ddl_uses_precreate_and_post_bind_steps() { + let sequence = PostgresOwnedSequence { + name: "it_quick_entry_id_seq".to_string(), + owner_table: "it_quick_entry".to_string(), + owner_column: "id".to_string(), + }; + let create_sql = + format!("CREATE SEQUENCE IF NOT EXISTS {}", postgres_sequence_qualified_name("public", &sequence.name)); + let owner_sql = format!( + "ALTER SEQUENCE {} OWNED BY {}.{}", + postgres_sequence_qualified_name("public", &sequence.name), + qualified_table(&sequence.owner_table, "public", &DatabaseType::Postgres), + quote_identifier(&sequence.owner_column, &DatabaseType::Postgres) + ); + + assert_eq!(create_sql, "CREATE SEQUENCE IF NOT EXISTS \"public\".\"it_quick_entry_id_seq\"".to_string()); + assert_eq!( + owner_sql, + "ALTER SEQUENCE \"public\".\"it_quick_entry_id_seq\" OWNED BY \"public\".\"it_quick_entry\".\"id\"" + .to_string() + ); + } + + #[test] + fn postgres_owned_sequence_state_detects_conflicting_existing_sequence() { + let source = PostgresOwnedSequence { + name: "it_quick_entry_id_seq".to_string(), + owner_table: "it_quick_entry".to_string(), + owner_column: "id".to_string(), + }; + + let conflicting = PostgresSequenceSnapshot { + name: "it_quick_entry_id_seq".to_string(), + owner_table: Some("other_table".to_string()), + owner_column: Some("id".to_string()), + }; + + let error = validate_existing_postgres_sequence(&source, Some(&conflicting), "archive").unwrap_err(); + + assert!(error.contains("\"archive\".\"it_quick_entry_id_seq\"")); + assert!(error.contains("already exists with incompatible ownership")); + } + + #[test] + fn postgres_transfer_reused_table_ddl_preserves_serial_sequence_dependencies() { + let columns = vec![ + db::ColumnInfo { + name: "id".to_string(), + data_type: "integer".to_string(), + is_nullable: false, + column_default: Some("nextval('public.it_quick_entry_id_seq'::regclass)".to_string()), + is_primary_key: true, + extra: None, + comment: None, + numeric_precision: None, + numeric_scale: None, + character_maximum_length: None, + }, + db::ColumnInfo { + name: "name".to_string(), + data_type: "text".to_string(), + is_nullable: false, + column_default: None, + is_primary_key: false, + extra: None, + comment: None, + numeric_precision: None, + numeric_scale: None, + character_maximum_length: None, + }, + ]; + let source_ddl = crate::schema::render_postgres_table_ddl("public", "it_quick_entry", &columns, &[], &[]); + let rewritten = rewrite_transfer_source_table_ddl( + &source_ddl, + "public", + "archive", + &DatabaseType::Postgres, + &DatabaseType::Postgres, + ); + let sequence = PostgresOwnedSequence { + name: "it_quick_entry_id_seq".to_string(), + owner_table: "it_quick_entry".to_string(), + owner_column: "id".to_string(), + }; + let create_sql = + format!("CREATE SEQUENCE IF NOT EXISTS {}", postgres_sequence_qualified_name("archive", &sequence.name)); + let owner_sql = format!( + "ALTER SEQUENCE {} OWNED BY {}.{}", + postgres_sequence_qualified_name("archive", &sequence.name), + qualified_table(&sequence.owner_table, "archive", &DatabaseType::Postgres), + quote_identifier(&sequence.owner_column, &DatabaseType::Postgres) + ); + let sequence_sync_sql = generate_postgres_sequence_sync_sql(&columns, "it_quick_entry", "archive"); + + assert!(source_ddl.starts_with("CREATE TABLE \"public\".\"it_quick_entry\"")); + assert!(!source_ddl.contains("CREATE SEQUENCE")); + assert!(rewritten.contains("CREATE TABLE \"archive\".\"it_quick_entry\"")); + assert!(rewritten.contains("nextval('\"archive\".it_quick_entry_id_seq'::regclass)")); + assert_eq!(create_sql, "CREATE SEQUENCE IF NOT EXISTS \"archive\".\"it_quick_entry_id_seq\"".to_string()); + assert_eq!( + owner_sql, + "ALTER SEQUENCE \"archive\".\"it_quick_entry_id_seq\" OWNED BY \"archive\".\"it_quick_entry\".\"id\"" + .to_string() + ); + assert_eq!( + sequence_sync_sql, + vec![ + "SELECT setval(pg_get_serial_sequence('\"archive\".\"it_quick_entry\"', 'id'), GREATEST(COALESCE(MAX(\"id\"), 0), 1), MAX(\"id\") IS NOT NULL) FROM \"archive\".\"it_quick_entry\"".to_string() + ] + ); + } + #[test] fn postgres_routine_schema_rewrite_targets_destination_schema() { let rewritten = rewrite_postgres_routine_schema(