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

laminar_db/
handle.rs
