crates/knut-bifrost/tests/integration_osm_filesource.rs
//! OSM GeoParquet → FileSource → graph, with an Airflow asset event.
//!
//! There are two tests here, by design:
//!
//! 1. [`osm_geoparquet_slice_becomes_graph_vertices`] — the **synthetic,
//! always-run** test. It builds an OSM-`relations`-shaped GeoParquet chunk
//! *in-repo* (id / WKB `geometry` / `tags` / `version`), drops it into a
//! watched directory, and lets bifröst's [`FileSource`] materialise it
//! through the real mapping pipeline into an in-memory recording sink — no
//! container, no real planet extract. It injects a known row count and
//! asserts the synced rows, the ingested vertices, and the emitted asset
//! event payload all match. Runs by default and on release.
//!
//! 2. [`osm_relations_become_falkordb_nodes`] — the **heavy, opt-in** test
//! (`#[ignore]`). Same path, but against a *real* multi-GB planet
//! `relations.parquet` and a real FalkorDB container. It is gated behind
//! `--ignored` and **panics loudly** if its prerequisites are missing, so it
//! can never silently pass.
//!
//! ## Prerequisites (only for the `#[ignore]` test)
//!
//! - Podman socket: `systemctl --user start podman.socket`
//! - FalkorDB image built locally:
//! ```bash
//! podman build -t knut-falkordb-valkey9 -f containers/falkordb-valkey9/Containerfile .
//! ```
//! - OSM GeoParquet path in `KNUT_OSM_RELATIONS`, e.g.
//! `/some/path/planet.geoparquet/relations.parquet`
//!
//! ## Run
//!
//! ```bash
//! cargo test -p knut-bifrost --test integration_osm_filesource # synthetic
//! KNUT_OSM_RELATIONS=/path/relations.parquet \
//! cargo test -p knut-bifrost --test integration_osm_filesource -- --ignored --nocapture
//! ```
use std::sync::{Arc, Mutex};
use std::time::Duration;
use arrow_array::{BinaryArray, Int32Array, Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use async_trait::async_trait;
use knut_bifrost::{
AssetEmitter, AssetEvent, GraphMutation, Listener, MappingSpec, SinkReport, SyncState,
TransformRegistry, emit_round, sink::GraphSink, source::FileSource,
};
use testcontainers::{GenericImage, ImageExt, core::ContainerPort, runners::AsyncRunner};
/// Emit a functional-status row into the nornir test matrix. Feature ON →
/// recorded; OFF → inlined no-op (and the crate isn't built into this target).
#[cfg(feature = "testmatrix")]
fn fs(check: &str, ok: bool, detail: &str) {
nornir_testmatrix::functional_status("knut-bifrost-osm", check, ok, detail);
}
#[cfg(not(feature = "testmatrix"))]
fn fs(_check: &str, _ok: bool, _detail: &str) {}
/// Rows of OSM relations to ingest. A few thousand is enough to prove the
/// path while reading only the first row group of a (multi-GB) parquet.
const ROWS: usize = 5_000;
const MAPPING_YAML: &str = r#"
version: 1
tables:
- table: lake.osm.relations
mode: incremental
pipeline:
- vertex:
label: Relation
id_column: id
properties: [version, tags]
"#;
fn ensure_docker_host() {
if std::env::var("DOCKER_HOST").is_err() {
let uid = std::process::Command::new("id")
.arg("-u")
.output()
.ok()
.and_then(|o| String::from_utf8(o.stdout).ok())
.and_then(|s| s.trim().parse::<u32>().ok())
.unwrap_or(1000);
let sock = format!("/run/user/{uid}/podman/podman.sock");
if std::path::Path::new(&sock).exists() {
unsafe { std::env::set_var("DOCKER_HOST", format!("unix://{sock}")) };
}
}
}
/// Build a synthetic OSM-`relations`-shaped GeoParquet slice in-repo.
///
/// Same scalar columns the mapping reads — `id` (i64), `tags` (utf8 JSON),
/// `version` (i32) — plus a `geometry` WKB-point binary column to mirror the
/// real planet table's shape (and prove the binary column is harmlessly
/// skipped by the mapping). Deterministic: same `n` → identical content.
fn synthetic_relations(n: usize) -> RecordBatch {
let ids: Vec<i64> = (1..=n as i64).collect();
let tags: Vec<String> = ids
.iter()
.map(|i| format!(r#"{{"type":"multipolygon","seq":"{i}"}}"#))
.collect();
let versions: Vec<i32> = ids.iter().map(|i| (*i % 7 + 1) as i32).collect();
// Minimal WKB point (little-endian, type=1) with deterministic coords, so
// the column is real binary the reader must carry but the mapping ignores.
let geoms: Vec<Vec<u8>> = ids
.iter()
.map(|i| {
let lon = 11.0 + (*i as f64 % 100.0) * 0.001;
let lat = 57.0 + (*i as f64 % 100.0) * 0.001;
let mut b = Vec::with_capacity(21);
b.push(1u8); // little-endian
b.extend_from_slice(&1u32.to_le_bytes()); // wkbPoint
b.extend_from_slice(&lon.to_le_bytes());
b.extend_from_slice(&lat.to_le_bytes());
b
})
.collect();
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("geometry", DataType::Binary, true),
Field::new("tags", DataType::Utf8, true),
Field::new("version", DataType::Int32, true),
])),
vec![
Arc::new(Int64Array::from(ids)),
Arc::new(BinaryArray::from(
geoms.iter().map(|g| g.as_slice()).collect::<Vec<_>>(),
)),
Arc::new(StringArray::from(tags)),
Arc::new(Int32Array::from(versions)),
],
)
.unwrap()
}
/// Read the first `ROWS` rows of a real OSM relations parquet, projecting only
/// the scalar columns we map (`id`, `tags`, `version` — skipping `geometry`
/// binary and `members`). Reads row group 0 only, so it never scans the
/// whole multi-GB file.
fn read_osm_slice(path: &std::path::Path) -> Vec<RecordBatch> {
use parquet::arrow::ProjectionMask;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
let file = std::fs::File::open(path).unwrap_or_else(|e| panic!("open {}: {e}", path.display()));
let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
// relations schema: id(0) geometry(1) tags(2) members(3) version(4).
let mask = ProjectionMask::roots(builder.parquet_schema(), [0, 2, 4]);
let reader = builder
.with_row_groups(vec![0])
.with_projection(mask)
.with_limit(ROWS)
.build()
.unwrap();
reader.map(|b| b.unwrap()).collect()
}
#[derive(Default)]
struct RecordingEmitter {
seen: Arc<Mutex<Vec<AssetEvent>>>,
}
#[async_trait]
impl AssetEmitter for RecordingEmitter {
async fn emit(&self, events: &[AssetEvent]) -> knut_bifrost::Result<()> {
self.seen.lock().unwrap().extend_from_slice(events);
Ok(())
}
}
/// An in-memory graph sink — the "ingested" side, no FalkorDB needed.
#[derive(Clone, Default)]
struct RecordingSink {
vertices: Arc<Mutex<usize>>,
}
#[async_trait]
impl GraphSink for RecordingSink {
async fn apply(&mut self, mutation: &GraphMutation) -> knut_bifrost::Result<SinkReport> {
let v: usize = mutation.vertices.iter().map(|s| s.batch.num_rows()).sum();
*self.vertices.lock().unwrap() += v;
let e: usize = mutation.edges.iter().map(|s| s.batch.num_rows()).sum();
let d: usize = mutation.deletes.iter().map(|s| s.batch.num_rows()).sum();
Ok(SinkReport {
vertices_written: v,
edges_written: e,
deleted: d,
})
}
}
/// **Synthetic, always-run.** OSM-relations-shaped GeoParquet → FileSource →
/// mapping → in-memory sink → asset event. Real input rows, real assertions.
#[tokio::test]
async fn osm_geoparquet_slice_becomes_graph_vertices() {
const N: usize = 1_234; // a known, non-round count we assert end-to-end.
// ── stage the synthetic OSM slice as a chunk in a watched directory ──────
let drops = tempfile::tempdir().unwrap();
let table_dir = drops.path().join("lake.osm.relations");
std::fs::create_dir_all(&table_dir).unwrap();
let batch = synthetic_relations(N);
assert_eq!(batch.num_rows(), N, "fixture has the injected row count");
graphar::io::write_chunk(
table_dir.join("chunk0.parquet"),
std::slice::from_ref(&batch),
&graphar::FileType::Parquet,
)
.unwrap();
// ── bifröst sync: FileSource → mapping → in-memory sink ──────────────────
let sink = RecordingSink::default();
let state_dir = tempfile::tempdir().unwrap();
let mut listener = Listener::new(
MappingSpec::from_yaml(MAPPING_YAML).unwrap(),
TransformRegistry::new(),
FileSource::new(drops.path()),
sink.clone(),
SyncState::open(state_dir.path().join("state.json")).unwrap(),
);
let round = listener.sync_once().await.unwrap();
assert_eq!(round.failed, vec![], "no sync failures");
assert_eq!(round.synced.len(), 1, "the relations table synced");
assert_eq!(round.synced[0].0, "lake.osm.relations", "right table");
assert_eq!(round.synced[0].1, N, "all injected rows shipped");
assert_eq!(
*sink.vertices.lock().unwrap(),
N,
"every relation became a graph vertex"
);
// ── announce the round as an Airflow asset event ─────────────────────────
let emitter = RecordingEmitter::default();
emit_round(&emitter, "graph://osm/", &round).await.unwrap();
let events = emitter.seen.lock().unwrap();
assert_eq!(events.len(), 1, "one event for the one synced table");
assert_eq!(events[0].uri, "graph://osm/lake.osm.relations");
assert_eq!(events[0].extra["rows"], N, "event carries the row count");
assert_eq!(
events[0].extra["vertices_written"], N,
"event carries the vertex tally"
);
drop(events);
// ── idle round: no new files → nothing re-shipped ────────────────────────
let idle = listener.sync_once().await.unwrap();
assert!(idle.synced.is_empty(), "no new files → idle round");
let v = *sink.vertices.lock().unwrap();
assert_eq!(v, N, "idle round did not re-ingest");
fs(
"osm_geoparquet_slice_becomes_graph_vertices",
round.synced.len() == 1 && round.synced[0].1 == N && v == N && idle.synced.is_empty(),
&format!("{N} OSM relation rows → {v} vertices + 1 asset event; idle re-ship none"),
);
}
async fn count(conn: &mut redis::aio::MultiplexedConnection, graph: &str, q: &str) -> i64 {
fn walk(v: &redis::Value) -> Option<i64> {
match v {
redis::Value::Int(i) => Some(*i),
redis::Value::Array(items) | redis::Value::Set(items) => items.iter().find_map(walk),
redis::Value::BulkString(b) => String::from_utf8_lossy(b).parse().ok(),
_ => None,
}
}
let v: redis::Value = redis::cmd("GRAPH.QUERY")
.arg(graph)
.arg(q)
.query_async(conn)
.await
.unwrap();
walk(&v).unwrap_or(-1)
}
/// **Heavy, opt-in (`#[ignore]`).** Same path against a *real* planet
/// `relations.parquet` + a real FalkorDB. Set `KNUT_OSM_RELATIONS` to the
/// parquet path; the test **panics** (does not silently pass) if it is missing.
#[tokio::test]
#[ignore = "needs podman + FalkorDB image + real OSM geoparquet (KNUT_OSM_RELATIONS); run with --ignored"]
async fn osm_relations_become_falkordb_nodes() {
use knut_bifrost::sink::FalkorDbSink;
ensure_docker_host();
// No hardcoded local path: the real dataset is supplied via env var, and a
// missing file is a hard failure — this opt-in test never silently passes.
let osm_path = std::env::var("KNUT_OSM_RELATIONS").expect(
"set KNUT_OSM_RELATIONS to a real OSM relations.parquet to run this --ignored test",
);
let osm_path = std::path::PathBuf::from(osm_path);
assert!(
osm_path.exists(),
"KNUT_OSM_RELATIONS points at a missing file: {}",
osm_path.display()
);
// ── stage the OSM slice as a chunk in a watched directory ────────────────
let drops = tempfile::tempdir().unwrap();
let table_dir = drops.path().join("lake.osm.relations");
std::fs::create_dir_all(&table_dir).unwrap();
let batches = read_osm_slice(&osm_path);
let staged: usize = batches.iter().map(|b| b.num_rows()).sum();
assert!(staged > 0, "read some OSM rows");
graphar::io::write_chunk(
table_dir.join("chunk0.parquet"),
&batches,
&graphar::FileType::Parquet,
)
.unwrap();
// ── FalkorDB container ───────────────────────────────────────────────────
let falkor_image = std::env::var("KNUT_FALKORDB_IMAGE")
.unwrap_or_else(|_| "localhost/knut-falkordb-valkey9:latest".into());
let (fname, ftag) = falkor_image.rsplit_once(':').unwrap();
let falkor = GenericImage::new(fname, ftag)
.with_exposed_port(ContainerPort::Tcp(6379))
.with_startup_timeout(Duration::from_secs(120))
.start()
.await
.expect("start falkordb (build the image first)");
let falkor_port = falkor.get_host_port_ipv4(6379).await.unwrap();
let falkor_url = format!("redis://127.0.0.1:{falkor_port}");
// ── bifröst sync: FileSource → fabric → FalkorDB ─────────────────────────
let state_dir = tempfile::tempdir().unwrap();
let mut listener = Listener::new(
MappingSpec::from_yaml(MAPPING_YAML).unwrap(),
TransformRegistry::new(),
FileSource::new(drops.path()),
FalkorDbSink::connect(&falkor_url, "osm").await.unwrap(),
SyncState::open(state_dir.path().join("state.json")).unwrap(),
);
let round = listener.sync_once().await.unwrap();
assert_eq!(round.failed, vec![], "sync failures");
assert_eq!(round.synced.len(), 1, "the relations table synced");
assert_eq!(round.synced[0].1, staged, "all staged rows shipped");
// ── announce the round as an Airflow asset event ─────────────────────────
let emitter = RecordingEmitter::default();
emit_round(&emitter, "graph://osm/", &round).await.unwrap();
let events = emitter.seen.lock().unwrap();
assert_eq!(events.len(), 1);
assert_eq!(events[0].uri, "graph://osm/lake.osm.relations");
assert_eq!(events[0].extra["rows"], staged);
drop(events);
// ── verify the graph ─────────────────────────────────────────────────────
let client = redis::Client::open(falkor_url.clone()).unwrap();
let mut conn = client.get_multiplexed_async_connection().await.unwrap();
assert_eq!(
count(&mut conn, "osm", "MATCH (n:Relation) RETURN count(n)").await,
staged as i64,
"every OSM relation became a node",
);
// Idle round: no new files → nothing re-shipped.
let idle = listener.sync_once().await.unwrap();
assert!(idle.synced.is_empty(), "no new files → idle");
}