feat(schema-diff): preserve dialect-specific DDL semantics

crates/dbx-core/src/two_phase_commit.rs

@@ -0,0 +1,678 @@
+use async_trait::async_trait;
+use serde::{Deserialize, Serialize};
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+
+use crate::state_persistence::{StateBackend, StateMachine};
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum TransactionStatus {
+    Preparing,
+    Prepared,
+    Committing,
+    Committed,
+    RollingBack,
+    RolledBack,
+    Mixed,
+    Unknown,
+}
+
+impl TransactionStatus {
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            Self::Preparing => "preparing",
+            Self::Prepared => "prepared",
+            Self::Committing => "committing",
+            Self::Committed => "committed",
+            Self::RollingBack => "rolling_back",
+            Self::RolledBack => "rolled_back",
+            Self::Mixed => "mixed",
+            Self::Unknown => "unknown",
+        }
+    }
+
+    #[allow(clippy::should_implement_trait)]
+    pub fn from_str(s: &str) -> Self {
+        match s {
+            "preparing" => Self::Preparing,
+            "prepared" => Self::Prepared,
+            "committing" => Self::Committing,
+            "committed" => Self::Committed,
+            "rolling_back" => Self::RollingBack,
+            "rolled_back" => Self::RolledBack,
+            "mixed" => Self::Mixed,
+            _ => Self::Unknown,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct VoteResult {
+    pub participant_id: String,
+    pub vote: bool,
+    pub reason: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ParticipantCommitResult {
+    pub participant_id: String,
+    pub committed: bool,
+    pub error: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ParticipantInfo {
+    pub id: String,
+    pub name: String,
+    pub role: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TransactionLog {
+    pub transaction_id: String,
+    pub status: String,
+    pub participants: Vec<ParticipantInfo>,
+    pub created_at: String,
+    pub updated_at: String,
+    pub metadata: serde_json::Value,
+}
+
+impl TransactionLog {
+    fn new(transaction_id: String, participants: &[ParticipantInfo], metadata: serde_json::Value) -> Self {
+        Self {
+            transaction_id,
+            status: TransactionStatus::Preparing.as_str().to_string(),
+            participants: participants.to_vec(),
+            created_at: chrono::Utc::now().to_rfc3339(),
+            updated_at: chrono::Utc::now().to_rfc3339(),
+            metadata,
+        }
+    }
+}
+
+#[async_trait]
+pub trait Participant: Send + Sync {
+    fn id(&self) -> &str;
+    fn name(&self) -> &str;
+    fn role(&self) -> &str;
+
+    async fn prepare(&self, transaction_id: &str) -> Result<(), String>;
+    async fn commit(&self, transaction_id: &str) -> Result<(), String>;
+    async fn rollback(&self, transaction_id: &str) -> Result<(), String>;
+}
+
+pub struct TwoPhaseCommit {
+    backend: Arc<dyn StateBackend>,
+    state_machine: StateMachine,
+}
+
+impl TwoPhaseCommit {
+    pub fn new(backend: Arc<dyn StateBackend>) -> Self {
+        let sm = StateMachine::new(backend.clone());
+        Self { backend, state_machine: sm }
+    }
+
+    fn log_key(transaction_id: &str) -> String {
+        format!("2pc:tx:{transaction_id}")
+    }
+
+    async fn save_log(&self, log: &TransactionLog) -> Result<(), String> {
+        let data = serde_json::to_vec(log).map_err(|e| e.to_string())?;
+        self.backend.save(&Self::log_key(&log.transaction_id), &data).await
+    }
+
+    async fn load_log(&self, transaction_id: &str) -> Result<Option<TransactionLog>, String> {
+        let raw = self.backend.load(&Self::log_key(transaction_id)).await?;
+        match raw {
+            Some(data) => {
+                let log: TransactionLog = serde_json::from_slice(&data).map_err(|e| e.to_string())?;
+                Ok(Some(log))
+            }
+            None => Ok(None),
+        }
+    }
+
+    async fn update_status(&self, transaction_id: &str, status: TransactionStatus) -> Result<TransactionLog, String> {
+        let mut log =
+            self.load_log(transaction_id).await?.ok_or_else(|| format!("Transaction not found: {transaction_id}"))?;
+        log.status = status.as_str().to_string();
+        log.updated_at = chrono::Utc::now().to_rfc3339();
+        self.save_log(&log).await?;
+        Ok(log)
+    }
+
+    pub async fn execute(
+        &self,
+        transaction_id: &str,
+        participants: &[Arc<dyn Participant>],
+        metadata: serde_json::Value,
+    ) -> Result<TransactionLog, String> {
+        let infos: Vec<ParticipantInfo> = participants
+            .iter()
+            .map(|p| ParticipantInfo { id: p.id().to_string(), name: p.name().to_string(), role: p.role().to_string() })
+            .collect();
+
+        let log = TransactionLog::new(transaction_id.to_string(), &infos, metadata);
+        self.save_log(&log).await?;
+
+        let _state = self
+            .state_machine
+            .create_state(&format!("2pc_{transaction_id}"), serde_json::json!({"phase": "begin"}))
+            .await?;
+
+        let vote_results = self.prepare_phase(transaction_id, participants).await?;
+        let all_agreed = vote_results.iter().all(|v| v.vote);
+
+        if all_agreed {
+            self.update_status(transaction_id, TransactionStatus::Committing).await?;
+            let commit_results = self.commit_phase(transaction_id, participants, 2).await;
+            let all_committed =
+                participants.is_empty() || (!commit_results.is_empty() && commit_results.iter().all(|r| r.committed));
+            if all_committed {
+                let log = self.update_status(transaction_id, TransactionStatus::Committed).await?;
+                let _ = self
+                    .state_machine
+                    .transition(&format!("2pc_{transaction_id}"), crate::state_persistence::StateTransition::Completed)
+                    .await;
+                Ok(log)
+            } else {
+                let _ = self.update_status(transaction_id, TransactionStatus::RollingBack).await?;
+                let _ = self.rollback_phase(transaction_id, participants).await;
+
+                // Use commit-phase outcomes only — never re-invoke commit() to probe.
+                let has_committed = commit_results.iter().any(|r| r.committed);
+                let final_status = if has_committed { TransactionStatus::Mixed } else { TransactionStatus::RolledBack };
+
+                let log = self.update_status(transaction_id, final_status).await?;
+                let _ = self
+                    .state_machine
+                    .transition(&format!("2pc_{transaction_id}"), crate::state_persistence::StateTransition::Failed)
+                    .await;
+                let failures: Vec<String> = commit_results
+                    .iter()
+                    .filter(|r| !r.committed)
+                    .map(|r| format!("{}: {}", r.participant_id, r.error.as_deref().unwrap_or("unknown")))
+                    .collect();
+                Err(format!("Commit phase failed: {}. Status: {}", failures.join("; "), log.status))
+            }
+        } else {
+            let _ = self.update_status(transaction_id, TransactionStatus::RollingBack).await?;
+            let _ = self.rollback_phase(transaction_id, participants).await;
+            let log = self.update_status(transaction_id, TransactionStatus::RolledBack).await?;
+            let _ = self
+                .state_machine
+                .transition(&format!("2pc_{transaction_id}"), crate::state_persistence::StateTransition::Failed)
+                .await;
+
+            let rejections: Vec<String> = vote_results
+                .iter()
+                .filter_map(|v| {
+                    if !v.vote {
+                        Some(format!(
+                            "{}: {}",
+                            v.participant_id,
+                            v.reason.clone().unwrap_or_else(|| "no reason".to_string())
+                        ))
+                    } else {
+                        None
+                    }
+                })
+                .collect();
+
+            Err(format!(
+                "Prepare phase failed: participants did not agree. Rejections: [{}]. Status: {}",
+                rejections.join(", "),
+                log.status
+            ))
+        }
+    }
+
+    async fn prepare_phase(
+        &self,
+        transaction_id: &str,
+        participants: &[Arc<dyn Participant>],
+    ) -> Result<Vec<VoteResult>, String> {
+        self.update_status(transaction_id, TransactionStatus::Preparing).await?;
+
+        let mut results = Vec::new();
+        for p in participants {
+            match p.prepare(transaction_id).await {
+                Ok(()) => {
+                    results.push(VoteResult { participant_id: p.id().to_string(), vote: true, reason: None });
+                }
+                Err(e) => {
+                    results.push(VoteResult {
+                        participant_id: p.id().to_string(),
+                        vote: false,
+                        reason: Some(e.clone()),
+                    });
+                    return Ok(results);
+                }
+            }
+        }
+
+        self.update_status(transaction_id, TransactionStatus::Prepared).await?;
+        Ok(results)
+    }
+
+    /// Run commit for all participants; always returns per-participant outcomes
+    /// (including partial success) so callers can mark Mixed without re-committing.
+    async fn commit_phase(
+        &self,
+        transaction_id: &str,
+        participants: &[Arc<dyn Participant>],
+        max_retries: u32,
+    ) -> Vec<ParticipantCommitResult> {
+        let mut results: Vec<ParticipantCommitResult> = Vec::new();
+        let mut pending: Vec<&Arc<dyn Participant>> = participants.iter().collect();
+
+        for attempt in 0..=max_retries {
+            let mut next_pending = Vec::new();
+            for p in pending {
+                match p.commit(transaction_id).await {
+                    Ok(()) => {
+                        results.push(ParticipantCommitResult {
+                            participant_id: p.id().to_string(),
+                            committed: true,
+                            error: None,
+                        });
+                    }
+                    Err(e) => {
+                        if attempt < max_retries {
+                            log::warn!(
+                                "2PC commit retry {}/{} for participant {}: {}",
+                                attempt + 1,
+                                max_retries,
+                                p.id(),
+                                e
+                            );
+                            next_pending.push(p);
+                        } else {
+                            results.push(ParticipantCommitResult {
+                                participant_id: p.id().to_string(),
+                                committed: false,
+                                error: Some(e),
+                            });
+                        }
+                    }
+                }
+            }
+            pending = next_pending;
+            if pending.is_empty() {
+                break;
+            }
+        }
+
+        results
+    }
+
+    async fn rollback_phase(&self, transaction_id: &str, participants: &[Arc<dyn Participant>]) -> Result<(), String> {
+        let mut errors = Vec::new();
+        for p in participants {
+            if let Err(e) = p.rollback(transaction_id).await {
+                errors.push(format!("{}: {}", p.id(), e));
+            }
+        }
+        if errors.is_empty() {
+            Ok(())
+        } else {
+            Err(errors.join("; "))
+        }
+    }
+
+    pub async fn recover(
+        &self,
+        transaction_id: &str,
+        participants: &[Arc<dyn Participant>],
+    ) -> Result<TransactionLog, String> {
+        let log = self
+            .load_log(transaction_id)
+            .await?
+            .ok_or_else(|| format!("No transaction log found for: {transaction_id}"))?;
+
+        let status = TransactionStatus::from_str(&log.status);
+        match status {
+            TransactionStatus::Preparing | TransactionStatus::Prepared => {
+                self.rollback_phase(transaction_id, participants).await?;
+                self.update_status(transaction_id, TransactionStatus::RolledBack).await
+            }
+            TransactionStatus::Committing => {
+                let commit_results = self.commit_phase(transaction_id, participants, 2).await;
+                let all_committed = participants.is_empty()
+                    || (!commit_results.is_empty() && commit_results.iter().all(|r| r.committed));
+                if all_committed {
+                    self.update_status(transaction_id, TransactionStatus::Committed).await
+                } else {
+                    self.rollback_phase(transaction_id, participants).await?;
+                    let has_committed = commit_results.iter().any(|r| r.committed);
+                    let final_status =
+                        if has_committed { TransactionStatus::Mixed } else { TransactionStatus::RolledBack };
+                    self.update_status(transaction_id, final_status).await?;
+                    let failures: Vec<String> = commit_results
+                        .iter()
+                        .filter(|r| !r.committed)
+                        .map(|r| format!("{}: {}", r.participant_id, r.error.as_deref().unwrap_or("unknown")))
+                        .collect();
+                    Err(format!("Recovery failed during commit: {}", failures.join("; ")))
+                }
+            }
+            TransactionStatus::RollingBack => {
+                self.rollback_phase(transaction_id, participants).await?;
+                self.update_status(transaction_id, TransactionStatus::RolledBack).await
+            }
+            TransactionStatus::Committed | TransactionStatus::RolledBack | TransactionStatus::Mixed => Ok(log),
+            TransactionStatus::Unknown => {
+                Err(format!("Unknown transaction status for {transaction_id}: cannot recover"))
+            }
+        }
+    }
+}
+
+/// Participant adapter for database statement execution via TwoPhaseCommit.
+/// The caller provides async closures for the actual database operations.
+#[allow(clippy::type_complexity)]
+pub struct FnParticipant {
+    id: String,
+    name: String,
+    role: String,
+    prepare_fn: Box<dyn Fn() -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>> + Send + Sync>,
+    commit_fn: Box<dyn Fn() -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>> + Send + Sync>,
+    rollback_fn: Box<dyn Fn() -> Pin<Box<dyn Future<Output = Result<(), String>> + Send>> + Send + Sync>,
+}
+
+impl FnParticipant {
+    pub fn new(
+        id: String,
+        name: String,
+        role: String,
+        prepare_fn: impl Fn() -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), String>> + Send>>
+            + Send
+            + Sync
+            + 'static,
+        commit_fn: impl Fn() -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), String>> + Send>>
+            + Send
+            + Sync
+            + 'static,
+        rollback_fn: impl Fn() -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), String>> + Send>>
+            + Send
+            + Sync
+            + 'static,
+    ) -> Self {
+        Self {
+            id,
+            name,
+            role,
+            prepare_fn: Box::new(prepare_fn),
+            commit_fn: Box::new(commit_fn),
+            rollback_fn: Box::new(rollback_fn),
+        }
+    }
+}
+
+#[async_trait]
+impl Participant for FnParticipant {
+    fn id(&self) -> &str {
+        &self.id
+    }
+    fn name(&self) -> &str {
+        &self.name
+    }
+    fn role(&self) -> &str {
+        &self.role
+    }
+
+    async fn prepare(&self, _transaction_id: &str) -> Result<(), String> {
+        (self.prepare_fn)().await
+    }
+
+    async fn commit(&self, _transaction_id: &str) -> Result<(), String> {
+        (self.commit_fn)().await
+    }
+
+    async fn rollback(&self, _transaction_id: &str) -> Result<(), String> {
+        (self.rollback_fn)().await
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::Mutex;
+
+    struct MockParticipant {
+        id: String,
+        name: String,
+        role: String,
+        should_fail_prepare: bool,
+        should_fail_commit: bool,
+        prepare_called: Mutex<bool>,
+        commit_called: Mutex<bool>,
+        rollback_called: Mutex<bool>,
+    }
+
+    impl MockParticipant {
+        fn new(id: &str, name: &str) -> Self {
+            Self {
+                id: id.to_string(),
+                name: name.to_string(),
+                role: "worker".to_string(),
+                should_fail_prepare: false,
+                should_fail_commit: false,
+                prepare_called: Mutex::new(false),
+                commit_called: Mutex::new(false),
+                rollback_called: Mutex::new(false),
+            }
+        }
+
+        fn with_prepare_failure(mut self) -> Self {
+            self.should_fail_prepare = true;
+            self
+        }
+
+        fn with_commit_failure(mut self) -> Self {
+            self.should_fail_commit = true;
+            self
+        }
+
+        #[allow(dead_code)]
+        fn prepare_was_called(&self) -> bool {
+            *self.prepare_called.lock().unwrap()
+        }
+
+        #[allow(dead_code)]
+        fn commit_was_called(&self) -> bool {
+            *self.commit_called.lock().unwrap()
+        }
+
+        #[allow(dead_code)]
+        fn rollback_was_called(&self) -> bool {
+            *self.rollback_called.lock().unwrap()
+        }
+    }
+
+    #[async_trait]
+    impl Participant for MockParticipant {
+        fn id(&self) -> &str {
+            &self.id
+        }
+        fn name(&self) -> &str {
+            &self.name
+        }
+        fn role(&self) -> &str {
+            &self.role
+        }
+
+        async fn prepare(&self, _transaction_id: &str) -> Result<(), String> {
+            *self.prepare_called.lock().unwrap() = true;
+            if self.should_fail_prepare {
+                Err(format!("{}: prepare failed", self.id))
+            } else {
+                Ok(())
+            }
+        }
+
+        async fn commit(&self, _transaction_id: &str) -> Result<(), String> {
+            *self.commit_called.lock().unwrap() = true;
+            if self.should_fail_commit {
+                Err(format!("{}: commit failed", self.id))
+            } else {
+                Ok(())
+            }
+        }
+
+        async fn rollback(&self, _transaction_id: &str) -> Result<(), String> {
+            *self.rollback_called.lock().unwrap() = true;
+            Ok(())
+        }
+    }
+
+    async fn make_2pc() -> (TwoPhaseCommit, std::path::PathBuf) {
+        use crate::storage::Storage;
+        let path = std::env::temp_dir().join(format!("test_2pc_{}.db", uuid::Uuid::new_v4()));
+        let storage = Arc::new(Storage::open(&path).await.unwrap());
+        let backend: Arc<dyn StateBackend> = Arc::new(crate::state_persistence::LocalBackend::new(storage));
+        (TwoPhaseCommit::new(backend), path)
+    }
+
+    #[tokio::test]
+    async fn two_phase_commit_success() {
+        let (coordinator, _tmp) = make_2pc().await;
+        let p1: Arc<dyn Participant> = Arc::new(MockParticipant::new("p1", "worker1"));
+        let p2: Arc<dyn Participant> = Arc::new(MockParticipant::new("p2", "worker2"));
+
+        let result =
+            coordinator.execute("tx_success", &[p1.clone(), p2.clone()], serde_json::json!({"op": "test"})).await;
+
+        assert!(result.is_ok());
+        let log = result.unwrap();
+        assert_eq!(log.status, TransactionStatus::Committed.as_str());
+    }
+
+    #[tokio::test]
+    async fn two_phase_commit_prepare_failure_triggers_rollback() {
+        let (coordinator, _tmp) = make_2pc().await;
+        let p1: Arc<dyn Participant> = Arc::new(MockParticipant::new("p1", "good").with_commit_failure());
+        let p2: Arc<dyn Participant> = Arc::new(MockParticipant::new("p2", "bad").with_prepare_failure());
+
+        let result = coordinator.execute("tx_prepare_fail", &[p1.clone(), p2.clone()], serde_json::json!({})).await;
+
+        assert!(result.is_err());
+        let err = result.unwrap_err();
+        assert!(err.contains("Prepare phase failed"));
+        assert!(err.contains("p2"));
+    }
+
+    #[tokio::test]
+    async fn two_phase_commit_commit_failure_triggers_rollback() {
+        let (coordinator, _tmp) = make_2pc().await;
+        let failing: Arc<dyn Participant> = Arc::new(MockParticipant::new("p1", "fail_commit").with_commit_failure());
+        let passing: Arc<dyn Participant> = Arc::new(MockParticipant::new("p2", "ok"));
+
+        let result =
+            coordinator.execute("tx_commit_fail", &[failing.clone(), passing.clone()], serde_json::json!({})).await;
+
+        assert!(result.is_err());
+        let log = coordinator.load_log("tx_commit_fail").await.unwrap().unwrap();
+        assert_eq!(log.status, TransactionStatus::Mixed.as_str());
+    }
+
+    #[tokio::test]
+    async fn two_phase_commit_empty_participants() {
+        let (coordinator, _tmp) = make_2pc().await;
+        let participants: Vec<Arc<dyn Participant>> = vec![];
+
+        let result = coordinator.execute("tx_empty", &participants, serde_json::json!({})).await;
+
+        assert!(result.is_ok());
+        let log = result.unwrap();
+        assert_eq!(log.status, TransactionStatus::Committed.as_str());
+    }
+
+    #[tokio::test]
+    async fn recovery_from_preparing_state() {
+        let (coordinator, _tmp) = make_2pc().await;
+        let p1: Arc<dyn Participant> = Arc::new(MockParticipant::new("p1", "w1"));
+
+        let log = TransactionLog {
+            transaction_id: "tx_recover_preparing".to_string(),
+            status: TransactionStatus::Preparing.as_str().to_string(),
+            participants: vec![ParticipantInfo {
+                id: "p1".to_string(),
+                name: "w1".to_string(),
+                role: "worker".to_string(),
+            }],
+            created_at: chrono::Utc::now().to_rfc3339(),
+            updated_at: chrono::Utc::now().to_rfc3339(),
+            metadata: serde_json::json!({}),
+        };
+        coordinator.save_log(&log).await.unwrap();
+
+        let result = coordinator.recover("tx_recover_preparing", &[Arc::clone(&p1)]).await;
+
+        assert!(result.is_ok());
+        let recovered = result.unwrap();
+        assert_eq!(recovered.status, TransactionStatus::RolledBack.as_str());
+    }
+
+    #[tokio::test]
+    async fn recovery_from_committing_state() {
+        let (coordinator, _tmp) = make_2pc().await;
+        let p1: Arc<dyn Participant> = Arc::new(MockParticipant::new("p1", "w1"));
+
+        let log = TransactionLog {
+            transaction_id: "tx_recover_committing".to_string(),
+            status: TransactionStatus::Committing.as_str().to_string(),
+            participants: vec![ParticipantInfo {
+                id: "p1".to_string(),
+                name: "w1".to_string(),
+                role: "worker".to_string(),
+            }],
+            created_at: chrono::Utc::now().to_rfc3339(),
+            updated_at: chrono::Utc::now().to_rfc3339(),
+            metadata: serde_json::json!({}),
+        };
+        coordinator.save_log(&log).await.unwrap();
+
+        let result = coordinator.recover("tx_recover_committing", &[Arc::clone(&p1)]).await;
+
+        assert!(result.is_ok());
+        let recovered = result.unwrap();
+        assert_eq!(recovered.status, TransactionStatus::Committed.as_str());
+    }
+
+    #[tokio::test]
+    async fn recovery_already_completed_noop() {
+        let (coordinator, _tmp) = make_2pc().await;
+        let p1: Arc<dyn Participant> = Arc::new(MockParticipant::new("p1", "w1"));
+
+        let log = TransactionLog {
+            transaction_id: "tx_done".to_string(),
+            status: TransactionStatus::Committed.as_str().to_string(),
+            participants: vec![ParticipantInfo {
+                id: "p1".to_string(),
+                name: "w1".to_string(),
+                role: "worker".to_string(),
+            }],
+            created_at: chrono::Utc::now().to_rfc3339(),
+            updated_at: chrono::Utc::now().to_rfc3339(),
+            metadata: serde_json::json!({}),
+        };
+        coordinator.save_log(&log).await.unwrap();
+
+        let result = coordinator.recover("tx_done", &[p1]).await;
+        assert!(result.is_ok());
+        let recovered = result.unwrap();
+        assert_eq!(recovered.status, TransactionStatus::Committed.as_str());
+    }
+
+    #[tokio::test]
+    async fn recovery_unknown_transaction_errors() {
+        let (coordinator, _tmp) = make_2pc().await;
+        let p1: Arc<dyn Participant> = Arc::new(MockParticipant::new("p1", "w1"));
+
+        let result = coordinator.recover("tx_nonexistent", &[p1]).await;
+        assert!(result.is_err());
+    }
+}