fix(export): read JSONL from one snapshot

src/export.rs

@@ -1,4 +1,4 @@
-use std::io::Write;
+use std::io::{BufWriter, Seek, Write};
 
 use anyhow::{Result, anyhow};
 use serde::Serialize;
@@ -14,6 +14,7 @@ use crate::types::{
 
 pub(crate) const RECORD_SCHEMA_VERSION: u32 = 5;
 const RECORD_TYPE: &str = "session";
+const EXPORT_IN_MEMORY_DB_LIMIT: usize = 8 * 1024 * 1024;
 
 #[derive(Clone, Copy)]
 pub(crate) struct ExportIncludes {
@@ -175,6 +176,7 @@ pub(crate) fn write_jsonl<W: Write>(
     options: &ExportOptions,
     mut writer: W,
 ) -> Result<()> {
+    let snapshot = store.conn.unchecked_transaction()?;
     let sessions = if options.session_ids.is_empty() {
         store.list_export_sessions(
             options.sources.as_deref(),
@@ -194,7 +196,26 @@ pub(crate) fn write_jsonl<W: Write>(
         sessions
     };
 
-    write_jsonl_for_sessions(store, sessions, options.includes, &mut writer)
+    let (records, mut spool) = collect_session_records(
+        store,
+        sessions,
+        options.includes,
+        !options.session_ids.is_empty(),
+    )?;
+    if let Some(spool) = spool.as_mut() {
+        spool.flush()?;
+    }
+    snapshot.commit()?;
+    if let Some(spool) = spool {
+        let mut file = spool.into_inner()?;
+        file.rewind()?;
+        std::io::copy(&mut file, &mut writer)?;
+    } else {
+        for record in records {
+            write_session_record(&mut writer, &record)?;
+        }
+    }
+    Ok(())
 }
 
 pub(crate) fn session_record_value(
@@ -213,12 +234,19 @@ pub(crate) fn session_record_value(
     ))?)
 }
 
-fn write_jsonl_for_sessions<W: Write>(
+fn collect_session_records(
     store: &Store,
     sessions: Vec<Session>,
     includes: ExportIncludes,
-    mut writer: W,
-) -> Result<()> {
+    force_spool: bool,
+) -> Result<(Vec<ExportSessionRecord>, Option<BufWriter<std::fs::File>>)> {
+    let page_count: usize = store.conn.query_row("PRAGMA page_count", [], |row| row.get(0))?;
+    let page_size: usize = store.conn.query_row("PRAGMA page_size", [], |row| row.get(0))?;
+    let mut records = Vec::new();
+    let mut spool = (force_spool
+        || page_count.saturating_mul(page_size) > EXPORT_IN_MEMORY_DB_LIMIT)
+        .then(|| tempfile::tempfile().map(BufWriter::new))
+        .transpose()?;
     for session in sessions {
         let topology = store.session_topology(&session.id)?;
         let messages =
@@ -234,10 +262,18 @@ fn write_jsonl_for_sessions<W: Write>(
             Vec::new()
         };
         let record = build_session_record(session, topology, messages, usage_events, events);
-        serde_json::to_writer(&mut writer, &record)?;
-        writer.write_all(b"\n")?;
+        if let Some(file) = spool.as_mut() {
+            write_session_record(file, &record)?;
+        } else {
+            records.push(record);
+        }
     }
+    Ok((records, spool))
+}
 
+fn write_session_record<W: Write>(mut writer: W, record: &ExportSessionRecord) -> Result<()> {
+    serde_json::to_writer(&mut writer, record)?;
+    writer.write_all(b"\n")?;
     Ok(())
 }
 

src/integration/regression.rs

@@ -358,6 +358,101 @@ fn export_jsonl_emits_session_messages_and_usage_events() {
     assert_eq!(value["events"][0]["attrs_json"], r#"{"path":"src/main.rs"}"#);
 }
 
+#[test]
+fn export_jsonl_reads_every_record_from_one_snapshot() {
+    struct VersionSwitchWriter {
+        output: Vec<u8>,
+        writer: rusqlite::Connection,
+        switched: bool,
+        checkpoint_busy: Option<i64>,
+    }
+
+    impl std::io::Write for VersionSwitchWriter {
+        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+            self.output.extend_from_slice(buf);
+            if !self.switched && buf.contains(&b'\n') {
+                self.writer
+                    .execute_batch(
+                        "BEGIN IMMEDIATE;
+                         UPDATE sessions SET title = 'version-b' WHERE id = 's2';
+                         UPDATE messages SET content = 'version-b' WHERE session_id = 's2';
+                         COMMIT;",
+                    )
+                    .map_err(std::io::Error::other)?;
+                self.checkpoint_busy = Some(
+                    self.writer
+                        .query_row("PRAGMA wal_checkpoint(TRUNCATE)", [], |row| row.get(0))
+                        .map_err(std::io::Error::other)?,
+                );
+                self.switched = true;
+            }
+            Ok(buf.len())
+        }
+
+        fn flush(&mut self) -> std::io::Result<()> {
+            Ok(())
+        }
+    }
+
+    schema::register_sqlite_vec();
+    let root = tempfile::tempdir().unwrap();
+    let db_path = root.path().join("recall.db");
+    let conn = rusqlite::Connection::open(&db_path).unwrap();
+    conn.execute_batch(
+        "PRAGMA journal_mode=WAL;
+         PRAGMA busy_timeout=5000;
+         PRAGMA foreign_keys=ON;",
+    )
+    .unwrap();
+    schema::init(&conn).unwrap();
+    let store = Store { conn };
+
+    let mut first = make_session("s1", "codex", "raw1", "version-a");
+    first.started_at = 2;
+    let mut second = make_session("s2", "codex", "raw2", "version-a");
+    second.started_at = 1;
+    for session in [&first, &second] {
+        store.insert_session(session).unwrap();
+        store.insert_messages(&[make_message(&session.id, Role::User, "version-a", 0)]).unwrap();
+    }
+    let mut third = make_session("s3", "codex", "raw3", "large");
+    third.started_at = 0;
+    store.insert_session(&third).unwrap();
+    let large_message = "x".repeat(9 * 1024 * 1024);
+    store.insert_messages(&[make_message(&third.id, Role::User, &large_message, 0)]).unwrap();
+    drop(large_message);
+
+    let writer_conn = rusqlite::Connection::open(&db_path).unwrap();
+    writer_conn.execute_batch("PRAGMA busy_timeout=0; PRAGMA foreign_keys=ON;").unwrap();
+    let mut writer = VersionSwitchWriter {
+        output: Vec::new(),
+        writer: writer_conn,
+        switched: false,
+        checkpoint_busy: None,
+    };
+    let options = ExportOptions {
+        session_ids: Vec::new(),
+        sources: None,
+        time_range: TimeRange::All,
+        scope: ProjectScope::Global,
+        thread_role: None,
+        limit: None,
+        includes: ExportIncludes { messages: true, usage: false, events: false },
+    };
+
+    write_jsonl(&store, &options, &mut writer).unwrap();
+
+    let records = String::from_utf8(std::mem::take(&mut writer.output))
+        .unwrap()
+        .lines()
+        .map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
+        .collect::<Vec<_>>();
+    assert_eq!(writer.checkpoint_busy, Some(0));
+    assert_eq!(records[1]["session"]["title"], "version-a");
+    assert_eq!(records[1]["messages"][0]["content"], "version-a");
+    assert_eq!(records[2]["messages"][0]["content"].as_str().unwrap().len(), 9 * 1024 * 1024);
+}
+
 #[test]
 fn export_jsonl_applies_include_projection() {
     let store = setup();
@@ -415,7 +510,7 @@ fn export_jsonl_can_select_sessions_by_id() {
     }
 
     let options = ExportOptions {
-        session_ids: vec!["s3".to_string(), "s1".to_string()],
+        session_ids: vec!["s3".to_string(), "s1".to_string(), "s3".to_string()],
         sources: None,
         time_range: TimeRange::All,
         scope: ProjectScope::Global,
@@ -428,11 +523,13 @@ fn export_jsonl_can_select_sessions_by_id() {
 
     let text = String::from_utf8(out).unwrap();
     let lines: Vec<_> = text.lines().collect();
-    assert_eq!(lines.len(), 2);
+    assert_eq!(lines.len(), 3);
     let first: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
     let second: serde_json::Value = serde_json::from_str(lines[1]).unwrap();
+    let third: serde_json::Value = serde_json::from_str(lines[2]).unwrap();
     assert_eq!(first["session"]["id"], "s3");
     assert_eq!(second["session"]["id"], "s1");
+    assert_eq!(third["session"]["id"], "s3");
 }
 
 #[test]