Skip to main content

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;