laminar_db/catalog/
mod.rs

1//! Source and sink catalog for tracking registered streaming objects.
2#![allow(clippy::disallowed_types)] // cold path
3
4use std::collections::{HashMap, VecDeque};
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::sync::Arc;
7use std::time::Duration;
8
9use arrow::array::RecordBatch;
10use arrow::datatypes::{Schema, SchemaRef};
11use parking_lot::RwLock;
12use tokio::sync::Notify;
13
14use laminar_core::streaming::{
15    self, BackpressureStrategy, SourceConfig, StreamingError, WaitStrategy,
16};
17
18pub(crate) fn schema_has_reserved_mutation_columns(schema: &Schema) -> bool {
19    schema.fields().iter().any(|field| {
20        ["_op", "__op", laminar_core::changelog::WEIGHT_COLUMN]
21            .iter()
22            .any(|reserved| field.name().eq_ignore_ascii_case(reserved))
23    })
24}
25
26pub(crate) fn validate_source_batch(
27    source_name: &str,
28    expected_schema: &SchemaRef,
29    primary_key: &[String],
30    primary_key_indices: &[usize],
31    batch: &RecordBatch,
32) -> Result<(), StreamingError> {
33    let actual_schema = batch.schema();
34    if !Arc::ptr_eq(&actual_schema, expected_schema)
35        && actual_schema.as_ref() != expected_schema.as_ref()
36    {
37        return Err(StreamingError::SchemaMismatch {
38            expected: expected_schema
39                .fields()
40                .iter()
41                .map(|field| field.name().clone())
42                .collect(),
43            actual: actual_schema
44                .fields()
45                .iter()
46                .map(|field| field.name().clone())
47                .collect(),
48        });
49    }
50    if primary_key.len() != primary_key_indices.len() {
51        return Err(StreamingError::InvalidConfig(format!(
52            "source '{source_name}' primary-key metadata is inconsistent"
53        )));
54    }
55    for (column, &index) in primary_key.iter().zip(primary_key_indices) {
56        let null_count = batch.column(index).null_count();
57        if null_count != 0 {
58            return Err(StreamingError::InvalidConfig(format!(
59                "source '{source_name}' primary-key column '{column}' contains {null_count} null value(s)"
60            )));
61        }
62    }
63    Ok(())
64}
65
66/// Record type for Arrow-based streaming subscriptions.
67#[derive(Clone, Debug)]
68pub struct ArrowRecord {
69    pub(crate) batch: RecordBatch,
70}
71
72impl laminar_core::streaming::Record for ArrowRecord {
73    fn schema() -> SchemaRef {
74        // This is a placeholder; the actual schema is on the SourceEntry.
75        // Source intake uses push_arrow; the count-bounded query-output bridge uses Record.
76        Arc::new(arrow::datatypes::Schema::empty())
77    }
78
79    fn to_record_batch(&self) -> RecordBatch {
80        self.batch.clone()
81    }
82}
83
84/// Recent snapshot history, independently bounded by count and retained Arrow bytes.
85struct SnapshotRing {
86    batches: VecDeque<(RecordBatch, usize)>,
87    bytes: usize,
88    capacity: usize,
89    max_bytes: usize,
90}
91
92impl SnapshotRing {
93    fn new(capacity: usize, max_bytes: usize) -> Self {
94        Self {
95            batches: VecDeque::new(),
96            bytes: 0,
97            capacity: capacity.max(1),
98            max_bytes,
99        }
100    }
101
102    fn push(&mut self, batch: RecordBatch) {
103        let bytes = streaming::retained_arrow_bytes(&batch);
104        // INVARIANT: successful source admission already checked this batch against the
105        // same byte limit. Eviction changes history only after queue admission succeeds.
106        debug_assert!(bytes <= self.max_bytes);
107        while self.batches.len() >= self.capacity || self.bytes > self.max_bytes - bytes {
108            if let Some((_, evicted_bytes)) = self.batches.pop_front() {
109                self.bytes -= evicted_bytes;
110            }
111        }
112        self.batches.push_back((batch, bytes));
113        self.bytes += bytes;
114    }
115
116    fn snapshot(&self) -> Vec<RecordBatch> {
117        self.batches
118            .iter()
119            .map(|(batch, _)| batch.clone())
120            .collect()
121    }
122}
123
124/// A registered source in the catalog.
125pub struct SourceEntry {
126    /// Source name.
127    pub name: String,
128    /// Arrow schema.
129    pub schema: SchemaRef,
130    /// Primary-key columns in declaration order.
131    pub primary_key: Vec<String>,
132    primary_key_indices: Vec<usize>,
133    /// Watermark column name, if configured.
134    pub watermark_column: Option<String>,
135    /// Maximum out-of-orderness for watermark generation.
136    pub max_out_of_orderness: Option<Duration>,
137    /// Whether this source uses `PROCTIME()` watermarks.
138    pub is_processing_time: std::sync::atomic::AtomicBool,
139    pub(crate) source: streaming::Source<ArrowRecord>,
140    /// The queue's single consumer while no pipeline bridge holds it; `None` while leased
141    /// and in schema-planning catalogs.
142    ingress: parking_lot::Mutex<Option<streaming::Sink<ArrowRecord>>>,
143    buffer: parking_lot::Mutex<SnapshotRing>,
144    /// Wakeup handle for `db.insert()` event-driven notification.
145    data_notify: Arc<Notify>,
146}
147
148impl SourceEntry {
149    /// Push a batch to both the channel and the snapshot ring.
150    pub(crate) fn push_and_buffer(
151        &self,
152        batch: RecordBatch,
153    ) -> Result<(), laminar_core::streaming::StreamingError> {
154        validate_source_batch(
155            &self.name,
156            &self.schema,
157            &self.primary_key,
158            &self.primary_key_indices,
159            &batch,
160        )?;
161        // Serialize admission and history publication so concurrent producers cannot leave
162        // accepted batches waiting outside either owner or publish snapshots out of order.
163        let mut buffer = self.buffer.lock();
164        self.source.push_arrow(batch.clone())?;
165        buffer.push(batch);
166        drop(buffer);
167        self.data_notify.notify_one();
168        Ok(())
169    }
170
171    pub(crate) fn snapshot(&self) -> Vec<RecordBatch> {
172        self.buffer.lock().snapshot()
173    }
174
175    pub(crate) fn data_notify(&self) -> Arc<Notify> {
176        Arc::clone(&self.data_notify)
177    }
178
179    /// Leases the queue consumer to one pipeline bridge until the lease drops.
180    pub(crate) fn lease_ingress(self: &Arc<Self>) -> Result<IngressLease, crate::DbError> {
181        let sink = self.ingress.lock().take().ok_or_else(|| {
182            crate::DbError::Pipeline(format!(
183                "source '{}' input queue is already consumed by another pipeline bridge",
184                self.name
185            ))
186        })?;
187        Ok(IngressLease {
188            entry: Arc::clone(self),
189            sink: Some(sink),
190        })
191    }
192
193    pub(crate) fn is_backpressured(&self) -> bool {
194        crate::metrics::is_backpressured(self.source.pending(), self.source.capacity())
195            || crate::metrics::is_backpressured(
196                self.source.queued_arrow_bytes(),
197                self.source.max_queued_bytes(),
198            )
199    }
200}
201
202/// A source queue consumer leased to one pipeline bridge. Dropping the lease returns the
203/// consumer with any undelivered batches still queued, so the next bridge receives them.
204pub(crate) struct IngressLease {
205    entry: Arc<SourceEntry>,
206    sink: Option<streaming::Sink<ArrowRecord>>,
207}
208
209impl IngressLease {
210    pub(crate) fn poll(&mut self) -> Option<RecordBatch> {
211        self.sink.as_mut().and_then(streaming::Sink::poll)
212    }
213}
214
215impl Drop for IngressLease {
216    fn drop(&mut self) {
217        if let Some(sink) = self.sink.take() {
218            *self.entry.ingress.lock() = Some(sink);
219        }
220    }
221}
222
223pub(crate) struct SinkEntry {
224    pub(crate) input: String,
225}
226
227pub(crate) struct QueryEntry {
228    pub(crate) id: u64,
229    pub(crate) sql: String,
230    pub(crate) active: bool,
231}
232
233pub(crate) struct StreamEntry {
234    pub(crate) name: String,
235    emitted_rows: AtomicU64,
236}
237
238impl StreamEntry {
239    pub(crate) fn record_emitted_rows(&self, rows: u64) {
240        self.emitted_rows.fetch_add(rows, Ordering::Relaxed);
241    }
242
243    pub(crate) fn emitted_rows(&self) -> u64 {
244        self.emitted_rows.load(Ordering::Relaxed)
245    }
246}
247
248/// Central registry of sources, sinks, streams, and queries.
249pub struct SourceCatalog {
250    sources: RwLock<HashMap<String, Arc<SourceEntry>>>,
251    sinks: RwLock<HashMap<String, SinkEntry>>,
252    streams: RwLock<HashMap<String, Arc<StreamEntry>>>,
253    queries: RwLock<HashMap<u64, QueryEntry>>,
254    next_query_id: AtomicU64,
255    default_buffer_size: usize,
256    default_backpressure: BackpressureStrategy,
257    push_source_max_bytes: usize,
258    // Only the unstarted private topology planner uses schema-only queues. This flag is checked
259    // during catalog registration, never during record intake, push, poll, or subscription.
260    #[cfg(feature = "cluster")]
261    topology_planning: bool,
262}
263
264impl SourceCatalog {
265    /// Create a catalog with the given defaults for new sources.
266    #[must_use]
267    pub fn new(buffer_size: usize, backpressure: BackpressureStrategy) -> Self {
268        Self {
269            sources: RwLock::new(HashMap::new()),
270            sinks: RwLock::new(HashMap::new()),
271            streams: RwLock::new(HashMap::new()),
272            queries: RwLock::new(HashMap::new()),
273            next_query_id: AtomicU64::new(1),
274            default_buffer_size: buffer_size,
275            default_backpressure: backpressure,
276            push_source_max_bytes: streaming::DEFAULT_SOURCE_MAX_QUEUED_BYTES,
277            #[cfg(feature = "cluster")]
278            topology_planning: false,
279        }
280    }
281
282    pub(crate) fn from_config(config: &crate::LaminarConfig) -> Self {
283        Self {
284            push_source_max_bytes: config.push_source_max_bytes,
285            ..Self::new(config.default_buffer_size, config.default_backpressure)
286        }
287    }
288
289    #[cfg(feature = "cluster")]
290    pub(crate) fn for_topology_planning(config: &crate::LaminarConfig) -> Self {
291        Self {
292            topology_planning: true,
293            ..Self::from_config(config)
294        }
295    }
296
297    #[allow(clippy::too_many_arguments)]
298    pub(crate) fn register_source(
299        &self,
300        name: &str,
301        schema: SchemaRef,
302        primary_key: Vec<String>,
303        watermark_column: Option<String>,
304        max_out_of_orderness: Option<Duration>,
305        buffer_size: Option<usize>,
306        backpressure: Option<BackpressureStrategy>,
307    ) -> Result<Arc<SourceEntry>, crate::DbError> {
308        let mut sources = self.sources.write();
309        if sources.contains_key(name) {
310            return Err(crate::DbError::SourceAlreadyExists(name.to_string()));
311        }
312
313        let mut primary_key_indices = Vec::with_capacity(primary_key.len());
314        for column in &primary_key {
315            let index = schema.index_of(column).map_err(|_| {
316                crate::DbError::InvalidOperation(format!(
317                    "source '{name}' primary-key column '{column}' is absent from its schema"
318                ))
319            })?;
320            if primary_key_indices.contains(&index) {
321                return Err(crate::DbError::InvalidOperation(format!(
322                    "source '{name}' primary key repeats column '{column}'"
323                )));
324            }
325            if schema.field(index).is_nullable() {
326                return Err(crate::DbError::InvalidOperation(format!(
327                    "source '{name}' primary-key column '{column}' must be non-nullable"
328                )));
329            }
330            primary_key_indices.push(index);
331        }
332
333        let buf_size = buffer_size.unwrap_or(self.default_buffer_size);
334        #[cfg(feature = "cluster")]
335        let buf_size = if self.topology_planning { 1 } else { buf_size };
336        let bp = backpressure.unwrap_or(self.default_backpressure);
337
338        // Channel buffer is at least 1024 to avoid blocking on small snapshot rings.
339        let channel_buf = buf_size.max(1024);
340        let config = SourceConfig {
341            channel: streaming::ChannelConfig {
342                buffer_size: channel_buf,
343                backpressure: bp,
344                wait_strategy: WaitStrategy::SpinYield,
345                track_stats: false,
346            },
347            name: Some(name.to_string()),
348            max_queued_bytes: self.push_source_max_bytes,
349        };
350
351        #[cfg(feature = "cluster")]
352        let (source, ingress) = if self.topology_planning {
353            (
354                streaming::create_for_schema_planning::<ArrowRecord>(config),
355                None,
356            )
357        } else {
358            let (source, sink) = streaming::create_with_config::<ArrowRecord>(config);
359            (source, Some(sink))
360        };
361        #[cfg(not(feature = "cluster"))]
362        let (source, ingress) = {
363            let (source, sink) = streaming::create_with_config::<ArrowRecord>(config);
364            (source, Some(sink))
365        };
366
367        let entry = Arc::new(SourceEntry {
368            name: name.to_string(),
369            schema,
370            primary_key,
371            primary_key_indices,
372            watermark_column,
373            max_out_of_orderness,
374            is_processing_time: std::sync::atomic::AtomicBool::new(false),
375            source,
376            ingress: parking_lot::Mutex::new(ingress),
377            buffer: parking_lot::Mutex::new(SnapshotRing::new(
378                buf_size,
379                self.push_source_max_bytes,
380            )),
381            data_notify: Arc::new(Notify::new()),
382        });
383
384        sources.insert(name.to_string(), Arc::clone(&entry));
385        Ok(entry)
386    }
387
388    #[cfg(test)]
389    pub(crate) fn register_source_or_replace(
390        &self,
391        name: &str,
392        schema: SchemaRef,
393        primary_key: Vec<String>,
394        watermark_column: Option<String>,
395        max_out_of_orderness: Option<Duration>,
396        buffer_size: Option<usize>,
397        backpressure: Option<BackpressureStrategy>,
398    ) -> Arc<SourceEntry> {
399        // Remove existing if present
400        self.sources.write().remove(name);
401        // Safe to unwrap since we just removed any conflict
402        self.register_source(
403            name,
404            schema,
405            primary_key,
406            watermark_column,
407            max_out_of_orderness,
408            buffer_size,
409            backpressure,
410        )
411        .unwrap()
412    }
413
414    /// Look up a registered source by name.
415    pub fn get_source(&self, name: &str) -> Option<Arc<SourceEntry>> {
416        self.sources.read().get(name).cloned()
417    }
418
419    /// Returns `true` if the source existed.
420    pub fn drop_source(&self, name: &str) -> bool {
421        self.sources.write().remove(name).is_some()
422    }
423
424    pub(crate) fn register_sink(&self, name: &str, input: &str) -> Result<(), crate::DbError> {
425        let mut sinks = self.sinks.write();
426        if sinks.contains_key(name) {
427            return Err(crate::DbError::SinkAlreadyExists(name.to_string()));
428        }
429        sinks.insert(
430            name.to_string(),
431            SinkEntry {
432                input: input.to_string(),
433            },
434        );
435        Ok(())
436    }
437
438    /// Returns `true` if the sink existed.
439    pub fn drop_sink(&self, name: &str) -> bool {
440        self.sinks.write().remove(name).is_some()
441    }
442
443    pub(crate) fn register_stream(&self, name: &str) -> Result<(), crate::DbError> {
444        let mut streams = self.streams.write();
445        if streams.contains_key(name) {
446            return Err(crate::DbError::StreamAlreadyExists(name.to_string()));
447        }
448
449        streams.insert(
450            name.to_string(),
451            Arc::new(StreamEntry {
452                name: name.to_string(),
453                emitted_rows: AtomicU64::new(0),
454            }),
455        );
456        Ok(())
457    }
458
459    pub(crate) fn get_stream_entry(&self, name: &str) -> Option<Arc<StreamEntry>> {
460        self.streams.read().get(name).cloned()
461    }
462
463    /// Returns `true` if the stream existed.
464    pub fn drop_stream(&self, name: &str) -> bool {
465        self.streams.write().remove(name).is_some()
466    }
467
468    /// All registered stream names.
469    pub fn list_streams(&self) -> Vec<String> {
470        self.streams.read().keys().cloned().collect()
471    }
472
473    /// All registered source names.
474    pub fn list_sources(&self) -> Vec<String> {
475        self.sources.read().keys().cloned().collect()
476    }
477
478    /// All registered sink names.
479    pub fn list_sinks(&self) -> Vec<String> {
480        self.sinks.read().keys().cloned().collect()
481    }
482
483    /// Input source/table name for a sink, if registered.
484    pub fn get_sink_input(&self, name: &str) -> Option<String> {
485        self.sinks.read().get(name).map(|e| e.input.clone())
486    }
487
488    pub(crate) fn register_query(&self, sql: &str) -> u64 {
489        let id = self.next_query_id.fetch_add(1, Ordering::Relaxed);
490        let mut queries = self.queries.write();
491        queries.insert(
492            id,
493            QueryEntry {
494                id,
495                sql: sql.to_string(),
496                active: true,
497            },
498        );
499        id
500    }
501
502    pub(crate) fn deactivate_query(&self, id: u64) -> bool {
503        // Cap retained deactivated queries so finished SELECTs can't accumulate
504        // unboundedly; over the cap, the oldest (lowest id) is dropped.
505        const MAX_INACTIVE_QUERIES: usize = 100;
506        let mut queries = self.queries.write();
507        if let Some(entry) = queries.get_mut(&id) {
508            let was_active = entry.active;
509            entry.active = false;
510            if was_active {
511                let inactive_count = queries.values().filter(|q| !q.active).count();
512                if inactive_count > MAX_INACTIVE_QUERIES {
513                    let oldest_inactive_id =
514                        queries.values().filter(|q| !q.active).map(|q| q.id).min();
515                    if let Some(oldest_id) = oldest_inactive_id {
516                        queries.remove(&oldest_id);
517                    }
518                }
519            }
520            true
521        } else {
522            false
523        }
524    }
525
526    pub(crate) fn list_queries(&self) -> Vec<(u64, String, bool)> {
527        self.queries
528            .read()
529            .values()
530            .map(|q| (q.id, q.sql.clone(), q.active))
531            .collect()
532    }
533
534    /// Schema for DESCRIBE queries.
535    pub fn describe_source(&self, name: &str) -> Option<SchemaRef> {
536        self.sources.read().get(name).map(|e| e.schema.clone())
537    }
538}
539
540#[cfg(test)]
541mod tests;
542
543#[cfg(test)]
544mod admission_tests;
laminar_db/catalog/mod.rs

laminar_db/catalog/
mod.rs
