laminar_sql/datafusion/execute/mod.rs
1//! Streaming SQL execution via DataFusion.
2
3use datafusion::execution::SendableRecordBatchStream;
4use datafusion::prelude::SessionContext;
5
6use crate::parser::parse_streaming_sql;
7use crate::planner::{QueryPlan, StreamingPlan, StreamingPlanner};
8use crate::Error;
9
10/// Result of executing a streaming SQL statement.
11#[derive(Debug)]
12pub enum StreamingSqlResult {
13 /// DDL statement result (CREATE SOURCE, CREATE SINK)
14 Ddl(DdlResult),
15 /// Query execution result with optional streaming metadata
16 Query(QueryResult),
17}
18
19/// Result of a DDL statement execution.
20#[derive(Debug)]
21pub struct DdlResult {
22 /// The streaming plan describing what was created or registered
23 pub plan: StreamingPlan,
24}
25
26/// Result of a query execution.
27///
28/// Contains both the `DataFusion` record batch stream and optional
29/// streaming metadata (window config, join config, emit clause) from
30/// the `QueryPlan`. Ring 0 operators use the `query_plan` to configure
31/// windowing and join behavior.
32pub struct QueryResult {
33 /// Record batch stream from `DataFusion` execution
34 pub stream: SendableRecordBatchStream,
35 /// Streaming query metadata (window config, join config, etc.)
36 ///
37 /// `None` for standard SQL pass-through queries.
38 /// `Some` for queries with streaming features (windows, joins).
39 pub query_plan: Option<QueryPlan>,
40}
41
42impl std::fmt::Debug for QueryResult {
43 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44 f.debug_struct("QueryResult")
45 .field("query_plan", &self.query_plan)
46 .field("stream", &"<SendableRecordBatchStream>")
47 .finish()
48 }
49}
50
51/// Executes a streaming SQL statement end-to-end.
52///
53/// This function performs the full pipeline:
54/// 1. Parse SQL with streaming extensions (CREATE SOURCE/SINK, windows, etc.)
55/// 2. Plan via [`StreamingPlanner`]
56/// 3. For DDL: return the streaming plan as [`DdlResult`]
57/// 4. For queries with streaming features: create `LogicalPlan` via
58/// `DataFusion`, execute, and return stream + [`QueryPlan`] metadata
59/// 5. For standard SQL: pass through to `DataFusion` directly
60///
61/// # Arguments
62///
63/// * `sql` - The SQL statement to execute
64/// * `ctx` - `DataFusion` session context (should have streaming functions registered)
65/// * `planner` - Streaming planner with registered sources/sinks
66///
67/// # Errors
68///
69/// Returns [`Error`] if parsing, planning, or execution fails.
70pub async fn execute_streaming_sql(
71 sql: &str,
72 ctx: &SessionContext,
73 planner: &mut StreamingPlanner,
74) -> std::result::Result<StreamingSqlResult, Error> {
75 let statements = parse_streaming_sql(sql)?;
76
77 if statements.is_empty() {
78 return Err(Error::ParseError(
79 crate::parser::ParseError::StreamingError("Empty SQL statement".to_string()),
80 ));
81 }
82
83 // Process the first statement
84 let statement = &statements[0];
85 let plan = planner.plan(statement)?;
86
87 match plan {
88 StreamingPlan::Query(query_plan) => {
89 let logical_plan = planner.to_logical_plan(&query_plan, ctx).await?;
90 let df = ctx.execute_logical_plan(logical_plan).await?;
91 let stream = df.execute_stream().await?;
92
93 Ok(StreamingSqlResult::Query(QueryResult {
94 stream,
95 query_plan: Some(query_plan),
96 }))
97 }
98 StreamingPlan::Standard(stmt) => {
99 let sql_str = stmt.to_string();
100 let df = ctx.sql(&sql_str).await?;
101 let stream = df.execute_stream().await?;
102
103 Ok(StreamingSqlResult::Query(QueryResult {
104 stream,
105 query_plan: None,
106 }))
107 }
108 StreamingPlan::RegisterSource(_)
109 | StreamingPlan::RegisterSink(_)
110 | StreamingPlan::RegisterLookupTable(_)
111 | StreamingPlan::DropLookupTable { .. } => Ok(StreamingSqlResult::Ddl(DdlResult { plan })),
112 }
113}
114
115#[cfg(test)]
116mod tests;