lc/
database.rs

1use anyhow::Result;
2use chrono::{DateTime, Utc};
3use rusqlite::{params, Connection};
4use std::path::PathBuf;
5use std::sync::{Arc, Mutex};
6
7#[derive(Debug, Clone)]
8pub struct ChatEntry {
9    pub chat_id: String,
10    pub model: String,
11    pub question: String,
12    pub response: String,
13    pub timestamp: DateTime<Utc>,
14    pub input_tokens: Option<i32>,
15    pub output_tokens: Option<i32>,
16}
17
18#[derive(Debug)]
19pub struct DatabaseStats {
20    pub total_entries: usize,
21    pub unique_sessions: usize,
22    pub file_size_bytes: u64,
23    pub date_range: Option<(DateTime<Utc>, DateTime<Utc>)>,
24    pub model_usage: Vec<(String, i64)>,
25}
26
27// Connection pool for reusing database connections
28pub struct ConnectionPool {
29    connections: Arc<Mutex<Vec<Connection>>>,
30    max_connections: usize,
31    db_path: PathBuf,
32}
33
34impl ConnectionPool {
35    pub fn new(db_path: PathBuf, max_connections: usize) -> Result<Self> {
36        let mut connections = Vec::with_capacity(max_connections);
37
38        // Pre-create initial connections
39        for _ in 0..std::cmp::min(2, max_connections) {
40            let conn = Connection::open(&db_path)?;
41            Self::configure_connection(&conn)?;
42            connections.push(conn);
43        }
44
45        Ok(Self {
46            connections: Arc::new(Mutex::new(connections)),
47            max_connections,
48            db_path,
49        })
50    }
51
52    fn configure_connection(conn: &Connection) -> Result<()> {
53        // Enable WAL mode for better concurrent performance
54        conn.pragma_update(None, "journal_mode", "WAL")?;
55        // Increase cache size for better performance
56        conn.pragma_update(None, "cache_size", 10000)?;
57        // Enable foreign keys
58        conn.pragma_update(None, "foreign_keys", true)?;
59        // Set synchronous to NORMAL for better performance
60        conn.pragma_update(None, "synchronous", "NORMAL")?;
61        Ok(())
62    }
63
64    pub fn get_connection(&self) -> Result<PooledConnection> {
65        let mut connections = self.connections.lock().unwrap();
66
67        if let Some(conn) = connections.pop() {
68            Ok(PooledConnection {
69                conn: Some(conn),
70                pool: self.connections.clone(),
71            })
72        } else if connections.len() < self.max_connections {
73            // Create new connection if under limit
74            let conn = Connection::open(&self.db_path)?;
75            Self::configure_connection(&conn)?;
76            Ok(PooledConnection {
77                conn: Some(conn),
78                pool: self.connections.clone(),
79            })
80        } else {
81            // Wait for a connection to become available
82            // In a real implementation, you might want to use a condition variable
83            // For now, create a new temporary connection
84            let conn = Connection::open(&self.db_path)?;
85            Self::configure_connection(&conn)?;
86            Ok(PooledConnection {
87                conn: Some(conn),
88                pool: self.connections.clone(),
89            })
90        }
91    }
92}
93
94// RAII wrapper for pooled connections
95pub struct PooledConnection {
96    conn: Option<Connection>,
97    pool: Arc<Mutex<Vec<Connection>>>,
98}
99
100impl PooledConnection {
101    pub fn execute(
102        &self,
103        sql: &str,
104        params: impl rusqlite::Params,
105    ) -> Result<usize, rusqlite::Error> {
106        self.conn.as_ref().unwrap().execute(sql, params)
107    }
108
109    pub fn query_row<T, P, F>(&self, sql: &str, params: P, f: F) -> Result<T, rusqlite::Error>
110    where
111        P: rusqlite::Params,
112        F: FnOnce(&rusqlite::Row<'_>) -> Result<T, rusqlite::Error>,
113    {
114        self.conn.as_ref().unwrap().query_row(sql, params, f)
115    }
116}
117
118impl Drop for PooledConnection {
119    fn drop(&mut self) {
120        if let Some(conn) = self.conn.take() {
121            let mut connections = self.pool.lock().unwrap();
122            connections.push(conn);
123        }
124    }
125}
126
127// Optimized Database struct with connection pooling
128pub struct Database {
129    pool: ConnectionPool,
130}
131
132impl Database {
133    pub fn new() -> Result<Self> {
134        let db_path = Self::database_path()?;
135        let pool = ConnectionPool::new(db_path, 5)?; // Max 5 connections
136
137        // Initialize database schema
138        let conn = pool.get_connection()?;
139        Self::initialize_schema(&conn)?;
140
141        Ok(Database { pool })
142    }
143
144    fn initialize_schema(conn: &PooledConnection) -> Result<()> {
145        // Create chat_logs table with optimized schema
146        conn.execute(
147            "CREATE TABLE IF NOT EXISTS chat_logs (
148                id INTEGER PRIMARY KEY AUTOINCREMENT,
149                chat_id TEXT NOT NULL,
150                model TEXT NOT NULL,
151                question TEXT NOT NULL,
152                response TEXT NOT NULL,
153                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
154                input_tokens INTEGER,
155                output_tokens INTEGER
156            )",
157            [],
158        )?;
159
160        // Add token columns to existing table if they don't exist (migration)
161        let _ = conn.execute("ALTER TABLE chat_logs ADD COLUMN input_tokens INTEGER", []);
162        let _ = conn.execute("ALTER TABLE chat_logs ADD COLUMN output_tokens INTEGER", []);
163
164        // Create session_state table for tracking current session
165        conn.execute(
166            "CREATE TABLE IF NOT EXISTS session_state (
167                key TEXT PRIMARY KEY,
168                value TEXT NOT NULL
169            )",
170            [],
171        )?;
172
173        // Create optimized indexes for better performance
174        conn.execute(
175            "CREATE INDEX IF NOT EXISTS idx_chat_logs_chat_id ON chat_logs(chat_id)",
176            [],
177        )?;
178
179        conn.execute(
180            "CREATE INDEX IF NOT EXISTS idx_chat_logs_timestamp ON chat_logs(timestamp DESC)",
181            [],
182        )?;
183
184        // Additional index for model statistics
185        conn.execute(
186            "CREATE INDEX IF NOT EXISTS idx_chat_logs_model ON chat_logs(model)",
187            [],
188        )?;
189
190        Ok(())
191    }
192
193    pub fn save_chat_entry_with_tokens(
194        &self,
195        chat_id: &str,
196        model: &str,
197        question: &str,
198        response: &str,
199        input_tokens: Option<i32>,
200        output_tokens: Option<i32>,
201    ) -> Result<()> {
202        let conn = self.pool.get_connection()?;
203
204        conn.execute(
205            "INSERT INTO chat_logs (chat_id, model, question, response, timestamp, input_tokens, output_tokens)
206             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
207            params![chat_id, model, question, response, Utc::now(), input_tokens, output_tokens]
208        )?;
209        Ok(())
210    }
211
212    pub fn get_chat_history(&self, chat_id: &str) -> Result<Vec<ChatEntry>> {
213        let conn = self.pool.get_connection()?;
214
215        let mut stmt = conn.conn.as_ref().unwrap().prepare(
216            "SELECT id, chat_id, model, question, response, timestamp, input_tokens, output_tokens
217             FROM chat_logs
218             WHERE chat_id = ?1
219             ORDER BY timestamp ASC",
220        )?;
221
222        let rows = stmt.query_map([chat_id], |row| {
223            Ok(ChatEntry {
224                chat_id: row.get(1)?,
225                model: row.get(2)?,
226                question: row.get(3)?,
227                response: row.get(4)?,
228                timestamp: row.get(5)?,
229                input_tokens: row.get(6).ok(),
230                output_tokens: row.get(7).ok(),
231            })
232        })?;
233
234        let mut entries = Vec::new();
235        for row in rows {
236            entries.push(row?);
237        }
238
239        Ok(entries)
240    }
241
242    // Optimized version with LIMIT for better performance on large datasets
243    pub fn get_all_logs(&self) -> Result<Vec<ChatEntry>> {
244        self.get_recent_logs(None)
245    }
246
247    pub fn get_recent_logs(&self, limit: Option<usize>) -> Result<Vec<ChatEntry>> {
248        let conn = self.pool.get_connection()?;
249
250        let sql = if let Some(limit) = limit {
251            format!(
252                "SELECT id, chat_id, model, question, response, timestamp, input_tokens, output_tokens
253                 FROM chat_logs
254                 ORDER BY timestamp DESC
255                 LIMIT {}",
256                limit
257            )
258        } else {
259            "SELECT id, chat_id, model, question, response, timestamp, input_tokens, output_tokens
260             FROM chat_logs
261             ORDER BY timestamp DESC"
262                .to_string()
263        };
264
265        let mut stmt = conn.conn.as_ref().unwrap().prepare(&sql)?;
266
267        let rows = stmt.query_map([], |row| {
268            Ok(ChatEntry {
269                chat_id: row.get(1)?,
270                model: row.get(2)?,
271                question: row.get(3)?,
272                response: row.get(4)?,
273                timestamp: row.get(5)?,
274                input_tokens: row.get(6).ok(),
275                output_tokens: row.get(7).ok(),
276            })
277        })?;
278
279        let mut entries = Vec::new();
280        for row in rows {
281            entries.push(row?);
282        }
283
284        Ok(entries)
285    }
286
287    pub fn set_current_session_id(&self, session_id: &str) -> Result<()> {
288        let conn = self.pool.get_connection()?;
289
290        conn.execute(
291            "INSERT OR REPLACE INTO session_state (key, value) VALUES ('current_session', ?1)",
292            [session_id],
293        )?;
294        Ok(())
295    }
296
297    pub fn get_current_session_id(&self) -> Result<Option<String>> {
298        let conn = self.pool.get_connection()?;
299
300        let mut stmt = conn
301            .conn
302            .as_ref()
303            .unwrap()
304            .prepare("SELECT value FROM session_state WHERE key = 'current_session'")?;
305
306        let mut rows = stmt.query_map([], |row| Ok(row.get::<_, String>(0)?))?;
307
308        if let Some(row) = rows.next() {
309            Ok(Some(row?))
310        } else {
311            Ok(None)
312        }
313    }
314
315    pub fn purge_all_logs(&self) -> Result<()> {
316        let conn = self.pool.get_connection()?;
317
318        // Use transaction for atomic operation
319        conn.execute("BEGIN TRANSACTION", [])?;
320
321        match (|| -> Result<()> {
322            conn.execute("DELETE FROM chat_logs", [])?;
323            conn.execute("DELETE FROM session_state", [])?;
324            Ok(())
325        })() {
326            Ok(_) => {
327                conn.execute("COMMIT", [])?;
328                Ok(())
329            }
330            Err(e) => {
331                conn.execute("ROLLBACK", [])?;
332                Err(e)
333            }
334        }
335    }
336
337    /// Purge logs based on age (older than specified days)
338    pub fn purge_logs_by_age(&self, days: u32) -> Result<usize> {
339        let conn = self.pool.get_connection()?;
340
341        let cutoff_date = chrono::Utc::now() - chrono::Duration::days(days as i64);
342
343        let deleted_count =
344            conn.execute("DELETE FROM chat_logs WHERE timestamp < ?1", [cutoff_date])?;
345
346        Ok(deleted_count)
347    }
348
349    /// Purge logs to keep only the most recent N entries
350    pub fn purge_logs_keep_recent(&self, keep_count: usize) -> Result<usize> {
351        let conn = self.pool.get_connection()?;
352
353        // First, get the total count
354        let total_count: i64 =
355            conn.query_row("SELECT COUNT(*) FROM chat_logs", [], |row| row.get(0))?;
356
357        if total_count <= keep_count as i64 {
358            return Ok(0); // Nothing to purge
359        }
360
361        let to_delete = total_count - keep_count as i64;
362
363        let deleted_count = conn.execute(
364            "DELETE FROM chat_logs WHERE id IN (
365                SELECT id FROM chat_logs
366                ORDER BY timestamp ASC
367                LIMIT ?1
368            )",
369            [to_delete],
370        )?;
371
372        Ok(deleted_count)
373    }
374
375    /// Purge logs when database size exceeds threshold (in MB)
376    pub fn purge_logs_by_size(&self, max_size_mb: u64) -> Result<usize> {
377        let db_path = Self::database_path()?;
378        let current_size = std::fs::metadata(&db_path).map(|m| m.len()).unwrap_or(0);
379
380        let max_size_bytes = max_size_mb * 1024 * 1024;
381
382        if current_size <= max_size_bytes {
383            return Ok(0); // No purging needed
384        }
385
386        // Purge oldest 25% of entries to get under the size limit
387        let conn = self.pool.get_connection()?;
388        let total_count: i64 =
389            conn.query_row("SELECT COUNT(*) FROM chat_logs", [], |row| row.get(0))?;
390
391        let to_delete = (total_count as f64 * 0.25) as i64;
392
393        if to_delete > 0 {
394            let deleted_count = conn.execute(
395                "DELETE FROM chat_logs WHERE id IN (
396                    SELECT id FROM chat_logs
397                    ORDER BY timestamp ASC
398                    LIMIT ?1
399                )",
400                [to_delete],
401            )?;
402
403            // Run VACUUM to reclaim space
404            conn.execute("VACUUM", [])?;
405
406            Ok(deleted_count)
407        } else {
408            Ok(0)
409        }
410    }
411
412    /// Smart purge with configurable thresholds
413    pub fn smart_purge(
414        &self,
415        max_age_days: Option<u32>,
416        max_entries: Option<usize>,
417        max_size_mb: Option<u64>,
418    ) -> Result<usize> {
419        let mut total_deleted = 0;
420
421        // Purge by age first
422        if let Some(days) = max_age_days {
423            total_deleted += self.purge_logs_by_age(days)?;
424        }
425
426        // Then purge by count
427        if let Some(max_count) = max_entries {
428            total_deleted += self.purge_logs_keep_recent(max_count)?;
429        }
430
431        // Finally check size
432        if let Some(max_mb) = max_size_mb {
433            total_deleted += self.purge_logs_by_size(max_mb)?;
434        }
435
436        Ok(total_deleted)
437    }
438
439    pub fn clear_session(&self, session_id: &str) -> Result<()> {
440        let conn = self.pool.get_connection()?;
441
442        conn.execute("DELETE FROM chat_logs WHERE chat_id = ?1", [session_id])?;
443        Ok(())
444    }
445
446    pub fn get_stats(&self) -> Result<DatabaseStats> {
447        let conn = self.pool.get_connection()?;
448
449        // Use single query with subqueries for better performance
450        let total_entries: i64 =
451            conn.query_row("SELECT COUNT(*) FROM chat_logs", [], |row| row.get(0))?;
452
453        let unique_sessions: i64 =
454            conn.query_row("SELECT COUNT(DISTINCT chat_id) FROM chat_logs", [], |row| {
455                row.get(0)
456            })?;
457
458        // Get database file size
459        let db_path = Self::database_path()?;
460        let file_size = std::fs::metadata(&db_path).map(|m| m.len()).unwrap_or(0);
461
462        // Get date range in single query
463        let date_range = if total_entries > 0 {
464            let (earliest, latest): (Option<DateTime<Utc>>, Option<DateTime<Utc>>) = conn
465                .query_row(
466                    "SELECT MIN(timestamp), MAX(timestamp) FROM chat_logs",
467                    [],
468                    |row| Ok((row.get(0).ok(), row.get(1).ok())),
469                )?;
470
471            match (earliest, latest) {
472                (Some(e), Some(l)) => Some((e, l)),
473                _ => None,
474            }
475        } else {
476            None
477        };
478
479        // Get model usage statistics
480        let mut stmt = conn.conn.as_ref().unwrap().prepare(
481            "SELECT model, COUNT(*) as count FROM chat_logs GROUP BY model ORDER BY count DESC",
482        )?;
483
484        let model_stats = stmt
485            .query_map([], |row| {
486                Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)?))
487            })?
488            .collect::<Result<Vec<_>, _>>()?;
489
490        Ok(DatabaseStats {
491            total_entries: total_entries as usize,
492            unique_sessions: unique_sessions as usize,
493            file_size_bytes: file_size,
494            date_range,
495            model_usage: model_stats,
496        })
497    }
498
499    fn database_path() -> Result<PathBuf> {
500        // Use data_local_dir for cross-platform data storage
501        // On macOS: ~/Library/Application Support/lc
502        // On Linux: ~/.local/share/lc
503        // On Windows: %LOCALAPPDATA%/lc
504        let data_dir = dirs::data_local_dir()
505            .ok_or_else(|| anyhow::anyhow!("Could not find data directory"))?
506            .join("lc");
507        std::fs::create_dir_all(&data_dir)?;
508        Ok(data_dir.join("logs.db"))
509    }
510}
511
512// Thread-safe singleton for global database access
513
514#[cfg(test)]
515mod tests {
516    use super::*;
517    use tempfile::tempdir;
518
519    #[test]
520    fn test_connection_pool() {
521        let temp_dir = tempdir().unwrap();
522        let db_path = temp_dir.path().join("test.db");
523
524        let pool = ConnectionPool::new(db_path, 3).unwrap();
525
526        // Test getting multiple connections
527        let conn1 = pool.get_connection().unwrap();
528        let conn2 = pool.get_connection().unwrap();
529        let conn3 = pool.get_connection().unwrap();
530
531        // All connections should be valid
532        assert!(conn1.query_row("SELECT 1", [], |_| Ok(())).is_ok());
533        assert!(conn2.query_row("SELECT 1", [], |_| Ok(())).is_ok());
534        assert!(conn3.query_row("SELECT 1", [], |_| Ok(())).is_ok());
535    }
536
537    #[test]
538    fn test_optimized_database() {
539        let temp_dir = tempdir().unwrap();
540        std::env::set_var("HOME", temp_dir.path());
541
542        let db = Database::new().unwrap();
543
544        // Test saving and retrieving
545        db.save_chat_entry_with_tokens(
546            "test_session",
547            "test_model",
548            "test question",
549            "test response",
550            Some(100),
551            Some(50),
552        )
553        .unwrap();
554
555        let history = db.get_chat_history("test_session").unwrap();
556        assert_eq!(history.len(), 1);
557        assert_eq!(history[0].question, "test question");
558        assert_eq!(history[0].input_tokens, Some(100));
559        assert_eq!(history[0].output_tokens, Some(50));
560    }
561}
lc/database.rs

lc/
database.rs
