nordisk/knut

crates/knut-bifrost/src/airflow.rs

//! Airflow 3 asset events (AIP-82) — the *emit* half of the story.
//!
//! After a sync round lands, bifröst can announce "table X is now in the
//! graph" as an Airflow **asset event**, so downstream DAGs (analytics,
//! exports, alerts) schedule on the asset instead of a cron. The *listen*
//! half — reacting to file drops — is [`FileSource`](crate::source::FileSource),
//! whose [`TableState::done_files`](crate::state::TableState::done_files)
//! ledger is exactly the bookkeeping AIP-82's poll-based triggers expect.
//!
//! The event model and JSON body are dependency-free and unit-tested; the
//! HTTP transport ([`AirflowEmitter`]) is behind the `airflow` feature.

use async_trait::async_trait;

use crate::{error::Result, listener::RoundReport};

/// One asset event: an asset URI plus a free-form `extra` payload that
/// downstream DAGs can read (rows synced, graph elements written).
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub struct AssetEvent {
    /// The asset's URI, e.g. `graph://social/lake.social.users`.
    pub uri: String,
    /// Arbitrary metadata attached to the event.
    pub extra: serde_json::Value,
}

/// Somewhere asset events are delivered. The blanket file/Spark sources don't
/// know about Airflow; emission is a separate, swappable concern.
#[async_trait]
pub trait AssetEmitter: Send + Sync {
    async fn emit(&self, events: &[AssetEvent]) -> Result<()>;
}

/// Turn a completed [`RoundReport`] into one [`AssetEvent`] per table that
/// actually synced. The URI is `{uri_prefix}{table}`; the `extra` carries the
/// row count and the sink's vertex/edge tallies. Failed tables emit nothing —
/// their state did not advance, so the asset is not yet current.
pub fn round_to_events(uri_prefix: &str, report: &RoundReport) -> Vec<AssetEvent> {
    report
        .synced
        .iter()
        .map(|(table, rows, sink)| AssetEvent {
            uri: format!("{uri_prefix}{table}"),
            extra: serde_json::json!({
                "table": table,
                "rows": rows,
                "vertices_written": sink.vertices_written,
                "edges_written": sink.edges_written,
            }),
        })
        .collect()
}

/// The JSON body for Airflow 3's `POST /api/v2/assets/events`. Kept pure so
/// the wire shape can be asserted without standing up a server.
pub fn asset_event_body(event: &AssetEvent) -> serde_json::Value {
    serde_json::json!({
        "asset_uri": event.uri,
        "extra": event.extra,
    })
}

/// Emit every event derived from `report` through `emitter`. Convenience for
/// the common "sync round → announce assets" wiring in an `on_round` callback.
pub async fn emit_round(
    emitter: &impl AssetEmitter,
    uri_prefix: &str,
    report: &RoundReport,
) -> Result<()> {
    let events = round_to_events(uri_prefix, report);
    if events.is_empty() {
        return Ok(());
    }
    emitter.emit(&events).await
}

#[cfg(feature = "airflow")]
pub use http::AirflowEmitter;

#[cfg(feature = "airflow")]
mod http {
    use super::*;
    use crate::error::BifrostError;

    /// Posts asset events to an Airflow 3 REST API
    /// (`POST {base_url}/api/v2/assets/events`), one request per event, with an
    /// optional bearer token.
    pub struct AirflowEmitter {
        client: reqwest::Client,
        events_url: String,
        token: Option<String>,
    }

    impl AirflowEmitter {
        /// `base_url` is the Airflow API root, e.g. `http://airflow:8080`.
        pub fn new(base_url: impl AsRef<str>, token: Option<String>) -> Self {
            let base = base_url.as_ref().trim_end_matches('/');
            Self {
                client: super::tls_roots(reqwest::Client::builder()).build().expect("reqwest client"),
                events_url: format!("{base}/api/v2/assets/events"),
                token,
            }
        }
    }

    #[async_trait]
    impl AssetEmitter for AirflowEmitter {
        async fn emit(&self, events: &[AssetEvent]) -> Result<()> {
            for event in events {
                let mut req = self
                    .client
                    .post(&self.events_url)
                    .json(&asset_event_body(event));
                if let Some(token) = &self.token {
                    req = req.bearer_auth(token);
                }
                let resp = req
                    .send()
                    .await
                    .map_err(|e| BifrostError::Source(format!("airflow emit: {e}")))?;
                if !resp.status().is_success() {
                    return Err(BifrostError::Source(format!(
                        "airflow emit: {} for {}",
                        resp.status(),
                        event.uri
                    )));
                }
            }
            Ok(())
        }
    }
}

/// reqwest 0.13 with reqwest 0.12 `rustls-tls` behaviour (2026-10-02): the ring
/// provider and the bundled Mozilla roots (webpki), not 0.13's aws-lc-rs + OS
/// trust store. Roots added later on the builder are merged in, as in 0.12.
#[cfg(feature = "airflow")]
#[allow(dead_code)]
fn tls_roots(b: reqwest::ClientBuilder) -> reqwest::ClientBuilder {
    let _ = rustls::crypto::ring::default_provider().install_default();
    b.tls_certs_only(
        webpki_root_certs::TLS_SERVER_ROOT_CERTS
            .iter()
            .filter_map(|c| reqwest::Certificate::from_der(c).ok()),
    )
}