diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/db.rs | 130 | ||||
| -rw-r--r-- | src/files.rs | 108 | ||||
| -rw-r--r-- | src/main.rs | 17 |
3 files changed, 116 insertions, 139 deletions
@@ -1,8 +1,6 @@ use anyhow::{Result, anyhow}; use directories::BaseDirs; -use r2d2::{Pool, PooledConnection}; -use r2d2_sqlite::SqliteConnectionManager; -use rusqlite::params; +use rusqlite::{Connection, params}; use std::{ path::{Path, PathBuf}, time::UNIX_EPOCH, @@ -22,36 +20,36 @@ const DEDUPE_DB: &str = "DELETE FROM flacs WHERE rowid NOT IN (SELECT MAX(rowid) FROM flacs GROUP BY path)"; const GET_MODTIME: &str = "SELECT modtime FROM flacs WHERE path = ?1"; -pub fn open_db( - path: Option<impl AsRef<Path>>, - threads: usize, -) -> Result<Pool<SqliteConnectionManager>> { - if let Some(file) = path { - let manager = SqliteConnectionManager::file(file); - let pool = Pool::builder().max_size(threads as u32).build(manager)?; - let conn = pool.get()?; - conn.execute(TABLE_CREATE, ())?; - Ok(pool) - } else if let Some(base_dir) = BaseDirs::new() { - let file = Path::new(base_dir.data_dir()).join("reencoder.db"); - let manager = SqliteConnectionManager::file(file); - let pool = Pool::builder().max_size(threads as u32).build(manager)?; - let conn = pool.get()?; - conn.execute(TABLE_CREATE, ())?; - Ok(pool) - } else { - Err(anyhow!("Failed to locate data directory")) - } +pub trait Database { + type Conn; + fn new(path: Option<impl AsRef<Path>>) -> Result<Self::Conn>; + fn insert_file(&self, filename: impl AsRef<Path>) -> Result<()>; + fn update_file(&self, filename: impl AsRef<Path>) -> Result<()>; + fn check_file(&self, filename: impl AsRef<Path>) -> Result<bool>; + fn init_clean_files(&self) -> Result<Vec<PathBuf>, rusqlite::Error>; + fn remove_file(&self, filename: impl AsRef<Path>) -> Result<()>; + fn get_toencode_files(&self) -> Result<Vec<PathBuf>, rusqlite::Error>; + fn get_toencode_number(&self) -> Result<u64, rusqlite::Error>; + fn get_modtime(&self, file: impl AsRef<Path>) -> Result<u64>; + fn vacuum(&self) -> Result<()>; } -pub struct Database(pub PooledConnection<SqliteConnectionManager>); - -impl Database { - pub fn new(conn: PooledConnection<SqliteConnectionManager>) -> Self { - Database(conn) +impl Database for Connection { + type Conn = Connection; + fn new(path: Option<impl AsRef<Path>>) -> Result<Self> { + let conn = if let Some(file) = path { + Connection::open(file)? + } else if let Some(base_dir) = BaseDirs::new() { + let file = Path::new(base_dir.data_dir()).join("reencoder.db"); + Connection::open(file)? + } else { + return Err(anyhow!("Failed to locate data directory")); + }; + conn.execute(TABLE_CREATE, ())?; + Ok(conn) } - pub fn insert_file(&self, filename: impl AsRef<Path>) -> Result<()> { + fn insert_file(&self, filename: impl AsRef<Path>) -> Result<()> { let toencode = !matches!(get_vendor(&filename)?.as_str(), CURRENT_VENDOR); let modtime = filename @@ -61,7 +59,7 @@ impl Database { .duration_since(UNIX_EPOCH)? .as_secs(); - self.0.execute( + self.execute( ADD_NEW_ITEM, params![filename.as_ref().to_str().unwrap(), toencode, modtime], )?; @@ -69,7 +67,7 @@ impl Database { Ok(()) } - pub fn update_file(&self, filename: impl AsRef<Path>) -> Result<()> { + fn update_file(&self, filename: impl AsRef<Path>) -> Result<()> { let modtime = filename .as_ref() .metadata()? @@ -77,7 +75,7 @@ impl Database { .duration_since(UNIX_EPOCH)? .as_secs(); - self.0.execute( + self.execute( REPLACE_ITEM, params![filename.as_ref().to_str().unwrap(), false, modtime], )?; @@ -85,8 +83,8 @@ impl Database { Ok(()) } - pub fn check_file(&self, filename: impl AsRef<Path>) -> Result<bool> { - if self.0.query_one( + fn check_file(&self, filename: impl AsRef<Path>) -> Result<bool> { + if self.query_one( CHECK_FILE, params!(filename.as_ref().to_str().unwrap()), |row| { @@ -100,9 +98,9 @@ impl Database { } } - pub fn init_clean_files(&self) -> Result<Vec<PathBuf>, rusqlite::Error> { - self.0.execute(DEDUPE_DB, ())?; - let mut stmt = self.0.prepare(FETCH_FILES)?; + fn init_clean_files(&self) -> Result<Vec<PathBuf>, rusqlite::Error> { + self.execute(DEDUPE_DB, ())?; + let mut stmt = self.prepare(FETCH_FILES)?; let mut rows = stmt.query(())?; let mut files = Vec::new(); while let Ok(Some(row)) = rows.next() { @@ -112,14 +110,13 @@ impl Database { Ok(files) } - pub fn remove_file(&self, filename: impl AsRef<Path>) -> Result<()> { - self.0 - .execute(REMOVE_FILE, params!(filename.as_ref().to_str().unwrap()))?; + fn remove_file(&self, filename: impl AsRef<Path>) -> Result<()> { + self.execute(REMOVE_FILE, params!(filename.as_ref().to_str().unwrap()))?; Ok(()) } - pub fn get_toencode_files(&self) -> Result<Vec<PathBuf>, rusqlite::Error> { - let mut stmt = self.0.prepare(TOENCODE_QUERY)?; + fn get_toencode_files(&self) -> Result<Vec<PathBuf>, rusqlite::Error> { + let mut stmt = self.prepare(TOENCODE_QUERY)?; let mut rows = stmt.query(())?; let mut files: Vec<PathBuf> = Vec::new(); while let Ok(Some(row)) = rows.next() { @@ -129,15 +126,15 @@ impl Database { Ok(files) } - pub fn get_toencode_number(&self) -> Result<u64, rusqlite::Error> { - self.0.query_one(TOENCODE_NUMBER, (), |row| { + fn get_toencode_number(&self) -> Result<u64, rusqlite::Error> { + self.query_one(TOENCODE_NUMBER, (), |row| { let num: u64 = row.get(0)?; Ok(num) }) } - pub fn get_modtime(&self, file: impl AsRef<Path>) -> Result<u64> { - Ok(self.0.query_one( + fn get_modtime(&self, file: impl AsRef<Path>) -> Result<u64> { + Ok(self.query_one( GET_MODTIME, params![file.as_ref().to_str().unwrap()], |row| { @@ -147,8 +144,8 @@ impl Database { )?) } - pub fn vaccum(&self) -> Result<()> { - self.0.execute("VACUUM", ())?; + fn vacuum(&self) -> Result<()> { + self.execute("VACUUM", ())?; Ok(()) } } @@ -163,12 +160,11 @@ mod tests { let dbname = String::from("temp1.db"); let filenames = ["16bit.flac", "24bit.flac", "32bit.flac"]; let mut counter = 0; - let pool = open_db(Some(&dbname), 10).unwrap(); - let conn = Database::new(pool.get().unwrap()); + let conn = Connection::new(Some(&dbname)).unwrap(); for file in filenames { conn.insert_file(&file.to_string()).unwrap(); } - let mut stmt = conn.0.prepare(TOENCODE_QUERY).unwrap(); + let mut stmt = conn.prepare(TOENCODE_QUERY).unwrap(); let mut returned = stmt.query(()).unwrap(); while let Ok(Some(_)) = returned.next() { @@ -182,27 +178,25 @@ mod tests { fn check_update() { let dbname = String::from("temp2.db"); let filenames = ["16bit.flac", "24bit.flac", "32bit.flac"]; - let pool = open_db(Some(&dbname), 10).unwrap(); - let conn = Database::new(pool.get().unwrap()); + let conn = Connection::new(Some(&dbname)).unwrap(); for file in filenames { conn.insert_file(Path::new(file).canonicalize().unwrap()) .unwrap(); } - conn.0 - .execute( - REPLACE_ITEM, - params![ - Path::new("16bit.flac") - .canonicalize() - .unwrap() - .to_str() - .unwrap(), - true, - "" - ], - ) - .unwrap(); + conn.execute( + REPLACE_ITEM, + params![ + Path::new("16bit.flac") + .canonicalize() + .unwrap() + .to_str() + .unwrap(), + true, + "" + ], + ) + .unwrap(); conn.update_file( Path::new("16bit.flac") @@ -213,7 +207,7 @@ mod tests { ) .unwrap(); - let mut stmt = conn.0.prepare(TOENCODE_QUERY).unwrap(); + let mut stmt = conn.prepare(TOENCODE_QUERY).unwrap(); let mut returned = stmt.query(()).unwrap(); let mut counter = 0; while let Ok(Some(_)) = returned.next() { diff --git a/src/files.rs b/src/files.rs index 46d7727..6b5c20e 100644 --- a/src/files.rs +++ b/src/files.rs @@ -1,15 +1,13 @@ use anyhow::{Result, anyhow}; #[cfg(not(test))] use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle}; -use r2d2::Pool; -use r2d2_sqlite::SqliteConnectionManager; -use rayon::prelude::*; +use rusqlite::Connection; use std::{ error::Error, fmt::Display, path::{Path, PathBuf}, sync::{ - Arc, + Arc, Mutex, atomic::{AtomicBool, Ordering}, }, time::UNIX_EPOCH, @@ -51,7 +49,7 @@ impl Display for FileError { impl Error for FileError {} -fn handle_file(file: impl AsRef<Path>, conn: &Database) -> Result<()> { +fn handle_file(file: impl AsRef<Path>, conn: &Connection) -> Result<()> { if conn.check_file(&file)? { let modtime = file .as_ref() @@ -77,7 +75,7 @@ fn handle_file(file: impl AsRef<Path>, conn: &Database) -> Result<()> { pub fn index_files_recursively( path: impl AsRef<Path>, - pool: &Pool<SqliteConnectionManager>, + conn: &Connection, handler: Arc<AtomicBool>, ) -> Result<()> { if !path.as_ref().is_dir() { @@ -105,7 +103,6 @@ pub fn index_files_recursively( } } - let conn = Database::new(pool.get()?); for entry in WalkDir::new(abspath) { if handler.load(Ordering::SeqCst) { let path = entry.unwrap().into_path(); @@ -113,7 +110,7 @@ pub fn index_files_recursively( continue; } if path.extension().is_some_and(|x| x == "flac") { - if let Err(error) = handle_file(&path, &conn) { + if let Err(error) = handle_file(&path, conn) { eprintln!("{}", FileError::new(path, error)); } else { #[cfg(not(test))] @@ -136,11 +133,7 @@ pub fn index_files_recursively( Ok(()) } -pub fn reencode_files( - pool: &Pool<SqliteConnectionManager>, - handler: Arc<AtomicBool>, -) -> Result<()> { - let conn = Database::new(pool.get()?); +pub fn reencode_files(conn: Connection, handler: Arc<AtomicBool>, threads: usize) -> Result<()> { #[cfg(not(test))] let bar = ProgressBar::with_draw_target( Some(conn.get_toencode_number()?), @@ -150,33 +143,42 @@ pub fn reencode_files( .with_message("Reencoding"); let files = conn.get_toencode_files()?; - drop(conn); - files.par_iter().for_each(|file| { - if handler.load(Ordering::SeqCst) { - let conn = match pool.get() { - Ok(conn) => Database::new(conn), - Err(error) => { - eprintln!("{}", FileError::new(file, error.into())); - return; - } - }; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build()?; - if !file.exists() { - let _ = conn.remove_file(file); - #[cfg(not(test))] - bar.dec_length(1); - return; - } + let lock = Arc::new(Mutex::new(conn)); + + pool.scope(|scope| { + for file in files { + if handler.load(Ordering::SeqCst) { + scope.spawn(|_| { + let newconn = if let Ok(conn) = lock.lock() { + conn + } else { + eprintln!("Error setting up lock on file:\t{}", file.to_string_lossy()); + return; + }; + if !file.exists() { + let _ = newconn.remove_file(&file); + #[cfg(not(test))] + bar.dec_length(1); + return; + } - if let Err(error) = handle_encode(file) { - eprintln!("{}", FileError::new(file, error)); + if let Err(error) = handle_encode(&file) { + eprintln!("{}", FileError::new(&file, error)); + } else { + if let Err(error) = newconn.update_file(&file) { + eprintln!("{}", FileError::new(file, error)); + } + #[cfg(not(test))] + bar.inc(1) + } + }); } else { - if let Err(error) = conn.update_file(file) { - eprintln!("{}", FileError::new(file, error)); - } - #[cfg(not(test))] - bar.inc(1) + break; } } }); @@ -192,10 +194,8 @@ pub fn reencode_files( Ok(()) } -pub fn clean_files(pool: &Pool<SqliteConnectionManager>, handler: Arc<AtomicBool>) -> Result<()> { - let conn = Database::new(pool.get()?); +pub fn clean_files(conn: &Connection, handler: Arc<AtomicBool>) -> Result<()> { let files = conn.init_clean_files()?; - drop(conn); #[cfg(not(test))] let spinner = ProgressBar::with_draw_target(None, ProgressDrawTarget::stdout_with_hz(60)) @@ -203,13 +203,6 @@ pub fn clean_files(pool: &Pool<SqliteConnectionManager>, handler: Arc<AtomicBool files.iter().for_each(|file| { if handler.load(Ordering::SeqCst) && !file.exists() { - let conn = match pool.get() { - Ok(conn) => Database::new(conn), - Err(error) => { - eprintln!("{}", FileError::new(file, error.into())); - return; - } - }; if let Err(error) = conn.remove_file(file) { eprintln!("{}", FileError::new(file, error)) }; @@ -220,8 +213,7 @@ pub fn clean_files(pool: &Pool<SqliteConnectionManager>, handler: Arc<AtomicBool #[cfg(not(test))] spinner.finish(); - let conn = Database::new(pool.get()?); - conn.vaccum()?; + conn.vacuum()?; Ok(()) } @@ -229,14 +221,13 @@ pub fn clean_files(pool: &Pool<SqliteConnectionManager>, handler: Arc<AtomicBool #[cfg(test)] mod tests { use super::*; - use crate::db::*; #[test] fn test_index_lots_of_files() { let dbname = "temp3.db"; let handler = Arc::new(AtomicBool::new(true)); - let pool = open_db(Some(dbname), 10).unwrap(); - index_files_recursively(Path::new("./testfiles"), &pool, handler).unwrap(); + let conn = Connection::new(Some(&dbname)).unwrap(); + index_files_recursively(Path::new("./testfiles"), &conn, handler).unwrap(); std::fs::remove_file(dbname).unwrap(); } @@ -244,8 +235,7 @@ mod tests { fn test_clean_files() { let dbname = "temp4.db"; let handler = Arc::new(AtomicBool::new(true)); - let pool = open_db(Some(dbname), 10).unwrap(); - let conn = Database::new(pool.get().unwrap()); + let conn = Connection::new(Some(&dbname)).unwrap(); let filenames = ["16bit.flac", "24bit.flac", "32bit.flac", "nonexisting.flac"]; std::fs::copy("32bit.flac", "nonexisting.flac").unwrap(); for file in filenames { @@ -254,7 +244,7 @@ mod tests { std::fs::remove_file("nonexisting.flac").unwrap(); - clean_files(&pool, handler).unwrap(); + clean_files(&conn, handler).unwrap(); let counter = conn.init_clean_files().unwrap().len(); std::fs::remove_file(dbname).unwrap(); assert!(counter == 3) @@ -264,14 +254,12 @@ mod tests { fn test_reencode_lots_of_files() { let dbname = "temp5.db"; let handler = Arc::new(AtomicBool::new(true)); - let pool = open_db(Some(dbname), 10).unwrap(); + let conn = Connection::new(Some(&dbname)).unwrap(); let temp = handler.clone(); - index_files_recursively(Path::new("./testfiles"), &pool, temp).unwrap(); - let conn = Database::new(pool.get().unwrap()); + index_files_recursively(Path::new("./testfiles"), &conn, temp).unwrap(); println!("\n{}", conn.get_toencode_number().unwrap()); - drop(conn); - reencode_files(&pool, handler).unwrap(); - let conn = Database::new(pool.get().unwrap()); + reencode_files(conn, handler, 2).unwrap(); + let conn = Connection::new(Some(&dbname)).unwrap(); println!("\n{}", conn.get_toencode_number().unwrap()); std::fs::remove_file(dbname).unwrap(); } diff --git a/src/main.rs b/src/main.rs index 1363a77..d1e4562 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,6 +5,7 @@ use anyhow::Result; use clap::{Arg, ArgAction, Command, ValueHint, command, value_parser}; use clap_complete::{Generator, Shell, generate}; use console::style; +use rusqlite::Connection; use std::{ path::PathBuf, sync::{ @@ -92,36 +93,30 @@ fn main() -> Result<()> { r.store(false, Ordering::SeqCst); })?; - let threads = *args.get_one::<usize>("threads").unwrap(); - - let dbpool = db::open_db(args.get_one::<PathBuf>("db"), threads)?; + let conn = Connection::new(args.get_one::<PathBuf>("db"))?; let path = args.get_one::<PathBuf>("path"); if path.is_none() && !args.get_flag("clean") && !args.get_flag("doit") { - let conn = Database::new(dbpool.get()?); let count = conn.get_toencode_number()?; println!("Files to reencode:\t{}", style(count).green()); return Ok(()); } - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(threads) - .build()?; - if let Some(realpath) = path { let hanlder = running.clone(); - files::index_files_recursively(realpath, &dbpool, hanlder)?; + files::index_files_recursively(realpath, &conn, hanlder)?; } if args.get_flag("clean") { let handler = running.clone(); - pool.install(|| files::clean_files(&dbpool, handler))?; + files::clean_files(&conn, handler)?; } if args.get_flag("doit") { let hanlder = running.clone(); - pool.install(|| files::reencode_files(&dbpool, hanlder))?; + let threads = *args.get_one::<usize>("threads").unwrap(); + files::reencode_files(conn, hanlder, threads)?; } Ok::<(), anyhow::Error>(()) } |