gian/imap-migration-tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Ensure migration resumption functionality with cache tracking (#28)

Browse Source 
* Implement migration resumption functionality with cache tracking and tests

* Remove unused pytest import from migration resumption tests
This commit is contained in:
Javier Callico 2026-02-08 14:16:01 -05:00 committed by GitHub
parent f0961c14ea
commit 13db9b2a30
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 305 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

95
src/migrate_imap_emails.py
View File

@ -388,13 +388,13 @@ def process_batch(
dest = imap_session.get_thread_connection(thread_local, "dest", dest_conf)
if not src or not dest:
safe_print("Error: Could not establish connections in worker thread.")
return
return False, 0
try:
src.select(f'"{folder_name}"', readonly=False)
except Exception as e:
safe_print(f"Error selecting folder {folder_name} in worker: {e}")
return
return False, 0
if not gmail_mode:
try:
@ -402,16 +402,26 @@ def process_batch(
dest.select(f'"{folder_name}"')
except Exception as e:
safe_print(f"Error selecting folder {folder_name} in worker: {e}")
return
return False, 0
# Extract info for cache update if needed
dest_host = dest_conf.get("host")
dest_user = dest_conf.get("user")
deleted_count = 0
max_uid_processed = 0
for uid in uids:
uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
# Track max UID seen in this batch
try:
uid_int = int(uid_str)
if uid_int > max_uid_processed:
max_uid_processed = uid_int
except ValueError:
pass
max_retries = 2
for attempt in range(max_retries):
@ -419,20 +429,20 @@ def process_batch(
thread_local.src = src
if not src_ok:
safe_print(f"[{folder_name}] ERROR: Source connection/folder lost for UID {uid_str}")
return
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
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
return False, 0
success, src, dest, deleted = process_single_uid(
src,
@ -473,6 +483,8 @@ def process_batch(
except Exception as e:
safe_print(f"[{folder_name}] ERROR Expunge: {e}")
return True, max_uid_processed
def migrate_folder(
src,
@ -546,6 +558,28 @@ def migrate_folder(
safe_print(f"Skipping {folder_name}: {e}")
return
# Get Current UIDVALIDITY (for resume check)
current_validity = None
try:
typ, val_data = src.response("UIDVALIDITY")
if val_data:
current_validity = int(val_data[0])
except Exception:
pass
start_uid = 0
if not full_migrate and current_validity and progress_cache_data:
try:
src_data = restore_cache.get_source_data(progress_cache_data, folder_name)
if src_data and src_data.get("uid_validity") == current_validity:
start_uid = src_data.get("last_uid", 0) or 0
if start_uid > 0:
safe_print(
f"Resuming {folder_name} from Source UID > {start_uid} (UIDVALIDITY: {current_validity})"
)
except Exception:
pass
# Get UIDs
# Search for UNDELETED to avoid processing messages marked for deletion but not yet expunged
resp, data = src.uid("search", None, "UNDELETED")
@ -555,10 +589,23 @@ def migrate_folder(
uids = data[0].split()
total = len(uids)
if total == 0:
safe_print(f"Folder {folder_name} is empty.")
# Even if empty, we might need to delete from dest
if dest_delete:
# Filter UIDs if resuming
if start_uid > 0:
filtered = []
for u in uids:
try:
if int(u) > start_uid:
filtered.append(u)
except ValueError:
filtered.append(u)
uids = filtered
safe_print(f"Filtered to {len(uids)} new UIDs (was {total}).")
if not uids:
safe_print(f"Folder {folder_name} is up to date (no new UIDs).")
# Do not perform dest_delete if we filtered or found nothing via resume,
# as we don't have the full source picture.
if dest_delete and start_uid == 0:
safe_print("Checking destination for orphan emails to delete...")
imap_common.delete_orphan_emails(dest, folder_name, set())
return
@ -587,9 +634,13 @@ def migrate_folder(
if not uids_to_process:
safe_print(f"No new messages to migrate in {folder_name}.")
if dest_delete and src_msg_ids is not None:
# Only support orphan delete if we have full user list (no start_uid resume)
if dest_delete and src_msg_ids is not None and start_uid == 0:
safe_print("Syncing destination: removing emails not in source...")
imap_common.delete_orphan_emails(dest, folder_name, src_msg_ids, dest_uid_to_msgid)
# Still deleting skipped duplicates from source is valid?
# Only if we aren't resuming? If skipping duplicates, we found duplicates.
# But we return here.
return
# Create batches from filtered UIDs
@ -623,12 +674,28 @@ def migrate_folder(
)
)
# Wait for all batches to complete
for future in concurrent.futures.as_completed(futures):
# Wait for all batches to complete and update watermark
should_update_watermark = True
for future in futures:
try:
future.result()
success, batch_max_uid = future.result()
if success:
if should_update_watermark and current_validity and progress_cache_data and batch_max_uid > 0:
restore_cache.record_source_progress(
folder_name=folder_name,
uid_validity=current_validity,
last_uid=batch_max_uid,
progress_cache_path=cache_file,
progress_cache_data=progress_cache_data,
progress_cache_lock=progress_cache_lock,
log_fn=safe_print,
)
else:
should_update_watermark = False
except Exception as e:
safe_print(f"Batch Error: {e}")
should_update_watermark = False
except KeyboardInterrupt:
safe_print("\n\n!!! Migration interrupted by user. Shutting down threads... !!!\n")
executor.shutdown(wait=False, cancel_futures=True)

59
src/restore_cache.py
View File

@ -252,3 +252,62 @@ def record_progress(
progress_cache_lock,
log_fn=log_fn,
)
def get_source_data(
cache_data: dict,
folder_name: str,
) -> dict:
"""Retrieve source tracking data for a folder."""
folders = cache_data.get("folders", {})
if not isinstance(folders, dict):
return {}
folder_entry = folders.get(folder_name, {})
if not isinstance(folder_entry, dict):
return {}
return folder_entry.get("source_state", {})
def set_source_data(
cache_data: dict,
folder_name: str,
uid_validity: int,
last_uid: int,
) -> None:
"""Update source tracking data for a folder."""
if "folders" not in cache_data or not isinstance(cache_data["folders"], dict):
cache_data["folders"] = {}
if folder_name not in cache_data["folders"] or not isinstance(cache_data["folders"][folder_name], dict):
cache_data["folders"][folder_name] = {}
# Ensure folder entry is a dict
if not isinstance(cache_data["folders"][folder_name], dict):
cache_data["folders"][folder_name] = {}
cache_data["folders"][folder_name]["source_state"] = {"uid_validity": uid_validity, "last_uid": last_uid}
def record_source_progress(
*,
folder_name: str,
uid_validity: int,
last_uid: int,
progress_cache_path: str | None,
progress_cache_data: dict | None,
progress_cache_lock: threading.Lock | None,
log_fn: Callable[[str], None] | None = None,
) -> None:
"""Record source progress (last processed UID) to the cache."""
if progress_cache_path and progress_cache_data is not None and progress_cache_lock is not None:
with progress_cache_lock:
set_source_data(progress_cache_data, folder_name, uid_validity, last_uid)
# Force save to ensure watermark is persistent
maybe_save_dest_index_cache(
progress_cache_path,
progress_cache_data,
progress_cache_lock,
force=True,
log_fn=log_fn,
)

165
test/test_migrate_resume.py Normal file
View File

@ -0,0 +1,165 @@
"""Tests for migration resumption functionality using cache."""
import json
import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
import migrate_imap_emails
import restore_cache
from conftest import temp_argv, temp_env
def _run_migrate(cache_dir, src_port, dest_port):
env = {
"SRC_IMAP_HOST": f"imap://localhost:{src_port}",
"SRC_IMAP_USERNAME": "src",
"SRC_IMAP_PASSWORD": "p",
"DEST_IMAP_HOST": f"imap://localhost:{dest_port}",
"DEST_IMAP_USERNAME": "dest",
"DEST_IMAP_PASSWORD": "p",
"MAX_WORKERS": "1",
}
argv = [
"migrate_imap_emails.py",
"--src-host",
f"imap://localhost:{src_port}",
"--src-user",
"src",
"--src-pass",
"p",
"--dest-host",
f"imap://localhost:{dest_port}",
"--dest-user",
"dest",
"--dest-pass",
"p",
"--migrate-cache",
str(cache_dir),
"--workers",
"1",
]
with temp_env(env), temp_argv(argv):
migrate_imap_emails.main()
class TestMigrationResumption:
"""Tests that migration resumes from the last known UID using source state tracking."""
def test_migrate_resumes_using_last_uid(self, mock_server_factory, tmp_path):
# Setup source with 3 messages
msg1 = b"Subject: Msg1\r\nMessage-ID: <msg1@test>\r\n\r\nBody1"
msg2 = b"Subject: Msg2\r\nMessage-ID: <msg2@test>\r\n\r\nBody2"
msg3 = b"Subject: Msg3\r\nMessage-ID: <msg3@test>\r\n\r\nBody3"
src_data = {"INBOX": [msg1, msg2, msg3]} # UIDs will be 1, 2, 3
dest_data = {"INBOX": []}
_src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
cache_dir = tmp_path / "cache"
cache_dir.mkdir()
# Pre-seed cache to simulate a prior run that finished up to UID 2
# Mock server uses UIDVALIDITY=1 by default
dest_host = f"imap://localhost:{p2}"
dest_user = "dest"
cache_path = restore_cache.get_dest_index_cache_path(str(cache_dir), dest_host, dest_user)
cache_data = {
"version": 1,
"dest": {"host": dest_host, "user": dest_user},
"folders": {
"INBOX": {
# Simulate that we've already seen msg1 and msg2
"message_ids": ["<msg1@test>", "<msg2@test>"],
"source_state": {"uid_validity": 1, "last_uid": 2},
}
},
"_meta": {},
}
# Update cache data for the robust test
# We delete msg1 from cache ID list but keep last_uid=2.
# If it resumes from 2, it won't see msg1 (UID 1).
cache_data["folders"]["INBOX"]["message_ids"] = ["<msg2@test>"] # Removed msg1
with open(cache_path, "w") as f:
json.dump(cache_data, f)
# Run without patch - relies on real localhost tcp connection
_run_migrate(cache_dir, p1, p2)
msgs_in_dest = dest_server.folders["INBOX"]
assert len(msgs_in_dest) == 1
assert b"Msg3" in msgs_in_dest[0]["content"]
# Check that cache was updated with new last_uid
with open(cache_path) as f:
new_cache = json.load(f)
src_state = new_cache["folders"]["INBOX"]["source_state"]
assert src_state["last_uid"] == 3
# Should also have added msg3 to message_ids
assert "<msg3@test>" in new_cache["folders"]["INBOX"]["message_ids"]
def test_migrate_ignores_resume_on_uidvalidity_mismatch(self, mock_server_factory, tmp_path):
# Setup source with 2 messages
msg1 = b"Subject: Msg1\r\nMessage-ID: <msg1@test>\r\n\r\nBody1"
msg2 = b"Subject: Msg2\r\nMessage-ID: <msg2@test>\r\n\r\nBody2"
src_data = {"INBOX": [msg1, msg2]} # UIDs 1, 2
dest_data = {"INBOX": []}
_src_server, dest_server, p1, p2 = mock_server_factory(src_data, dest_data)
cache_dir = tmp_path / "cache"
cache_dir.mkdir()
dest_host = f"imap://localhost:{p2}"
dest_user = "dest"
cache_path = restore_cache.get_dest_index_cache_path(str(cache_dir), dest_host, dest_user)
# Pre-seed cache with mismatching UIDVALIDITY (999 vs 1)
# But claim we processed everything (last_uid=2)
# Also remove msg1 from cache to detect if it gets re-processed
cache_data = {
"version": 1,
"dest": {"host": dest_host, "user": dest_user},
"folders": {
"INBOX": {
"message_ids": ["<msg2@test>"],
"source_state": {
"uid_validity": 999, # Mismatch
"last_uid": 2,
},
}
},
"_meta": {},
}
with open(cache_path, "w") as f:
json.dump(cache_data, f)
# Run without patch - relies on real localhost tcp connection
_run_migrate(cache_dir, p1, p2)
# Since UIDVALIDITY mismatched, it should ignore last_uid=2 and rescan from start (UID 1 and 2).
# UID 1 is NOT in cache -> should be copied.
# UID 2 IS in cache -> should be skipped.
msgs_in_dest = dest_server.folders["INBOX"]
assert len(msgs_in_dest) == 1
assert b"Msg1" in msgs_in_dest[0]["content"]
# Verify it updated the cache to the NEW UIDVALIDITY (1) and last_uid (2)
with open(cache_path) as f:
new_cache = json.load(f)
src_state = new_cache["folders"]["INBOX"]["source_state"]
assert src_state["uid_validity"] == 1
# last_uid is 1 because UID 2 was skipped as a duplicate (pre-filtered),
# so the batch processor only saw UID 1.
# This is acceptable (conservative) behavior.
assert src_state["last_uid"] == 1