feat: Add a new database client to store shared state for persistence

Author: dormant-user

Cargo.toml

@@ -41,3 +41,4 @@ uuid = { version = "1.23.1", features = ["v4"] }
 dotenv = "0.15.0"
 url = "2.5.8"
 regex = "1.12.3"
+rusqlite = { version = "0.31", features = ["bundled"] }

src/database.rs

@@ -0,0 +1,109 @@
+use crate::settings::{PutItem, RsyncTrack, Status};
+use rusqlite::{params, Connection};
+use std::collections::HashMap;
+
+/// Opens (or creates) the SQLite database and ensures the schema exists.
+pub fn open() -> Connection {
+    let conn = Connection::open("rutorrent.db").expect("Failed to open database");
+    conn.execute_batch(
+        "CREATE TABLE IF NOT EXISTS state (
+            hash        TEXT PRIMARY KEY,
+            name        TEXT NOT NULL,
+            status      TEXT NOT NULL,
+            progress    REAL NOT NULL DEFAULT 0.0,
+            url         TEXT NOT NULL,
+            save_path   TEXT NOT NULL,
+            remote_host TEXT NOT NULL,
+            remote_user TEXT NOT NULL,
+            remote_path TEXT NOT NULL,
+            delete_after_copy INTEGER NOT NULL DEFAULT 0
+        );",
+    )
+        .expect("Failed to create schema");
+    conn
+}
+
+/// Inserts or replaces a tracked torrent entry.
+pub fn upsert(conn: &Connection, hash: &str, entry: &RsyncTrack) {
+    let (status, progress) = encode_status(&entry.status);
+    conn.execute(
+        "INSERT OR REPLACE INTO state
+            (hash, name, status, progress, url, save_path, remote_host, remote_user, remote_path, delete_after_copy)
+         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
+        params![
+            hash,
+            entry.name,
+            status,
+            progress,
+            entry.put_item.url,
+            entry.put_item.save_path,
+            entry.put_item.remote_host,
+            entry.put_item.remote_username,
+            entry.put_item.remote_path,
+            entry.put_item.delete_after_copy as i32,
+        ],
+    )
+        .expect("Failed to upsert state");
+}
+
+/// Removes a torrent entry by hash.
+pub fn remove(conn: &Connection, hash: &str) {
+    conn.execute("DELETE FROM state WHERE hash = ?1", params![hash])
+        .expect("Failed to remove state");
+}
+
+/// Loads all persisted entries back into a HashMap on startup.
+pub fn load_all(conn: &Connection) -> HashMap<String, RsyncTrack> {
+    let mut stmt = conn
+        .prepare("SELECT hash, name, status, progress, url, save_path, remote_host, remote_user, remote_path, delete_after_copy FROM state")
+        .expect("Failed to prepare load query");
+
+    stmt.query_map([], |row| {
+        let hash: String = row.get(0)?;
+        let name: String = row.get(1)?;
+        let status_str: String = row.get(2)?;
+        let progress: f64 = row.get(3)?;
+        let url: String = row.get(4)?;
+        let save_path: String = row.get(5)?;
+        let remote_host: String = row.get(6)?;
+        let remote_username: String = row.get(7)?;
+        let remote_path: String = row.get(8)?;
+        let delete_after_copy: i32 = row.get(9)?;
+
+        let status = decode_status(&status_str, progress);
+        let put_item = PutItem {
+            url,
+            name: None,
+            hash: None,
+            trackers: None,
+            save_path,
+            remote_host,
+            remote_username,
+            remote_path,
+            delete_after_copy: delete_after_copy != 0,
+        };
+
+        Ok((hash, RsyncTrack { name, status, put_item }))
+    })
+        .expect("Failed to query state")
+        .filter_map(|r| r.ok())
+        .collect()
+}
+
+fn encode_status(status: &Status) -> (&'static str, f64) {
+    match status {
+        Status::Downloading(p) => ("Downloading", *p),
+        Status::Copying(p)     => ("Copying", *p),
+        Status::Completed      => ("Completed", 1.0),
+        Status::Failed         => ("Failed", 0.0),
+    }
+}
+
+fn decode_status(s: &str, progress: f64) -> Status {
+    match s {
+        "Copying"   => Status::Copying(progress),
+        "Completed" => Status::Completed,
+        "Failed"    => Status::Failed,
+        _           => Status::Downloading(progress),
+    }
+}

src/lib.rs

@@ -14,6 +14,7 @@ mod settings;
 mod squire;
 mod swagger;
 mod telegram;
+mod database;
 
 /// Contains entrypoint and initializer settings to trigger the asynchronous `HTTPServer`
 ///
@@ -29,13 +30,16 @@ pub async fn start() -> std::io::Result<()> {
     squire::load_env_file();
     let config = settings::Config::new();
     logger::init_logger(config.utc_logger, config.log_level);
-    let state: settings::SharedState = Arc::new(RwLock::new(HashMap::new()));
+    let db_conn = database::open();
+    let initial_state = database::load_all(&db_conn);
+    log::info!("Loaded {} entries from database", initial_state.len());
+    let state: settings::SharedState = Arc::new(RwLock::new(initial_state));
     let pending: settings::PendingMap = Arc::new(RwLock::new(HashMap::new()));
 
     let client = qb::client(&config)
         .await
         .expect("Failed to authenticate qBittorrent");
-    squire::spawn_worker(client, state.clone(), pending.clone(), config.clone());
+    squire::spawn_worker(client, state.clone(), pending.clone(), config.clone(), db_conn);
 
     let host = config.host.clone();
     let port = config.port;
@@ -61,6 +65,7 @@ pub async fn start() -> std::io::Result<()> {
             .route("/torrent", web::delete().to(api::delete_torrent))
             .route("/swagger", web::get().to(swagger::redirector))
             .route("/ui", web::get().to(swagger::redirector))
+            .route("/", web::get().to(swagger::redirector))
             .service(swagger::service())
     })
     .bind((host, port))?

src/squire.rs

@@ -1,4 +1,4 @@
-use crate::{ntfy, qb, rsync, settings, telegram};
+use crate::{database, ntfy, qb, rsync, settings, telegram};
 use regex::Regex;
 use reqwest::Client;
 use serde_json::Value;
@@ -124,7 +124,9 @@ pub fn spawn_worker(
     state: settings::SharedState,
     pending: settings::PendingMap,
     config: settings::Config,
+    db_conn: rusqlite::Connection,
 ) {
+    let db_conn = std::sync::Mutex::new(db_conn);
     tokio::spawn(async move {
         log::info!("Worker started");
 
@@ -204,6 +206,9 @@ pub fn spawn_worker(
                 if !returned.contains(h.as_str()) {
                     log::info!("Torrent removed from QBitAPI, dropping from state: {}", h);
                     db.remove(h);
+                    if let Ok(conn) = db_conn.lock() {
+                        database::remove(&conn, h);
+                    }
                 }
             });
 
@@ -229,6 +234,9 @@ pub fn spawn_worker(
                             config_cloned,
                         );
                         db.remove(&hash);
+                        if let Ok(conn) = db_conn.lock() {
+                            database::remove(&conn, &hash);
+                        }
                     }
 
                     settings::Status::Completed => {
@@ -266,6 +274,9 @@ pub fn spawn_worker(
                             }
                         }
                         db.remove(&hash);
+                        if let Ok(conn) = db_conn.lock() {
+                            database::remove(&conn, &hash);
+                        }
                     }
 
                     settings::Status::Downloading(_) => {