Skip to main content

mitosis/
diff.rs

1//! `mitosis diff`: what a clone changed against the template generation it was made from.
2//!
3//! The base side is a private, socket-only Postgres on a copy-on-write clone of the generation folder; a
4//! guard stops it and deletes the folder on every path. Both sides are read with `psql`.
5
6use crate::clone_backend::CloneBackend;
7use crate::clones;
8use crate::ctx::Ctx;
9use crate::error::{ErrorCode, MitosisError, Result};
10use crate::names::random_hex;
11use crate::pg::{psql, start, stop, wait_ready, PgInstall, RunningPg, StartOptions};
12use crate::protocol::{Count, DiffReport, RowChange, RowDiff, SchemaChange, TableCount};
13use serde_json::Value;
14use std::collections::{BTreeMap, HashMap, HashSet};
15use std::path::PathBuf;
16use std::sync::Arc;
17use std::time::Duration;
18
19/// Rows shown per kind when `--limit` is not given.
20pub const DEFAULT_LIMIT: u32 = 20;
21/// The most rows per kind `--limit` accepts.
22pub const MAX_LIMIT: u32 = 1000;
23/// `--table` refuses tables above this many rows (by the planner's estimate).
24const MAX_TABLE_ROWS: i64 = 1_000_000;
25/// The socket file name only; the base server listens on no port.
26const BASE_PORT: u16 = 5432;
27
28fn invalid(msg: impl Into<String>) -> MitosisError {
29    MitosisError::new(ErrorCode::InvalidRequest, msg)
30}
31
32fn internal(msg: impl Into<String>) -> MitosisError {
33    MitosisError::new(ErrorCode::Internal, msg)
34}
35
36async fn blocking<T, F>(f: F) -> Result<T>
37where
38    F: FnOnce() -> Result<T> + Send + 'static,
39    T: Send + 'static,
40{
41    tokio::task::spawn_blocking(f)
42        .await
43        .map_err(|e| internal(format!("background task failed: {e}")))?
44}
45
46/// Refuses a clone that does not belong to the MCP session `parent_pid`.
47pub fn check_session(ctx: &Ctx, id_or_name: &str, parent_pid: i32) -> Result<()> {
48    let row = ctx.store().get_clone(id_or_name)?.ok_or_else(|| {
49        MitosisError::new(
50            ErrorCode::CloneNotFound,
51            format!("clone {id_or_name} not found"),
52        )
53        .with_hint("see this session's clones with the clone_list tool")
54    })?;
55    if row.parent_pid != Some(parent_pid) {
56        return Err(invalid(format!(
57            "clone {} was not created by this session; it may be another agent's database",
58            row.id
59        ))
60        .with_hint("this tool only reads clones from clone_create in this session"));
61    }
62    Ok(())
63}
64
65/// The private base server and the folders around it. `close` (or, on a cancelled request, drop)
66/// stops the server and removes everything it made; a failure is written to `daemon.log` with the
67/// leaked folder, and `Drop` tries once more.
68struct BaseServer {
69    backend: Arc<dyn CloneBackend>,
70    install: PgInstall,
71    dir: PathBuf,
72    socket_dir: PathBuf,
73    log: PathBuf,
74    daemon_log: PathBuf,
75    pg: Option<RunningPg>,
76    timeout: Duration,
77    done: bool,
78}
79
80impl BaseServer {
81    /// Stops the server first; the folder is removed only once nothing runs in it.
82    fn cleanup(&mut self) -> Result<()> {
83        if self.dir.exists() {
84            stop(&self.install, &self.dir, self.pg, self.timeout)?;
85            self.backend.remove_dir(&self.dir)?;
86        }
87        let _ = std::fs::remove_dir_all(&self.socket_dir);
88        let _ = std::fs::remove_file(&self.log);
89        self.done = true;
90        Ok(())
91    }
92
93    fn cleanup_logged(&mut self) {
94        if let Err(e) = self.cleanup() {
95            use std::io::Write;
96            if let Ok(mut f) = std::fs::OpenOptions::new()
97                .create(true)
98                .append(true)
99                .open(&self.daemon_log)
100            {
101                let _ = writeln!(
102                    f,
103                    "[{}] diff: could not clean up the private server in {}: {}",
104                    crate::templates::now(),
105                    self.dir.display(),
106                    e.message
107                );
108            }
109        }
110    }
111
112    /// Cleans up off the runtime. Never fails: the diff's own result stands.
113    async fn close(mut self) {
114        let _ = tokio::task::spawn_blocking(move || {
115            self.cleanup_logged();
116            self // dropping it here retries a cleanup that failed
117        })
118        .await;
119    }
120}
121
122// shortcut: a drop on a runtime thread blocks on pg_ctl; it only runs for a cancelled request
123// or a failed cleanup. If `start` was still running then, `pg` is unknown and the postmaster can
124// outlive the folder removal.
125impl Drop for BaseServer {
126    fn drop(&mut self) {
127        if !self.done {
128            self.cleanup_logged();
129        }
130    }
131}
132
133/// One server a query can reach over a socket.
134#[derive(Clone)]
135struct Side {
136    install: PgInstall,
137    socket_dir: PathBuf,
138    port: u16,
139    db: String,
140}
141
142impl Side {
143    async fn query(&self, sql: String) -> Result<String> {
144        let s = self.clone();
145        blocking(move || psql(&s.install, &s.socket_dir, s.port, &s.db, &sql)).await
146    }
147}
148
149pub async fn diff(
150    ctx: &Arc<Ctx>,
151    id_or_name: &str,
152    table: Option<&str>,
153    limit: u32,
154) -> Result<DiffReport> {
155    if limit > MAX_LIMIT {
156        return Err(invalid(format!("limit {limit} is above {MAX_LIMIT}"))
157            .with_hint(format!("ask for at most {MAX_LIMIT} rows per kind")));
158    }
159    let row = clones::ensure_running(ctx, id_or_name).await?;
160    let label = row.name.clone().unwrap_or_else(|| row.id.clone());
161    let pruned = || {
162        invalid(format!(
163            "generation {} of template {} that clone {label} was made from was pruned",
164            row.generation, row.template
165        ))
166        .with_hint("the base is gone; nothing to compare against. Raise the template's `keep` before pruning next time")
167    };
168
169    // Held from reading the generation until its folder is cloned, so no prune moves it.
170    let swap = ctx.template_swap.read().await;
171    let gen = ctx
172        .store()
173        .get_generation(&row.template, row.generation)?
174        .filter(|g| g.dir.is_dir())
175        .ok_or_else(pruned)?;
176    let install = ctx.install_major(gen.pg_major).await?;
177
178    let tag = random_hex(4);
179    let mut base = BaseServer {
180        backend: ctx.backend.clone(),
181        install: install.clone(),
182        dir: ctx.home.templates_dir().join(format!(".diff-{tag}")),
183        // Short: Unix socket paths are limited to about 100 bytes.
184        socket_dir: ctx.home.sockets_dir().join(format!("d{tag}")),
185        log: ctx.home.logs_dir().join(format!("diff-{tag}.log")),
186        daemon_log: ctx.home.logs_dir().join("daemon.log"),
187        pg: None,
188        timeout: ctx.config().stop_timeout,
189        done: false,
190    };
191    let (backend, src, dst) = (ctx.backend.clone(), gen.dir.clone(), base.dir.clone());
192    let cloned = blocking(move || backend.clone_dir(&src, &dst)).await;
193    drop(swap);
194
195    let run = async {
196        cloned?;
197        let opts = StartOptions {
198            port: BASE_PORT,
199            socket_dir: base.socket_dir.clone(),
200            listen_tcp: false,
201            durable: false,
202            log_file: base.log.clone(),
203        };
204        let (i, d, o) = (install.clone(), base.dir.clone(), opts.clone());
205        let pg = blocking(move || start(&i, &d, &o)).await?;
206        base.pg = Some(pg);
207        wait_ready(&install, &pg, &opts, ctx.config().ready_timeout).await?;
208        let before = Side {
209            install: install.clone(),
210            socket_dir: base.socket_dir.clone(),
211            port: BASE_PORT,
212            db: row.database.clone(),
213        };
214        let after = Side {
215            install: install.clone(),
216            socket_dir: row.socket_dir.clone(),
217            port: row.port,
218            db: row.database.clone(),
219        };
220        compare(&before, &after, table, limit).await
221    }
222    .await;
223    base.close().await;
224    let (schema, tables, rows) = run?;
225    Ok(DiffReport {
226        clone: label,
227        template: row.template,
228        generation: row.generation,
229        schema,
230        tables,
231        rows,
232    })
233}
234
235type Compared = (Vec<SchemaChange>, Vec<TableCount>, Option<RowDiff>);
236
237async fn compare(before: &Side, after: &Side, table: Option<&str>, limit: u32) -> Result<Compared> {
238    let (b, a) = tokio::try_join!(snapshot(before), snapshot(after))?;
239    let schema = schema_changes(&b.objects, &a.objects);
240    let tables = table_counts(before, after, &b.tables, &a.tables).await?;
241    let rows = match table {
242        Some(t) => Some(row_diff(before, after, t, limit as usize).await?),
243        None => None,
244    };
245    Ok((schema, tables, rows))
246}
247
248// ---- schema -------------------------------------------------------------------------------
249
250/// (kind, name, definition, owning table).
251type Obj = (String, String, Option<String>, Option<String>);
252
253/// (schema, table, planner estimate).
254type TableEntry = (String, String, i64);
255
256struct Snapshot {
257    objects: Vec<Obj>,
258    tables: Vec<TableEntry>,
259}
260
261const KINDS: [&str; 8] = [
262    "table",
263    "column",
264    "index",
265    "constraint",
266    "view",
267    "function",
268    "sequence",
269    "extension",
270];
271
272/// Catalog reads in user schemas only. Two statements, one JSON line each.
273const SNAPSHOT_SQL: &str = r#"
274WITH ns AS (
275  SELECT oid, nspname FROM pg_namespace
276  WHERE nspname NOT IN ('pg_catalog', 'information_schema')
277    AND nspname NOT LIKE 'pg\_toast%' AND nspname NOT LIKE 'pg\_temp\_%'
278    AND nspname NOT LIKE 'pg\_toast\_temp\_%'
279)
280SELECT coalesce(jsonb_agg(jsonb_build_array(kind, name, def, tbl)), '[]'::jsonb) FROM (
281  SELECT 'table'::text AS kind, n.nspname || '.' || c.relname AS name, NULL::text AS def, NULL::text AS tbl
282    FROM pg_class c JOIN ns n ON n.oid = c.relnamespace WHERE c.relkind IN ('r', 'p', 'f')
283  UNION ALL
284  SELECT 'column', n.nspname || '.' || c.relname || '.' || a.attname,
285         format_type(a.atttypid, a.atttypmod)
286           || CASE WHEN a.attnotnull THEN ' not null' ELSE '' END
287           || CASE WHEN a.attidentity <> '' THEN ' identity ' || a.attidentity::text ELSE '' END
288           || CASE WHEN a.attgenerated <> '' THEN ' generated as (' || pg_get_expr(d.adbin, d.adrelid) || ')'
289                   WHEN d.adbin IS NOT NULL THEN ' default ' || pg_get_expr(d.adbin, d.adrelid)
290                   ELSE '' END,
291         n.nspname || '.' || c.relname
292    FROM pg_attribute a JOIN pg_class c ON c.oid = a.attrelid JOIN ns n ON n.oid = c.relnamespace
293    LEFT JOIN pg_attrdef d ON d.adrelid = a.attrelid AND d.adnum = a.attnum
294    WHERE c.relkind IN ('r', 'p', 'f') AND a.attnum > 0 AND NOT a.attisdropped
295  UNION ALL
296  SELECT 'index', n.nspname || '.' || ic.relname, pg_get_indexdef(i.indexrelid), n.nspname || '.' || tc.relname
297    FROM pg_index i JOIN pg_class ic ON ic.oid = i.indexrelid JOIN pg_class tc ON tc.oid = i.indrelid
298    JOIN ns n ON n.oid = ic.relnamespace
299  UNION ALL
300  SELECT 'constraint', n.nspname || '.' || c.relname || '.' || k.conname, pg_get_constraintdef(k.oid),
301         n.nspname || '.' || c.relname
302    FROM pg_constraint k JOIN pg_class c ON c.oid = k.conrelid JOIN ns n ON n.oid = c.relnamespace
303    WHERE k.contype IN ('p', 'u', 'f', 'c', 'x')
304  UNION ALL
305  SELECT 'view', n.nspname || '.' || c.relname, pg_get_viewdef(c.oid), NULL
306    FROM pg_class c JOIN ns n ON n.oid = c.relnamespace WHERE c.relkind IN ('v', 'm')
307  UNION ALL
308  SELECT 'function', n.nspname || '.' || p.proname || '(' || pg_get_function_identity_arguments(p.oid) || ')',
309         pg_get_functiondef(p.oid), NULL
310    FROM pg_proc p JOIN ns n ON n.oid = p.pronamespace
311    WHERE p.prokind IN ('f', 'p')
312      AND NOT EXISTS (SELECT 1 FROM pg_depend d WHERE d.classid = 'pg_proc'::regclass AND d.objid = p.oid AND d.deptype = 'e')
313  UNION ALL
314  SELECT 'sequence', n.nspname || '.' || c.relname,
315         format_type(s.seqtypid, NULL) || ' start ' || s.seqstart || ' increment ' || s.seqincrement
316           || ' min ' || s.seqmin || ' max ' || s.seqmax || CASE WHEN s.seqcycle THEN ' cycle' ELSE '' END,
317         (SELECT rn.nspname || '.' || rc.relname FROM pg_depend d
318            JOIN pg_class rc ON rc.oid = d.refobjid JOIN pg_namespace rn ON rn.oid = rc.relnamespace
319            WHERE d.classid = 'pg_class'::regclass AND d.objid = c.oid
320              AND d.refclassid = 'pg_class'::regclass AND d.deptype IN ('a', 'i') LIMIT 1)
321    FROM pg_class c JOIN pg_sequence s ON s.seqrelid = c.oid JOIN ns n ON n.oid = c.relnamespace
322  UNION ALL
323  SELECT 'extension', extname::text, extversion, NULL FROM pg_extension
324) x;
325WITH ns AS (
326  SELECT oid, nspname FROM pg_namespace
327  WHERE nspname NOT IN ('pg_catalog', 'information_schema')
328    AND nspname NOT LIKE 'pg\_toast%' AND nspname NOT LIKE 'pg\_temp\_%'
329    AND nspname NOT LIKE 'pg\_toast\_temp\_%'
330)
331SELECT coalesce(jsonb_agg(jsonb_build_array(n.nspname, c.relname, c.reltuples::bigint)
332                          ORDER BY n.nspname, c.relname), '[]'::jsonb)
333  FROM pg_class c JOIN ns n ON n.oid = c.relnamespace WHERE c.relkind IN ('r', 'p');
334"#;
335
336/// `public.t.id` -> `t.id`: the default schema is left off.
337fn short(name: &str) -> String {
338    name.strip_prefix("public.").unwrap_or(name).to_string()
339}
340
341async fn snapshot(side: &Side) -> Result<Snapshot> {
342    let out = side.query(SNAPSHOT_SQL.to_string()).await?;
343    let mut lines = out.lines().filter(|l| !l.trim().is_empty());
344    let bad = |what: &str| internal(format!("unexpected catalog reply ({what})"));
345    let objects: Vec<Obj> = serde_json::from_str(lines.next().ok_or_else(|| bad("objects"))?)?;
346    let tables: Vec<TableEntry> = serde_json::from_str(lines.next().ok_or_else(|| bad("tables"))?)?;
347    let objects = objects
348        .into_iter()
349        .map(|(k, n, d, t)| (k, short(&n), d, t.map(|t| short(&t))))
350        .collect();
351    Ok(Snapshot { objects, tables })
352}
353
354/// Added, removed and changed objects. The columns, indexes, constraints and sequences of an added
355/// or removed table are left out: the table says it.
356fn schema_changes(before: &[Obj], after: &[Obj]) -> Vec<SchemaChange> {
357    let rank = |k: &str| KINDS.iter().position(|x| *x == k).unwrap_or(KINDS.len());
358    let map = |objs: &[Obj]| -> BTreeMap<(usize, String), (Option<String>, Option<String>)> {
359        objs.iter()
360            .map(|(k, n, d, t)| ((rank(k), n.clone()), (d.clone(), t.clone())))
361            .collect()
362    };
363    let (b, a) = (map(before), map(after));
364    let table_rank = rank("table");
365    let names = |m: &BTreeMap<(usize, String), _>| -> HashSet<String> {
366        m.keys()
367            .filter(|(r, _)| *r == table_rank)
368            .map(|(_, n)| n.clone())
369            .collect()
370    };
371    let (bt, at) = (names(&b), names(&a));
372    let whole: HashSet<&String> = bt.symmetric_difference(&at).collect();
373
374    let keys: std::collections::BTreeSet<&(usize, String)> = b.keys().chain(a.keys()).collect();
375    let mut out = Vec::new();
376    for key in keys {
377        let (bv, av) = (b.get(key), a.get(key));
378        let owner = av.or(bv).and_then(|(_, t)| t.as_ref());
379        if owner.is_some_and(|t| whole.contains(t)) {
380            continue;
381        }
382        let (change, bdef, adef) = match (bv, av) {
383            (None, Some((d, _))) => ("added", None, d.clone()),
384            (Some((d, _)), None) => ("removed", d.clone(), None),
385            (Some((bd, _)), Some((ad, _))) if bd != ad => ("changed", bd.clone(), ad.clone()),
386            _ => continue,
387        };
388        out.push(SchemaChange {
389            kind: KINDS.get(key.0).unwrap_or(&"table").to_string(),
390            name: key.1.clone(),
391            change: change.into(),
392            before: bdef,
393            after: adef,
394        });
395    }
396    out
397}
398
399// ---- row counts ---------------------------------------------------------------------------
400
401fn qi(s: &str) -> String {
402    format!("\"{}\"", s.replace('"', "\"\""))
403}
404
405fn ql(s: &str) -> String {
406    format!("'{}'", s.replace('\'', "''"))
407}
408
409fn display_table(schema: &str, rel: &str) -> String {
410    if schema == "public" {
411        rel.to_string()
412    } else {
413        format!("{schema}.{rel}")
414    }
415}
416
417/// Exact counts under a 5 s statement timeout; a table that times out keeps its estimate.
418async fn counts(side: &Side, tables: &[TableEntry]) -> Result<HashMap<String, Count>> {
419    let mut sql = String::from("\\set ON_ERROR_STOP off\nSET statement_timeout = '5s';\n");
420    for (i, (s, r, _)) in tables.iter().enumerate() {
421        sql.push_str(&format!(
422            "SELECT jsonb_build_array({i}, count(*)) FROM {}.{};\n",
423            qi(s),
424            qi(r)
425        ));
426    }
427    let out = side.query(sql).await?;
428    let mut exact: HashMap<usize, u64> = HashMap::new();
429    for line in out.lines().filter(|l| l.starts_with('[')) {
430        if let Ok((i, n)) = serde_json::from_str::<(usize, u64)>(line) {
431            exact.insert(i, n);
432        }
433    }
434    Ok(tables
435        .iter()
436        .enumerate()
437        .map(|(i, (s, r, est))| {
438            let count = match exact.get(&i) {
439                Some(n) => Count {
440                    n: *n,
441                    estimate: false,
442                },
443                None => Count {
444                    n: (*est).max(0) as u64,
445                    estimate: true,
446                },
447            };
448            (display_table(s, r), count)
449        })
450        .collect())
451}
452
453async fn table_counts(
454    before: &Side,
455    after: &Side,
456    bt: &[TableEntry],
457    at: &[TableEntry],
458) -> Result<Vec<TableCount>> {
459    let (b, a) = tokio::try_join!(counts(before, bt), counts(after, at))?;
460    let names: std::collections::BTreeSet<&String> = b.keys().chain(a.keys()).collect();
461    Ok(names
462        .into_iter()
463        .map(|t| TableCount {
464            table: t.clone(),
465            before: b.get(t).cloned(),
466            after: a.get(t).cloned(),
467        })
468        .collect())
469}
470
471// ---- rows ---------------------------------------------------------------------------------
472
473struct TableInfo {
474    pk: Vec<String>,
475    estimate: i64,
476}
477
478/// The key columns and row estimate of `schema.rel`; none when the table does not exist.
479async fn table_info(side: &Side, schema: &str, rel: &str) -> Result<Option<TableInfo>> {
480    let qualified = format!("{}.{}", qi(schema), qi(rel));
481    let sql = format!(
482        "SELECT coalesce((SELECT jsonb_build_object(
483           'pk', (SELECT coalesce(jsonb_agg(a.attname ORDER BY array_position(i.indkey::int2[], a.attnum)), '[]'::jsonb)
484                    FROM pg_index i JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY (i.indkey::int2[])
485                    WHERE i.indrelid = c.oid AND i.indisprimary),
486           'est', CASE WHEN c.reltuples >= 0 THEN c.reltuples::bigint ELSE pg_relation_size(c.oid) / 8192 * 255 END)
487           FROM pg_class c WHERE c.oid = to_regclass({})), 'null'::jsonb);",
488        ql(&qualified)
489    );
490    let v: Value = serde_json::from_str(side.query(sql).await?.trim())?;
491    if v.is_null() {
492        return Ok(None);
493    }
494    let pk = serde_json::from_value(v["pk"].clone())?;
495    Ok(Some(TableInfo {
496        pk,
497        estimate: v["est"].as_i64().unwrap_or(0),
498    }))
499}
500
501fn pk_list(pk: &[String]) -> String {
502    pk.iter().map(|c| qi(c)).collect::<Vec<_>>().join(", ")
503}
504
505/// Sorted by the key's text in byte order, so the merge below can compare keys as plain strings.
506fn hash_sql(q: &str, pk: &[String]) -> String {
507    format!(
508        "SELECT md5(row_to_json(mitosis_row)::text) || ' ' || jsonb_build_array({cols})::text \
509         FROM {q} AS mitosis_row ORDER BY jsonb_build_array({cols})::text COLLATE \"C\";",
510        cols = pk_list(pk)
511    )
512}
513
514/// The lines of one psql query, read as they arrive: `(md5 of the row, key text)`.
515struct RowStream {
516    child: std::process::Child,
517    lines: std::io::Lines<std::io::BufReader<std::process::ChildStdout>>,
518    writer: Option<std::thread::JoinHandle<std::io::Result<()>>>,
519}
520
521impl RowStream {
522    fn open(side: &Side, sql: String) -> Result<RowStream> {
523        use std::io::Write;
524        use std::process::{Command, Stdio};
525        let mut child = Command::new(side.install.bin("psql"))
526            .arg("-h")
527            .arg(&side.socket_dir)
528            .args([
529                "-p",
530                &side.port.to_string(),
531                "-U",
532                "postgres",
533                "-d",
534                &side.db,
535            ])
536            .args(["-v", "ON_ERROR_STOP=1", "-X", "-A", "-t", "-q", "-f", "-"])
537            .env("PGCONNECT_TIMEOUT", "5")
538            .stdin(Stdio::piped())
539            .stdout(Stdio::piped())
540            .stderr(Stdio::piped())
541            .spawn()?;
542        let mut stdin = child.stdin.take().expect("psql stdin is piped");
543        let stdout = child.stdout.take().expect("psql stdout is piped");
544        let writer = std::thread::spawn(move || stdin.write_all(sql.as_bytes()));
545        Ok(RowStream {
546            child,
547            lines: std::io::BufRead::lines(std::io::BufReader::new(stdout)),
548            writer: Some(writer),
549        })
550    }
551
552    fn next_row(&mut self) -> Result<Option<(String, String)>> {
553        loop {
554            match self.lines.next() {
555                None => return Ok(None),
556                Some(Err(e)) => return Err(e.into()),
557                Some(Ok(l)) => {
558                    if let Some((h, k)) = l.split_once(' ') {
559                        return Ok(Some((h.to_string(), k.to_string())));
560                    }
561                }
562            }
563        }
564    }
565
566    /// Waits for psql after its output is read to the end; a failed query is an error.
567    fn finish(&mut self) -> Result<()> {
568        use std::io::Read;
569        let mut err = String::new();
570        if let Some(mut e) = self.child.stderr.take() {
571            let _ = e.read_to_string(&mut err);
572        }
573        let status = self.child.wait()?;
574        let sent = self.writer.take().map(|w| {
575            w.join()
576                .unwrap_or_else(|_| Err(std::io::Error::other("writer panicked")))
577        });
578        if !status.success() {
579            return Err(internal(format!("psql failed: {}", err.trim())));
580        }
581        if let Some(Err(e)) = sent {
582            return Err(internal(format!("could not send SQL to psql: {e}")));
583        }
584        Ok(())
585    }
586}
587
588impl Drop for RowStream {
589    fn drop(&mut self) {
590        let _ = self.child.kill();
591        let _ = self.child.wait();
592    }
593}
594
595#[derive(Default)]
596struct Merged {
597    inserted: Vec<String>,
598    deleted: Vec<String>,
599    updated: Vec<String>,
600    inserted_total: u64,
601    deleted_total: u64,
602    updated_total: u64,
603}
604
605impl Merged {
606    fn push(list: &mut Vec<String>, total: &mut u64, key: String, limit: usize) {
607        *total += 1;
608        if list.len() < limit {
609            list.push(key);
610        }
611    }
612}
613
614/// Merges two key-ordered streams; memory is the `limit` keys kept per kind, not the table.
615fn merge(
616    mut before: Option<RowStream>,
617    mut after: Option<RowStream>,
618    limit: usize,
619) -> Result<Merged> {
620    let next = |s: &mut Option<RowStream>| -> Result<Option<(String, String)>> {
621        match s {
622            Some(s) => s.next_row(),
623            None => Ok(None),
624        }
625    };
626    let mut m = Merged::default();
627    let (mut b, mut a) = (next(&mut before)?, next(&mut after)?);
628    loop {
629        match (b.take(), a.take()) {
630            (None, None) => break,
631            (Some((_, k)), None) => {
632                Merged::push(&mut m.deleted, &mut m.deleted_total, k, limit);
633                b = next(&mut before)?;
634            }
635            (None, Some((_, k))) => {
636                Merged::push(&mut m.inserted, &mut m.inserted_total, k, limit);
637                a = next(&mut after)?;
638            }
639            (Some(bv), Some(av)) => match bv.1.cmp(&av.1) {
640                std::cmp::Ordering::Less => {
641                    Merged::push(&mut m.deleted, &mut m.deleted_total, bv.1, limit);
642                    b = next(&mut before)?;
643                    a = Some(av);
644                }
645                std::cmp::Ordering::Greater => {
646                    Merged::push(&mut m.inserted, &mut m.inserted_total, av.1, limit);
647                    a = next(&mut after)?;
648                    b = Some(bv);
649                }
650                std::cmp::Ordering::Equal => {
651                    if bv.0 != av.0 {
652                        Merged::push(&mut m.updated, &mut m.updated_total, av.1, limit);
653                    }
654                    b = next(&mut before)?;
655                    a = next(&mut after)?;
656                }
657            },
658        }
659    }
660    for s in [&mut before, &mut after].into_iter().flatten() {
661        s.finish()?;
662    }
663    Ok(m)
664}
665
666fn canon(text: &str) -> Result<Value> {
667    Ok(serde_json::from_str(text)?)
668}
669
670/// Whole rows as JSON.
671const ROW_JSON: &str = "to_jsonb(mitosis_row)";
672/// Per column, the md5 of its JSON text: compares values exactly, as the row hash does.
673const COLUMN_HASHES: &str =
674    "(SELECT jsonb_object_agg(k, md5(v::text)) FROM json_each(row_to_json(mitosis_row)) AS e(k, v))";
675
676/// `expr` of the rows with the given key texts, keyed by their canonical key.
677async fn fetch_rows(
678    side: &Side,
679    q: &str,
680    pk: &[String],
681    keys: &[&str],
682    expr: &str,
683) -> Result<HashMap<String, Value>> {
684    if keys.is_empty() {
685        return Ok(HashMap::new());
686    }
687    let lits = keys.iter().map(|k| ql(k)).collect::<Vec<_>>().join(", ");
688    let cols = pk_list(pk);
689    let sql = format!(
690        "SELECT jsonb_build_array(jsonb_build_array({cols}), {expr})::text \
691         FROM {q} AS mitosis_row WHERE jsonb_build_array({cols}) = ANY (ARRAY[{lits}]::jsonb[]);"
692    );
693    let out = side.query(sql).await?;
694    let mut map = HashMap::new();
695    for line in out.lines().filter(|l| l.starts_with('[')) {
696        let Value::Array(mut pair) = canon(line)? else {
697            continue;
698        };
699        if pair.len() == 2 {
700            let row = pair.pop().unwrap_or(Value::Null);
701            map.insert(pair[0].to_string(), row);
702        }
703    }
704    Ok(map)
705}
706
707/// A one-column key is shown as its value, a composite key as an array.
708fn pk_value(key: &str) -> Result<Value> {
709    Ok(match canon(key)? {
710        Value::Array(mut a) if a.len() == 1 => a.pop().unwrap_or(Value::Null),
711        v => v,
712    })
713}
714
715fn changed_columns(b: &Value, a: &Value) -> Vec<String> {
716    let (Some(b), Some(a)) = (b.as_object(), a.as_object()) else {
717        return Vec::new();
718    };
719    let mut cols: Vec<String> = b
720        .keys()
721        .chain(a.keys().filter(|k| !b.contains_key(*k)))
722        .filter(|k| b.get(*k) != a.get(*k))
723        .cloned()
724        .collect();
725    cols.sort();
726    cols
727}
728
729async fn row_diff(before: &Side, after: &Side, table: &str, limit: usize) -> Result<RowDiff> {
730    let (schema, rel) = table.split_once('.').unwrap_or(("public", table));
731    let q = format!("{}.{}", qi(schema), qi(rel));
732    let (bi, ai) = tokio::try_join!(
733        table_info(before, schema, rel),
734        table_info(after, schema, rel)
735    )?;
736    let pk = match (&bi, &ai) {
737        (None, None) => {
738            return Err(invalid(format!("table {table} does not exist"))
739                .with_hint("use `schema.table` for a table outside the public schema"))
740        }
741        (Some(b), Some(a)) if b.pk != a.pk => {
742            return Err(invalid(format!(
743                "the primary key of {table} differs between the template and the clone"
744            ))
745            .with_hint("a row diff needs the same key on both sides"))
746        }
747        (_, Some(i)) | (Some(i), None) => i.pk.clone(),
748    };
749    if pk.is_empty() {
750        return Err(invalid(format!("table {table} has no primary key"))
751            .with_hint("rows are matched by primary key; use the row counts instead"));
752    }
753    let biggest = bi
754        .iter()
755        .chain(ai.iter())
756        .map(|i| i.estimate)
757        .max()
758        .unwrap_or(0);
759    if biggest > MAX_TABLE_ROWS {
760        return Err(invalid(format!(
761            "table {table} has about {biggest} rows; a row diff stops at {MAX_TABLE_ROWS}"
762        ))
763        .with_hint("diff a smaller table, or compare the row counts"));
764    }
765
766    let (bsql, asql) = (hash_sql(&q, &pk), hash_sql(&q, &pk));
767    let (bs, as_) = (
768        bi.is_some().then(|| (before.clone(), bsql)),
769        ai.is_some().then(|| (after.clone(), asql)),
770    );
771    let m = blocking(move || {
772        let open = |s: Option<(Side, String)>| {
773            s.map(|(side, sql)| RowStream::open(&side, sql)).transpose()
774        };
775        merge(open(bs)?, open(as_)?, limit)
776    })
777    .await?;
778
779    let (del, ins, upd) = (
780        m.deleted.iter().map(String::as_str).collect::<Vec<_>>(),
781        m.inserted.iter().map(String::as_str).collect::<Vec<_>>(),
782        m.updated.iter().map(String::as_str).collect::<Vec<_>>(),
783    );
784    let (bcols, acols) = tokio::try_join!(
785        fetch_rows(before, &q, &pk, &upd, COLUMN_HASHES),
786        fetch_rows(after, &q, &pk, &upd, COLUMN_HASHES)
787    )?;
788    let (brows, arows) = tokio::try_join!(
789        fetch_rows(before, &q, &pk, &del, ROW_JSON),
790        fetch_rows(after, &q, &pk, &ins, ROW_JSON)
791    )?;
792    let by_key = |m: &HashMap<String, Value>, k: &str| -> Result<Value> {
793        Ok(m.get(&canon(k)?.to_string())
794            .cloned()
795            .unwrap_or(Value::Null))
796    };
797    let mut out = RowDiff {
798        table: table.to_string(),
799        inserted: Vec::new(),
800        deleted: Vec::new(),
801        updated: Vec::new(),
802        inserted_total: m.inserted_total,
803        deleted_total: m.deleted_total,
804        updated_total: m.updated_total,
805    };
806    for k in &ins {
807        out.inserted.push(by_key(&arows, k)?);
808    }
809    for k in &del {
810        out.deleted.push(by_key(&brows, k)?);
811    }
812    for k in &upd {
813        out.updated.push(RowChange {
814            pk: pk_value(k)?,
815            columns: changed_columns(&by_key(&bcols, k)?, &by_key(&acols, k)?),
816        });
817    }
818    Ok(out)
819}
820
821// ---- human output -------------------------------------------------------------------------
822
823fn count_text(c: &Option<Count>) -> String {
824    match c {
825        Some(c) => format!("{}{}", if c.estimate { "~" } else { "" }, c.n),
826        None => "(none)".into(),
827    }
828}
829
830/// The report as text: Schema, Row counts and (with `--table`) Rows.
831pub fn render(r: &DiffReport) -> String {
832    let mut o = format!(
833        "clone {} against template {} generation {}\n\nSchema\n",
834        r.clone, r.template, r.generation
835    );
836    if r.schema.is_empty() {
837        o.push_str("  no changes\n");
838    }
839    for c in &r.schema {
840        let mark = match c.change.as_str() {
841            "added" => '+',
842            "removed" => '-',
843            _ => '~',
844        };
845        o.push_str(&format!("  {mark} {} {}\n", c.kind, c.name));
846        if c.change == "changed" {
847            for (label, v) in [("before", &c.before), ("after", &c.after)] {
848                if let Some(v) = v {
849                    o.push_str(&format!(
850                        "      {label}: {}\n",
851                        v.replace('\n', "\n        ")
852                    ));
853                }
854            }
855        } else if c.kind == "column" {
856            if let Some(v) = c.after.as_ref().or(c.before.as_ref()) {
857                o.push_str(&format!("      {v}\n"));
858            }
859        }
860    }
861
862    o.push_str("\nRow counts (~ is an estimate)\n");
863    let shown: Vec<&TableCount> = r
864        .tables
865        .iter()
866        .filter(|t| {
867            t.before != t.after
868                || t.before.as_ref().is_some_and(|c| c.estimate)
869                || t.after.as_ref().is_some_and(|c| c.estimate)
870        })
871        .collect();
872    if shown.is_empty() {
873        o.push_str("  no changes\n");
874    }
875    for t in &shown {
876        o.push_str(&format!(
877            "  {}  {} -> {}\n",
878            t.table,
879            count_text(&t.before),
880            count_text(&t.after)
881        ));
882    }
883    let same = r.tables.len() - shown.len();
884    if same > 0 {
885        o.push_str(&format!("  ({same} tables unchanged)\n"));
886    }
887
888    if let Some(d) = &r.rows {
889        o.push_str(&format!(
890            "\nRows of {}: {} inserted, {} deleted, {} updated\n",
891            d.table, d.inserted_total, d.deleted_total, d.updated_total
892        ));
893        for row in &d.inserted {
894            o.push_str(&format!("  + {row}\n"));
895        }
896        for row in &d.deleted {
897            o.push_str(&format!("  - {row}\n"));
898        }
899        for u in &d.updated {
900            o.push_str(&format!("  ~ {}: {}\n", u.pk, u.columns.join(", ")));
901        }
902        let shown = (d.inserted.len() + d.deleted.len() + d.updated.len()) as u64;
903        if shown < d.inserted_total + d.deleted_total + d.updated_total {
904            o.push_str("  (more rows exist; raise --limit)\n");
905        }
906    }
907    o
908}
909
910#[cfg(test)]
911mod tests {
912    use super::*;
913
914    fn obj(k: &str, n: &str, d: Option<&str>, t: Option<&str>) -> Obj {
915        (k.into(), n.into(), d.map(Into::into), t.map(Into::into))
916    }
917
918    #[test]
919    fn added_table_hides_its_columns_but_a_changed_index_shows() {
920        let before = vec![
921            obj("table", "a", None, None),
922            obj("index", "a_i", Some("on x"), Some("a")),
923        ];
924        let after = vec![
925            obj("table", "a", None, None),
926            obj("table", "b", None, None),
927            obj("column", "b.id", Some("integer"), Some("b")),
928            obj("index", "a_i", Some("on y"), Some("a")),
929        ];
930        let got = schema_changes(&before, &after);
931        let got: Vec<(&str, &str, &str)> = got
932            .iter()
933            .map(|c| (c.kind.as_str(), c.name.as_str(), c.change.as_str()))
934            .collect();
935        assert_eq!(got, [("table", "b", "added"), ("index", "a_i", "changed")]);
936    }
937
938    #[test]
939    fn short_drops_public_only() {
940        assert_eq!(short("public.t.id"), "t.id");
941        assert_eq!(short("app.t"), "app.t");
942    }
943
944    #[test]
945    fn changed_columns_lists_differences_and_missing() {
946        let b = serde_json::json!({"id": 1, "a": "x", "gone": 1});
947        let a = serde_json::json!({"id": 1, "a": "y", "new": 2});
948        assert_eq!(changed_columns(&b, &a), ["a", "gone", "new"]);
949    }
950}