crates/knut-pipelines/src/airflow.rs
//! Airflow 3 (AIP-82) asset-event sink (feature `airflow`).
//!
//! A run's completion events map cleanly onto Airflow **asset events**: when an
//! SDP output finishes refreshing, a downstream Airflow DAG can schedule on it
//! instead of a cron. [`AirflowAssetSink`] is a [`PipelineEventSink`] that
//! watches a run's event stream and, each time a dataset reaches
//! [`NodeStatus::Done`], records one [`AssetEvent`] (URI + row count) — then
//! forwards the batch through bifröst's [`AssetEmitter`].
//!
//! This is the roadmap's Bifröst hook, now wired: it **reuses**
//! `knut_bifrost::airflow`'s `AssetEvent` + `AssetEmitter` (no duplicate model)
//! and this crate's [`ViewerState`] for attribution + status + row counts (no
//! duplicate event parsing). Completion detection, URI shape and `extra`
//! payload are pure and unit-tested; the actual HTTP transport is bifröst's
//! `AirflowEmitter` (its own `airflow` feature) or any [`AssetEmitter`].
//!
//! The sink's collection is synchronous (it fits the sync [`PipelineEventSink`]
//! seam); emission is a separate `async` step, exactly as bifröst's own
//! `emit_round` splits "derive events" from "POST them".
use std::collections::BTreeSet;
use knut_bifrost::airflow::{AssetEmitter, AssetEvent};
use crate::error::PipelineError;
use crate::model::{DataflowGraph, PipelineRunEvent};
use crate::progress::{NodeStatus, ViewerState, ViewerUpdate};
use crate::viewer::PipelineEventSink;
/// A [`PipelineEventSink`] that turns each SDP output's completion into an
/// Airflow (AIP-82) [`AssetEvent`].
///
/// Feed it a run's [`PipelineRunEvent`]s (it is the sink `start_run` pushes
/// into); it folds them into an internal [`ViewerState`] — reusing the same
/// flow→dataset attribution, status derivation and `wrote N rows` parsing the
/// viewer uses — and appends one asset event the first time each dataset turns
/// [`NodeStatus::Done`]. A failed dataset advances no state, so it emits
/// nothing (its asset is not current) — mirroring `round_to_events`.
pub struct AirflowAssetSink {
/// The viewer state the events are folded into (attribution + rows).
state: ViewerState,
/// Prefix prepended to a dataset name to form the asset URI, e.g.
/// `sdp://sales/` → `sdp://sales/gold_daily`.
uri_prefix: String,
/// Datasets already announced, so each asset event fires at most once.
emitted: BTreeSet<String>,
/// Asset events collected so far, ready to [`AirflowAssetSink::emit`].
pending: Vec<AssetEvent>,
}
impl AirflowAssetSink {
/// A sink over `graph`, forming asset URIs as `{uri_prefix}{dataset}`.
pub fn new(graph: DataflowGraph, uri_prefix: impl Into<String>) -> Self {
AirflowAssetSink {
state: ViewerState::new(graph),
uri_prefix: uri_prefix.into(),
emitted: BTreeSet::new(),
pending: Vec::new(),
}
}
/// The asset events collected so far (one per completed dataset).
pub fn events(&self) -> &[AssetEvent] {
&self.pending
}
/// Take the collected asset events, leaving the sink empty.
pub fn take_events(&mut self) -> Vec<AssetEvent> {
std::mem::take(&mut self.pending)
}
/// Forward every collected asset event through `emitter` (e.g. bifröst's
/// `AirflowEmitter`). A no-op when nothing completed. Leaves the collected
/// events in place so a caller may re-emit or inspect them.
pub async fn emit(&self, emitter: &impl AssetEmitter) -> crate::Result<()> {
if self.pending.is_empty() {
return Ok(());
}
emitter
.emit(&self.pending)
.await
.map_err(|e| PipelineError::Other(format!("airflow asset emit failed: {e}")))
}
/// Append an asset event for each dataset that has newly reached
/// [`NodeStatus::Done`] since the last scan.
fn scan_completed(&mut self) {
let mut fresh = Vec::new();
for d in &self.state.graph.datasets {
if self.state.status_of(&d.name) == NodeStatus::Done && !self.emitted.contains(&d.name)
{
let rows = self.state.rows_written.get(&d.name).copied();
fresh.push((d.name.clone(), d.output_type.label(), rows));
}
}
for (name, output_type, rows) in fresh {
self.emitted.insert(name.clone());
self.pending.push(AssetEvent {
uri: format!("{}{}", self.uri_prefix, name),
extra: serde_json::json!({
"dataset": name,
"rows": rows,
"output_type": output_type,
"source": "spark-sdp",
}),
});
}
}
}
impl PipelineEventSink for AirflowAssetSink {
fn on_event(&mut self, event: &PipelineRunEvent) -> crate::Result<()> {
self.state.apply(ViewerUpdate::Event(event.clone()));
self.scan_completed();
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model::{Dataset, EventLevel, Flow, OutputType};
fn medallion() -> DataflowGraph {
DataflowGraph::new()
.with_dataset(Dataset::new("bronze", OutputType::Table))
.with_dataset(Dataset::new("silver", OutputType::MaterializedView))
.with_flow(Flow::new("ingest", "bronze", Vec::<String>::new()))
.with_flow(Flow::new("clean", "silver", ["bronze"]))
}
fn info(msg: &str) -> PipelineRunEvent {
PipelineRunEvent::info(msg)
}
#[test]
fn synthetic_run_emits_one_asset_event_per_completed_output() {
let mut sink = AirflowAssetSink::new(medallion(), "sdp://demo/");
// A synthetic run: flows named, attributed to their target datasets.
for ev in [
info("Run started."),
info("Flow 'ingest' is STARTING."),
info("Flow 'ingest' has COMPLETED and wrote 1000 rows."),
info("Flow 'clean' is RUNNING."),
info("Flow 'clean' has COMPLETED and wrote 500 rows."),
] {
sink.on_event(&ev).unwrap();
}
sink.on_finish().unwrap();
let events = sink.events();
assert_eq!(events.len(), 2, "one asset event per completed dataset");
// bronze completed first (via flow `ingest`), silver second (via `clean`).
assert_eq!(events[0].uri, "sdp://demo/bronze");
assert_eq!(events[0].extra["rows"], 1000);
assert_eq!(events[0].extra["output_type"], "table");
assert_eq!(events[0].extra["source"], "spark-sdp");
assert_eq!(events[1].uri, "sdp://demo/silver");
assert_eq!(events[1].extra["rows"], 500);
assert_eq!(events[1].extra["output_type"], "materialized_view");
// The AIP-82 wire body is bifröst's — assert we produce its shape.
let body = knut_bifrost::airflow::asset_event_body(&events[1]);
assert_eq!(body["asset_uri"], "sdp://demo/silver");
assert_eq!(body["extra"]["dataset"], "silver");
}
#[test]
fn each_output_emits_at_most_once() {
let mut sink = AirflowAssetSink::new(medallion(), "sdp://demo/");
// Two completion events for the same flow must not double-emit.
sink.on_event(&info("Flow 'ingest' has COMPLETED and wrote 10 rows."))
.unwrap();
sink.on_event(&info("Flow 'ingest' has COMPLETED and wrote 10 rows."))
.unwrap();
assert_eq!(sink.events().len(), 1);
}
#[test]
fn failed_output_emits_nothing() {
let mut sink = AirflowAssetSink::new(medallion(), "sdp://demo/");
sink.on_event(&PipelineRunEvent {
timestamp: None,
level: EventLevel::Error,
message: "Flow 'clean' FAILED: boom".into(),
element: None,
})
.unwrap();
assert!(
sink.events().is_empty(),
"a failed output is not a current asset"
);
}
#[test]
fn take_events_drains() {
let mut sink = AirflowAssetSink::new(medallion(), "asset://");
sink.on_event(&info("Flow 'ingest' has COMPLETED."))
.unwrap();
assert_eq!(sink.take_events().len(), 1);
assert!(sink.events().is_empty());
}
/// A mock [`AssetEmitter`] that records the events it was handed, so the
/// test can prove [`AirflowAssetSink::emit`] forwards the full batch it
/// folded from the run — the sink → bifröst emit seam, with no HTTP.
#[derive(Default)]
struct RecordingEmitter {
seen: std::sync::Mutex<Vec<AssetEvent>>,
calls: std::sync::atomic::AtomicUsize,
}
#[async_trait::async_trait]
impl AssetEmitter for RecordingEmitter {
async fn emit(&self, events: &[AssetEvent]) -> knut_bifrost::error::Result<()> {
self.calls
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
self.seen.lock().unwrap().extend_from_slice(events);
Ok(())
}
}
#[tokio::test]
async fn emit_forwards_the_folded_run_to_the_emitter() {
// Fold a synthetic COMPLETED run into the sink…
let mut sink = AirflowAssetSink::new(medallion(), "sdp://demo/");
for ev in [
info("Flow 'ingest' has COMPLETED and wrote 1000 rows."),
info("Flow 'clean' has COMPLETED and wrote 500 rows."),
] {
sink.on_event(&ev).unwrap();
}
sink.on_finish().unwrap();
// …then emit through the mock and assert the exact batch arrived.
let emitter = RecordingEmitter::default();
sink.emit(&emitter).await.unwrap();
assert_eq!(emitter.calls.load(std::sync::atomic::Ordering::Relaxed), 1);
let seen = emitter.seen.lock().unwrap();
assert_eq!(seen.len(), 2, "both completed datasets were forwarded");
assert_eq!(seen[0].uri, "sdp://demo/bronze");
assert_eq!(seen[0].extra["rows"], 1000);
assert_eq!(seen[1].uri, "sdp://demo/silver");
assert_eq!(seen[1].extra["rows"], 500);
}
#[tokio::test]
async fn emit_is_a_no_op_when_nothing_completed() {
// No completion events → nothing to announce → the emitter is never hit.
let sink = AirflowAssetSink::new(medallion(), "sdp://demo/");
let emitter = RecordingEmitter::default();
sink.emit(&emitter).await.unwrap();
assert_eq!(emitter.calls.load(std::sync::atomic::Ordering::Relaxed), 0);
assert!(emitter.seen.lock().unwrap().is_empty());
}
}