diff --git a/src/export.rs b/src/export.rs index 5895c85..00be013 100644 --- a/src/export.rs +++ b/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( 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( 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( +fn collect_session_records( store: &Store, sessions: Vec, includes: ExportIncludes, - mut writer: W, -) -> Result<()> { + force_spool: bool, +) -> Result<(Vec, Option>)> { + 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( 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(mut writer: W, record: &ExportSessionRecord) -> Result<()> { + serde_json::to_writer(&mut writer, record)?; + writer.write_all(b"\n")?; Ok(()) } diff --git a/src/integration/regression.rs b/src/integration/regression.rs index 15e408a..8d93c42 100644 --- a/src/integration/regression.rs +++ b/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, + writer: rusqlite::Connection, + switched: bool, + checkpoint_busy: Option, + } + + impl std::io::Write for VersionSwitchWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + 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::(line).unwrap()) + .collect::>(); + 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]