durable_rust/

ctx.rs

use sea_orm::{ConnectionTrait, DatabaseConnection, DbBackend, Statement};
use serde::Serialize;
use serde::de::DeserializeOwned;
use std::sync::atomic::{AtomicI32, Ordering};
use uuid::Uuid;

use crate::error::DurableError;

/// Context threaded through every workflow and step.
///
/// Users never create or manage task IDs. The SDK handles everything
/// via `(parent_id, name)` lookups — the unique constraint in the schema
/// guarantees exactly-once step creation.
pub struct Ctx {
    db: DatabaseConnection,
    task_id: Uuid,
    sequence: AtomicI32,
}

impl Ctx {
    // ── Workflow lifecycle (user-facing) ──────────────────────────

    /// Start or resume a root workflow by name.
    ///
    /// ```ignore
    /// let ctx = Ctx::start(&db, "ingest", json!({"crawl": "CC-2026"})).await?;
    /// ```
    pub async fn start(
        db: &DatabaseConnection,
        name: &str,
        input: Option<serde_json::Value>,
    ) -> Result<Self, DurableError> {
        let task_id = find_or_create_task(db, None, None, name, "WORKFLOW", input).await?;
        set_status(db, task_id, "RUNNING").await?;
        Ok(Self {
            db: db.clone(),
            task_id,
            sequence: AtomicI32::new(0),
        })
    }

    /// Run a step. If already completed, returns saved output. Otherwise executes
    /// the closure, saves the result, and returns it.
    ///
    /// ```ignore
    /// let count: u32 = ctx.step("fetch_count", || async { api.get_count().await }).await?;
    /// ```
    pub async fn step<T, F, Fut>(&self, name: &str, f: F) -> Result<T, DurableError>
    where
        T: Serialize + DeserializeOwned,
        F: FnOnce() -> Fut,
        Fut: std::future::Future<Output = Result<T, DurableError>>,
    {
        let seq = self.sequence.fetch_add(1, Ordering::SeqCst);

        // Find or create — idempotent via UNIQUE(parent_id, name)
        let step_id =
            find_or_create_task(&self.db, Some(self.task_id), Some(seq), name, "STEP", None)
                .await?;

        // If already completed, replay
        if let Some(output) = get_output(&self.db, step_id).await? {
            let val: T = serde_json::from_value(output)?;
            tracing::debug!(step = name, seq, "replaying saved output");
            return Ok(val);
        }

        // Execute
        set_status(&self.db, step_id, "RUNNING").await?;
        match f().await {
            Ok(val) => {
                let json = serde_json::to_value(&val)?;
                complete_task(&self.db, step_id, json).await?;
                tracing::debug!(step = name, seq, "step completed");
                Ok(val)
            }
            Err(e) => {
                fail_task(&self.db, step_id, &e.to_string()).await?;
                Err(e)
            }
        }
    }

    /// Start or resume a child workflow. Returns a new `Ctx` scoped to the child.
    ///
    /// ```ignore
    /// let child_ctx = ctx.child("embed_batch", json!({"vectors": 1000})).await?;
    /// // use child_ctx.step(...) for steps inside the child
    /// child_ctx.complete(json!({"done": true})).await?;
    /// ```
    pub async fn child(
        &self,
        name: &str,
        input: Option<serde_json::Value>,
    ) -> Result<Self, DurableError> {
        let seq = self.sequence.fetch_add(1, Ordering::SeqCst);
        let child_id = find_or_create_task(
            &self.db,
            Some(self.task_id),
            Some(seq),
            name,
            "WORKFLOW",
            input,
        )
        .await?;

        // If child already completed, return a Ctx that will replay
        // (the caller should check is_completed() or just run steps which will replay)
        set_status(&self.db, child_id, "RUNNING").await?;

        Ok(Self {
            db: self.db.clone(),
            task_id: child_id,
            sequence: AtomicI32::new(0),
        })
    }

    /// Check if this workflow/child was already completed (for skipping in parent).
    pub async fn is_completed(&self) -> Result<bool, DurableError> {
        let status = get_status(&self.db, self.task_id).await?;
        Ok(status.as_deref() == Some("COMPLETED"))
    }

    /// Get the saved output if this task is completed.
    pub async fn get_output<T: DeserializeOwned>(&self) -> Result<Option<T>, DurableError> {
        match get_output(&self.db, self.task_id).await? {
            Some(val) => Ok(Some(serde_json::from_value(val)?)),
            None => Ok(None),
        }
    }

    /// Mark this workflow as completed with an output value.
    pub async fn complete<T: Serialize>(&self, output: &T) -> Result<(), DurableError> {
        let json = serde_json::to_value(output)?;
        complete_task(&self.db, self.task_id, json).await
    }

    /// Mark this workflow as failed.
    pub async fn fail(&self, error: &str) -> Result<(), DurableError> {
        fail_task(&self.db, self.task_id, error).await
    }

    // ── Accessors ────────────────────────────────────────────────

    pub fn db(&self) -> &DatabaseConnection {
        &self.db
    }

    pub fn task_id(&self) -> Uuid {
        self.task_id
    }

    pub fn next_sequence(&self) -> i32 {
        self.sequence.fetch_add(1, Ordering::SeqCst)
    }
}

// ── Internal SQL helpers ─────────────────────────────────────────────

/// Find an existing task by (parent_id, name) or create a new one.
/// This is the core idempotency mechanism — the UNIQUE(parent_id, name)
/// constraint in the schema makes this safe.
async fn find_or_create_task(
    db: &DatabaseConnection,
    parent_id: Option<Uuid>,
    sequence: Option<i32>,
    name: &str,
    kind: &str,
    input: Option<serde_json::Value>,
) -> Result<Uuid, DurableError> {
    let parent_sql = match parent_id {
        Some(p) => format!("'{p}'"),
        None => "NULL".to_string(),
    };

    // Try to find existing
    let find_sql = format!(
        "SELECT id FROM durable.task WHERE parent_id {eq} AND name = '{name}'",
        eq = match parent_id {
            Some(p) => format!("= '{p}'"),
            None => "IS NULL".to_string(),
        }
    );
    let row = db
        .query_one(Statement::from_string(DbBackend::Postgres, find_sql))
        .await?;

    if let Some(row) = row {
        let id: Uuid = row
            .try_get_by_index(0)
            .map_err(|e| DurableError::custom(e.to_string()))?;
        return Ok(id);
    }

    // Create new
    let id = Uuid::new_v4();
    let seq_sql = match sequence {
        Some(s) => s.to_string(),
        None => "NULL".to_string(),
    };
    let input_sql = match &input {
        Some(v) => format!("'{}'", serde_json::to_string(v)?),
        None => "NULL".to_string(),
    };

    let sql = format!(
        "INSERT INTO durable.task (id, parent_id, sequence, name, kind, status, input) \
         VALUES ('{id}', {parent_sql}, {seq_sql}, '{name}', '{kind}', 'PENDING', {input_sql})"
    );
    db.execute(Statement::from_string(DbBackend::Postgres, sql))
        .await?;

    Ok(id)
}

async fn get_output(
    db: &DatabaseConnection,
    task_id: Uuid,
) -> Result<Option<serde_json::Value>, DurableError> {
    let sql =
        format!("SELECT output FROM durable.task WHERE id = '{task_id}' AND status = 'COMPLETED'");
    let row = db
        .query_one(Statement::from_string(DbBackend::Postgres, sql))
        .await?;

    match row {
        Some(r) => Ok(r.try_get_by_index(0).ok()),
        None => Ok(None),
    }
}

async fn get_status(
    db: &DatabaseConnection,
    task_id: Uuid,
) -> Result<Option<String>, DurableError> {
    let sql = format!("SELECT status FROM durable.task WHERE id = '{task_id}'");
    let row = db
        .query_one(Statement::from_string(DbBackend::Postgres, sql))
        .await?;

    match row {
        Some(r) => Ok(r.try_get_by_index(0).ok()),
        None => Ok(None),
    }
}

async fn set_status(
    db: &DatabaseConnection,
    task_id: Uuid,
    status: &str,
) -> Result<(), DurableError> {
    let sql = format!(
        "UPDATE durable.task SET status = '{status}', started_at = COALESCE(started_at, now()) \
         WHERE id = '{task_id}'"
    );
    db.execute(Statement::from_string(DbBackend::Postgres, sql))
        .await?;
    Ok(())
}

async fn complete_task(
    db: &DatabaseConnection,
    task_id: Uuid,
    output: serde_json::Value,
) -> Result<(), DurableError> {
    let sql = format!(
        "UPDATE durable.task SET status = 'COMPLETED', output = '{}', completed_at = now() \
         WHERE id = '{task_id}'",
        output
    );
    db.execute(Statement::from_string(DbBackend::Postgres, sql))
        .await?;
    Ok(())
}

async fn fail_task(
    db: &DatabaseConnection,
    task_id: Uuid,
    error: &str,
) -> Result<(), DurableError> {
    let sql = format!(
        "UPDATE durable.task SET status = 'FAILED', error = '{}', completed_at = now() \
         WHERE id = '{task_id}'",
        error.replace('\'', "''")
    );
    db.execute(Statement::from_string(DbBackend::Postgres, sql))
        .await?;
    Ok(())
}