gemini_adk_rs/session/
database.rs

1//! Database-backed session service.
2//!
3//! Feature-gated behind `database-sessions`. Provides a real `sqlx`-backed
4//! [`SessionService`] implementation that supports both SQLite and (when the
5//! `postgres-sessions` feature is enabled) PostgreSQL, chosen by the
6//! connection-URL scheme.
7//!
8//! Storage layout (portable across drivers):
9//! - `sessions(id TEXT PRIMARY KEY, app_name TEXT, user_id TEXT, data TEXT)`
10//!   where `data` is the full JSON of the [`Session`] *without* its events.
11//! - `events(session_id TEXT, seq INTEGER, data TEXT)` where `data` is the
12//!   full JSON of one [`Event`], ordered by `seq`.
13//!
14//! Whole structs are serialized with `serde_json`, avoiding column drift and
15//! keeping the schema driver-portable.
16
17#[cfg(feature = "database-sessions")]
18use async_trait::async_trait;
19#[cfg(feature = "database-sessions")]
20use tokio::sync::Mutex;
21
22#[cfg(feature = "database-sessions")]
23use sqlx::SqlitePool;
24#[cfg(feature = "database-sessions")]
25use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
26
27#[cfg(feature = "postgres-sessions")]
28use sqlx::PgPool;
29#[cfg(feature = "postgres-sessions")]
30use sqlx::postgres::PgPoolOptions;
31
32#[cfg(feature = "database-sessions")]
33use super::{Session, SessionError, SessionId, SessionService};
34
35#[cfg(feature = "database-sessions")]
36use crate::events::Event;
37
38/// Internal connection pool, chosen by URL scheme.
39#[cfg(feature = "database-sessions")]
40enum Pool {
41    /// SQLite pool (file or in-memory).
42    Sqlite(SqlitePool),
43    /// PostgreSQL pool.
44    #[cfg(feature = "postgres-sessions")]
45    Postgres(PgPool),
46}
47
48/// SQL database-backed session service.
49///
50/// Supports SQLite and (with `postgres-sessions`) PostgreSQL via connection
51/// URL. The pool is opened lazily on first use and cached, so [`new`] stays
52/// synchronous and cheap.
53///
54/// [`new`]: DatabaseSessionService::new
55#[cfg(feature = "database-sessions")]
56pub struct DatabaseSessionService {
57    connection_url: String,
58    /// Max pool connections (PostgreSQL). SQLite always uses 1 so an
59    /// in-memory database persists across calls.
60    max_connections: Option<u32>,
61    pool: Mutex<Option<std::sync::Arc<Pool>>>,
62}
63
64#[cfg(feature = "database-sessions")]
65impl DatabaseSessionService {
66    /// Create a new database session service.
67    ///
68    /// The connection URL determines the backend by scheme:
69    /// - `sqlite:` / `sqlite::memory:` → SQLite
70    /// - `postgres:` / `postgresql:` → PostgreSQL (requires `postgres-sessions`)
71    ///
72    /// This does not open a connection; the pool is created lazily on first
73    /// use (or eagerly via [`initialize`](Self::initialize)).
74    pub fn new(connection_url: impl Into<String>) -> Self {
75        Self {
76            connection_url: connection_url.into(),
77            max_connections: None,
78            pool: Mutex::new(None),
79        }
80    }
81
82    /// Set the maximum number of pool connections (PostgreSQL only).
83    ///
84    /// SQLite ignores this and always uses a single connection so that an
85    /// in-memory database survives across calls.
86    pub fn with_max_connections(mut self, max: u32) -> Self {
87        self.max_connections = Some(max);
88        self
89    }
90
91    /// Returns the connection URL this service was configured with.
92    pub fn connection_url(&self) -> &str {
93        &self.connection_url
94    }
95
96    /// Open the connection pool for the configured URL.
97    async fn open_pool(&self) -> Result<Pool, SessionError> {
98        let url = self.connection_url.as_str();
99        if is_postgres_url(url) {
100            #[cfg(feature = "postgres-sessions")]
101            {
102                let mut opts = PgPoolOptions::new();
103                if let Some(max) = self.max_connections {
104                    opts = opts.max_connections(max);
105                }
106                let pool = opts
107                    .connect(url)
108                    .await
109                    .map_err(|e| SessionError::Storage(e.to_string()))?;
110                return Ok(Pool::Postgres(pool));
111            }
112            #[cfg(not(feature = "postgres-sessions"))]
113            {
114                return Err(SessionError::Storage(format!(
115                    "PostgreSQL URL '{url}' requires the 'postgres-sessions' feature"
116                )));
117            }
118        }
119
120        // SQLite (default). Support `sqlite::memory:`, `sqlite:path`,
121        // `sqlite://path`, and bare paths.
122        let opts = sqlite_connect_options(url)?;
123        // max_connections(1) is CRITICAL for `:memory:` so that the single
124        // in-memory database persists across calls instead of every new
125        // connection getting a fresh empty DB.
126        let pool = SqlitePoolOptions::new()
127            .max_connections(1)
128            .connect_with(opts)
129            .await
130            .map_err(|e| SessionError::Storage(e.to_string()))?;
131        Ok(Pool::Sqlite(pool))
132    }
133
134    /// Lazily ensure the pool is connected and the schema exists, returning a
135    /// shared handle to it.
136    async fn pool(&self) -> Result<std::sync::Arc<Pool>, SessionError> {
137        let mut guard = self.pool.lock().await;
138        if let Some(p) = guard.as_ref() {
139            return Ok(p.clone());
140        }
141        let pool = self.open_pool().await?;
142        create_schema(&pool).await?;
143        let arc = std::sync::Arc::new(pool);
144        *guard = Some(arc.clone());
145        Ok(arc)
146    }
147
148    /// Initialize the database: open the pool and create the schema.
149    ///
150    /// Safe to call multiple times. Subsequent CRUD calls reuse the cached
151    /// pool.
152    pub async fn initialize(&self) -> Result<(), SessionError> {
153        self.pool().await?;
154        Ok(())
155    }
156}
157
158/// Returns true if a sqlx error is a unique-constraint violation.
159fn is_unique_violation(e: &sqlx::Error) -> bool {
160    matches!(e, sqlx::Error::Database(db) if db.is_unique_violation())
161}
162
163/// Returns true if `url` denotes a PostgreSQL connection.
164#[cfg(feature = "database-sessions")]
165fn is_postgres_url(url: &str) -> bool {
166    url.starts_with("postgres:") || url.starts_with("postgresql:")
167}
168
169/// Build SQLite connect options from a connection URL, creating the file if
170/// missing.
171#[cfg(feature = "database-sessions")]
172fn sqlite_connect_options(url: &str) -> Result<SqliteConnectOptions, SessionError> {
173    use std::str::FromStr;
174    // Normalize bare paths to a `sqlite:` URL form sqlx understands.
175    let normalized = if url.starts_with("sqlite:") {
176        url.to_string()
177    } else {
178        format!("sqlite://{url}")
179    };
180    SqliteConnectOptions::from_str(&normalized)
181        .map(|o| o.create_if_missing(true))
182        .map_err(|e| SessionError::Storage(e.to_string()))
183}
184
185/// Create the portable schema if it does not already exist.
186#[cfg(feature = "database-sessions")]
187async fn create_schema(pool: &Pool) -> Result<(), SessionError> {
188    match pool {
189        Pool::Sqlite(p) => {
190            sqlx::query(
191                "CREATE TABLE IF NOT EXISTS sessions (\
192                    id TEXT PRIMARY KEY, \
193                    app_name TEXT NOT NULL, \
194                    user_id TEXT NOT NULL, \
195                    data TEXT NOT NULL)",
196            )
197            .execute(p)
198            .await
199            .map_err(|e| SessionError::Storage(e.to_string()))?;
200            sqlx::query(
201                "CREATE TABLE IF NOT EXISTS events (\
202                    session_id TEXT NOT NULL, \
203                    seq INTEGER NOT NULL, \
204                    data TEXT NOT NULL, \
205                    PRIMARY KEY (session_id, seq))",
206            )
207            .execute(p)
208            .await
209            .map_err(|e| SessionError::Storage(e.to_string()))?;
210            sqlx::query(
211                "CREATE INDEX IF NOT EXISTS idx_sessions_app_user \
212                    ON sessions (app_name, user_id)",
213            )
214            .execute(p)
215            .await
216            .map_err(|e| SessionError::Storage(e.to_string()))?;
217            sqlx::query("CREATE INDEX IF NOT EXISTS idx_events_session ON events (session_id)")
218                .execute(p)
219                .await
220                .map_err(|e| SessionError::Storage(e.to_string()))?;
221            Ok(())
222        }
223        #[cfg(feature = "postgres-sessions")]
224        Pool::Postgres(p) => {
225            sqlx::query(
226                "CREATE TABLE IF NOT EXISTS sessions (\
227                    id TEXT PRIMARY KEY, \
228                    app_name TEXT NOT NULL, \
229                    user_id TEXT NOT NULL, \
230                    data TEXT NOT NULL)",
231            )
232            .execute(p)
233            .await
234            .map_err(|e| SessionError::Storage(e.to_string()))?;
235            sqlx::query(
236                "CREATE TABLE IF NOT EXISTS events (\
237                    session_id TEXT NOT NULL, \
238                    seq BIGINT NOT NULL, \
239                    data TEXT NOT NULL, \
240                    PRIMARY KEY (session_id, seq))",
241            )
242            .execute(p)
243            .await
244            .map_err(|e| SessionError::Storage(e.to_string()))?;
245            sqlx::query(
246                "CREATE INDEX IF NOT EXISTS idx_sessions_app_user \
247                    ON sessions (app_name, user_id)",
248            )
249            .execute(p)
250            .await
251            .map_err(|e| SessionError::Storage(e.to_string()))?;
252            sqlx::query("CREATE INDEX IF NOT EXISTS idx_events_session ON events (session_id)")
253                .execute(p)
254                .await
255                .map_err(|e| SessionError::Storage(e.to_string()))?;
256            Ok(())
257        }
258    }
259}
260
261/// Serialize a session WITHOUT its events to a JSON string.
262#[cfg(feature = "database-sessions")]
263fn session_to_data(session: &Session) -> Result<String, SessionError> {
264    let mut bare = session.clone();
265    bare.events = Vec::new();
266    serde_json::to_string(&bare).map_err(|e| SessionError::Storage(e.to_string()))
267}
268
269#[cfg(feature = "database-sessions")]
270#[async_trait]
271impl SessionService for DatabaseSessionService {
272    async fn create_session(&self, app_name: &str, user_id: &str) -> Result<Session, SessionError> {
273        let pool = self.pool().await?;
274        let session = Session::new(app_name, user_id);
275        let id = session.id.as_str().to_string();
276        let data = session_to_data(&session)?;
277
278        match pool.as_ref() {
279            Pool::Sqlite(p) => {
280                sqlx::query(
281                    "INSERT INTO sessions (id, app_name, user_id, data) VALUES (?, ?, ?, ?)",
282                )
283                .bind(&id)
284                .bind(app_name)
285                .bind(user_id)
286                .bind(&data)
287                .execute(p)
288                .await
289                .map_err(|e| SessionError::Storage(e.to_string()))?;
290            }
291            #[cfg(feature = "postgres-sessions")]
292            Pool::Postgres(p) => {
293                sqlx::query(
294                    "INSERT INTO sessions (id, app_name, user_id, data) VALUES ($1, $2, $3, $4)",
295                )
296                .bind(&id)
297                .bind(app_name)
298                .bind(user_id)
299                .bind(&data)
300                .execute(p)
301                .await
302                .map_err(|e| SessionError::Storage(e.to_string()))?;
303            }
304        }
305        Ok(session)
306    }
307
308    async fn get_session(&self, id: &SessionId) -> Result<Option<Session>, SessionError> {
309        let pool = self.pool().await?;
310        let id_str = id.as_str();
311
312        let data: Option<String> = match pool.as_ref() {
313            Pool::Sqlite(p) => sqlx::query_scalar("SELECT data FROM sessions WHERE id = ?")
314                .bind(id_str)
315                .fetch_optional(p)
316                .await
317                .map_err(|e| SessionError::Storage(e.to_string()))?,
318            #[cfg(feature = "postgres-sessions")]
319            Pool::Postgres(p) => sqlx::query_scalar("SELECT data FROM sessions WHERE id = $1")
320                .bind(id_str)
321                .fetch_optional(p)
322                .await
323                .map_err(|e| SessionError::Storage(e.to_string()))?,
324        };
325
326        match data {
327            None => Ok(None),
328            Some(json) => {
329                let mut session: Session = serde_json::from_str(&json)
330                    .map_err(|e| SessionError::Storage(e.to_string()))?;
331                // Match in-memory behavior: load events alongside the session.
332                session.events = self.get_events(id).await?;
333                Ok(Some(session))
334            }
335        }
336    }
337
338    async fn list_sessions(
339        &self,
340        app_name: &str,
341        user_id: &str,
342    ) -> Result<Vec<Session>, SessionError> {
343        let pool = self.pool().await?;
344
345        let rows: Vec<String> = match pool.as_ref() {
346            Pool::Sqlite(p) => {
347                sqlx::query_scalar("SELECT data FROM sessions WHERE app_name = ? AND user_id = ?")
348                    .bind(app_name)
349                    .bind(user_id)
350                    .fetch_all(p)
351                    .await
352                    .map_err(|e| SessionError::Storage(e.to_string()))?
353            }
354            #[cfg(feature = "postgres-sessions")]
355            Pool::Postgres(p) => {
356                sqlx::query_scalar("SELECT data FROM sessions WHERE app_name = $1 AND user_id = $2")
357                    .bind(app_name)
358                    .bind(user_id)
359                    .fetch_all(p)
360                    .await
361                    .map_err(|e| SessionError::Storage(e.to_string()))?
362            }
363        };
364
365        let mut sessions = Vec::with_capacity(rows.len());
366        for json in rows {
367            let session: Session =
368                serde_json::from_str(&json).map_err(|e| SessionError::Storage(e.to_string()))?;
369            sessions.push(session);
370        }
371        Ok(sessions)
372    }
373
374    async fn delete_session(&self, id: &SessionId) -> Result<(), SessionError> {
375        let pool = self.pool().await?;
376        let id_str = id.as_str();
377
378        match pool.as_ref() {
379            Pool::Sqlite(p) => {
380                sqlx::query("DELETE FROM events WHERE session_id = ?")
381                    .bind(id_str)
382                    .execute(p)
383                    .await
384                    .map_err(|e| SessionError::Storage(e.to_string()))?;
385                sqlx::query("DELETE FROM sessions WHERE id = ?")
386                    .bind(id_str)
387                    .execute(p)
388                    .await
389                    .map_err(|e| SessionError::Storage(e.to_string()))?;
390            }
391            #[cfg(feature = "postgres-sessions")]
392            Pool::Postgres(p) => {
393                sqlx::query("DELETE FROM events WHERE session_id = $1")
394                    .bind(id_str)
395                    .execute(p)
396                    .await
397                    .map_err(|e| SessionError::Storage(e.to_string()))?;
398                sqlx::query("DELETE FROM sessions WHERE id = $1")
399                    .bind(id_str)
400                    .execute(p)
401                    .await
402                    .map_err(|e| SessionError::Storage(e.to_string()))?;
403            }
404        }
405        Ok(())
406    }
407
408    async fn append_event(&self, id: &SessionId, event: Event) -> Result<(), SessionError> {
409        let pool = self.pool().await?;
410        let id_str = id.as_str();
411
412        // Ensure the session exists (matches in-memory NotFound semantics).
413        let exists: Option<String> = match pool.as_ref() {
414            Pool::Sqlite(p) => sqlx::query_scalar("SELECT id FROM sessions WHERE id = ?")
415                .bind(id_str)
416                .fetch_optional(p)
417                .await
418                .map_err(|e| SessionError::Storage(e.to_string()))?,
419            #[cfg(feature = "postgres-sessions")]
420            Pool::Postgres(p) => sqlx::query_scalar("SELECT id FROM sessions WHERE id = $1")
421                .bind(id_str)
422                .fetch_optional(p)
423                .await
424                .map_err(|e| SessionError::Storage(e.to_string()))?,
425        };
426        if exists.is_none() {
427            return Err(SessionError::NotFound(id.clone()));
428        }
429
430        let data =
431            serde_json::to_string(&event).map_err(|e| SessionError::Storage(e.to_string()))?;
432
433        // Allocate the sequence number and insert in a single atomic statement
434        // (`INSERT ... SELECT COALESCE(MAX(seq), -1) + 1`) so concurrent
435        // appenders can't read the same `MAX(seq)` and collide. On Postgres,
436        // two transactions under READ COMMITTED can still both observe the
437        // pre-insert snapshot and one will hit the `(session_id, seq)` unique
438        // constraint — retry a bounded number of times in that case. SQLite's
439        // single-connection pool serializes writes, so it never retries.
440        const MAX_ATTEMPTS: u32 = 8;
441        for attempt in 0..MAX_ATTEMPTS {
442            let result = match pool.as_ref() {
443                Pool::Sqlite(p) => sqlx::query(
444                    "INSERT INTO events (session_id, seq, data) \
445                     SELECT ?, COALESCE(MAX(seq), -1) + 1, ? FROM events WHERE session_id = ?",
446                )
447                .bind(id_str)
448                .bind(&data)
449                .bind(id_str)
450                .execute(p)
451                .await
452                .map(|_| ()),
453                #[cfg(feature = "postgres-sessions")]
454                Pool::Postgres(p) => sqlx::query(
455                    "INSERT INTO events (session_id, seq, data) \
456                     SELECT $1, COALESCE(MAX(seq), -1) + 1, $2 FROM events WHERE session_id = $1",
457                )
458                .bind(id_str)
459                .bind(&data)
460                .execute(p)
461                .await
462                .map(|_| ()),
463            };
464            match result {
465                Ok(_) => return Ok(()),
466                Err(e) if is_unique_violation(&e) && attempt + 1 < MAX_ATTEMPTS => continue,
467                Err(e) => return Err(SessionError::Storage(e.to_string())),
468            }
469        }
470        Err(SessionError::Storage(
471            "append_event: exhausted retries allocating event sequence number".into(),
472        ))
473    }
474
475    async fn get_events(&self, id: &SessionId) -> Result<Vec<Event>, SessionError> {
476        let pool = self.pool().await?;
477        let id_str = id.as_str();
478
479        let rows: Vec<String> = match pool.as_ref() {
480            Pool::Sqlite(p) => {
481                sqlx::query_scalar("SELECT data FROM events WHERE session_id = ? ORDER BY seq ASC")
482                    .bind(id_str)
483                    .fetch_all(p)
484                    .await
485                    .map_err(|e| SessionError::Storage(e.to_string()))?
486            }
487            #[cfg(feature = "postgres-sessions")]
488            Pool::Postgres(p) => {
489                sqlx::query_scalar("SELECT data FROM events WHERE session_id = $1 ORDER BY seq ASC")
490                    .bind(id_str)
491                    .fetch_all(p)
492                    .await
493                    .map_err(|e| SessionError::Storage(e.to_string()))?
494            }
495        };
496
497        let mut events = Vec::with_capacity(rows.len());
498        for json in rows {
499            let event: Event =
500                serde_json::from_str(&json).map_err(|e| SessionError::Storage(e.to_string()))?;
501            events.push(event);
502        }
503        Ok(events)
504    }
505}
506
507#[cfg(all(test, feature = "database-sessions"))]
508mod tests {
509    use super::*;
510
511    #[test]
512    fn construction() {
513        let svc = DatabaseSessionService::new("sqlite::memory:");
514        assert_eq!(svc.connection_url(), "sqlite::memory:");
515    }
516
517    #[test]
518    fn construction_with_postgres_url() {
519        let svc = DatabaseSessionService::new("postgres://localhost/mydb");
520        assert_eq!(svc.connection_url(), "postgres://localhost/mydb");
521    }
522
523    #[tokio::test]
524    async fn initialize_succeeds_for_sqlite_memory() {
525        let svc = DatabaseSessionService::new("sqlite::memory:");
526        svc.initialize().await.unwrap();
527    }
528
529    #[tokio::test]
530    async fn create_session_persists() {
531        let svc = DatabaseSessionService::new("sqlite::memory:");
532        svc.initialize().await.unwrap();
533        let session = svc.create_session("app", "user").await.unwrap();
534        assert_eq!(session.app_name, "app");
535        assert_eq!(session.user_id, "user");
536    }
537
538    /// Regression for the event-sequence race: concurrent `append_event`
539    /// calls must each get a distinct seq and none may be dropped on the
540    /// `(session_id, seq)` primary key.
541    #[tokio::test]
542    async fn concurrent_appends_allocate_distinct_sequences() {
543        let svc = std::sync::Arc::new(DatabaseSessionService::new("sqlite::memory:"));
544        svc.initialize().await.unwrap();
545        let session = svc.create_session("app", "user").await.unwrap();
546
547        const N: usize = 25;
548        let mut handles = Vec::new();
549        for i in 0..N {
550            let svc = svc.clone();
551            let id = session.id.clone();
552            handles.push(tokio::spawn(async move {
553                svc.append_event(&id, Event::new("user", Some(format!("msg-{i}"))))
554                    .await
555            }));
556        }
557        for h in handles {
558            h.await.unwrap().unwrap();
559        }
560
561        // All N events landed (no drops), and ordering is stable.
562        let events = svc.get_events(&session.id).await.unwrap();
563        assert_eq!(events.len(), N, "every concurrent append must persist");
564    }
565
566    #[tokio::test]
567    async fn trait_impl_is_object_safe() {
568        let svc = DatabaseSessionService::new("sqlite::memory:");
569        let _dyn_ref: &dyn SessionService = &svc;
570    }
571
572    /// Full round-trip against an in-memory SQLite database, mirroring the
573    /// in-memory service's test expectations as the oracle.
574    #[tokio::test]
575    async fn full_round_trip() {
576        let svc = DatabaseSessionService::new("sqlite::memory:");
577        svc.initialize().await.unwrap();
578
579        // create
580        let session = svc.create_session("my-app", "user-1").await.unwrap();
581        assert_eq!(session.app_name, "my-app");
582        assert_eq!(session.user_id, "user-1");
583
584        // get -> Some
585        let fetched = svc.get_session(&session.id).await.unwrap();
586        assert!(fetched.is_some());
587        let fetched = fetched.unwrap();
588        assert_eq!(fetched.id, session.id);
589        assert!(fetched.events.is_empty());
590
591        // append_event x2
592        svc.append_event(&session.id, Event::new("user", Some("Hello!".to_string())))
593            .await
594            .unwrap();
595        svc.append_event(
596            &session.id,
597            Event::new("assistant", Some("Hi there".to_string())),
598        )
599        .await
600        .unwrap();
601
602        // get_events -> ordered
603        let events = svc.get_events(&session.id).await.unwrap();
604        assert_eq!(events.len(), 2);
605        assert_eq!(events[0].author, "user");
606        assert_eq!(events[0].content.as_deref(), Some("Hello!"));
607        assert_eq!(events[1].author, "assistant");
608
609        // get_session now loads events alongside
610        let with_events = svc.get_session(&session.id).await.unwrap().unwrap();
611        assert_eq!(with_events.events.len(), 2);
612
613        // list_sessions filters by app + user
614        svc.create_session("my-app", "user-1").await.unwrap();
615        svc.create_session("my-app", "user-2").await.unwrap();
616        svc.create_session("other-app", "user-1").await.unwrap();
617        let list = svc.list_sessions("my-app", "user-1").await.unwrap();
618        assert_eq!(list.len(), 2);
619
620        // delete -> get None
621        svc.delete_session(&session.id).await.unwrap();
622        let gone = svc.get_session(&session.id).await.unwrap();
623        assert!(gone.is_none());
624        // events cleared too
625        let no_events = svc.get_events(&session.id).await.unwrap();
626        assert!(no_events.is_empty());
627    }
628
629    #[tokio::test]
630    async fn append_to_missing_session_is_not_found() {
631        let svc = DatabaseSessionService::new("sqlite::memory:");
632        svc.initialize().await.unwrap();
633        let id = SessionId::new();
634        let result = svc
635            .append_event(&id, Event::new("user", Some("Hi".to_string())))
636            .await;
637        assert!(matches!(result, Err(SessionError::NotFound(_))));
638    }
639
640    #[tokio::test]
641    async fn lazy_connect_without_explicit_initialize() {
642        // CRUD should work without an explicit initialize() call.
643        let svc = DatabaseSessionService::new("sqlite::memory:");
644        let session = svc.create_session("app", "user").await.unwrap();
645        let fetched = svc.get_session(&session.id).await.unwrap();
646        assert!(fetched.is_some());
647    }
648}
gemini_adk_rs/session/database.rs

gemini_adk_rs/session/
database.rs
