crates/knut-pipelines/src/client.rs
//! Spark Connect gRPC client for the pipelines API (feature `grpc`).
//!
//! This is the live half of the viewer. It connects to a Spark Connect
//! server (`sc://host:port`), drives the `spark.pipelines` command family
//! (`CreateDataflowGraph`, `DefineOutput`, `DefineFlow`, `StartRun`), and
//! streams `PipelineEvent`s back, converting each into the crate's
//! dependency-light [`crate::model::PipelineRunEvent`] and feeding it to a
//! [`crate::viewer::PipelineEventSink`].
//!
//! All pipelines commands ride on `Command.pipeline_command` (commands.proto
//! field 19) inside a `Plan`, executed via the server-streaming
//! `SparkConnectService.ExecutePlan` RPC; events arrive as
//! `ExecutePlanResponse.pipeline_event_result` (base.proto field 21) and
//! command results as `pipeline_command_result` (field 22).
//!
//! The generated bindings come from the vendored `proto/` set via `build.rs`;
//! they carry their own tonic 0.14 tree, isolated from the workspace's tonic
//! 0.13 by the `grpc` feature gate.
use futures_core::Stream;
use crate::error::PipelineError;
use crate::model::{
DataflowGraph, Dataset, EventLevel, Flow, OutputType, PipelineRun, PipelineRunEvent,
};
use crate::progress::{ExecutionProgress, StageProgress, ViewerUpdate, parse_event_message};
use crate::viewer::PipelineEventSink;
/// The generated Spark Connect protobuf + gRPC bindings (`spark.connect`
/// package). Kept private; the public surface is [`PipelinesClient`] and the
/// model types.
pub(crate) mod proto {
// The generated set is the full Spark Connect surface; we only use the
// pipelines slice, so unused types/dead code are expected and silenced.
#![allow(clippy::all, missing_docs, dead_code, rustdoc::all)]
tonic::include_proto!("spark.connect");
}
use proto::execute_plan_response::ResponseType;
use proto::pipeline_command::{self, CreateDataflowGraph, DefineFlow, DefineOutput, StartRun};
use proto::plan::OpType;
use proto::relation::RelType;
use proto::spark_connect_service_client::SparkConnectServiceClient;
use proto::{
Command, ExecutePlanRequest, OutputType as ProtoOutputType, Plan, Relation, Sql, UserContext,
};
/// Map our model's [`OutputType`] onto the generated enum's i32 value.
fn output_type_to_proto(t: OutputType) -> i32 {
let v = match t {
OutputType::MaterializedView => ProtoOutputType::MaterializedView,
OutputType::Table => ProtoOutputType::Table,
OutputType::TemporaryView => ProtoOutputType::TemporaryView,
OutputType::Sink => ProtoOutputType::Sink,
OutputType::Unspecified => ProtoOutputType::Unspecified,
};
v as i32
}
/// Convert a generated `PipelineEvent` into the model event a sink consumes.
fn event_from_proto(ev: proto::PipelineEvent) -> PipelineRunEvent {
let message = ev.message.unwrap_or_default();
// The proto event has no severity; derive a coarse level from the text so
// a viewer can colour the stream until Spark surfaces a real level.
let lower = message.to_ascii_lowercase();
let level = if lower.contains("error") || lower.contains("failed") {
EventLevel::Error
} else if lower.contains("warn") {
EventLevel::Warn
} else {
EventLevel::Info
};
let timestamp = ev.timestamp.map(|t| {
// google.protobuf.Timestamp -> RFC3339 if chrono can represent it.
chrono::DateTime::from_timestamp(t.seconds, t.nanos as u32)
.map(|dt| dt.to_rfc3339())
.unwrap_or_else(|| format!("{}.{:09}", t.seconds, t.nanos))
});
// The proto event has no typed element field; recover the quoted element
// name (a flow or dataset) from the message so a viewer can attribute the
// line to a node without a message scan. ViewerState resolves a flow name
// to its target dataset.
let element = parse_event_message(&message).element;
PipelineRunEvent {
timestamp,
level,
message,
element,
}
}
/// Convert a generated `ExecutionProgress` (Spark job progress) into the
/// model snapshot the viewer renders.
fn progress_from_proto(p: proto::execute_plan_response::ExecutionProgress) -> ExecutionProgress {
ExecutionProgress {
stages: p
.stages
.into_iter()
.map(|s| StageProgress {
stage_id: s.stage_id,
num_tasks: s.num_tasks,
num_completed_tasks: s.num_completed_tasks,
input_bytes_read: s.input_bytes_read,
done: s.done,
})
.collect(),
num_inflight_tasks: p.num_inflight_tasks,
}
}
/// Map one `ExecutePlanResponse` to a [`ViewerUpdate`], if it carries one we
/// care about (a pipeline event or a job-progress snapshot). Other response
/// types (arrow batches, sql results, …) are ignored by the viewer.
fn response_to_update(resp: proto::ExecutePlanResponse) -> Option<ViewerUpdate> {
match resp.response_type {
Some(ResponseType::PipelineEventResult(per)) => per
.event
.map(|ev| ViewerUpdate::Event(event_from_proto(ev))),
Some(ResponseType::ExecutionProgress(p)) => {
Some(ViewerUpdate::Progress(progress_from_proto(p)))
}
_ => None,
}
}
/// Build a Spark Connect `Relation` from a SQL string. Shared by the SQL job
/// path (`execute_sql`) and runnable pipeline flows (`define_flow`).
fn sql_relation(query: &str) -> Relation {
// `args`/`pos_args` are deprecated in the proto but prost still emits them
// as struct fields we must populate; silence the lint locally.
#[allow(deprecated)]
let sql = Sql {
query: query.to_string(),
args: Default::default(),
pos_args: Vec::new(),
named_arguments: Default::default(),
pos_arguments: Vec::new(),
};
Relation {
common: None,
rel_type: Some(RelType::Sql(sql)),
}
}
/// A `'static`, boxed stream of viewer updates. Boxing keeps the public
/// streaming methods free of return-position-impl-Trait lifetime capture, so
/// callers can hold the stream while the arguments that produced it drop.
pub type UpdateStream = std::pin::Pin<Box<dyn Stream<Item = crate::Result<ViewerUpdate>> + Send>>;
/// Wrap a raw `ExecutePlanResponse` tonic stream as an [`UpdateStream`],
/// keeping only the responses the viewer cares about.
fn updates_of(stream: tonic::Streaming<proto::ExecutePlanResponse>) -> UpdateStream {
use futures_util::StreamExt;
Box::pin(stream.filter_map(|resp| async move {
match resp {
Ok(r) => {
let update = response_to_update(r);
// Skip building the per-event log fields entirely when the
// `knutpipe::run` target is below the active level (L).
let log_on = tracing::enabled!(target: "knutpipe::run", tracing::Level::INFO);
match &update {
Some(ViewerUpdate::Event(ev)) if log_on => tracing::info!(
target: "knutpipe::run", event = "pipeline_event",
level = ?ev.level, message = %ev.message, "run event",
),
Some(ViewerUpdate::Progress(p)) if log_on => tracing::info!(
target: "knutpipe::run", event = "job_progress",
fraction = p.fraction(), tasks = p.completed_tasks(),
total = p.total_tasks(), inflight = p.num_inflight_tasks,
"job progress",
),
_ => {}
}
update.map(Ok)
}
Err(s) => Some(Err(PipelineError::Transport(s.to_string()))),
}
}))
}
/// A connected Spark Connect pipelines client.
///
/// Owns a gRPC channel and a per-connection session id. Construct with
/// [`PipelinesClient::connect`].
pub struct PipelinesClient {
inner: SparkConnectServiceClient<
tonic::service::interceptor::InterceptedService<tonic::transport::Channel, AuthInterceptor>,
>,
session_id: String,
user_id: String,
/// The URL this client connected to — tagged onto [`crate::SparkRuntimeInfo`]
/// so the Spark view knows which endpoint it is reporting on.
connect_url: String,
}
impl PipelinesClient {
/// Connect to a Spark Connect server.
///
/// Accepts either a `sc://host:port[/;k=v…]` URL (the Spark Connect scheme)
/// or a plain `http://host:port`. The `sc://` session params are threaded
/// through: `user_id` sets the `UserContext.user_id` on every request,
/// `session_id` pins the request session (else a fresh UUID is generated),
/// and `token` is attached as a `Bearer authorization` metadata via an
/// interceptor. Params on a plain `http(s)://` URL are not expressible, so
/// that path uses a generated session and no auth.
pub async fn connect(url: &str) -> crate::Result<Self> {
let parsed = parse_sc_url(url)?;
let mut ep = tonic::transport::Endpoint::from_shared(parsed.endpoint.clone())
.map_err(|e| PipelineError::Transport(e.to_string()))?;
// **`use_ssl=true` has to reach the CHANNEL, not just the scheme string.**
// `Endpoint::from_shared("https://…")` alone does not enable TLS in tonic — it
// needs a `tls_config`, and without one the connect fails with a scheme error
// rather than silently downgrading. Asking for it explicitly is what makes the
// parsed `https` mean what an operator wrote.
if parsed.endpoint.starts_with("https://") {
ep = ep
.tls_config(tonic::transport::ClientTlsConfig::new().with_enabled_roots())
.map_err(|e| PipelineError::Transport(format!("TLS config for {url}: {e}")))?;
} else if parsed.token.is_some() {
// A bearer token about to ride a PLAINTEXT channel. Not fatal — a local
// `sc://localhost:15002` with a dev token is a real and reasonable setup —
// but it must never be silent, because the failure mode is a credential on
// the wire that nobody chose.
tracing::warn!(
target: "knutpipe::client", event = "token_over_cleartext",
url = %url,
"a `token` is set but the channel is PLAINTEXT (no `use_ssl=true`) — the \
bearer credential will cross the wire in clear",
);
}
let channel = ep
.connect()
.await
.map_err(|e| PipelineError::Transport(e.to_string()))?;
let session_id = parsed.session_id.clone().unwrap_or_else(new_uuid);
let user_id = parsed.user_id.clone().unwrap_or_default();
let interceptor = AuthInterceptor {
token: parsed.token.clone(),
};
tracing::info!(
target: "knutpipe::client", event = "connected",
url = %url, session = %session_id, user = %user_id,
authenticated = parsed.token.is_some(), "connected to Spark Connect",
);
Ok(PipelinesClient {
inner: SparkConnectServiceClient::with_interceptor(channel, interceptor),
session_id,
user_id,
connect_url: url.to_string(),
})
}
fn user_context(&self) -> UserContext {
UserContext {
user_id: self.user_id.clone(),
user_name: String::new(),
extensions: Vec::new(),
}
}
/// Wrap a `PipelineCommand` into an `ExecutePlanRequest` carrying a
/// command plan.
fn pipeline_request(&self, cmd: pipeline_command::CommandType) -> ExecutePlanRequest {
let command = Command {
command_type: Some(proto::command::CommandType::PipelineCommand(
proto::PipelineCommand {
command_type: Some(cmd),
},
)),
};
ExecutePlanRequest {
session_id: self.session_id.clone(),
user_context: Some(self.user_context()),
plan: Some(Plan {
op_type: Some(OpType::Command(command)),
}),
client_type: Some(concat!("knut-pipelines/", env!("CARGO_PKG_VERSION")).to_string()),
..Default::default()
}
}
/// Create a new dataflow graph on the server, returning its id.
pub async fn create_dataflow_graph(
&mut self,
default_catalog: Option<&str>,
default_database: Option<&str>,
) -> crate::Result<String> {
let cmd = pipeline_command::CommandType::CreateDataflowGraph(CreateDataflowGraph {
default_catalog: default_catalog.map(str::to_string),
default_database: default_database.map(str::to_string),
sql_conf: Default::default(),
});
let req = self.pipeline_request(cmd);
let mut stream = self
.inner
.execute_plan(req)
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.into_inner();
// Pull the command result off the response stream.
while let Some(resp) = stream
.message()
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
{
if let Some(ResponseType::PipelineCommandResult(r)) = resp.response_type
&& let Some(proto::pipeline_command_result::ResultType::CreateDataflowGraphResult(
g,
)) = r.result_type
{
return g.dataflow_graph_id.ok_or_else(|| {
PipelineError::Protocol("CreateDataflowGraph returned no id".into())
});
}
}
Err(PipelineError::Protocol(
"CreateDataflowGraph: no result on stream".into(),
))
}
/// Define one output (dataset node) against a graph. Drains the response
/// stream (the define is fire-and-confirm).
pub async fn define_output(&mut self, graph_id: &str, dataset: &Dataset) -> crate::Result<()> {
let table_details = proto::pipeline_command::define_output::TableDetails {
// `BTreeMap` -> the proto's `map<string, string>` (a prost HashMap).
// Order is irrelevant on the wire; the BTreeMap exists so the *spec*
// side stays diff-stable.
table_properties: dataset.properties.clone().into_iter().collect(),
partition_cols: dataset.partition_cols.clone(),
format: dataset.format.clone(),
clustering_columns: Vec::new(),
schema: dataset.schema.clone().map(|s| {
proto::pipeline_command::define_output::table_details::Schema::SchemaString(s)
}),
};
let define = DefineOutput {
dataflow_graph_id: Some(graph_id.to_string()),
output_name: Some(dataset.name.clone()),
output_type: Some(output_type_to_proto(dataset.output_type)),
comment: dataset.comment.clone(),
source_code_location: None,
details: Some(
proto::pipeline_command::define_output::Details::TableDetails(table_details),
),
};
self.run_to_completion(pipeline_command::CommandType::DefineOutput(define))
.await
}
/// Define one flow (edge) against a graph.
pub async fn define_flow(&mut self, graph_id: &str, flow: &Flow) -> crate::Result<()> {
let define = DefineFlow {
dataflow_graph_id: Some(graph_id.to_string()),
flow_name: Some(flow.name.clone()),
target_dataset_name: Some(flow.target.clone()),
sql_conf: Default::default(),
client_id: None,
source_code_location: None,
// Only set the `once` option for genuine one-time flows — Spark
// rejects the option being present at all for normal flows
// (DEFINE_FLOW_ONCE_OPTION_NOT_SUPPORTED).
once: flow.once.then_some(true),
// A flow with a query is *runnable* — its SQL relation is what
// StartRun executes to materialise the target dataset.
details: flow.query.as_deref().map(|q| {
proto::pipeline_command::define_flow::Details::RelationFlowDetails(
proto::pipeline_command::define_flow::WriteRelationFlowDetails {
relation: Some(sql_relation(q)),
},
)
}),
};
self.run_to_completion(pipeline_command::CommandType::DefineFlow(define))
.await
}
/// Register every dataset and flow of an in-memory [`DataflowGraph`] on a
/// freshly created server graph, returning the server graph id.
pub async fn define_graph(&mut self, graph: &DataflowGraph) -> crate::Result<String> {
graph.validate()?;
let id = self
.create_dataflow_graph(
graph.default_catalog.as_deref(),
graph.default_database.as_deref(),
)
.await?;
for d in &graph.datasets {
self.define_output(&id, d).await?;
}
for f in &graph.flows {
self.define_flow(&id, f).await?;
}
Ok(id)
}
/// Start a run and stream its [`PipelineRunEvent`]s into `sink` until the
/// stream ends.
pub async fn start_run<S: PipelineEventSink>(
&mut self,
run: &PipelineRun,
sink: &mut S,
) -> crate::Result<()> {
let cmd = pipeline_command::CommandType::StartRun(StartRun {
dataflow_graph_id: run.graph_id.clone(),
full_refresh_selection: Vec::new(),
full_refresh_all: Some(run.full_refresh_all),
refresh_selection: run.refresh_selection.clone(),
dry: Some(run.dry),
storage: run.storage.clone(),
});
let req = self.pipeline_request(cmd);
let mut stream = self
.inner
.execute_plan(req)
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.into_inner();
let result = async {
while let Some(resp) = stream
.message()
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
{
if let Some(ResponseType::PipelineEventResult(r)) = resp.response_type
&& let Some(ev) = r.event
{
sink.on_event(&event_from_proto(ev))?;
}
}
Ok(())
}
.await;
sink.on_finish()?;
result
}
/// A [`PipelineRunEvent`] stream for callers that want to drive the loop
/// themselves (e.g. an egui/ratatui frontend selecting over channels)
/// rather than push into a [`PipelineEventSink`].
pub async fn run_event_stream(
&mut self,
run: &PipelineRun,
) -> crate::Result<impl Stream<Item = crate::Result<PipelineRunEvent>>> {
use futures_util::StreamExt;
let cmd = pipeline_command::CommandType::StartRun(StartRun {
dataflow_graph_id: run.graph_id.clone(),
full_refresh_selection: Vec::new(),
full_refresh_all: Some(run.full_refresh_all),
refresh_selection: run.refresh_selection.clone(),
dry: Some(run.dry),
storage: run.storage.clone(),
});
let req = self.pipeline_request(cmd);
let stream = self
.inner
.execute_plan(req)
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.into_inner();
Ok(stream.filter_map(|resp| async move {
match resp {
Ok(r) => match r.response_type {
Some(ResponseType::PipelineEventResult(per)) => {
per.event.map(|ev| Ok(event_from_proto(ev)))
}
_ => None,
},
Err(s) => Some(Err(PipelineError::Transport(s.to_string()))),
}
}))
}
/// Unified live stream for a **pipeline run**: yields both pipeline
/// [`ViewerUpdate::Event`]s and Spark [`ViewerUpdate::Progress`] snapshots
/// as they arrive. This is what the GUI/CLI watch loop consumes so a run
/// shows its log *and* its job progress bars at once.
pub async fn run_updates(&mut self, run: &PipelineRun) -> crate::Result<UpdateStream> {
let cmd = pipeline_command::CommandType::StartRun(StartRun {
dataflow_graph_id: run.graph_id.clone(),
full_refresh_selection: Vec::new(),
full_refresh_all: Some(run.full_refresh_all),
refresh_selection: run.refresh_selection.clone(),
dry: Some(run.dry),
storage: run.storage.clone(),
});
let req = self.pipeline_request(cmd);
let stream = self
.inner
.execute_plan(req)
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.into_inner();
Ok(updates_of(stream))
}
/// Execute an arbitrary **Spark SQL** query over Spark Connect and stream
/// its job progress as [`ViewerUpdate::Progress`] snapshots. This is the
/// *plain Spark jobs* path — no pipeline needed — so the viewer doubles as
/// a live progress monitor for any Spark Connect query.
///
/// Result rows (arrow batches) are intentionally not surfaced; the viewer
/// is about job *progress*, not result sets.
///
/// **Progress is server-driven, not requested here.** Spark Connect pushes
/// `ExecutePlanResponse.execution_progress` (base.proto field 18)
/// unconditionally while a plan runs, paced by the server-side
/// `spark.connect.progress.reportInterval` config — there is *no*
/// per-request toggle. `ExecutePlanRequest.request_options` (field 5) is
/// documented "currently unused" and only carries `ReattachOptions` /
/// `ResultChunkingOptions` / an extension `Any` — none of which control
/// progress. So we leave it at its default (empty) and simply consume the
/// `ExecutionProgress` frames the server volunteers.
pub async fn execute_sql(&mut self, sql: &str) -> crate::Result<UpdateStream> {
let relation = sql_relation(sql);
let req = ExecutePlanRequest {
session_id: self.session_id.clone(),
user_context: Some(self.user_context()),
plan: Some(Plan {
op_type: Some(OpType::Root(relation)),
}),
client_type: Some(concat!("knut-pipelines/", env!("CARGO_PKG_VERSION")).to_string()),
..Default::default()
};
let stream = self
.inner
.execute_plan(req)
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.into_inner();
Ok(updates_of(stream))
}
/// Count the rows in a table or view by summing the `row_count` of every
/// returned `ArrowBatch`.
///
/// Spark's SDP run events carry **no** per-flow output count — the
/// `PipelineEvent` proto is `{timestamp, message}` only (SPARK-51727) — so a
/// viewer that wants row counts must count back after the run. This reads
/// the proto `ArrowBatch.row_count` field directly and never decodes the
/// Arrow payload, so it stays dependency-free (no `arrow` crate). It does
/// stream the rows, so use it for the modest materialised views a run
/// produces, not arbitrarily large tables.
pub async fn count_rows(&mut self, table: &str) -> crate::Result<u64> {
let relation = sql_relation(&format!("SELECT * FROM {table}"));
let req = ExecutePlanRequest {
session_id: self.session_id.clone(),
user_context: Some(self.user_context()),
plan: Some(Plan {
op_type: Some(OpType::Root(relation)),
}),
client_type: Some(concat!("knut-pipelines/", env!("CARGO_PKG_VERSION")).to_string()),
..Default::default()
};
let mut stream = self
.inner
.execute_plan(req)
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.into_inner();
let mut total: u64 = 0;
while let Some(resp) = stream
.message()
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
{
if let Some(ResponseType::ArrowBatch(b)) = resp.response_type {
if b.row_count > 0 {
total += b.row_count as u64;
}
}
}
Ok(total)
}
// ── Spark runtime + Structured-Streaming VIEW (the korp Spark pane) ──────
//
// These drive the streaming query manager / streaming query commands
// (commands.proto) over the same `ExecutePlan` RPC the pipeline commands
// use, and fold the proto results into the neutral `spark_view` model korp
// renders. Read verbs (list/status/recent-progress) + the one control verb
// the Connect protocol exposes (stop). Cluster executor/worker topology is
// NOT on the Connect wire (Spark's REST UI only) — deliberately not faked.
/// Wrap an arbitrary top-level [`Command`] into an `ExecutePlanRequest`.
fn command_request(&self, command_type: proto::command::CommandType) -> ExecutePlanRequest {
ExecutePlanRequest {
session_id: self.session_id.clone(),
user_context: Some(self.user_context()),
plan: Some(Plan {
op_type: Some(OpType::Command(Command {
command_type: Some(command_type),
})),
}),
client_type: Some(concat!("knut-pipelines/", env!("CARGO_PKG_VERSION")).to_string()),
..Default::default()
}
}
/// The Spark version the connected server reports — the coarse runtime fact
/// a Connect session can answer (`AnalyzePlan` `SparkVersion`, no arrow
/// decode). Returns a [`SparkRuntimeInfo`] tagged with this connection's URL.
pub async fn runtime_info(&mut self) -> crate::Result<crate::SparkRuntimeInfo> {
use proto::analyze_plan_request::{Analyze, SparkVersion};
let req = proto::AnalyzePlanRequest {
session_id: self.session_id.clone(),
user_context: Some(self.user_context()),
client_type: Some(concat!("knut-pipelines/", env!("CARGO_PKG_VERSION")).to_string()),
analyze: Some(Analyze::SparkVersion(SparkVersion {})),
..Default::default()
};
let resp = self
.inner
.analyze_plan(req)
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.into_inner();
let version = match resp.result {
Some(proto::analyze_plan_response::Result::SparkVersion(v)) => v.version,
_ => {
return Err(PipelineError::Protocol(
"AnalyzePlan(SparkVersion) returned no version".into(),
));
}
};
Ok(crate::SparkRuntimeInfo::new(
version,
self.connect_url.clone(),
))
}
/// List the **active Structured Streaming queries** on the server
/// (`StreamingQueryManager.active()`), folded into the neutral
/// [`StreamingQueryInfo`] model.
pub async fn list_active_streams(&mut self) -> crate::Result<Vec<crate::StreamingQueryInfo>> {
use proto::streaming_query_manager_command::Command as MgrCmd;
let cmd = proto::command::CommandType::StreamingQueryManagerCommand(
proto::StreamingQueryManagerCommand {
command: Some(MgrCmd::Active(true)),
},
);
let mut stream = self
.inner
.execute_plan(self.command_request(cmd))
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.into_inner();
while let Some(resp) = stream
.message()
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
{
if let Some(ResponseType::StreamingQueryManagerCommandResult(r)) = resp.response_type
&& let Some(proto::streaming_query_manager_command_result::ResultType::Active(a)) =
r.result_type
{
return Ok(a
.active_queries
.into_iter()
.map(|q| {
let (id, run_id) = q.id.map(|i| (i.id, i.run_id)).unwrap_or_default();
crate::StreamingQueryInfo::new(id, run_id, q.name)
})
.collect());
}
}
Err(PipelineError::Protocol(
"StreamingQueryManagerCommand(active): no result on stream".into(),
))
}
/// The live [`StreamingQueryStatus`] of one query (`StreamingQuery.status()`).
pub async fn stream_status(
&mut self,
id: &str,
run_id: &str,
) -> crate::Result<crate::StreamingQueryStatus> {
use proto::streaming_query_command::Command as QCmd;
let cmd =
proto::command::CommandType::StreamingQueryCommand(proto::StreamingQueryCommand {
query_id: Some(proto::StreamingQueryInstanceId {
id: id.to_string(),
run_id: run_id.to_string(),
}),
command: Some(QCmd::Status(true)),
});
let mut stream = self
.inner
.execute_plan(self.command_request(cmd))
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.into_inner();
while let Some(resp) = stream
.message()
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
{
if let Some(ResponseType::StreamingQueryCommandResult(r)) = resp.response_type
&& let Some(proto::streaming_query_command_result::ResultType::Status(s)) =
r.result_type
{
return Ok(crate::StreamingQueryStatus::new(
s.status_message,
s.is_active,
s.is_data_available,
s.is_trigger_active,
));
}
}
Err(PipelineError::Protocol(
"StreamingQueryCommand(status): no result on stream".into(),
))
}
/// The recent per-micro-batch progress reports of one query as JSON strings
/// (`StreamingQuery.recentProgress()`) — the throughput/lag timeline korp
/// charts. Left as JSON (not decoded) to stay arrow/dependency-free here.
pub async fn stream_recent_progress(
&mut self,
id: &str,
run_id: &str,
) -> crate::Result<Vec<String>> {
use proto::streaming_query_command::Command as QCmd;
let cmd =
proto::command::CommandType::StreamingQueryCommand(proto::StreamingQueryCommand {
query_id: Some(proto::StreamingQueryInstanceId {
id: id.to_string(),
run_id: run_id.to_string(),
}),
command: Some(QCmd::RecentProgress(true)),
});
let mut stream = self
.inner
.execute_plan(self.command_request(cmd))
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.into_inner();
while let Some(resp) = stream
.message()
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
{
if let Some(ResponseType::StreamingQueryCommandResult(r)) = resp.response_type
&& let Some(proto::streaming_query_command_result::ResultType::RecentProgress(p)) =
r.result_type
{
return Ok(p.recent_progress_json);
}
}
Err(PipelineError::Protocol(
"StreamingQueryCommand(recent_progress): no result on stream".into(),
))
}
/// **Stop** a running streaming query (`StreamingQuery.stop()`) — the one
/// control verb the Connect protocol exposes for streaming queries. Drains
/// the confirmation stream.
pub async fn stop_stream(&mut self, id: &str, run_id: &str) -> crate::Result<()> {
use proto::streaming_query_command::Command as QCmd;
let cmd =
proto::command::CommandType::StreamingQueryCommand(proto::StreamingQueryCommand {
query_id: Some(proto::StreamingQueryInstanceId {
id: id.to_string(),
run_id: run_id.to_string(),
}),
command: Some(QCmd::Stop(true)),
});
let mut stream = self
.inner
.execute_plan(self.command_request(cmd))
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.into_inner();
while stream
.message()
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.is_some()
{}
Ok(())
}
/// Drive a fire-and-confirm pipeline command, draining its stream.
async fn run_to_completion(&mut self, cmd: pipeline_command::CommandType) -> crate::Result<()> {
let req = self.pipeline_request(cmd);
let mut stream = self
.inner
.execute_plan(req)
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.into_inner();
while stream
.message()
.await
.map_err(|s| PipelineError::Transport(s.to_string()))?
.is_some()
{}
Ok(())
}
}
/// A parsed Spark Connect connection string: the tonic HTTP endpoint plus the
/// session parameters carried in the `sc://` URL's `;k=v` tail.
///
/// The Spark Connect URL grammar is
/// `sc://host:port/;user_id=<u>;session_id=<s>;token=<t>;use_ssl=<b>;…` — the
/// params live after the host:port, introduced by `/;` and separated by `;`.
/// We recognise the ones that shape a session (`user_id`, `session_id`, `token`)
/// and the one that shapes the TRANSPORT (`use_ssl`).
///
/// # `use_ssl` used to be dropped, and that put a bearer token on the wire in clear
///
/// Every unrecognised param fell into a `_ => {}` arm and the endpoint was hard-coded
/// `http://`. So `sc://host:15002/;token=SECRET;use_ssl=true` — an operator asking for
/// TLS, in Spark's own documented grammar — built a **plaintext** channel and attached
/// the token to it as an `authorization` header anyway. Nothing warned. The unit test
/// beside this even asserted the endpoint came back `http://`, with a comment calling
/// `use_ssl` "an unrecognised param which is ignored": the defect was pinned as
/// intended behaviour.
///
/// Silently downgrading a caller's explicit TLS request is worse than refusing it.
#[derive(Debug, Clone, PartialEq, Eq, Default)]
struct ScConnect {
/// `http://host:port`, or `https://host:port` when `use_ssl=true`.
endpoint: String,
/// `user_id` param → the `UserContext.user_id` on every request.
user_id: Option<String>,
/// `session_id` param → the request `session_id` (else a fresh UUID).
session_id: Option<String>,
/// `token` param → a `Bearer` `authorization` metadata on every request.
token: Option<String>,
}
/// Parse a `sc://host:port[/;k=v…]` Spark Connect URL into an [`ScConnect`].
/// Plain `http(s)://host:port` passes through as a bare endpoint (no params).
fn parse_sc_url(url: &str) -> crate::Result<ScConnect> {
if let Some(rest) = url.strip_prefix("sc://") {
// host:port is everything up to the first '/' or ';'.
let end = rest.find(['/', ';']).unwrap_or(rest.len());
let host_port = &rest[..end];
if host_port.is_empty() {
return Err(PipelineError::Transport(format!("empty host in {url}")));
}
let mut out = ScConnect::default();
let mut use_ssl = false;
// The remainder holds the `;`-separated `k=v` params (after the `/;`).
for kv in rest[end..].split(['/', ';']).filter(|s| !s.is_empty()) {
if let Some((k, v)) = kv.split_once('=') {
match k {
"user_id" => out.user_id = Some(v.to_string()),
"session_id" => out.session_id = Some(v.to_string()),
"token" => out.token = Some(v.to_string()),
// Spark's own spelling is `true`/`false`; anything else is a
// malformed value and must NOT quietly read as "no TLS" —
// see the refusal below.
"use_ssl" => match v {
"true" => use_ssl = true,
"false" => use_ssl = false,
other => {
return Err(PipelineError::Transport(format!(
"use_ssl={other} in {url} is neither `true` nor `false` — \
refusing rather than guessing, because guessing `false` \
would put any token in this URL on the wire in clear"
)));
}
},
_ => {} // user_agent, … — not session- or transport-shaping here
}
}
}
out.endpoint = format!("{}://{host_port}", if use_ssl { "https" } else { "http" });
Ok(out)
} else if url.starts_with("http://") || url.starts_with("https://") {
Ok(ScConnect {
endpoint: url.to_string(),
..Default::default()
})
} else {
Err(PipelineError::Transport(format!(
"unrecognised Spark Connect URL scheme: {url}"
)))
}
}
/// A tonic interceptor that attaches a `Bearer` token as the `authorization`
/// metadata on every request, when the `sc://` URL carried a `token` param.
/// A `None` token is a pass-through, so unauthenticated servers are unaffected.
#[derive(Clone)]
struct AuthInterceptor {
token: Option<String>,
}
impl tonic::service::Interceptor for AuthInterceptor {
fn call(
&mut self,
mut request: tonic::Request<()>,
) -> std::result::Result<tonic::Request<()>, tonic::Status> {
if let Some(t) = &self.token {
let value = tonic::metadata::MetadataValue::try_from(format!("Bearer {t}"))
.map_err(|_| tonic::Status::invalid_argument("invalid sc:// token"))?;
request.metadata_mut().insert("authorization", value);
}
Ok(request)
}
}
/// Minimal UUIDv4-ish session id (random hex in the canonical shape). Avoids a
/// `uuid` dependency; the server only needs a stable, unique-per-session
/// string.
fn new_uuid() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
let a = nanos as u64;
let b = (nanos >> 64) as u64 ^ (std::process::id() as u64).rotate_left(17);
format!(
"{:08x}-{:04x}-4{:03x}-{:04x}-{:012x}",
(a & 0xffff_ffff) as u32,
((a >> 32) & 0xffff) as u16,
((a >> 48) & 0x0fff) as u16,
(0x8000 | (b & 0x3fff)) as u16,
b & 0xffff_ffff_ffff,
)
}
#[cfg(test)]
mod grpc_tests {
use super::*;
/// Endpoint translation: `sc://` → `http://` (params dropped from the
/// endpoint), plain http passes through, unknown schemes error.
#[test]
fn sc_url_parsing() {
let ep = |u| parse_sc_url(u).map(|p| p.endpoint);
assert_eq!(ep("sc://spark:15002").unwrap(), "http://spark:15002");
assert_eq!(ep("sc://h:15002/;user_id=x").unwrap(), "http://h:15002");
assert_eq!(
ep("http://localhost:15002").unwrap(),
"http://localhost:15002"
);
assert!(parse_sc_url("grpc://nope").is_err());
}
#[test]
fn sc_url_session_params_are_parsed() {
// The single-param case the backlog pins: user_id threads through.
let p = parse_sc_url("sc://h:15002/;user_id=alice").unwrap();
assert_eq!(p.endpoint, "http://h:15002");
assert_eq!(p.user_id.as_deref(), Some("alice"));
assert_eq!(p.session_id, None);
assert_eq!(p.token, None);
// The full `;`-separated param list, INCLUDING `use_ssl` — which this
// assertion used to pin the wrong way round. It read
// `assert_eq!(p.endpoint, "http://spark:15002")` under a comment calling
// `use_ssl` "an unrecognised param which is ignored", so the test froze the
// defect as intended behaviour: an operator asking for TLS in Spark's own
// documented grammar got a PLAINTEXT channel with their bearer token attached
// to it, and nothing said so.
let p =
parse_sc_url("sc://spark:15002/;user_id=bob;session_id=sess-7;token=xyz;use_ssl=true")
.unwrap();
assert_eq!(p.endpoint, "https://spark:15002", "use_ssl=true means TLS");
assert_eq!(p.user_id.as_deref(), Some("bob"));
assert_eq!(p.session_id.as_deref(), Some("sess-7"));
assert_eq!(p.token.as_deref(), Some("xyz"));
// …and `use_ssl=false` is plaintext, explicitly — the default is unchanged, so
// this cannot pass by the scheme having become unconditionally `https`.
let p = parse_sc_url("sc://spark:15002/;token=xyz;use_ssl=false").unwrap();
assert_eq!(p.endpoint, "http://spark:15002");
assert_eq!(p.token.as_deref(), Some("xyz"));
// A MALFORMED value refuses rather than falling back to plaintext. `use_ssl=1`,
// `use_ssl=yes`, a typo — every one of them is an operator who believed they
// asked for TLS, and reading them as `false` would put the token beside them on
// the wire in clear. Guessing is the failure; refusing is the fix.
for bad in ["1", "yes", "TRUE", ""] {
let e = parse_sc_url(&format!("sc://spark:15002/;token=xyz;use_ssl={bad}"))
.expect_err("a malformed use_ssl must refuse");
let msg = format!("{e}");
assert!(msg.contains("use_ssl"), "names the parameter: {msg}");
assert!(msg.contains("clear"), "says WHY refusing beats guessing: {msg}");
}
// No params → all None, endpoint still recovered.
let p = parse_sc_url("sc://spark:15002").unwrap();
assert_eq!(p.endpoint, "http://spark:15002");
assert!(p.user_id.is_none() && p.session_id.is_none() && p.token.is_none());
// A plain http endpoint carries no params.
let p = parse_sc_url("http://localhost:15002").unwrap();
assert_eq!(p.endpoint, "http://localhost:15002");
assert!(p.user_id.is_none());
}
#[test]
fn auth_interceptor_sets_bearer_only_when_token_present() {
use tonic::service::Interceptor;
// With a token, every request gets a `Bearer` authorization metadata.
let mut with = AuthInterceptor {
token: Some("xyz".into()),
};
let req = with.call(tonic::Request::new(())).unwrap();
assert_eq!(req.metadata().get("authorization").unwrap(), "Bearer xyz");
// Without a token, the request is untouched (unauthenticated servers).
let mut without = AuthInterceptor { token: None };
let req = without.call(tonic::Request::new(())).unwrap();
assert!(req.metadata().get("authorization").is_none());
}
#[test]
fn uuid_shape() {
let u = new_uuid();
assert_eq!(u.len(), 36);
assert_eq!(u.as_bytes()[14], b'4'); // version nibble
}
#[test]
fn output_type_maps_to_proto() {
// Each model OutputType maps onto the matching generated enum's i32.
for (t, want) in [
(
OutputType::MaterializedView,
ProtoOutputType::MaterializedView,
),
(OutputType::Table, ProtoOutputType::Table),
(OutputType::TemporaryView, ProtoOutputType::TemporaryView),
(OutputType::Sink, ProtoOutputType::Sink),
(OutputType::Unspecified, ProtoOutputType::Unspecified),
] {
assert_eq!(output_type_to_proto(t), want as i32, "mapping {t:?}");
}
}
#[test]
fn progress_converts_from_proto() {
let p = proto::execute_plan_response::ExecutionProgress {
stages: vec![
proto::execute_plan_response::execution_progress::StageInfo {
stage_id: 3,
num_tasks: 8,
num_completed_tasks: 5,
input_bytes_read: 4096,
done: false,
},
],
num_inflight_tasks: 3,
};
let m = progress_from_proto(p);
assert_eq!(m.stages.len(), 1);
assert_eq!(m.stages[0].stage_id, 3);
assert_eq!(m.completed_tasks(), 5);
assert_eq!(m.num_inflight_tasks, 3);
}
#[test]
fn response_progress_becomes_update() {
let resp = proto::ExecutePlanResponse {
response_type: Some(ResponseType::ExecutionProgress(
proto::execute_plan_response::ExecutionProgress {
stages: vec![],
num_inflight_tasks: 1,
},
)),
..Default::default()
};
match response_to_update(resp) {
Some(ViewerUpdate::Progress(p)) => assert_eq!(p.num_inflight_tasks, 1),
other => panic!("expected Progress update, got {other:?}"),
}
}
#[test]
fn event_level_derived_from_message() {
let ev = event_from_proto(proto::PipelineEvent {
timestamp: None,
message: Some("Flow failed: boom".into()),
});
assert_eq!(ev.level, EventLevel::Error);
}
}