crates/knut-pipelines/src/model.rs
//! The Spark Declarative Pipelines (SDP) data model, hand-written and
//! dependency-light.
//!
//! These types are a faithful, ergonomic mirror of the wire model defined in
//! `spark/connect/pipelines.proto` (SPARK-51727, the `spark.pipelines`
//! feature added in Spark 4.x). They are deliberately *independent* of any
//! generated gRPC code so the default build compiles with no `protoc` and no
//! network: the [`crate::client`] module (feature `grpc`) converts the
//! generated protobuf structs into these.
//!
//! ## Vocabulary
//!
//! A **dataflow graph** is the compiled pipeline. Its nodes are
//! [`Dataset`]s — the *outputs* a pipeline materialises:
//!
//! * a **materialized view** ([`OutputType::MaterializedView`]) — a query
//! whose result is recomputed and published to the catalog,
//! * a **streaming table** / table ([`OutputType::Table`]) — an output
//! continuously appended to from streaming flows,
//! * a **temporary view** ([`OutputType::TemporaryView`]) — internal, not
//! published, and
//! * a **sink** ([`OutputType::Sink`]) — a streaming write to an external
//! system, not published to the catalog.
//!
//! The edges are [`Flow`]s: each flow reads some upstream datasets and
//! *writes to exactly one* target dataset. The set of `(flow.reads ->
//! flow.target)` pairs is the DAG the viewer renders.
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
/// The kind of output a [`Dataset`] node represents.
///
/// Mirrors `spark.connect.OutputType`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum OutputType {
/// A query result recomputed and published to the catalog.
MaterializedView,
/// A (streaming) table published to the catalog.
Table,
/// An internal view, not published to the catalog.
TemporaryView,
/// A streaming write to an external system, not published.
Sink,
/// The proto's safe default / forward-compat unknown value.
Unspecified,
}
impl OutputType {
/// A short, stable label for text/JSON rendering.
pub fn label(self) -> &'static str {
match self {
OutputType::MaterializedView => "materialized_view",
OutputType::Table => "table",
OutputType::TemporaryView => "temporary_view",
OutputType::Sink => "sink",
OutputType::Unspecified => "unspecified",
}
}
}
/// Where a graph element was declared in user source code.
///
/// Mirrors `spark.connect.SourceCodeLocation`. Useful for a viewer to deep
/// link an event or node back to the `.py`/`.sql` that defined it.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct SourceLocation {
/// File the element was defined in, if known.
pub file_name: Option<String>,
/// 1-based line number, if known.
pub line_number: Option<i32>,
/// Path of the top-level pipeline file at runtime, if known.
pub definition_path: Option<String>,
}
impl SourceLocation {
/// `true` when no location information at all is present.
pub fn is_empty(&self) -> bool {
self.file_name.is_none() && self.line_number.is_none() && self.definition_path.is_none()
}
}
/// A node in the dataflow graph: one declared output (table / MV / view /
/// sink).
///
/// Mirrors the salient fields of `PipelineCommand.DefineOutput`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Dataset {
/// Output name, possibly qualified (`catalog.db.name`).
pub name: String,
/// What kind of output this is.
pub output_type: OutputType,
/// Optional human comment attached at definition.
#[serde(skip_serializing_if = "Option::is_none", default)]
pub comment: Option<String>,
/// Storage format (e.g. `delta`, `iceberg`), if specified.
#[serde(skip_serializing_if = "Option::is_none", default)]
pub format: Option<String>,
/// Partition columns, if any.
#[serde(skip_serializing_if = "Vec::is_empty", default)]
pub partition_cols: Vec<String>,
/// A schema string, if the schema was declared rather than inferred.
#[serde(skip_serializing_if = "Option::is_none", default)]
pub schema: Option<String>,
/// Table properties (SQL `TBLPROPERTIES`), mirroring
/// `DefineOutput.TableDetails.table_properties`.
///
/// A `BTreeMap` rather than a `HashMap` so the serialized order is stable —
/// these round-trip through JSON/RON specs that get diffed and historized,
/// and a randomized map order would churn every write.
#[serde(skip_serializing_if = "BTreeMap::is_empty", default)]
pub properties: BTreeMap<String, String>,
/// Where this output was declared in source.
#[serde(skip_serializing_if = "SourceLocation::is_empty", default)]
pub source_location: SourceLocation,
}
impl Dataset {
/// A new dataset with just a name and type; everything else defaulted.
pub fn new(name: impl Into<String>, output_type: OutputType) -> Self {
Dataset {
name: name.into(),
output_type,
comment: None,
format: None,
partition_cols: Vec::new(),
schema: None,
properties: BTreeMap::new(),
source_location: SourceLocation::default(),
}
}
/// Attach a human comment to this dataset, returning `self` for chaining.
pub fn with_comment(mut self, comment: impl Into<String>) -> Self {
self.comment = Some(comment.into());
self
}
}
/// An edge in the dataflow graph: a flow that writes to one target dataset,
/// reading zero or more upstream datasets.
///
/// Mirrors the salient fields of `PipelineCommand.DefineFlow`. The upstream
/// `reads` are resolved from the flow's relation/query — the proto carries an
/// unresolved `Relation`; here we keep the resolved dependency names that the
/// viewer needs to draw edges.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Flow {
/// Flow name (single-part for standalone flows).
pub name: String,
/// Dataset this flow writes to (an edge head).
pub target: String,
/// Upstream datasets this flow reads (edge tails). May be empty for a
/// source flow that reads from outside the graph.
#[serde(skip_serializing_if = "Vec::is_empty", default)]
pub reads: Vec<String>,
/// One-time flow (e.g. a backfill) rather than continuously running.
#[serde(default)]
pub once: bool,
/// The SQL query that defines this flow — what it reads and writes to its
/// target. Required for the flow to actually *run* in a pipeline
/// (`define_graph` → `start_run`); `None` registers a structure-only edge.
#[serde(skip_serializing_if = "Option::is_none", default)]
pub query: Option<String>,
/// Where this flow was declared in source.
#[serde(skip_serializing_if = "SourceLocation::is_empty", default)]
pub source_location: SourceLocation,
}
impl Flow {
/// A flow named `name` writing `target`, reading `reads`.
pub fn new(
name: impl Into<String>,
target: impl Into<String>,
reads: impl IntoIterator<Item = impl Into<String>>,
) -> Self {
Flow {
name: name.into(),
target: target.into(),
reads: reads.into_iter().map(Into::into).collect(),
once: false,
query: None,
source_location: SourceLocation::default(),
}
}
/// Attach the SQL query that defines this flow — making it *runnable* by a
/// pipeline run, not just a structural edge.
pub fn with_query(mut self, sql: impl Into<String>) -> Self {
self.query = Some(sql.into());
self
}
}
/// A compiled Spark Declarative Pipeline: its datasets (nodes) and flows
/// (edges), plus the server-assigned graph id.
///
/// Mirrors the result of `CreateDataflowGraph` + the `DefineOutput` /
/// `DefineFlow` commands registered against it.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct DataflowGraph {
/// Server-assigned graph id (`dataflow_graph_id`), if known.
#[serde(skip_serializing_if = "Option::is_none", default)]
pub graph_id: Option<String>,
/// Default catalog for unqualified names.
#[serde(skip_serializing_if = "Option::is_none", default)]
pub default_catalog: Option<String>,
/// Default database for unqualified names.
#[serde(skip_serializing_if = "Option::is_none", default)]
pub default_database: Option<String>,
/// The output nodes, in declaration order.
pub datasets: Vec<Dataset>,
/// The flow edges, in declaration order.
pub flows: Vec<Flow>,
}
impl DataflowGraph {
/// An empty graph.
pub fn new() -> Self {
Self::default()
}
/// Add a dataset node, returning `self` for chaining.
pub fn with_dataset(mut self, d: Dataset) -> Self {
self.datasets.push(d);
self
}
/// Add a flow edge, returning `self` for chaining.
pub fn with_flow(mut self, f: Flow) -> Self {
self.flows.push(f);
self
}
/// Look up a dataset by exact name.
pub fn dataset(&self, name: &str) -> Option<&Dataset> {
self.datasets.iter().find(|d| d.name == name)
}
/// All flows that write to `dataset`.
pub fn flows_into<'a>(&'a self, dataset: &'a str) -> impl Iterator<Item = &'a Flow> {
self.flows.iter().filter(move |f| f.target == dataset)
}
/// Validate internal consistency: unique dataset/flow names and every
/// flow target resolving to a defined dataset.
///
/// Reads (upstream sources) are intentionally *not* required to be defined
/// datasets — a flow may read an external table outside the graph.
pub fn validate(&self) -> crate::Result<()> {
use crate::error::PipelineError;
use std::collections::HashSet;
let mut seen_ds = HashSet::new();
for d in &self.datasets {
if !seen_ds.insert(d.name.as_str()) {
return Err(PipelineError::DuplicateName {
kind: "dataset",
name: d.name.clone(),
});
}
}
let mut seen_flow = HashSet::new();
for f in &self.flows {
if !seen_flow.insert(f.name.as_str()) {
return Err(PipelineError::DuplicateName {
kind: "flow",
name: f.name.clone(),
});
}
if !seen_ds.contains(f.target.as_str()) {
return Err(PipelineError::DanglingFlow {
flow: f.name.clone(),
target: f.target.clone(),
});
}
}
Ok(())
}
/// Populate every flow's [`Flow::reads`] by walking its defining relation
/// (its SQL `query`), so the DAG's edges are complete **without replaying
/// the pipeline definitions** against a server.
///
/// For each flow that carries a `query` but has no explicit `reads`, this
/// scans the query for references to other in-graph datasets (an
/// identifier-boundary, case-insensitive match on the dataset name — and,
/// for a qualified `catalog.db.name`, on its unqualified tail too). A flow
/// never reads its own target. External sources (e.g. `range(...)`, a table
/// outside the graph) match no dataset and are correctly left out.
///
/// Returns the number of flows whose `reads` were filled in. Flows that
/// already carry explicit `reads`, or that have no `query`, are untouched —
/// so a hand-authored spec's edges are never overwritten.
pub fn populate_reads(&mut self) -> usize {
// Snapshot dataset names up front (immutable borrow) so the flow loop
// can mutate `self.flows`.
let names: Vec<String> = self.datasets.iter().map(|d| d.name.clone()).collect();
let mut filled = 0;
for flow in &mut self.flows {
if !flow.reads.is_empty() {
continue;
}
let Some(query) = flow.query.as_deref() else {
continue;
};
let mut reads = Vec::new();
for name in &names {
if *name == flow.target {
continue; // a flow never reads its own target
}
let tail = name.rsplit('.').next().unwrap_or(name);
let hit = sql_references(query, name)
|| (tail != name.as_str() && sql_references(query, tail));
if hit && !reads.contains(name) {
reads.push(name.clone());
}
}
if !reads.is_empty() {
flow.reads = reads;
filled += 1;
}
}
filled
}
/// A deterministic topological order of dataset names following the
/// flow edges (a flow's `reads` precede its `target`). Returns `None` if
/// the graph (restricted to in-graph datasets) contains a cycle.
///
/// Edges to/from datasets not present in the graph (external reads) are
/// ignored for ordering.
pub fn topo_order(&self) -> Option<Vec<String>> {
use std::collections::{BTreeMap, BTreeSet, VecDeque};
let nodes: BTreeSet<&str> = self.datasets.iter().map(|d| d.name.as_str()).collect();
// adjacency: read -> target, plus in-degree per node.
let mut indeg: BTreeMap<&str, usize> = nodes.iter().map(|n| (*n, 0)).collect();
let mut adj: BTreeMap<&str, BTreeSet<&str>> = BTreeMap::new();
for f in &self.flows {
if !nodes.contains(f.target.as_str()) {
continue;
}
for r in &f.reads {
if !nodes.contains(r.as_str()) {
continue; // external source
}
if adj.entry(r.as_str()).or_default().insert(f.target.as_str()) {
*indeg.get_mut(f.target.as_str()).unwrap() += 1;
}
}
}
// Kahn's algorithm; BTree iteration gives a stable order.
let mut queue: VecDeque<&str> = indeg
.iter()
.filter(|(_, d)| **d == 0)
.map(|(n, _)| *n)
.collect();
let mut out = Vec::with_capacity(nodes.len());
while let Some(n) = queue.pop_front() {
out.push(n.to_string());
if let Some(succ) = adj.get(n) {
for &m in succ {
let d = indeg.get_mut(m).unwrap();
*d -= 1;
if *d == 0 {
queue.push_back(m);
}
}
}
}
(out.len() == nodes.len()).then_some(out)
}
}
/// A run of a pipeline: the unit `StartRun` kicks off, against one graph.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct PipelineRun {
/// Graph this run executes.
#[serde(skip_serializing_if = "Option::is_none", default)]
pub graph_id: Option<String>,
/// Full refresh of the whole graph rather than incremental update.
#[serde(default)]
pub full_refresh_all: bool,
/// Subset of datasets selected for refresh (empty = all).
#[serde(skip_serializing_if = "Vec::is_empty", default)]
pub refresh_selection: Vec<String>,
/// Storage location for pipeline checkpoints + metadata. Spark requires
/// this to actually *execute* a run (materialise outputs); `None` only
/// works for a `dry` validation.
#[serde(skip_serializing_if = "Option::is_none", default)]
pub storage: Option<String>,
/// Validate-only: do not actually execute any flows.
#[serde(default)]
pub dry: bool,
}
/// Severity of a [`PipelineRunEvent`].
///
/// The proto's `PipelineEvent` currently carries only a timestamp and a
/// message; severity is parsed/derived by the viewer so it can colour the
/// stream. Defaults to [`EventLevel::Info`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
#[derive(Default)]
pub enum EventLevel {
#[default]
Info,
Warn,
Error,
}
/// One event emitted while a pipeline run progresses — the atom of the live
/// stream viewer.
///
/// Mirrors `spark.connect.PipelineEvent` (timestamp + message), enriched with
/// a derived [`EventLevel`] and an optional originating element name so a
/// viewer can attribute the line to a node/edge in the DAG.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PipelineRunEvent {
/// RFC3339 timestamp string, if the server set one.
#[serde(skip_serializing_if = "Option::is_none", default)]
pub timestamp: Option<String>,
/// Derived severity.
#[serde(default)]
pub level: EventLevel,
/// The user-facing message.
pub message: String,
/// Dataset/flow this event is about, if attributable.
#[serde(skip_serializing_if = "Option::is_none", default)]
pub element: Option<String>,
}
impl PipelineRunEvent {
/// A bare info event with a message.
pub fn info(message: impl Into<String>) -> Self {
PipelineRunEvent {
timestamp: None,
level: EventLevel::Info,
message: message.into(),
element: None,
}
}
}
/// Does `sql` reference the identifier `name` at identifier boundaries?
///
/// Case-insensitive (Spark identifiers are). A match must not be flanked by
/// identifier characters (`[A-Za-z0-9_]`) so `bronze` does not match inside
/// `bronze_raw`. Kept dependency-free (no regex, no SQL parser): a viewer only
/// needs to know *which known datasets a query touches* to draw edges, not to
/// fully parse the query.
pub(crate) fn sql_references(sql: &str, name: &str) -> bool {
if name.is_empty() {
return false;
}
let hay = sql.to_ascii_lowercase();
let needle = name.to_ascii_lowercase();
let bytes = hay.as_bytes();
let is_ident = |c: u8| c.is_ascii_alphanumeric() || c == b'_';
let mut from = 0;
while let Some(pos) = hay[from..].find(&needle) {
let start = from + pos;
let end = start + needle.len();
let before_ok = start == 0 || !is_ident(bytes[start - 1]);
let after_ok = end >= bytes.len() || !is_ident(bytes[end]);
if before_ok && after_ok {
return true;
}
from = start + 1;
}
false
}
#[cfg(test)]
mod model_tests {
use super::*;
fn medallion() -> DataflowGraph {
// bronze <- range() (external); silver <- bronze; gold <- silver, bronze.
// Flows carry only a query — no explicit `reads` — so edges must be
// recovered by walking the relations.
DataflowGraph::new()
.with_dataset(Dataset::new("bronze", OutputType::Table))
.with_dataset(Dataset::new("silver", OutputType::MaterializedView))
.with_dataset(Dataset::new("gold", OutputType::MaterializedView))
.with_flow(
Flow::new("ingest", "bronze", Vec::<String>::new())
.with_query("SELECT id FROM range(0, 1000)"),
)
.with_flow(
Flow::new("clean", "silver", Vec::<String>::new())
.with_query("SELECT id FROM bronze WHERE id % 2 = 0"),
)
.with_flow(
Flow::new("rollup", "gold", Vec::<String>::new())
.with_query("SELECT s.id FROM silver s JOIN bronze b ON s.id = b.id"),
)
}
#[test]
fn sql_references_respects_identifier_boundaries() {
assert!(sql_references("SELECT * FROM bronze", "bronze"));
assert!(sql_references("select * from BRONZE", "bronze")); // case-insensitive
assert!(sql_references(
"FROM catalog.db.bronze x",
"catalog.db.bronze"
));
// Must not match a substring inside a longer identifier.
assert!(!sql_references("SELECT * FROM bronze_raw", "bronze"));
assert!(!sql_references("SELECT * FROM my_bronze", "bronze"));
}
#[test]
fn populate_reads_walks_relations_to_complete_edges() {
let mut g = medallion();
let filled = g.populate_reads();
// silver + gold get reads; bronze reads range() → external → untouched.
assert_eq!(filled, 2);
let ingest = g.flows.iter().find(|f| f.name == "ingest").unwrap();
assert!(ingest.reads.is_empty()); // external source only
let clean = g.flows.iter().find(|f| f.name == "clean").unwrap();
assert_eq!(clean.reads, vec!["bronze".to_string()]);
let rollup = g.flows.iter().find(|f| f.name == "rollup").unwrap();
// Graph order (bronze, silver, gold) preserved; target `gold` excluded.
assert_eq!(
rollup.reads,
vec!["bronze".to_string(), "silver".to_string()]
);
// The recovered edges give a real topo order without a server.
assert_eq!(
g.topo_order(),
Some(vec!["bronze".into(), "silver".into(), "gold".into()])
);
}
#[test]
fn populate_reads_preserves_explicit_edges() {
// An explicit `reads` (a hand-authored spec) is never overwritten, even
// if the query would resolve differently.
let mut g = DataflowGraph::new()
.with_dataset(Dataset::new("a", OutputType::Table))
.with_dataset(Dataset::new("b", OutputType::Table))
.with_flow(Flow::new("f", "b", ["a"]).with_query("SELECT 1"));
assert_eq!(g.populate_reads(), 0);
assert_eq!(g.flows[0].reads, vec!["a".to_string()]);
}
#[test]
fn populate_reads_qualified_tail_match() {
// A qualified dataset name is matched by its unqualified tail too.
let mut g = DataflowGraph::new()
.with_dataset(Dataset::new(
"spark_catalog.default.bronze",
OutputType::Table,
))
.with_dataset(Dataset::new(
"spark_catalog.default.silver",
OutputType::Table,
))
.with_flow(
Flow::new(
"clean",
"spark_catalog.default.silver",
Vec::<String>::new(),
)
.with_query("SELECT * FROM bronze"),
);
assert_eq!(g.populate_reads(), 1);
assert_eq!(
g.flows[0].reads,
vec!["spark_catalog.default.bronze".to_string()]
);
}
}