use std::sync::{Mutex, Arc}; use std::ffi::OsString; use std::fs::{File,read_to_string}; use std::io::{Read, Seek, SeekFrom}; use std::time::UNIX_EPOCH; use std::collections::HashMap; use sha2::{Sha256, Digest}; use walkdir::DirEntry; use rusqlite::{params, Connection}; use crate::document_extraction::extract_document_text; #[derive(Debug, Clone)] pub struct FileMetadata { pub path: String, pub size: u64, pub moddate: u64, pub hash: Vec, } #[derive(Debug)] pub struct BatchUpdate { pub files_to_update: Vec<(DirEntry, FileMetadata)>, pub files_to_insert: Vec, } const HASHLEN:usize = 1024 * 8; const MAXIMUM_TEXT_SIZE:usize = 1024 * 512; const MAXIMUM_FILE_SIZE:u64 = 1024 * 1024 * 50; const PLAINTEXT_EXTENSIONS_LIST: [&'static str; 86] = ["","txt","rtf","log", // Text Documents "csv", // Spreadsheet "sh","bat","cmd","bash","ps1","psm1","psd1","pssc","psrc", // Scripts "c","cpp","i","cs","csx","caki", // C# "cpp","cc","cxx","c++","hpp","hh","hxx","h","ii", // C++ "tex","bib","bbx","cbx", // LaTeX "css","xml","md","json","yaml","yml", // Markup Languages and others "html","htm","shtml","xhtml","xht","mdoc","jsp","asp","aspx","jshtm", // HTML "js","cjs","mjs","es6","es","jsx","ts","tsx", // Javascript and TypeScript "cfg","conf","ini","gitattributes","gitignore", // Config and related files "java","jav", // Java "pl","pm","pod","t","psgi", // Perl "php","php4","php5","phtml","ctp", // PHP "py","rpy","pyw","cpy","gyp","gypi","pyi","ipy","pyt","ipynb", // Python "wasm","wat", // Web Assembly ]; const SUPPORTED_DOCUMENT_EXTENSIONS_LIST: [&'static str; 9] = ["odt", "docx", "doc", // Office Documents "ppt", "pptx", "odp", // Presentation "xls", "xlsx", "ods"]; // Spreadsheet /// Load existing file metadata from database indexed by path pub fn load_existing_files(conn: &Connection) -> Result, rusqlite::Error> { let mut existing_files = HashMap::new(); let mut stmt = conn.prepare("SELECT path, size, moddate, hash FROM files")?; let rows = stmt.query_map([], |row| { Ok(FileMetadata { path: row.get(0)?, size: row.get(1)?, moddate: row.get(2)?, hash: row.get(3)?, }) })?; for row in rows { let metadata = row?; existing_files.insert(metadata.path.clone(), metadata); } Ok(existing_files) } /// Analyze files and determine which need updates vs inserts pub fn analyze_files_for_batch_update( entries: &[DirEntry], existing_files: &HashMap ) -> BatchUpdate { let mut files_to_update = Vec::new(); let mut files_to_insert = Vec::new(); for entry in entries { let meta = match entry.metadata() { Ok(m) if !m.is_dir() => m, _ => continue, }; let fpath = match entry.path().canonicalize() { Ok(fp) => fp.to_string_lossy().to_string(), Err(_) => continue, }; let fmodified = match meta.modified() .ok() .and_then(|m| m.duration_since(UNIX_EPOCH).ok().map(|d| d.as_secs())) { Some(time) => time, None => continue, }; if let Some(existing_metadata) = existing_files.get(&fpath) { // File exists in database, check if modification date changed if existing_metadata.moddate != fmodified { files_to_update.push((entry.clone(), existing_metadata.clone())); } // If moddate is same, skip processing this file entirely } else { // New file, needs to be inserted files_to_insert.push(entry.clone()); } } BatchUpdate { files_to_update, files_to_insert, } } /// Get a hash of a file by reading the first and last HASHLEN bytes of the file fn get_file_hash(size: u64, path: OsString) -> Result, std::io::Error> { let mut hasher = Sha256::new(); let mut f: File = File::open(path)?; hasher.update(&size.to_le_bytes()); if size > HASHLEN as u64 { let mut file_start_block = [0u8; HASHLEN]; f.read_exact(&mut file_start_block)?; hasher.update(file_start_block); f.seek(SeekFrom::End(0 - HASHLEN as i64))?; let mut file_end_block = [0u8; HASHLEN]; f.read_exact(&mut file_end_block)?; hasher.update(file_end_block); } else if size > 0 { let mut file_block = Vec::new(); f.read_to_end(&mut file_block)?; hasher.update(file_block); } drop(f); Ok(hasher.finalize().to_vec()) } /// Process updated files in batch with transaction pub fn process_batch_updates( conn_mutex: &Arc>, files_to_update: &[(DirEntry, FileMetadata)], stop_flag: &Arc>, status_callback: Option> ) -> Result<(), String> { if files_to_update.is_empty() { return Ok(()); } const BATCH_SIZE: usize = 1000; let total_files = files_to_update.len(); // Process files in batches of BATCH_SIZE for (batch_idx, batch) in files_to_update.chunks(BATCH_SIZE).enumerate() { // Check stop flag at the start of each batch if *stop_flag.lock().unwrap() { return Ok(()); } let conn = conn_mutex.lock().unwrap(); let tx = conn.unchecked_transaction().map_err(|e| format!("Failed to begin transaction: {}", e))?; for (i, (entry, _old_metadata)) in batch.iter().enumerate() { let global_index = batch_idx * BATCH_SIZE + i + 1; // Check stop flag if *stop_flag.lock().unwrap() { drop(tx); drop(conn); return Ok(()); } // Update status with current file if let Some(ref callback) = status_callback { let filename = entry.path().file_name() .and_then(|n| n.to_str()) .unwrap_or("unknown"); callback(&format!("Updating file {}/{}: {}", global_index, total_files, filename)); } let meta = entry.metadata().map_err(|e| format!("Failed to get metadata: {}", e))?; if meta.is_dir() { continue; } let fpath = match entry.path().canonicalize() { Ok(fp) => fp.into_os_string(), Err(_) => continue, }; let fsize = meta.len(); let fmodified = meta.modified() .map_err(|e| format!("Failed to get modified time: {}", e))? .duration_since(UNIX_EPOCH) .map_err(|e| format!("Failed to calculate duration: {}", e))? .as_secs(); let fhash = get_file_hash(fsize, fpath.clone()) .map_err(|e| format!("Failed to calculate hash: {}", e))?; // Check stop flag after hash calculation if *stop_flag.lock().unwrap() { drop(tx); drop(conn); return Ok(()); } // Update files table tx.execute( "UPDATE files SET size = ?1, moddate = ?2, hash = ?3 WHERE path = ?4", params![fsize, fmodified, fhash, fpath.to_string_lossy()] ).map_err(|e| format!("Failed to update file record: {}", e))?; // Delete old searchable text entry tx.execute( "DELETE FROM searchabletext WHERE path = ?1", params![fpath.to_string_lossy()] ).map_err(|e| format!("Failed to delete old searchable text: {}", e))?; // Insert new searchable text if applicable if fsize <= MAXIMUM_FILE_SIZE { let default_ext = OsString::new(); let file_extension = entry.path().extension().unwrap_or(&default_ext) .to_ascii_lowercase().to_str().unwrap_or("").to_string(); let ext_str = file_extension.as_str(); if PLAINTEXT_EXTENSIONS_LIST.contains(&ext_str) { if let Ok(file_string) = read_to_string(&fpath) { let trimmed_file_string = if file_string.len() > MAXIMUM_TEXT_SIZE { file_string[..MAXIMUM_TEXT_SIZE].to_string() } else { file_string }; let fname = entry.path().file_name().unwrap().to_os_string(); tx.execute( "INSERT INTO searchabletext VALUES (?1, ?2, ?3)", params![fname.to_str(), fpath.to_string_lossy(), trimmed_file_string] ).map_err(|e| format!("Failed to insert searchable text: {}", e))?; } } else if SUPPORTED_DOCUMENT_EXTENSIONS_LIST.contains(&ext_str) { if let Ok(extracted_text) = extract_document_text(&fpath, ext_str) { if !extracted_text.trim().is_empty() { let trimmed_file_string = if extracted_text.len() > MAXIMUM_TEXT_SIZE { extracted_text[..MAXIMUM_TEXT_SIZE].to_string() } else { extracted_text }; let fname = entry.path().file_name().unwrap().to_os_string(); tx.execute( "INSERT INTO searchabletext VALUES (?1, ?2, ?3)", params![fname.to_str(), fpath.to_string_lossy(), trimmed_file_string] ).map_err(|e| format!("Failed to insert document text: {}", e))?; } } } } } tx.commit().map_err(|e| format!("Failed to commit transaction: {}", e))?; } Ok(()) } /// Process new files in batch with transaction pub fn process_batch_inserts( conn_mutex: &Arc>, files_to_insert: &[DirEntry], stop_flag: &Arc>, status_callback: Option> ) -> Result<(), String> { if files_to_insert.is_empty() { return Ok(()); } const BATCH_SIZE: usize = 1000; let total_files = files_to_insert.len(); // Process files in batches of BATCH_SIZE for (batch_idx, batch) in files_to_insert.chunks(BATCH_SIZE).enumerate() { // Check stop flag at the start of each batch if *stop_flag.lock().unwrap() { return Ok(()); } let conn = conn_mutex.lock().unwrap(); let tx = conn.unchecked_transaction().map_err(|e| format!("Failed to begin transaction: {}", e))?; for (i, entry) in batch.iter().enumerate() { let global_index = batch_idx * BATCH_SIZE + i + 1; // Check stop flag if *stop_flag.lock().unwrap() { drop(tx); drop(conn); return Ok(()); } // Update status with current file if let Some(ref callback) = status_callback { let filename = entry.path().file_name() .and_then(|n| n.to_str()) .unwrap_or("unknown"); callback(&format!("Indexing file {}/{}: {}", global_index, total_files, filename)); } let meta = entry.metadata().map_err(|e| format!("Failed to get metadata: {}", e))?; if meta.is_dir() { continue; } let fpath = match entry.path().canonicalize() { Ok(fp) => fp.into_os_string(), Err(_) => continue, }; let fsize = meta.len(); let fmodified = meta.modified() .map_err(|e| format!("Failed to get modified time: {}", e))? .duration_since(UNIX_EPOCH) .map_err(|e| format!("Failed to calculate duration: {}", e))? .as_secs(); let fhash = get_file_hash(fsize, fpath.clone()) .map_err(|e| format!("Failed to calculate hash: {}", e))?; let fname = entry.path().file_name().unwrap().to_os_string(); // Insert into files table tx.execute( "INSERT INTO files VALUES (?1, ?2, ?3, ?4, ?5)", params![fname.to_str(), fpath.to_string_lossy(), fsize, fmodified, fhash] ).map_err(|e| format!("Failed to insert file record: {}", e))?; // Insert searchable text if applicable if fsize <= MAXIMUM_FILE_SIZE { let default_ext = OsString::new(); let file_extension = entry.path().extension().unwrap_or(&default_ext) .to_ascii_lowercase().to_str().unwrap_or("").to_string(); let ext_str = file_extension.as_str(); if PLAINTEXT_EXTENSIONS_LIST.contains(&ext_str) { if let Ok(file_string) = read_to_string(&fpath) { let trimmed_file_string = if file_string.len() > MAXIMUM_TEXT_SIZE { file_string[..MAXIMUM_TEXT_SIZE].to_string() } else { file_string }; tx.execute( "INSERT INTO searchabletext VALUES (?1, ?2, ?3)", params![fname.to_str(), fpath.to_string_lossy(), trimmed_file_string] ).map_err(|e| format!("Failed to insert searchable text: {}", e))?; } } else if SUPPORTED_DOCUMENT_EXTENSIONS_LIST.contains(&ext_str) { if let Ok(extracted_text) = extract_document_text(&fpath, ext_str) { if !extracted_text.trim().is_empty() { let trimmed_file_string = if extracted_text.len() > MAXIMUM_TEXT_SIZE { extracted_text[..MAXIMUM_TEXT_SIZE].to_string() } else { extracted_text }; tx.execute( "INSERT INTO searchabletext VALUES (?1, ?2, ?3)", params![fname.to_str(), fpath.to_string_lossy(), trimmed_file_string] ).map_err(|e| format!("Failed to insert document text: {}", e))?; } } } } } tx.commit().map_err(|e| format!("Failed to commit transaction: {}", e))?; } Ok(()) } pub fn process_entry(conn_mutex: &Arc>, entry: DirEntry) { let meta = entry.metadata().unwrap(); if !meta.is_dir() { // let fpath = entry.path().canonicalize()?.into_os_string(); let fpath_result = entry.path().canonicalize(); let fpath = match fpath_result { Ok(fp) => fp.into_os_string(), Err(error) => { println!("Error converting fpath: {:?}", error); return; }, }; // Get basic file properties let fsize = meta.len(); let fmodified = meta.modified().unwrap().duration_since(UNIX_EPOCH).unwrap().as_secs(); // Get file hash let fhash_result = get_file_hash(fsize, fpath.clone()); let fhash = match fhash_result { Ok(fh) => fh, Err(error) => { println!("Error digesting hash: {:?}", error); return; }, }; // let fhash = b""; // Insert results into database let query = "INSERT INTO files VALUES (?1,?2,?3,?4,?5)"; let conn = conn_mutex.lock().unwrap(); let mut stmt = conn.prepare_cached(query).unwrap(); let fname = entry.path().file_name().unwrap().to_os_string(); let stmt_result = stmt.execute(params![fname.to_str(), fpath.to_string_lossy(), fsize, fmodified, fhash]); match stmt_result { Ok(us) => us, Err(error) => { println!("Error with sqlite transaction: {:?}", error); return; }, }; std::mem::drop(stmt); // Free the mutex lock so that other threads can access the database std::mem::drop(conn); // Generate searchable plain text for file if applicable if fsize <= MAXIMUM_FILE_SIZE { let default_ext = OsString::new(); let file_extension = entry.path().extension().unwrap_or(&default_ext).to_ascii_lowercase().to_str().unwrap().to_string(); let ext_str = file_extension.as_str(); if PLAINTEXT_EXTENSIONS_LIST.contains(&ext_str) { let file_contents_result = read_to_string(fpath.clone()); let file_string = match file_contents_result { Ok(fs) => fs, Err(_error) => { // println!("Error reading file to string: {:?}", error); return; }, }; let trimmed_file_string; // Trim file contents if too large if file_string.len() > MAXIMUM_TEXT_SIZE { trimmed_file_string = file_string[..MAXIMUM_TEXT_SIZE].to_string(); } else { trimmed_file_string = file_string; } // Insert file contents into database let query2 = "INSERT INTO searchabletext VALUES (?1,?2,?3)"; let conn2 = conn_mutex.lock().unwrap(); let mut stmt2 = conn2.prepare_cached(query2).unwrap(); let stmt_result2 = stmt2.execute(params![fname.to_str(), fpath.to_string_lossy(), trimmed_file_string]); match stmt_result2 { Ok(us) => us, Err(error) => { println!("Error with sqlite transaction 2: {:?}", error); return; }, }; } else if SUPPORTED_DOCUMENT_EXTENSIONS_LIST.contains(&ext_str) { // Extract text from office documents match extract_document_text(&fpath, ext_str) { Ok(extracted_text) => { if !extracted_text.trim().is_empty() { let trimmed_file_string; // Trim file contents if too large if extracted_text.len() > MAXIMUM_TEXT_SIZE { trimmed_file_string = extracted_text[..MAXIMUM_TEXT_SIZE].to_string(); } else { trimmed_file_string = extracted_text; } // Insert file contents into database let query2 = "INSERT INTO searchabletext VALUES (?1,?2,?3)"; let conn2 = conn_mutex.lock().unwrap(); let mut stmt2 = conn2.prepare_cached(query2).unwrap(); let stmt_result2 = stmt2.execute(params![fname.to_str(), fpath.to_string_lossy(), trimmed_file_string]); match stmt_result2 { Ok(_) => {}, Err(error) => { println!("Error with sqlite transaction for document: {:?}", error); }, }; } } Err(error) => { println!("Error extracting text from document {}: {:?}", fpath.to_string_lossy(), error); } } } } } }