Skip to main content

task_graph_mcp/db/
mod.rs

1//! Database layer for the Task Graph MCP Server.
2
3pub mod agents;
4pub mod attachments;
5pub mod deps;
6pub mod locks;
7pub mod schema;
8pub mod search;
9pub mod state_transitions;
10pub mod stats;
11pub mod tasks;
12
13pub use search::{AttachmentMatch, SearchResult};
14
15use anyhow::Result;
16use rusqlite::Connection;
17use std::path::Path;
18use std::sync::{Arc, Mutex};
19
20mod embedded {
21    use refinery::embed_migrations;
22    embed_migrations!("migrations");
23}
24
25/// Database handle wrapping a SQLite connection.
26#[derive(Clone)]
27pub struct Database {
28    conn: Arc<Mutex<Connection>>,
29}
30
31impl Database {
32    /// Open or create the database at the given path.
33    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
34        let conn = Connection::open(path)?;
35
36        // Enable WAL mode for concurrent access
37        conn.execute_batch(
38            "PRAGMA journal_mode=WAL;
39             PRAGMA foreign_keys=ON;
40             PRAGMA busy_timeout=5000;",
41        )?;
42
43        let db = Self {
44            conn: Arc::new(Mutex::new(conn)),
45        };
46
47        db.run_migrations()?;
48
49        Ok(db)
50    }
51
52    /// Open an in-memory database (for testing).
53    #[allow(dead_code)]
54    pub fn open_in_memory() -> Result<Self> {
55        let conn = Connection::open_in_memory()?;
56
57        conn.execute_batch("PRAGMA foreign_keys=ON;")?;
58
59        let db = Self {
60            conn: Arc::new(Mutex::new(conn)),
61        };
62
63        db.run_migrations()?;
64
65        Ok(db)
66    }
67
68    /// Run database migrations.
69    fn run_migrations(&self) -> Result<()> {
70        let mut conn = self.conn.lock().unwrap();
71        embedded::migrations::runner().run(&mut *conn)?;
72        Ok(())
73    }
74
75    /// Execute a function with exclusive access to the connection.
76    pub fn with_conn<F, T>(&self, f: F) -> Result<T>
77    where
78        F: FnOnce(&Connection) -> Result<T>,
79    {
80        let conn = self.conn.lock().unwrap();
81        f(&conn)
82    }
83
84    /// Execute a function with mutable access to the connection (for transactions).
85    pub fn with_conn_mut<F, T>(&self, f: F) -> Result<T>
86    where
87        F: FnOnce(&mut Connection) -> Result<T>,
88    {
89        let mut conn = self.conn.lock().unwrap();
90        f(&mut conn)
91    }
92}
93
94/// Get the current timestamp in milliseconds.
95pub fn now_ms() -> i64 {
96    chrono::Utc::now().timestamp_millis()
97}