Implement #201

imessage-exporter/src/app/error.rs

@@ -7,14 +7,15 @@ use std::{
     io::Error as IoError,
 };
 
-use imessage_database::error::table::TableError;
+use imessage_database::{error::table::TableError, util::size::format_file_size};
 
 /// Errors that can happen during the application's runtime
 #[derive(Debug)]
 pub enum RuntimeError {
     InvalidOptions(String),
     DiskError(IoError),
     DatabaseError(TableError),
+    NotEnoughAvailableSpace(u64, u64),
 }
 
 impl Display for RuntimeError {
@@ -23,6 +24,14 @@ impl Display for RuntimeError {
             RuntimeError::InvalidOptions(why) => write!(fmt, "Invalid options!\n{why}"),
             RuntimeError::DiskError(why) => write!(fmt, "{why}"),
             RuntimeError::DatabaseError(why) => write!(fmt, "{why}"),
+            RuntimeError::NotEnoughAvailableSpace(estimated_bytes, available_bytes) => {
+                write!(
+                    fmt, 
+                    "Not enough free disk space!\nEstimated export size: {}\nDisk space available: {}",
+                    format_file_size(*estimated_bytes),
+                    format_file_size(*available_bytes)
+                )
+            }
         }
     }
 }

imessage-exporter/src/app/options.rs

@@ -29,6 +29,7 @@ pub const OPTION_END_DATE: &str = "end-date";
 pub const OPTION_DISABLE_LAZY_LOADING: &str = "no-lazy";
 pub const OPTION_CUSTOM_NAME: &str = "custom-name";
 pub const OPTION_PLATFORM: &str = "platform";
+pub const OPTION_BYPASS_FREE_SPACE_CHECK: &str = "ignore-disk-warning";
 
 // Other CLI Text
 pub const SUPPORTED_FILE_TYPES: &str = "txt, html";
@@ -62,6 +63,8 @@ pub struct Options {
     pub custom_name: Option<String>,
     /// The database source's platform
     pub platform: Platform,
+    /// If true, disable the free disk space check
+    pub ignore_disk_space: bool,
 }
 
 impl Options {
@@ -77,6 +80,7 @@ impl Options {
         let no_lazy = args.get_flag(OPTION_DISABLE_LAZY_LOADING);
         let custom_name: Option<&String> = args.get_one(OPTION_CUSTOM_NAME);
         let platform_type: Option<&String> = args.get_one(OPTION_PLATFORM);
+        let ignore_disk_space = args.get_flag(OPTION_BYPASS_FREE_SPACE_CHECK);
 
         // Build the export type
         let export_type: Option<ExportType> = match export_file_type {
@@ -213,6 +217,7 @@ impl Options {
             no_lazy,
             custom_name: custom_name.cloned(),
             platform,
+            ignore_disk_space,
         })
     }
 
@@ -232,15 +237,20 @@ fn validate_path(
     export_path: Option<&String>,
     export_type: &Option<&ExportType>,
 ) -> Result<PathBuf, RuntimeError> {
+    // Build a path from the user-provided data or the default location
     let resolved_path =
         PathBuf::from(export_path.unwrap_or(&format!("{}/{DEFAULT_OUTPUT_DIR}", home())));
+
+    // If there is an export type selected, ensure we do not overwrite files of the same type
     if let Some(export_type) = export_type {
         if resolved_path.exists() {
+            // Get the word to use if there is a problem with the specified path
             let path_word = match export_path {
                 Some(_) => "Specified",
                 None => "Default",
             };
 
+            // Ensure the directory exists and does not contain files of the same export type
             match resolved_path.read_dir() {
                 Ok(files) => {
                     let export_type_extension = export_type.to_string();
@@ -311,7 +321,7 @@ fn get_command() -> Command {
             Arg::new(OPTION_ATTACHMENT_ROOT)
                 .short('r')
                 .long(OPTION_ATTACHMENT_ROOT)
-                .help(format!("Specify an optional custom path to look for attachments in (macOS only).\nOnly use this if attachments are stored separately from the database's default location.\nThe default location is {DEFAULT_ATTACHMENT_ROOT}\n"))
+                .help(format!("Specify an optional custom path to look for attachments in (macOS only)\nOnly use this if attachments are stored separately from the database's default location\nThe default location is {DEFAULT_ATTACHMENT_ROOT}\n"))
                 .display_order(4)
                 .value_name("path/to/attachments"),
         )
@@ -335,15 +345,15 @@ fn get_command() -> Command {
             Arg::new(OPTION_START_DATE)
                 .short('s')
                 .long(OPTION_START_DATE)
-                .help("The start date filter. Only messages sent on or after this date will be included\n")
+                .help("The start date filter\nOnly messages sent on or after this date will be included\n")
                 .display_order(7)
                 .value_name("YYYY-MM-DD"),
         )
         .arg(
             Arg::new(OPTION_END_DATE)
                 .short('e')
                 .long(OPTION_END_DATE)
-                .help("The end date filter. Only messages sent before this date will be included\n")
+                .help("The end date filter\nOnly messages sent before this date will be included\n")
                 .display_order(8)
                 .value_name("YYYY-MM-DD"),
         )
@@ -362,6 +372,14 @@ fn get_command() -> Command {
                 .help("Specify an optional custom name for the database owner's messages in exports\n")
                 .display_order(10)
         )
+        .arg(
+            Arg::new(OPTION_BYPASS_FREE_SPACE_CHECK)
+                .short('b')
+                .long(OPTION_BYPASS_FREE_SPACE_CHECK)
+                .help("Bypass the disk space check when exporting data\nBy default, exports will not run if there is not enough free disk space\n")
+                .action(ArgAction::SetTrue)
+                .display_order(11)
+        )
 }
 
 /// Parse arguments from the command line
@@ -403,6 +421,7 @@ mod arg_tests {
             no_lazy: false,
             custom_name: None,
             platform: Platform::default(),
+            ignore_disk_space: false,
         };
 
         assert_eq!(actual, expected);
@@ -476,25 +495,27 @@ mod arg_tests {
     #[test]
     fn can_build_option_export_html() {
         // Get matches from sample args
-        let cli_args: Vec<&str> = vec!["imessage-exporter", "-f", "html"];
+        let cli_args: Vec<&str> = vec!["imessage-exporter", "-f", "html", "-o", "/tmp"];
         let command = get_command();
         let args = command.get_matches_from(cli_args);
 
         // Build the Options
         let actual = Options::from_args(&args).unwrap();
 
         // Expected data
+        let tmp_dir = String::from("/tmp");
         let expected = Options {
             db_path: default_db_path(),
             attachment_root: None,
             attachment_manager: AttachmentManager::default(),
             diagnostic: false,
             export_type: Some(ExportType::HTML),
-            export_path: validate_path(None, &None).unwrap(),
+            export_path: validate_path(Some(&tmp_dir), &None).unwrap(),
             query_context: QueryContext::default(),
             no_lazy: false,
             custom_name: None,
             platform: Platform::default(),
+            ignore_disk_space: false,
         };
 
         assert_eq!(actual, expected);
@@ -522,6 +543,7 @@ mod arg_tests {
             no_lazy: true,
             custom_name: None,
             platform: Platform::default(),
+            ignore_disk_space: false,
         };
 
         assert_eq!(actual, expected);

imessage-exporter/src/app/runtime.rs

@@ -5,6 +5,7 @@ use std::{
     path::PathBuf,
 };
 
+use fs2::available_space;
 use rusqlite::Connection;
 
 use crate::{
@@ -24,11 +25,11 @@ use imessage_database::{
         handle::Handle,
         messages::Message,
         table::{
-            get_connection, Cacheable, Deduplicate, Diagnostic, ATTACHMENTS_DIR, MAX_LENGTH, ME,
-            ORPHANED, UNKNOWN,
+            get_connection, get_db_size, Cacheable, Deduplicate, Diagnostic, ATTACHMENTS_DIR,
+            MAX_LENGTH, ME, ORPHANED, UNKNOWN,
         },
     },
-    util::{dates::get_offset, export_type::ExportType},
+    util::{dates::get_offset, export_type::ExportType, size::format_file_size},
 };
 
 /// Stores the application state and handles application lifecycle
@@ -218,6 +219,47 @@ impl Config {
         })
     }
 
+    /// Ensure there is available disk space for the requested export
+    fn ensure_free_space(&self) -> Result<(), RuntimeError> {
+        // Export size is usually about 6% the size of the db; we divide by 10 to over-estimate about 10% of the total size
+        // for some safe headroom
+        let total_db_size =
+            get_db_size(&self.options.db_path).map_err(RuntimeError::DatabaseError)?;
+        let mut estimated_export_size = total_db_size / 10;
+
+        let free_space_at_location = available_space(&self.options.export_path).unwrap();
+
+        // Validate that there is enough disk space free to write the export
+        match self.options.attachment_manager {
+            AttachmentManager::Disabled => {
+                if estimated_export_size >= free_space_at_location {
+                    return Err(RuntimeError::NotEnoughAvailableSpace(
+                        estimated_export_size,
+                        free_space_at_location,
+                    ));
+                }
+            }
+            _ => {
+                let total_attachment_size = Attachment::get_total_attachment_bytes(&self.db)
+                    .map_err(RuntimeError::DatabaseError)?;
+                estimated_export_size += total_attachment_size;
+                if (estimated_export_size + total_attachment_size) >= free_space_at_location {
+                    return Err(RuntimeError::NotEnoughAvailableSpace(
+                        estimated_export_size + total_attachment_size,
+                        free_space_at_location,
+                    ));
+                }
+            }
+        };
+
+        println!(
+            "Estimated export size: {}",
+            format_file_size(estimated_export_size)
+        );
+
+        Ok(())
+    }
+
     /// Handles diagnostic tests for database
     fn run_diagnostic(&self) -> Result<(), TableError> {
         println!("\niMessage Database Diagnostics\n");
@@ -227,18 +269,27 @@ impl Config {
         ChatToHandle::run_diagnostic(&self.db)?;
 
         // Global Diagnostics
+        println!("Global diagnostic data:");
+
+        let total_db_size = get_db_size(&self.options.db_path)?;
+        println!(
+            "    Total database size: {}",
+            format_file_size(total_db_size)
+        );
+
         let unique_handles: HashSet<i32> =
             HashSet::from_iter(self.real_participants.values().cloned());
         let duplicated_handles = self.participants.len() - unique_handles.len();
         if duplicated_handles > 0 {
-            println!("Duplicated contacts: {duplicated_handles}");
+            println!("    Duplicated contacts: {duplicated_handles}");
         }
 
         let unique_chats: HashSet<i32> = HashSet::from_iter(self.real_chatrooms.values().cloned());
         let duplicated_chats = self.chatrooms.len() - unique_chats.len();
         if duplicated_chats > 0 {
-            println!("Duplicated chats: {duplicated_chats}");
+            println!("    Duplicated chats: {duplicated_chats}");
         }
+
         Ok(())
     }
 
@@ -264,11 +315,17 @@ impl Config {
         } else if let Some(export_type) = &self.options.export_type {
             // Ensure the path we want to export to exists
             create_dir_all(&self.options.export_path).map_err(RuntimeError::DiskError)?;
+
             // Ensure the path we want to copy attachments to exists, if requested
             if !matches!(self.options.attachment_manager, AttachmentManager::Disabled) {
                 create_dir_all(self.attachment_path()).map_err(RuntimeError::DiskError)?;
             }
 
+            // Ensure there is enough free disk space to write the export
+            if self.options.ignore_disk_space {
+                self.ensure_free_space()?;
+            }
+
             // Create exporter, pass it data we care about, then kick it off
             match export_type {
                 ExportType::HTML => {
@@ -324,6 +381,7 @@ mod filename_tests {
             no_lazy: false,
             custom_name: None,
             platform: Platform::macOS,
+            ignore_disk_space: false,
         }
     }
 
@@ -553,6 +611,7 @@ mod who_tests {
             no_lazy: false,
             custom_name: None,
             platform: Platform::macOS,
+            ignore_disk_space: false,
         }
     }
 
@@ -775,6 +834,7 @@ mod directory_tests {
             no_lazy: false,
             custom_name: None,
             platform: Platform::macOS,
+            ignore_disk_space: false,
         }
     }
 

imessage-exporter/src/exporters/html.rs

@@ -1298,6 +1298,7 @@ mod tests {
             no_lazy: false,
             custom_name: None,
             platform: Platform::macOS,
+            ignore_disk_space: false,
         }
     }
 

imessage-exporter/src/exporters/txt.rs

@@ -858,6 +858,7 @@ mod tests {
             no_lazy: false,
             custom_name: None,
             platform: Platform::macOS,
+            ignore_disk_space: false,
         }
     }