graph_flow/
storage_postgres.rs

1use async_trait::async_trait;
2use serde_json;
3use sqlx::{postgres::PgPoolOptions, Pool, Postgres};
4
5use crate::{Session, error::{Result, GraphError}, storage::SessionStorage};
6
7/// PostgreSQL-backed [`SessionStorage`] implementation.
8///
9/// Note: the `sessions` table stores ids as `UUID`, so [`Session::id`] must be
10/// a valid UUID string - other values fail at runtime with a cast error.
11///
12/// Sessions are saved with last-write-wins semantics: there is no optimistic
13/// locking, so avoid running the same session concurrently from multiple
14/// workers (the later save silently overwrites the earlier one).
15pub struct PostgresSessionStorage {
16    pool: Pool<Postgres>,
17}
18
19impl PostgresSessionStorage {
20    /// Connect with default pool settings (5 max connections) and run the
21    /// schema migration. Use [`PostgresSessionStorage::with_pool`] for custom
22    /// pool configuration.
23    pub async fn connect(database_url: &str) -> Result<Self> {
24        let pool = PgPoolOptions::new()
25            .max_connections(5)
26            .connect(database_url)
27            .await
28            .map_err(|e| GraphError::StorageError(format!("Failed to connect to Postgres: {e}")))?;
29
30        Self::with_pool(pool).await
31    }
32
33    /// Build the storage from a pre-configured pool (custom connection limits,
34    /// timeouts, etc.) and run the schema migration.
35    pub async fn with_pool(pool: Pool<Postgres>) -> Result<Self> {
36        Self::migrate(&pool).await?;
37        Ok(Self { pool })
38    }
39
40    async fn migrate(pool: &Pool<Postgres>) -> Result<()> {
41        sqlx::query(
42            r#"
43            CREATE TABLE IF NOT EXISTS sessions (
44                id UUID PRIMARY KEY,
45                graph_id TEXT NOT NULL,
46                current_task_id TEXT NOT NULL,
47                status_message TEXT,
48                context JSONB NOT NULL,
49                created_at TIMESTAMPTZ DEFAULT NOW(),
50                updated_at TIMESTAMPTZ DEFAULT NOW()
51            );
52            "#,
53        )
54        .execute(pool)
55        .await
56        .map_err(|e| GraphError::StorageError(format!("Migration failed: {e}")))?;
57        Ok(())
58    }
59}
60
61#[async_trait]
62impl SessionStorage for PostgresSessionStorage {
63    async fn save(&self, session: Session) -> Result<()> {
64        let context_json = serde_json::to_value(&session.context)
65            .map_err(|e| GraphError::StorageError(format!("Context serialization failed: {e}")))?;
66
67        // Single upsert - already atomic, no explicit transaction needed.
68        sqlx::query(
69            r#"
70            INSERT INTO sessions (id, graph_id, current_task_id, status_message, context, updated_at)
71            VALUES ($1::uuid, $2, $3, $4, $5, NOW())
72            ON CONFLICT (id) DO UPDATE
73            SET graph_id = EXCLUDED.graph_id,
74                current_task_id = EXCLUDED.current_task_id,
75                status_message = EXCLUDED.status_message,
76                context = EXCLUDED.context,
77                updated_at = NOW()
78            "#,
79        )
80        .bind(&session.id)
81        .bind(&session.graph_id)
82        .bind(&session.current_task_id)
83        .bind(&session.status_message)
84        .bind(&context_json)
85        .execute(&self.pool)
86        .await
87        .map_err(|e| GraphError::StorageError(format!("Failed to save session: {e}")))?;
88
89        Ok(())
90    }
91
92    async fn get(&self, id: &str) -> Result<Option<Session>> {
93        let row = sqlx::query_as::<_, (String, String, String, Option<String>, serde_json::Value)>(
94            r#"
95            SELECT id::text, graph_id, current_task_id, status_message, context
96            FROM sessions
97            WHERE id = $1::uuid
98            "#,
99        )
100        .bind(id)
101        .fetch_optional(&self.pool)
102        .await
103        .map_err(|e| GraphError::StorageError(format!("Failed to fetch session: {e}")))?;
104
105        if let Some((session_id, graph_id, current_task_id, status_message, context_json)) = row {
106            let context: crate::Context = serde_json::from_value(context_json)
107                .map_err(|e| GraphError::StorageError(format!("Context deserialization failed: {e}")))?;
108            Ok(Some(Session {
109                id: session_id,
110                graph_id,
111                current_task_id,
112                status_message,
113                context,
114            }))
115        } else {
116            Ok(None)
117        }
118    }
119
120    async fn delete(&self, id: &str) -> Result<()> {
121        sqlx::query(
122            r#"
123            DELETE FROM sessions WHERE id = $1::uuid
124            "#,
125        )
126        .bind(id)
127        .execute(&self.pool)
128        .await
129        .map_err(|e| GraphError::StorageError(format!("Failed to delete session: {e}")))?;
130        Ok(())
131    }
132}
graph_flow/storage_postgres.rs

graph_flow/
storage_postgres.rs
