fix(migrate): large-database & managed-Postgres hardening for export/import/copy

crates/codex-migrate/src/archive.rs

@@ -20,6 +20,7 @@ use crate::manifest::{
     ARCHIVE_FORMAT_VERSION, ArtifactEntry, ArtifactGroup, Manifest, TableCount, backend_name,
     schema_version,
 };
+use crate::progress::Progress;
 use crate::reroot::{self, RerootStats};
 use crate::{TransferReport, registry};
 
@@ -45,6 +46,9 @@ pub struct ImportOutcome {
     pub manifest: Manifest,
     pub report: TransferReport,
     pub reroot: RerootStats,
+    /// Per-table canonical-content mismatches, when full verification was
+    /// requested; empty otherwise.
+    pub full_verify: Vec<crate::full_verify::FullMismatch>,
 }
 
 /// Export `conn` and the given artifact trees into a `.tar.gz` at `out_path`.
@@ -54,12 +58,13 @@ pub async fn export_archive(
     conn: &DatabaseConnection,
     out_path: &Path,
     artifacts: &[ArtifactSource],
+    progress: Progress,
 ) -> Result<Manifest> {
     let staging = tempfile::tempdir().context("failed to create export staging dir")?;
     let db_dir = staging.path().join("db");
     std::fs::create_dir_all(&db_dir)?;
 
-    let counts = registry::dump_all_to_dir(conn, &db_dir)
+    let counts = registry::dump_all_to_dir(conn, &db_dir, progress)
         .await
         .context("failed to dump database tables")?;
     let total_rows = counts.iter().map(|c| c.rows).sum();
@@ -111,11 +116,14 @@ pub async fn import_archive(
     conn: &DatabaseConnection,
     in_path: &Path,
     targets: &[ArtifactTarget],
+    progress: Progress,
+    full_verify: bool,
 ) -> Result<ImportOutcome> {
     let staging = tempfile::tempdir().context("failed to create import staging dir")?;
     let manifest = extract(in_path, staging.path())?;
+    let db_dir = staging.path().join("db");
 
-    let report = crate::load_from_dir(conn, &staging.path().join("db"))
+    let report = crate::load_from_dir(conn, &db_dir, progress)
         .await
         .context("failed to load database from archive")?;
 
@@ -160,10 +168,25 @@ pub async fn import_archive(
     .await
     .context("failed to re-root artifact paths")?;
 
+    // Optional deep check: compare the canonical content of every row in the
+    // archive against the loaded target.
+    let full_verify_mismatches = if full_verify {
+        let source = registry::digest_all_from_ndjson_dir(&db_dir)
+            .await
+            .context("failed to digest archive contents")?;
+        let target = registry::digest_all_from_conn(conn)
+            .await
+            .context("failed to digest imported data")?;
+        crate::full_verify::compare_digests(&source, &target)
+    } else {
+        Vec::new()
+    };
+
     Ok(ImportOutcome {
         manifest,
         report,
         reroot,
+        full_verify: full_verify_mismatches,
     })
 }
 

crates/codex-migrate/src/engine.rs

@@ -19,6 +19,9 @@ use sea_orm::{
 use serde::Serialize;
 use serde::de::DeserializeOwned;
 
+use crate::full_verify::{DigestAccumulator, TableDigest};
+use crate::progress::{Progress, ROW_REPORT_INTERVAL};
+
 /// Default insert batch size. Bounds memory and round-trips while staying well
 /// under parameter limits for the widest tables.
 pub const DEFAULT_BATCH_SIZE: usize = 1000;
@@ -78,34 +81,48 @@ fn iden_string<I: Iden>(iden: &I) -> String {
 
 /// Stream every row of `E` and write it as one NDJSON line to `out`.
 /// Returns the number of rows written.
-pub async fn dump_table<E, C, W>(conn: &C, out: &mut W) -> Result<u64>
+pub async fn dump_table<E, C, W>(conn: &C, out: &mut W, progress: Progress) -> Result<u64>
 where
     E: EntityTrait,
     E::Model: Serialize,
     C: ConnectionTrait + StreamTrait,
     W: Write,
 {
+    let table = iden_string(&E::default());
     let mut stream = open_source_stream::<E, C>(conn).await?;
     let mut count = 0u64;
+    let mut next_report = ROW_REPORT_INTERVAL;
     while let Some(model) = stream.try_next().await? {
         serde_json::to_writer(&mut *out, &model)?;
         out.write_all(b"\n")?;
         count += 1;
+        if count >= next_report {
+            progress.table_rows(&table, count);
+            next_report += ROW_REPORT_INTERVAL;
+        }
     }
     Ok(count)
 }
 
 /// Read NDJSON lines from `reader`, deserialize each into `E::Model`, and
 /// insert them into `conn` in batches of `batch_size`. Returns rows inserted.
-pub async fn load_table<E, C, R>(conn: &C, reader: R, batch_size: usize) -> Result<u64>
+pub async fn load_table<E, C, R>(
+    conn: &C,
+    reader: R,
+    batch_size: usize,
+    progress: Progress,
+) -> Result<u64>
 where
     E: EntityTrait,
     E::Model: DeserializeOwned + IntoActiveModel<E::ActiveModel>,
     C: ConnectionTrait,
     R: BufRead,
 {
+    let table = iden_string(&E::default());
+    let batch_size = safe_batch_size::<E>(conn.get_database_backend(), batch_size);
     let mut batch: Vec<E::ActiveModel> = Vec::with_capacity(batch_size);
     let mut count = 0u64;
+    let mut next_report = ROW_REPORT_INTERVAL;
     for line in reader.lines() {
         let line = line?;
         if line.trim().is_empty() {
@@ -115,6 +132,10 @@ where
         batch.push(model.into_active_model());
         if batch.len() >= batch_size {
             count += insert_batch::<E, C>(conn, std::mem::take(&mut batch)).await?;
+            if count >= next_report {
+                progress.table_rows(&table, count);
+                next_report += ROW_REPORT_INTERVAL;
+            }
         }
     }
     if !batch.is_empty() {
@@ -126,20 +147,32 @@ where
 /// Stream every row of `E` directly from `src` into `dst` in batches, without
 /// an intermediate serialized form. This is the path used by the direct
 /// database-to-database `copy`.
-pub async fn copy_table<E, S, D>(src: &S, dst: &D, batch_size: usize) -> Result<u64>
+pub async fn copy_table<E, S, D>(
+    src: &S,
+    dst: &D,
+    batch_size: usize,
+    progress: Progress,
+) -> Result<u64>
 where
     E: EntityTrait,
     E::Model: IntoActiveModel<E::ActiveModel>,
     S: ConnectionTrait + StreamTrait,
     D: ConnectionTrait,
 {
+    let table = iden_string(&E::default());
+    let batch_size = safe_batch_size::<E>(dst.get_database_backend(), batch_size);
     let mut stream = open_source_stream::<E, S>(src).await?;
     let mut batch: Vec<E::ActiveModel> = Vec::with_capacity(batch_size);
     let mut count = 0u64;
+    let mut next_report = ROW_REPORT_INTERVAL;
     while let Some(model) = stream.try_next().await? {
         batch.push(model.into_active_model());
         if batch.len() >= batch_size {
             count += insert_batch::<E, D>(dst, std::mem::take(&mut batch)).await?;
+            if count >= next_report {
+                progress.table_rows(&table, count);
+                next_report += ROW_REPORT_INTERVAL;
+            }
         }
     }
     if !batch.is_empty() {
@@ -158,6 +191,43 @@ where
     Ok(E::find().count(conn).await?)
 }
 
+/// Compute the canonical content digest of `E` read from a connection.
+pub async fn digest_table_from_conn<E, C>(conn: &C) -> Result<TableDigest>
+where
+    E: EntityTrait,
+    E::Model: Serialize,
+    C: ConnectionTrait + StreamTrait,
+{
+    let table = iden_string(&E::default());
+    let mut stream = open_source_stream::<E, C>(conn).await?;
+    let mut acc = DigestAccumulator::default();
+    while let Some(model) = stream.try_next().await? {
+        acc.add(&model);
+    }
+    Ok(acc.finish(table))
+}
+
+/// Compute the canonical content digest of `E` read from an NDJSON reader (the
+/// archive payload).
+pub fn digest_table_from_ndjson<E, R>(reader: R) -> Result<TableDigest>
+where
+    E: EntityTrait,
+    E::Model: DeserializeOwned + Serialize,
+    R: BufRead,
+{
+    let table = iden_string(&E::default());
+    let mut acc = DigestAccumulator::default();
+    for line in reader.lines() {
+        let line = line?;
+        if line.trim().is_empty() {
+            continue;
+        }
+        let model: E::Model = serde_json::from_str(&line)?;
+        acc.add(&model);
+    }
+    Ok(acc.finish(table))
+}
+
 /// Delete every row of `E`. Used by `--replace` before a load. Callers must
 /// have foreign-key enforcement disabled (see [`crate::fk`]).
 pub async fn truncate_table<E, C>(conn: &C) -> Result<u64>
@@ -168,6 +238,19 @@ where
     Ok(E::delete_many().exec(conn).await?.rows_affected)
 }
 
+/// Cap the batch so a multi-row insert can't exceed the destination's bind
+/// parameter limit. A batch binds `rows × columns` parameters; PostgreSQL caps
+/// a statement at 65535 and SQLite at 32766, so a wide table (e.g.
+/// `book_metadata`, ~66 columns) overflows a naive 1000-row batch.
+fn safe_batch_size<E: EntityTrait>(backend: DatabaseBackend, requested: usize) -> usize {
+    let columns = <E::Column as Iterable>::iter().count().max(1);
+    let param_limit = match backend {
+        DatabaseBackend::Postgres | DatabaseBackend::MySql => 65535,
+        DatabaseBackend::Sqlite => 32766,
+    };
+    requested.min((param_limit / columns).max(1))
+}
+
 /// Insert one batch, using `exec_without_returning` to avoid a per-row
 /// RETURNING clause on bulk loads. No-op (and no DB round-trip) when empty.
 async fn insert_batch<E, C>(conn: &C, batch: Vec<E::ActiveModel>) -> Result<u64>
@@ -183,3 +266,29 @@ where
     E::insert_many(batch).exec_without_returning(conn).await?;
     Ok(n)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::safe_batch_size;
+    use sea_orm::DatabaseBackend;
+
+    #[test]
+    fn caps_wide_table_under_postgres_param_limit() {
+        // book_metadata has 66 columns; a naive 1000-row batch would bind
+        // 66000 parameters, over PostgreSQL's 65535 limit.
+        let n = safe_batch_size::<codex_db::entities::book_metadata::Entity>(
+            DatabaseBackend::Postgres,
+            1000,
+        );
+        assert!(n < 1000, "wide table should be capped, got {n}");
+        assert!(n * 66 <= 65535, "batch {n} still exceeds the limit");
+    }
+
+    #[test]
+    fn keeps_requested_size_for_narrow_table() {
+        // genres has 4 columns; 1000 rows is well within the limit.
+        let n =
+            safe_batch_size::<codex_db::entities::genres::Entity>(DatabaseBackend::Postgres, 1000);
+        assert_eq!(n, 1000);
+    }
+}