diff --git a/jobs/gmail-bulk-import/Dockerfile b/jobs/gmail-bulk-import/Dockerfile deleted file mode 100644 index a81cd62..0000000 --- a/jobs/gmail-bulk-import/Dockerfile +++ /dev/null @@ -1,11 +0,0 @@ -FROM python:3.11-slim - -WORKDIR /app - -COPY packages/kb-mail/ /packages/kb-mail/ -RUN pip install --no-cache-dir /packages/kb-mail/ - -COPY jobs/gmail-bulk-import/ /app/ -RUN pip install --no-cache-dir /app/ - -ENTRYPOINT ["gmail-bulk-import"] diff --git a/jobs/gmail-bulk-import/docker-compose.yml b/jobs/gmail-bulk-import/docker-compose.yml deleted file mode 100644 index d90fc01..0000000 --- a/jobs/gmail-bulk-import/docker-compose.yml +++ /dev/null @@ -1,13 +0,0 @@ -services: - gmail-bulk-import: - build: - context: ../.. - dockerfile: jobs/gmail-bulk-import/Dockerfile - restart: "no" - volumes: - - ${MBOX_PATH:?set MBOX_PATH=/path/to/allmail.mbox}:/data/input.mbox:ro - - ${ARCHIVE_ROOT:?set ARCHIVE_ROOT=/path/to/archive}:/data/archive - command: - - --mbox=/data/input.mbox - - --archive=/data/archive - - --dsn=${KB_DSN:-} diff --git a/jobs/gmail-bulk-import/src/gmail_bulk_import/importer.py b/jobs/gmail-bulk-import/src/gmail_bulk_import/importer.py index 7afee38..65735f7 100644 --- a/jobs/gmail-bulk-import/src/gmail_bulk_import/importer.py +++ b/jobs/gmail-bulk-import/src/gmail_bulk_import/importer.py @@ -1,10 +1,28 @@ """Gmail Takeout bulk importer — one-shot job. -Reads a .mbox file (Google Takeout "All Mail" export) and imports each -message into the append-only .eml archive and optionally into kb-postgres. +Install (from repo root): + pip install -e packages/kb-mail/ + pip install -e jobs/gmail-bulk-import/ -Idempotent: re-running against the same archive skips already-imported -messages (FileExistsError from archive → skipped; ON CONFLICT DO NOTHING in DB). +Usage: + # Archive only (no DB): + gmail-bulk-import --mbox ~/takeout/allmail.mbox --archive /data/kb/archive + + # With kb-postgres on PIHA (locally on PIHA): + gmail-bulk-import --mbox ~/takeout/allmail.mbox --archive /data/kb/archive \\ + --dsn postgresql://kb:@localhost:5433/kb + + # From another host pointing at PIHA: + gmail-bulk-import ... --dsn postgresql://kb:@piha:5433/kb + + # Dry run — count and parse, no writes: + gmail-bulk-import --mbox ~/takeout/allmail.mbox --archive /data/kb/archive --dry-run + + # Sample first 100 messages before full import: + gmail-bulk-import --mbox ~/takeout/allmail.mbox --archive /data/kb/archive --limit 100 + +Run at low priority on PIHA: + ionice -c 3 nice -n 19 gmail-bulk-import --mbox ... --archive ... """ from __future__ import annotations @@ -12,6 +30,7 @@ import argparse import asyncio import email.policy import hashlib +import json import mailbox import sys from datetime import datetime, timezone @@ -23,11 +42,22 @@ import asyncpg import structlog from kb_mail.archive import save_eml -from kb_mail.db import insert_envelope from kb_mail.envelope import Envelope _log = structlog.get_logger(__name__) +BATCH_SIZE = 500 +_EPOCH = datetime(1970, 1, 1, tzinfo=timezone.utc) + +# Mirrors archive._UNSAFE — used to reconstruct raw_ref for already-archived messages. +_ARCHIVE_UNSAFE = str.maketrans({"/": "_", "\\": "_", ":": "_", "<": "", ">": ""}) + + +def _eml_ref(envelope_id: str, ts: datetime) -> str: + """Compute the archive-relative .eml path without writing (mirrors save_eml logic).""" + safe_id = envelope_id.translate(_ARCHIVE_UNSAFE) + return f"gmail/{ts.year:04d}/{ts.month:02d}/{safe_id}.eml" + def _message_id(msg: mailbox.mboxMessage) -> str: """Return a stable envelope id from Message-ID header or SHA-256 content hash.""" @@ -38,7 +68,7 @@ def _message_id(msg: mailbox.mboxMessage) -> str: def _parse_date(msg: mailbox.mboxMessage) -> datetime: - """Parse Date header to UTC-aware datetime. Falls back to epoch if unparseable.""" + """Parse Date header to UTC-aware datetime. Returns epoch as last resort.""" date_str = msg.get("Date", "") if date_str: try: @@ -47,34 +77,70 @@ def _parse_date(msg: mailbox.mboxMessage) -> datetime: dt = dt.replace(tzinfo=timezone.utc) return dt.astimezone(timezone.utc) except Exception: - _log.warning("unparseable_date", header=date_str[:120]) - return datetime(1970, 1, 1, tzinfo=timezone.utc) + pass + return _EPOCH -async def _import_one( - msg: mailbox.mboxMessage, - archive_root: Path, - conn: Optional[asyncpg.Connection], -) -> str: - """Import a single message. Returns 'imported', 'skipped', or 'error'.""" - try: - envelope_id = _message_id(msg) - ts = _parse_date(msg) - raw = msg.as_bytes(policy=email.policy.compat32) +def _parse_attachments(msg: mailbox.mboxMessage) -> list[dict]: + """Return attachment descriptors from MIME structure. + + A part is an attachment if it has a filename, Content-Disposition: attachment, + or is not inline text/plain or text/html. + Errors on individual parts are logged and skipped — never propagated. + """ + attachments = [] + for part in msg.walk(): + if part.get_content_maintype() == "multipart": + continue + content_type = part.get_content_type() + disposition = part.get_content_disposition() or "" + filename = part.get_filename() + + if ( + content_type in ("text/plain", "text/html") + and disposition != "attachment" + and not filename + ): + continue try: - raw_ref = await save_eml(archive_root, envelope_id, "gmail", ts, raw) - except FileExistsError: - return "skipped" + payload = part.get_payload(decode=True) + if payload is None: + continue + attachments.append({ + "type": "attachment", + "filename": filename, + "content_type": content_type, + "size": len(payload), + "sha256": hashlib.sha256(payload).hexdigest(), + }) + except Exception: + _log.warning("attachment_parse_failed", + filename=filename, content_type=content_type) + return attachments - if conn is not None: - env = Envelope(id=envelope_id, source="gmail", ts=ts, raw_ref=raw_ref) - await insert_envelope(conn, env) - return "imported" - except Exception: - _log.exception("message_failed", id=msg.get("Message-ID", "")) - return "error" +async def _insert_batch(conn: asyncpg.Connection, envs: list[Envelope]) -> None: + rows = [ + ( + env.id, + env.source, + env.ts, + json.dumps(env.geo) if env.geo is not None else None, + env.raw_ref, + json.dumps(env.entities), + ) + for env in envs + ] + await conn.executemany( + """ + INSERT INTO envelope (id, source, ts, geo, raw_ref, entities) + VALUES ($1, $2, $3, $4::jsonb, $5, $6::jsonb) + ON CONFLICT (id) DO NOTHING + """, + rows, + ) + _log.debug("batch_flushed", count=len(envs)) async def run_import( @@ -82,34 +148,88 @@ async def run_import( archive_root: Path, dsn: Optional[str] = None, dry_run: bool = False, + limit: Optional[int] = None, ) -> dict[str, int]: """Import all messages from an mbox file. - Returns stats: {processed, imported, skipped, errors}. - With dry_run=True counts messages without writing anything. + Returns stats: {processed, imported, skipped, errors, epoch_fallback, + msgs_with_attachments, total_attachments, total_attachment_bytes}. + + dry_run=True parses and counts everything without writing. + limit=N stops after processing N messages (for sampling). """ - stats: dict[str, int] = {"processed": 0, "imported": 0, "skipped": 0, "errors": 0} + stats: dict[str, int] = { + "processed": 0, + "imported": 0, + "skipped": 0, + "errors": 0, + "epoch_fallback": 0, + "msgs_with_attachments": 0, + "total_attachments": 0, + "total_attachment_bytes": 0, + } conn: Optional[asyncpg.Connection] = None if not dry_run and dsn: conn = await asyncpg.connect(dsn) + pending: list[Envelope] = [] + + async def _flush() -> None: + if conn is not None and pending: + await _insert_batch(conn, pending) + pending.clear() + try: mbox = mailbox.mbox(str(mbox_path), create=False) for msg in mbox: + if limit is not None and stats["processed"] >= limit: + break + stats["processed"] += 1 + + envelope_id = _message_id(msg) + ts = _parse_date(msg) + + if ts == _EPOCH: + stats["epoch_fallback"] += 1 + _log.warning("epoch_fallback", id=envelope_id) + + attachments = _parse_attachments(msg) + stats["total_attachments"] += len(attachments) + stats["total_attachment_bytes"] += sum(a["size"] for a in attachments) + if attachments: + stats["msgs_with_attachments"] += 1 + if dry_run: stats["imported"] += 1 - else: - outcome = await _import_one(msg, archive_root, conn) - if outcome == "imported": + continue + + try: + raw = msg.as_bytes(policy=email.policy.compat32) + try: + raw_ref = await save_eml(archive_root, envelope_id, "gmail", ts, raw) stats["imported"] += 1 - elif outcome == "skipped": + except FileExistsError: + raw_ref = _eml_ref(envelope_id, ts) stats["skipped"] += 1 - else: - stats["errors"] += 1 + + pending.append( + Envelope(id=envelope_id, source="gmail", ts=ts, + raw_ref=raw_ref, entities=attachments) + ) + if len(pending) >= BATCH_SIZE: + await _flush() + + except Exception: + _log.exception("message_failed", id=envelope_id) + stats["errors"] += 1 + if stats["processed"] % 500 == 0: _log.info("progress", **stats) + + await _flush() + finally: if conn: await conn.close() @@ -127,9 +247,12 @@ def main() -> None: parser.add_argument("--archive", required=True, type=Path, metavar="DIR", help="Archive root directory") parser.add_argument("--dsn", metavar="DSN", - help="asyncpg DSN for kb-postgres (omit for archive-only mode)") + help="asyncpg DSN for kb-postgres " + "(e.g. postgresql://kb:@localhost:5433/kb on PIHA)") parser.add_argument("--dry-run", action="store_true", - help="Count messages only; no writes to archive or DB") + help="Parse and count only; no writes to archive or DB") + parser.add_argument("--limit", type=int, metavar="N", + help="Process at most N messages (sample before full import)") args = parser.parse_args() if not args.mbox.is_file(): @@ -137,7 +260,7 @@ def main() -> None: sys.exit(1) stats = asyncio.run( - run_import(args.mbox, args.archive, args.dsn or None, args.dry_run) + run_import(args.mbox, args.archive, args.dsn or None, args.dry_run, args.limit) ) sys.exit(1 if stats["errors"] > 0 else 0) diff --git a/jobs/gmail-bulk-import/tests/test_importer.py b/jobs/gmail-bulk-import/tests/test_importer.py index e29d442..5a0de47 100644 --- a/jobs/gmail-bulk-import/tests/test_importer.py +++ b/jobs/gmail-bulk-import/tests/test_importer.py @@ -1,13 +1,23 @@ """Unit tests for Gmail bulk importer — no DB, no external services required.""" from __future__ import annotations +import hashlib import mailbox from datetime import datetime, timezone +from email import encoders +from email.mime.base import MIMEBase +from email.mime.multipart import MIMEMultipart +from email.mime.text import MIMEText from pathlib import Path import pytest -from gmail_bulk_import.importer import _message_id, _parse_date, run_import +from gmail_bulk_import.importer import ( + _message_id, + _parse_attachments, + _parse_date, + run_import, +) def _make_msg(headers: dict) -> mailbox.mboxMessage: @@ -34,6 +44,36 @@ def _msg_with_id(mid: str = "") -> mailbox.mboxMessage: }) +def _mime_with_attachments( + mid: str = "", + payloads: list | None = None, +) -> mailbox.mboxMessage: + """Build a multipart mboxMessage with binary attachments. + + payloads: list of (filename, content_type, bytes). Defaults to PDF + PNG. + """ + if payloads is None: + payloads = [ + ("report.pdf", "application/pdf", b"PDF content here"), + ("photo.png", "image/png", b"PNG bytes"), + ] + mime = MIMEMultipart() + mime["From"] = "alice@example.com" + mime["To"] = "bob@example.com" + mime["Subject"] = "Files attached" + mime["Message-ID"] = mid + mime["Date"] = "Tue, 10 Jun 2025 12:00:00 +0000" + mime.attach(MIMEText("See attachments", "plain")) + for filename, ctype, data in payloads: + maintype, subtype = ctype.split("/", 1) + att = MIMEBase(maintype, subtype) + att.set_payload(data) + encoders.encode_base64(att) + att.add_header("Content-Disposition", "attachment", filename=filename) + mime.attach(att) + return mailbox.mboxMessage(mime) + + class TestMessageId: def test_uses_message_id_header(self): msg = _make_msg({"Message-ID": ""}) @@ -78,6 +118,30 @@ class TestParseDate: assert ts == datetime(2025, 6, 10, 12, 0, 0, tzinfo=timezone.utc) +class TestParseAttachments: + def test_plain_text_message_has_no_attachments(self): + assert _parse_attachments(_msg_with_id()) == [] + + def test_two_attachments_manifest(self): + atts = _parse_attachments(_mime_with_attachments()) + assert len(atts) == 2 + + pdf = next(a for a in atts if a["filename"] == "report.pdf") + assert pdf["type"] == "attachment" + assert pdf["content_type"] == "application/pdf" + assert pdf["size"] == len(b"PDF content here") + assert pdf["sha256"] == hashlib.sha256(b"PDF content here").hexdigest() + + png = next(a for a in atts if a["filename"] == "photo.png") + assert png["content_type"] == "image/png" + assert png["size"] == len(b"PNG bytes") + assert png["sha256"] == hashlib.sha256(b"PNG bytes").hexdigest() + + def test_attachment_descriptors_have_required_keys(self): + for att in _parse_attachments(_mime_with_attachments()): + assert {"type", "filename", "content_type", "size", "sha256"} <= att.keys() + + class TestRunImport: async def test_dry_run_counts_without_writing(self, tmp_path): mbox_path = tmp_path / "mail.mbox" @@ -100,8 +164,7 @@ class TestRunImport: assert stats["imported"] == 1 assert stats["skipped"] == 0 assert stats["errors"] == 0 - eml_files = list(archive.rglob("*.eml")) - assert len(eml_files) == 1 + assert len(list(archive.rglob("*.eml"))) == 1 async def test_eml_stored_under_gmail_source(self, tmp_path): mbox_path = tmp_path / "mail.mbox" @@ -137,7 +200,6 @@ class TestRunImport: assert len(list(archive.rglob("*.eml"))) == 2 async def test_archive_only_mode_no_dsn(self, tmp_path): - """Without a DSN, job runs archive-only (no DB connection attempted).""" mbox_path = tmp_path / "mail.mbox" archive = tmp_path / "archive" _make_mbox(mbox_path, [_msg_with_id()]) @@ -146,3 +208,76 @@ class TestRunImport: assert stats["imported"] == 1 assert list(archive.rglob("*.eml")) + + async def test_limit_truncates_processing(self, tmp_path): + mbox_path = tmp_path / "mail.mbox" + archive = tmp_path / "archive" + _make_mbox(mbox_path, [_msg_with_id(f"") for i in range(5)]) + + stats = await run_import(mbox_path, archive, limit=3) + + assert stats["processed"] == 3 + assert stats["imported"] == 3 + + async def test_limit_with_dry_run(self, tmp_path): + mbox_path = tmp_path / "mail.mbox" + archive = tmp_path / "archive" + _make_mbox(mbox_path, [_msg_with_id(f"") for i in range(10)]) + + stats = await run_import(mbox_path, archive, dry_run=True, limit=4) + + assert stats["processed"] == 4 + assert stats["imported"] == 4 + assert not archive.exists() + + async def test_epoch_fallback_counted(self, tmp_path): + mbox_path = tmp_path / "mail.mbox" + archive = tmp_path / "archive" + _make_mbox(mbox_path, [ + _make_msg({"Message-ID": "", "Date": "not-a-date"}), + _make_msg({"Message-ID": ""}), + _msg_with_id(""), + ]) + + stats = await run_import(mbox_path, archive) + + assert stats["epoch_fallback"] == 2 + assert stats["processed"] == 3 + + async def test_attachment_stats_aggregated(self, tmp_path): + mbox_path = tmp_path / "mail.mbox" + archive = tmp_path / "archive" + _make_mbox(mbox_path, [ + _msg_with_id(""), + _mime_with_attachments(""), + ]) + + stats = await run_import(mbox_path, archive) + + assert stats["msgs_with_attachments"] == 1 + assert stats["total_attachments"] == 2 + assert stats["total_attachment_bytes"] == ( + len(b"PDF content here") + len(b"PNG bytes") + ) + + async def test_batch_flush_multiple_messages(self, tmp_path): + """BATCH_SIZE=500; 3 messages stay in pending until end-of-loop flush.""" + mbox_path = tmp_path / "mail.mbox" + archive = tmp_path / "archive" + _make_mbox(mbox_path, [_msg_with_id(f"") for i in range(3)]) + + stats = await run_import(mbox_path, archive) + + assert stats["imported"] == 3 + assert stats["errors"] == 0 + assert len(list(archive.rglob("*.eml"))) == 3 + + async def test_attachment_entities_in_stats(self, tmp_path): + mbox_path = tmp_path / "mail.mbox" + archive = tmp_path / "archive" + _make_mbox(mbox_path, [_mime_with_attachments()]) + + stats = await run_import(mbox_path, archive) + + assert stats["imported"] == 1 + assert stats["total_attachments"] == 2