nordisk/knut

crates/knut-pipelines/src/spark_view.rs

//! **The Spark runtime + Structured-Streaming VIEW model** — the backend API a
//! consumer (korp's Spark pane) calls to *see* a live Spark Connect runtime:
//! its version, the Structured Streaming queries currently running, each query's
//! status + recent progress, and — where the protocol allows — to STOP one.
//!
//! # Why a neutral model here
//!
//! The live wiring rides on Spark Connect's `StreamingQueryManagerCommand` /
//! `StreamingQueryCommand` (commands.proto), which are only compiled under the
//! `grpc` feature (they carry the isolated tonic-0.14 + generated-proto tree).
//! These **neutral** types carry none of that — they are plain data korp can
//! depend on without pulling gRPC in — and the `grpc`-gated methods on
//! [`crate::PipelinesClient`] fold the proto results into them (see
//! `client.rs`). So the view is *one small model*, engine details never leak
//! into korp, and the proto→model fold is unit-testable here without a live
//! server.
//!
//! # What the protocol allows
//!
//! Spark Connect exposes streaming queries read-side in full (list active,
//! status, recent-progress JSON) and one control verb we surface —
//! [`stop`](crate::PipelinesClient) — matching `StreamingQueryCommand.stop`.
//! Cluster-executor/worker topology is **not** on the Connect wire (that lives
//! in the Spark REST UI, a separate endpoint); this model deliberately covers
//! only what Connect can answer, and says so, rather than inventing fields the
//! protocol cannot fill.

use serde::{Deserialize, Serialize};

/// One running (or recently-run) Structured Streaming query, as the streaming
/// query manager's `active()` lists it.
///
/// `id` persists across restarts-from-checkpoint; `run_id` changes each
/// start/restart (mirrors Spark's `StreamingQueryInstanceId`).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct StreamingQueryInfo {
    /// Stable query id (persists across checkpoint restarts).
    pub id: String,
    /// This run's id (changes on every start/restart).
    pub run_id: String,
    /// The query's name, if it was given one.
    pub name: Option<String>,
}

impl StreamingQueryInfo {
    /// Construct from the three fields Spark returns.
    pub fn new(id: impl Into<String>, run_id: impl Into<String>, name: Option<String>) -> Self {
        StreamingQueryInfo {
            id: id.into(),
            run_id: run_id.into(),
            name: name.filter(|n| !n.is_empty()),
        }
    }
}

/// A streaming query's live status — the `status()` API's four facts.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct StreamingQueryStatus {
    /// Human-readable status line (`"Waiting for data to arrive"`, …).
    pub message: String,
    /// Whether the query is currently running.
    pub is_active: bool,
    /// Whether data is available to be processed right now.
    pub is_data_available: bool,
    /// Whether a trigger (micro-batch) is currently executing.
    pub is_trigger_active: bool,
}

impl StreamingQueryStatus {
    /// Construct from the four `StatusResult` fields.
    pub fn new(
        message: impl Into<String>,
        is_active: bool,
        is_data_available: bool,
        is_trigger_active: bool,
    ) -> Self {
        StreamingQueryStatus {
            message: message.into(),
            is_active,
            is_data_available,
            is_trigger_active,
        }
    }
}

/// Coarse Spark runtime facts a Connect session can answer without the REST UI.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SparkRuntimeInfo {
    /// The Spark version the connected server reports (`AnalyzePlan` SparkVersion).
    pub version: String,
    /// The Spark Connect endpoint this view is bound to.
    pub connect_url: String,
}

impl SparkRuntimeInfo {
    /// Construct a runtime-info record.
    pub fn new(version: impl Into<String>, connect_url: impl Into<String>) -> Self {
        SparkRuntimeInfo {
            version: version.into(),
            connect_url: connect_url.into(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn streaming_query_info_drops_empty_name() {
        let named = StreamingQueryInfo::new("q1", "r1", Some("clicks".into()));
        assert_eq!(named.name.as_deref(), Some("clicks"));
        // An empty name from the proto (`optional string` absent → "") folds to None.
        let anon = StreamingQueryInfo::new("q2", "r2", Some(String::new()));
        assert_eq!(anon.name, None, "an empty name is normalized to None");
        let none = StreamingQueryInfo::new("q3", "r3", None);
        assert_eq!(none.name, None);
    }

    #[test]
    fn status_and_runtime_round_trip_json() {
        let s = StreamingQueryStatus::new("Processing new data", true, true, true);
        let j = serde_json::to_string(&s).unwrap();
        assert_eq!(serde_json::from_str::<StreamingQueryStatus>(&j).unwrap(), s);

        let r = SparkRuntimeInfo::new("4.1.2", "sc://spark:15002");
        assert_eq!(r.version, "4.1.2");
        let j = serde_json::to_string(&r).unwrap();
        assert_eq!(serde_json::from_str::<SparkRuntimeInfo>(&j).unwrap(), r);
    }
}