Prevent redundant IMAP commands (#32)
* Pre-fetch folder Message-IDs for migrate and restore scripts Both scripts were issuing an IMAP SELECT + SEARCH for every single email to check for duplicates. Now we fetch all Message-IDs per folder once via a shared load_folder_msg_ids() in imap_common and use in-memory set lookups. This also applies to Gmail label folders, cutting IMAP round-trips dramatically. * Skip redundant folder CREATE calls during migration and restore load_folder_msg_ids() already creates each folder on first access, so the per-email CREATE in upload_email(), the per-label CREATE in append_email(), and the batch-level ensure_folder_exists + select in process_batch() were issuing unnecessary IMAP commands. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Remove exception suppression and propagate auth errors in IMAP ops * Remove duplicate checking from upload_email function * Simplify destination connection handling in process_batch * Rename source_msg_ids to cached_dest_msg_ids for clarity * Refactor destination message ID caching logic in migration code * Skip duplicate filtering during full restore for flag/label sync * Replace mock-based tests with e2e tests for load_folder_msg_ids Swap 6 mock-heavy unit tests for 4 e2e tests using the mock IMAP server, testing actual input/output behavior (fetch, cache, empty folder, folder creation) instead of internal implementation details. --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Nathan Moinvaziri
GitHub
parent
c6a568857f
commit
02c70a1e5c
@ -463,28 +463,6 @@ def parse_message_id_and_subject_from_bytes(raw_message):
|
||||
return None, "(No Subject)"
|
||||
|
||||
|
||||
def message_exists_in_folder(dest_conn, msg_id):
|
||||
"""
|
||||
Checks if a message with the given Message-ID exists in the CURRENTLY SELECTED folder of dest_conn.
|
||||
Only considers non-deleted messages (UNDELETED) so that messages pending expunge
|
||||
don't block uploads of fresh copies.
|
||||
Returns True if found, False otherwise.
|
||||
"""
|
||||
if not msg_id:
|
||||
return False
|
||||
|
||||
clean_id = msg_id.replace('"', '\\"')
|
||||
try:
|
||||
typ, data = dest_conn.search(None, f'UNDELETED (HEADER Message-ID "{clean_id}")')
|
||||
if typ != "OK":
|
||||
return False
|
||||
|
||||
dest_ids = data[0].split()
|
||||
return len(dest_ids) > 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def get_message_ids_in_folder(imap_conn):
|
||||
"""
|
||||
Fetches all Message-IDs from the currently selected folder.
|
||||
@ -706,6 +684,55 @@ def delete_orphan_emails(imap_conn, folder_name, source_msg_ids, dest_uid_to_msg
|
||||
return deleted_count
|
||||
|
||||
|
||||
def load_folder_msg_ids(
|
||||
imap_conn,
|
||||
folder_name,
|
||||
msg_ids_by_folder,
|
||||
msg_ids_lock,
|
||||
progress_cache_data=None,
|
||||
progress_cache_lock=None,
|
||||
dest_host=None,
|
||||
dest_user=None,
|
||||
):
|
||||
"""Load Message-IDs for a folder, fetching from server if not yet cached.
|
||||
|
||||
On first call for a folder, SELECTs the folder and fetches all Message-IDs
|
||||
from the server (merged with progress cache). Subsequent calls return the
|
||||
cached set without any IMAP operations.
|
||||
|
||||
Returns the set of Message-IDs, or None if tracking is disabled.
|
||||
"""
|
||||
if msg_ids_by_folder is None or msg_ids_lock is None:
|
||||
return None
|
||||
|
||||
with msg_ids_lock:
|
||||
existing = msg_ids_by_folder.get(folder_name)
|
||||
|
||||
if existing is not None:
|
||||
return existing
|
||||
|
||||
# Build from progress cache
|
||||
built: set[str] = set()
|
||||
if is_progress_cache_ready(progress_cache_data, progress_cache_lock) and dest_host and dest_user:
|
||||
built = restore_cache.get_cached_message_ids(
|
||||
progress_cache_data,
|
||||
progress_cache_lock,
|
||||
dest_host,
|
||||
dest_user,
|
||||
folder_name,
|
||||
)
|
||||
|
||||
# Fetch from server for a comprehensive set (one-time cost per folder)
|
||||
ensure_folder_exists(imap_conn, folder_name)
|
||||
imap_conn.select(f'"{folder_name}"')
|
||||
server_ids = set(get_message_ids_in_folder(imap_conn).values())
|
||||
built.update(server_ids)
|
||||
|
||||
with msg_ids_lock:
|
||||
msg_ids_by_folder.setdefault(folder_name, built)
|
||||
return msg_ids_by_folder[folder_name]
|
||||
|
||||
|
||||
def load_manifest(local_path, filename):
|
||||
"""Load a manifest JSON file from a backup directory.
|
||||
|
||||
|
||||
@ -194,7 +194,7 @@ def process_single_uid(
|
||||
label_index,
|
||||
check_duplicate,
|
||||
full_migrate: bool = False,
|
||||
existing_dest_msg_ids: Optional[set[str]] = None,
|
||||
existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = None,
|
||||
existing_dest_msg_ids_lock: Optional[threading.Lock] = None,
|
||||
progress_cache_path: Optional[str] = None,
|
||||
progress_cache_data: Optional[dict] = None,
|
||||
@ -255,28 +255,38 @@ def process_single_uid(
|
||||
else:
|
||||
target_folder = folder_name
|
||||
|
||||
imap_common.ensure_folder_exists(dest, target_folder)
|
||||
dest.select(f'"{target_folder}"')
|
||||
|
||||
is_cached = False
|
||||
is_duplicate = False
|
||||
|
||||
# check cache first if available
|
||||
# Fast path: check source folder cache to skip without any IMAP ops
|
||||
cached_dest_msg_ids = (
|
||||
existing_dest_msg_ids_by_folder.get(folder_name) if existing_dest_msg_ids_by_folder else None
|
||||
)
|
||||
cache_hit = False
|
||||
if not full_migrate and msg_id and existing_dest_msg_ids is not None:
|
||||
if not full_migrate and msg_id and cached_dest_msg_ids is not None:
|
||||
if existing_dest_msg_ids_lock is not None:
|
||||
with existing_dest_msg_ids_lock:
|
||||
cache_hit = msg_id in existing_dest_msg_ids
|
||||
cache_hit = msg_id in cached_dest_msg_ids
|
||||
else:
|
||||
cache_hit = msg_id in existing_dest_msg_ids
|
||||
cache_hit = msg_id in cached_dest_msg_ids
|
||||
|
||||
if cache_hit:
|
||||
is_cached = True
|
||||
is_duplicate = True
|
||||
safe_print(f"[{target_folder}] SKIP (cached) | {size_str:<8} | {subject[:40]}")
|
||||
elif bool(msg_id and check_duplicate and imap_common.message_exists_in_folder(dest, msg_id)):
|
||||
is_duplicate = True
|
||||
safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {subject[:40]}")
|
||||
elif msg_id and check_duplicate:
|
||||
# Check target folder using pre-fetched Message-ID set (one-time fetch per folder)
|
||||
target_msg_ids = imap_common.load_folder_msg_ids(
|
||||
dest,
|
||||
target_folder,
|
||||
existing_dest_msg_ids_by_folder,
|
||||
existing_dest_msg_ids_lock,
|
||||
progress_cache_data,
|
||||
progress_cache_lock,
|
||||
dest_host,
|
||||
dest_user,
|
||||
)
|
||||
if target_msg_ids is not None and msg_id in target_msg_ids:
|
||||
is_duplicate = True
|
||||
safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {subject[:40]}")
|
||||
|
||||
if is_duplicate:
|
||||
if preserve_flags and flags and msg_id:
|
||||
@ -301,10 +311,13 @@ def process_single_uid(
|
||||
|
||||
# Update cache if processed effectively (copied or duplicate)
|
||||
if msg_id:
|
||||
cached_dest_msg_ids = (
|
||||
existing_dest_msg_ids_by_folder.get(folder_name) if existing_dest_msg_ids_by_folder else None
|
||||
)
|
||||
restore_cache.record_progress(
|
||||
message_id=msg_id,
|
||||
folder_name=folder_name,
|
||||
existing_dest_msg_ids=existing_dest_msg_ids,
|
||||
existing_dest_msg_ids=cached_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,
|
||||
@ -327,9 +340,22 @@ def process_single_uid(
|
||||
):
|
||||
continue
|
||||
try:
|
||||
imap_common.ensure_folder_exists(dest, label_folder)
|
||||
dest.select(f'"{label_folder}"')
|
||||
if not imap_common.message_exists_in_folder(dest, msg_id):
|
||||
# Get or fetch Message-IDs for label folder (one-time server fetch per folder)
|
||||
label_folder_msg_ids = imap_common.load_folder_msg_ids(
|
||||
dest,
|
||||
label_folder,
|
||||
existing_dest_msg_ids_by_folder,
|
||||
existing_dest_msg_ids_lock,
|
||||
progress_cache_data,
|
||||
progress_cache_lock,
|
||||
dest_host,
|
||||
dest_user,
|
||||
)
|
||||
|
||||
# Check duplicate using pre-fetched set instead of per-message SEARCH
|
||||
label_already_exists = label_folder_msg_ids is not None and msg_id in label_folder_msg_ids
|
||||
|
||||
if not label_already_exists:
|
||||
valid_flags = f"({flags})" if (preserve_flags and flags) else None
|
||||
if imap_common.append_email(
|
||||
dest,
|
||||
@ -339,6 +365,10 @@ def process_single_uid(
|
||||
valid_flags,
|
||||
ensure_folder=False,
|
||||
):
|
||||
# Update in-memory set so subsequent emails see this one
|
||||
if label_folder_msg_ids is not None and existing_dest_msg_ids_lock is not None:
|
||||
with existing_dest_msg_ids_lock:
|
||||
label_folder_msg_ids.add(msg_id)
|
||||
safe_print(f" -> Applied label: {label}")
|
||||
if preserve_flags and flags:
|
||||
for flag in flags.split():
|
||||
@ -348,6 +378,8 @@ def process_single_uid(
|
||||
elif preserve_flags and flags:
|
||||
imap_common.sync_flags_on_existing(dest, label_folder, msg_id, flags, size)
|
||||
except Exception as e:
|
||||
if imap_oauth2.is_auth_error(e):
|
||||
raise
|
||||
safe_print(f" -> Error applying label {label}: {e}")
|
||||
|
||||
deleted = 0
|
||||
@ -383,7 +415,7 @@ def process_batch(
|
||||
label_index=None,
|
||||
check_duplicate=True,
|
||||
full_migrate=False,
|
||||
existing_dest_msg_ids: Optional[set[str]] = None,
|
||||
existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = None,
|
||||
existing_dest_msg_ids_lock: Optional[threading.Lock] = None,
|
||||
progress_cache_path: Optional[str] = None,
|
||||
progress_cache_data: Optional[dict] = None,
|
||||
@ -401,14 +433,6 @@ def process_batch(
|
||||
safe_print(f"Error selecting folder {folder_name} in worker: {e}")
|
||||
return False, 0
|
||||
|
||||
if not gmail_mode:
|
||||
try:
|
||||
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}")
|
||||
return False, 0
|
||||
|
||||
# Extract info for cache update if needed
|
||||
dest_host = dest_conf.get("host")
|
||||
dest_user = dest_conf.get("user")
|
||||
@ -436,18 +460,11 @@ def process_batch(
|
||||
safe_print(f"[{folder_name}] ERROR: Source connection/folder lost for UID {uid_str}")
|
||||
return False, 0
|
||||
|
||||
if not gmail_mode:
|
||||
dest, dest_ok = imap_session.ensure_folder_session(dest, dest_conf, folder_name, readonly=False)
|
||||
thread_local.dest = dest
|
||||
if not dest_ok:
|
||||
safe_print(f"[{folder_name}] ERROR: Dest connection/folder lost for UID {uid_str}")
|
||||
return False, 0
|
||||
else:
|
||||
dest = imap_session.ensure_connection(dest, dest_conf)
|
||||
thread_local.dest = dest
|
||||
if not dest:
|
||||
safe_print(f"[{folder_name}] ERROR: Dest connection lost for UID {uid_str}")
|
||||
return False, 0
|
||||
dest = imap_session.ensure_connection(dest, dest_conf)
|
||||
thread_local.dest = dest
|
||||
if not dest:
|
||||
safe_print(f"[{folder_name}] ERROR: Dest connection lost for UID {uid_str}")
|
||||
return False, 0
|
||||
|
||||
success, src, dest, deleted = process_single_uid(
|
||||
src,
|
||||
@ -461,7 +478,7 @@ def process_batch(
|
||||
label_index,
|
||||
check_duplicate,
|
||||
full_migrate,
|
||||
existing_dest_msg_ids=existing_dest_msg_ids,
|
||||
existing_dest_msg_ids_by_folder=existing_dest_msg_ids_by_folder,
|
||||
existing_dest_msg_ids_lock=existing_dest_msg_ids_lock,
|
||||
progress_cache_path=progress_cache_path,
|
||||
progress_cache_data=progress_cache_data,
|
||||
@ -512,8 +529,8 @@ def migrate_folder(
|
||||
safe_print(f"--- Preparing Folder: {folder_name} ---")
|
||||
|
||||
# Load cache if provided
|
||||
existing_dest_msg_ids = None
|
||||
existing_dest_msg_ids_lock = None
|
||||
existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = {}
|
||||
existing_dest_msg_ids_lock: Optional[threading.Lock] = threading.Lock()
|
||||
dest_host = dest_conf.get("host")
|
||||
dest_user = dest_conf.get("user")
|
||||
cache_file = progress_cache_file
|
||||
@ -532,19 +549,18 @@ def migrate_folder(
|
||||
|
||||
if imap_common.is_progress_cache_ready(progress_cache_data, progress_cache_lock):
|
||||
try:
|
||||
existing_dest_msg_ids = restore_cache.get_cached_message_ids(
|
||||
cache_ids = restore_cache.get_cached_message_ids(
|
||||
progress_cache_data,
|
||||
progress_cache_lock,
|
||||
dest_host,
|
||||
dest_user,
|
||||
folder_name,
|
||||
)
|
||||
existing_dest_msg_ids_lock = threading.Lock()
|
||||
safe_print(f"Cache has {len(existing_dest_msg_ids)} Message-IDs for this folder.")
|
||||
existing_dest_msg_ids_by_folder[folder_name] = cache_ids
|
||||
safe_print(f"Cache has {len(cache_ids)} Message-IDs for this folder.")
|
||||
except Exception as e:
|
||||
safe_print(f"Warning: Failed to read cache for folder '{folder_name}': {e}")
|
||||
existing_dest_msg_ids = set()
|
||||
existing_dest_msg_ids_lock = threading.Lock()
|
||||
existing_dest_msg_ids_by_folder[folder_name] = set()
|
||||
|
||||
# Maintain folder structure (skip in Gmail mode; worker will create/select target label folders)
|
||||
if not gmail_mode:
|
||||
@ -621,8 +637,12 @@ def migrate_folder(
|
||||
safe_print(f"Pre-fetching destination Message-IDs for {folder_name}...")
|
||||
dest_uid_to_msgid = imap_common.get_message_ids_in_folder(dest)
|
||||
dest_msg_ids = set(dest_uid_to_msgid.values())
|
||||
if existing_dest_msg_ids and not full_migrate:
|
||||
dest_msg_ids.update(existing_dest_msg_ids)
|
||||
# Update destination Message-IDs with what we processed
|
||||
if not full_migrate:
|
||||
existing_dest_msg_ids_by_folder.setdefault(folder_name, set()).update(dest_msg_ids)
|
||||
else:
|
||||
existing_dest_msg_ids_by_folder[folder_name] = dest_msg_ids
|
||||
dest_msg_ids = existing_dest_msg_ids_by_folder[folder_name]
|
||||
safe_print(f"Found {len(dest_msg_ids)} existing messages in destination (server + cache).")
|
||||
|
||||
# Pre-fetch source Message-IDs and filter out duplicates before processing
|
||||
@ -671,7 +691,7 @@ def migrate_folder(
|
||||
label_index,
|
||||
check_duplicate,
|
||||
full_migrate,
|
||||
existing_dest_msg_ids,
|
||||
existing_dest_msg_ids_by_folder,
|
||||
existing_dest_msg_ids_lock,
|
||||
cache_file,
|
||||
progress_cache_data,
|
||||
|
||||
@ -155,40 +155,18 @@ def get_eml_files(folder_path):
|
||||
return eml_files
|
||||
|
||||
|
||||
def email_exists_in_folder(imap_conn, message_id):
|
||||
"""
|
||||
Check if an email with the given Message-ID exists in the currently selected folder.
|
||||
"""
|
||||
if not message_id:
|
||||
return False
|
||||
|
||||
try:
|
||||
return imap_common.message_exists_in_folder(imap_conn, message_id)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def upload_email(dest, folder_name, raw_content, date_str, message_id, flags=None, check_duplicate=True):
|
||||
def upload_email(dest, folder_name, raw_content, date_str, flags=None):
|
||||
"""
|
||||
Upload a single email to the destination folder.
|
||||
Returns UploadResult enum: SUCCESS, ALREADY_EXISTS, or FAILURE.
|
||||
Returns UploadResult enum: SUCCESS or FAILURE.
|
||||
|
||||
Callers are expected to check for duplicates via load_folder_msg_ids()
|
||||
before calling this function, and to ensure the folder exists.
|
||||
|
||||
Args:
|
||||
flags: Optional string of IMAP flags like "\\Seen" for read emails.
|
||||
check_duplicate: Whether to check for duplicates before uploading.
|
||||
"""
|
||||
try:
|
||||
# Ensure folder exists
|
||||
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 UploadResult.ALREADY_EXISTS
|
||||
|
||||
# Upload with original date and flags
|
||||
success = imap_common.append_email(
|
||||
dest,
|
||||
folder_name,
|
||||
@ -285,47 +263,39 @@ def process_restore_batch(
|
||||
target_folder = folder_name
|
||||
remaining_labels = labels
|
||||
|
||||
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()
|
||||
existing_dest_msg_ids = imap_common.load_folder_msg_ids(
|
||||
dest,
|
||||
target_folder,
|
||||
existing_dest_msg_ids_by_folder,
|
||||
existing_dest_msg_ids_lock,
|
||||
progress_cache_data,
|
||||
progress_cache_lock,
|
||||
dest_host,
|
||||
dest_user,
|
||||
)
|
||||
|
||||
with existing_dest_msg_ids_lock:
|
||||
existing_dest_msg_ids = existing_dest_msg_ids_by_folder.get(target_folder)
|
||||
# Check if email already exists on destination using pre-fetched set
|
||||
email_already_on_dest = (
|
||||
message_id and existing_dest_msg_ids is not None and message_id in existing_dest_msg_ids
|
||||
)
|
||||
|
||||
# Lazy-load per-folder progress cache (no destination scan).
|
||||
if existing_dest_msg_ids is None:
|
||||
built: set[str] = set()
|
||||
if (
|
||||
imap_common.is_progress_cache_ready(progress_cache_data, progress_cache_lock)
|
||||
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
|
||||
):
|
||||
# Incremental default: skip entirely if already present
|
||||
if email_already_on_dest and not full_restore:
|
||||
safe_print(f"[{target_folder}] SKIP (already present) | {size_str:<8} | {display_subject}")
|
||||
continue # Skip to next file
|
||||
|
||||
# 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)
|
||||
if email_already_on_dest:
|
||||
# Full restore: treat as existing for flag sync
|
||||
upload_result = UploadResult.ALREADY_EXISTS
|
||||
else:
|
||||
# Upload — skip per-message SEARCH since we have a pre-fetched set
|
||||
upload_result = upload_email(
|
||||
dest,
|
||||
target_folder,
|
||||
raw_content,
|
||||
date_str,
|
||||
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.
|
||||
@ -377,12 +347,24 @@ def process_restore_batch(
|
||||
continue
|
||||
|
||||
try:
|
||||
# Ensure label folder exists
|
||||
imap_common.ensure_folder_exists(dest, label_folder)
|
||||
# Get or fetch Message-IDs for label folder (one-time server fetch per folder)
|
||||
label_folder_msg_ids = imap_common.load_folder_msg_ids(
|
||||
dest,
|
||||
label_folder,
|
||||
existing_dest_msg_ids_by_folder,
|
||||
existing_dest_msg_ids_lock,
|
||||
progress_cache_data,
|
||||
progress_cache_lock,
|
||||
dest_host,
|
||||
dest_user,
|
||||
)
|
||||
|
||||
# Select and check for duplicate
|
||||
dest.select(f'"{label_folder}"')
|
||||
if not email_exists_in_folder(dest, message_id):
|
||||
# Check duplicate using pre-fetched set instead of per-message SEARCH
|
||||
label_already_exists = (
|
||||
label_folder_msg_ids is not None and message_id and message_id in label_folder_msg_ids
|
||||
)
|
||||
|
||||
if not label_already_exists:
|
||||
append_success = imap_common.append_email(
|
||||
dest,
|
||||
label_folder,
|
||||
@ -392,34 +374,6 @@ def process_restore_batch(
|
||||
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 (
|
||||
imap_common.is_progress_cache_ready(
|
||||
progress_cache_data,
|
||||
progress_cache_lock,
|
||||
)
|
||||
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,
|
||||
@ -439,6 +393,8 @@ def process_restore_batch(
|
||||
elif full_restore and apply_flags and flags:
|
||||
imap_common.sync_flags_on_existing(dest, label_folder, message_id, flags, size)
|
||||
except Exception as e:
|
||||
if imap_oauth2.is_auth_error(e):
|
||||
raise
|
||||
safe_print(f" -> Error applying label {label}: {e}")
|
||||
|
||||
except Exception as e:
|
||||
@ -565,8 +521,16 @@ def restore_folder(
|
||||
|
||||
safe_print(f"{len(dest_msg_ids)} existing messages in destination.")
|
||||
|
||||
# Pre-filter files to skip duplicates
|
||||
if dest_msg_ids:
|
||||
# Update destination Message-IDs with what we processed
|
||||
if not full_restore:
|
||||
with existing_dest_msg_ids_lock:
|
||||
existing_dest_msg_ids_by_folder.setdefault(folder_name, set()).update(dest_msg_ids)
|
||||
else:
|
||||
existing_dest_msg_ids_by_folder[folder_name] = dest_msg_ids
|
||||
dest_msg_ids = existing_dest_msg_ids_by_folder[folder_name]
|
||||
|
||||
# Pre-filter files to skip duplicates (skip in full_restore: need to process all for flag/label sync)
|
||||
if dest_msg_ids and not full_restore:
|
||||
files_to_restore = pre_filter_eml_files(eml_files, dest_msg_ids)
|
||||
|
||||
if not files_to_restore:
|
||||
|
||||
@ -320,48 +320,6 @@ class TestDetectTrashFolder:
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestMessageExistsInFolder:
|
||||
"""Tests for message_exists_in_folder function."""
|
||||
|
||||
def test_no_message_id(self):
|
||||
"""Test returns False when message_id is None."""
|
||||
mock_conn = Mock()
|
||||
result = imap_common.message_exists_in_folder(mock_conn, None)
|
||||
assert result is False
|
||||
|
||||
def test_search_fails(self):
|
||||
"""Test returns False when search fails."""
|
||||
mock_conn = Mock()
|
||||
mock_conn.search.return_value = ("NO", [])
|
||||
|
||||
result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
|
||||
assert result is False
|
||||
|
||||
def test_no_matches(self):
|
||||
"""Test returns False when no matches found."""
|
||||
mock_conn = Mock()
|
||||
mock_conn.search.return_value = ("OK", [b""])
|
||||
|
||||
result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
|
||||
assert result is False
|
||||
|
||||
def test_match_found(self):
|
||||
"""Test returns True when message with same ID found."""
|
||||
mock_conn = Mock()
|
||||
mock_conn.search.return_value = ("OK", [b"1"])
|
||||
|
||||
result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
|
||||
assert result is True
|
||||
|
||||
def test_search_exception(self):
|
||||
"""Test returns False when search raises an exception."""
|
||||
mock_conn = Mock()
|
||||
mock_conn.search.side_effect = Exception("Connection error")
|
||||
|
||||
result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
|
||||
assert result is False
|
||||
|
||||
|
||||
class TestExtractMessageId:
|
||||
"""Tests for extract_message_id function."""
|
||||
|
||||
@ -807,3 +765,77 @@ class TestGetMessageIdsInFolder:
|
||||
|
||||
result = imap_common.get_message_ids_in_folder(mock_conn)
|
||||
assert set(result.values()) == {"<valid@example.com>"}
|
||||
|
||||
|
||||
class TestLoadFolderMsgIds:
|
||||
"""Tests for load_folder_msg_ids() using the mock IMAP server."""
|
||||
|
||||
def test_returns_none_when_disabled(self):
|
||||
"""Returns None if dict or lock is None (tracking disabled)."""
|
||||
mock_conn = Mock()
|
||||
import threading
|
||||
|
||||
lock = threading.Lock()
|
||||
assert imap_common.load_folder_msg_ids(mock_conn, "INBOX", None, lock) is None
|
||||
assert imap_common.load_folder_msg_ids(mock_conn, "INBOX", {}, None) is None
|
||||
assert imap_common.load_folder_msg_ids(mock_conn, "INBOX", None, None) is None
|
||||
|
||||
def test_fetches_message_ids_from_server(self, single_mock_server):
|
||||
"""Fetches Message-IDs from server and caches them."""
|
||||
import imaplib
|
||||
import threading
|
||||
|
||||
server_data = {
|
||||
"TestFolder": [
|
||||
b"Subject: A\r\nMessage-ID: <a@test.com>\r\n\r\nBody A",
|
||||
b"Subject: B\r\nMessage-ID: <b@test.com>\r\n\r\nBody B",
|
||||
]
|
||||
}
|
||||
_server, port = single_mock_server(server_data)
|
||||
conn = imaplib.IMAP4("localhost", port)
|
||||
conn.login("user", "pass")
|
||||
|
||||
lock = threading.Lock()
|
||||
by_folder: dict[str, set[str]] = {}
|
||||
|
||||
result = imap_common.load_folder_msg_ids(conn, "TestFolder", by_folder, lock)
|
||||
assert result == {"<a@test.com>", "<b@test.com>"}
|
||||
assert "TestFolder" in by_folder
|
||||
|
||||
# Second call returns the same cached set object
|
||||
result2 = imap_common.load_folder_msg_ids(conn, "TestFolder", by_folder, lock)
|
||||
assert result2 is result
|
||||
conn.logout()
|
||||
|
||||
def test_empty_folder(self, single_mock_server):
|
||||
"""Returns empty set for a folder with no messages."""
|
||||
import imaplib
|
||||
import threading
|
||||
|
||||
_server, port = single_mock_server({"INBOX": []})
|
||||
conn = imaplib.IMAP4("localhost", port)
|
||||
conn.login("user", "pass")
|
||||
|
||||
lock = threading.Lock()
|
||||
by_folder: dict[str, set[str]] = {}
|
||||
|
||||
result = imap_common.load_folder_msg_ids(conn, "INBOX", by_folder, lock)
|
||||
assert result == set()
|
||||
conn.logout()
|
||||
|
||||
def test_creates_folder_if_missing(self, single_mock_server):
|
||||
"""Creates folder on first access if it doesn't exist."""
|
||||
import imaplib
|
||||
import threading
|
||||
|
||||
server, port = single_mock_server({"INBOX": []})
|
||||
conn = imaplib.IMAP4("localhost", port)
|
||||
conn.login("user", "pass")
|
||||
|
||||
lock = threading.Lock()
|
||||
by_folder: dict[str, set[str]] = {}
|
||||
|
||||
result = imap_common.load_folder_msg_ids(conn, "NewFolder", by_folder, lock)
|
||||
assert result == set()
|
||||
assert "NewFolder" in server.folders
|
||||
conn.logout()
|
||||
|
||||
@ -285,7 +285,7 @@ class TestCacheHitWithoutLock:
|
||||
|
||||
src.select('"INBOX"', readonly=False)
|
||||
|
||||
existing_dest_msg_ids = {msg_id}
|
||||
existing_dest_msg_ids_by_folder = {"INBOX": {msg_id}}
|
||||
success, _src, _dest, deleted = migrate_imap_emails.process_single_uid(
|
||||
src,
|
||||
dest,
|
||||
@ -298,7 +298,7 @@ class TestCacheHitWithoutLock:
|
||||
None,
|
||||
True,
|
||||
False,
|
||||
existing_dest_msg_ids=existing_dest_msg_ids,
|
||||
existing_dest_msg_ids_by_folder=existing_dest_msg_ids_by_folder,
|
||||
existing_dest_msg_ids_lock=None,
|
||||
progress_cache_path=None,
|
||||
progress_cache_data=None,
|
||||
|
||||
@ -509,45 +509,12 @@ class TestRestoreE2EHelpers:
|
||||
"INBOX",
|
||||
b"Subject: Upload\r\nMessage-ID: <up@test>\r\n\r\nBody",
|
||||
'"15-Jan-2024 10:30:00 +0000"',
|
||||
"<up@test>",
|
||||
)
|
||||
|
||||
assert result == restore_imap_emails.UploadResult.SUCCESS
|
||||
assert len(server.folders["INBOX"]) == 1
|
||||
conn.logout()
|
||||
|
||||
def test_upload_email_duplicate(self, single_mock_server):
|
||||
dest_data = {"INBOX": [b"Subject: Dup\r\nMessage-ID: <dup@test>\r\n\r\nBody"]}
|
||||
server, port = single_mock_server(dest_data)
|
||||
|
||||
conn = imaplib.IMAP4("localhost", port)
|
||||
conn.login("user", "pass")
|
||||
|
||||
result = restore_imap_emails.upload_email(
|
||||
conn,
|
||||
"INBOX",
|
||||
b"Subject: Dup\r\nMessage-ID: <dup@test>\r\n\r\nBody",
|
||||
'"15-Jan-2024 10:30:00 +0000"',
|
||||
"<dup@test>",
|
||||
check_duplicate=True,
|
||||
)
|
||||
|
||||
assert result == restore_imap_emails.UploadResult.ALREADY_EXISTS
|
||||
assert len(server.folders["INBOX"]) == 1
|
||||
conn.logout()
|
||||
|
||||
def test_email_exists_in_folder(self, single_mock_server):
|
||||
dest_data = {"INBOX": [b"Subject: Exists\r\nMessage-ID: <exists@test>\r\n\r\nBody"]}
|
||||
_server, port = single_mock_server(dest_data)
|
||||
|
||||
conn = imaplib.IMAP4("localhost", port)
|
||||
conn.login("user", "pass")
|
||||
conn.select('"INBOX"')
|
||||
|
||||
assert restore_imap_emails.email_exists_in_folder(conn, "<exists@test>") is True
|
||||
assert restore_imap_emails.email_exists_in_folder(conn, "<missing@test>") is False
|
||||
conn.logout()
|
||||
|
||||
def test_sync_flags_on_existing(self, single_mock_server):
|
||||
dest_data = {
|
||||
"INBOX": [{"uid": 1, "flags": set(), "content": b"Subject: Flag\r\nMessage-ID: <flag@test>\r\n\r\nBody"}]
|
||||
|
||||
Loading…
Reference in New Issue
Block a user
