laminar_db/
handle.rs

1//! Handle types for query results, source access, and subscriptions.
2
3use std::marker::PhantomData;
4use std::sync::Arc;
5use std::time::Duration;
6
7use arrow::array::RecordBatch;
8use arrow::datatypes::SchemaRef;
9
10use laminar_core::streaming::{Record, Subscription};
11
12use crate::catalog::{ArrowRecord, SourceEntry};
13use crate::DbError;
14
15/// Result of executing a SQL statement.
16#[derive(Debug)]
17pub enum ExecuteResult {
18    /// DDL statement completed (CREATE, DROP, ALTER).
19    Ddl(DdlInfo),
20    /// Query is running, subscribe to results.
21    Query(QueryHandle),
22    /// Rows were affected (INSERT INTO).
23    RowsAffected(u64),
24    /// Metadata result (SHOW, DESCRIBE).
25    Metadata(RecordBatch),
26}
27
28impl ExecuteResult {
29    /// Convert to `QueryHandle`, or error if this is not a query result.
30    ///
31    /// # Errors
32    /// Returns `DbError::InvalidOperation` if result is not a query.
33    pub fn into_query(self) -> Result<QueryHandle, DbError> {
34        match self {
35            Self::Query(q) => Ok(q),
36            _ => Err(DbError::InvalidOperation(
37                "Expected a query result".to_string(),
38            )),
39        }
40    }
41}
42
43/// Information about a completed DDL statement.
44#[derive(Debug, Clone)]
45pub struct DdlInfo {
46    /// e.g. `"CREATE SOURCE"`.
47    pub statement_type: String,
48    /// Object affected.
49    pub object_name: String,
50    /// Whether this statement changed catalog/runtime state. `IF NOT EXISTS`
51    /// success uses `false` so the attempted definition cannot replace durable DDL.
52    pub(crate) applied: bool,
53}
54
55/// Handle to a running streaming query.
56#[derive(Debug)]
57pub struct QueryHandle {
58    pub(crate) id: u64,
59    pub(crate) schema: SchemaRef,
60    pub(crate) sql: String,
61    pub(crate) subscription: Option<Subscription<ArrowRecord>>,
62    pub(crate) active: bool,
63    pub(crate) cancel_token: tokio_util::sync::CancellationToken,
64}
65
66impl QueryHandle {
67    /// Output schema.
68    #[must_use]
69    pub fn schema(&self) -> &SchemaRef {
70        &self.schema
71    }
72
73    /// SQL text.
74    #[must_use]
75    pub fn sql(&self) -> &str {
76        &self.sql
77    }
78
79    /// Query ID.
80    #[must_use]
81    pub fn id(&self) -> u64 {
82        self.id
83    }
84
85    /// Whether the query is still active.
86    #[must_use]
87    pub fn is_active(&self) -> bool {
88        self.active
89    }
90
91    /// Take the raw `RecordBatch` subscription for this query.
92    ///
93    /// The subscription is consumed on first call; subsequent calls error.
94    ///
95    /// # Errors
96    /// Returns `DbError::InvalidOperation` if the subscription was already consumed.
97    pub fn subscribe_raw(&mut self) -> Result<Subscription<ArrowRecord>, DbError> {
98        self.subscription
99            .take()
100            .ok_or_else(|| DbError::InvalidOperation("Subscription already consumed".to_string()))
101    }
102
103    /// Cancel the query and drop its subscription.
104    pub fn cancel(&mut self) {
105        self.active = false;
106        self.subscription = None;
107        self.cancel_token.cancel();
108    }
109}
110
111impl Drop for QueryHandle {
112    fn drop(&mut self) {
113        self.cancel_token.cancel();
114    }
115}
116
117/// Deserialize rows from a `RecordBatch`. Auto-generated by `#[derive(FromRecordBatch)]`.
118pub trait FromBatch: Sized {
119    /// Single row.
120    fn from_batch(batch: &RecordBatch, row: usize) -> Self;
121    /// All rows.
122    fn from_batch_all(batch: &RecordBatch) -> Vec<Self>;
123}
124
125/// A typed frame from a named subscription.
126#[derive(Debug)]
127pub enum TypedSubscriptionFrame<T> {
128    /// Rows emitted by one processing cycle.
129    Rows {
130        /// Deserialized rows in the shared-log entry.
131        rows: Vec<T>,
132        /// Portal-local delivery sequence; not a durable or cluster-wide resume position.
133        sequence: u64,
134    },
135    /// Durable progress frontier for this checkpoint.
136    Barrier {
137        /// Portal-local delivery sequence; not a cluster-wide ordering position.
138        sequence: u64,
139        /// Engine checkpoint epoch.
140        epoch: u64,
141        /// Engine checkpoint identifier.
142        checkpoint_id: u64,
143        /// Local-log cut, or a gateway-local compatibility cut in cluster mode.
144        through_sequence: u64,
145    },
146}
147
148/// Typed compatibility frame plus optional durable cluster partition metadata.
149#[derive(Debug)]
150pub struct TypedSubscriptionEnvelope<T> {
151    /// Existing typed rows or checkpoint progress.
152    pub frame: TypedSubscriptionFrame<T>,
153    /// Present for committed cluster data and progress frames.
154    pub cluster: Option<crate::subscription::ClusterSubscriptionFrameMetadata>,
155}
156
157/// Terminal failure from a named subscription.
158#[derive(Debug, thiserror::Error, PartialEq, Eq)]
159pub enum SubscriptionError {
160    /// The consumer fell behind the bounded shared log.
161    #[error("subscription fell behind by {skipped} entries")]
162    Lagged {
163        /// Number of entries no longer available.
164        skipped: u64,
165    },
166    /// The subscribed object or its delivery path failed.
167    #[error("subscription failed: {message}")]
168    Failed {
169        /// Failure detail.
170        message: String,
171    },
172}
173
174/// Typed subscription that deserializes named-stream `RecordBatch` rows.
175pub struct TypedSubscription<T: FromBatch> {
176    portal: crate::subscription::SubscriptionPortal,
177    _phantom: PhantomData<T>,
178}
179
180impl<T: FromBatch> TypedSubscription<T> {
181    pub(crate) fn new(portal: crate::subscription::SubscriptionPortal) -> Self {
182        Self {
183            portal,
184            _phantom: PhantomData,
185        }
186    }
187
188    /// Schema resolved while the portal was attached under the topology lock.
189    #[must_use]
190    pub fn schema(&self) -> SchemaRef {
191        self.portal.schema()
192    }
193
194    /// Wait for the next data or checkpoint frame.
195    ///
196    /// # Errors
197    /// Returns a terminal error if delivery has a gap or the object fails.
198    pub async fn next_frame(
199        &mut self,
200    ) -> Result<Option<TypedSubscriptionFrame<T>>, SubscriptionError> {
201        Ok(self.next_envelope().await?.map(|envelope| envelope.frame))
202    }
203
204    /// Wait for the next frame with optional durable cluster partition metadata.
205    ///
206    /// # Errors
207    /// Returns a terminal error if delivery has a gap or the object fails.
208    pub async fn next_envelope(
209        &mut self,
210    ) -> Result<Option<TypedSubscriptionEnvelope<T>>, SubscriptionError> {
211        let Some(envelope) = self.portal.next_envelope().await else {
212            return Ok(None);
213        };
214        Ok(Some(TypedSubscriptionEnvelope {
215            frame: convert_typed_frame(envelope.frame)?,
216            cluster: envelope.cluster,
217        }))
218    }
219
220    /// Return the next immediately available data or checkpoint frame.
221    ///
222    /// # Errors
223    /// Returns a terminal error if delivery has a gap or the object fails.
224    pub fn try_next_frame(
225        &mut self,
226    ) -> Result<Option<TypedSubscriptionFrame<T>>, SubscriptionError> {
227        Ok(self.try_next_envelope()?.map(|envelope| envelope.frame))
228    }
229
230    /// Return the next immediately available frame and optional cluster metadata.
231    ///
232    /// # Errors
233    /// Returns a terminal error if delivery has a gap or the object fails.
234    pub fn try_next_envelope(
235        &mut self,
236    ) -> Result<Option<TypedSubscriptionEnvelope<T>>, SubscriptionError> {
237        let Some(envelope) = self.portal.try_next_envelope() else {
238            return Ok(None);
239        };
240        Ok(Some(TypedSubscriptionEnvelope {
241            frame: convert_typed_frame(envelope.frame)?,
242            cluster: envelope.cluster,
243        }))
244    }
245
246    /// Return the next immediately available row batch, skipping checkpoint markers.
247    ///
248    /// # Errors
249    /// Returns a terminal error if delivery has a gap or the object fails.
250    pub fn poll(&mut self) -> Result<Option<Vec<T>>, SubscriptionError> {
251        loop {
252            match self.try_next_frame()? {
253                Some(TypedSubscriptionFrame::Rows { rows, .. }) => return Ok(Some(rows)),
254                Some(TypedSubscriptionFrame::Barrier { .. }) => {}
255                None => return Ok(None),
256            }
257        }
258    }
259}
260
261fn convert_typed_frame<T: FromBatch>(
262    frame: crate::subscription::PortalFrame,
263) -> Result<TypedSubscriptionFrame<T>, SubscriptionError> {
264    match frame {
265        crate::subscription::PortalFrame::Batch {
266            batch,
267            sequence,
268            lease: _lease,
269        } => Ok(TypedSubscriptionFrame::Rows {
270            rows: T::from_batch_all(&batch),
271            sequence,
272        }),
273        crate::subscription::PortalFrame::Barrier {
274            sequence,
275            epoch,
276            checkpoint_id,
277            through_sequence,
278        } => Ok(TypedSubscriptionFrame::Barrier {
279            sequence,
280            epoch,
281            checkpoint_id,
282            through_sequence,
283        }),
284        crate::subscription::PortalFrame::Lagged(skipped) => {
285            Err(SubscriptionError::Lagged { skipped })
286        }
287        crate::subscription::PortalFrame::Error { message } => {
288            Err(SubscriptionError::Failed { message })
289        }
290    }
291}
292
293impl<T: FromBatch> std::fmt::Debug for TypedSubscription<T> {
294    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
295        f.debug_struct("TypedSubscription")
296            .field("portal", &self.portal)
297            .finish()
298    }
299}
300
301#[cfg(test)]
302mod typed_subscription_tests {
303    use super::*;
304
305    #[derive(Debug)]
306    struct Row;
307
308    impl FromBatch for Row {
309        fn from_batch(_batch: &RecordBatch, _row: usize) -> Self {
310            Self
311        }
312
313        fn from_batch_all(_batch: &RecordBatch) -> Vec<Self> {
314            Vec::new()
315        }
316    }
317
318    #[test]
319    fn delivery_gap_is_not_converted_to_empty_poll() {
320        let error =
321            convert_typed_frame::<Row>(crate::subscription::PortalFrame::Lagged(7)).unwrap_err();
322        assert_eq!(error, SubscriptionError::Lagged { skipped: 7 });
323    }
324
325    #[test]
326    fn delivery_failure_preserves_detail() {
327        let error = convert_typed_frame::<Row>(crate::subscription::PortalFrame::Error {
328            message: "object dropped".into(),
329        })
330        .unwrap_err();
331        assert_eq!(
332            error,
333            SubscriptionError::Failed {
334                message: "object dropped".into()
335            }
336        );
337    }
338}
339
340/// Typed handle for pushing data into a registered source.
341pub struct SourceHandle<T: Record> {
342    entry: Arc<SourceEntry>,
343    _phantom: PhantomData<T>,
344}
345
346impl<T: Record> SourceHandle<T> {
347    /// Validates that `T`'s schema matches the source schema.
348    pub(crate) fn new(entry: Arc<SourceEntry>) -> Result<Self, DbError> {
349        let rust_schema = T::schema();
350        let sql_schema = &entry.schema;
351
352        // Validate field count matches
353        if rust_schema.fields().len() != sql_schema.fields().len() {
354            return Err(DbError::SchemaMismatch(format!(
355                "Rust type has {} fields but source '{}' has {} columns",
356                rust_schema.fields().len(),
357                entry.name,
358                sql_schema.fields().len()
359            )));
360        }
361
362        Ok(Self {
363            entry,
364            _phantom: PhantomData,
365        })
366    }
367
368    /// Push a single record.
369    ///
370    /// # Errors
371    /// Returns `StreamingError` if the pipeline is not running.
372    #[allow(clippy::needless_pass_by_value)]
373    pub fn push(&self, record: T) -> Result<(), laminar_core::streaming::StreamingError> {
374        let batch = record.to_record_batch();
375        self.entry.push_and_buffer(batch)
376    }
377
378    /// Push multiple records, returns count successfully sent.
379    pub fn push_batch(&self, records: impl IntoIterator<Item = T>) -> usize {
380        const BATCH_SIZE: usize = 1024;
381        let mut count = 0;
382        let mut buffer = Vec::with_capacity(BATCH_SIZE);
383
384        for record in records {
385            buffer.push(record);
386            if buffer.len() >= BATCH_SIZE {
387                let batch = T::to_record_batch_from_iter(buffer.drain(..));
388                if self.push_arrow(batch).is_err() {
389                    return count;
390                }
391                count += BATCH_SIZE;
392            }
393        }
394
395        if !buffer.is_empty() {
396            let len = buffer.len();
397            let batch = T::to_record_batch_from_iter(buffer);
398            if self.push_arrow(batch).is_ok() {
399                count += len;
400            }
401        }
402        count
403    }
404
405    /// Push a raw `RecordBatch` (sent to pipeline and buffered for snapshots).
406    ///
407    /// # Errors
408    /// Returns `StreamingError` if the pipeline is not running.
409    pub fn push_arrow(
410        &self,
411        batch: RecordBatch,
412    ) -> Result<(), laminar_core::streaming::StreamingError> {
413        self.entry.push_and_buffer(batch)
414    }
415
416    /// Emit a watermark.
417    pub fn watermark(&self, timestamp: i64) {
418        self.entry.source.watermark(timestamp);
419    }
420
421    /// Current watermark.
422    #[must_use]
423    pub fn current_watermark(&self) -> i64 {
424        self.entry.source.current_watermark()
425    }
426
427    /// Buffered record count.
428    #[must_use]
429    pub fn pending(&self) -> usize {
430        self.entry.source.pending()
431    }
432
433    /// Buffer capacity.
434    #[must_use]
435    pub fn capacity(&self) -> usize {
436        self.entry.source.capacity()
437    }
438
439    /// True when buffer is >80% full.
440    #[must_use]
441    pub fn is_backpressured(&self) -> bool {
442        crate::metrics::is_backpressured(self.pending(), self.capacity())
443    }
444
445    /// Source name.
446    #[must_use]
447    pub fn name(&self) -> &str {
448        &self.entry.name
449    }
450
451    /// Schema.
452    #[must_use]
453    pub fn schema(&self) -> &SchemaRef {
454        &self.entry.schema
455    }
456
457    /// Max out-of-orderness, if configured.
458    #[must_use]
459    pub fn max_out_of_orderness(&self) -> Option<Duration> {
460        self.entry.max_out_of_orderness
461    }
462
463    /// Set the event-time column for watermark-based late-row filtering.
464    pub fn set_event_time_column(&self, column: &str) {
465        self.entry.source.set_event_time_column(column);
466    }
467}
468
469impl<T: Record> std::fmt::Debug for SourceHandle<T> {
470    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
471        f.debug_struct("SourceHandle")
472            .field("name", &self.entry.name)
473            .field("pending", &self.pending())
474            .finish()
475    }
476}
477
478/// Untyped handle for pushing raw `RecordBatch` data.
479pub struct UntypedSourceHandle {
480    entry: Arc<SourceEntry>,
481}
482
483impl UntypedSourceHandle {
484    pub(crate) fn new(entry: Arc<SourceEntry>) -> Self {
485        Self { entry }
486    }
487
488    /// Push a raw `RecordBatch` (sent to pipeline and buffered for snapshots).
489    ///
490    /// # Errors
491    /// Returns `StreamingError` if the pipeline is not running.
492    pub fn push_arrow(
493        &self,
494        batch: RecordBatch,
495    ) -> Result<(), laminar_core::streaming::StreamingError> {
496        self.entry.push_and_buffer(batch)
497    }
498
499    /// Emit a watermark.
500    pub fn watermark(&self, timestamp: i64) {
501        self.entry.source.watermark(timestamp);
502    }
503
504    /// Current watermark.
505    #[must_use]
506    pub fn current_watermark(&self) -> i64 {
507        self.entry.source.current_watermark()
508    }
509
510    /// Buffered record count.
511    #[must_use]
512    pub fn pending(&self) -> usize {
513        self.entry.source.pending()
514    }
515
516    /// Buffer capacity.
517    #[must_use]
518    pub fn capacity(&self) -> usize {
519        self.entry.source.capacity()
520    }
521
522    /// True when buffer is >80% full.
523    #[must_use]
524    pub fn is_backpressured(&self) -> bool {
525        crate::metrics::is_backpressured(self.pending(), self.capacity())
526    }
527
528    /// Source name.
529    #[must_use]
530    pub fn name(&self) -> &str {
531        &self.entry.name
532    }
533
534    /// Schema.
535    #[must_use]
536    pub fn schema(&self) -> &SchemaRef {
537        &self.entry.schema
538    }
539
540    /// Max out-of-orderness, if configured.
541    #[must_use]
542    pub fn max_out_of_orderness(&self) -> Option<Duration> {
543        self.entry.max_out_of_orderness
544    }
545
546    /// Set the event-time column for watermark-based late-row filtering.
547    pub fn set_event_time_column(&self, column: &str) {
548        self.entry.source.set_event_time_column(column);
549    }
550}
551
552impl std::fmt::Debug for UntypedSourceHandle {
553    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
554        f.debug_struct("UntypedSourceHandle")
555            .field("name", &self.entry.name)
556            .finish()
557    }
558}
559
560/// Type of a node in the pipeline topology.
561#[derive(Debug, Clone, Copy, PartialEq, Eq)]
562pub enum PipelineNodeType {
563    /// A data source (CREATE SOURCE).
564    Source,
565    /// A continuous stream (CREATE STREAM).
566    Stream,
567    /// A data sink (CREATE SINK).
568    Sink,
569}
570
571/// A node in the pipeline topology graph.
572#[derive(Debug, Clone)]
573pub struct PipelineNode {
574    /// Node name.
575    pub name: String,
576    /// Node type.
577    pub node_type: PipelineNodeType,
578    /// Arrow schema (sources only).
579    pub schema: Option<SchemaRef>,
580    /// SQL definition (streams only).
581    pub sql: Option<String>,
582}
583
584/// A directed edge in the pipeline topology graph.
585#[derive(Debug, Clone)]
586pub struct PipelineEdge {
587    /// Source node.
588    pub from: String,
589    /// Target node.
590    pub to: String,
591}
592
593/// The complete pipeline topology: nodes and edges.
594#[derive(Debug, Clone)]
595pub struct PipelineTopology {
596    /// Nodes.
597    pub nodes: Vec<PipelineNode>,
598    /// Edges.
599    pub edges: Vec<PipelineEdge>,
600}
601
602/// Metadata about a registered stream.
603#[derive(Debug, Clone)]
604pub struct StreamInfo {
605    /// Name.
606    pub name: String,
607    /// Defining SQL query.
608    pub sql: Option<String>,
609}
610
611/// Information about a registered source.
612#[derive(Debug, Clone)]
613pub struct SourceInfo {
614    /// Name.
615    pub name: String,
616    /// Schema.
617    pub schema: SchemaRef,
618    /// Watermark column, if configured.
619    pub watermark_column: Option<String>,
620}
621
622/// Information about a registered sink.
623#[derive(Debug, Clone)]
624pub struct SinkInfo {
625    /// Name.
626    pub name: String,
627}
628
629/// JSON-serializable projection of a `MaterializedView` for the control-plane
630/// HTTP API.
631#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
632pub struct MaterializedViewInfo {
633    /// View name.
634    pub name: String,
635    /// SQL definition.
636    pub sql: String,
637    /// Current execution state (e.g. `"Running"`, `"Dropping"`).
638    pub state: String,
639}
640
641impl From<&laminar_core::mv::MaterializedView> for MaterializedViewInfo {
642    fn from(view: &laminar_core::mv::MaterializedView) -> Self {
643        Self {
644            name: view.name.clone(),
645            sql: view.sql.clone(),
646            state: format!("{:?}", view.state),
647        }
648    }
649}
650
651/// Information about a running query.
652#[derive(Debug, Clone)]
653pub struct QueryInfo {
654    /// ID.
655    pub id: u64,
656    /// SQL text.
657    pub sql: String,
658    /// Active.
659    pub active: bool,
660}
laminar_db/handle.rs

laminar_db/
handle.rs
