From 74d2156a752ba42c0a57f615611407f6309b0f5d Mon Sep 17 00:00:00 2001 From: Chris Czub Date: Wed, 22 Jul 2026 12:12:27 -0400 Subject: [PATCH 1/2] Apply rustfmt to web.rs (whitespace only, no functional change) cargo fmt --check was already failing on main for this file; fixing it separately so the importer change stays clean. --- src/bin/web.rs | 33 ++++++++++++++++----------------- 1 file changed, 16 insertions(+), 17 deletions(-) diff --git a/src/bin/web.rs b/src/bin/web.rs index 9e0b59b..626868c 100644 --- a/src/bin/web.rs +++ b/src/bin/web.rs @@ -176,27 +176,26 @@ async fn build_summary(pool: &SqlitePool) -> Result Response { let detail = err.to_string(); let lower = detail.to_lowercase(); - let (status, kind, message) = if lower.contains("unable to open database file") - || lower.contains("no such table") - { - ( + let (status, kind, message) = + if lower.contains("unable to open database file") || lower.contains("no such table") { + ( StatusCode::SERVICE_UNAVAILABLE, "missing_db", "stats.db (or its tables) not found — set up the database and run the scanner first", ) - } else if lower.contains("database is locked") || lower.contains("database is busy") { - ( - StatusCode::SERVICE_UNAVAILABLE, - "busy", - "database busy — scanner is writing, retry shortly", - ) - } else { - ( - StatusCode::INTERNAL_SERVER_ERROR, - "internal", - "internal error", - ) - }; + } else if lower.contains("database is locked") || lower.contains("database is busy") { + ( + StatusCode::SERVICE_UNAVAILABLE, + "busy", + "database busy — scanner is writing, retry shortly", + ) + } else { + ( + StatusCode::INTERNAL_SERVER_ERROR, + "internal", + "internal error", + ) + }; ( status, Json(json!({ "error": kind, "message": message, "detail": detail })), From 9c03697a2e169c897a633a213740a821d9241a36 Mon Sep 17 00:00:00 2001 From: Chris Czub Date: Wed, 22 Jul 2026 12:12:46 -0400 Subject: [PATCH 2/2] Add Takeout mbox importer, ingest run tracking, and ingestion coordination Phase A of the ingestion design in #26; implements the importer from #25. New 'import ' subcommand: a streaming, header-only mbox parser (bounded memory regardless of file size, mboxrd >From escaping, CRLF, capped header lines, malformed-region recovery with skip counts). Messages dedupe into seen_mails under the 'mid:' namespace so re-imports add zero counts and RFC ids can never collide with Gmail API ids. Sender extraction shares the scan's From/Return-Path priority and cleanup, so counts are consistent across sources. Ingestion coordination, shared by scan and import: - ingest_runs table (migration 1) records one row per run: source, state (running/done/failed/cancelled/abandoned), pid, timestamps with a ~2s heartbeat, messages_seen/new, bytes_done/bytes_total for imports (the byte offset doubles as the resume point), resume_token for scans, and error kind/message on failure. All writes ride the existing single db_writer channel. - Schema versioning via PRAGMA user_version with guarded migrations; the pre-existing tables are the version-0 baseline and old databases upgrade transparently on startup. Databases newer than the binary are refused. - Exclusive kernel flock on .ingest.lock before touching the database; a second ingester exits promptly with a clear message, and the next ingester marks rows orphaned by a dead one as abandoned (liveness is judged by the lock, never by pids or heartbeats). flock comes from rustix (fs feature only): safe wrappers, no unsafe in this crate. - SIGTERM/SIGINT stop fetching/parsing, drain the writer, persist resume state, mark the row cancelled, and exit 0. 'import --resume ' continues from the recorded offset after validating a file fingerprint (size + mtime + hash of the first 64 KiB), falling back to a full re-parse (still correct via dedupe) if the file changed. - --db/--credentials/--tokens flags with GMAIL_STATS_* env fallbacks and the previous cwd-relative defaults; bare 'cargo run' remains the OAuth scan. Output is now quiet-ish by default (periodic progress lines); --verbose restores the per-message sender lines, which were also a privacy leak, and --quiet reduces output to errors plus the summary. The scan's credential/auth failures now land as failed run rows with an error kind instead of panics, and it records a best-effort total_estimate from users.getProfile for future progress display. Tests cover the parser fixtures (escaping, CRLF, folding, malformed regions, truncation, resume offsets, ~50MB on-disk and ~64MB generated streams), import idempotency, resume with matching and stale fingerprints, cancellation, the flock refusing a second ingester, and migration of a version-0 database. README documents both ingestion modes, the Takeout how-to including the ~2 day Advanced Protection delay, the new flags, and the interim double-counting caveat when mixing sources in one database. --- .gitignore | 1 + Cargo.lock | 61 +++ Cargo.toml | 8 +- README.md | 109 +++- src/ingest.rs | 679 +++++++++++++++++++++++ src/main.rs | 1450 +++++++++++++++++++++++++++++++++++++++---------- src/mbox.rs | 657 ++++++++++++++++++++++ 7 files changed, 2677 insertions(+), 288 deletions(-) create mode 100644 src/ingest.rs create mode 100644 src/mbox.rs diff --git a/.gitignore b/.gitignore index a64310c..45cdc1c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ /target tokencache.json stats.db* +*.ingest.lock credentials.json *.swp _site/ diff --git a/Cargo.lock b/Cargo.lock index 9674d27..c460cd2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -418,6 +418,16 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + [[package]] name = "etcetera" version = "0.11.0" @@ -439,6 +449,12 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -619,9 +635,11 @@ dependencies = [ "google-gmail1", "lazy_static", "regex", + "rustix", "serde", "serde_json", "sqlx", + "tempfile", "tokio", "yup-oauth2", ] @@ -964,6 +982,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "lock_api" version = "0.4.14" @@ -1226,6 +1250,19 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.23.42" @@ -1484,6 +1521,16 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + [[package]] name = "slab" version = "0.4.7" @@ -1753,6 +1800,19 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "thiserror" version = "2.0.19" @@ -1830,6 +1890,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "windows-sys 0.61.2", diff --git a/Cargo.toml b/Cargo.toml index 14230e6..ea5789b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,8 +12,14 @@ futures = "0.3" google-gmail1 = "7.0" lazy_static = "1.5" regex = "1.13" +# Kernel flock() for the cross-process ingest lock: safe wrappers (no unsafe in +# this crate), tiny with only the fs feature enabled. Unix-only, by design. +rustix = { version = "1.0", features = ["fs"] } serde = "1.0" serde_json = "1.0" sqlx = { version = "0.9", features = ["runtime-tokio", "sqlite"] } -tokio = { version = "1.53", features = ["macros", "rt-multi-thread", "sync", "time"] } +tokio = { version = "1.53", features = ["macros", "rt-multi-thread", "signal", "sync", "time"] } yup-oauth2 = "12" + +[dev-dependencies] +tempfile = "3" diff --git a/README.md b/README.md index 74118da..0952064 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,32 @@ access to manage my entire inbox, which I didn't feel comfortable with. So, I spent a few hours and hacked this together. The code's not great, it's slow and very basic, but it works. # How? + +There are two ways to get your mail into the database. Both fill the same +tables and use the same sender normalization, so the resulting stats look the +same either way: + +| | **Gmail API scan** (default) | **Takeout mbox import** | +|---|---|---| +| Command | `cargo run` | `cargo run -- import path/to/All-mail.mbox` | +| Needs | A Google Cloud OAuth client + browser consent | A Google Takeout export file | +| Network | Yes — fetches every message's headers (slow for big inboxes, rate-limited) | No — purely local, fast | +| Works with Advanced Protection | **No** — Google blocks Gmail OAuth scopes for APP accounts outright (`Error 400: policy_enforced`) | **Yes** | +| Stays current | Re-run any time to pick up new mail | Snapshot as of the export | + +Choose the **API scan** if you can OAuth and want to re-scan incrementally; +choose the **Takeout import** if you're enrolled in Advanced Protection, don't +want to create a Cloud project, or already have an export lying around. + +> **Warning: don't mix both modes into one database (yet).** The scan keys +> messages by Gmail's internal id and the importer by RFC `Message-ID`, so a +> message ingested by both is counted twice. Each mode on its own is fully +> idempotent — re-running never double-counts. Cross-source dedupe is planned +> (Phase D of [#26](https://github.com/zbuc/gmail-stats/issues/26)); until it +> lands, use separate `--db` files if you want both. + +## Option A: Gmail API scan (OAuth) + First follow the Google Workspace [setup instructions](https://developers.google.com/workspace/guides/get-started) to get OAuth credentials associated with a Google Cloud Project. @@ -34,10 +60,85 @@ Finally, run the application: $ cargo run ``` +This will trigger an OAuth flow which you launch in your browser, after which the access credentials are stored on disk (in `tokencache.json` by default). +Please be aware that these are credentials that would allow anyone to read the contents of your email inbox, so you probably want to `rm tokencache.json` +after you're done. + +A scan can be interrupted at any time with Ctrl-C (or SIGTERM): it drains +pending writes, records its progress, and exits cleanly. Because every message +is deduplicated, simply re-running the scan continues where it makes sense and +never double-counts. + +## Option B: Google Takeout mbox import + +No Google Cloud project, no OAuth, no network access — and the only option +that works for accounts enrolled in Google's Advanced Protection Program. + +1. Go to [takeout.google.com](https://takeout.google.com), deselect + everything except **Mail**, and request the export. Gmail exports arrive in + **mbox** format. + + > **Advanced Protection note:** for exactly the accounts that need this + > path, the export is slow — with Advanced Protection enabled, Google + > deliberately waits **about 2 days** before delivering a Takeout export. + > Plan ahead; once downloaded, the import itself is local and fast. + +2. Download and extract the archive; you're looking for the large `.mbox` + file (often named `All mail Including Spam and Trash.mbox`). + +3. Import it: + + ```console + $ cargo run -- import "path/to/All mail Including Spam and Trash.mbox" + ``` + +The importer streams the file (multi-GB archives are fine — memory use stays +flat), parses only the message headers, and skips unparseable regions with a +count in the final summary. Re-running the same import adds zero new counts: +messages are deduplicated by their `Message-ID` header. + +An interrupted import (Ctrl-C, SIGTERM, crash) can be resumed from where it +stopped — the cancel message prints the exact command, e.g.: + +```console +$ cargo run -- import path/to/mail.mbox --resume 3 +``` + +Resume validates that the file hasn't changed (size, mtime, content +fingerprint); if it has, it safely falls back to re-parsing from the start, +which stays correct thanks to the dedupe. + +## Command line + +``` +gmail_stats [scan] [OPTIONS] scan Gmail over the API (OAuth; the default) +gmail_stats import [OPTIONS] import a Google Takeout mbox export + +--db SQLite database path [env: GMAIL_STATS_DB] [default: stats.db] +--credentials OAuth client secret (scan) [env: GMAIL_STATS_CREDENTIALS] [default: credentials.json] +--tokens OAuth token cache (scan) [env: GMAIL_STATS_TOKENS] [default: tokencache.json] +--resume import only: resume from the recorded byte offset +--quiet only errors and the final summary +--verbose per-message detail (prints every sender) +``` + +With `cargo run`, pass options after `--`: `cargo run -- import mail.mbox --db elsewhere.db`. + +By default output is a periodic progress line; `--verbose` restores the old +per-message `sender: ...` lines (be aware they print inbox metadata to the +terminal), and `--quiet` reduces output to errors and the final summary. + +Only one ingester (scan **or** import) can run against a database at a time: +a kernel lock on `.ingest.lock` makes a second one exit immediately with a +clear message. Each run is also recorded in an `ingest_runs` table (state, +message counts, progress, timestamps), which you can inspect with `sqlite3` +and which the web viewer will use for live progress in a later phase of +[#26](https://github.com/zbuc/gmail-stats/issues/26). + ## Configuration Two optional environment variables tune the Gmail API request rate for your -Google Cloud project's per-user quota: +Google Cloud project's per-user quota (scan mode only): * `GMAIL_STATS_FETCH_CONCURRENCY` — how many `messages.get` calls may be in flight at once (default: 8, minimum: 1). @@ -50,11 +151,9 @@ For example, to run gently against a project with a low quota: $ GMAIL_STATS_FETCH_CONCURRENCY=2 GMAIL_STATS_RATE_LIMIT_MS=100 cargo run ``` -This will trigger an OAuth flow which you launch in your browser, after which the access credentials are stored on disk in the local directory. -Please be aware that these are credentials that would allow anyone to read the contents of your email inbox, so you probably want to `rm tokencache.json` -after you're done. +## Viewing the results -When the script finishes running, you can view the statistics on senders in the DB: +When the run finishes, you can view the statistics on senders in the DB: ```console $ sqlite3 stats.db diff --git a/src/ingest.rs b/src/ingest.rs new file mode 100644 index 0000000..04feccb --- /dev/null +++ b/src/ingest.rs @@ -0,0 +1,679 @@ +//! Ingestion coordination shared by the Gmail API scan and the Takeout mbox +//! importer: schema migrations, the cross-process ingest lock, `ingest_runs` +//! bookkeeping, and the single DB writer task. + +use std::path::{Path, PathBuf}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use anyhow::Context; +use futures::TryStreamExt; +use sqlx::sqlite::SqliteConnectOptions; +use sqlx::{Connection, Row, Sqlite, SqliteConnection, SqliteExecutor, Transaction}; +use tokio::sync::mpsc; +use tokio::time::MissedTickBehavior; + +/// The schema version this binary knows how to produce, tracked in +/// `PRAGMA user_version`. The pre-versioning tables (`seen_mails`, `senders`) +/// are the version-0 baseline; migration 1 adds `ingest_runs`. +pub const SCHEMA_VERSION: i64 = 1; + +pub fn now_unix() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0) +} + +/// Bring the database up to [`SCHEMA_VERSION`]. +/// +/// The baseline (pre-versioning) schema is applied with guarded, idempotent +/// DDL exactly as before schema versioning existed, so an old database — or a +/// brand-new file — needs no manual steps. Versioned migrations then run one +/// by one inside transactions, each bumping `PRAGMA user_version`, and are +/// guarded so they never run twice and a database from a *newer* binary is +/// refused instead of mangled. +pub async fn migrate(conn: &mut SqliteConnection) -> anyhow::Result<()> { + // Baseline: create the original tables if they don't exist yet and enforce + // uniqueness of seen_mails.mail_id so that replaying a message (e.g. after + // a mid-run retry) is idempotent instead of double-counting. Pre-existing + // duplicate rows from earlier versions are collapsed before the unique + // index is created. + sqlx::query("CREATE TABLE IF NOT EXISTS seen_mails (mail_id string)") + .execute(&mut *conn) + .await?; + sqlx::query("CREATE TABLE IF NOT EXISTS senders (sender string, mails_sent int)") + .execute(&mut *conn) + .await?; + sqlx::query( + "DELETE FROM seen_mails WHERE rowid NOT IN \ + (SELECT MIN(rowid) FROM seen_mails GROUP BY mail_id)", + ) + .execute(&mut *conn) + .await?; + sqlx::query( + "CREATE UNIQUE INDEX IF NOT EXISTS seen_mails_mail_id_unique ON seen_mails (mail_id)", + ) + .execute(&mut *conn) + .await?; + + let version: i64 = sqlx::query("PRAGMA user_version") + .fetch_one(&mut *conn) + .await? + .try_get(0)?; + anyhow::ensure!( + version <= SCHEMA_VERSION, + "database schema is version {version}, newer than this binary supports \ + ({SCHEMA_VERSION}); upgrade gmail_stats" + ); + + if version < 1 { + // Migration 1: the ingest_runs table (one row per scan/import run). + let mut tx = conn.begin().await?; + sqlx::query( + "CREATE TABLE IF NOT EXISTS ingest_runs ( + run_id INTEGER PRIMARY KEY AUTOINCREMENT, + source TEXT NOT NULL, + state TEXT NOT NULL, + pid INTEGER, + started_at_unix INTEGER NOT NULL, + updated_at_unix INTEGER NOT NULL, + finished_at_unix INTEGER, + messages_seen INTEGER NOT NULL DEFAULT 0, + messages_new INTEGER NOT NULL DEFAULT 0, + total_estimate INTEGER, + bytes_total INTEGER, + bytes_done INTEGER, + resume_token TEXT, + mbox_path TEXT, + mbox_fingerprint TEXT, + error_kind TEXT, + error TEXT, + auth_url TEXT + )", + ) + .execute(&mut *tx) + .await?; + sqlx::query("PRAGMA user_version = 1") + .execute(&mut *tx) + .await?; + tx.commit().await?; + } + + Ok(()) +} + +/// Where the ingest lockfile for a given database lives: right next to it. +pub fn ingest_lock_path(db_path: &Path) -> PathBuf { + let mut name = db_path.as_os_str().to_owned(); + name.push(".ingest.lock"); + PathBuf::from(name) +} + +/// Held for the whole life of the process; the kernel releases the flock when +/// the process exits, however it dies, so there are no stale locks. +#[derive(Debug)] +pub struct IngestLock { + _file: std::fs::File, +} + +/// Take an exclusive kernel `flock` on the lockfile next to the database. +/// This is the authoritative "one ingester per database" gate: a second +/// scanner or importer (from a terminal or, later, spawned by the web viewer) +/// fails here promptly, before touching the database. +pub fn acquire_ingest_lock(db_path: &Path) -> anyhow::Result { + let lock_path = ingest_lock_path(db_path); + let file = std::fs::OpenOptions::new() + .create(true) + .truncate(false) + .write(true) + .open(&lock_path) + .with_context(|| format!("opening ingest lockfile {}", lock_path.display()))?; + match rustix::fs::flock(&file, rustix::fs::FlockOperation::NonBlockingLockExclusive) { + Ok(()) => Ok(IngestLock { _file: file }), + Err(rustix::io::Errno::WOULDBLOCK) => anyhow::bail!( + "another ingester is already running against {} \ + (could not acquire exclusive lock on {}); wait for it to finish or stop it first", + db_path.display(), + lock_path.display() + ), + Err(e) => Err(anyhow::Error::new(e).context(format!("locking {}", lock_path.display()))), + } +} + +/// Mark any run rows left open by a dead ingester as abandoned. Callers hold +/// the ingest lock, so an open row cannot belong to a live process: liveness +/// is judged by the flock, never by heartbeats or pids. +pub async fn abandon_stale_runs(conn: &mut SqliteConnection) -> anyhow::Result { + let now = now_unix(); + let result = sqlx::query( + "UPDATE ingest_runs SET state = 'abandoned', finished_at_unix = ?, updated_at_unix = ? \ + WHERE state NOT IN ('done', 'failed', 'cancelled', 'abandoned')", + ) + .bind(now) + .bind(now) + .execute(conn) + .await?; + Ok(result.rows_affected()) +} + +/// mbox-specific details recorded on an import run's row. +pub struct MboxRunInfo<'a> { + pub path: &'a Path, + pub fingerprint: &'a str, + pub bytes_total: u64, + pub start_offset: u64, +} + +/// Insert the row for a new run and return its run_id. +pub async fn create_run( + conn: &mut SqliteConnection, + source: &str, + mbox: Option>, +) -> anyhow::Result { + let now = now_unix(); + let (path, fingerprint, bytes_total, bytes_done) = match &mbox { + Some(info) => ( + Some(info.path.display().to_string()), + Some(info.fingerprint.to_string()), + Some(info.bytes_total as i64), + Some(info.start_offset as i64), + ), + None => (None, None, None, None), + }; + let result = sqlx::query( + "INSERT INTO ingest_runs \ + (source, state, pid, started_at_unix, updated_at_unix, \ + bytes_total, bytes_done, mbox_path, mbox_fingerprint) \ + VALUES (?, 'running', ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(source) + .bind(std::process::id() as i64) + .bind(now) + .bind(now) + .bind(bytes_total) + .bind(bytes_done) + .bind(path) + .bind(fingerprint) + .execute(conn) + .await?; + Ok(result.last_insert_rowid()) +} + +/// Write a run's terminal state. Opens its own short-lived connection: this +/// runs strictly after the writer task has drained and exited (we hold the +/// ingest lock, so no other writer exists), including on error paths where +/// the writer's connection is no longer available. +pub async fn finish_run( + options: &SqliteConnectOptions, + run_id: i64, + state: &str, + error_kind: Option<&str>, + error: Option<&str>, +) -> anyhow::Result<()> { + let mut conn = SqliteConnection::connect_with(options).await?; + let now = now_unix(); + sqlx::query( + "UPDATE ingest_runs SET state = ?, finished_at_unix = ?, updated_at_unix = ?, \ + error_kind = ?, error = ? WHERE run_id = ?", + ) + .bind(state) + .bind(now) + .bind(now) + .bind(error_kind) + .bind(error) + .bind(run_id) + .execute(&mut conn) + .await?; + conn.close().await.ok(); + Ok(()) +} + +/// A fetched/parsed message's result, sent to the single DB writer task. +pub struct SeenMail { + pub message_id: String, + pub sender: String, +} + +/// Everything that flows over the writer channel; all durable writes ride +/// through here so single-writer discipline holds. +pub enum WriteMsg { + Seen(SeenMail), + Progress { + messages_seen: u64, + /// Import only: committed parse offset (doubles as the resume point). + /// Sent after the Seen messages it covers — the channel is FIFO, so by + /// the time it is applied those messages are committed. + bytes_done: Option, + /// Scan only: last fully-processed page token. + resume_token: Option, + /// Scan only: the mailbox's total message count, once known. + total_estimate: Option, + }, +} + +/// The single DB writer: receives results over the channel and commits them +/// one transaction at a time on its own dedicated connection. +/// +/// The insert-and-count pair is idempotent per message id: the sender counter +/// (and the run's messages_new) is only incremented when the INSERT OR IGNORE +/// actually inserts the row, so a message that slips through a read-side seen +/// check twice is still counted exactly once. Between messages the writer +/// heartbeats the run row (~2s) so watchers can distinguish "slow" from +/// "stalled". +pub async fn db_writer( + mut conn: SqliteConnection, + mut rx: mpsc::Receiver, + run_id: i64, +) -> anyhow::Result<()> { + let mut heartbeat = tokio::time::interval(Duration::from_secs(2)); + heartbeat.set_missed_tick_behavior(MissedTickBehavior::Delay); + loop { + tokio::select! { + msg = rx.recv() => match msg { + None => break, + Some(WriteMsg::Seen(mail)) => { + let mut tx = conn.begin().await?; + let newly_seen = mark_seen(&mail.message_id, &mut *tx).await?; + if newly_seen { + increment_sender_mails(&mail.sender, &mut tx).await?; + sqlx::query( + "UPDATE ingest_runs SET messages_new = messages_new + 1, \ + updated_at_unix = ? WHERE run_id = ?", + ) + .bind(now_unix()) + .bind(run_id) + .execute(&mut *tx) + .await?; + } + tx.commit().await?; + } + Some(WriteMsg::Progress { messages_seen, bytes_done, resume_token, total_estimate }) => { + sqlx::query( + "UPDATE ingest_runs SET messages_seen = ?, \ + bytes_done = COALESCE(?, bytes_done), \ + resume_token = COALESCE(?, resume_token), \ + total_estimate = COALESCE(?, total_estimate), \ + updated_at_unix = ? WHERE run_id = ?", + ) + .bind(messages_seen as i64) + .bind(bytes_done.map(|b| b as i64)) + .bind(resume_token) + .bind(total_estimate) + .bind(now_unix()) + .bind(run_id) + .execute(&mut conn) + .await?; + } + }, + _ = heartbeat.tick() => { + sqlx::query("UPDATE ingest_runs SET updated_at_unix = ? WHERE run_id = ?") + .bind(now_unix()) + .bind(run_id) + .execute(&mut conn) + .await?; + } + } + } + Ok(()) +} + +pub async fn seen_mail( + message_id: &str, + executor: impl SqliteExecutor<'_>, +) -> anyhow::Result { + let mut res = sqlx::query("SELECT count(1) AS ct FROM seen_mails WHERE mail_id = ?") + .bind(message_id) + .fetch(executor); + while let Some(row) = res.try_next().await? { + let count: u32 = row.try_get("ct")?; + if count > 0 { + return Ok(true); + } + } + Ok(false) +} + +/// Record a message id as seen. Returns whether the id was newly inserted; +/// false means it was already present and must not be counted again. +async fn mark_seen(message_id: &str, executor: impl SqliteExecutor<'_>) -> anyhow::Result { + let result = sqlx::query("INSERT OR IGNORE INTO seen_mails (mail_id) VALUES (?)") + .bind(message_id) + .execute(executor) + .await?; + Ok(result.rows_affected() > 0) +} + +async fn increment_sender_mails( + sender: &str, + tx: &mut Transaction<'_, Sqlite>, +) -> anyhow::Result<()> { + let row = sqlx::query("SELECT mails_sent FROM senders WHERE sender = ?") + .bind(sender) + .fetch_optional(&mut **tx) + .await?; + if row.is_none() { + // no match + sqlx::query("INSERT INTO senders (sender, mails_sent) VALUES (?, 1)") + .bind(sender) + .execute(&mut **tx) + .await?; + + return Ok(()); + } + + let row = row.unwrap(); + let mut mails_sent = 0; + let count = row.try_get("mails_sent"); + + let count = count?; + + if count > 0 { + mails_sent = count; + } + + mails_sent += 1; + sqlx::query("UPDATE senders SET mails_sent = ? WHERE sender = ?") + .bind(mails_sent) + .bind(sender) + .execute(&mut **tx) + .await?; + + Ok(()) +} + +/// Fingerprint an mbox file for resume validation: size, mtime, and an +/// FNV-1a hash of the first 64 KiB. Cheap to compute, and catches both +/// truncation/replacement (size, prefix) and in-place edits (mtime) without +/// hashing a multi-GB file. +pub fn mbox_fingerprint(path: &Path) -> anyhow::Result { + use std::io::Read; + let meta = + std::fs::metadata(path).with_context(|| format!("cannot stat {}", path.display()))?; + let mtime = meta + .modified() + .ok() + .and_then(|t| t.duration_since(UNIX_EPOCH).ok()) + .map(|d| d.as_secs()) + .unwrap_or(0); + let mut file = + std::fs::File::open(path).with_context(|| format!("cannot open {}", path.display()))?; + let mut buf = vec![0u8; 64 * 1024]; + let mut filled = 0; + loop { + let n = file.read(&mut buf[filled..])?; + if n == 0 { + break; + } + filled += n; + if filled == buf.len() { + break; + } + } + Ok(format!( + "{}:{}:{:016x}", + meta.len(), + mtime, + fnv1a(&buf[..filled]) + )) +} + +fn fnv1a(bytes: &[u8]) -> u64 { + let mut hash: u64 = 0xcbf2_9ce4_8422_2325; + for &b in bytes { + hash ^= u64::from(b); + hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + } + hash +} + +#[cfg(test)] +mod tests { + use super::*; + use sqlx::sqlite::{SqliteJournalMode, SqliteSynchronous}; + use std::io::Write; + use std::str::FromStr; + + async fn memory_conn() -> SqliteConnection { + let options = SqliteConnectOptions::from_str("sqlite::memory:").unwrap(); + SqliteConnection::connect_with(&options).await.unwrap() + } + + fn file_options(path: &Path) -> SqliteConnectOptions { + SqliteConnectOptions::new() + .filename(path) + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal) + .synchronous(SqliteSynchronous::Normal) + .busy_timeout(Duration::from_secs(5)) + } + + async fn table_exists(conn: &mut SqliteConnection, name: &str) -> bool { + let count: i64 = + sqlx::query("SELECT count(1) AS ct FROM sqlite_master WHERE type='table' AND name=?") + .bind(name) + .fetch_one(conn) + .await + .unwrap() + .try_get("ct") + .unwrap(); + count > 0 + } + + async fn user_version(conn: &mut SqliteConnection) -> i64 { + sqlx::query("PRAGMA user_version") + .fetch_one(conn) + .await + .unwrap() + .try_get(0) + .unwrap() + } + + #[tokio::test] + async fn migrate_initializes_a_fresh_database() { + let mut conn = memory_conn().await; + migrate(&mut conn).await.unwrap(); + assert!(table_exists(&mut conn, "seen_mails").await); + assert!(table_exists(&mut conn, "senders").await); + assert!(table_exists(&mut conn, "ingest_runs").await); + assert_eq!(user_version(&mut conn).await, SCHEMA_VERSION); + } + + #[tokio::test] + async fn migrate_upgrades_a_version_0_database_with_existing_data() { + let mut conn = memory_conn().await; + // Simulate a pre-versioning database: baseline tables, some data, + // user_version still 0, no ingest_runs. + sqlx::query("CREATE TABLE seen_mails (mail_id string)") + .execute(&mut conn) + .await + .unwrap(); + sqlx::query("CREATE TABLE senders (sender string, mails_sent int)") + .execute(&mut conn) + .await + .unwrap(); + sqlx::query("INSERT INTO seen_mails (mail_id) VALUES ('m1'), ('m1'), ('m2')") + .execute(&mut conn) + .await + .unwrap(); + sqlx::query("INSERT INTO senders (sender, mails_sent) VALUES ('a@example.com', 7)") + .execute(&mut conn) + .await + .unwrap(); + assert_eq!(user_version(&mut conn).await, 0); + + migrate(&mut conn).await.unwrap(); + + assert_eq!(user_version(&mut conn).await, 1); + assert!(table_exists(&mut conn, "ingest_runs").await); + // Existing data survives; pre-index duplicates are collapsed. + let seen: i64 = sqlx::query("SELECT count(1) AS ct FROM seen_mails") + .fetch_one(&mut conn) + .await + .unwrap() + .try_get("ct") + .unwrap(); + assert_eq!(seen, 2); + let mails: i64 = sqlx::query("SELECT mails_sent FROM senders WHERE sender='a@example.com'") + .fetch_one(&mut conn) + .await + .unwrap() + .try_get("mails_sent") + .unwrap(); + assert_eq!(mails, 7); + + // Running again is a no-op. + migrate(&mut conn).await.unwrap(); + assert_eq!(user_version(&mut conn).await, 1); + } + + #[tokio::test] + async fn migrate_refuses_a_newer_database() { + let mut conn = memory_conn().await; + sqlx::query("PRAGMA user_version = 99") + .execute(&mut conn) + .await + .unwrap(); + let err = migrate(&mut conn).await.unwrap_err(); + assert!(err.to_string().contains("newer"), "unexpected error: {err}"); + } + + #[test] + fn second_ingester_is_refused_by_the_flock() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("stats.db"); + let first = acquire_ingest_lock(&db).expect("first lock"); + let err = acquire_ingest_lock(&db).expect_err("second lock must fail"); + assert!( + err.to_string() + .contains("another ingester is already running"), + "unexpected error: {err}" + ); + drop(first); + // Once the first ingester is gone the lock is free again. + acquire_ingest_lock(&db).expect("lock after release"); + } + + #[tokio::test] + async fn run_row_lifecycle() { + let dir = tempfile::tempdir().unwrap(); + let options = file_options(&dir.path().join("stats.db")); + let mut conn = SqliteConnection::connect_with(&options).await.unwrap(); + migrate(&mut conn).await.unwrap(); + + let run_id = create_run(&mut conn, "gmail_api", None).await.unwrap(); + let row = sqlx::query("SELECT state, pid, source FROM ingest_runs WHERE run_id = ?") + .bind(run_id) + .fetch_one(&mut conn) + .await + .unwrap(); + assert_eq!(row.try_get::("state").unwrap(), "running"); + assert_eq!(row.try_get::("source").unwrap(), "gmail_api"); + assert_eq!( + row.try_get::("pid").unwrap(), + std::process::id() as i64 + ); + + // A later ingester (holding the freed lock) marks the open row abandoned. + let abandoned = abandon_stale_runs(&mut conn).await.unwrap(); + assert_eq!(abandoned, 1); + let state: String = sqlx::query("SELECT state FROM ingest_runs WHERE run_id = ?") + .bind(run_id) + .fetch_one(&mut conn) + .await + .unwrap() + .try_get("state") + .unwrap(); + assert_eq!(state, "abandoned"); + + // Closed rows are left alone. + finish_run(&options, run_id, "failed", Some("io"), Some("boom")) + .await + .unwrap(); + assert_eq!(abandon_stale_runs(&mut conn).await.unwrap(), 0); + let row = sqlx::query( + "SELECT state, error_kind, error, finished_at_unix FROM ingest_runs WHERE run_id = ?", + ) + .bind(run_id) + .fetch_one(&mut conn) + .await + .unwrap(); + assert_eq!(row.try_get::("state").unwrap(), "failed"); + assert_eq!(row.try_get::("error_kind").unwrap(), "io"); + assert_eq!(row.try_get::("error").unwrap(), "boom"); + assert!(row.try_get::("finished_at_unix").unwrap() > 0); + } + + #[tokio::test] + async fn db_writer_dedupes_and_tracks_progress() { + let dir = tempfile::tempdir().unwrap(); + let options = file_options(&dir.path().join("stats.db")); + let mut conn = SqliteConnection::connect_with(&options).await.unwrap(); + migrate(&mut conn).await.unwrap(); + let run_id = create_run(&mut conn, "mbox", None).await.unwrap(); + + let (tx, rx) = mpsc::channel(16); + let writer = tokio::spawn(db_writer(conn, rx, run_id)); + for _ in 0..2 { + tx.send(WriteMsg::Seen(SeenMail { + message_id: "mid:".into(), + sender: "a@example.com".into(), + })) + .await + .unwrap(); + } + tx.send(WriteMsg::Seen(SeenMail { + message_id: "mid:".into(), + sender: "a@example.com".into(), + })) + .await + .unwrap(); + tx.send(WriteMsg::Progress { + messages_seen: 3, + bytes_done: Some(4096), + resume_token: None, + total_estimate: None, + }) + .await + .unwrap(); + drop(tx); + writer.await.unwrap().unwrap(); + + let mut conn = SqliteConnection::connect_with(&options).await.unwrap(); + let row = sqlx::query( + "SELECT messages_seen, messages_new, bytes_done FROM ingest_runs WHERE run_id = ?", + ) + .bind(run_id) + .fetch_one(&mut conn) + .await + .unwrap(); + assert_eq!(row.try_get::("messages_seen").unwrap(), 3); + assert_eq!(row.try_get::("messages_new").unwrap(), 2); + assert_eq!(row.try_get::("bytes_done").unwrap(), 4096); + let mails: i64 = sqlx::query("SELECT mails_sent FROM senders WHERE sender='a@example.com'") + .fetch_one(&mut conn) + .await + .unwrap() + .try_get("mails_sent") + .unwrap(); + assert_eq!(mails, 2); + } + + #[test] + fn fingerprint_changes_when_the_file_changes() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("mail.mbox"); + std::fs::write(&path, b"From a@example.com\nFrom: a@example.com\n\nbody\n").unwrap(); + let fp1 = mbox_fingerprint(&path).unwrap(); + let fp2 = mbox_fingerprint(&path).unwrap(); + assert_eq!(fp1, fp2, "fingerprint must be stable for an unchanged file"); + + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + file.write_all(b"From b@example.com\n\nmore\n").unwrap(); + drop(file); + let fp3 = mbox_fingerprint(&path).unwrap(); + assert_ne!(fp1, fp3, "fingerprint must change when the file grows"); + } +} diff --git a/src/main.rs b/src/main.rs index 2069809..05b6315 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,10 @@ +use std::fs::File; +use std::io::{BufReader, Read, Seek, SeekFrom}; +use std::path::{Path, PathBuf}; use std::str::FromStr; +use std::sync::atomic::{AtomicBool, AtomicU8, Ordering}; use std::sync::Arc; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use anyhow::Context; use futures::{stream, TryStreamExt}; @@ -9,11 +13,16 @@ use google_gmail1::{api::Scope, hyper_rustls, hyper_util, yup_oauth2, Gmail}; use lazy_static::lazy_static; use regex::Regex; use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous}; -use sqlx::{Connection, Pool, Row, Sqlite, SqliteConnection, SqliteExecutor, Transaction}; +use sqlx::{Connection, Pool, Row, Sqlite, SqliteConnection}; use tokio::sync::{mpsc, Mutex}; use tokio::task; use tokio::time::{Interval, MissedTickBehavior}; +mod ingest; +mod mbox; + +use ingest::{seen_mail, SeenMail, WriteMsg}; + type GmailHub = Gmail>; @@ -30,10 +39,167 @@ const DEFAULT_RATE_LIMIT_MS: u64 = 25; /// on a transient error (rate limit, 5xx, or network failure). const MAX_FETCH_RETRIES: u32 = 5; -/// A fetched message's result, sent to the single DB writer task. -struct SeenMail { - message_id: String, - sender: String, +const USAGE: &str = "\ +gmail_stats - collect per-sender mail counts into a local SQLite database + +USAGE: + gmail_stats [scan] [OPTIONS] scan Gmail over the API (OAuth; the default) + gmail_stats import [OPTIONS] import a Google Takeout mbox export + +OPTIONS: + --db SQLite database path + [env: GMAIL_STATS_DB] [default: stats.db] + --credentials OAuth client secret file (scan only) + [env: GMAIL_STATS_CREDENTIALS] [default: credentials.json] + --tokens OAuth token cache file (scan only) + [env: GMAIL_STATS_TOKENS] [default: tokencache.json] + --resume import only: resume a cancelled/failed/abandoned + import from its recorded byte offset; falls back to + a full re-parse if the file changed (dedupe keeps + the counts correct either way) + --quiet only errors and the final summary + --verbose per-message detail (prints every sender) + -h, --help show this help +"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +enum Verbosity { + Quiet = 0, + Normal = 1, + Verbose = 2, +} + +/// Process-wide output level, set once at startup. Normal prints periodic +/// progress; Verbose restores the historical per-message `sender: ...` lines +/// (off by default — they spray inbox metadata into the terminal). +static VERBOSITY: AtomicU8 = AtomicU8::new(Verbosity::Normal as u8); + +fn set_verbosity(v: Verbosity) { + VERBOSITY.store(v as u8, Ordering::Relaxed); +} + +fn verbose() -> bool { + VERBOSITY.load(Ordering::Relaxed) >= Verbosity::Verbose as u8 +} + +/// True unless --quiet: periodic progress and operational notes are printed. +fn chatty() -> bool { + VERBOSITY.load(Ordering::Relaxed) >= Verbosity::Normal as u8 +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum Mode { + Scan, + Import { path: PathBuf, resume: Option }, +} + +#[derive(Debug)] +struct Config { + db: PathBuf, + credentials: PathBuf, + tokens: PathBuf, + verbosity: Verbosity, + mode: Mode, +} + +#[derive(Debug)] +enum ParsedArgs { + Run(Config), + Help, +} + +fn parse_args( + args: impl IntoIterator, + env: impl Fn(&str) -> Option, +) -> Result { + let mut db = None; + let mut credentials = None; + let mut tokens = None; + let mut quiet = false; + let mut verbose = false; + let mut resume = None; + let mut positionals: Vec = Vec::new(); + + let mut it = args.into_iter(); + while let Some(arg) = it.next() { + let (flag, inline) = match arg.split_once('=') { + Some((f, v)) if f.starts_with("--") => (f.to_string(), Some(v.to_string())), + _ => (arg.clone(), None), + }; + let mut take_value = |name: &str| -> Result { + if let Some(v) = inline.clone() { + return Ok(v); + } + it.next().ok_or_else(|| format!("{name} requires a value")) + }; + match flag.as_str() { + "--db" => db = Some(take_value("--db")?), + "--credentials" => credentials = Some(take_value("--credentials")?), + "--tokens" => tokens = Some(take_value("--tokens")?), + "--resume" => { + let v = take_value("--resume")?; + resume = Some( + v.parse::() + .map_err(|_| format!("--resume expects a run id, got {v:?}"))?, + ); + } + "--quiet" => quiet = true, + "--verbose" => verbose = true, + "-h" | "--help" => return Ok(ParsedArgs::Help), + _ if flag.starts_with('-') => return Err(format!("unknown option {flag:?}")), + _ => positionals.push(arg), + } + } + + if quiet && verbose { + return Err("--quiet and --verbose are mutually exclusive".to_string()); + } + + let mode = match positionals.first().map(String::as_str) { + None | Some("scan") => { + if positionals.len() > 1 { + return Err(format!("unexpected argument {:?}", positionals[1])); + } + if resume.is_some() { + return Err("--resume is only valid with the import subcommand".to_string()); + } + Mode::Scan + } + Some("import") => { + let path = positionals + .get(1) + .ok_or_else(|| "import requires the path to an mbox file".to_string())?; + if positionals.len() > 2 { + return Err(format!("unexpected argument {:?}", positionals[2])); + } + Mode::Import { + path: PathBuf::from(path), + resume, + } + } + Some(other) => return Err(format!("unknown subcommand {other:?}")), + }; + + let pick = |flag: Option, env_name: &str, default: &str| -> PathBuf { + flag.map(PathBuf::from) + .or_else(|| env(env_name).map(PathBuf::from)) + .unwrap_or_else(|| PathBuf::from(default)) + }; + + Ok(ParsedArgs::Run(Config { + db: pick(db, "GMAIL_STATS_DB", "stats.db"), + credentials: pick(credentials, "GMAIL_STATS_CREDENTIALS", "credentials.json"), + tokens: pick(tokens, "GMAIL_STATS_TOKENS", "tokencache.json"), + verbosity: if quiet { + Verbosity::Quiet + } else if verbose { + Verbosity::Verbose + } else { + Verbosity::Normal + }, + mode, + })) } lazy_static! { @@ -42,8 +208,415 @@ lazy_static! { static ref EMAIL_RE_2: Regex = Regex::new(r"^([\w\-\.]+@([\w-]+\.)+[\w-]{2,4})$").unwrap(); } +fn db_connect_options(db: &Path) -> SqliteConnectOptions { + SqliteConnectOptions::new() + .filename(db) + // Create the database file on first run; migrate() creates the + // tables, so a fresh install needs no manual setup. + .create_if_missing(true) + // WAL mode allows the seen-mail reads (and the read-only web viewer) + // to proceed concurrently with the single writer task's commits. + .journal_mode(SqliteJournalMode::Wal) + // Synchronous mode is OK because a transaction may roll back during a + // crash, however both ingestion modes are re-runnable and idempotent. + .synchronous(SqliteSynchronous::Normal) + .busy_timeout(Duration::from_secs(5)) +} + +/// Flip a shared flag on SIGTERM/SIGINT so both ingestion modes can stop +/// fetching/parsing, drain the writer channel, persist resume state, and mark +/// their run row cancelled. A second signal aborts immediately. +fn spawn_signal_handler() -> Arc { + let cancel = Arc::new(AtomicBool::new(false)); + let flag = cancel.clone(); + tokio::spawn(async move { + use tokio::signal::unix::{signal, SignalKind}; + let mut sigterm = signal(SignalKind::terminate()).expect("installing SIGTERM handler"); + let mut sigint = signal(SignalKind::interrupt()).expect("installing SIGINT handler"); + tokio::select! { + _ = sigterm.recv() => {} + _ = sigint.recv() => {} + } + eprintln!( + "shutting down: draining queued writes and saving resume state \ + (signal again to abort immediately)" + ); + flag.store(true, Ordering::Relaxed); + tokio::select! { + _ = sigterm.recv() => {} + _ = sigint.recv() => {} + } + eprintln!("aborting"); + std::process::exit(130); + }); + cancel +} + #[tokio::main] async fn main() -> anyhow::Result<()> { + let cfg = match parse_args(std::env::args().skip(1), |name| std::env::var(name).ok()) { + Ok(ParsedArgs::Run(cfg)) => cfg, + Ok(ParsedArgs::Help) => { + print!("{USAGE}"); + return Ok(()); + } + Err(msg) => { + eprintln!("error: {msg}\n\n{USAGE}"); + std::process::exit(2); + } + }; + set_verbosity(cfg.verbosity); + + // Mutual exclusion comes first: never touch the database while another + // ingester (scan or import, terminal- or web-launched) is running. The + // kernel releases the flock whenever this process dies, so there are no + // stale locks to clean up. + let _ingest_lock = ingest::acquire_ingest_lock(&cfg.db)?; + + let options = db_connect_options(&cfg.db); + // The writer task owns the only connection that ever writes. + let mut writer_conn = SqliteConnection::connect_with(&options).await?; + ingest::migrate(&mut writer_conn).await?; + // We hold the lock, so any still-open run row belongs to a dead ingester. + let abandoned = ingest::abandon_stale_runs(&mut writer_conn).await?; + if abandoned > 0 && chatty() { + println!("marked {abandoned} interrupted ingest run(s) as abandoned"); + } + + let cancel = spawn_signal_handler(); + + match &cfg.mode { + Mode::Scan => run_scan(&cfg, &options, writer_conn, cancel).await, + Mode::Import { path, resume } => { + let summary = run_import(&options, writer_conn, path, *resume, cancel).await?; + report_import_summary(&summary, path); + Ok(()) + } + } +} + +// --------------------------------------------------------------------------- +// Takeout mbox import +// --------------------------------------------------------------------------- + +struct ImportOutcome { + messages_seen: u64, + skipped: u64, + cancelled: bool, +} + +#[derive(Debug)] +struct ImportSummary { + run_id: i64, + messages_seen: u64, + messages_new: u64, + skipped: u64, + cancelled: bool, +} + +fn report_import_summary(summary: &ImportSummary, mbox_path: &Path) { + if summary.cancelled { + println!( + "import cancelled cleanly after {} message(s) ({} new); resume with: \ + gmail_stats import {} --resume {}", + summary.messages_seen, + summary.messages_new, + mbox_path.display(), + summary.run_id + ); + } else if summary.skipped > 0 { + println!( + "import finished: {} message(s) parsed, {} new, \ + {} unparseable message(s) skipped (malformed or missing Message-ID)", + summary.messages_seen, summary.messages_new, summary.skipped + ); + } else { + println!( + "import finished: {} message(s) parsed, {} new", + summary.messages_seen, summary.messages_new + ); + } +} + +async fn run_import( + options: &SqliteConnectOptions, + mut writer_conn: SqliteConnection, + mbox_path: &Path, + resume: Option, + cancel: Arc, +) -> anyhow::Result { + let meta = std::fs::metadata(mbox_path) + .with_context(|| format!("cannot read mbox file {}", mbox_path.display()))?; + anyhow::ensure!( + meta.is_file(), + "{} is not a regular file", + mbox_path.display() + ); + { + // Sanity sniff: refuse files that clearly aren't mbox, so a wrong path + // can't slurp arbitrary file contents into the sender stats. + let mut file = + File::open(mbox_path).with_context(|| format!("opening {}", mbox_path.display()))?; + let mut magic = [0u8; 5]; + let looks_like_mbox = file.read_exact(&mut magic).is_ok() && &magic == b"From "; + anyhow::ensure!( + looks_like_mbox, + "{} does not look like an mbox file (expected it to start with a `From ` separator)", + mbox_path.display() + ); + } + + let fingerprint = ingest::mbox_fingerprint(mbox_path)?; + let start_offset = match resume { + None => 0, + Some(run_id) => resolve_resume_offset(&mut writer_conn, run_id, &fingerprint).await?, + }; + let run_id = ingest::create_run( + &mut writer_conn, + "mbox", + Some(ingest::MboxRunInfo { + path: mbox_path, + fingerprint: &fingerprint, + bytes_total: meta.len(), + start_offset, + }), + ) + .await?; + if chatty() { + if start_offset > 0 { + println!( + "importing {} from byte offset {start_offset} (run {run_id})", + mbox_path.display() + ); + } else { + println!( + "importing {} ({:.1} MiB, run {run_id})", + mbox_path.display(), + meta.len() as f64 / (1024.0 * 1024.0) + ); + } + } + + let (write_tx, write_rx) = mpsc::channel::(1024); + let writer_handle = task::spawn(ingest::db_writer(writer_conn, write_rx, run_id)); + let path = mbox_path.to_path_buf(); + let bytes_total = meta.len(); + let cancel_flag = cancel.clone(); + let producer = task::spawn_blocking(move || { + parse_and_send(&path, start_offset, bytes_total, write_tx, &cancel_flag) + }); + + let outcome = producer + .await + .map_err(|e| anyhow::anyhow!("import worker panicked: {e:?}"))?; + let writer_result = writer_handle + .await + .map_err(|e| anyhow::anyhow!("DB writer task panicked: {e:?}"))?; + + let outcome = match (outcome, writer_result) { + (Ok(outcome), Ok(())) => outcome, + (_, Err(writer_err)) => { + let msg = format!("{writer_err:#}"); + ingest::finish_run(options, run_id, "failed", Some("db"), Some(&msg)) + .await + .ok(); + return Err(writer_err.context("DB writer task failed")); + } + (Err(e), Ok(())) => { + let msg = format!("{e:#}"); + ingest::finish_run(options, run_id, "failed", Some("io"), Some(&msg)) + .await + .ok(); + return Err(e); + } + }; + + let state = if outcome.cancelled { + "cancelled" + } else { + "done" + }; + ingest::finish_run(options, run_id, state, None, None).await?; + + let mut conn = SqliteConnection::connect_with(options).await?; + let messages_new: i64 = sqlx::query("SELECT messages_new FROM ingest_runs WHERE run_id = ?") + .bind(run_id) + .fetch_one(&mut conn) + .await? + .try_get("messages_new")?; + conn.close().await.ok(); + + Ok(ImportSummary { + run_id, + messages_seen: outcome.messages_seen, + messages_new: messages_new as u64, + skipped: outcome.skipped, + cancelled: outcome.cancelled, + }) +} + +/// Decide where to start `--resume `: at the recorded byte offset if +/// the file still matches the run's fingerprint, else from the beginning +/// (which stays correct — every message id is deduped in seen_mails). +async fn resolve_resume_offset( + conn: &mut SqliteConnection, + run_id: i64, + fingerprint: &str, +) -> anyhow::Result { + let row = sqlx::query( + "SELECT source, mbox_fingerprint, bytes_done FROM ingest_runs WHERE run_id = ?", + ) + .bind(run_id) + .fetch_optional(&mut *conn) + .await?; + let Some(row) = row else { + anyhow::bail!("no ingest run {run_id} to resume"); + }; + let source: String = row.try_get("source")?; + anyhow::ensure!( + source == "mbox", + "run {run_id} is a {source} run, not an mbox import" + ); + let recorded: Option = row.try_get("mbox_fingerprint")?; + let bytes_done: Option = row.try_get("bytes_done")?; + match (recorded, bytes_done) { + (Some(recorded), Some(offset)) if recorded == fingerprint && offset > 0 => { + if chatty() { + println!("resuming run {run_id} from byte offset {offset}"); + } + Ok(offset as u64) + } + _ => { + if chatty() { + println!( + "mbox file changed since run {run_id} (or it recorded no progress); \ + re-importing from the start - dedupe keeps the counts correct" + ); + } + Ok(0) + } + } +} + +/// The blocking import producer: streams the mbox with the header-only parser +/// and feeds the writer channel. Memory stays bounded no matter the file size +/// (one BufReader block plus one capped line buffer). Runs on a blocking +/// thread; checks the cancel flag between messages. +fn parse_and_send( + path: &Path, + start_offset: u64, + bytes_total: u64, + write_tx: mpsc::Sender, + cancel: &AtomicBool, +) -> anyhow::Result { + let mut file = File::open(path).with_context(|| format!("opening {}", path.display()))?; + if start_offset > 0 { + file.seek(SeekFrom::Start(start_offset)).with_context(|| { + format!( + "seeking to resume offset {start_offset} in {}", + path.display() + ) + })?; + } + let mut reader = mbox::MboxReader::new(BufReader::with_capacity(64 * 1024, file), start_offset); + + let mut messages_seen = 0u64; + let mut missing_id = 0u64; + let mut committed_offset = start_offset; + let mut cancelled = false; + let mut last_progress = Instant::now(); + + loop { + if cancel.load(Ordering::Relaxed) { + cancelled = true; + break; + } + let msg = reader + .next_message() + .with_context(|| format!("reading {}", path.display()))?; + let Some(msg) = msg else { break }; + messages_seen += 1; + committed_offset = msg.end_offset; + match msg.message_id { + None => { + missing_id += 1; + if verbose() { + println!( + "skipping message without Message-ID (ends at byte {})", + msg.end_offset + ); + } + } + Some(mid) => { + // Same sender priority (From, then Return-Path) and cleanup as + // the API scan, so counts are consistent across sources. + let sender = + cleanup_sender(pick_sender(msg.from.as_deref(), msg.return_path.as_deref())); + if verbose() { + println!("sender: {:?}", sender); + } + // RFC Message-IDs live in the `mid:` namespace of seen_mails so + // they can never collide with Gmail API ids, and re-imports + // stay idempotent. This namespacing is the permanent keyspace + // rule for cross-source dedupe (issue #26). + write_tx + .blocking_send(WriteMsg::Seen(SeenMail { + message_id: format!("mid:{mid}"), + sender, + })) + .map_err(|_| anyhow::anyhow!("DB writer task closed unexpectedly"))?; + } + } + if last_progress.elapsed() >= Duration::from_secs(2) { + last_progress = Instant::now(); + // Sent after the Seen messages it covers (the channel is FIFO), so + // the recorded bytes_done only ever points past committed work. + write_tx + .blocking_send(WriteMsg::Progress { + messages_seen, + bytes_done: Some(committed_offset), + resume_token: None, + total_estimate: None, + }) + .map_err(|_| anyhow::anyhow!("DB writer task closed unexpectedly"))?; + if chatty() { + let pct = (committed_offset.min(bytes_total) * 100) + .checked_div(bytes_total) + .unwrap_or(100); + println!( + "imported {messages_seen} message(s) ({pct}% of {:.1} MiB)", + bytes_total as f64 / (1024.0 * 1024.0) + ); + } + } + } + + // Final progress (best-effort: the writer may already be gone on error + // paths, and the run row is finalized by the caller regardless). + let _ = write_tx.blocking_send(WriteMsg::Progress { + messages_seen, + bytes_done: Some(committed_offset), + resume_token: None, + total_estimate: None, + }); + + Ok(ImportOutcome { + messages_seen, + skipped: reader.skipped() + missing_id, + cancelled, + }) +} + +// --------------------------------------------------------------------------- +// Gmail API scan +// --------------------------------------------------------------------------- + +async fn run_scan( + cfg: &Config, + options: &SqliteConnectOptions, + mut writer_conn: SqliteConnection, + cancel: Arc, +) -> anyhow::Result<()> { let fetch_concurrency: usize = env_or("GMAIL_STATS_FETCH_CONCURRENCY", DEFAULT_FETCH_CONCURRENCY)?; anyhow::ensure!( @@ -56,31 +629,17 @@ async fn main() -> anyhow::Result<()> { "GMAIL_STATS_RATE_LIMIT_MS must be at least 1" ); - let options = SqliteConnectOptions::from_str("sqlite://./stats.db")? - // Create the database file on first run; init_schema below creates the - // tables, so a fresh install needs no manual setup. - .create_if_missing(true) - // WAL mode allows the seen-mail reads to proceed concurrently with the - // single writer task's commits. - .journal_mode(SqliteJournalMode::Wal) - // Synchronous mode is OK because a transaction may roll back during a crash, however - // all mail listings are re-fetched during each run. - .synchronous(SqliteSynchronous::Normal) - .busy_timeout(Duration::from_secs(5)); - // Small pool used only for the read-side seen-mail checks; all writes go - // through the single writer task below, which eliminates writer-vs-writer + // through the single writer task, which eliminates writer-vs-writer // deadlocks entirely. let pool = SqlitePoolOptions::new() .max_connections(fetch_concurrency as u32) .connect_with(options.clone()) .await?; - // The writer task owns the only connection that ever writes. - let mut writer_conn = SqliteConnection::connect_with(&options).await?; - init_schema(&mut writer_conn).await?; - let (mut write_tx, write_rx) = mpsc::channel::(100); - let mut writer_handle = task::spawn(db_writer(writer_conn, write_rx)); + let run_id = ingest::create_run(&mut writer_conn, "gmail_api", None).await?; + let (mut write_tx, write_rx) = mpsc::channel::(100); + let mut writer_handle = task::spawn(ingest::db_writer(writer_conn, write_rx, run_id)); // Simple client-side rate limiter: each API call waits for the next tick. let mut interval = tokio::time::interval(Duration::from_millis(rate_limit_ms)); @@ -88,22 +647,59 @@ async fn main() -> anyhow::Result<()> { let rate_limiter = Arc::new(Mutex::new(interval)); // Read application OAuth secret from a file. - let secret = yup_oauth2::read_application_secret("credentials.json") - .await - .expect("credentials.json"); + let secret = match yup_oauth2::read_application_secret(&cfg.credentials).await { + Ok(secret) => secret, + Err(e) => { + drop(write_tx); + let _ = writer_handle.await; + let msg = format!( + "reading OAuth client secret from {}: {e}", + cfg.credentials.display() + ); + ingest::finish_run( + options, + run_id, + "failed", + Some("missing_credentials"), + Some(&msg), + ) + .await + .ok(); + return Err(anyhow::Error::new(e).context(format!( + "reading OAuth client secret from {} (see the README for OAuth setup, or \ + use `gmail_stats import` with a Google Takeout mbox export instead)", + cfg.credentials.display() + ))); + } + }; // Create an authenticator that uses an InstalledFlow to authenticate. The - // authentication tokens are persisted to a file named tokencache.json. The - // authenticator takes care of caching tokens to disk and refreshing tokens once - // they've expired. - let auth = yup_oauth2::InstalledFlowAuthenticator::builder( + // authentication tokens are persisted to the token cache file, so they are + // cached to disk and refreshed once they've expired. + let auth = match yup_oauth2::InstalledFlowAuthenticator::builder( secret, yup_oauth2::InstalledFlowReturnMethod::HTTPRedirect, ) - .persist_tokens_to_disk("tokencache.json") + .persist_tokens_to_disk(cfg.tokens.clone()) .build() .await - .unwrap(); + { + Ok(auth) => auth, + Err(e) => { + drop(write_tx); + let _ = writer_handle.await; + ingest::finish_run( + options, + run_id, + "failed", + Some("auth_required"), + Some(&e.to_string()), + ) + .await + .ok(); + return Err(anyhow::Error::new(e).context("building the OAuth authenticator")); + } + }; let hub = Gmail::new( hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::new()).build( @@ -116,6 +712,27 @@ async fn main() -> anyhow::Result<()> { auth, ); + // Best-effort total for progress display; ignored on failure (the scan + // itself will surface any real auth/API problem). + if let Ok((_, profile)) = hub + .users() + .get_profile("me") + .add_scope(Scope::Readonly) + .doit() + .await + { + if let Some(total) = profile.messages_total { + let _ = write_tx + .send(WriteMsg::Progress { + messages_seen: 0, + bytes_done: None, + resume_token: None, + total_estimate: Some(i64::from(total)), + }) + .await; + } + } + // Retry a failed run a few times, backing off exponentially between // attempts (see the sleep at the bottom of the loop body). Only transient // failures (rate limits and server errors that outlasted the per-call @@ -126,26 +743,36 @@ async fn main() -> anyhow::Result<()> { // never accumulate into an abort — only repeated failures with no // progress in between exhaust the budget. let mut retries = 0; - let mut pages_done: u64 = 0; - // The page token to resume from, advanced only once a page has been - // fully processed: a retry re-lists the page that failed instead of - // restarting the whole mailbox scan from page one. - let mut resume_token: Option = None; + let mut progress = ScanProgress { + resume_token: None, + pages_done: 0, + messages_listed: 0, + }; loop { - let pages_before = pages_done; - match work( - &pool, - &hub, - &write_tx, - &rate_limiter, - fetch_concurrency, - &mut resume_token, - &mut pages_done, + let pages_before = progress.pages_done; + let attempt = work( + &ScanCtx { + pool: &pool, + hub: &hub, + write_tx: &write_tx, + rate_limiter: &rate_limiter, + fetch_concurrency, + cancel: cancel.as_ref(), + }, + &mut progress, ) - .await - { + .await; + match attempt { Ok(()) => break, Err(e) => { + if cancel.load(Ordering::Relaxed) { + // Shutdown was requested; the error is almost certainly + // fallout from stopping mid-flight. Save state and leave. + if chatty() { + println!("stopping scan for shutdown ({e:#})"); + } + break; + } // Quiesce the writer before deciding what to do next: close // the channel and wait for the writer to drain and commit // every queued result. Without this barrier a retry's @@ -170,7 +797,7 @@ async fn main() -> anyhow::Result<()> { // processed at least one page) before failing, the previous // errors were intermittent, not persistent — start the budget // over so MAX_RETRIES bounds *consecutive* failures. - if pages_done > pages_before { + if progress.pages_done > pages_before { retries = 0; } retries += 1; @@ -181,22 +808,32 @@ async fn main() -> anyhow::Result<()> { } else { e.context("giving up: error is not retryable") }; - return match writer_err { - None => Err(final_err), + let final_err = match writer_err { + None => final_err, Some(writer_err) => { - Err(final_err.context(format!("DB writer task failed: {writer_err:?}"))) + final_err.context(format!("DB writer task failed: {writer_err:?}")) } }; + ingest::finish_run( + options, + run_id, + "failed", + Some(classify_scan_error(&final_err)), + Some(&format!("{final_err:#}")), + ) + .await + .ok(); + return Err(final_err); } if let Some(writer_err) = &writer_err { println!("DB writer task failed, restarting it: {writer_err:?}"); } // Spawn a fresh writer so the retry has a chance of working. - let writer_conn = SqliteConnection::connect_with(&options).await?; - let (tx, rx) = mpsc::channel::(100); + let writer_conn = SqliteConnection::connect_with(options).await?; + let (tx, rx) = mpsc::channel::(100); write_tx = tx; - writer_handle = task::spawn(db_writer(writer_conn, rx)); + writer_handle = task::spawn(ingest::db_writer(writer_conn, rx, run_id)); // Back off before retrying: errors that reach this loop are // typically sustained conditions (e.g. quota exhaustion that @@ -207,7 +844,7 @@ async fn main() -> anyhow::Result<()> { println!( "Transient error (attempt {retries}/{MAX_RETRIES}), retrying in {delay:?}: {e:?}" ); - tokio::time::sleep(delay).await; + sleep_cancellable(delay, &cancel).await; } } } @@ -215,7 +852,45 @@ async fn main() -> anyhow::Result<()> { // Close the channel so the writer task drains its queue and exits, then // surface any error it hit. drop(write_tx); - writer_handle.await??; + match writer_handle.await { + Err(join_err) => { + let msg = format!("DB writer task panicked: {join_err:?}"); + ingest::finish_run(options, run_id, "failed", Some("db"), Some(&msg)) + .await + .ok(); + anyhow::bail!("{msg}"); + } + Ok(Err(writer_err)) => { + ingest::finish_run( + options, + run_id, + "failed", + Some("db"), + Some(&format!("{writer_err:#}")), + ) + .await + .ok(); + return Err(writer_err.context("DB writer task failed")); + } + Ok(Ok(())) => {} + } + + if cancel.load(Ordering::Relaxed) { + ingest::finish_run(options, run_id, "cancelled", None, None).await?; + println!("scan cancelled cleanly; run {run_id} keeps the resume state"); + } else { + ingest::finish_run(options, run_id, "done", None, None).await?; + let row = + sqlx::query("SELECT messages_seen, messages_new FROM ingest_runs WHERE run_id = ?") + .bind(run_id) + .fetch_one(&pool) + .await?; + println!( + "scan finished: {} message(s) listed, {} new", + row.try_get::("messages_seen")?, + row.try_get::("messages_new")? + ); + } Ok(()) } @@ -234,6 +909,18 @@ fn backoff_delay(attempt: u32) -> Duration { base + Duration::from_millis(jitter_ms) } +/// Sleep that wakes early (within ~250ms) once the cancel flag is set. +async fn sleep_cancellable(duration: Duration, cancel: &AtomicBool) { + let deadline = Instant::now() + duration; + while !cancel.load(Ordering::Relaxed) { + let now = Instant::now(); + if now >= deadline { + break; + } + tokio::time::sleep((deadline - now).min(Duration::from_millis(250))).await; + } +} + /// Read a value from an environment variable, falling back to a default when /// the variable is unset and failing loudly when it is set but unparseable. fn env_or(name: &str, default: T) -> anyhow::Result @@ -249,86 +936,70 @@ where } } -/// Create the tables if they don't exist yet and enforce uniqueness of -/// seen_mails.mail_id so that replaying a message (e.g. after a mid-run retry) -/// is idempotent instead of double-counting. Pre-existing duplicate rows from -/// earlier versions are collapsed before the unique index is created. -async fn init_schema(conn: &mut SqliteConnection) -> anyhow::Result<()> { - sqlx::query("CREATE TABLE IF NOT EXISTS seen_mails (mail_id string)") - .execute(&mut *conn) - .await?; - sqlx::query("CREATE TABLE IF NOT EXISTS senders (sender string, mails_sent int)") - .execute(&mut *conn) - .await?; - sqlx::query( - "DELETE FROM seen_mails WHERE rowid NOT IN \ - (SELECT MIN(rowid) FROM seen_mails GROUP BY mail_id)", - ) - .execute(&mut *conn) - .await?; - sqlx::query( - "CREATE UNIQUE INDEX IF NOT EXISTS seen_mails_mail_id_unique ON seen_mails (mail_id)", - ) - .execute(&mut *conn) - .await?; - Ok(()) +struct ScanCtx<'a> { + pool: &'a Pool, + hub: &'a GmailHub, + write_tx: &'a mpsc::Sender, + rate_limiter: &'a Arc>, + fetch_concurrency: usize, + cancel: &'a AtomicBool, } -/// The single DB writer: receives fetched results over the channel and commits -/// them one transaction at a time on its own dedicated connection. -/// -/// The insert-and-count pair is idempotent per message id: the sender counter -/// is only incremented when the INSERT OR IGNORE actually inserts the row, so -/// a message that slips through the read-side seen check twice (its first -/// result still queued and uncommitted when it is re-listed) is still counted -/// exactly once. -async fn db_writer( - mut conn: SqliteConnection, - mut rx: mpsc::Receiver, -) -> anyhow::Result<()> { - while let Some(mail) = rx.recv().await { - let mut tx = conn.begin().await?; - let newly_seen = mark_seen(&mail.message_id, &mut *tx).await?; - if newly_seen { - increment_sender_mails(&mail.sender, &mut tx).await?; - } - tx.commit().await?; - } - Ok(()) +/// Progress that survives across retry attempts. resume_token is the page +/// token to resume from, advanced only once a page has been fully processed: +/// a retry re-lists the page that failed instead of restarting the whole +/// mailbox scan from page one. It is also persisted to the run row after +/// every page, so a future scan could pick it up best-effort. +struct ScanProgress { + resume_token: Option, + pages_done: u64, + messages_listed: u64, } -async fn work( - pool: &Pool, - hub: &GmailHub, - write_tx: &mpsc::Sender, - rate_limiter: &Arc>, - fetch_concurrency: usize, - resume_token: &mut Option, - pages_done: &mut u64, -) -> anyhow::Result<()> { +async fn work(ctx: &ScanCtx<'_>, progress: &mut ScanProgress) -> anyhow::Result<()> { // Fetch 500 messages at a time, starting from the last page we fully - // processed (so a retry doesn't re-list the whole mailbox from page one). + // processed. loop { - let result = list_messages(hub, resume_token.as_deref(), rate_limiter).await?; + if ctx.cancel.load(Ordering::Relaxed) { + return Ok(()); + } + let result = + list_messages(ctx.hub, progress.resume_token.as_deref(), ctx.rate_limiter).await?; let next_page_token = result.next_page_token; - parse_messages( - pool, - result.messages.unwrap_or_default(), - hub, - write_tx, - rate_limiter, - fetch_concurrency, - ) - .await?; + let messages = result.messages.unwrap_or_default(); + let listed = messages.len() as u64; + parse_messages(ctx, messages).await?; + + if ctx.cancel.load(Ordering::Relaxed) { + // Fetches for this page were (partially) skipped for shutdown; + // don't advance the resume point past them. + return Ok(()); + } // Completing a page is forward progress: only now advance the resume // point and the page counter the caller uses to reset its retry // budget, so a long scan isn't aborted by transient errors // accumulated across otherwise-successful attempts. - *pages_done += 1; - *resume_token = next_page_token; - if resume_token.is_none() { + progress.pages_done += 1; + progress.messages_listed += listed; + progress.resume_token = next_page_token.clone(); + ctx.write_tx + .send(WriteMsg::Progress { + messages_seen: progress.messages_listed, + bytes_done: None, + resume_token: next_page_token, + total_estimate: None, + }) + .await + .map_err(|_| anyhow::anyhow!("DB writer task closed unexpectedly"))?; + if chatty() { + println!( + "page {} done ({} messages listed)", + progress.pages_done, progress.messages_listed + ); + } + if progress.resume_token.is_none() { return Ok(()); } } @@ -362,10 +1033,12 @@ async fn list_messages( Ok((_, response)) => return Ok(response), Err(ref e) if attempts < MAX_FETCH_RETRIES && is_transient_gmail(e) => { attempts += 1; - println!( - "transient error listing messages (attempt {}), backing off for {:?}: {}", - attempts, backoff, e - ); + if chatty() { + println!( + "transient error listing messages (attempt {}), backing off for {:?}: {}", + attempts, backoff, e + ); + } tokio::time::sleep(backoff).await; backoff *= 2; } @@ -374,33 +1047,31 @@ async fn list_messages( } } -async fn parse_messages( - pool: &Pool, - messages: Vec, - hub: &GmailHub, - write_tx: &mpsc::Sender, - rate_limiter: &Arc>, - fetch_concurrency: usize, -) -> anyhow::Result<()> { +async fn parse_messages(ctx: &ScanCtx<'_>, messages: Vec) -> anyhow::Result<()> { // Fetch each individual message concurrently (bounded), then hand the // result to the writer task to increment the counter for the sender. stream::iter(messages.into_iter().map(Ok::<_, anyhow::Error>)) - .try_for_each_concurrent(fetch_concurrency, |message_meta| { - let hub = hub.clone(); - let write_tx = write_tx.clone(); - let rate_limiter = rate_limiter.clone(); + .try_for_each_concurrent(ctx.fetch_concurrency, |message_meta| { + let hub = ctx.hub.clone(); + let write_tx = ctx.write_tx.clone(); + let rate_limiter = ctx.rate_limiter.clone(); async move { + if ctx.cancel.load(Ordering::Relaxed) { + return Ok(()); + } let message_id = message_meta.id.expect("message missing id"); - if seen_mail(&message_id, pool).await? { + if seen_mail(&message_id, ctx.pool).await? { return Ok(()); } let message = fetch_message(&hub, &message_id, &rate_limiter).await?; let sender = cleanup_sender(get_sender(&message)?); - println!("sender: {:?}", sender); + if verbose() { + println!("sender: {:?}", sender); + } write_tx - .send(SeenMail { message_id, sender }) + .send(WriteMsg::Seen(SeenMail { message_id, sender })) .await .map_err(|_| anyhow::anyhow!("DB writer task closed unexpectedly"))?; @@ -434,10 +1105,12 @@ async fn fetch_message( Ok((_, message)) => return Ok(message), Err(ref e) if attempts < MAX_FETCH_RETRIES && is_transient_gmail(e) => { attempts += 1; - println!( - "transient error fetching {} (attempt {}), backing off for {:?}: {}", - message_id, attempts, backoff, e - ); + if chatty() { + println!( + "transient error fetching {} (attempt {}), backing off for {:?}: {}", + message_id, attempts, backoff, e + ); + } tokio::time::sleep(backoff).await; backoff *= 2; } @@ -534,65 +1207,41 @@ fn is_transient_gmail(err: &google_gmail1::Error) -> bool { } } -async fn seen_mail(message_id: &str, executor: impl SqliteExecutor<'_>) -> anyhow::Result { - let mut res = sqlx::query("SELECT count(1) AS ct FROM seen_mails WHERE mail_id = ?") - .bind(message_id) - .fetch(executor); - while let Some(row) = res.try_next().await? { - let count: u32 = row.try_get("ct")?; - if count > 0 { - return Ok(true); +/// Coarse error classification recorded on failed scan runs so a watcher (or +/// the future web UI) can key remediation off it. +fn classify_scan_error(err: &anyhow::Error) -> &'static str { + if format!("{err:#}").contains("policy_enforced") { + return "policy_enforced"; + } + for cause in err.chain() { + if let Some(io_err) = cause.downcast_ref::() { + if io_err.kind() == std::io::ErrorKind::NotFound { + return "missing_credentials"; + } + return "io"; + } + if cause.downcast_ref::().is_some() { + return "db"; + } + if let Some(gmail_err) = cause.downcast_ref::() { + return if is_transient_gmail(gmail_err) { + "rate_limited" + } else { + "gmail_api" + }; } } - Ok(false) + "other" } -/// Record a message id as seen. Returns whether the id was newly inserted; -/// false means it was already present and must not be counted again. -async fn mark_seen(message_id: &str, executor: impl SqliteExecutor<'_>) -> anyhow::Result { - let result = sqlx::query("INSERT OR IGNORE INTO seen_mails (mail_id) VALUES (?)") - .bind(message_id) - .execute(executor) - .await?; - Ok(result.rows_affected() > 0) -} - -async fn increment_sender_mails( - sender: &str, - tx: &mut Transaction<'_, Sqlite>, -) -> anyhow::Result<()> { - let row = sqlx::query("SELECT mails_sent FROM senders WHERE sender = ?") - .bind(sender) - .fetch_optional(&mut **tx) - .await?; - if row.is_none() { - // no match - sqlx::query("INSERT INTO senders (sender, mails_sent) VALUES (?, 1)") - .bind(sender) - .execute(&mut **tx) - .await?; - - return Ok(()); - } - - let row = row.unwrap(); - let mut mails_sent = 0; - let count = row.try_get("mails_sent"); +// --------------------------------------------------------------------------- +// Sender extraction, shared by both ingestion modes +// --------------------------------------------------------------------------- - let count = count?; - - if count > 0 { - mails_sent = count; - } - - mails_sent += 1; - sqlx::query("UPDATE senders SET mails_sent = ? WHERE sender = ?") - .bind(mails_sent) - .bind(sender) - .execute(&mut **tx) - .await?; - - Ok(()) +/// The shared sender priority: From first, then Return-Path. Used by both the +/// API scan (get_sender) and the mbox importer so counts are consistent. +fn pick_sender(from: Option<&str>, return_path: Option<&str>) -> String { + from.or(return_path).unwrap_or_default().to_string() } // Attempt to extract a formatted email address, or just return the original value @@ -612,28 +1261,28 @@ fn cleanup_sender(sender: String) -> String { } fn get_sender(message: &Message) -> anyhow::Result { - // Header names are case-insensitive (RFC 5322); check candidates in priority order. - const CANDIDATES: [&str; 2] = ["from", "return-path"]; + // Header names are case-insensitive (RFC 5322); check candidates in + // priority order via the shared pick_sender. let headers = message .payload .as_ref() .and_then(|p| p.headers.as_deref()) .unwrap_or(&[]); - - for candidate in CANDIDATES { - if let Some(value) = headers.iter().find_map(|header| { + let find = |name: &str| { + headers.iter().find_map(|header| { header .name .as_deref() - .filter(|name| name.eq_ignore_ascii_case(candidate)) + .filter(|n| n.eq_ignore_ascii_case(name)) .and(header.value.as_deref()) - }) { - return Ok(value.to_string()); - } - } + }) + }; - println!("weird email without from header: {:?}", message.id); - Ok(String::new()) + let sender = pick_sender(find("from"), find("return-path")); + if sender.is_empty() && verbose() { + println!("weird email without from header: {:?}", message.id); + } + Ok(sender) } #[cfg(test)] @@ -778,7 +1427,17 @@ mod tests { assert_eq!(cleanup_sender(String::new()), ""); } - // --- get_sender --- + // --- pick_sender / get_sender --- + + #[test] + fn pick_sender_prefers_from_over_return_path() { + assert_eq!( + pick_sender(Some("a@example.com"), Some("b@example.com")), + "a@example.com" + ); + assert_eq!(pick_sender(None, Some("b@example.com")), "b@example.com"); + assert_eq!(pick_sender(None, None), ""); + } #[test] fn get_sender_prefers_from_over_return_path() { @@ -824,98 +1483,325 @@ mod tests { assert_eq!(get_sender(&message).unwrap(), ""); } - // --- SQLite counting logic --- + // --- argument parsing --- - async fn test_pool() -> Pool { - let options = SqliteConnectOptions::from_str("sqlite::memory:").unwrap(); - // A single connection so every query sees the same in-memory database. - let pool = SqlitePoolOptions::new() - .max_connections(1) - .connect_with(options) - .await - .unwrap(); + fn no_env(_: &str) -> Option { + None + } - // Run the production DDL so the tests exercise the same schema the - // real database has, including the `string` type name (SQLite gives - // such columns NUMERIC affinity rather than TEXT) and the unique - // index on seen_mails.mail_id that makes mark_seen idempotent. - let mut conn = pool.acquire().await.unwrap(); - init_schema(&mut conn).await.unwrap(); - drop(conn); + fn run_cfg(args: &[&str], env: impl Fn(&str) -> Option) -> Config { + match parse_args(args.iter().map(|s| s.to_string()), env) { + Ok(ParsedArgs::Run(cfg)) => cfg, + other => panic!("expected a runnable config, got {:?}", other.err()), + } + } - pool + #[test] + fn bare_invocation_is_the_scan_with_default_paths() { + let cfg = run_cfg(&[], no_env); + assert_eq!(cfg.mode, Mode::Scan); + assert_eq!(cfg.db, PathBuf::from("stats.db")); + assert_eq!(cfg.credentials, PathBuf::from("credentials.json")); + assert_eq!(cfg.tokens, PathBuf::from("tokencache.json")); + assert_eq!(cfg.verbosity, Verbosity::Normal); } - #[tokio::test] - async fn seen_mail_dedup() { - let pool = test_pool().await; + #[test] + fn explicit_scan_subcommand_is_equivalent() { + assert_eq!(run_cfg(&["scan"], no_env).mode, Mode::Scan); + } + + #[test] + fn import_subcommand_takes_a_path_and_resume() { + let cfg = run_cfg(&["import", "/tmp/All mail.mbox", "--resume", "7"], no_env); + assert_eq!( + cfg.mode, + Mode::Import { + path: PathBuf::from("/tmp/All mail.mbox"), + resume: Some(7) + } + ); + } - assert!(!seen_mail("mail-1", &pool).await.unwrap()); + #[test] + fn path_flags_override_env_which_overrides_defaults() { + let env = |name: &str| match name { + "GMAIL_STATS_DB" => Some("/env/stats.db".to_string()), + "GMAIL_STATS_CREDENTIALS" => Some("/env/creds.json".to_string()), + _ => None, + }; + let cfg = run_cfg(&["--db", "/flag/stats.db"], env); + assert_eq!(cfg.db, PathBuf::from("/flag/stats.db")); + assert_eq!(cfg.credentials, PathBuf::from("/env/creds.json")); + assert_eq!(cfg.tokens, PathBuf::from("tokencache.json")); + } - // First sighting inserts the row and reports it as newly seen. - assert!(mark_seen("mail-1", &pool).await.unwrap()); + #[test] + fn inline_flag_values_are_accepted() { + let cfg = run_cfg(&["--db=inline.db", "--quiet"], no_env); + assert_eq!(cfg.db, PathBuf::from("inline.db")); + assert_eq!(cfg.verbosity, Verbosity::Quiet); + } - assert!(seen_mail("mail-1", &pool).await.unwrap()); - // Other ids remain unseen. - assert!(!seen_mail("mail-2", &pool).await.unwrap()); + #[test] + fn verbose_flag_selects_verbose_output() { + assert_eq!( + run_cfg(&["--verbose"], no_env).verbosity, + Verbosity::Verbose + ); + } - // Replaying the same id is ignored and reports it as already seen. - assert!(!mark_seen("mail-1", &pool).await.unwrap()); - let rows: i64 = sqlx::query("SELECT count(1) AS ct FROM seen_mails") - .fetch_one(&pool) + #[test] + fn invalid_invocations_are_rejected() { + for args in [ + &["--quiet", "--verbose"][..], + &["import"][..], + &["--resume", "3"][..], + &["--resume", "x", "import", "f.mbox"][..], + &["frobnicate"][..], + &["--frobnicate"][..], + &["scan", "extra"][..], + &["import", "a.mbox", "extra"][..], + &["--db"][..], + ] { + let result = parse_args(args.iter().map(|s| s.to_string()), no_env); + assert!(result.is_err(), "expected {args:?} to be rejected"); + } + } + + #[test] + fn help_flag_wins() { + assert!(matches!( + parse_args(["--help".to_string()], no_env), + Ok(ParsedArgs::Help) + )); + assert!(matches!( + parse_args(["import".to_string(), "-h".to_string()], no_env), + Ok(ParsedArgs::Help) + )); + } + + // --- end-to-end import --- + + /// Three messages: normal, Return-Path-only with an mboxrd-escaped body + /// line, and one without a Message-ID (skipped). + const IMPORT_FIXTURE: &str = "From a@example.com Thu Jan 1 00:00:00 2020\n\ + From: Alice \n\ + Message-ID: \n\ + Subject: one\n\ + \n\ + body one\n\ + >From an escaped body line\n\ + From b@example.com Thu Jan 1 00:00:00 2020\n\ + Return-Path: \n\ + Message-ID: \n\ + \n\ + body two\n\ + From c@example.com Thu Jan 1 00:00:00 2020\n\ + From: Alice \n\ + Subject: no message id\n\ + \n\ + body three\n"; + + async fn migrated_conn(options: &SqliteConnectOptions) -> SqliteConnection { + let mut conn = SqliteConnection::connect_with(options).await.unwrap(); + ingest::migrate(&mut conn).await.unwrap(); + conn + } + + async fn sender_count(options: &SqliteConnectOptions, sender: &str) -> Option { + let mut conn = SqliteConnection::connect_with(options).await.unwrap(); + sqlx::query("SELECT mails_sent FROM senders WHERE sender = ?") + .bind(sender) + .fetch_optional(&mut conn) .await .unwrap() - .try_get("ct") - .unwrap(); - assert_eq!(rows, 1); + .map(|row| row.try_get("mails_sent").unwrap()) + } + + fn unset_cancel() -> Arc { + Arc::new(AtomicBool::new(false)) } #[tokio::test] - async fn increment_sender_mails_inserts_then_increments() { - let pool = test_pool().await; - - // First mail from a sender inserts a row with mails_sent = 1. - let mut tx = pool.begin().await.unwrap(); - increment_sender_mails( - &cleanup_sender("Jane ".to_string()), - &mut tx, + async fn import_populates_the_db_and_reimport_adds_zero_counts() { + let dir = tempfile::tempdir().unwrap(); + let mbox = dir.path().join("mail.mbox"); + std::fs::write(&mbox, IMPORT_FIXTURE).unwrap(); + let options = db_connect_options(&dir.path().join("stats.db")); + + let conn = migrated_conn(&options).await; + let s1 = run_import(&options, conn, &mbox, None, unset_cancel()) + .await + .unwrap(); + assert_eq!(s1.messages_seen, 3); + assert_eq!(s1.messages_new, 2); + assert_eq!(s1.skipped, 1); + assert!(!s1.cancelled); + assert_eq!(sender_count(&options, "a@example.com").await, Some(1)); + assert_eq!(sender_count(&options, "b@example.com").await, Some(1)); + + // Message-IDs are namespaced so they can never collide with Gmail ids. + let mut conn = SqliteConnection::connect_with(&options).await.unwrap(); + let namespaced: i64 = + sqlx::query("SELECT count(1) AS ct FROM seen_mails WHERE mail_id LIKE 'mid:<%'") + .fetch_one(&mut conn) + .await + .unwrap() + .try_get("ct") + .unwrap(); + assert_eq!(namespaced, 2); + + // Re-importing the same mbox is a no-op for the counts. + let conn = migrated_conn(&options).await; + let s2 = run_import(&options, conn, &mbox, None, unset_cancel()) + .await + .unwrap(); + assert_eq!(s2.messages_seen, 3); + assert_eq!(s2.messages_new, 0, "re-import must add zero new counts"); + assert_eq!(sender_count(&options, "a@example.com").await, Some(1)); + assert_eq!(sender_count(&options, "b@example.com").await, Some(1)); + + // Both runs were recorded as done, with progress totals. + let mut conn = SqliteConnection::connect_with(&options).await.unwrap(); + let done: i64 = sqlx::query( + "SELECT count(1) AS ct FROM ingest_runs \ + WHERE source='mbox' AND state='done' AND bytes_done = bytes_total", ) + .fetch_one(&mut conn) .await + .unwrap() + .try_get("ct") .unwrap(); - tx.commit().await.unwrap(); + assert_eq!(done, 2); + } - let count: i64 = - sqlx::query("SELECT mails_sent FROM senders WHERE sender = 'jane@example.com'") - .fetch_one(&pool) - .await - .unwrap() - .try_get("mails_sent") - .unwrap(); - assert_eq!(count, 1); + /// Insert a half-finished import run row the way a cancelled run leaves it. + async fn plant_resumable_run( + conn: &mut SqliteConnection, + mbox: &Path, + fingerprint: &str, + bytes_done: u64, + ) -> i64 { + sqlx::query( + "INSERT INTO ingest_runs \ + (source, state, pid, started_at_unix, updated_at_unix, \ + bytes_total, bytes_done, mbox_path, mbox_fingerprint) \ + VALUES ('mbox', 'cancelled', 0, 0, 0, ?, ?, ?, ?)", + ) + .bind(IMPORT_FIXTURE.len() as i64) + .bind(bytes_done as i64) + .bind(mbox.display().to_string()) + .bind(fingerprint) + .execute(conn) + .await + .unwrap() + .last_insert_rowid() + } - // A second mail from the same (cleaned-up) sender increments it. - let mut tx = pool.begin().await.unwrap(); - increment_sender_mails(&cleanup_sender("jane@example.com".to_string()), &mut tx) + fn first_message_end_offset() -> u64 { + let mut reader = mbox::MboxReader::new(std::io::Cursor::new(IMPORT_FIXTURE.as_bytes()), 0); + reader.next_message().unwrap().unwrap().end_offset + } + + #[tokio::test] + async fn resume_with_matching_fingerprint_continues_from_the_offset() { + let dir = tempfile::tempdir().unwrap(); + let mbox = dir.path().join("mail.mbox"); + std::fs::write(&mbox, IMPORT_FIXTURE).unwrap(); + let options = db_connect_options(&dir.path().join("stats.db")); + + let mut conn = migrated_conn(&options).await; + let fingerprint = ingest::mbox_fingerprint(&mbox).unwrap(); + let offset = first_message_end_offset(); + let run_id = plant_resumable_run(&mut conn, &mbox, &fingerprint, offset).await; + + let summary = run_import(&options, conn, &mbox, Some(run_id), unset_cancel()) .await .unwrap(); - tx.commit().await.unwrap(); + // Only the two messages after the offset are parsed; Alice's message + // was before the resume point and is never touched. + assert_eq!(summary.messages_seen, 2); + assert_eq!(summary.messages_new, 1); + assert_eq!(summary.skipped, 1); + assert_eq!(sender_count(&options, "a@example.com").await, None); + assert_eq!(sender_count(&options, "b@example.com").await, Some(1)); + } - let count: i64 = - sqlx::query("SELECT mails_sent FROM senders WHERE sender = 'jane@example.com'") - .fetch_one(&pool) - .await - .unwrap() - .try_get("mails_sent") - .unwrap(); - assert_eq!(count, 2); + #[tokio::test] + async fn resume_with_stale_fingerprint_falls_back_to_a_full_parse() { + let dir = tempfile::tempdir().unwrap(); + let mbox = dir.path().join("mail.mbox"); + std::fs::write(&mbox, IMPORT_FIXTURE).unwrap(); + let options = db_connect_options(&dir.path().join("stats.db")); - // Only one row exists for the sender. - let rows: i64 = sqlx::query("SELECT count(1) AS ct FROM senders") - .fetch_one(&pool) + let mut conn = migrated_conn(&options).await; + let offset = first_message_end_offset(); + let run_id = plant_resumable_run(&mut conn, &mbox, "stale:fingerprint:0000", offset).await; + + let summary = run_import(&options, conn, &mbox, Some(run_id), unset_cancel()) + .await + .unwrap(); + assert_eq!(summary.messages_seen, 3, "must re-parse from the start"); + assert_eq!(summary.messages_new, 2); + assert_eq!(sender_count(&options, "a@example.com").await, Some(1)); + assert_eq!(sender_count(&options, "b@example.com").await, Some(1)); + } + + #[tokio::test] + async fn resuming_a_scan_run_is_rejected() { + let dir = tempfile::tempdir().unwrap(); + let mbox = dir.path().join("mail.mbox"); + std::fs::write(&mbox, IMPORT_FIXTURE).unwrap(); + let options = db_connect_options(&dir.path().join("stats.db")); + + let mut conn = migrated_conn(&options).await; + let run_id = ingest::create_run(&mut conn, "gmail_api", None) + .await + .unwrap(); + let err = run_import(&options, conn, &mbox, Some(run_id), unset_cancel()) + .await + .unwrap_err(); + assert!(err.to_string().contains("not an mbox import")); + } + + #[tokio::test] + async fn cancelled_import_marks_the_run_cancelled() { + let dir = tempfile::tempdir().unwrap(); + let mbox = dir.path().join("mail.mbox"); + std::fs::write(&mbox, IMPORT_FIXTURE).unwrap(); + let options = db_connect_options(&dir.path().join("stats.db")); + + let conn = migrated_conn(&options).await; + let cancel = Arc::new(AtomicBool::new(true)); // cancel before the first message + let summary = run_import(&options, conn, &mbox, None, cancel) + .await + .unwrap(); + assert!(summary.cancelled); + assert_eq!(summary.messages_seen, 0); + + let mut conn = SqliteConnection::connect_with(&options).await.unwrap(); + let state: String = sqlx::query("SELECT state FROM ingest_runs WHERE run_id = ?") + .bind(summary.run_id) + .fetch_one(&mut conn) .await .unwrap() - .try_get("ct") + .try_get("state") .unwrap(); - assert_eq!(rows, 1); + assert_eq!(state, "cancelled"); + } + + #[tokio::test] + async fn import_refuses_a_non_mbox_file() { + let dir = tempfile::tempdir().unwrap(); + let bogus = dir.path().join("notes.txt"); + std::fs::write(&bogus, "definitely not mail\n").unwrap(); + let options = db_connect_options(&dir.path().join("stats.db")); + + let conn = migrated_conn(&options).await; + let err = run_import(&options, conn, &bogus, None, unset_cancel()) + .await + .unwrap_err(); + assert!(err.to_string().contains("does not look like an mbox file")); } } diff --git a/src/mbox.rs b/src/mbox.rs new file mode 100644 index 0000000..b39234e --- /dev/null +++ b/src/mbox.rs @@ -0,0 +1,657 @@ +//! Streaming, header-only parser for mbox files (what Google Takeout exports). +//! +//! Design constraints, in service of multi-GB archives: +//! +//! * **Bounded memory**: the file is read through a `BufRead` one line at a +//! time into a single reusable buffer. Lines longer than [`MAX_LINE_BYTES`] +//! are truncated in memory but always fully consumed from the stream, so +//! message framing survives pathological input. Only the three headers the +//! importer needs (`From`, `Return-Path`, `Message-ID`) are retained, each +//! capped at [`MAX_HEADER_VALUE_BYTES`] after unfolding. +//! * **Header-only**: bodies are skipped line-by-line without interpretation. +//! * **mboxrd `>From ` escaping**: Takeout escapes body lines that would look +//! like separators as `>From `; such lines do not match the `From ` prefix +//! and are naturally treated as body content. An *unescaped* `From ` line +//! is, per the format, a new message separator. +//! * **CRLF**: a trailing `\r` is stripped from every line, so CRLF archives +//! parse identically to LF ones (offsets still count the real bytes). +//! * **Malformed regions**: garbage before the first separator, or a header +//! section containing a line that is neither a header, a continuation, nor +//! a separator, causes a skip to the next valid `From ` line. Each skipped +//! region increments [`MboxReader::skipped`]; the caller additionally counts +//! messages without a `Message-ID` as skipped. +//! * **Resumability**: every returned message carries `end_offset`, the +//! absolute offset of the byte after the message (i.e. the start of the next +//! `From ` separator, or EOF). A reader constructed with that offset over a +//! file seeked to it continues the parse exactly. + +use std::io::{self, BufRead}; + +/// Longest raw line kept in memory. Longer lines are truncated (but fully +/// consumed). Sender and Message-ID headers fit comfortably within this. +const MAX_LINE_BYTES: usize = 8 * 1024; +/// Cap on an unfolded (multi-line) header value. +const MAX_HEADER_VALUE_BYTES: usize = 8 * 1024; + +/// The handful of headers the importer cares about, parsed out of one message. +#[derive(Debug)] +pub struct RawMessage { + pub message_id: Option, + pub from: Option, + pub return_path: Option, + /// Absolute offset of the first byte after this message: the start of the + /// next message's `From ` separator line, or EOF. Once every message up to + /// here has been committed, this is a safe byte offset to resume from. + pub end_offset: u64, +} + +#[derive(Clone, Copy, Debug)] +enum Interest { + From, + ReturnPath, + MessageId, +} + +/// The header currently being accumulated (headers may be folded over +/// multiple lines, so a header is only committed when the next one starts). +enum Current { + /// Not inside any header yet (start of the header section). + None, + /// Inside a header the importer does not care about. + Boring, + /// Inside one of the interesting headers. + Keep(Interest, String), +} + +pub struct MboxReader { + inner: R, + /// Absolute offset of the next unread byte. + offset: u64, + /// Reusable line buffer (capped at MAX_LINE_BYTES). + line: Vec, + /// The next message's `From ` separator line was already consumed while + /// scanning for the end of the previous message. + pending_from: bool, + /// Malformed regions skipped so far. + skipped: u64, +} + +impl MboxReader { + /// `start_offset` is the absolute file offset the reader `inner` is + /// positioned at (0 for a whole file, or a resume offset for a seeked + /// file); all reported offsets are absolute. + pub fn new(inner: R, start_offset: u64) -> Self { + MboxReader { + inner, + offset: start_offset, + line: Vec::with_capacity(1024), + pending_from: false, + skipped: 0, + } + } + + /// Number of malformed regions skipped so far (junk before a separator, or + /// a header section that could not be parsed). + pub fn skipped(&self) -> u64 { + self.skipped + } + + fn read_line(&mut self) -> io::Result { + read_line_capped(&mut self.inner, &mut self.line) + } + + /// Parse the next message's headers; `Ok(None)` at EOF. + pub fn next_message(&mut self) -> io::Result> { + 'message: loop { + if !self.at_separator()? { + return Ok(None); + } + + let mut msg = RawMessage { + message_id: None, + from: None, + return_path: None, + end_offset: 0, + }; + let mut current = Current::None; + + // Header section: runs until a blank line (body follows), a new + // `From ` separator (message without a body), or EOF. + loop { + let n = self.read_line()?; + if n == 0 { + commit(&mut current, &mut msg); + msg.end_offset = self.offset; + return Ok(Some(msg)); + } + self.offset += n; + + if self.line.is_empty() { + commit(&mut current, &mut msg); + // Body: skip lines until the next separator or EOF. An + // mboxrd-escaped `>From ` line does not match the prefix + // and is skipped like any other body line. + loop { + let n = self.read_line()?; + if n == 0 { + msg.end_offset = self.offset; + return Ok(Some(msg)); + } + self.offset += n; + if self.line.starts_with(b"From ") { + self.pending_from = true; + msg.end_offset = self.offset - n; + return Ok(Some(msg)); + } + } + } + + if self.line.starts_with(b"From ") { + // Message with no blank line / body before the next one. + commit(&mut current, &mut msg); + self.pending_from = true; + msg.end_offset = self.offset - n; + return Ok(Some(msg)); + } + + if self.line[0] == b' ' || self.line[0] == b'\t' { + // Folded continuation of the previous header line. + match &mut current { + Current::None => { + // Continuation with no header before it: malformed. + self.skipped += 1; + self.seek_next_separator()?; + continue 'message; + } + Current::Boring => {} + Current::Keep(_, value) => { + let folded = String::from_utf8_lossy(&self.line); + let folded = folded.trim(); + if !folded.is_empty() + && value.len() + folded.len() < MAX_HEADER_VALUE_BYTES + { + if !value.is_empty() { + value.push(' '); + } + value.push_str(folded); + } + } + } + continue; + } + + match self.line.iter().position(|&b| b == b':') { + None | Some(0) => { + // Not a header, a continuation, or a separator: give + // up on this message and skip to the next one. + self.skipped += 1; + self.seek_next_separator()?; + continue 'message; + } + Some(colon) => { + commit(&mut current, &mut msg); + let name = self.line[..colon].trim_ascii(); + current = match interest_for(name) { + Some(interest) => { + let value = String::from_utf8_lossy(&self.line[colon + 1..]) + .trim() + .to_string(); + Current::Keep(interest, value) + } + None => Current::Boring, + }; + } + } + } + } + } + + /// Position the parser just after a `From ` separator line. Returns false + /// at EOF. Counts any non-blank junk encountered on the way as one + /// skipped region. + fn at_separator(&mut self) -> io::Result { + if self.pending_from { + self.pending_from = false; + return Ok(true); + } + let mut saw_junk = false; + loop { + let n = self.read_line()?; + if n == 0 { + if saw_junk { + self.skipped += 1; + } + return Ok(false); + } + self.offset += n; + if self.line.starts_with(b"From ") { + if saw_junk { + self.skipped += 1; + } + return Ok(true); + } + if !self.line.is_empty() { + saw_junk = true; + } + } + } + + /// After a malformed region: discard lines until the next `From ` + /// separator (left pending) or EOF. + fn seek_next_separator(&mut self) -> io::Result<()> { + loop { + let n = self.read_line()?; + if n == 0 { + return Ok(()); + } + self.offset += n; + if self.line.starts_with(b"From ") { + self.pending_from = true; + return Ok(()); + } + } + } +} + +fn commit(current: &mut Current, msg: &mut RawMessage) { + if let Current::Keep(interest, value) = std::mem::replace(current, Current::None) { + let slot = match interest { + Interest::From => &mut msg.from, + Interest::ReturnPath => &mut msg.return_path, + Interest::MessageId => &mut msg.message_id, + }; + // First occurrence wins; empty values are as good as absent. + if slot.is_none() && !value.is_empty() { + *slot = Some(value); + } + } +} + +fn interest_for(name: &[u8]) -> Option { + if name.eq_ignore_ascii_case(b"from") { + Some(Interest::From) + } else if name.eq_ignore_ascii_case(b"return-path") { + Some(Interest::ReturnPath) + } else if name.eq_ignore_ascii_case(b"message-id") { + Some(Interest::MessageId) + } else { + None + } +} + +/// Read one line (through the next `\n` or EOF) into `buf`, which is cleared +/// first. At most [`MAX_LINE_BYTES`] are kept, but the full line is always +/// consumed from the stream. The newline is not stored; a trailing `\r` is +/// stripped. Returns the number of bytes consumed; 0 means EOF. +fn read_line_capped(reader: &mut R, buf: &mut Vec) -> io::Result { + buf.clear(); + let mut consumed: u64 = 0; + loop { + let (done, used) = { + let available = match reader.fill_buf() { + Ok(a) => a, + Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue, + Err(e) => return Err(e), + }; + if available.is_empty() { + break; + } + match available.iter().position(|&b| b == b'\n') { + Some(i) => { + if buf.len() < MAX_LINE_BYTES { + let end = i.min(MAX_LINE_BYTES - buf.len()); + buf.extend_from_slice(&available[..end]); + } + (true, i + 1) + } + None => { + if buf.len() < MAX_LINE_BYTES { + let end = available.len().min(MAX_LINE_BYTES - buf.len()); + buf.extend_from_slice(&available[..end]); + } + (false, available.len()) + } + } + }; + reader.consume(used); + consumed += used as u64; + if done { + break; + } + } + if buf.last() == Some(&b'\r') { + buf.pop(); + } + Ok(consumed) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::{BufReader, Cursor, Read, Write}; + + fn read_all(data: &[u8]) -> (Vec, u64) { + let mut reader = MboxReader::new(Cursor::new(data), 0); + let mut messages = Vec::new(); + while let Some(msg) = reader.next_message().unwrap() { + messages.push(msg); + } + (messages, reader.skipped()) + } + + const TWO_MESSAGES: &str = "From a@example.com Thu Jan 1 00:00:00 2020\n\ + From: Alice \n\ + Message-ID: \n\ + Subject: hello\n\ + \n\ + body line one\n\ + body line two\n\ + From b@example.com Thu Jan 1 00:00:00 2020\n\ + Return-Path: \n\ + Message-ID: \n\ + \n\ + body\n"; + + #[test] + fn parses_normal_messages() { + let (msgs, skipped) = read_all(TWO_MESSAGES.as_bytes()); + assert_eq!(skipped, 0); + assert_eq!(msgs.len(), 2); + assert_eq!(msgs[0].from.as_deref(), Some("Alice ")); + assert_eq!(msgs[0].message_id.as_deref(), Some("")); + assert_eq!(msgs[1].from, None); + assert_eq!(msgs[1].return_path.as_deref(), Some("")); + assert_eq!(msgs[1].message_id.as_deref(), Some("")); + // end_offset of a message is the start of the next separator / EOF. + let second_start = TWO_MESSAGES.find("From b@").unwrap() as u64; + assert_eq!(msgs[0].end_offset, second_start); + assert_eq!(msgs[1].end_offset, TWO_MESSAGES.len() as u64); + } + + #[test] + fn parses_crlf_line_endings() { + let crlf = TWO_MESSAGES.replace('\n', "\r\n"); + let (msgs, skipped) = read_all(crlf.as_bytes()); + assert_eq!(skipped, 0); + assert_eq!(msgs.len(), 2); + assert_eq!(msgs[0].from.as_deref(), Some("Alice ")); + assert_eq!(msgs[1].message_id.as_deref(), Some("")); + // Offsets count the real (CRLF) bytes. + assert_eq!(msgs[1].end_offset, crlf.len() as u64); + } + + #[test] + fn mboxrd_escaped_from_lines_stay_in_the_body() { + let data = "From a@example.com Thu Jan 1 00:00:00 2020\n\ + From: a@example.com\n\ + Message-ID: \n\ + \n\ + >From the archives, an escaped line\n\ + >>From twice-escaped\n\ + still message one\n"; + let (msgs, skipped) = read_all(data.as_bytes()); + assert_eq!(skipped, 0); + assert_eq!(msgs.len(), 1, "escaped From lines must not split messages"); + assert_eq!(msgs[0].end_offset, data.len() as u64); + } + + #[test] + fn unescaped_from_line_in_body_is_a_separator() { + // mboxrd semantics: an unescaped `From ` line *is* a separator. The + // resulting bogus "message" has no parseable headers and is counted + // as a skipped region. + let data = "From a@example.com Thu Jan 1 00:00:00 2020\n\ + From: a@example.com\n\ + Message-ID: \n\ + \n\ + From here on, an unescaped body line\n\ + more body\n"; + let (msgs, skipped) = read_all(data.as_bytes()); + assert_eq!(msgs.len(), 1); + assert_eq!(skipped, 1); + } + + #[test] + fn missing_message_id_is_reported_as_none() { + let data = "From a@example.com Thu Jan 1 00:00:00 2020\n\ + From: a@example.com\n\ + Subject: no message id\n\ + \n\ + body\n"; + let (msgs, skipped) = read_all(data.as_bytes()); + assert_eq!(skipped, 0); + assert_eq!(msgs.len(), 1); + assert_eq!(msgs[0].message_id, None); + assert_eq!(msgs[0].from.as_deref(), Some("a@example.com")); + } + + #[test] + fn malformed_header_region_skips_to_next_separator() { + let data = "From a@example.com Thu Jan 1 00:00:00 2020\n\ + From: a@example.com\n\ + this line is not a header and not a continuation\n\ + trailing garbage\n\ + From b@example.com Thu Jan 1 00:00:00 2020\n\ + From: b@example.com\n\ + Message-ID: \n\ + \n\ + body\n"; + let (msgs, skipped) = read_all(data.as_bytes()); + assert_eq!(skipped, 1); + assert_eq!(msgs.len(), 1, "the malformed message is dropped"); + assert_eq!(msgs[0].message_id.as_deref(), Some("")); + } + + #[test] + fn junk_before_first_separator_is_skipped_and_counted() { + let data = "random junk that is not mbox\nmore junk\n\ + From a@example.com Thu Jan 1 00:00:00 2020\n\ + From: a@example.com\n\ + Message-ID: \n\ + \n\ + body\n"; + let (msgs, skipped) = read_all(data.as_bytes()); + assert_eq!(skipped, 1); + assert_eq!(msgs.len(), 1); + assert_eq!(msgs[0].message_id.as_deref(), Some("")); + } + + #[test] + fn folded_headers_are_unfolded() { + let data = "From a@example.com Thu Jan 1 00:00:00 2020\n\ + From: A Very Long Display Name\n\ + \t\n\ + Message-ID:\n\ + \x20\n\ + \n\ + body\n"; + let (msgs, skipped) = read_all(data.as_bytes()); + assert_eq!(skipped, 0); + assert_eq!( + msgs[0].from.as_deref(), + Some("A Very Long Display Name ") + ); + assert_eq!( + msgs[0].message_id.as_deref(), + Some("") + ); + } + + #[test] + fn header_names_are_case_insensitive_and_first_wins() { + let data = "From a@example.com Thu Jan 1 00:00:00 2020\n\ + RETURN-PATH: \n\ + return-path: \n\ + MESSAGE-id: \n\ + \n\ + body\n"; + let (msgs, _) = read_all(data.as_bytes()); + assert_eq!(msgs[0].return_path.as_deref(), Some("")); + assert_eq!(msgs[0].message_id.as_deref(), Some("")); + } + + #[test] + fn overlong_lines_are_truncated_but_framing_survives() { + let long_value = "a".repeat(3 * MAX_LINE_BYTES); + let data = format!( + "From a@example.com Thu Jan 1 00:00:00 2020\n\ + X-Long: {long_value}\n\ + Message-ID: \n\ + \n\ + {long_value}\n\ + From b@example.com Thu Jan 1 00:00:00 2020\n\ + Message-ID: \n\ + \n\ + body\n" + ); + let (msgs, skipped) = read_all(data.as_bytes()); + assert_eq!(skipped, 0); + assert_eq!(msgs.len(), 2); + assert_eq!(msgs[0].message_id.as_deref(), Some("")); + assert_eq!(msgs[1].message_id.as_deref(), Some("")); + assert_eq!(msgs[1].end_offset, data.len() as u64); + } + + #[test] + fn message_without_body_or_trailing_newline() { + let data = "From a@example.com Thu Jan 1 00:00:00 2020\n\ + From: a@example.com\n\ + Message-ID: "; + let (msgs, skipped) = read_all(data.as_bytes()); + assert_eq!(skipped, 0); + assert_eq!(msgs.len(), 1); + assert_eq!(msgs[0].message_id.as_deref(), Some("")); + assert_eq!(msgs[0].end_offset, data.len() as u64); + } + + #[test] + fn separator_directly_after_headers_ends_the_message() { + let data = "From a@example.com Thu Jan 1 00:00:00 2020\n\ + Message-ID: \n\ + From b@example.com Thu Jan 1 00:00:00 2020\n\ + Message-ID: \n\ + \n\ + body\n"; + let (msgs, skipped) = read_all(data.as_bytes()); + assert_eq!(skipped, 0); + assert_eq!(msgs.len(), 2); + assert_eq!(msgs[0].message_id.as_deref(), Some("")); + assert_eq!(msgs[1].message_id.as_deref(), Some("")); + } + + #[test] + fn resume_from_end_offset_continues_the_parse() { + let (all, _) = read_all(TWO_MESSAGES.as_bytes()); + let offset = all[0].end_offset; + let rest = &TWO_MESSAGES.as_bytes()[offset as usize..]; + let mut reader = MboxReader::new(Cursor::new(rest), offset); + let msg = reader.next_message().unwrap().unwrap(); + assert_eq!(msg.message_id.as_deref(), Some("")); + assert_eq!(msg.end_offset, TWO_MESSAGES.len() as u64); + assert!(reader.next_message().unwrap().is_none()); + assert_eq!(reader.skipped(), 0); + } + + /// An infinite-source style reader: generates mbox messages on the fly, so + /// the input provably never exists in memory as a whole. Combined with the + /// capped line buffer this demonstrates the parser streams in bounded + /// memory regardless of input size. + struct SyntheticMbox { + next_index: u64, + total: u64, + chunk: Vec, + pos: usize, + } + + fn synthetic_message(i: u64, body_bytes: usize) -> Vec { + let body = "x".repeat(72); + let mut out = format!( + "From sender{i}@example.com Thu Jan 1 00:00:00 2020\n\ + From: Sender {i} \n\ + Message-ID: \n\ + Subject: synthetic message {i}\n\ + \n", + i % 997 + ) + .into_bytes(); + let mut written = 0; + while written < body_bytes { + out.extend_from_slice(body.as_bytes()); + out.push(b'\n'); + written += body.len() + 1; + } + out + } + + impl Read for SyntheticMbox { + fn read(&mut self, out: &mut [u8]) -> io::Result { + if self.pos == self.chunk.len() { + if self.next_index == self.total { + return Ok(0); + } + self.chunk = synthetic_message(self.next_index, 1024); + self.pos = 0; + self.next_index += 1; + } + let n = (self.chunk.len() - self.pos).min(out.len()); + out[..n].copy_from_slice(&self.chunk[self.pos..self.pos + n]); + self.pos += n; + Ok(n) + } + } + + #[test] + fn streams_a_large_generated_mailbox() { + // ~64 MB of mbox, generated on the fly (never materialized). + let total = 55_000u64; + let source = SyntheticMbox { + next_index: 0, + total, + chunk: Vec::new(), + pos: 0, + }; + let mut reader = MboxReader::new(BufReader::with_capacity(64 * 1024, source), 0); + let mut count = 0u64; + let mut last_offset = 0u64; + while let Some(msg) = reader.next_message().unwrap() { + count += 1; + assert!(msg.message_id.is_some()); + assert!(msg.end_offset > last_offset); + last_offset = msg.end_offset; + } + assert_eq!(count, total); + assert_eq!(reader.skipped(), 0); + } + + #[test] + fn parses_a_large_mbox_file_from_disk() { + // ~50 MB real file in a temp dir, parsed through the same streaming + // path the importer uses (BufReader over File). + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("large.mbox"); + let total = 12_000u64; + { + let mut file = std::io::BufWriter::new(std::fs::File::create(&path).unwrap()); + for i in 0..total { + file.write_all(&synthetic_message(i, 4096)).unwrap(); + } + } + let size = std::fs::metadata(&path).unwrap().len(); + assert!( + size > 45 * 1024 * 1024, + "fixture should be ~50MB, got {size}" + ); + + let file = std::fs::File::open(&path).unwrap(); + let mut reader = MboxReader::new(BufReader::with_capacity(64 * 1024, file), 0); + let mut count = 0u64; + let mut final_offset = 0u64; + while let Some(msg) = reader.next_message().unwrap() { + count += 1; + final_offset = msg.end_offset; + } + assert_eq!(count, total); + assert_eq!(reader.skipped(), 0); + assert_eq!(final_offset, size); + } +}