Skip to main content

tally_cli/
database.rs

1use crate::models::Counter;
2use anyhow::Result;
3use fs2::FileExt;
4use sqlite::ConnectionThreadSafe;
5use std::fs::File;
6
7pub struct Connection {
8    conn: ConnectionThreadSafe,
9    lock_file: File,
10}
11
12impl Connection {
13    pub fn get(&self) -> &ConnectionThreadSafe {
14        &self.conn
15    }
16
17    pub fn new(name: &str) -> Result<Connection> {
18        // Acquire file lock
19        let lock_path = format!("{}.lock", name);
20        let lock_file = File::create(&lock_path)?;
21
22        if lock_file.try_lock_exclusive().is_err() {
23            eprintln!("tally: waiting for another instance to release {lock_path}");
24            lock_file.lock_exclusive()?;
25        }
26
27        let mut connection = sqlite::Connection::open_thread_safe(name)?;
28        connection.set_busy_timeout(5_000)?;
29        connection.execute("PRAGMA journal_mode = WAL;")?;
30
31        let mut conn = Connection {
32            conn: connection,
33            lock_file,
34        };
35        conn.init_database()?;
36        Ok(conn)
37    }
38
39    fn init_database(&mut self) -> Result<()> {
40        // create the default table
41        self.conn.execute(
42            "
43            CREATE TABLE IF NOT EXISTS counters (
44                name TEXT PRIMARY KEY,
45                count INTEGER NOT NULL,
46                step INTEGER NOT NULL,
47                template TEXT NOT NULL
48            );
49            ",
50        )?;
51
52        self.conn.execute(
53            "
54            CREATE TABLE IF NOT EXISTS default_counter (
55                name TEXT NOT NULL,
56                timestamp DATETIME NOT NULL,
57                FOREIGN KEY (name) REFERENCES counters(name)
58            );
59            ",
60        )?;
61
62        // Setup default counter
63        let mut stmt = self.conn.prepare("SELECT COUNT(*) FROM counters;")?;
64        if let Some(row) = stmt.iter().next() {
65            if let sqlite::Value::Integer(count) = &row?[0] {
66                if *count == 0 {
67                    let default = Counter::new("tally");
68                    default.insert(&self.conn)?;
69                    default.set_default(&self.conn)?;
70                }
71            }
72        }
73
74        Ok(())
75    }
76}
77
78impl Drop for Connection {
79    fn drop(&mut self) {
80        if let Err(e) = fs2::FileExt::unlock(&self.lock_file) {
81            eprintln!("Warning: Failed to unlock file: {}", e);
82        }
83    }
84}