crates/knut-flight/tests/live_data.rs
//! End-to-end proof of the **live-data seam** (`Catalog::from_parquet_dir`): the
//! server serving a REAL graph read off disk, not the in-RAM bench generator.
//!
//! The flow mirrors what korp's `iceberg2graph` step does — land the source
//! tables as parquet, then serve the graph the rel2graph engine builds from them:
//!
//! 1. generate the seeded 12-table ERP database and write its source tables to
//! `<tmp>/<table>.parquet` (`write_source_parquet` — the export bin's `--parquet`);
//! 2. build a `Catalog` with `from_parquet_dir` over that directory + the bench
//! mapping spec — this READS the parquet back and runs the SAME
//! `knut_bifrost::rel2graph::build_mutation` engine over it;
//! 3. assert the served graph is byte-for-byte the shape the bench catalog serves
//! (35 datasets, 1106 nodes, 2955 resolved edges, no REPORTS_TO self-loop);
//! 4. stand the server up on an ephemeral port and pull `nodes`/`edges`/a source
//! table with a real `FlightClient` — proving a client pull of real data.
use std::net::SocketAddr;
use std::path::PathBuf;
use arrow_array::{Int64Array, StringArray};
use arrow_flight::flight_service_server::FlightServiceServer;
use arrow_flight::{FlightClient, Ticket};
use futures::TryStreamExt;
use knut_flight::{Catalog, GraphFlightService};
use knut_rel2graph_bench::{DEFAULT_SEED, generate, mapping_spec, write_source_parquet};
use tonic::transport::server::TcpIncoming;
use tonic::transport::{Channel, Server};
/// A unique scratch dir under the OS temp dir; removed if it already exists.
fn scratch(tag: &str) -> PathBuf {
let dir = std::env::temp_dir().join(format!("knut-flight-live-{tag}-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).expect("create scratch dir");
dir
}
/// Write the scale-1 bench source tables to parquet and build a live catalog over
/// them — the shared setup for the tests below.
fn live_catalog(dir: &std::path::Path) -> Catalog {
let db = generate(1, DEFAULT_SEED);
let paths = write_source_parquet(&db, dir).expect("write source parquet");
assert_eq!(paths.len(), 12, "12 source tables written as parquet");
Catalog::from_parquet_dir(&mapping_spec(), dir).expect("build catalog from parquet export")
}
async fn spawn(catalog: Catalog) -> FlightClient {
let incoming = TcpIncoming::bind("127.0.0.1:0".parse::<SocketAddr>().unwrap())
.expect("bind ephemeral port");
let addr: SocketAddr = incoming.local_addr().unwrap();
let service = GraphFlightService::new(catalog);
tokio::spawn(async move {
Server::builder()
.add_service(FlightServiceServer::new(service))
.serve_with_incoming(incoming)
.await
.unwrap();
});
let channel = Channel::from_shared(format!("http://{addr}"))
.unwrap()
.connect()
.await
.expect("connect to flight server");
FlightClient::new(channel)
}
#[test]
fn parquet_catalog_equals_bench_catalog_shape() {
let dir = scratch("shape");
let live = live_catalog(&dir);
let bench = Catalog::from_bench(1);
// Same dataset count and same key datasets — the live path is the bench path
// with the data swapped underneath the single seam.
assert_eq!(
live.len(),
bench.len(),
"same dataset count as the bench catalog"
);
assert_eq!(live.len(), 35, "2 aggregate + 10 node + 11 edge + 12 table");
let nodes = live.get("nodes").expect("nodes flight");
let node_rows: usize = nodes.batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(node_rows, 1106, "scale-1 node total from parquet");
let edges = live.get("edges").expect("edges flight");
let edge_rows: usize = edges.batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(edge_rows, 2955, "scale-1 resolved-edge total from parquet");
assert!(
live.get("edge/Employee_REPORTS_TO_Employee").is_some(),
"self-ref edge flight present from real data"
);
assert!(
live.get("table/erp.hr.employees").is_some(),
"source table served from parquet"
);
let _ = std::fs::remove_dir_all(&dir);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn client_pulls_real_graph_from_parquet() {
let dir = scratch("pull");
let mut client = spawn(live_catalog(&dir)).await;
// nodes over the parquet-backed catalog.
let nodes = client
.do_get(Ticket::new("nodes".as_bytes().to_vec()))
.await
.expect("do_get nodes")
.try_collect::<Vec<_>>()
.await
.expect("collect node batches");
let node_rows: usize = nodes.iter().map(|b| b.num_rows()).sum();
assert_eq!(node_rows, 1106, "client pulled 1106 real nodes");
// edges + the a7208dc self-ref proof, over real parquet data.
let edges = client
.do_get(Ticket::new("edges".as_bytes().to_vec()))
.await
.expect("do_get edges")
.try_collect::<Vec<_>>()
.await
.expect("collect edge batches");
let edge_rows: usize = edges.iter().map(|b| b.num_rows()).sum();
assert_eq!(edge_rows, 2955, "client pulled 2955 real resolved edges");
let mut reports_to = 0;
for b in &edges {
let rel = b
.column(b.schema().index_of("rel").unwrap())
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
let src = b
.column(b.schema().index_of("src").unwrap())
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let dst = b
.column(b.schema().index_of("dst").unwrap())
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
for i in 0..b.num_rows() {
if rel.value(i) == "REPORTS_TO" {
reports_to += 1;
assert_ne!(src.value(i), dst.value(i), "REPORTS_TO is not a self-loop");
}
}
}
assert_eq!(
reports_to, 21,
"REPORTS_TO edges from real data (24 employees - 3 roots)"
);
// A source table pulled back — the raw relational rows that came off disk.
let table = client
.do_get(Ticket::new("table/erp.hr.employees".as_bytes().to_vec()))
.await
.expect("do_get table")
.try_collect::<Vec<_>>()
.await
.expect("collect table batches");
let rows: usize = table.iter().map(|b| b.num_rows()).sum();
assert_eq!(rows, 24, "24 employees read back from parquet");
assert!(
table[0].schema().index_of("manager_id").is_ok(),
"raw relational schema"
);
let _ = std::fs::remove_dir_all(&dir);
}