crates/knut-pipelines/tests/integration_spark_connect.rs
//! The pipelines viewer's **live path** against a real Spark Connect server.
//!
//! Starts a Spark Connect container, connects with [`PipelinesClient`], runs a
//! Spark SQL job through `execute_sql`, and folds the streamed
//! `ExecutionProgress` into a [`ViewerState`] — the exact path `knutpipe watch
//! --sql …` drives. Job progress is what makes the viewer a *Spark jobs*
//! monitor, so this exercises the part that works on any Spark Connect 4.x
//! (the `spark.pipelines` define/run API is only wired in newer builds and is
//! probed leniently below).
//!
//! ## Prerequisites
//! - Podman socket: `systemctl --user start podman.socket`
//! - A Spark Connect image (override with `KNUT_SPARK_IMAGE`):
//! `podman build -t knut-spark-iceberg -f containers/spark-iceberg/Containerfile .`
//!
//! ## Run
//! ```bash
//! cargo test -p knut-pipelines --features grpc --test integration_spark_connect \
//! -- --ignored --nocapture
//! ```
#![cfg(feature = "grpc")]
use std::time::Duration;
use futures_util::StreamExt;
use knut_pipelines::{
DataflowGraph, Dataset, Flow, OutputType, PipelineEventSink, PipelineRun, PipelineRunEvent,
PipelinesClient, ViewerState, ViewerUpdate,
};
use testcontainers::{GenericImage, ImageExt, core::ContainerPort, runners::AsyncRunner};
/// Collects a pipeline run's events so the test can assert it actually ran.
#[derive(Default)]
struct CollectingSink {
events: Vec<PipelineRunEvent>,
}
impl PipelineEventSink for CollectingSink {
fn on_event(&mut self, event: &PipelineRunEvent) -> knut_pipelines::Result<()> {
self.events.push(event.clone());
Ok(())
}
}
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}")) };
}
}
}
/// Drain a SQL job's update stream into a `ViewerState`, returning how many
/// progress snapshots arrived. Errors propagate (a transport failure fails the
/// run); a clean end means the job executed.
async fn run_sql(
client: &mut PipelinesClient,
sql: &str,
) -> knut_pipelines::Result<(ViewerState, usize)> {
let mut stream = client.execute_sql(sql).await?;
let mut state = ViewerState::default();
let mut progress = 0usize;
while let Some(item) = stream.next().await {
let update = item?;
if matches!(update, ViewerUpdate::Progress(_)) {
progress += 1;
}
state.apply(update);
}
Ok((state, progress))
}
#[tokio::test]
#[ignore = "needs podman + a Spark Connect image; run with --ignored"]
async fn knutpipe_live_path_streams_spark_job_progress() {
ensure_docker_host();
// ── setup: Spark Connect container ──────────────────────────────────────
let image = std::env::var("KNUT_SPARK_IMAGE")
.unwrap_or_else(|_| "localhost/knut-spark-iceberg:latest".into());
let (name, tag) = image.rsplit_once(':').unwrap_or((image.as_str(), "latest"));
let spark = GenericImage::new(name, tag)
.with_exposed_port(ContainerPort::Tcp(15002))
.with_startup_timeout(Duration::from_secs(300))
.start()
.await
.expect("start spark connect (build knut-spark-iceberg or set KNUT_SPARK_IMAGE)");
let port = spark.get_host_port_ipv4(15002).await.unwrap();
let url = format!("sc://127.0.0.1:{port}");
// ── wait until Spark Connect answers a real query ───────────────────────
// (the channel connects before the JVM is ready, so probe with SQL)
let mut client = {
let mut attempt = 0;
loop {
let ok = match PipelinesClient::connect(&url).await {
Ok(mut c) => run_sql(&mut c, "SELECT 1").await.map(|_| c).ok(),
Err(_) => None,
};
match ok {
Some(c) => break c,
None if attempt < 60 => {
attempt += 1;
tokio::time::sleep(Duration::from_secs(3)).await;
}
None => panic!("Spark Connect never became ready after 180s"),
}
}
};
// ── the live path: run a real Spark job, stream its progress ────────────
// The server pushes `ExecutionProgress` (base.proto field 18) unconditionally
// while a plan runs, paced by `spark.connect.progress.reportInterval` (default
// ~2s) — there is no per-request toggle (see `execute_sql`'s doc). So to make
// "at least one progress snapshot" deterministic we run a genuinely
// multi-second job: a shuffle aggregation over 400M rows across 32 partitions
// forces stages/tasks that span several report intervals.
let (state, progress) = run_sql(
&mut client,
"SELECT count(*) AS n FROM \
(SELECT id % 1000 AS g, count(*) AS c \
FROM range(0, 400000000, 1, 32) GROUP BY id % 1000)",
)
.await
.expect("SQL job runs over Spark Connect");
eprintln!(
"live run: {progress} progress snapshot(s); final = {:?}",
state.progress
);
// The server volunteers ExecutionProgress while the (deliberately long) job
// runs, so at least one snapshot must have arrived and been folded in.
assert!(
progress >= 1,
"expected ≥1 ViewerUpdate::Progress from a multi-second Spark job, got {progress}"
);
let p = state
.progress
.as_ref()
.expect("a progress snapshot was folded into the ViewerState");
assert!(p.total_tasks() >= 0, "progress snapshot is well-formed");
// ── #1: define a *runnable* dataflow graph and start a run ──────────────
// One materialized view fed by a source flow whose SQL is what StartRun
// executes — exercising define_graph (CreateDataflowGraph + DefineOutput +
// DefineFlow{WriteRelationFlowDetails}) end to end.
let graph = DataflowGraph::new()
.with_dataset(Dataset::new("knut_nums", OutputType::MaterializedView))
.with_flow(
Flow::new("knut_nums", "knut_nums", Vec::<String>::new())
.with_query("SELECT id FROM range(0, 1000)"),
);
let graph_id = client
.define_graph(&graph)
.await
.expect("define a runnable dataflow graph");
assert!(!graph_id.is_empty(), "server assigned a graph id");
eprintln!("defined runnable dataflow graph {graph_id}");
let run = PipelineRun {
graph_id: Some(graph_id),
full_refresh_all: true,
// Spark requires a storage location to execute (materialise) a run;
// a path inside the container works for this single-MV pipeline.
storage: Some("file:///tmp/knut_pipeline_storage".into()),
..Default::default()
};
let mut sink = CollectingSink::default();
match client.start_run(&run, &mut sink).await {
Ok(()) => {
eprintln!("run streamed {} event(s):", sink.events.len());
for ev in sink.events.iter().take(8) {
eprintln!(" [{:?}] {}", ev.level, ev.message);
}
assert!(
!sink.events.is_empty(),
"a real run emits at least one event"
);
}
// Materialising a MV can hit image-specific storage limits; the define
// path (above) is the #1 contract, so a run error is reported, not fatal.
Err(e) => eprintln!("start_run ended early (define path still proven): {e}"),
}
// teardown: container drops here.
}