nordisk/nornir

crates/mimir-ingest/examples/keys.rs

//! Every OPEN identity of a lake, as text: `keys <lake-dir> [--universe]` →
//! one line per open row — `table ⇥ repo ⇥ owner ⇥ key columns…`, sorted — and
//! a count per (table, repo) on stderr.
//!
//! Two lakes consumed from the same commit hold the same lines whoever dumped
//! them (the driver on this box, the driver in the toolchain image): `diff`
//! of two outputs is the comparison. `--universe` adds the ring (dependencies
//! and std), without it only the repos' own rows are printed.

use std::collections::BTreeMap;
use std::io::Write;

use anyhow::Result;
use mimir_ingest::{EventLog, FoldScope, Kind};
use skade::arrow_array::{Array, StringArray};
use skade::arrow_schema::DataType;

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<()> {
    let a: Vec<String> = std::env::args().skip(1).collect();
    anyhow::ensure!(!a.is_empty(), "usage: keys <lake-dir> [--universe]");
    let universe = a.iter().any(|x| x == "--universe");
    let log = EventLog::open_lake(std::path::Path::new(&a[0])).await?;
    let repos: BTreeMap<String, i64> = log.offsets().await?.into_iter().map(|(r, o)| (r, o.seq)).collect();
    let scope = FoldScope { repos, universe };
    let mut lines: Vec<String> = Vec::new();
    let mut counts: BTreeMap<(String, String), usize> = BTreeMap::new();
    for kind in Kind::ALL {
        let table = EventLog::table_name(kind);
        let rows = log.fold_rows(kind, &scope, &[]).await?;
        let keys = log.spec(kind).key_columns.clone();
        let text = |b: &skade::arrow_array::RecordBatch, name: &str| -> Result<StringArray> {
            let c = skade::arrow_cast::cast(b.column_by_name(name).ok_or_else(|| anyhow::anyhow!("{table}: no `{name}`"))?, &DataType::Utf8)?;
            Ok(c.as_any().downcast_ref::<StringArray>().expect("cast to utf8").clone())
        };
        let mut cols = Vec::new();
        for b in &rows.batches {
            let mut v = vec![text(b, "repo")?, text(b, "krate")?];
            for k in &keys {
                v.push(text(b, k)?);
            }
            cols.push(v);
        }
        for &(b, r) in &rows.open {
            let c = &cols[b as usize];
            let cell = |a: &StringArray| if a.is_null(r as usize) { String::new() } else { a.value(r as usize).to_string() };
            *counts.entry((table.to_string(), cell(&c[0]))).or_default() += 1;
            lines.push(std::iter::once(table.to_string()).chain(c.iter().map(cell)).collect::<Vec<_>>().join("\t"));
        }
    }
    lines.sort_unstable();
    let mut out = std::io::BufWriter::new(std::io::stdout().lock());
    for l in &lines {
        writeln!(out, "{l}")?;
    }
    for ((table, repo), n) in counts {
        eprintln!("{n:>9}  {table}  {repo}");
    }
    Ok(())
}