crates/knut-bifrost/tests/integration_spark_falkordb.rs
//! TOP GUN: the full bifröst pipeline against live containers.
//!
//! Spark 4.2+ Connect + Iceberg 1.11.0 ──sync──► FalkorDB (Valkey 9)
//!
//! Rows inserted into an Iceberg table through Spark SQL come out as nodes
//! and relationships in FalkorDB, twice: initial sync + incremental delta.
//!
//! ## Prerequisites
//!
//! - Podman socket: `systemctl --user start podman.socket`
//! - Both images built locally first:
//! ```bash
//! podman build -t knut-spark-iceberg -f containers/spark-iceberg/Containerfile .
//! podman build -t knut-falkordb-valkey9 -f containers/falkordb-valkey9/Containerfile .
//! ```
//!
//! ## Run
//!
//! ```bash
//! cargo test -p knut-bifrost --features spark \
//! --test integration_spark_falkordb -- --ignored --nocapture
//! ```
#![cfg(feature = "spark")]
use std::time::Duration;
use knut_bifrost::{
Listener, MappingSpec, SyncState, TransformRegistry, sink::FalkorDbSink,
source::SparkConnectSource,
};
use testcontainers::{GenericImage, ImageExt, core::ContainerPort, runners::AsyncRunner};
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}")) };
}
}
}
const MAPPING_YAML: &str = r#"
version: 1
tables:
- table: lake.itest.users
mode: incremental
pipeline:
- vertex:
label: Person
id_column: user_id
properties: [name]
- fk_edges:
type: WORKS_AT
src: { label: Person, column: user_id }
dst: { label: Company, column: company_id }
- table: lake.itest.companies
mode: full
pipeline:
- vertex:
label: Company
id_column: company_id
properties: [company_name]
"#;
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)
}
#[tokio::test]
#[ignore = "needs podman + locally built images; run with --ignored"]
async fn iceberg_rows_become_falkordb_graph() {
ensure_docker_host();
// ── containers ──────────────────────────────────────────────────────────
let spark = GenericImage::new("localhost/knut-spark-iceberg", "latest")
.with_exposed_port(ContainerPort::Tcp(15002))
.with_startup_timeout(Duration::from_secs(300))
.start()
.await
.expect("start spark-iceberg (build it first: podman build -t knut-spark-iceberg -f containers/spark-iceberg/Containerfile .)");
let spark_port = spark.get_host_port_ipv4(15002).await.unwrap();
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 it first: podman build -t knut-falkordb-valkey9 -f containers/falkordb-valkey9/Containerfile .)");
let falkor_port = falkor.get_host_port_ipv4(6379).await.unwrap();
let falkor_url = format!("redis://127.0.0.1:{falkor_port}");
// ── Spark Connect: wait until the server answers a real query ───────────
// (building a session succeeds before the JVM is up, so probe with SQL)
let source = {
let url = format!("sc://127.0.0.1:{spark_port}");
let mut attempt = 0;
loop {
let probe = match SparkConnectSource::connect(&url).await {
Ok(s) => match s.sql("SELECT 1").await {
Ok(_) => Some(s),
Err(_) => None,
},
Err(_) => None,
};
match probe {
Some(s) => break s,
None if attempt < 60 => {
attempt += 1;
tokio::time::sleep(Duration::from_secs(3)).await;
}
None => panic!("Spark Connect never became ready after 180s"),
}
}
};
// ── Iceberg fixtures through Spark SQL ──────────────────────────────────
source
.sql("CREATE NAMESPACE IF NOT EXISTS lake.itest")
.await
.unwrap();
source
.sql("CREATE TABLE lake.itest.users (user_id BIGINT, name STRING, company_id BIGINT) USING iceberg")
.await
.unwrap();
source
.sql("CREATE TABLE lake.itest.companies (company_id BIGINT, company_name STRING) USING iceberg")
.await
.unwrap();
source
.sql("INSERT INTO lake.itest.companies VALUES (10, 'acme'), (11, 'globex')")
.await
.unwrap();
source
.sql("INSERT INTO lake.itest.users VALUES (1, 'alice', 10), (2, 'bob', 10), (3, 'carol', 11)")
.await
.unwrap();
// ── bifröst sync round 1 ────────────────────────────────────────────────
let state_dir = tempfile::tempdir().unwrap();
let sink = FalkorDbSink::connect(&falkor_url, "itest").await.unwrap();
let mut listener = Listener::new(
MappingSpec::from_yaml(MAPPING_YAML).unwrap(),
TransformRegistry::new(),
source,
sink,
SyncState::open(state_dir.path().join("state.json")).unwrap(),
);
let round = listener.sync_once().await.unwrap();
eprintln!("round 1: {round:?}");
assert_eq!(round.failed, vec![], "round 1 failures");
assert_eq!(round.synced.len(), 2, "both tables synced");
let client = redis::Client::open(falkor_url.clone()).unwrap();
let mut conn = client.get_multiplexed_async_connection().await.unwrap();
assert_eq!(
count(&mut conn, "itest", "MATCH (n:Person) RETURN count(n)").await,
3
);
assert_eq!(
count(&mut conn, "itest", "MATCH (n:Company) RETURN count(n)").await,
2
);
assert_eq!(
count(
&mut conn,
"itest",
"MATCH (:Person)-[r:WORKS_AT]->(:Company) RETURN count(r)"
)
.await,
3
);
// ── idle round: nothing changed, nothing re-shipped ─────────────────────
let round = listener.sync_once().await.unwrap();
assert!(round.synced.is_empty(), "no snapshot change → no sync");
// ── incremental round: one appended row flows through ───────────────────
// (the listener owns the source now; reuse raw SQL through a 2nd session)
let url = format!("sc://127.0.0.1:{spark_port}");
let admin = SparkConnectSource::connect(&url).await.unwrap();
admin
.sql("INSERT INTO lake.itest.users VALUES (4, 'dave', 11)")
.await
.unwrap();
let round = listener.sync_once().await.unwrap();
assert_eq!(round.failed, vec![], "round 3 failures");
let users_round: Vec<_> = round
.synced
.iter()
.filter(|(t, _, _)| t == "lake.itest.users")
.collect();
assert_eq!(users_round.len(), 1);
assert_eq!(users_round[0].1, 1, "incremental read ships only the delta");
assert_eq!(
count(&mut conn, "itest", "MATCH (n:Person) RETURN count(n)").await,
4
);
assert_eq!(
count(
&mut conn,
"itest",
"MATCH (:Person)-[r:WORKS_AT]->(:Company) RETURN count(r)"
)
.await,
4
);
}