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        /// Physical sequence within this in-memory object incarnation; not a durable resume token.
133        sequence: u64,
134    },
135    /// Durable progress frontier for this checkpoint.
136    Barrier {
137        /// Physical sequence of this progress entry within the current object incarnation.
138        sequence: u64,
139        /// Engine checkpoint epoch.
140        epoch: u64,
141        /// Engine checkpoint identifier.
142        checkpoint_id: u64,
143        /// Every logical entry below this sequence is covered by the cut.
144        through_sequence: u64,
145    },
146}
147
148/// Terminal failure from a named subscription.
149#[derive(Debug, thiserror::Error, PartialEq, Eq)]
150pub enum SubscriptionError {
151    /// The consumer fell behind the bounded shared log.
152    #[error("subscription fell behind by {skipped} entries")]
153    Lagged {
154        /// Number of entries no longer available.
155        skipped: u64,
156    },
157    /// The subscribed object or its delivery path failed.
158    #[error("subscription failed: {message}")]
159    Failed {
160        /// Failure detail.
161        message: String,
162    },
163}
164
165/// Typed subscription that deserializes named-stream `RecordBatch` rows.
166pub struct TypedSubscription<T: FromBatch> {
167    portal: crate::subscription::SubscriptionPortal,
168    _phantom: PhantomData<T>,
169}
170
171impl<T: FromBatch> TypedSubscription<T> {
172    pub(crate) fn new(portal: crate::subscription::SubscriptionPortal) -> Self {
173        Self {
174            portal,
175            _phantom: PhantomData,
176        }
177    }
178
179    /// Schema resolved while the portal was attached under the topology lock.
180    #[must_use]
181    pub fn schema(&self) -> SchemaRef {
182        self.portal.schema()
183    }
184
185    /// Wait for the next data or checkpoint frame.
186    ///
187    /// # Errors
188    /// Returns a terminal error if delivery has a gap or the object fails.
189    pub async fn next_frame(
190        &mut self,
191    ) -> Result<Option<TypedSubscriptionFrame<T>>, SubscriptionError> {
192        self.portal
193            .next_frame()
194            .await
195            .map(convert_typed_frame)
196            .transpose()
197    }
198
199    /// Return the next immediately available data or checkpoint frame.
200    ///
201    /// # Errors
202    /// Returns a terminal error if delivery has a gap or the object fails.
203    pub fn try_next_frame(
204        &mut self,
205    ) -> Result<Option<TypedSubscriptionFrame<T>>, SubscriptionError> {
206        self.portal
207            .try_next_frame()
208            .map(convert_typed_frame)
209            .transpose()
210    }
211
212    /// Return the next immediately available row batch, skipping checkpoint markers.
213    ///
214    /// # Errors
215    /// Returns a terminal error if delivery has a gap or the object fails.
216    pub fn poll(&mut self) -> Result<Option<Vec<T>>, SubscriptionError> {
217        loop {
218            match self.try_next_frame()? {
219                Some(TypedSubscriptionFrame::Rows { rows, .. }) => return Ok(Some(rows)),
220                Some(TypedSubscriptionFrame::Barrier { .. }) => {}
221                None => return Ok(None),
222            }
223        }
224    }
225}
226
227fn convert_typed_frame<T: FromBatch>(
228    frame: crate::subscription::PortalFrame,
229) -> Result<TypedSubscriptionFrame<T>, SubscriptionError> {
230    match frame {
231        crate::subscription::PortalFrame::Batch {
232            batch,
233            sequence,
234            lease: _lease,
235        } => Ok(TypedSubscriptionFrame::Rows {
236            rows: T::from_batch_all(&batch),
237            sequence,
238        }),
239        crate::subscription::PortalFrame::Barrier {
240            sequence,
241            epoch,
242            checkpoint_id,
243            through_sequence,
244        } => Ok(TypedSubscriptionFrame::Barrier {
245            sequence,
246            epoch,
247            checkpoint_id,
248            through_sequence,
249        }),
250        crate::subscription::PortalFrame::Lagged(skipped) => {
251            Err(SubscriptionError::Lagged { skipped })
252        }
253        crate::subscription::PortalFrame::Error { message } => {
254            Err(SubscriptionError::Failed { message })
255        }
256    }
257}
258
259impl<T: FromBatch> std::fmt::Debug for TypedSubscription<T> {
260    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
261        f.debug_struct("TypedSubscription")
262            .field("portal", &self.portal)
263            .finish()
264    }
265}
266
267#[cfg(test)]
268mod typed_subscription_tests {
269    use super::*;
270
271    #[derive(Debug)]
272    struct Row;
273
274    impl FromBatch for Row {
275        fn from_batch(_batch: &RecordBatch, _row: usize) -> Self {
276            Self
277        }
278
279        fn from_batch_all(_batch: &RecordBatch) -> Vec<Self> {
280            Vec::new()
281        }
282    }
283
284    #[test]
285    fn delivery_gap_is_not_converted_to_empty_poll() {
286        let error =
287            convert_typed_frame::<Row>(crate::subscription::PortalFrame::Lagged(7)).unwrap_err();
288        assert_eq!(error, SubscriptionError::Lagged { skipped: 7 });
289    }
290
291    #[test]
292    fn delivery_failure_preserves_detail() {
293        let error = convert_typed_frame::<Row>(crate::subscription::PortalFrame::Error {
294            message: "object dropped".into(),
295        })
296        .unwrap_err();
297        assert_eq!(
298            error,
299            SubscriptionError::Failed {
300                message: "object dropped".into()
301            }
302        );
303    }
304}
305
306/// Typed handle for pushing data into a registered source.
307pub struct SourceHandle<T: Record> {
308    entry: Arc<SourceEntry>,
309    _phantom: PhantomData<T>,
310}
311
312impl<T: Record> SourceHandle<T> {
313    /// Validates that `T`'s schema matches the source schema.
314    pub(crate) fn new(entry: Arc<SourceEntry>) -> Result<Self, DbError> {
315        let rust_schema = T::schema();
316        let sql_schema = &entry.schema;
317
318        // Validate field count matches
319        if rust_schema.fields().len() != sql_schema.fields().len() {
320            return Err(DbError::SchemaMismatch(format!(
321                "Rust type has {} fields but source '{}' has {} columns",
322                rust_schema.fields().len(),
323                entry.name,
324                sql_schema.fields().len()
325            )));
326        }
327
328        Ok(Self {
329            entry,
330            _phantom: PhantomData,
331        })
332    }
333
334    /// Push a single record.
335    ///
336    /// # Errors
337    /// Returns `StreamingError` if the pipeline is not running.
338    #[allow(clippy::needless_pass_by_value)]
339    pub fn push(&self, record: T) -> Result<(), laminar_core::streaming::StreamingError> {
340        let batch = record.to_record_batch();
341        self.entry.push_and_buffer(batch)
342    }
343
344    /// Push multiple records, returns count successfully sent.
345    pub fn push_batch(&self, records: impl IntoIterator<Item = T>) -> usize {
346        const BATCH_SIZE: usize = 1024;
347        let mut count = 0;
348        let mut buffer = Vec::with_capacity(BATCH_SIZE);
349
350        for record in records {
351            buffer.push(record);
352            if buffer.len() >= BATCH_SIZE {
353                let batch = T::to_record_batch_from_iter(buffer.drain(..));
354                if self.push_arrow(batch).is_err() {
355                    return count;
356                }
357                count += BATCH_SIZE;
358            }
359        }
360
361        if !buffer.is_empty() {
362            let len = buffer.len();
363            let batch = T::to_record_batch_from_iter(buffer);
364            if self.push_arrow(batch).is_ok() {
365                count += len;
366            }
367        }
368        count
369    }
370
371    /// Push a raw `RecordBatch` (sent to pipeline and buffered for snapshots).
372    ///
373    /// # Errors
374    /// Returns `StreamingError` if the pipeline is not running.
375    pub fn push_arrow(
376        &self,
377        batch: RecordBatch,
378    ) -> Result<(), laminar_core::streaming::StreamingError> {
379        self.entry.push_and_buffer(batch)
380    }
381
382    /// Emit a watermark.
383    pub fn watermark(&self, timestamp: i64) {
384        self.entry.source.watermark(timestamp);
385    }
386
387    /// Current watermark.
388    #[must_use]
389    pub fn current_watermark(&self) -> i64 {
390        self.entry.source.current_watermark()
391    }
392
393    /// Buffered record count.
394    #[must_use]
395    pub fn pending(&self) -> usize {
396        self.entry.source.pending()
397    }
398
399    /// Buffer capacity.
400    #[must_use]
401    pub fn capacity(&self) -> usize {
402        self.entry.source.capacity()
403    }
404
405    /// True when buffer is >80% full.
406    #[must_use]
407    pub fn is_backpressured(&self) -> bool {
408        crate::metrics::is_backpressured(self.pending(), self.capacity())
409    }
410
411    /// Source name.
412    #[must_use]
413    pub fn name(&self) -> &str {
414        &self.entry.name
415    }
416
417    /// Schema.
418    #[must_use]
419    pub fn schema(&self) -> &SchemaRef {
420        &self.entry.schema
421    }
422
423    /// Max out-of-orderness, if configured.
424    #[must_use]
425    pub fn max_out_of_orderness(&self) -> Option<Duration> {
426        self.entry.max_out_of_orderness
427    }
428
429    /// Set the event-time column for watermark-based late-row filtering.
430    pub fn set_event_time_column(&self, column: &str) {
431        self.entry.source.set_event_time_column(column);
432    }
433}
434
435impl<T: Record> std::fmt::Debug for SourceHandle<T> {
436    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
437        f.debug_struct("SourceHandle")
438            .field("name", &self.entry.name)
439            .field("pending", &self.pending())
440            .finish()
441    }
442}
443
444/// Untyped handle for pushing raw `RecordBatch` data.
445pub struct UntypedSourceHandle {
446    entry: Arc<SourceEntry>,
447}
448
449impl UntypedSourceHandle {
450    pub(crate) fn new(entry: Arc<SourceEntry>) -> Self {
451        Self { entry }
452    }
453
454    /// Push a raw `RecordBatch` (sent to pipeline and buffered for snapshots).
455    ///
456    /// # Errors
457    /// Returns `StreamingError` if the pipeline is not running.
458    pub fn push_arrow(
459        &self,
460        batch: RecordBatch,
461    ) -> Result<(), laminar_core::streaming::StreamingError> {
462        self.entry.push_and_buffer(batch)
463    }
464
465    /// Emit a watermark.
466    pub fn watermark(&self, timestamp: i64) {
467        self.entry.source.watermark(timestamp);
468    }
469
470    /// Current watermark.
471    #[must_use]
472    pub fn current_watermark(&self) -> i64 {
473        self.entry.source.current_watermark()
474    }
475
476    /// Buffered record count.
477    #[must_use]
478    pub fn pending(&self) -> usize {
479        self.entry.source.pending()
480    }
481
482    /// Buffer capacity.
483    #[must_use]
484    pub fn capacity(&self) -> usize {
485        self.entry.source.capacity()
486    }
487
488    /// True when buffer is >80% full.
489    #[must_use]
490    pub fn is_backpressured(&self) -> bool {
491        crate::metrics::is_backpressured(self.pending(), self.capacity())
492    }
493
494    /// Source name.
495    #[must_use]
496    pub fn name(&self) -> &str {
497        &self.entry.name
498    }
499
500    /// Schema.
501    #[must_use]
502    pub fn schema(&self) -> &SchemaRef {
503        &self.entry.schema
504    }
505
506    /// Max out-of-orderness, if configured.
507    #[must_use]
508    pub fn max_out_of_orderness(&self) -> Option<Duration> {
509        self.entry.max_out_of_orderness
510    }
511
512    /// Set the event-time column for watermark-based late-row filtering.
513    pub fn set_event_time_column(&self, column: &str) {
514        self.entry.source.set_event_time_column(column);
515    }
516}
517
518impl std::fmt::Debug for UntypedSourceHandle {
519    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
520        f.debug_struct("UntypedSourceHandle")
521            .field("name", &self.entry.name)
522            .finish()
523    }
524}
525
526/// Type of a node in the pipeline topology.
527#[derive(Debug, Clone, Copy, PartialEq, Eq)]
528pub enum PipelineNodeType {
529    /// A data source (CREATE SOURCE).
530    Source,
531    /// A continuous stream (CREATE STREAM).
532    Stream,
533    /// A data sink (CREATE SINK).
534    Sink,
535}
536
537/// A node in the pipeline topology graph.
538#[derive(Debug, Clone)]
539pub struct PipelineNode {
540    /// Node name.
541    pub name: String,
542    /// Node type.
543    pub node_type: PipelineNodeType,
544    /// Arrow schema (sources only).
545    pub schema: Option<SchemaRef>,
546    /// SQL definition (streams only).
547    pub sql: Option<String>,
548}
549
550/// A directed edge in the pipeline topology graph.
551#[derive(Debug, Clone)]
552pub struct PipelineEdge {
553    /// Source node.
554    pub from: String,
555    /// Target node.
556    pub to: String,
557}
558
559/// The complete pipeline topology: nodes and edges.
560#[derive(Debug, Clone)]
561pub struct PipelineTopology {
562    /// Nodes.
563    pub nodes: Vec<PipelineNode>,
564    /// Edges.
565    pub edges: Vec<PipelineEdge>,
566}
567
568/// Metadata about a registered stream.
569#[derive(Debug, Clone)]
570pub struct StreamInfo {
571    /// Name.
572    pub name: String,
573    /// Defining SQL query.
574    pub sql: Option<String>,
575}
576
577/// Information about a registered source.
578#[derive(Debug, Clone)]
579pub struct SourceInfo {
580    /// Name.
581    pub name: String,
582    /// Schema.
583    pub schema: SchemaRef,
584    /// Watermark column, if configured.
585    pub watermark_column: Option<String>,
586}
587
588/// Information about a registered sink.
589#[derive(Debug, Clone)]
590pub struct SinkInfo {
591    /// Name.
592    pub name: String,
593}
594
595/// JSON-serializable projection of a `MaterializedView` for the control-plane
596/// HTTP API.
597#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
598pub struct MaterializedViewInfo {
599    /// View name.
600    pub name: String,
601    /// SQL definition.
602    pub sql: String,
603    /// Current execution state (e.g. `"Running"`, `"Dropping"`).
604    pub state: String,
605}
606
607impl From<&laminar_core::mv::MaterializedView> for MaterializedViewInfo {
608    fn from(view: &laminar_core::mv::MaterializedView) -> Self {
609        Self {
610            name: view.name.clone(),
611            sql: view.sql.clone(),
612            state: format!("{:?}", view.state),
613        }
614    }
615}
616
617/// Information about a running query.
618#[derive(Debug, Clone)]
619pub struct QueryInfo {
620    /// ID.
621    pub id: u64,
622    /// SQL text.
623    pub sql: String,
624    /// Active.
625    pub active: bool,
626}
laminar_db/handle.rs

laminar_db/
handle.rs
