diff --git a/src/db/connection.rs b/src/db/connection.rs index dff1de5..b974cf2 100644 --- a/src/db/connection.rs +++ b/src/db/connection.rs @@ -56,7 +56,9 @@ pub fn discover_database_path(custom_path: Option<&Path>) -> Option { default_database_candidates(custom_path) .into_iter() .find(|candidate| { - candidate.exists() && database_has_expected_tables(candidate).unwrap_or(false) + candidate.exists() + && (database_has_v1_tables(candidate).unwrap_or(false) + || database_has_v2_tables(candidate).unwrap_or(false)) }) } @@ -65,11 +67,10 @@ pub fn open_database(path: &Path) -> Result { .map_err(|e| Error::database_open(path, e)) } -pub fn database_has_expected_tables(path: &Path) -> Result { - let conn = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY) - .map_err(|e| Error::database_open(path, e))?; +pub fn database_has_tables(path: &Path, check: Vec<&str>) -> Result { + let conn = open_database(path)?; - for table in ["session", "message", "project"] { + for table in &check { let exists = conn .query_row( "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?1)", @@ -85,6 +86,14 @@ pub fn database_has_expected_tables(path: &Path) -> Result { Ok(true) } +pub fn database_has_v1_tables(path: &Path) -> Result { + database_has_tables(path, vec!["session", "message", "project"]) +} + +pub fn database_has_v2_tables(path: &Path) -> Result { + database_has_tables(path, vec!["session_v2", "session_message", "project"]) +} + fn dedupe_preserve_order(paths: Vec) -> Vec { let mut seen = std::collections::HashSet::new(); paths @@ -95,7 +104,7 @@ fn dedupe_preserve_order(paths: Vec) -> Vec { #[cfg(test)] mod tests { - use super::{database_has_expected_tables, default_database_candidates}; + use super::{database_has_v1_tables, default_database_candidates}; use std::fs; use std::path::Path; use std::time::{SystemTime, UNIX_EPOCH}; @@ -121,7 +130,7 @@ mod tests { .unwrap(); drop(conn); - assert!(!database_has_expected_tables(&db_path).unwrap()); + assert!(!database_has_v1_tables(&db_path).unwrap()); let _ = fs::remove_file(db_path); } } diff --git a/src/db/models.rs b/src/db/models.rs index 057964a..ecb62d7 100644 --- a/src/db/models.rs +++ b/src/db/models.rs @@ -266,7 +266,7 @@ pub struct JsonMessageRecord { pub struct JsonModelRecord { #[serde(rename = "providerID")] pub provider_id: Option, - #[serde(rename = "modelID")] + #[serde(rename = "modelID", alias = "id")] pub model_id: Option, } diff --git a/src/db/queries.rs b/src/db/queries.rs index 8ffee1b..506a4ec 100644 --- a/src/db/queries.rs +++ b/src/db/queries.rs @@ -4,7 +4,7 @@ use std::path::{Path, PathBuf}; use chrono::{DateTime, Local}; -use crate::db::connection::{self, discover_database_path, open_database}; +use crate::db::connection::{self, database_has_v2_tables, discover_database_path, open_database}; use crate::db::errors::{Error, Result}; use crate::db::models::{ AppData, DataSourceKind, ImportStats, InputOptions, JsonMessageRecord, MessageRecord, @@ -40,17 +40,22 @@ pub fn load_app_data(options: &InputOptions) -> Result { pub fn load_from_sqlite(db_path: &Path) -> Result { let conn = open_database(db_path)?; - let mut session_stmt = conn - .prepare( - " + let has_v2_tables = database_has_v2_tables(db_path)?; + let query = format!( + " SELECT s.id, s.parent_id, s.title, s.time_created, s.time_updated, s.time_archived, p.name as project_name, p.worktree as project_worktree - FROM session s + FROM {} s LEFT JOIN project p ON s.project_id = p.id ORDER BY s.time_created DESC ", - ) - .map_err(Error::database_query)?; + if has_v2_tables { + "session_v2" + } else { + "session" + } + ); + let mut session_stmt = conn.prepare(&query).map_err(Error::database_query)?; let sessions_iter = session_stmt .query_map([], |row| { @@ -84,7 +89,7 @@ pub fn load_from_sqlite(db_path: &Path) -> Result { .iter() .map(|session| (session.id.as_str(), session)) .collect::>(); - let sqlite_messages = load_messages_sqlite(&conn, &session_lookup)?; + let sqlite_messages = load_messages_sqlite(&conn, &has_v2_tables, &session_lookup)?; let mut all_events = Vec::new(); let mut all_messages = Vec::new(); @@ -129,17 +134,24 @@ struct SqliteMessageLoad { fn load_messages_sqlite( conn: &rusqlite::Connection, + has_v2_tables: &bool, sessions: &BTreeMap<&str, &SessionRow>, ) -> Result { - let mut stmt = conn - .prepare( - " - SELECT session_id, data - FROM message - ORDER BY session_id ASC, time_created ASC - ", - ) - .map_err(Error::database_query)?; + let query = if *has_v2_tables { + " + SELECT session_id, data, type AS role + FROM session_message + WHERE type IN ('user', 'assistant') + ORDER BY session_id ASC, time_created ASC + " + } else { + " + SELECT session_id, data, NULL AS role + FROM message + ORDER BY session_id ASC, time_created ASC + " + }; + let mut stmt = conn.prepare(&query).map_err(Error::database_query)?; let mut messages = Vec::new(); let mut skipped_messages = 0usize; @@ -159,7 +171,7 @@ fn load_messages_sqlite( } }; - let record: JsonMessageRecord = match serde_json::from_slice(data) { + let mut record: JsonMessageRecord = match serde_json::from_slice(data) { Ok(record) => record, Err(_) => { skipped_messages = skipped_messages.saturating_add(1); @@ -167,6 +179,12 @@ fn load_messages_sqlite( } }; + if record.role.is_none() + && let Some(role_column) = row.get("role").map_err(Error::database_query)? + { + record.role = Some(role_column); + } + if let Some(message) = parse_json_record(record, session, DataSourceKind::Sqlite) { messages.push(message); }