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/// Session ids are stored as `TEXT`, so any string id works (pre-0.6 the
10/// column was `UUID`; the migration converts it automatically).
11///
12/// Saves are protected by optimistic locking on [`Session::version`]: a save
13/// whose version does not match the stored row fails with
14/// [`GraphError::SessionConflict`] instead of silently overwriting newer data.
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        let statements = [
42            r#"
43            CREATE TABLE IF NOT EXISTS sessions (
44                id TEXT 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                version BIGINT NOT NULL DEFAULT 0,
50                created_at TIMESTAMPTZ DEFAULT NOW(),
51                updated_at TIMESTAMPTZ DEFAULT NOW()
52            );
53            "#,
54            // Upgrade pre-0.6 tables in place: UUID ids become TEXT (a no-op
55            // if already TEXT) and the version column is added.
56            r#"ALTER TABLE sessions ALTER COLUMN id TYPE TEXT;"#,
57            r#"ALTER TABLE sessions ADD COLUMN IF NOT EXISTS version BIGINT NOT NULL DEFAULT 0;"#,
58        ];
59
60        for statement in statements {
61            sqlx::query(statement)
62                .execute(pool)
63                .await
64                .map_err(|e| GraphError::StorageError(format!("Migration failed: {e}")))?;
65        }
66        Ok(())
67    }
68}
69
70#[async_trait]
71impl SessionStorage for PostgresSessionStorage {
72    async fn save(&self, session: Session) -> Result<()> {
73        let context_json = serde_json::to_value(&session.context)
74            .map_err(|e| GraphError::StorageError(format!("Context serialization failed: {e}")))?;
75        let version = i64::try_from(session.version)
76            .map_err(|e| GraphError::StorageError(format!("Session version overflow: {e}")))?;
77
78        // Optimistic locking: the UPDATE only applies when the stored version
79        // still matches the version this session was loaded with.
80        let result = sqlx::query(
81            r#"
82            INSERT INTO sessions (id, graph_id, current_task_id, status_message, context, version, updated_at)
83            VALUES ($1, $2, $3, $4, $5, $6 + 1, NOW())
84            ON CONFLICT (id) DO UPDATE
85            SET graph_id = EXCLUDED.graph_id,
86                current_task_id = EXCLUDED.current_task_id,
87                status_message = EXCLUDED.status_message,
88                context = EXCLUDED.context,
89                version = sessions.version + 1,
90                updated_at = NOW()
91            WHERE sessions.version = $6
92            "#,
93        )
94        .bind(&session.id)
95        .bind(&session.graph_id)
96        .bind(&session.current_task_id)
97        .bind(&session.status_message)
98        .bind(&context_json)
99        .bind(version)
100        .execute(&self.pool)
101        .await
102        .map_err(|e| GraphError::StorageError(format!("Failed to save session: {e}")))?;
103
104        if result.rows_affected() == 0 {
105            return Err(GraphError::SessionConflict(format!(
106                "Session '{}' was modified concurrently (attempted save from version {})",
107                session.id, session.version
108            )));
109        }
110
111        Ok(())
112    }
113
114    async fn get(&self, id: &str) -> Result<Option<Session>> {
115        let row = sqlx::query_as::<_, (String, String, String, Option<String>, serde_json::Value, i64)>(
116            r#"
117            SELECT id, graph_id, current_task_id, status_message, context, version
118            FROM sessions
119            WHERE id = $1
120            "#,
121        )
122        .bind(id)
123        .fetch_optional(&self.pool)
124        .await
125        .map_err(|e| GraphError::StorageError(format!("Failed to fetch session: {e}")))?;
126
127        if let Some((session_id, graph_id, current_task_id, status_message, context_json, version)) = row {
128            let context: crate::Context = serde_json::from_value(context_json)
129                .map_err(|e| GraphError::StorageError(format!("Context deserialization failed: {e}")))?;
130            Ok(Some(Session {
131                id: session_id,
132                graph_id,
133                current_task_id,
134                status_message,
135                context,
136                version: version as u64,
137            }))
138        } else {
139            Ok(None)
140        }
141    }
142
143    async fn delete(&self, id: &str) -> Result<()> {
144        sqlx::query(
145            r#"
146            DELETE FROM sessions WHERE id = $1
147            "#,
148        )
149        .bind(id)
150        .execute(&self.pool)
151        .await
152        .map_err(|e| GraphError::StorageError(format!("Failed to delete session: {e}")))?;
153        Ok(())
154    }
155}
graph_flow/storage_postgres.rs

graph_flow/
storage_postgres.rs
