fix(postgres): preserve owned sequences during transfer

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<String>,
+    owner_column: Option<String>,
+}
+
+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<bool, String> {
+    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<Vec<PostgresOwnedSequence>, 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<Vec<PostgresSequenceSnapshot>, 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<String>>(1),
+            owner_column: row.get::<_, Option<String>>(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<Vec<PostgresOwnedSequence>, 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::<HashMap<_, _>>();
+
+    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(