Implement Gmail labels preservation feature and update tests

This commit is contained in:
Javier Callico 2026-01-26 19:41:18 -05:00
parent bfac175c0b
commit 2c230dcf64
4 changed files with 570 additions and 1 deletions

1
.gitignore vendored
View File

@ -6,3 +6,4 @@ __pycache__/
env/
venv/
.env
.coverage

View File

@ -38,6 +38,7 @@ This repository contains a set of Python scripts designed to migrate emails betw
- **Format**: Saves emails as individual `.eml` files (RFC 5322), compatible with Outlook, Thunderbird, and Apple Mail.
- **Structure**: Replicates the IMAP folder hierarchy locally.
- **Incremental**: Skips emails that have already been downloaded (based on UID) so you can run it periodically to fetch new messages.
- **Gmail Labels Preservation**: Creates a `labels_manifest.json` file mapping each email's Message-ID to its Gmail labels, enabling proper restoration with labels intact.
## Getting Started
@ -197,6 +198,60 @@ python3 backup_imap_emails.py --dest-path "/Users/jdoe/Documents/Emails"
python3 backup_imap_emails.py --dest-path "./my_backup" "[Gmail]/Sent Mail"
```
### 6. Gmail Backup with Labels Preservation
When backing up a Gmail account, use `--preserve-labels` to create a manifest file that maps each email to its Gmail labels. This is essential for restoring emails to another Gmail account with labels intact.
```bash
# Recommended: Backup All Mail with labels manifest
python3 backup_imap_emails.py \
--src-host "imap.gmail.com" \
--src-user "you@gmail.com" \
--src-pass "your-app-password" \
--dest-path "./gmail_backup" \
--preserve-labels \
"[Gmail]/All Mail"
```
**For large accounts (100K+ emails)**, you can build the manifest first to test:
```bash
# Step 1: Build manifest only (fast, no download)
python3 backup_imap_emails.py \
--src-host "imap.gmail.com" \
--src-user "you@gmail.com" \
--src-pass "your-app-password" \
--dest-path "./gmail_backup" \
--manifest-only
# Step 2: Download emails (can run later, manifest already exists)
python3 backup_imap_emails.py \
--src-host "imap.gmail.com" \
--src-user "you@gmail.com" \
--src-pass "your-app-password" \
--dest-path "./gmail_backup" \
"[Gmail]/All Mail"
```
**How it works:**
1. The script scans ALL folders in your Gmail account to identify which emails have which labels
2. Creates a `labels_manifest.json` file mapping each email's `Message-ID` to its labels
3. Downloads all emails from `[Gmail]/All Mail` (contains every email once, no duplicates)
**Example `labels_manifest.json`:**
```json
{
"<CAExample123@mail.gmail.com>": ["INBOX", "Work", "Projects/2024"],
"<CAExample456@mail.gmail.com>": ["Sent Mail", "Personal", "Starred"]
}
```
**Benefits:**
- ✅ No duplicate emails on disk (each email saved once)
- ✅ Labels are preserved for restoration
- ✅ Includes system labels (INBOX, Sent Mail, Starred) and user labels
- ✅ Progress reporting for large accounts
**Note:** Gmail labels like "Important" that are auto-managed by Gmail are excluded from the manifest as they cannot be reliably restored.
## Troubleshooting
- **"Too many simultaneous connections"**:

View File

@ -10,6 +10,7 @@ Features:
- Filename Sanitization: Saves files as "{UID}_{Subject}.eml" with unsafe characters removed.
- Folder Replication: Recreates the IMAP folder structure locally.
- Parallel Processing: Uses multithreading for fast downloads.
- Gmail Labels Preservation: Creates a manifest mapping Message-IDs to Gmail labels for restoration.
Configuration:
SRC_IMAP_HOST, SRC_IMAP_USERNAME, SRC_IMAP_PASSWORD: Source credentials.
@ -17,10 +18,16 @@ Configuration:
Usage:
python3 backup_imap_emails.py --dest-path "./my_backup"
Gmail Labels:
python3 backup_imap_emails.py --dest-path "./my_backup" --preserve-labels "[Gmail]/All Mail"
This backs up all emails from [Gmail]/All Mail and creates a labels_manifest.json
file that maps each email's Message-ID to its Gmail labels for later restoration.
"""
import argparse
import concurrent.futures
import json
import os
import sys
import threading
@ -35,6 +42,16 @@ BATCH_SIZE = 10
thread_local = threading.local()
print_lock = threading.Lock()
# Gmail-specific folders to exclude from label mapping
GMAIL_SYSTEM_FOLDERS = {
"[Gmail]/All Mail",
"[Gmail]/Spam",
"[Gmail]/Trash",
"[Gmail]/Drafts",
"[Gmail]/Bin",
"[Gmail]/Important", # This is actually a label, but often system-managed
}
def safe_print(message):
t_name = threading.current_thread().name
@ -142,6 +159,212 @@ def get_existing_uids(local_path):
return existing
def is_gmail_label_folder(folder_name):
"""
Determines if a folder represents a Gmail label (user-created or system label
that should be preserved).
Excludes system folders like All Mail, Spam, Trash, Drafts.
"""
# Exclude system folders that aren't really "labels"
if folder_name in GMAIL_SYSTEM_FOLDERS:
return False
# INBOX is a special case - it's a label in Gmail
if folder_name == "INBOX":
return True
# [Gmail]/Sent Mail and [Gmail]/Starred are labels worth preserving
if folder_name in ("[Gmail]/Sent Mail", "[Gmail]/Starred"):
return True
# Any folder NOT under [Gmail]/ is a user label
if not folder_name.startswith("[Gmail]/"):
return True
return False
def get_message_ids_in_folder(imap_conn, folder_name, progress_callback=None):
"""
Returns a set of Message-IDs for all emails in a given folder.
Optional progress_callback(current, total) for progress reporting.
"""
message_ids = set()
try:
imap_conn.select(f'"{folder_name}"', readonly=True)
except Exception as e:
safe_print(f"Could not select folder {folder_name}: {e}")
return message_ids
try:
resp, data = imap_conn.uid("search", None, "ALL")
if resp != "OK" or not data or not data[0]:
return message_ids
uids = data[0].split()
if not uids:
return message_ids
total_uids = len(uids)
# Fetch Message-IDs in batches - use larger batch for header-only fetches
batch_size = 200
for i in range(0, len(uids), batch_size):
batch = uids[i : i + batch_size]
uid_range = b",".join(batch)
# Report progress
if progress_callback:
progress_callback(min(i + batch_size, total_uids), total_uids)
try:
resp, items = imap_conn.uid("fetch", uid_range, "(BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])")
if resp != "OK":
continue
for item in items:
if isinstance(item, tuple) and len(item) >= 2:
header_data = item[1]
if isinstance(header_data, bytes):
header_str = header_data.decode("utf-8", errors="ignore")
# Extract Message-ID from header
for line in header_str.split("\n"):
if line.lower().startswith("message-id:"):
msg_id = line.split(":", 1)[1].strip()
if msg_id:
message_ids.add(msg_id)
break
except Exception as e:
safe_print(f"Error fetching batch in {folder_name}: {e}")
# Try to keep connection alive
try:
imap_conn.noop()
except Exception:
pass
continue
except Exception as e:
safe_print(f"Error searching folder {folder_name}: {e}")
return message_ids
def build_labels_manifest(imap_conn, local_path):
"""
Builds a manifest mapping Message-IDs to their Gmail labels.
Scans all folders (labels) in the account and records which Message-IDs
appear in each label.
Returns a dict: { "message-id": ["Label1", "Label2", ...], ... }
Saves the manifest to labels_manifest.json in the backup directory.
"""
import time
manifest = {}
start_time = time.time()
total_emails_scanned = 0
safe_print("--- Building Gmail Labels Manifest ---")
# Get all folders
try:
typ, folders = imap_conn.list()
if typ != "OK":
safe_print("Error: Could not list folders for label mapping.")
return manifest
except Exception as e:
safe_print(f"Error listing folders: {e}")
return manifest
# Parse folder names and filter to label folders
label_folders = []
for f_info in folders:
name = imap_common.normalize_folder_name(f_info)
if is_gmail_label_folder(name):
label_folders.append(name)
total_folders = len(label_folders)
safe_print(f"Found {total_folders} label folders to scan.\n")
# Scan each label folder
for folder_idx, folder_name in enumerate(label_folders, 1):
folder_start = time.time()
# Progress callback for this folder
def progress_cb(current, total):
elapsed = time.time() - start_time
print(
f"\r Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s) ",
end="",
flush=True,
)
safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
message_ids = get_message_ids_in_folder(imap_conn, folder_name, progress_cb)
print() # New line after progress
# Keep connection alive between folders
try:
imap_conn.noop()
except Exception:
pass
# Determine the label name to store
# For [Gmail]/Sent Mail -> "Sent Mail"
# For [Gmail]/Starred -> "Starred"
# For INBOX -> "INBOX"
# For user folders -> folder name as-is
if folder_name.startswith("[Gmail]/"):
label_name = folder_name[8:] # Remove "[Gmail]/" prefix
else:
label_name = folder_name
for msg_id in message_ids:
if msg_id not in manifest:
manifest[msg_id] = []
if label_name not in manifest[msg_id]:
manifest[msg_id].append(label_name)
folder_elapsed = time.time() - folder_start
total_emails_scanned += len(message_ids)
safe_print(f" -> {len(message_ids)} emails with label '{label_name}' ({folder_elapsed:.1f}s)")
# Summary
total_elapsed = time.time() - start_time
safe_print("\nManifest building complete:")
safe_print(f" - Folders scanned: {total_folders}")
safe_print(f" - Total email-label mappings: {total_emails_scanned}")
safe_print(f" - Unique emails with labels: {len(manifest)}")
safe_print(f" - Time elapsed: {total_elapsed:.1f}s")
# Save manifest
manifest_path = os.path.join(local_path, "labels_manifest.json")
try:
with open(manifest_path, "w", encoding="utf-8") as f:
json.dump(manifest, f, indent=2, ensure_ascii=False)
safe_print(f"\nLabels manifest saved to: {manifest_path}")
except Exception as e:
safe_print(f"Error saving manifest: {e}")
return manifest
def load_labels_manifest(local_path):
"""
Loads an existing labels manifest from the backup directory.
Returns the manifest dict or empty dict if not found.
"""
manifest_path = os.path.join(local_path, "labels_manifest.json")
if os.path.exists(manifest_path):
try:
with open(manifest_path, encoding="utf-8") as f:
return json.load(f)
except Exception:
pass
return {}
def backup_folder(src_main, folder_name, local_base_path, src_conf):
safe_print(f"--- Processing Folder: {folder_name} ---")
@ -233,6 +456,19 @@ def main():
# Config
parser.add_argument("--workers", type=int, default=int(os.getenv("MAX_WORKERS", 10)), help="Thread count")
parser.add_argument("--batch", type=int, default=int(os.getenv("BATCH_SIZE", 10)), help="Emails per batch")
# Gmail Labels
parser.add_argument(
"--preserve-labels",
action="store_true",
help="Gmail only: Create a labels_manifest.json mapping Message-IDs to labels for restoration",
)
parser.add_argument(
"--manifest-only",
action="store_true",
help="Gmail only: Build the labels manifest and exit without downloading emails",
)
parser.add_argument("folder", nargs="?", help="Specific folder to backup")
args = parser.parse_args()
@ -275,8 +511,12 @@ def main():
print(f"Source Host : {args.src_host}")
print(f"Source User : {args.src_user}")
print(f"Destination Path: {local_path}")
if args.folder:
if args.manifest_only:
print("Mode : Manifest Only (no email download)")
elif args.folder:
print(f"Target Folder : {args.folder}")
if args.preserve_labels or args.manifest_only:
print("Preserve Labels : Yes (Gmail)")
print("-----------------------------\n")
try:
@ -284,6 +524,24 @@ def main():
if not src:
sys.exit(1)
# Build labels manifest BEFORE backing up emails
# This way we capture the label state at backup time
if args.preserve_labels or args.manifest_only:
print("Building Gmail labels manifest...")
print("This scans all folders to map Message-IDs to labels.\n")
build_labels_manifest(src, local_path)
print("") # Blank line after manifest building
# If manifest-only mode, we're done
if args.manifest_only:
src.logout()
manifest_path = os.path.join(local_path, "labels_manifest.json")
print("\nManifest-only mode complete.")
print(f"Labels manifest saved to: {manifest_path}")
print("\nTo download emails, run again without --manifest-only:")
print(f' python3 backup_imap_emails.py --dest-path "{local_path}" "[Gmail]/All Mail"')
sys.exit(0)
if args.folder:
backup_folder(src, args.folder, local_path, src_conf)
else:
@ -296,6 +554,11 @@ def main():
src.logout()
print("\nBackup completed successfully.")
if args.preserve_labels:
manifest_path = os.path.join(local_path, "labels_manifest.json")
print(f"\nGmail labels manifest saved to: {manifest_path}")
print("Use this file when restoring to reapply labels to emails.")
except KeyboardInterrupt:
print("\nBackup interrupted by user.")
sys.exit(0)

View File

@ -369,3 +369,253 @@ class TestBackupErrorHandling:
captured = capsys.readouterr()
assert "Error creating backup directory" in captured.out
class TestGmailLabelsPreservation:
"""Tests for Gmail labels manifest functionality."""
def test_is_gmail_label_folder_user_labels(self):
"""Test that user labels are correctly identified."""
assert backup_imap_emails.is_gmail_label_folder("Work") is True
assert backup_imap_emails.is_gmail_label_folder("Personal") is True
assert backup_imap_emails.is_gmail_label_folder("Projects/2024") is True
assert backup_imap_emails.is_gmail_label_folder("INBOX") is True
def test_is_gmail_label_folder_gmail_labels(self):
"""Test that Gmail system labels are correctly identified."""
assert backup_imap_emails.is_gmail_label_folder("[Gmail]/Sent Mail") is True
assert backup_imap_emails.is_gmail_label_folder("[Gmail]/Starred") is True
def test_is_gmail_label_folder_system_folders(self):
"""Test that Gmail system folders are excluded."""
assert backup_imap_emails.is_gmail_label_folder("[Gmail]/All Mail") is False
assert backup_imap_emails.is_gmail_label_folder("[Gmail]/Spam") is False
assert backup_imap_emails.is_gmail_label_folder("[Gmail]/Trash") is False
assert backup_imap_emails.is_gmail_label_folder("[Gmail]/Drafts") is False
assert backup_imap_emails.is_gmail_label_folder("[Gmail]/Bin") is False
def test_load_labels_manifest_nonexistent(self, tmp_path):
"""Test loading manifest when file doesn't exist."""
result = backup_imap_emails.load_labels_manifest(str(tmp_path))
assert result == {}
def test_load_labels_manifest_existing(self, tmp_path):
"""Test loading existing manifest file."""
import json
manifest_data = {
"<msg1@test.com>": ["INBOX", "Work"],
"<msg2@test.com>": ["Sent Mail", "Personal"],
}
manifest_path = tmp_path / "labels_manifest.json"
manifest_path.write_text(json.dumps(manifest_data))
result = backup_imap_emails.load_labels_manifest(str(tmp_path))
assert result == manifest_data
def test_load_labels_manifest_invalid_json(self, tmp_path):
"""Test loading invalid JSON manifest file."""
manifest_path = tmp_path / "labels_manifest.json"
manifest_path.write_text("not valid json {{{")
result = backup_imap_emails.load_labels_manifest(str(tmp_path))
assert result == {}
def test_get_message_ids_in_folder(self, monkeypatch):
"""Test extraction of message IDs from a folder."""
mock_conn = MagicMock()
mock_conn.select.return_value = ("OK", [b"1"])
mock_conn.uid.side_effect = [
("OK", [b"1 2 3"]), # search result
(
"OK",
[
(b"1 (BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: <msg1@test.com>\r\n"),
b")",
(b"2 (BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: <msg2@test.com>\r\n"),
b")",
(b"3 (BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: <msg3@test.com>\r\n"),
b")",
],
), # fetch result
]
result = backup_imap_emails.get_message_ids_in_folder(mock_conn, "INBOX", None)
assert "<msg1@test.com>" in result
assert "<msg2@test.com>" in result
assert "<msg3@test.com>" in result
def test_get_message_ids_in_folder_with_progress(self, monkeypatch):
"""Test extraction of message IDs with progress callback."""
mock_conn = MagicMock()
mock_conn.select.return_value = ("OK", [b"1"])
mock_conn.uid.side_effect = [
("OK", [b"1 2 3"]), # search result
(
"OK",
[
(b"1 (BODY[HEADER.FIELDS (MESSAGE-ID)] {30}", b"Message-ID: <msg1@test.com>\r\n"),
b")",
],
), # fetch result
]
progress_calls = []
def progress_cb(current, total):
progress_calls.append((current, total))
backup_imap_emails.get_message_ids_in_folder(mock_conn, "INBOX", progress_cb)
# Progress should have been called
assert len(progress_calls) > 0
# Last call should show completion
assert progress_calls[-1][0] == progress_calls[-1][1]
def test_get_message_ids_in_folder_select_error(self, monkeypatch):
"""Test handling of folder select error."""
mock_conn = MagicMock()
mock_conn.select.side_effect = Exception("Select failed")
result = backup_imap_emails.get_message_ids_in_folder(mock_conn, "INBOX")
assert result == set()
def test_get_message_ids_in_folder_empty(self, monkeypatch):
"""Test extraction from empty folder."""
mock_conn = MagicMock()
mock_conn.select.return_value = ("OK", [b"0"])
mock_conn.uid.return_value = ("OK", [b""])
result = backup_imap_emails.get_message_ids_in_folder(mock_conn, "INBOX", None)
assert result == set()
def test_build_labels_manifest(self, monkeypatch, tmp_path):
"""Test building labels manifest from mock folders."""
mock_conn = MagicMock()
# Mock folder list - simulate Gmail structure
mock_conn.list.return_value = (
"OK",
[
b'(\\HasNoChildren) "/" "INBOX"',
b'(\\HasNoChildren) "/" "Work"',
b'(\\HasNoChildren) "/" "[Gmail]/All Mail"',
b'(\\HasNoChildren) "/" "[Gmail]/Sent Mail"',
],
)
# Track which folder is selected
folder_data = {
"INBOX": {"<msg1@test.com>", "<msg2@test.com>"},
"Work": {"<msg1@test.com>"},
"[Gmail]/Sent Mail": {"<msg2@test.com>"},
}
def mock_get_message_ids(conn, folder, progress_cb=None):
return folder_data.get(folder, set())
monkeypatch.setattr(backup_imap_emails, "get_message_ids_in_folder", mock_get_message_ids)
result = backup_imap_emails.build_labels_manifest(mock_conn, str(tmp_path))
# Check manifest structure
assert "<msg1@test.com>" in result
assert "<msg2@test.com>" in result
assert "INBOX" in result["<msg1@test.com>"]
assert "Work" in result["<msg1@test.com>"]
assert "INBOX" in result["<msg2@test.com>"]
assert "Sent Mail" in result["<msg2@test.com>"]
# Check file was saved
manifest_path = tmp_path / "labels_manifest.json"
assert manifest_path.exists()
def test_build_labels_manifest_list_error(self, monkeypatch, tmp_path):
"""Test handling of folder list error."""
mock_conn = MagicMock()
mock_conn.list.return_value = ("NO", [])
result = backup_imap_emails.build_labels_manifest(mock_conn, str(tmp_path))
assert result == {}
def test_preserve_labels_flag_integration(self, single_mock_server, monkeypatch, tmp_path):
"""Test --preserve-labels flag creates manifest."""
src_data = {
"INBOX": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
"Work": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
"[Gmail]/All Mail": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
}
server, port = single_mock_server(src_data)
env = {
"SRC_IMAP_HOST": "localhost",
"SRC_IMAP_USERNAME": "user",
"SRC_IMAP_PASSWORD": "pass",
}
monkeypatch.setattr(os, "environ", env)
monkeypatch.setattr(
sys,
"argv",
["backup_imap_emails.py", "--dest-path", str(tmp_path), "--preserve-labels", "[Gmail]/All Mail"],
)
monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
backup_imap_emails.main()
# Check manifest was created
manifest_path = tmp_path / "labels_manifest.json"
assert manifest_path.exists()
def test_manifest_only_flag(self, single_mock_server, monkeypatch, tmp_path):
"""Test --manifest-only flag creates manifest without downloading emails."""
src_data = {
"INBOX": [b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody"],
"Work": [b"Subject: Work Email\r\nMessage-ID: <2@test>\r\n\r\nBody"],
"[Gmail]/All Mail": [
b"Subject: Inbox Email\r\nMessage-ID: <1@test>\r\n\r\nBody",
b"Subject: Work Email\r\nMessage-ID: <2@test>\r\n\r\nBody",
],
}
server, port = single_mock_server(src_data)
env = {
"SRC_IMAP_HOST": "localhost",
"SRC_IMAP_USERNAME": "user",
"SRC_IMAP_PASSWORD": "pass",
}
monkeypatch.setattr(os, "environ", env)
monkeypatch.setattr(
sys,
"argv",
["backup_imap_emails.py", "--dest-path", str(tmp_path), "--manifest-only"],
)
monkeypatch.setattr("imap_common.get_imap_connection", make_single_mock_connection(port))
# Should exit with code 0 after creating manifest
with pytest.raises(SystemExit) as exc_info:
backup_imap_emails.main()
assert exc_info.value.code == 0
# Check manifest was created
manifest_path = tmp_path / "labels_manifest.json"
assert manifest_path.exists()
# Check NO email folders were created (no download happened)
# Only the manifest file should exist
items = list(tmp_path.iterdir())
assert len(items) == 1
assert items[0].name == "labels_manifest.json"
def test_gmail_system_folders_constant(self):
"""Test that GMAIL_SYSTEM_FOLDERS contains expected entries."""
expected = {
"[Gmail]/All Mail",
"[Gmail]/Spam",
"[Gmail]/Trash",
"[Gmail]/Drafts",
"[Gmail]/Bin",
"[Gmail]/Important",
}
assert backup_imap_emails.GMAIL_SYSTEM_FOLDERS == expected