spark_connect/
streaming.rs

1//! Structured Streaming support mirroring `pyspark.sql.connect.streaming`.
2//!
3//! Provides DataStreamReader, DataStreamWriter, StreamingQuery, and StreamingQueryManager
4//! for building and executing streaming workloads.
5
6use std::collections::HashMap;
7use std::sync::{Arc, Mutex};
8use uuid;
9
10use spark_connect_core::client::ReattachableResponseStream;
11use spark_connect_core::error::{Result, SparkError};
12use spark_connect_core::runtime::block_on;
13use spark_connect_proto as proto;
14
15use crate::dataframe::DataFrame;
16use crate::plan::LogicalPlan;
17use crate::readwriter::ReadType;
18use crate::session::SparkSession;
19use crate::udf::PythonUDFPayload;
20
21/// DataStreamReader for reading streaming data from various sources.
22///
23/// Mirrors `pyspark.sql.connect.streaming.DataStreamReader`.
24pub struct DataStreamReader {
25    session: SparkSession,
26    format: Option<String>,
27    schema: String,
28    options: HashMap<String, String>,
29    source_name: Option<String>,
30}
31
32impl DataStreamReader {
33    /// Create a new DataStreamReader.
34    pub(crate) fn new(session: SparkSession) -> Self {
35        DataStreamReader {
36            session,
37            format: None,
38            schema: String::new(),
39            options: HashMap::new(),
40            source_name: None,
41        }
42    }
43
44    /// Set the format/source type (e.g., "rate", "socket", "kafka", "json", "parquet", "csv").
45    pub fn format(mut self, source: &str) -> Self {
46        self.format = Some(source.to_string());
47        self
48    }
49
50    /// Set the schema from a DDL string or JSON string.
51    pub fn schema(mut self, schema: impl Into<String>) -> Self {
52        self.schema = schema.into();
53        self
54    }
55
56    /// Set a single option key-value pair.
57    pub fn option(mut self, key: &str, value: &str) -> Self {
58        self.options.insert(key.to_string(), value.to_string());
59        self
60    }
61
62    /// Set multiple options.
63    pub fn options(mut self, options: HashMap<String, String>) -> Self {
64        self.options.extend(options);
65        self
66    }
67
68    /// Set the source name for checkpoint stability.
69    pub fn name(mut self, source_name: &str) -> Self {
70        self.source_name = Some(source_name.to_string());
71        self
72    }
73
74    /// Load streaming data from the specified path(s).
75    pub fn load(self, path: Option<&str>) -> DataFrame {
76        let paths = path.map(|p| vec![p.to_string()]);
77        let plan = LogicalPlan::Read {
78            read_type: ReadType::DataSource {
79                format: self.format.clone(),
80                schema: if self.schema.is_empty() {
81                    None
82                } else {
83                    Some(self.schema.clone())
84                },
85                options: self.options.clone(),
86                paths: paths.unwrap_or_default(),
87                predicates: vec![],
88                source_name: self.source_name.clone(),
89            },
90            is_streaming: true,
91        };
92        DataFrame::new(self.session, plan)
93    }
94
95    /// Read a named streaming table.
96    pub fn table(self, table_name: &str) -> DataFrame {
97        let plan = LogicalPlan::Read {
98            read_type: ReadType::NamedTable {
99                table_name: table_name.to_string(),
100                options: self.options.clone(),
101            },
102            is_streaming: true,
103        };
104        DataFrame::new(self.session, plan)
105    }
106
107    /// Read the streaming CDC changes of a named table. Mirrors
108    /// `DataStreamReader.changes` (`is_streaming = true`).
109    pub fn changes(self, table_name: &str) -> DataFrame {
110        let plan = LogicalPlan::RelationChanges {
111            table_name: table_name.to_string(),
112            options: self.options.clone(),
113            is_streaming: Some(true),
114        };
115        DataFrame::new(self.session, plan)
116    }
117
118    /// Read streaming JSON data from a path.
119    pub fn json(mut self, path: &str) -> DataFrame {
120        self.format = Some("json".to_string());
121        let paths = vec![path.to_string()];
122        let plan = LogicalPlan::Read {
123            read_type: ReadType::DataSource {
124                format: self.format.clone(),
125                schema: if self.schema.is_empty() {
126                    None
127                } else {
128                    Some(self.schema.clone())
129                },
130                options: self.options.clone(),
131                paths,
132                predicates: vec![],
133                source_name: self.source_name.clone(),
134            },
135            is_streaming: true,
136        };
137        DataFrame::new(self.session, plan)
138    }
139
140    /// Read streaming Parquet data from a path.
141    pub fn parquet(mut self, path: &str) -> DataFrame {
142        self.format = Some("parquet".to_string());
143        let paths = vec![path.to_string()];
144        let plan = LogicalPlan::Read {
145            read_type: ReadType::DataSource {
146                format: self.format.clone(),
147                schema: if self.schema.is_empty() {
148                    None
149                } else {
150                    Some(self.schema.clone())
151                },
152                options: self.options.clone(),
153                paths,
154                predicates: vec![],
155                source_name: self.source_name.clone(),
156            },
157            is_streaming: true,
158        };
159        DataFrame::new(self.session, plan)
160    }
161
162    /// Read streaming CSV data from a path.
163    pub fn csv(mut self, path: &str) -> DataFrame {
164        self.format = Some("csv".to_string());
165        let paths = vec![path.to_string()];
166        let plan = LogicalPlan::Read {
167            read_type: ReadType::DataSource {
168                format: self.format.clone(),
169                schema: if self.schema.is_empty() {
170                    None
171                } else {
172                    Some(self.schema.clone())
173                },
174                options: self.options.clone(),
175                paths,
176                predicates: vec![],
177                source_name: self.source_name.clone(),
178            },
179            is_streaming: true,
180        };
181        DataFrame::new(self.session, plan)
182    }
183
184    /// Read streaming ORC data from a path.
185    pub fn orc(mut self, path: &str) -> DataFrame {
186        self.format = Some("orc".to_string());
187        let paths = vec![path.to_string()];
188        let plan = LogicalPlan::Read {
189            read_type: ReadType::DataSource {
190                format: self.format.clone(),
191                schema: if self.schema.is_empty() {
192                    None
193                } else {
194                    Some(self.schema.clone())
195                },
196                options: self.options.clone(),
197                paths,
198                predicates: vec![],
199                source_name: self.source_name.clone(),
200            },
201            is_streaming: true,
202        };
203        DataFrame::new(self.session, plan)
204    }
205
206    /// Read streaming text data from a path.
207    pub fn text(mut self, path: &str) -> DataFrame {
208        self.format = Some("text".to_string());
209        let paths = vec![path.to_string()];
210        let plan = LogicalPlan::Read {
211            read_type: ReadType::DataSource {
212                format: self.format.clone(),
213                schema: if self.schema.is_empty() {
214                    None
215                } else {
216                    Some(self.schema.clone())
217                },
218                options: self.options.clone(),
219                paths,
220                predicates: vec![],
221                source_name: self.source_name.clone(),
222            },
223            is_streaming: true,
224        };
225        DataFrame::new(self.session, plan)
226    }
227}
228
229/// Trigger for streaming queries.
230#[derive(Debug, Clone)]
231pub enum Trigger {
232    /// Process every `interval` milliseconds or duration string (e.g., "10 seconds").
233    ProcessingTime(String),
234    /// Process only once.
235    Once,
236    /// Process as soon as available data arrives.
237    AvailableNow,
238    /// Continuous processing with checkpoint interval.
239    Continuous(String),
240}
241
242/// DataStreamWriter for writing streaming data to various sinks.
243///
244/// Mirrors `pyspark.sql.connect.streaming.DataStreamWriter`.
245pub struct DataStreamWriter {
246    session: SparkSession,
247    plan: LogicalPlan,
248    format: Option<String>,
249    output_mode: Option<String>,
250    options: HashMap<String, String>,
251    partitioning_columns: Vec<String>,
252    clustering_columns: Vec<String>,
253    query_name: Option<String>,
254    trigger: Option<Trigger>,
255    path: Option<String>,
256    table_name: Option<String>,
257    foreach_batch_payload: Option<PythonUDFPayload>,
258    foreach_payload: Option<PythonUDFPayload>,
259}
260
261impl DataStreamWriter {
262    /// Create a new DataStreamWriter.
263    pub(crate) fn new(session: SparkSession, plan: LogicalPlan) -> Self {
264        DataStreamWriter {
265            session,
266            plan,
267            format: None,
268            output_mode: None,
269            options: HashMap::new(),
270            partitioning_columns: vec![],
271            clustering_columns: vec![],
272            query_name: None,
273            trigger: None,
274            path: None,
275            table_name: None,
276            foreach_batch_payload: None,
277            foreach_payload: None,
278        }
279    }
280
281    /// Set the output mode ("append", "update", "complete").
282    pub fn output_mode(mut self, mode: &str) -> Self {
283        self.output_mode = Some(mode.to_string());
284        self
285    }
286
287    /// Set the format/sink type (e.g., "parquet", "json", "csv", "console", "noop", "kafka").
288    pub fn format(mut self, source: &str) -> Self {
289        self.format = Some(source.to_string());
290        self
291    }
292
293    /// Set a single option key-value pair.
294    pub fn option(mut self, key: &str, value: &str) -> Self {
295        self.options.insert(key.to_string(), value.to_string());
296        self
297    }
298
299    /// Set multiple options.
300    pub fn options(mut self, options: HashMap<String, String>) -> Self {
301        self.options.extend(options);
302        self
303    }
304
305    /// Set partitioning columns.
306    pub fn partition_by(mut self, columns: Vec<&str>) -> Self {
307        self.partitioning_columns = columns.iter().map(|s| s.to_string()).collect();
308        self
309    }
310
311    /// Set clustering columns.
312    pub fn cluster_by(mut self, columns: Vec<&str>) -> Self {
313        self.clustering_columns = columns.iter().map(|s| s.to_string()).collect();
314        self
315    }
316
317    /// Set the query name.
318    pub fn query_name(mut self, name: &str) -> Self {
319        self.query_name = Some(name.to_string());
320        self
321    }
322
323    /// Set the trigger type.
324    pub fn trigger(mut self, trigger: Trigger) -> Self {
325        self.trigger = Some(trigger);
326        self
327    }
328
329    /// Set a foreach batch function (PythonUDF payload).
330    pub fn foreach_batch(mut self, payload: PythonUDFPayload) -> Self {
331        self.foreach_batch_payload = Some(payload);
332        self
333    }
334
335    /// Set a foreach function (PythonUDF payload).
336    pub fn foreach(mut self, payload: PythonUDFPayload) -> Self {
337        self.foreach_payload = Some(payload);
338        self
339    }
340
341    /// Start the streaming query writing to a path, returning a StreamingQuery handle.
342    pub fn start(mut self, path: &str) -> Result<StreamingQuery> {
343        self.path = Some(path.to_string());
344        self._start_internal()
345    }
346
347    /// Start the streaming query writing to a table, returning a StreamingQuery handle.
348    pub fn to_table(mut self, table_name: &str) -> Result<StreamingQuery> {
349        self.table_name = Some(table_name.to_string());
350        self._start_internal()
351    }
352
353    /// Internal method to build and execute the write stream command.
354    fn _start_internal(self) -> Result<StreamingQuery> {
355        let mut write_op = proto::WriteStreamOperationStart::default();
356        write_op.input = Some(self.plan.to_proto());
357
358        if let Some(fmt) = &self.format {
359            write_op.format = fmt.clone();
360        }
361
362        if let Some(mode) = &self.output_mode {
363            write_op.output_mode = mode.clone();
364        }
365
366        write_op.options = self.options;
367        write_op.partitioning_column_names = self.partitioning_columns;
368        write_op.clustering_column_names = self.clustering_columns;
369
370        if let Some(name) = &self.query_name {
371            write_op.query_name = name.clone();
372        }
373
374        if let Some(trigger) = &self.trigger {
375            match trigger {
376                Trigger::ProcessingTime(interval) => {
377                    write_op.trigger = Some(
378                        proto::write_stream_operation_start::Trigger::ProcessingTimeInterval(
379                            interval.clone(),
380                        ),
381                    );
382                }
383                Trigger::Once => {
384                    write_op.trigger =
385                        Some(proto::write_stream_operation_start::Trigger::Once(true));
386                }
387                Trigger::AvailableNow => {
388                    write_op.trigger = Some(
389                        proto::write_stream_operation_start::Trigger::AvailableNow(true),
390                    );
391                }
392                Trigger::Continuous(interval) => {
393                    write_op.trigger = Some(
394                        proto::write_stream_operation_start::Trigger::ContinuousCheckpointInterval(
395                            interval.clone(),
396                        ),
397                    );
398                }
399            }
400        }
401
402        if let Some(path) = &self.path {
403            // A memory/console/foreach sink has no path; only set the destination for a
404            // real (non-empty) path so those sinks send an unset sink_destination.
405            if !path.is_empty() {
406                write_op.sink_destination = Some(
407                    proto::write_stream_operation_start::SinkDestination::Path(path.clone()),
408                );
409            }
410        }
411
412        if let Some(table) = &self.table_name {
413            write_op.sink_destination = Some(
414                proto::write_stream_operation_start::SinkDestination::TableName(table.clone()),
415            );
416        }
417
418        if let Some(foreach_batch) = &self.foreach_batch_payload {
419            let mut foreach_func = proto::StreamingForeachFunction::default();
420            foreach_func.function =
421                Some(proto::streaming_foreach_function::Function::PythonFunction(
422                    foreach_batch.to_proto(),
423                ));
424            write_op.foreach_batch = Some(foreach_func);
425        }
426
427        if let Some(foreach) = &self.foreach_payload {
428            let mut foreach_func = proto::StreamingForeachFunction::default();
429            foreach_func.function = Some(
430                proto::streaming_foreach_function::Function::PythonFunction(foreach.to_proto()),
431            );
432            write_op.foreach_writer = Some(foreach_func);
433        }
434
435        // Build the plan with the WriteStreamOperationStart command
436        let mut plan = proto::Plan::default();
437        let mut cmd = proto::Command::default();
438        cmd.command_type = Some(proto::command::CommandType::WriteStreamOperationStart(
439            write_op,
440        ));
441        plan.op_type = Some(proto::plan::OpType::Command(cmd));
442
443        // Create ExecutePlanRequest
444        let request = proto::ExecutePlanRequest {
445            session_id: self.session.client().session_id().to_string(),
446            user_context: Some(proto::UserContext::default()),
447            plan: Some(plan),
448            ..Default::default()
449        };
450
451        // Execute the plan and read the WriteStreamOperationStartResult, which carries
452        // the server-assigned query id / run id / name.
453        let mut response_stream = block_on(self.session.client().execute_plan(request))?;
454        let mut query_id = String::new();
455        let mut run_id = String::new();
456        let mut name = self.query_name.clone();
457        while let Some(resp) =
458            block_on(response_stream.message()).map_err(SparkError::from_grpc_status)?
459        {
460            if let Some(
461                proto::execute_plan_response::ResponseType::WriteStreamOperationStartResult(res),
462            ) = resp.response_type
463            {
464                if let Some(qid) = res.query_id {
465                    query_id = qid.id;
466                    run_id = qid.run_id;
467                }
468                if !res.name.is_empty() {
469                    name = Some(res.name);
470                }
471            }
472        }
473        if query_id.is_empty() {
474            return Err(SparkError::connect_msg(
475                "writeStream.start: server returned no WriteStreamOperationStartResult",
476            ));
477        }
478
479        Ok(StreamingQuery {
480            session: self.session,
481            query_id,
482            run_id,
483            name,
484        })
485    }
486}
487
488/// A handle to an active streaming query.
489///
490/// Mirrors `pyspark.sql.connect.streaming.StreamingQuery`.
491#[derive(Clone)]
492pub struct StreamingQuery {
493    session: SparkSession,
494    query_id: String,
495    run_id: String,
496    name: Option<String>,
497}
498
499impl StreamingQuery {
500    /// Get the query ID.
501    pub fn id(&self) -> &str {
502        &self.query_id
503    }
504
505    /// Get the run ID.
506    pub fn run_id(&self) -> &str {
507        &self.run_id
508    }
509
510    /// Get the query name.
511    pub fn name(&self) -> Option<&str> {
512        self.name.as_deref()
513    }
514
515    /// Check if the query is actively running.
516    pub fn is_active(&self) -> Result<bool> {
517        self._fetch_status().map(|status| status.is_active)
518    }
519
520    /// Get the current status of the query.
521    pub fn status(&self) -> Result<StreamingQueryStatus> {
522        self._fetch_status()
523    }
524
525    /// Stop the streaming query.
526    pub fn stop(&self) -> Result<()> {
527        let mut cmd = proto::StreamingQueryCommand::default();
528        cmd.command = Some(proto::streaming_query_command::Command::Stop(true));
529        self._execute_command(cmd)?;
530        Ok(())
531    }
532
533    /// Wait for the query to terminate with optional timeout in seconds.
534    pub fn await_termination(&self, timeout_sec: Option<f64>) -> Result<Option<bool>> {
535        let mut cmd = proto::StreamingQueryCommand::default();
536        let mut await_term = proto::streaming_query_command::AwaitTerminationCommand::default();
537
538        if let Some(timeout) = timeout_sec {
539            if timeout <= 0.0 {
540                return Err(SparkError::value(
541                    "INVALID_TIMEOUT",
542                    &[("value", &timeout.to_string())],
543                ));
544            }
545            await_term.timeout_ms = Some((timeout * 1000.0) as i64);
546        }
547
548        cmd.command = Some(proto::streaming_query_command::Command::AwaitTermination(
549            await_term,
550        ));
551
552        let result = self._execute_command(cmd)?;
553
554        if let Some(proto::streaming_query_command_result::ResultType::AwaitTermination(
555            await_result,
556        )) = result.result_type
557        {
558            if timeout_sec.is_some() {
559                Ok(Some(await_result.terminated))
560            } else {
561                Ok(None)
562            }
563        } else {
564            Ok(None)
565        }
566    }
567
568    /// Get the last streaming progress, if available.
569    pub fn last_progress(&self) -> Result<Option<String>> {
570        let mut cmd = proto::StreamingQueryCommand::default();
571        cmd.command = Some(proto::streaming_query_command::Command::LastProgress(true));
572
573        let result = self._execute_command(cmd)?;
574
575        if let Some(proto::streaming_query_command_result::ResultType::RecentProgress(progress)) =
576            result.result_type
577        {
578            if let Some(progress_result) = progress.recent_progress_json.last() {
579                return Ok(Some(progress_result.clone()));
580            }
581        }
582
583        Ok(None)
584    }
585
586    /// Get recent streaming progress results.
587    pub fn recent_progress(&self) -> Result<Vec<String>> {
588        let mut cmd = proto::StreamingQueryCommand::default();
589        cmd.command = Some(proto::streaming_query_command::Command::RecentProgress(
590            true,
591        ));
592
593        let result = self._execute_command(cmd)?;
594
595        if let Some(proto::streaming_query_command_result::ResultType::RecentProgress(progress)) =
596            result.result_type
597        {
598            Ok(progress.recent_progress_json)
599        } else {
600            Ok(vec![])
601        }
602    }
603
604    /// Process all available data in the streaming query.
605    pub fn process_all_available(&self) -> Result<()> {
606        let mut cmd = proto::StreamingQueryCommand::default();
607        cmd.command = Some(proto::streaming_query_command::Command::ProcessAllAvailable(true));
608        self._execute_command(cmd)?;
609        Ok(())
610    }
611
612    /// Print the execution plan of the streaming query.
613    pub fn explain(&self, extended: bool) -> Result<String> {
614        let mut cmd = proto::StreamingQueryCommand::default();
615        let mut explain = proto::streaming_query_command::ExplainCommand::default();
616        explain.extended = extended;
617        cmd.command = Some(proto::streaming_query_command::Command::Explain(explain));
618
619        let result = self._execute_command(cmd)?;
620
621        if let Some(proto::streaming_query_command_result::ResultType::Explain(explain_result)) =
622            result.result_type
623        {
624            Ok(explain_result.result)
625        } else {
626            Ok(String::new())
627        }
628    }
629
630    /// Get any exception that occurred in the streaming query.
631    pub fn exception(&self) -> Result<Option<StreamingQueryException>> {
632        let mut cmd = proto::StreamingQueryCommand::default();
633        cmd.command = Some(proto::streaming_query_command::Command::Exception(true));
634
635        let result = self._execute_command(cmd)?;
636
637        if let Some(proto::streaming_query_command_result::ResultType::Exception(exc)) =
638            result.result_type
639        {
640            if let Some(msg) = exc.exception_message {
641                if !msg.is_empty() {
642                    return Ok(Some(StreamingQueryException {
643                        message: msg,
644                        error_class: exc.error_class.unwrap_or_default(),
645                    }));
646                }
647            }
648        }
649
650        Ok(None)
651    }
652
653    /// Fetch the current status of the query.
654    fn _fetch_status(&self) -> Result<StreamingQueryStatus> {
655        let mut cmd = proto::StreamingQueryCommand::default();
656        cmd.command = Some(proto::streaming_query_command::Command::Status(true));
657
658        let result = self._execute_command(cmd)?;
659
660        if let Some(proto::streaming_query_command_result::ResultType::Status(status)) =
661            result.result_type
662        {
663            Ok(StreamingQueryStatus {
664                is_active: status.is_active,
665                status_message: status.status_message,
666                is_data_available: status.is_data_available,
667                is_trigger_active: status.is_trigger_active,
668            })
669        } else {
670            Err(SparkError::connect_msg(
671                "Missing status in StreamingQueryCommandResult",
672            ))
673        }
674    }
675
676    /// Execute a streaming query command and return the parsed result.
677    ///
678    /// The server replies on the execute-plan stream with a
679    /// `StreamingQueryCommandResult` in the `response_type` oneof; we drain the
680    /// stream (via the shared collector, so metrics/progress are captured too) and
681    /// return the first such result. Earlier this discarded the stream and returned
682    /// a default, so every status/isActive/explain/exception/progress call saw an
683    /// empty result — status/isActive then failed with "Missing status".
684    fn _execute_command(
685        &self,
686        mut cmd: proto::StreamingQueryCommand,
687    ) -> Result<proto::StreamingQueryCommandResult> {
688        let mut query_id = proto::StreamingQueryInstanceId::default();
689        query_id.id = self.query_id.clone();
690        query_id.run_id = self.run_id.clone();
691        cmd.query_id = Some(query_id);
692
693        let responses = crate::dataframe::execute_command_collect(
694            &self.session,
695            proto::command::CommandType::StreamingQueryCommand(cmd),
696        )?;
697
698        for resp in responses {
699            if let Some(proto::execute_plan_response::ResponseType::StreamingQueryCommandResult(
700                result,
701            )) = resp.response_type
702            {
703                return Ok(result);
704            }
705        }
706
707        Ok(proto::StreamingQueryCommandResult::default())
708    }
709}
710
711/// Status information for a streaming query.
712#[derive(Debug, Clone)]
713pub struct StreamingQueryStatus {
714    pub is_active: bool,
715    pub status_message: String,
716    pub is_data_available: bool,
717    pub is_trigger_active: bool,
718}
719
720/// Exception information for a streaming query.
721#[derive(Debug, Clone)]
722pub struct StreamingQueryException {
723    pub message: String,
724    pub error_class: String,
725}
726
727/// An iterator over streaming query listener events from the server.
728/// Yields events incrementally as they arrive, without buffering the entire stream.
729pub struct ListenerEventStream {
730    stream: ReattachableResponseStream,
731    buffered_events: std::vec::IntoIter<(i32, String)>,
732    done: bool,
733}
734
735impl Iterator for ListenerEventStream {
736    type Item = Result<(i32, String)>;
737
738    fn next(&mut self) -> Option<Self::Item> {
739        // First, yield any buffered events from the last response
740        if let Some(event) = self.buffered_events.next() {
741            return Some(Ok(event));
742        }
743
744        if self.done {
745            return None;
746        }
747
748        loop {
749            match block_on(self.stream.message()) {
750                Ok(Some(resp)) => {
751                    if let Some(
752                        proto::execute_plan_response::ResponseType::StreamingQueryListenerEventsResult(res),
753                    ) = resp.response_type
754                    {
755                        if !res.events.is_empty() {
756                            let mut events = vec![];
757                            for event in res.events {
758                                events.push((event.event_type, event.event_json));
759                            }
760                            self.buffered_events = events.into_iter();
761                            // Yield the first buffered event
762                            if let Some(event) = self.buffered_events.next() {
763                                return Some(Ok(event));
764                            }
765                        }
766                    }
767                    // Keep pulling for events if this response had none
768                }
769                Ok(None) => {
770                    self.done = true;
771                    return None;
772                }
773                Err(e) => {
774                    self.done = true;
775                    return Some(Err(e));
776                }
777            }
778        }
779    }
780}
781
782/// Event-type constants for listener events (mirror pyspark's values).
783pub const QUERY_PROGRESS_EVENT: i32 = 1;
784pub const QUERY_TERMINATED_EVENT: i32 = 2;
785pub const QUERY_IDLE_EVENT: i32 = 3;
786
787/// A streaming-query listener event delivered by the client-side listener bus.
788///
789/// `event_json` is the server-provided JSON for the event; `event_type` is one of
790/// [`QUERY_PROGRESS_EVENT`], [`QUERY_TERMINATED_EVENT`], [`QUERY_IDLE_EVENT`].
791#[derive(Debug, Clone)]
792pub struct StreamingQueryListenerEvent {
793    pub event_type: i32,
794    pub event_json: String,
795}
796
797/// A client-side listener for streaming-query events. Implement this trait (Rust
798/// clients) and register it with [`StreamingQueryManager::add_listener`]; the manager
799/// runs a background bus that streams events from the server and dispatches them here.
800pub trait StreamingQueryListener: Send + Sync {
801    fn on_event(&self, event: &StreamingQueryListenerEvent);
802}
803
804/// Shared state of the client-side listener bus: the registered listeners and the
805/// background dispatch thread (started with the first listener, stopped with the last).
806#[derive(Default)]
807struct ListenerBusState {
808    listeners: Vec<(String, Arc<dyn StreamingQueryListener>)>,
809    thread: Option<std::thread::JoinHandle<()>>,
810}
811
812/// The background dispatch loop: opens the server event stream and forwards each event
813/// to every currently-registered listener. Ends when the stream closes (the server was
814/// asked to stop via `RemoveListenerBusListener`) or all listeners are removed.
815fn run_listener_event_loop(session: SparkSession, bus: Arc<Mutex<ListenerBusState>>) {
816    let stream = match StreamingQueryManager::new(session).listener_event_stream() {
817        Ok(s) => s,
818        Err(_) => return,
819    };
820    for item in stream {
821        let listeners: Vec<Arc<dyn StreamingQueryListener>> = {
822            let st = bus.lock().unwrap();
823            if st.listeners.is_empty() {
824                break;
825            }
826            st.listeners.iter().map(|(_, l)| l.clone()).collect()
827        };
828        match item {
829            Ok((event_type, event_json)) => {
830                let ev = StreamingQueryListenerEvent {
831                    event_type,
832                    event_json,
833                };
834                for l in &listeners {
835                    l.on_event(&ev);
836                }
837            }
838            Err(_) => break,
839        }
840    }
841}
842
843/// Manager for active streaming queries.
844///
845/// Mirrors `pyspark.sql.connect.streaming.StreamingQueryManager`, including a native
846/// client-side listener bus (so Rust clients get the listener feature too).
847pub struct StreamingQueryManager {
848    session: SparkSession,
849    bus: Arc<Mutex<ListenerBusState>>,
850}
851
852impl StreamingQueryManager {
853    /// Create a new StreamingQueryManager.
854    pub(crate) fn new(session: SparkSession) -> Self {
855        StreamingQueryManager {
856            session,
857            bus: Arc::new(Mutex::new(ListenerBusState::default())),
858        }
859    }
860
861    /// Register a client-side listener. Returns an id that can be passed to
862    /// [`StreamingQueryManager::remove_listener`]. Starts the background dispatch
863    /// thread when it is the first listener. Mirrors `StreamingQueryManager.addListener`.
864    pub fn add_listener(&self, listener: Arc<dyn StreamingQueryListener>) -> Result<String> {
865        let id = uuid::Uuid::new_v4().to_string();
866        let mut st = self.bus.lock().unwrap();
867        st.listeners.push((id.clone(), listener));
868        if st.listeners.len() == 1 {
869            let session = self.session.clone();
870            let bus = self.bus.clone();
871            st.thread = Some(std::thread::spawn(move || {
872                run_listener_event_loop(session, bus);
873            }));
874        }
875        Ok(id)
876    }
877
878    /// Remove a client-side listener by id. Stops the background dispatch thread when
879    /// the last listener is removed. Mirrors `StreamingQueryManager.removeListener`.
880    pub fn remove_listener(&self, id: &str) -> Result<()> {
881        let now_empty = {
882            let mut st = self.bus.lock().unwrap();
883            st.listeners.retain(|(lid, _)| lid != id);
884            st.listeners.is_empty()
885        };
886        if now_empty {
887            self.stop_listener_bus();
888        }
889        Ok(())
890    }
891
892    /// Remove all client-side listeners and stop the dispatch thread. Mirrors
893    /// `StreamingQueryManager.close`.
894    pub fn close(&self) -> Result<()> {
895        let had_listeners = {
896            let mut st = self.bus.lock().unwrap();
897            let had = !st.listeners.is_empty();
898            st.listeners.clear();
899            had
900        };
901        if had_listeners {
902            self.stop_listener_bus();
903        }
904        Ok(())
905    }
906
907    /// Ask the server to stop streaming listener events (which ends the background
908    /// thread's stream), then join the thread.
909    fn stop_listener_bus(&self) {
910        let mut lb = proto::StreamingQueryListenerBusCommand::default();
911        lb.command = Some(
912            proto::streaming_query_listener_bus_command::Command::RemoveListenerBusListener(true),
913        );
914        let _ = crate::dataframe::execute_command_collect(
915            &self.session,
916            proto::command::CommandType::StreamingQueryListenerBusCommand(lb),
917        );
918        let handle = self.bus.lock().unwrap().thread.take();
919        if let Some(h) = handle {
920            let _ = h.join();
921        }
922    }
923
924    /// Get all active streaming queries.
925    pub fn active(&self) -> Result<Vec<StreamingQuery>> {
926        let mut cmd = proto::StreamingQueryManagerCommand::default();
927        cmd.command = Some(proto::streaming_query_manager_command::Command::Active(
928            true,
929        ));
930
931        let result = self._execute_manager_command(cmd)?;
932
933        if let Some(proto::streaming_query_manager_command_result::ResultType::Active(active)) =
934            result.result_type
935        {
936            let queries = active
937                .active_queries
938                .into_iter()
939                .map(|q| {
940                    let query_id = q.id.as_ref().map(|id| id.id.clone()).unwrap_or_default();
941                    let run_id =
942                        q.id.as_ref()
943                            .map(|id| id.run_id.clone())
944                            .unwrap_or_default();
945                    let name = q.name;
946
947                    StreamingQuery {
948                        session: self.session.clone(),
949                        query_id,
950                        run_id,
951                        name,
952                    }
953                })
954                .collect();
955            return Ok(queries);
956        }
957
958        Ok(vec![])
959    }
960
961    /// Get a specific streaming query by ID.
962    pub fn get(&self, id: &str) -> Result<Option<StreamingQuery>> {
963        let mut cmd = proto::StreamingQueryManagerCommand::default();
964        cmd.command = Some(proto::streaming_query_manager_command::Command::GetQuery(
965            id.to_string(),
966        ));
967
968        let result = self._execute_manager_command(cmd)?;
969
970        if let Some(proto::streaming_query_manager_command_result::ResultType::Query(query)) =
971            result.result_type
972        {
973            let query_id = query
974                .id
975                .as_ref()
976                .map(|id| id.id.clone())
977                .unwrap_or_default();
978            let run_id = query
979                .id
980                .as_ref()
981                .map(|id| id.run_id.clone())
982                .unwrap_or_default();
983            let name = query.name;
984
985            return Ok(Some(StreamingQuery {
986                session: self.session.clone(),
987                query_id,
988                run_id,
989                name,
990            }));
991        }
992
993        Ok(None)
994    }
995
996    /// Wait for any streaming query to terminate with optional timeout in seconds.
997    pub fn await_any_termination(&self, timeout_sec: Option<f64>) -> Result<Option<bool>> {
998        let mut cmd = proto::StreamingQueryManagerCommand::default();
999        let mut await_term =
1000            proto::streaming_query_manager_command::AwaitAnyTerminationCommand::default();
1001
1002        if let Some(timeout) = timeout_sec {
1003            if timeout <= 0.0 {
1004                return Err(SparkError::value(
1005                    "INVALID_TIMEOUT",
1006                    &[("value", &timeout.to_string())],
1007                ));
1008            }
1009            await_term.timeout_ms = Some((timeout * 1000.0) as i64);
1010        }
1011
1012        cmd.command =
1013            Some(proto::streaming_query_manager_command::Command::AwaitAnyTermination(await_term));
1014
1015        let result = self._execute_manager_command(cmd)?;
1016
1017        if let Some(
1018            proto::streaming_query_manager_command_result::ResultType::AwaitAnyTermination(
1019                await_result,
1020            ),
1021        ) = result.result_type
1022        {
1023            if timeout_sec.is_some() {
1024                return Ok(Some(await_result.terminated));
1025            } else {
1026                return Ok(None);
1027            }
1028        }
1029
1030        Ok(None)
1031    }
1032
1033    /// Reset terminated streaming queries.
1034    pub fn reset_terminated(&self) -> Result<()> {
1035        let mut cmd = proto::StreamingQueryManagerCommand::default();
1036        cmd.command = Some(proto::streaming_query_manager_command::Command::ResetTerminated(true));
1037
1038        self._execute_manager_command(cmd)?;
1039        Ok(())
1040    }
1041
1042    /// Register a server-side listener from a cloudpickled PythonUDF payload (the
1043    /// server runs it in a Python worker). This is distinct from the client-side
1044    /// listener bus ([`add_listener`](Self::add_listener)); it sends the
1045    /// `AddListener` manager command and returns the server listener id.
1046    pub fn register_python_listener(&self, payload: PythonUDFPayload) -> Result<String> {
1047        let listener_id = uuid::Uuid::new_v4().to_string();
1048        let mut cmd = proto::StreamingQueryManagerCommand::default();
1049        let mut listener_cmd =
1050            proto::streaming_query_manager_command::StreamingQueryListenerCommand::default();
1051        listener_cmd.python_listener_payload = Some(payload.to_proto());
1052        listener_cmd.id = listener_id.clone();
1053        cmd.command =
1054            Some(proto::streaming_query_manager_command::Command::AddListener(listener_cmd));
1055
1056        let result = self._execute_manager_command(cmd)?;
1057
1058        if let Some(proto::streaming_query_manager_command_result::ResultType::AddListener(true)) =
1059            result.result_type
1060        {
1061            Ok(listener_id)
1062        } else {
1063            Err(SparkError::connect_msg("Failed to add listener"))
1064        }
1065    }
1066
1067    /// Remove a server-side listener registered via
1068    /// [`register_python_listener`](Self::register_python_listener), by id.
1069    pub fn unregister_python_listener(&self, listener_id: &str) -> Result<()> {
1070        let mut cmd = proto::StreamingQueryManagerCommand::default();
1071        let mut listener_cmd =
1072            proto::streaming_query_manager_command::StreamingQueryListenerCommand::default();
1073        listener_cmd.id = listener_id.to_string();
1074        cmd.command =
1075            Some(proto::streaming_query_manager_command::Command::RemoveListener(listener_cmd));
1076
1077        self._execute_manager_command(cmd)?;
1078        Ok(())
1079    }
1080
1081    /// Stream listener events from the server incrementally (live).
1082    /// Returns a ListenerEventStream that yields events as they arrive.
1083    pub fn listener_event_stream(&self) -> Result<ListenerEventStream> {
1084        let mut cmd = proto::Command::default();
1085        let mut listener_bus_cmd = proto::StreamingQueryListenerBusCommand::default();
1086        // Subscribe to receive events via the oneof command field
1087        listener_bus_cmd.command = Some(
1088            proto::streaming_query_listener_bus_command::Command::AddListenerBusListener(true),
1089        );
1090        cmd.command_type =
1091            Some(proto::command::CommandType::StreamingQueryListenerBusCommand(listener_bus_cmd));
1092
1093        let mut plan = proto::Plan::default();
1094        plan.op_type = Some(proto::plan::OpType::Command(cmd));
1095
1096        let request = proto::ExecutePlanRequest {
1097            session_id: self.session.client().session_id().to_string(),
1098            user_context: Some(proto::UserContext::default()),
1099            plan: Some(plan),
1100            ..Default::default()
1101        };
1102
1103        let response_stream = block_on(self.session.client().execute_plan_reattachable(request))?;
1104
1105        Ok(ListenerEventStream {
1106            stream: response_stream,
1107            buffered_events: vec![].into_iter(),
1108            done: false,
1109        })
1110    }
1111
1112    /// Execute a streaming query manager command and return the parsed result.
1113    ///
1114    /// Like `StreamingQuery::_execute_command`, the server's reply carries a
1115    /// `StreamingQueryManagerCommandResult` in the response stream; drain it and
1116    /// return the first such result (was discarded before, so active/get/etc.
1117    /// always saw an empty result).
1118    fn _execute_manager_command(
1119        &self,
1120        cmd: proto::StreamingQueryManagerCommand,
1121    ) -> Result<proto::StreamingQueryManagerCommandResult> {
1122        let responses = crate::dataframe::execute_command_collect(
1123            &self.session,
1124            proto::command::CommandType::StreamingQueryManagerCommand(cmd),
1125        )?;
1126
1127        for resp in responses {
1128            if let Some(
1129                proto::execute_plan_response::ResponseType::StreamingQueryManagerCommandResult(
1130                    result,
1131                ),
1132            ) = resp.response_type
1133            {
1134                return Ok(result);
1135            }
1136        }
1137
1138        Ok(proto::StreamingQueryManagerCommandResult::default())
1139    }
1140}
1141
1142#[cfg(test)]
1143mod tests {
1144    use super::*;
1145    use crate::session::SparkSession;
1146
1147    fn session() -> SparkSession {
1148        SparkSession::builder()
1149            .remote("sc://localhost:15002")
1150            .get_or_create()
1151            .expect("failed to build session")
1152    }
1153
1154    #[test]
1155    fn stream_reader_format_option() {
1156        let spark = session();
1157        let reader = spark.read_stream();
1158        let reader = reader.format("kafka").option("brokers", "localhost:9092");
1159        assert_eq!(reader.format, Some("kafka".to_string()));
1160        assert_eq!(
1161            reader.options.get("brokers"),
1162            Some(&"localhost:9092".to_string())
1163        );
1164    }
1165
1166    #[test]
1167    fn stream_reader_schema() {
1168        let spark = session();
1169        let reader = spark.read_stream();
1170        let reader = reader.schema("id INT, name STRING".to_string());
1171        assert_eq!(reader.schema, "id INT, name STRING");
1172    }
1173
1174    #[test]
1175    fn stream_reader_source_name() {
1176        let spark = session();
1177        let reader = spark.read_stream();
1178        let reader = reader.name("my_source");
1179        assert_eq!(reader.source_name, Some("my_source".to_string()));
1180    }
1181
1182    #[test]
1183    fn stream_reader_load_creates_streaming_dataframe() {
1184        let spark = session();
1185        let reader = spark.read_stream().format("kafka");
1186        let df = reader.load(Some("/path/to/data"));
1187        assert!(matches!(
1188            &df.plan,
1189            crate::plan::LogicalPlan::Read {
1190                is_streaming: true,
1191                ..
1192            }
1193        ));
1194    }
1195
1196    #[test]
1197    fn stream_reader_json() {
1198        let spark = session();
1199        let reader = spark.read_stream();
1200        let df = reader.json("/path/to/json");
1201        // Verify that the plan has the streaming flag set
1202        match &df.plan {
1203            crate::plan::LogicalPlan::Read { is_streaming, .. } => {
1204                assert!(*is_streaming);
1205            }
1206            _ => panic!("expected Read plan with streaming"),
1207        }
1208    }
1209
1210    #[test]
1211    fn stream_writer_format_output_mode() {
1212        let spark = session();
1213        let df = spark.read_stream().format("kafka").load(None);
1214        let writer = df.write_stream();
1215        let writer = writer.format("parquet").output_mode("append");
1216        assert_eq!(writer.format, Some("parquet".to_string()));
1217        assert_eq!(writer.output_mode, Some("append".to_string()));
1218    }
1219
1220    #[test]
1221    fn stream_writer_partition_by() {
1222        let spark = session();
1223        let df = spark.read_stream().format("kafka").load(None);
1224        let writer = df.write_stream();
1225        let writer = writer.partition_by(vec!["date", "region"]);
1226        assert_eq!(writer.partitioning_columns, vec!["date", "region"]);
1227    }
1228
1229    #[test]
1230    fn stream_writer_cluster_by() {
1231        let spark = session();
1232        let df = spark.read_stream().format("kafka").load(None);
1233        let writer = df.write_stream();
1234        let writer = writer.cluster_by(vec!["user_id", "session_id"]);
1235        assert_eq!(writer.clustering_columns, vec!["user_id", "session_id"]);
1236    }
1237
1238    #[test]
1239    fn stream_writer_query_name() {
1240        let spark = session();
1241        let df = spark.read_stream().format("kafka").load(None);
1242        let writer = df.write_stream();
1243        let writer = writer.query_name("my_query");
1244        assert_eq!(writer.query_name, Some("my_query".to_string()));
1245    }
1246
1247    #[test]
1248    fn stream_writer_trigger_processing_time() {
1249        let spark = session();
1250        let df = spark.read_stream().format("kafka").load(None);
1251        let writer = df.write_stream();
1252        let trigger = Trigger::ProcessingTime("10 seconds".to_string());
1253        let writer = writer.trigger(trigger);
1254        assert!(writer.trigger.is_some());
1255    }
1256
1257    #[test]
1258    fn stream_writer_trigger_once() {
1259        let spark = session();
1260        let df = spark.read_stream().format("kafka").load(None);
1261        let writer = df.write_stream();
1262        let trigger = Trigger::Once;
1263        let writer = writer.trigger(trigger);
1264        assert!(writer.trigger.is_some());
1265    }
1266
1267    #[test]
1268    fn stream_writer_trigger_available_now() {
1269        let spark = session();
1270        let df = spark.read_stream().format("kafka").load(None);
1271        let writer = df.write_stream();
1272        let trigger = Trigger::AvailableNow;
1273        let writer = writer.trigger(trigger);
1274        assert!(writer.trigger.is_some());
1275    }
1276
1277    #[test]
1278    fn stream_writer_trigger_continuous() {
1279        let spark = session();
1280        let df = spark.read_stream().format("kafka").load(None);
1281        let writer = df.write_stream();
1282        let trigger = Trigger::Continuous("1 minute".to_string());
1283        let writer = writer.trigger(trigger);
1284        assert!(writer.trigger.is_some());
1285    }
1286
1287    #[test]
1288    fn stream_writer_option_options() {
1289        let spark = session();
1290        let df = spark.read_stream().format("kafka").load(None);
1291        let writer = df.write_stream();
1292        let writer = writer.option("key1", "val1");
1293        assert_eq!(writer.options.get("key1"), Some(&"val1".to_string()));
1294
1295        let mut opts = std::collections::HashMap::new();
1296        opts.insert("key2".to_string(), "val2".to_string());
1297        let writer = writer.options(opts);
1298        assert_eq!(writer.options.get("key2"), Some(&"val2".to_string()));
1299    }
1300}
spark_connect/streaming.rs

spark_connect/
streaming.rs
