diff --git a/README.md b/README.md index a9e4ca8..d9d0c7a 100644 --- a/README.md +++ b/README.md @@ -543,6 +543,9 @@ This creates a `flags_manifest.json` with the status of each email. ### 9. Restore Backup to IMAP Server Restore emails from a local backup to any IMAP server. +By default, restore runs in **incremental mode**: it uploads only emails that are not already present on the destination (based on `Message-ID`). +Use `--full-restore` to force the legacy behavior (process all emails and re-sync labels/flags for already-present messages). + ```bash # Restore all folders from backup python3 restore_imap_emails.py \ @@ -551,6 +554,14 @@ python3 restore_imap_emails.py \ --dest-user "you@gmail.com" \ --dest-pass "your-app-password" +# Force full restore (legacy behavior) +python3 restore_imap_emails.py \ + --src-path "./my_backup" \ + --dest-host "imap.gmail.com" \ + --dest-user "you@gmail.com" \ + --dest-pass "your-app-password" \ + --full-restore + # Restore with flags (read/starred status) python3 restore_imap_emails.py \ --src-path "./my_backup" \ diff --git a/src/imap_common.py b/src/imap_common.py index 7d03822..fb93905 100644 --- a/src/imap_common.py +++ b/src/imap_common.py @@ -4,6 +4,8 @@ IMAP Common Utilities Shared functionality for IMAP migration, counting, and comparison scripts. """ +from __future__ import annotations + import imaplib import os import re @@ -58,6 +60,55 @@ CMD_FETCH = "fetch" OP_ADD_FLAGS = "+FLAGS" +def ensure_folder_exists(imap_conn, folder_name: str) -> None: + """Best-effort create of a folder if it doesn't already exist. + + IMAP servers typically return an error if the mailbox exists; this helper + intentionally ignores those errors. + """ + try: + if folder_name and folder_name.upper() != FOLDER_INBOX: + imap_conn.create(f'"{folder_name}"') + except Exception: + # Folder may already exist, or server may restrict creation. + pass + + +def append_email( + imap_conn, + folder_name: str, + raw_content: bytes, + date_str: str, + flags: str | None = None, + *, + ensure_folder: bool = True, +) -> bool: + """Append an email message to a folder. + + This is intentionally a thin wrapper around IMAP APPEND; callers can + perform duplicate checks or folder selection separately. + + Args: + flags: Optional IMAP flags. If provided, they are normalized to a + parenthesized list before being passed to `imaplib.IMAP4.append`. + ensure_folder: If True, attempts to create the folder first (best-effort). + """ + if ensure_folder: + ensure_folder_exists(imap_conn, folder_name) + + normalized_flags = None + if flags: + stripped = str(flags).strip() + if stripped: + if stripped.startswith("(") and stripped.endswith(")"): + normalized_flags = stripped + else: + normalized_flags = f"({stripped})" + + resp, _ = imap_conn.append(f'"{folder_name}"', normalized_flags, date_str, raw_content) + return resp == "OK" + + def verify_env_vars(vars_list): """ Checks if all environment variables in the list are set. diff --git a/src/migrate_imap_emails.py b/src/migrate_imap_emails.py index 4770f22..adfe753 100644 --- a/src/migrate_imap_emails.py +++ b/src/migrate_imap_emails.py @@ -344,17 +344,10 @@ def process_batch( safe_print(f"Error selecting folder {folder_name} in worker: {e}") return - def ensure_dest_folder(folder): - try: - if folder.upper() != imap_common.FOLDER_INBOX: - dest.create(f'"{folder}"') - except Exception: - pass - # In non-Gmail-mode, we keep a selected destination folder for efficiency if not gmail_mode: try: - ensure_dest_folder(folder_name) + imap_common.ensure_folder_exists(dest, folder_name) dest.select(f'"{folder_name}"') except Exception as e: safe_print(f"Error selecting folder {folder_name} in worker: {e}") @@ -424,7 +417,7 @@ def process_batch( target_folder = folder_name remaining_labels = [] - ensure_dest_folder(target_folder) + imap_common.ensure_folder_exists(dest, target_folder) dest.select(f'"{target_folder}"') is_duplicate = bool(msg_id and check_duplicate and imap_common.message_exists_in_folder(dest, msg_id)) @@ -435,7 +428,14 @@ def process_batch( sync_flags_on_existing(dest, target_folder, msg_id, flags, size) else: valid_flags = f"({flags})" if (preserve_flags and flags) else None - dest.append(f'"{target_folder}"', valid_flags, date_str, msg_content) + imap_common.append_email( + dest, + target_folder, + msg_content, + date_str, + valid_flags, + ensure_folder=False, + ) safe_print(f"[{target_folder}] {'COPIED':<12} | {size_str:<8} | {subject[:40]}") if preserve_flags and flags: for flag in flags.split(): @@ -450,11 +450,18 @@ def process_batch( if label_folder in (imap_common.GMAIL_ALL_MAIL, imap_common.GMAIL_SPAM, imap_common.GMAIL_TRASH): continue try: - ensure_dest_folder(label_folder) + imap_common.ensure_folder_exists(dest, label_folder) dest.select(f'"{label_folder}"') if not imap_common.message_exists_in_folder(dest, msg_id): valid_flags = f"({flags})" if (preserve_flags and flags) else None - dest.append(f'"{label_folder}"', valid_flags, date_str, msg_content) + imap_common.append_email( + dest, + label_folder, + msg_content, + date_str, + valid_flags, + ensure_folder=False, + ) safe_print(f" -> Applied label: {label}") if preserve_flags and flags: for flag in flags.split(): diff --git a/src/restore_cache.py b/src/restore_cache.py new file mode 100644 index 0000000..dcce176 --- /dev/null +++ b/src/restore_cache.py @@ -0,0 +1,254 @@ +"""Restore progress cache. + +This module supports faster incremental restores by persisting, per destination+folder, +the set of Message-IDs already seen/processed by this tool. + +The caller decides where the cache file lives by passing a cache directory/root. +""" + +from __future__ import annotations + +import copy +import json +import os +import re +import threading +import time +from collections.abc import Callable + +RESTORE_CACHE_VERSION = 1 + +# Throttle disk writes so we can update frequently without rewriting a large JSON file +# on every single message. +_MIN_SECONDS_BETWEEN_SAVES = 2.0 +_MIN_PENDING_UPDATES_BEFORE_SAVE = 50 + + +def _safe_cache_component(value: str) -> str: + return re.sub(r"[^A-Za-z0-9_.-]+", "_", value or "") + + +def _prepare_cache_for_json(cache_data: dict) -> dict: + """Create a deep copy of cache_data suitable for JSON serialization. + Removes in-memory helper fields like _ids_set that should not be persisted. + """ + snapshot = copy.deepcopy(cache_data) + folders = snapshot.get("folders", {}) + if isinstance(folders, dict): + for folder_entry in folders.values(): + if isinstance(folder_entry, dict): + # Remove the in-memory set cache + folder_entry.pop("_ids_set", None) + return snapshot + + +def get_dest_index_cache_path(cache_root: str, dest_host: str, dest_user: str) -> str: + safe_host = _safe_cache_component(dest_host) + safe_user = _safe_cache_component(dest_user) + return os.path.join(cache_root, f"restore_cache_{safe_host}_{safe_user}.json") + + +def load_dest_index_cache(cache_path: str) -> dict: + try: + with open(cache_path, encoding="utf-8") as f: + data = json.load(f) + if not isinstance(data, dict): + return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}} + if data.get("version") != RESTORE_CACHE_VERSION: + return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}} + if not isinstance(data.get("folders"), dict): + data["folders"] = {} + if not isinstance(data.get("_meta"), dict): + data["_meta"] = {} + return data + except FileNotFoundError: + return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}} + except Exception: + return {"version": RESTORE_CACHE_VERSION, "folders": {}, "_meta": {}} + + +def save_dest_index_cache(cache_path: str, cache_data: dict) -> bool: + try: + tmp_path = f"{cache_path}.tmp" + with open(tmp_path, "w", encoding="utf-8") as f: + json.dump(cache_data, f, ensure_ascii=False) + os.replace(tmp_path, cache_path) + return True + except Exception: + # Cache is best-effort; do not fail restore. + return False + + +def _ensure_dest(cache_data: dict, dest_host: str, dest_user: str) -> None: + cache_dest = cache_data.get("dest") + if not isinstance(cache_dest, dict) or cache_dest.get("host") != dest_host or cache_dest.get("user") != dest_user: + cache_data.clear() + cache_data.update( + { + "version": RESTORE_CACHE_VERSION, + "dest": {"host": dest_host, "user": dest_user}, + "folders": {}, + "_meta": {}, + } + ) + + +def get_cached_message_ids( + cache_data: dict, + cache_lock: threading.Lock, + dest_host: str, + dest_user: str, + folder_name: str, +) -> set[str]: + """Return Message-IDs we've already seen/processed for this destination+folder.""" + with cache_lock: + _ensure_dest(cache_data, dest_host, dest_user) + folders = cache_data.setdefault("folders", {}) + entry = folders.get(folder_name) + if not isinstance(entry, dict): + return set() + ids = entry.get("message_ids") + if not isinstance(ids, list): + return set() + return {str(x) for x in ids if x} + + +def add_cached_message_id( + cache_data: dict, + cache_lock: threading.Lock, + dest_host: str, + dest_user: str, + folder_name: str, + message_id: str, +) -> bool: + """Add a Message-ID to the cache. Returns True if it was newly added.""" + if not message_id: + return False + msg_id = str(message_id).strip() + if not msg_id: + return False + + with cache_lock: + _ensure_dest(cache_data, dest_host, dest_user) + folders = cache_data.setdefault("folders", {}) + entry = folders.setdefault(folder_name, {}) + if not isinstance(entry, dict): + folders[folder_name] = {} + entry = folders[folder_name] + + ids = entry.get("message_ids") + if not isinstance(ids, list): + ids = [] + entry["message_ids"] = ids + + # Maintain entry-level in-memory set for O(1) lookups (not persisted to JSON) + ids_set = entry.get("_ids_set") + if ids_set is None: + ids_set = set(ids) + entry["_ids_set"] = ids_set + + if msg_id in ids_set: + return False + + ids.append(msg_id) + ids_set.add(msg_id) + + meta = cache_data.setdefault("_meta", {}) + if not isinstance(meta, dict): + meta = {} + cache_data["_meta"] = meta + meta["pending_updates"] = int(meta.get("pending_updates") or 0) + 1 + return True + + +def maybe_save_dest_index_cache( + cache_path: str, + cache_data: dict, + cache_lock: threading.Lock, + *, + force: bool = False, + log_fn: Callable[[str], None] | None = None, +) -> bool: + """Persist cache to disk if enough updates/time has accumulated.""" + now = time.time() + pending = 0 + with cache_lock: + meta = cache_data.setdefault("_meta", {}) + if not isinstance(meta, dict): + meta = {} + cache_data["_meta"] = meta + + pending = int(meta.get("pending_updates") or 0) + last_saved = float(meta.get("last_saved_ts") or 0.0) + + should_save = force or ( + pending > 0 + and (pending >= _MIN_PENDING_UPDATES_BEFORE_SAVE or (now - last_saved) >= _MIN_SECONDS_BETWEEN_SAVES) + ) + if not should_save: + return False + + # Update meta before writing so the on-disk file reflects the flush decision. + meta["pending_updates"] = 0 + meta["last_saved_ts"] = now + + # Prepare a clean snapshot for JSON serialization (removes in-memory helper fields) + snapshot = _prepare_cache_for_json(cache_data) + + did_write = save_dest_index_cache(cache_path, snapshot) + if did_write and log_fn is not None: + log_fn(f"Wrote restore cache ({pending} updates): {cache_path}") + return did_write + + +def record_progress( + *, + message_id: str | None, + folder_name: str, + existing_dest_msg_ids: set[str] | None, + existing_dest_msg_ids_lock: threading.Lock | None, + progress_cache_path: str | None, + progress_cache_data: dict | None, + progress_cache_lock: threading.Lock | None, + dest_host: str | None, + dest_user: str | None, + log_fn: Callable[[str], None] | None = None, +) -> None: + """Record a processed Message-ID for fast skipping on future incremental runs. + + Updates both: + - the in-memory set used by the current run, and + - the persisted progress cache on disk (throttled writes). + """ + if not message_id: + return + + msg_id = str(message_id).strip() + if not msg_id: + return + + if existing_dest_msg_ids is not None and existing_dest_msg_ids_lock is not None: + with existing_dest_msg_ids_lock: + existing_dest_msg_ids.add(msg_id) + + if ( + progress_cache_path + and progress_cache_data is not None + and progress_cache_lock is not None + and dest_host + and dest_user + ): + add_cached_message_id( + progress_cache_data, + progress_cache_lock, + dest_host, + dest_user, + folder_name, + msg_id, + ) + maybe_save_dest_index_cache( + progress_cache_path, + progress_cache_data, + progress_cache_lock, + log_fn=log_fn, + ) diff --git a/src/restore_imap_emails.py b/src/restore_imap_emails.py index aad0aaf..a29f195 100644 --- a/src/restore_imap_emails.py +++ b/src/restore_imap_emails.py @@ -56,9 +56,21 @@ import time from email import policy from email.parser import BytesParser from email.utils import parsedate_to_datetime +from enum import Enum +from typing import Optional import imap_common import imap_oauth2 +import restore_cache + + +class UploadResult(Enum): + """Result of an email upload operation.""" + + SUCCESS = "success" + ALREADY_EXISTS = "already_exists" + FAILURE = "failure" + # Defaults MAX_WORKERS = 4 # Lower default for restore to avoid rate limits @@ -293,7 +305,7 @@ def email_exists_in_folder(imap_conn, message_id): def upload_email(dest, folder_name, raw_content, date_str, message_id, flags=None, check_duplicate=True): """ Upload a single email to the destination folder. - Returns True on success, False if duplicate or on failure. + Returns UploadResult enum: SUCCESS, ALREADY_EXISTS, or FAILURE. Args: flags: Optional string of IMAP flags like "\\Seen" for read emails. @@ -301,26 +313,29 @@ def upload_email(dest, folder_name, raw_content, date_str, message_id, flags=Non """ try: # Ensure folder exists - if folder_name.upper() != imap_common.FOLDER_INBOX: - try: - dest.create(f'"{folder_name}"') - except Exception: - pass # Folder may already exist + imap_common.ensure_folder_exists(dest, folder_name) # Select folder dest.select(f'"{folder_name}"') # Check for duplicates if requested if check_duplicate and message_id and email_exists_in_folder(dest, message_id): - return False # Already exists + return UploadResult.ALREADY_EXISTS # Upload with original date and flags - resp, _ = dest.append(f'"{folder_name}"', flags, date_str, raw_content) - return resp == "OK" + success = imap_common.append_email( + dest, + folder_name, + raw_content, + date_str, + flags, + ensure_folder=False, + ) + return UploadResult.SUCCESS if success else UploadResult.FAILURE except Exception as e: safe_print(f"Error uploading to {folder_name}: {e}") - return False + return UploadResult.FAILURE def get_labels_from_manifest(manifest, message_id): @@ -351,7 +366,22 @@ def label_to_folder(label): return label -def process_restore_batch(eml_files, folder_name, dest_conf, manifest, apply_labels, apply_flags): +def process_restore_batch( + eml_files, + folder_name, + dest_conf, + manifest, + apply_labels, + apply_flags, + full_restore: bool = False, + existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = None, + existing_dest_msg_ids_lock: Optional[threading.Lock] = None, + progress_cache_data: Optional[dict] = None, + progress_cache_lock: Optional[threading.Lock] = None, + progress_cache_path: Optional[str] = None, + dest_host: Optional[str] = None, + dest_user: Optional[str] = None, +): """ Process a batch of .eml files for restoration. @@ -415,28 +445,78 @@ def process_restore_batch(eml_files, folder_name, dest_conf, manifest, apply_lab target_folder = folder_name remaining_labels = labels - # Upload to target folder - # In Gmail mode, check duplicates per-email since target folders vary - # In non-Gmail mode, duplicates were pre-filtered so no check needed - uploaded = upload_email( - dest, target_folder, raw_content, date_str, message_id, flags, check_duplicate=gmail_mode - ) + existing_dest_msg_ids: Optional[set[str]] = None + if existing_dest_msg_ids_by_folder is not None: + if existing_dest_msg_ids_lock is None: + existing_dest_msg_ids_lock = threading.Lock() - if not uploaded: + with existing_dest_msg_ids_lock: + existing_dest_msg_ids = existing_dest_msg_ids_by_folder.get(target_folder) + + # Lazy-load per-folder progress cache (no destination scan). + if existing_dest_msg_ids is None: + built: set[str] = set() + if progress_cache_data is not None and progress_cache_lock is not None and dest_host and dest_user: + built = restore_cache.get_cached_message_ids( + progress_cache_data, + progress_cache_lock, + dest_host, + dest_user, + target_folder, + ) + with existing_dest_msg_ids_lock: + existing_dest_msg_ids_by_folder.setdefault(target_folder, built) + existing_dest_msg_ids = existing_dest_msg_ids_by_folder[target_folder] + + # Incremental default: if we already know we processed it before, skip entirely. + if ( + not full_restore + and message_id + and existing_dest_msg_ids is not None + and message_id in existing_dest_msg_ids + ): + safe_print(f"[{target_folder}] SKIP (already present) | {size_str:<8} | {display_subject}") + continue + + # Upload to target folder. + # Keep server-side duplicate checks enabled to avoid creating duplicates for emails + # that exist on the destination but are not in our local progress cache. + upload_result = upload_email(dest, target_folder, raw_content, date_str, message_id, flags) + + # Only record progress when upload succeeds or email already exists. + # Failed uploads should not be marked as processed to allow retry on next run. + if upload_result in (UploadResult.SUCCESS, UploadResult.ALREADY_EXISTS): + restore_cache.record_progress( + message_id=message_id, + folder_name=target_folder, + existing_dest_msg_ids=existing_dest_msg_ids, + existing_dest_msg_ids_lock=existing_dest_msg_ids_lock, + progress_cache_path=progress_cache_path, + progress_cache_data=progress_cache_data, + progress_cache_lock=progress_cache_lock, + dest_host=dest_host, + dest_user=dest_user, + log_fn=safe_print, + ) + + if upload_result == UploadResult.ALREADY_EXISTS: safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {display_subject}") - # Even if skipped, sync flags on existing email if requested - if apply_flags and flags and message_id: + # Full restore preserves legacy behavior: sync flags on existing email if requested + if full_restore and apply_flags and flags and message_id: sync_flags_on_existing(dest, target_folder, message_id, flags, size) - else: + elif upload_result == UploadResult.SUCCESS: safe_print(f"[{target_folder}] UPLOADED | {size_str:<8} | {display_subject}") # Show applied flags in same style as labels if flags: for flag in flags.split(): safe_print(f" -> Applied flag: {flag}") + else: # FAILURE + safe_print(f"[{target_folder}] FAILED | {size_str:<8} | {display_subject}") - # Apply remaining Gmail labels (always, whether uploaded or skipped) - # This ensures labels are synced even for existing emails - if apply_labels and remaining_labels: + # Apply remaining Gmail labels: + # - Full restore: apply/sync labels even for existing emails + # - Incremental (default): apply labels only for newly uploaded emails + if apply_labels and remaining_labels and (upload_result == UploadResult.SUCCESS or full_restore): for label in remaining_labels: label_folder = label_to_folder(label) @@ -454,19 +534,63 @@ def process_restore_batch(eml_files, folder_name, dest_conf, manifest, apply_lab try: # Ensure label folder exists - if label_folder.upper() != imap_common.FOLDER_INBOX: - try: - dest.create(f'"{label_folder}"') - except Exception: - pass + imap_common.ensure_folder_exists(dest, label_folder) # Select and check for duplicate dest.select(f'"{label_folder}"') if not email_exists_in_folder(dest, message_id): - dest.append(f'"{label_folder}"', flags, date_str, raw_content) - safe_print(f" -> Applied label: {label}") - # If email exists in this label folder, sync flags - elif apply_flags and flags: + append_success = imap_common.append_email( + dest, + label_folder, + raw_content, + date_str, + flags, + ensure_folder=False, + ) + if append_success: + # Look up the correct set for label_folder + label_folder_msg_ids: Optional[set[str]] = None + if existing_dest_msg_ids_by_folder is not None: + with existing_dest_msg_ids_lock: + label_folder_msg_ids = existing_dest_msg_ids_by_folder.get(label_folder) + + # Lazy-load per-folder progress cache if not yet loaded + if label_folder_msg_ids is None: + built: set[str] = set() + if ( + progress_cache_data is not None + and progress_cache_lock is not None + and dest_host + and dest_user + ): + built = restore_cache.get_cached_message_ids( + progress_cache_data, + progress_cache_lock, + dest_host, + dest_user, + label_folder, + ) + with existing_dest_msg_ids_lock: + existing_dest_msg_ids_by_folder.setdefault(label_folder, built) + label_folder_msg_ids = existing_dest_msg_ids_by_folder[label_folder] + + restore_cache.record_progress( + message_id=message_id, + folder_name=label_folder, + existing_dest_msg_ids=label_folder_msg_ids, + existing_dest_msg_ids_lock=existing_dest_msg_ids_lock, + progress_cache_path=progress_cache_path, + progress_cache_data=progress_cache_data, + progress_cache_lock=progress_cache_lock, + dest_host=dest_host, + dest_user=dest_user, + log_fn=safe_print, + ) + safe_print(f" -> Applied label: {label}") + else: + safe_print(f" -> Failed to apply label {label} (will retry on next restore)") + # If email exists in this label folder, sync flags (full restore only) + elif full_restore and apply_flags and flags: sync_flags_on_existing(dest, label_folder, message_id, flags, size) except Exception as e: safe_print(f" -> Error applying label {label}: {e}") @@ -558,7 +682,17 @@ def delete_orphan_emails_from_dest(imap_conn, folder_name, local_msg_ids): return deleted_count -def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_labels, apply_flags, dest_delete=False): +def restore_folder( + folder_name, + local_folder_path, + dest_conf, + manifest, + apply_labels, + apply_flags, + dest_delete=False, + full_restore: bool = False, + cache_root: Optional[str] = None, +): """ Restore all emails from a local folder to the destination IMAP server. """ @@ -577,6 +711,30 @@ def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_la safe_print(f"Found {len(eml_files)} emails to restore.") + cache_root = cache_root or local_folder_path + cache_path = restore_cache.get_dest_index_cache_path(cache_root, dest_conf["host"], dest_conf["user"]) + cache_data: dict = restore_cache.load_dest_index_cache(cache_path) + cache_lock = threading.Lock() + + safe_print(f"Using progress cache: {cache_path}") + + # Incremental mode uses cached Message-IDs to skip already-processed emails. + existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = {folder_name: set()} + existing_dest_msg_ids_lock: Optional[threading.Lock] = threading.Lock() + try: + existing_dest_msg_ids_by_folder[folder_name] = restore_cache.get_cached_message_ids( + cache_data, + cache_lock, + dest_conf["host"], + dest_conf["user"], + folder_name, + ) + safe_print(f"Cache has {len(existing_dest_msg_ids_by_folder[folder_name])} Message-IDs for this folder.") + except Exception as e: + # Fall back to an empty cache for this folder if reading cached Message-IDs fails. + safe_print(f"Warning: Failed to load cached Message-IDs for folder '{folder_name}': {e}") + existing_dest_msg_ids_by_folder[folder_name] = set() + # If dest_delete enabled, get local Message-IDs for comparison local_msg_ids = None if dest_delete: @@ -641,6 +799,14 @@ def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_la manifest, apply_labels, apply_flags, + full_restore, + existing_dest_msg_ids_by_folder, + existing_dest_msg_ids_lock, + cache_data, + cache_lock, + cache_path, + dest_conf.get("host"), + dest_conf.get("user"), ) ) @@ -658,8 +824,11 @@ def restore_folder(folder_name, local_folder_path, dest_conf, manifest, apply_la delete_orphan_emails_from_dest(dest, folder_name, local_msg_ids) dest.logout() + # Force-flush progress cache at end. + restore_cache.maybe_save_dest_index_cache(cache_path, cache_data, cache_lock, force=True) -def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags): + +def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags, full_restore: bool = False): """ Special restoration mode for Gmail: Upload emails to their first label folder and then apply additional labels from the manifest. @@ -690,6 +859,19 @@ def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags): batches = [eml_files[i : i + BATCH_SIZE] for i in range(0, len(eml_files), BATCH_SIZE)] + cache_path = restore_cache.get_dest_index_cache_path(local_path, dest_conf["host"], dest_conf["user"]) + progress_cache_data: dict = restore_cache.load_dest_index_cache(cache_path) + progress_cache_lock = threading.Lock() + + safe_print(f"Using progress cache: {cache_path}") + safe_print( + "Cache will be populated as restore runs (no up-front destination indexing). " + "First run may still do per-message duplicate checks; subsequent runs will skip quickly." + ) + + existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = {} # lazily loaded per folder + existing_dest_msg_ids_lock: Optional[threading.Lock] = threading.Lock() + with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: futures = [] for batch in batches: @@ -705,6 +887,14 @@ def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags): manifest, True, # apply_labels apply_flags, # apply_flags + full_restore, + existing_dest_msg_ids_by_folder, + existing_dest_msg_ids_lock, + progress_cache_data, + progress_cache_lock, + cache_path, + dest_conf.get("host"), + dest_conf.get("user"), ) ) @@ -722,6 +912,9 @@ def restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags): elapsed = time.time() - start_time safe_print(f"\nRestore completed in {elapsed:.1f}s") + # Force-flush progress cache at end. + restore_cache.maybe_save_dest_index_cache(cache_path, progress_cache_data, progress_cache_lock, force=True) + def get_backup_folders(local_path): """ @@ -853,6 +1046,14 @@ def main(): help="Gmail restore mode: Upload to INBOX and apply labels + flags from manifest", ) + env_full_restore = os.getenv("FULL_RESTORE", "false").lower() == "true" + parser.add_argument( + "--full-restore", + action="store_true", + default=env_full_restore, + help="Force full restore (legacy): process all emails and sync labels/flags for already-present messages.", + ) + # Sync mode: delete from dest emails not in local backup env_dest_delete = os.getenv("DEST_DELETE", "false").lower() == "true" parser.add_argument( @@ -946,6 +1147,9 @@ def main(): print("Apply Flags : Yes (read/starred/answered/draft)") if args.dest_delete: print("Dest Delete : Yes (remove orphans from destination)") + print( + f"Restore Mode : {'Full (all emails)' if args.full_restore else 'Incremental (new emails only, use --full-restore to restore all)'}" + ) print("-----------------------------\n") try: @@ -958,14 +1162,24 @@ def main(): if args.gmail_mode: # Special Gmail mode - restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags) + restore_gmail_with_labels(local_path, dest_conf, manifest, apply_flags, full_restore=args.full_restore) elif args.folder: # Restore specific folder folder_path = os.path.join(local_path, args.folder.replace("/", os.sep)) if not os.path.exists(folder_path): print(f"Error: Folder not found: {folder_path}") sys.exit(1) - restore_folder(args.folder, folder_path, dest_conf, manifest, apply_labels, apply_flags, args.dest_delete) + restore_folder( + args.folder, + folder_path, + dest_conf, + manifest, + apply_labels, + apply_flags, + args.dest_delete, + full_restore=args.full_restore, + cache_root=local_path, + ) else: # Restore all folders folders = get_backup_folders(local_path) @@ -986,6 +1200,8 @@ def main(): apply_labels, apply_flags, args.dest_delete, + full_restore=args.full_restore, + cache_root=local_path, ) print("\nRestore completed successfully.") diff --git a/test/test_imap_common.py b/test/test_imap_common.py index 6d9e6a5..fca563e 100644 --- a/test/test_imap_common.py +++ b/test/test_imap_common.py @@ -350,6 +350,290 @@ class TestExtractMessageId: assert imap_common.extract_message_id(b"") is None +class TestEnsureFolderExists: + """Tests for ensure_folder_exists function.""" + + def test_creates_folder_successfully(self): + """Test folder is created when it doesn't exist.""" + mock_conn = Mock() + mock_conn.create.return_value = ("OK", []) + + # Should not raise any exception + imap_common.ensure_folder_exists(mock_conn, "TestFolder") + mock_conn.create.assert_called_once_with('"TestFolder"') + + def test_folder_already_exists(self): + """Test exception is suppressed when folder already exists.""" + mock_conn = Mock() + mock_conn.create.side_effect = Exception("Folder already exists") + + # Should not raise any exception + imap_common.ensure_folder_exists(mock_conn, "ExistingFolder") + mock_conn.create.assert_called_once_with('"ExistingFolder"') + + def test_inbox_folder_skipped(self): + """Test INBOX folder is not created.""" + mock_conn = Mock() + + imap_common.ensure_folder_exists(mock_conn, "INBOX") + mock_conn.create.assert_not_called() + + # Also test lowercase + imap_common.ensure_folder_exists(mock_conn, "inbox") + mock_conn.create.assert_not_called() + + def test_empty_folder_name(self): + """Test empty folder name is skipped.""" + mock_conn = Mock() + + imap_common.ensure_folder_exists(mock_conn, "") + mock_conn.create.assert_not_called() + + def test_none_folder_name(self): + """Test None folder name is skipped.""" + mock_conn = Mock() + + imap_common.ensure_folder_exists(mock_conn, None) + mock_conn.create.assert_not_called() + + def test_server_restriction_error_suppressed(self): + """Test server restriction errors are suppressed.""" + mock_conn = Mock() + mock_conn.create.side_effect = Exception("Permission denied") + + # Should not raise any exception + imap_common.ensure_folder_exists(mock_conn, "RestrictedFolder") + mock_conn.create.assert_called_once() + + +class TestAppendEmail: + """Tests for append_email function.""" + + def test_successful_append(self): + """Test successful email append returns True.""" + mock_conn = Mock() + mock_conn.append.return_value = ("OK", []) + + result = imap_common.append_email( + mock_conn, + "TestFolder", + b"email content", + "01-Jan-2024 12:00:00 +0000", + ensure_folder=False, + ) + + assert result is True + mock_conn.append.assert_called_once_with( + '"TestFolder"', + None, + "01-Jan-2024 12:00:00 +0000", + b"email content", + ) + + def test_append_with_flags_with_parentheses(self): + """Test flag normalization when flags already have parentheses.""" + mock_conn = Mock() + mock_conn.append.return_value = ("OK", []) + + result = imap_common.append_email( + mock_conn, + "TestFolder", + b"email content", + "01-Jan-2024 12:00:00 +0000", + flags="(\\Seen \\Flagged)", + ensure_folder=False, + ) + + assert result is True + # Flags should remain unchanged + mock_conn.append.assert_called_once_with( + '"TestFolder"', + "(\\Seen \\Flagged)", + "01-Jan-2024 12:00:00 +0000", + b"email content", + ) + + def test_append_with_flags_without_parentheses(self): + """Test flag normalization when flags lack parentheses.""" + mock_conn = Mock() + mock_conn.append.return_value = ("OK", []) + + result = imap_common.append_email( + mock_conn, + "TestFolder", + b"email content", + "01-Jan-2024 12:00:00 +0000", + flags="\\Seen", + ensure_folder=False, + ) + + assert result is True + # Flags should be wrapped in parentheses + mock_conn.append.assert_called_once_with( + '"TestFolder"', + "(\\Seen)", + "01-Jan-2024 12:00:00 +0000", + b"email content", + ) + + def test_append_with_none_flags(self): + """Test append with None flags.""" + mock_conn = Mock() + mock_conn.append.return_value = ("OK", []) + + result = imap_common.append_email( + mock_conn, + "TestFolder", + b"email content", + "01-Jan-2024 12:00:00 +0000", + flags=None, + ensure_folder=False, + ) + + assert result is True + mock_conn.append.assert_called_once_with( + '"TestFolder"', + None, + "01-Jan-2024 12:00:00 +0000", + b"email content", + ) + + def test_append_with_empty_flags(self): + """Test append with empty string flags.""" + mock_conn = Mock() + mock_conn.append.return_value = ("OK", []) + + result = imap_common.append_email( + mock_conn, + "TestFolder", + b"email content", + "01-Jan-2024 12:00:00 +0000", + flags="", + ensure_folder=False, + ) + + assert result is True + # Empty flags should result in None + mock_conn.append.assert_called_once_with( + '"TestFolder"', + None, + "01-Jan-2024 12:00:00 +0000", + b"email content", + ) + + def test_append_with_whitespace_only_flags(self): + """Test append with whitespace-only flags.""" + mock_conn = Mock() + mock_conn.append.return_value = ("OK", []) + + result = imap_common.append_email( + mock_conn, + "TestFolder", + b"email content", + "01-Jan-2024 12:00:00 +0000", + flags=" ", + ensure_folder=False, + ) + + assert result is True + # Whitespace-only flags should result in None + mock_conn.append.assert_called_once_with( + '"TestFolder"', + None, + "01-Jan-2024 12:00:00 +0000", + b"email content", + ) + + def test_append_with_ensure_folder_enabled(self): + """Test append with ensure_folder enabled.""" + mock_conn = Mock() + mock_conn.create.return_value = ("OK", []) + mock_conn.append.return_value = ("OK", []) + + result = imap_common.append_email( + mock_conn, + "TestFolder", + b"email content", + "01-Jan-2024 12:00:00 +0000", + ensure_folder=True, + ) + + assert result is True + mock_conn.create.assert_called_once_with('"TestFolder"') + mock_conn.append.assert_called_once() + + def test_append_with_ensure_folder_disabled(self): + """Test append with ensure_folder disabled.""" + mock_conn = Mock() + mock_conn.append.return_value = ("OK", []) + + result = imap_common.append_email( + mock_conn, + "TestFolder", + b"email content", + "01-Jan-2024 12:00:00 +0000", + ensure_folder=False, + ) + + assert result is True + mock_conn.create.assert_not_called() + mock_conn.append.assert_called_once() + + def test_append_failure_returns_false(self): + """Test append returns False on failure.""" + mock_conn = Mock() + mock_conn.append.return_value = ("NO", []) + + result = imap_common.append_email( + mock_conn, + "TestFolder", + b"email content", + "01-Jan-2024 12:00:00 +0000", + ensure_folder=False, + ) + + assert result is False + + def test_append_exception_propagates(self): + """Test append propagates exceptions so callers can handle/log.""" + mock_conn = Mock() + mock_conn.append.side_effect = Exception("Connection error") + + import pytest + + with pytest.raises(Exception, match="Connection error"): + imap_common.append_email( + mock_conn, + "TestFolder", + b"email content", + "01-Jan-2024 12:00:00 +0000", + ensure_folder=False, + ) + + def test_append_with_multiple_flags(self): + """Test append with multiple flags.""" + mock_conn = Mock() + mock_conn.append.return_value = ("OK", []) + + result = imap_common.append_email( + mock_conn, + "TestFolder", + b"email content", + "01-Jan-2024 12:00:00 +0000", + flags="\\Seen \\Flagged \\Answered", + ensure_folder=False, + ) + + assert result is True + # Multiple flags without parentheses should be wrapped + mock_conn.append.assert_called_once_with( + '"TestFolder"', + "(\\Seen \\Flagged \\Answered)", + "01-Jan-2024 12:00:00 +0000", + b"email content", + ) + + class TestGetMessageIdsInFolder: """Tests for get_message_ids_in_folder function.""" diff --git a/test/test_restore_imap_emails.py b/test/test_restore_imap_emails.py index b4d6087..c477316 100644 --- a/test/test_restore_imap_emails.py +++ b/test/test_restore_imap_emails.py @@ -19,6 +19,7 @@ import pytest sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src"))) import imap_common +import restore_cache import restore_imap_emails from conftest import make_single_mock_connection @@ -284,11 +285,11 @@ class TestUploadEmail: "", ) - assert result is True + assert result == restore_imap_emails.UploadResult.SUCCESS mock_conn.append.assert_called_once() def test_upload_email_duplicate(self, monkeypatch): - """Test upload returns False when message exists.""" + """Test upload returns ALREADY_EXISTS when message exists.""" mock_conn = MagicMock() mock_conn.select.return_value = ("OK", [b"1"]) @@ -304,7 +305,7 @@ class TestUploadEmail: check_duplicate=True, ) - assert result is False + assert result == restore_imap_emails.UploadResult.ALREADY_EXISTS mock_conn.append.assert_not_called() def test_upload_email_with_seen_flag(self, monkeypatch): @@ -326,10 +327,29 @@ class TestUploadEmail: flags="\\Seen", # Mark as read ) - assert result is True + assert result == restore_imap_emails.UploadResult.SUCCESS # Check that append was called with the \\Seen flag call_args = mock_conn.append.call_args - assert call_args[0][1] == "\\Seen" + assert call_args[0][1] == "(\\Seen)" + + def test_upload_email_failure(self, monkeypatch): + """Test upload returns FAILURE when an exception occurs.""" + mock_conn = MagicMock() + mock_conn.select.side_effect = Exception("Connection error") + + # Mock email_exists_in_folder to return False (not a duplicate) + monkeypatch.setattr(restore_imap_emails, "email_exists_in_folder", lambda *args: False) + + result = restore_imap_emails.upload_email( + mock_conn, + "INBOX", + b"raw email content", + '"15-Jan-2024 10:30:00 +0000"', + "", + ) + + assert result == restore_imap_emails.UploadResult.FAILURE + mock_conn.append.assert_not_called() class TestRestoreIntegration: @@ -371,6 +391,39 @@ Body content. # Run restore restore_imap_emails.main() + def test_restore_single_folder_full_restore_flag(self, single_mock_server, monkeypatch, tmp_path): + """Smoke test: --full-restore flag is accepted.""" + inbox = tmp_path / "INBOX" + inbox.mkdir() + + eml_content = b"""From: sender@test.com +To: recipient@test.com +Subject: Test Email +Message-ID: +Date: Mon, 15 Jan 2024 10:30:00 +0000 + +Body content. +""" + (inbox / "1_Test_Email.eml").write_bytes(eml_content) + + src_data = {"INBOX": []} + _server, port = single_mock_server(src_data) + + env = { + "DEST_IMAP_HOST": "localhost", + "DEST_IMAP_USERNAME": "user", + "DEST_IMAP_PASSWORD": "pass", + } + monkeypatch.setattr(os, "environ", env) + monkeypatch.setattr( + sys, + "argv", + ["restore_imap_emails.py", "--src-path", str(tmp_path), "--full-restore", "INBOX"], + ) + monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port)) + + restore_imap_emails.main() + def test_restore_with_labels_manifest(self, tmp_path): """Test that labels manifest is loaded correctly.""" # Create manifest @@ -416,6 +469,7 @@ class TestEmailExistsInFolder: def test_email_exists_exception(self, monkeypatch): """Test handling exception.""" mock_conn = MagicMock() + monkeypatch.setattr( "imap_common.message_exists_in_folder", lambda *args: (_ for _ in ()).throw(Exception("Error")), @@ -425,6 +479,46 @@ class TestEmailExistsInFolder: assert result is False +class TestRestoreProgressCache: + def test_progress_cache_add_get_persist(self, tmp_path): + import threading + + cache_path = restore_cache.get_dest_index_cache_path(str(tmp_path), "imap.example.com", "user@example.com") + cache_data = restore_cache.load_dest_index_cache(cache_path) + lock = threading.Lock() + + assert ( + restore_cache.get_cached_message_ids(cache_data, lock, "imap.example.com", "user@example.com", "INBOX") + == set() + ) + + assert ( + restore_cache.add_cached_message_id( + cache_data, lock, "imap.example.com", "user@example.com", "INBOX", "" + ) + is True + ) + assert ( + restore_cache.add_cached_message_id( + cache_data, lock, "imap.example.com", "user@example.com", "INBOX", "" + ) + is False + ) + assert ( + restore_cache.add_cached_message_id( + cache_data, lock, "imap.example.com", "user@example.com", "INBOX", "" + ) + is True + ) + + restore_cache.maybe_save_dest_index_cache(cache_path, cache_data, lock, force=True) + + cache_data2 = restore_cache.load_dest_index_cache(cache_path) + ids = restore_cache.get_cached_message_ids(cache_data2, lock, "imap.example.com", "user@example.com", "INBOX") + assert "" in ids + assert "" in ids + + class TestGetLabelsFromManifest: """Tests for get_labels_from_manifest function.""" @@ -504,7 +598,7 @@ class TestGmailModeDraftsFallbackRegression: def fake_upload_email(dest, folder_name, raw_content, date_str, message_id, flags=None, check_duplicate=True): captured["folder_name"] = folder_name - return True + return restore_imap_emails.UploadResult.SUCCESS def fake_parse_eml_file(_path): return ( @@ -739,3 +833,67 @@ class TestDestDeleteRestoreFunctionality: restore_imap_emails.main() assert len(server.folders["INBOX"]) == 0 + + +class TestAppendEmailReturnValueChecking: + """Tests that verify append_email return values are checked before recording progress.""" + + def test_record_progress_not_called_on_append_failure(self, tmp_path, monkeypatch): + """Test that record_progress is not called when append_email fails during label application.""" + from unittest.mock import Mock, patch + + # Create test email file + backup_root = tmp_path / "backup" + inbox = backup_root / "INBOX" + inbox.mkdir(parents=True) + + eml_content = b"""From: sender@test.com +To: recipient@test.com +Subject: Test Email +Message-ID: +Date: Mon, 15 Jan 2024 10:30:00 +0000 + +Body content. +""" + (inbox / "1_Test_Email.eml").write_bytes(eml_content) + + # Create labels manifest with multiple labels + manifest = {"": {"labels": ["INBOX", "Work"], "flags": []}} + manifest_path = backup_root / "labels_manifest.json" + manifest_path.write_text(json.dumps(manifest)) + + # Mock IMAP connection that fails on the second append (for label) + mock_conn = Mock() + mock_conn.select.return_value = ("OK", [b"0"]) + mock_conn.search.return_value = ("OK", [b""]) # No duplicates + + # First append succeeds (INBOX), second append fails (Work label) + mock_conn.append.side_effect = [ + ("OK", []), # INBOX upload succeeds + ("NO", []), # Work label append fails + ] + + # Track calls to record_progress + with patch("restore_cache.record_progress") as mock_record_progress: + # Set up environment + monkeypatch.setenv("DEST_IMAP_HOST", "localhost") + monkeypatch.setenv("DEST_IMAP_USERNAME", "user") + monkeypatch.setenv("DEST_IMAP_PASSWORD", "pass") + monkeypatch.setenv("MAX_WORKERS", "1") + + # Mock the connection + def mock_get_connection(*args, **kwargs): + return mock_conn + + monkeypatch.setattr("imap_common.get_imap_connection", mock_get_connection) + monkeypatch.setattr(sys, "argv", ["restore_imap_emails.py", "--src-path", str(backup_root), "INBOX"]) + + # Run restore + restore_imap_emails.main() + + # Verify record_progress was called only once (for successful INBOX upload) + # and NOT called for the failed Work label append + assert mock_record_progress.call_count == 1 + # Verify it was called for INBOX only + call_args = mock_record_progress.call_args + assert call_args[1]["folder_name"] == "INBOX" diff --git a/tools/mock_imap_server.py b/tools/mock_imap_server.py index 0fa7a4c..c5dc4c6 100644 --- a/tools/mock_imap_server.py +++ b/tools/mock_imap_server.py @@ -105,10 +105,21 @@ class MockIMAPHandler(socketserver.StreamRequestHandler): continue msgs = self.current_folders[self.selected_folder] + # Support UID range searches like: UID SEARCH UID 101:* + uid_min = None + try: + m = re.search(r"UID\s+(\d+)", sub_args, re.IGNORECASE) + if m: + uid_min = int(m.group(1)) + except Exception: + uid_min = None + # Handle UNDELETED # If args has UNDELETED, filter flags valid_uids = [] for m in msgs: + if uid_min is not None and m["uid"] < uid_min: + continue if "UNDELETED" in sub_args and "\\Deleted" in m["flags"]: continue valid_uids.append(str(m["uid"]))