Implement incremental restore with full restore option via restore cache. (#12)
* Implement incremental restore with full restore option and progress caching * Refactor restore cache documentation and simplify cache path retrieval * Add progress recording functionality for incremental restores * Refactor email folder handling by introducing helper functions for folder creation and email appending * Normalize flags format in append_email function for consistency * Remove unnecessary blank line before get_labels_from_manifest function * Enhance cache persistence with logging and error handling in save_dest_index_cache * Add future annotations import for type hinting support * Update src/restore_imap_emails.py When applying Gmail labels, emails uploaded to label folders are not recorded in the progress cache. This means on subsequent incremental restore runs, the code will need to check the server again for these label folders even though the emails were just uploaded. Consider calling restore_cache.record_progress after successfully appending to a label folder (after line 514) to cache the Message-ID for that label folder as well. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update src/restore_cache.py The condition for saving the cache will trigger a save based on the time threshold even when there are no pending updates (pending = 0). This can result in unnecessary disk writes. Consider adding a check to skip saving if there are no pending updates: force or (pending > 0 and (pending >= _MIN_PENDING_UPDATES_BEFORE_SAVE or (now - last_saved) >= _MIN_SECONDS_BETWEEN_SAVES)). This way, the cache is only saved when there are actually updates to persist, unless force=True. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update src/restore_imap_emails.py 'except' clause does nothing but pass and there is no explanatory comment. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Refactor condition in maybe_save_dest_index_cache for clarity * [WIP] WIP Address feedback on incremental and full restore implementation (#13) * Initial plan * Optimize add_cached_message_id for O(1) lookups using set Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> * Cache set at entry level to avoid O(n) set construction on each add Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> * Move copy import to top of file per Python conventions Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> * Improve code style: condense comments and docstring Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> * Add batch Message-ID fetching for faster duplicate detection (#9) * Add batch Message-ID fetching for faster duplicate detection * Remove duplicate get_message_ids_in_folder in migrate_imap_emails * Optimize orphan email deletion by reusing existing UID mappings * Remove unused get_msg_details function and related test code * Update src/restore_imap_emails.py The call to restore_cache.record_progress is missing required keyword arguments. The function signature requires parameters like existing_dest_msg_ids, existing_dest_msg_ids_lock, progress_cache_path, progress_cache_data, progress_cache_lock, dest_host, and dest_user. This call should match the pattern used at line 456-467 where all required parameters are properly passed as keyword arguments. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update src/imap_common.py The docstring states that the flags parameter is "Passed through to imaplib.IMAP4.append unchanged", but this is inaccurate. The function actually normalizes the flags by wrapping them in parentheses if they aren't already (lines 99-106). The docstring should be updated to reflect that the function normalizes flags to ensure they are in the correct IMAP format with parentheses. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Add unit tests for ensure_folder_exists and append_email helper functions (#14) * Initial plan * Add comprehensive unit tests for ensure_folder_exists and append_email Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> * Refactor folder existence checks to use imap_common.ensure_folder_exists directly * Initial plan (#15) Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> * Refactor process_restore_batch to improve argument clarity for restore_cache.record_progress * Update src/imap_common.py The append_email function catches all exceptions and returns False, but this makes it impossible for callers to distinguish between legitimate duplicate detection (email already exists) and actual errors (network failure, server error, etc.). Consider either raising exceptions for actual errors while returning False only for duplicates, or returning a more informative result (e.g., a tuple with status and reason, or an enum). This would allow restore_imap_emails.py to only record progress when emails are confirmed to be on the server. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Refactor test_append_exception to propagate exceptions for better error handling * Update src/imap_common.py The append_email function does not handle exceptions from imap_conn.append(). If the append call fails with an exception (e.g., connection error, permission error), the exception will propagate to the caller instead of returning False. This is inconsistent with the documented behavior that callers expect a boolean return value to indicate success or failure. The function should wrap the append call in a try-except block and return False on exception. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * [WIP] Address feedback on incremental restore implementation (#16) * Initial plan * Fix: Check append_email return value before recording progress for label application Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> * Add test to verify append_email return value is checked before recording progress Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> * Improve error message to clarify retry behavior for failed label applications Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> * [WIP] Address feedback on incremental restore implementation (#17) * Initial plan * Fix progress recording for failed uploads with three-state result enum Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> * Clean up whitespace and simplify labels manifest definition in tests * [WIP] Fix incremental restore return type consistency (#19) * Initial plan * Fix return type annotation for save_dest_index_cache to bool Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> * [WIP] Work in progress on incremental restore feedback (#20) * Initial plan * Fix label folder progress tracking to use correct message ID set Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> * Refactor conditional statement for clarity in process_restore_batch function * Remove exception handling in append_email function for cleaner flow --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: JCallico <335565+JCallico@users.noreply.github.com> Co-authored-by: Nathan Moinvaziri <nathan@nathanm.com>
This commit is contained in:
parent
68c741913b
commit
7e853bdf27
11
README.md
11
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" \
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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():
|
||||
|
||||
254
src/restore_cache.py
Normal file
254
src/restore_cache.py
Normal file
@ -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,
|
||||
)
|
||||
@ -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.")
|
||||
|
||||
@ -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."""
|
||||
|
||||
|
||||
@ -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:
|
||||
"<test@test.com>",
|
||||
)
|
||||
|
||||
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"',
|
||||
"<test@test.com>",
|
||||
)
|
||||
|
||||
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: <test123@test.com>
|
||||
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", "<a@test>"
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
restore_cache.add_cached_message_id(
|
||||
cache_data, lock, "imap.example.com", "user@example.com", "INBOX", "<a@test>"
|
||||
)
|
||||
is False
|
||||
)
|
||||
assert (
|
||||
restore_cache.add_cached_message_id(
|
||||
cache_data, lock, "imap.example.com", "user@example.com", "INBOX", "<b@test>"
|
||||
)
|
||||
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 "<a@test>" in ids
|
||||
assert "<b@test>" 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: <test123@test.com>
|
||||
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 = {"<test123@test.com>": {"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"
|
||||
|
||||
@ -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"]))
|
||||
|
||||
Loading…
Reference in New Issue
Block a user