stasis/infrastructure/runtime/

surreal_thread_store.rs

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use surrealdb::{engine::any::Any, Surreal};
use surrealdb_types::SurrealValue;

use crate::domain::errors::{Result, StasisError};
use crate::domain::runtime::thread::{
    NewThread, NewThreadEvent, ThreadEvent, ThreadSnapshot,
};
use crate::ports::outbound::runtime::thread_store::ThreadStore;

#[derive(Clone)]
pub struct SurrealThreadStore {
    db: Surreal<Any>,
    thread_table: String,
    event_table: String,
}

impl SurrealThreadStore {
    pub fn new(db: Surreal<Any>) -> Self {
        Self {
            db,
            thread_table: "thread".to_string(),
            event_table: "thread_event".to_string(),
        }
    }

    fn port_err(prefix: &str, err: impl std::fmt::Display) -> StasisError {
        StasisError::PortFailure(format!("{prefix}: {err}"))
    }
}

#[derive(Clone, Debug, Serialize, Deserialize, SurrealValue)]
struct ThreadRecordRow {
    thread_id: String,
    parent_thread_id: Option<String>,
    branch_label: Option<String>,
    created_at: DateTime<Utc>,
    updated_at: DateTime<Utc>,
}

#[derive(Clone, Debug, Serialize, Deserialize, SurrealValue)]
struct ThreadEventRow {
    event_id: String,
    thread_id: String,
    event_kind: String,
    payload_ref: String,
    occurred_at: DateTime<Utc>,
}

impl From<ThreadRecordRow> for ThreadSnapshot {
    fn from(row: ThreadRecordRow) -> Self {
        Self {
            thread_id: row.thread_id,
            parent_thread_id: row.parent_thread_id,
            branch_label: row.branch_label,
            created_at: row.created_at,
            updated_at: row.updated_at,
        }
    }
}

impl From<ThreadEventRow> for ThreadEvent {
    fn from(row: ThreadEventRow) -> Self {
        Self {
            event_id: row.event_id,
            thread_id: row.thread_id,
            event_kind: row.event_kind,
            payload_ref: row.payload_ref,
            occurred_at: row.occurred_at,
        }
    }
}

#[async_trait]
impl ThreadStore for SurrealThreadStore {
    async fn create_thread(&self, thread: NewThread) -> Result<ThreadSnapshot> {
        if let Some(parent_thread_id) = &thread.parent_thread_id {
            let parent = self.get_thread(parent_thread_id).await?;
            if parent.is_none() {
                return Err(StasisError::PortFailure(format!(
                    "parent thread not found: {}",
                    parent_thread_id
                )));
            }
        }

        let row = ThreadRecordRow {
            thread_id: thread.thread_id,
            parent_thread_id: thread.parent_thread_id,
            branch_label: thread.branch_label,
            created_at: thread.created_at,
            updated_at: thread.created_at,
        };

        self.db
            .query("CREATE type::record($table, $id) CONTENT $data")
            .bind(("table", self.thread_table.clone()))
            .bind(("id", row.thread_id.clone()))
            .bind(("data", row.clone()))
            .await
            .map_err(|e| Self::port_err("create thread", e))?;

        Ok(row.into())
    }

    async fn get_thread(&self, thread_id: &str) -> Result<Option<ThreadSnapshot>> {
        let mut response = self
            .db
            .query("SELECT * FROM type::record($table, $id)")
            .bind(("table", self.thread_table.clone()))
            .bind(("id", thread_id.to_string()))
            .await
            .map_err(|e| Self::port_err("get thread", e))?;

        let row: Option<ThreadRecordRow> = response
            .take(0)
            .map_err(|e| Self::port_err("decode thread", e))?;

        Ok(row.map(ThreadSnapshot::from))
    }

    async fn append_event(&self, event: NewThreadEvent) -> Result<ThreadEvent> {
        let Some(mut thread) = self.get_thread(&event.thread_id).await? else {
            return Err(StasisError::PortFailure(format!(
                "thread not found: {}",
                event.thread_id
            )));
        };

        let event_row = ThreadEventRow {
            event_id: event.event_id,
            thread_id: event.thread_id,
            event_kind: event.event_kind,
            payload_ref: event.payload_ref,
            occurred_at: event.occurred_at,
        };

        self.db
            .query("CREATE type::record($table, $id) CONTENT $data")
            .bind(("table", self.event_table.clone()))
            .bind(("id", event_row.event_id.clone()))
            .bind(("data", event_row.clone()))
            .await
            .map_err(|e| Self::port_err("append thread event", e))?;

        thread.updated_at = event_row.occurred_at;
        let thread_row = ThreadRecordRow {
            thread_id: thread.thread_id.clone(),
            parent_thread_id: thread.parent_thread_id.clone(),
            branch_label: thread.branch_label.clone(),
            created_at: thread.created_at,
            updated_at: thread.updated_at,
        };

        self.db
            .query("UPSERT type::record($table, $id) CONTENT $data")
            .bind(("table", self.thread_table.clone()))
            .bind(("id", thread_row.thread_id.clone()))
            .bind(("data", thread_row))
            .await
            .map_err(|e| Self::port_err("update thread metadata", e))?;

        Ok(event_row.into())
    }

    async fn list_events(&self, thread_id: &str) -> Result<Vec<ThreadEvent>> {
        let mut response = self
            .db
            .query("SELECT * FROM type::table($table) WHERE thread_id = $thread_id ORDER BY occurred_at ASC")
            .bind(("table", self.event_table.clone()))
            .bind(("thread_id", thread_id.to_string()))
            .await
            .map_err(|e| Self::port_err("list thread events", e))?;

        let rows: Vec<ThreadEventRow> = response
            .take(0)
            .map_err(|e| Self::port_err("decode thread events", e))?;

        Ok(rows.into_iter().map(ThreadEvent::from).collect())
    }

    async fn fork_thread(
        &self,
        parent_thread_id: &str,
        child_thread_id: &str,
        branch_label: Option<String>,
        created_at: DateTime<Utc>,
    ) -> Result<ThreadSnapshot> {
        self.create_thread(NewThread {
            thread_id: child_thread_id.to_string(),
            parent_thread_id: Some(parent_thread_id.to_string()),
            branch_label,
            created_at,
        })
        .await
    }

    async fn list_lineage(&self, thread_id: &str) -> Result<Vec<ThreadSnapshot>> {
        let mut lineage = Vec::new();
        let mut cursor = self.get_thread(thread_id).await?;

        while let Some(node) = cursor {
            cursor = if let Some(parent_thread_id) = &node.parent_thread_id {
                self.get_thread(parent_thread_id).await?
            } else {
                None
            };
            lineage.push(node);
        }

        lineage.reverse();
        Ok(lineage)
    }
}