laminar_db/

handle.rs

//! Handle types for query results, source access, and subscriptions.

use std::marker::PhantomData;
use std::sync::Arc;
use std::time::Duration;

use arrow::array::RecordBatch;
use arrow::datatypes::SchemaRef;

use laminar_core::streaming::{Record, Subscription};

use crate::catalog::{ArrowRecord, SourceEntry};
use crate::DbError;

/// Result of executing a SQL statement.
#[derive(Debug)]
pub enum ExecuteResult {
    /// DDL statement completed (CREATE, DROP, ALTER).
    Ddl(DdlInfo),
    /// Query is running, subscribe to results.
    Query(QueryHandle),
    /// Rows were affected (INSERT INTO).
    RowsAffected(u64),
    /// Metadata result (SHOW, DESCRIBE).
    Metadata(RecordBatch),
}

impl ExecuteResult {
    /// Convert to `QueryHandle`, returning an error if this is not a query result.
    ///
    /// # Errors
    ///
    /// Returns `DbError::InvalidOperation` if this is not a query result.
    pub fn into_query(self) -> Result<QueryHandle, DbError> {
        match self {
            Self::Query(q) => Ok(q),
            _ => Err(DbError::InvalidOperation(
                "Expected a query result".to_string(),
            )),
        }
    }
}

/// Information about a completed DDL statement.
#[derive(Debug, Clone)]
pub struct DdlInfo {
    /// The statement type (e.g., "CREATE SOURCE").
    pub statement_type: String,
    /// The object name affected.
    pub object_name: String,
}

/// Handle to a running streaming query.
#[derive(Debug)]
pub struct QueryHandle {
    /// Query identifier.
    pub(crate) id: u64,
    /// Output schema.
    pub(crate) schema: SchemaRef,
    /// The SQL text.
    pub(crate) sql: String,
    /// The subscription for receiving results.
    pub(crate) subscription: Option<Subscription<ArrowRecord>>,
    /// Whether the query is active.
    pub(crate) active: bool,
}

impl QueryHandle {
    /// Get the output schema.
    #[must_use]
    pub fn schema(&self) -> &SchemaRef {
        &self.schema
    }

    /// Get the query SQL text.
    #[must_use]
    pub fn sql(&self) -> &str {
        &self.sql
    }

    /// Get the query ID.
    #[must_use]
    pub fn id(&self) -> u64 {
        self.id
    }

    /// Check if the query is still active.
    #[must_use]
    pub fn is_active(&self) -> bool {
        self.active
    }

    /// Subscribe to raw `RecordBatch` results.
    pub(crate) fn subscribe_raw(&mut self) -> Result<Subscription<ArrowRecord>, DbError> {
        self.subscription
            .take()
            .ok_or_else(|| DbError::InvalidOperation("Subscription already consumed".to_string()))
    }

    /// Subscribe to typed results.
    ///
    /// The type `T` must implement `from_batch()` and `from_batch_all()` methods,
    /// which are generated by `#[derive(FromRecordBatch)]`.
    ///
    /// # Errors
    ///
    /// Returns `DbError::InvalidOperation` if the subscription was already consumed.
    pub fn subscribe<T: FromBatch>(&mut self) -> Result<TypedSubscription<T>, DbError> {
        let sub = self.subscribe_raw()?;
        Ok(TypedSubscription {
            inner: sub,
            _phantom: PhantomData,
        })
    }

    /// Cancel this query.
    pub fn cancel(&mut self) {
        self.active = false;
        self.subscription = None;
    }
}

/// Trait for types that can be deserialized from a `RecordBatch`.
///
/// Auto-generated by `#[derive(FromRecordBatch)]`.
pub trait FromBatch: Sized {
    /// Deserialize a single row from a `RecordBatch`.
    fn from_batch(batch: &RecordBatch, row: usize) -> Self;
    /// Deserialize all rows from a `RecordBatch`.
    fn from_batch_all(batch: &RecordBatch) -> Vec<Self>;
}

/// Typed subscription that deserializes `RecordBatch` rows.
pub struct TypedSubscription<T: FromBatch> {
    inner: Subscription<ArrowRecord>,
    _phantom: PhantomData<T>,
}

impl<T: FromBatch> TypedSubscription<T> {
    /// Create from a raw subscription.
    pub(crate) fn from_raw(sub: Subscription<ArrowRecord>) -> Self {
        Self {
            inner: sub,
            _phantom: PhantomData,
        }
    }

    /// Poll for the next batch of typed results (non-blocking).
    #[must_use]
    pub fn poll(&self) -> Option<Vec<T>> {
        self.inner.poll().map(|batch| T::from_batch_all(&batch))
    }

    /// Blocking receive.
    ///
    /// # Errors
    ///
    /// Returns `RecvError` if the channel is disconnected.
    pub fn recv(&self) -> Result<Vec<T>, laminar_core::streaming::RecvError> {
        self.inner.recv().map(|batch| T::from_batch_all(&batch))
    }

    /// Receive with timeout.
    ///
    /// # Errors
    ///
    /// Returns `RecvError` on timeout or if the channel is disconnected.
    pub fn recv_timeout(
        &self,
        timeout: Duration,
    ) -> Result<Vec<T>, laminar_core::streaming::RecvError> {
        self.inner
            .recv_timeout(timeout)
            .map(|batch| T::from_batch_all(&batch))
    }

    /// Zero-allocation callback-based consumption.
    ///
    /// Calls `f` for each deserialized record. Return `false` to stop.
    pub fn poll_each<F: FnMut(T) -> bool>(&self, max_batches: usize, mut f: F) -> usize {
        let mut count = 0;
        for _ in 0..max_batches {
            match self.inner.poll() {
                Some(batch) => {
                    let items = T::from_batch_all(&batch);
                    for item in items {
                        count += 1;
                        if !f(item) {
                            return count;
                        }
                    }
                }
                None => break,
            }
        }
        count
    }

    /// Get the underlying raw subscription.
    #[allow(dead_code)]
    pub(crate) fn into_raw(self) -> Subscription<ArrowRecord> {
        self.inner
    }
}

impl<T: FromBatch> std::fmt::Debug for TypedSubscription<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TypedSubscription").finish()
    }
}

/// Typed handle for pushing data into a registered source.
pub struct SourceHandle<T: Record> {
    entry: Arc<SourceEntry>,
    _phantom: PhantomData<T>,
}

impl<T: Record> SourceHandle<T> {
    /// Create a new source handle from a catalog entry.
    ///
    /// Validates that the Rust type's schema matches the source schema.
    pub(crate) fn new(entry: Arc<SourceEntry>) -> Result<Self, DbError> {
        let rust_schema = T::schema();
        let sql_schema = &entry.schema;

        // Validate field count matches
        if rust_schema.fields().len() != sql_schema.fields().len() {
            return Err(DbError::SchemaMismatch(format!(
                "Rust type has {} fields but source '{}' has {} columns",
                rust_schema.fields().len(),
                entry.name,
                sql_schema.fields().len()
            )));
        }

        Ok(Self {
            entry,
            _phantom: PhantomData,
        })
    }

    /// Push a single record.
    ///
    /// # Errors
    ///
    /// Returns `StreamingError` if the channel is full or closed.
    #[allow(clippy::needless_pass_by_value)]
    pub fn push(&self, record: T) -> Result<(), laminar_core::streaming::StreamingError> {
        let batch = record.to_record_batch();
        self.entry.push_and_buffer(batch)
    }

    /// Push a batch of records.
    pub fn push_batch(&self, records: impl IntoIterator<Item = T>) -> usize {
        const BATCH_SIZE: usize = 1024;
        let mut count = 0;
        let mut buffer = Vec::with_capacity(BATCH_SIZE);

        for record in records {
            buffer.push(record);
            if buffer.len() >= BATCH_SIZE {
                let batch = T::to_record_batch_from_iter(buffer.drain(..));
                if self.push_arrow(batch).is_err() {
                    return count;
                }
                count += BATCH_SIZE;
            }
        }

        if !buffer.is_empty() {
            let len = buffer.len();
            let batch = T::to_record_batch_from_iter(buffer);
            if self.push_arrow(batch).is_ok() {
                count += len;
            }
        }
        count
    }

    /// Push a raw `RecordBatch`.
    ///
    /// The batch is sent to the SPSC channel for pipeline processing and
    /// also buffered for ad-hoc `SELECT` snapshot queries.
    ///
    /// # Errors
    ///
    /// Returns `StreamingError` if the channel is full or closed.
    pub fn push_arrow(
        &self,
        batch: RecordBatch,
    ) -> Result<(), laminar_core::streaming::StreamingError> {
        self.entry.push_and_buffer(batch)
    }

    /// Emit a watermark.
    pub fn watermark(&self, timestamp: i64) {
        self.entry.source.watermark(timestamp);
    }

    /// Get current watermark.
    #[must_use]
    pub fn current_watermark(&self) -> i64 {
        self.entry.source.current_watermark()
    }

    /// Number of buffered records.
    #[must_use]
    pub fn pending(&self) -> usize {
        self.entry.source.pending()
    }

    /// Buffer capacity.
    #[must_use]
    pub fn capacity(&self) -> usize {
        self.entry.source.capacity()
    }

    /// Whether the source buffer is experiencing backpressure (>80% full).
    #[must_use]
    pub fn is_backpressured(&self) -> bool {
        crate::metrics::is_backpressured(self.pending(), self.capacity())
    }

    /// Get the source name.
    #[must_use]
    pub fn name(&self) -> &str {
        &self.entry.name
    }

    /// Get the schema.
    #[must_use]
    pub fn schema(&self) -> &SchemaRef {
        &self.entry.schema
    }

    /// Get the maximum out-of-orderness duration, if configured.
    #[must_use]
    pub fn max_out_of_orderness(&self) -> Option<Duration> {
        self.entry.max_out_of_orderness
    }

    /// Declare which column in the source data represents event time.
    ///
    /// When set, `source.watermark()` enables late-row filtering
    /// without a SQL `WATERMARK FOR` clause.
    pub fn set_event_time_column(&self, column: &str) {
        self.entry.source.set_event_time_column(column);
    }
}

impl<T: Record> std::fmt::Debug for SourceHandle<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SourceHandle")
            .field("name", &self.entry.name)
            .field("pending", &self.pending())
            .finish()
    }
}

/// Untyped handle for pushing raw `RecordBatch` data.
pub struct UntypedSourceHandle {
    entry: Arc<SourceEntry>,
}

impl UntypedSourceHandle {
    /// Create from a catalog entry.
    pub(crate) fn new(entry: Arc<SourceEntry>) -> Self {
        Self { entry }
    }

    /// Push a `RecordBatch`.
    ///
    /// The batch is sent to the SPSC channel for pipeline processing and
    /// also buffered for ad-hoc `SELECT` snapshot queries.
    ///
    /// # Errors
    ///
    /// Returns `StreamingError` if the channel is full or closed.
    pub fn push_arrow(
        &self,
        batch: RecordBatch,
    ) -> Result<(), laminar_core::streaming::StreamingError> {
        self.entry.push_and_buffer(batch)
    }

    /// Emit a watermark.
    pub fn watermark(&self, timestamp: i64) {
        self.entry.source.watermark(timestamp);
    }

    /// Get current watermark.
    #[must_use]
    pub fn current_watermark(&self) -> i64 {
        self.entry.source.current_watermark()
    }

    /// Number of buffered records.
    #[must_use]
    pub fn pending(&self) -> usize {
        self.entry.source.pending()
    }

    /// Buffer capacity.
    #[must_use]
    pub fn capacity(&self) -> usize {
        self.entry.source.capacity()
    }

    /// Whether the source buffer is experiencing backpressure (>80% full).
    #[must_use]
    pub fn is_backpressured(&self) -> bool {
        crate::metrics::is_backpressured(self.pending(), self.capacity())
    }

    /// Get the source name.
    #[must_use]
    pub fn name(&self) -> &str {
        &self.entry.name
    }

    /// Get the schema.
    #[must_use]
    pub fn schema(&self) -> &SchemaRef {
        &self.entry.schema
    }

    /// Get the maximum out-of-orderness duration, if configured.
    #[must_use]
    pub fn max_out_of_orderness(&self) -> Option<Duration> {
        self.entry.max_out_of_orderness
    }

    /// Declare which column in the source data represents event time.
    ///
    /// When set, `source.watermark()` enables late-row filtering
    /// without a SQL `WATERMARK FOR` clause.
    pub fn set_event_time_column(&self, column: &str) {
        self.entry.source.set_event_time_column(column);
    }
}

impl std::fmt::Debug for UntypedSourceHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("UntypedSourceHandle")
            .field("name", &self.entry.name)
            .finish()
    }
}

/// Type of a node in the pipeline topology.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PipelineNodeType {
    /// A data source (CREATE SOURCE).
    Source,
    /// A continuous stream (CREATE STREAM).
    Stream,
    /// A data sink (CREATE SINK).
    Sink,
}

/// A node in the pipeline topology graph.
#[derive(Debug, Clone)]
pub struct PipelineNode {
    /// Node name.
    pub name: String,
    /// Node type (source, stream, or sink).
    pub node_type: PipelineNodeType,
    /// Arrow schema, if available (sources have schemas).
    pub schema: Option<SchemaRef>,
    /// SQL definition, if applicable (streams have query SQL).
    pub sql: Option<String>,
}

/// A directed edge in the pipeline topology graph.
#[derive(Debug, Clone)]
pub struct PipelineEdge {
    /// Source node name.
    pub from: String,
    /// Target node name.
    pub to: String,
}

/// The complete pipeline topology: nodes and edges.
#[derive(Debug, Clone)]
pub struct PipelineTopology {
    /// All nodes in the pipeline.
    pub nodes: Vec<PipelineNode>,
    /// All edges (data flow connections).
    pub edges: Vec<PipelineEdge>,
}

/// Metadata about a registered stream.
#[derive(Debug, Clone)]
pub struct StreamInfo {
    /// Stream name.
    pub name: String,
    /// The SQL query that defines the stream.
    pub sql: Option<String>,
}

/// Information about a registered source.
#[derive(Debug, Clone)]
pub struct SourceInfo {
    /// Source name.
    pub name: String,
    /// Schema.
    pub schema: SchemaRef,
    /// Watermark column, if configured.
    pub watermark_column: Option<String>,
}

/// Information about a registered sink.
#[derive(Debug, Clone)]
pub struct SinkInfo {
    /// Sink name.
    pub name: String,
}

/// Information about a running query.
#[derive(Debug, Clone)]
pub struct QueryInfo {
    /// Query identifier.
    pub id: u64,
    /// SQL text.
    pub sql: String,
    /// Whether the query is active.
    pub active: bool,
}