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;
5use std::sync::atomic::{AtomicU64, AtomicUsize, 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        // ArrowRecord is only used as a type parameter; push_arrow bypasses this.
76        Arc::new(arrow::datatypes::Schema::empty())
77    }
78
79    fn to_record_batch(&self) -> RecordBatch {
80        self.batch.clone()
81    }
82}
83
84/// Bounded ring buffer for snapshot batches.
85///
86/// Concurrent `push()` calls each get a unique slot via atomic `fetch_add`.
87/// Per-slot mutex protects the actual read/write.
88struct SnapshotRing {
89    slots: Box<[parking_lot::Mutex<Option<RecordBatch>>]>,
90    tail: AtomicUsize,
91    capacity: usize,
92}
93
94impl SnapshotRing {
95    fn new(capacity: usize) -> Self {
96        let cap = capacity.max(1);
97        let slots: Vec<_> = (0..cap).map(|_| parking_lot::Mutex::new(None)).collect();
98        Self {
99            slots: slots.into_boxed_slice(),
100            tail: AtomicUsize::new(0),
101            capacity: cap,
102        }
103    }
104
105    fn push(&self, batch: RecordBatch) {
106        // fetch_add is atomic — concurrent pushers each get a unique slot.
107        let idx = self.tail.fetch_add(1, Ordering::Relaxed) % self.capacity;
108        *self.slots[idx].lock() = Some(batch);
109    }
110
111    fn snapshot(&self) -> Vec<RecordBatch> {
112        let tail = self.tail.load(Ordering::Acquire);
113        let count = tail.min(self.capacity);
114        // Read the most recent `count` slots, oldest first.
115        let start = if tail <= self.capacity {
116            0
117        } else {
118            tail % self.capacity
119        };
120        let mut result = Vec::with_capacity(count);
121        for i in 0..count {
122            let idx = (start + i) % self.capacity;
123            if let Some(batch) = self.slots[idx].lock().as_ref() {
124                result.push(batch.clone());
125            }
126        }
127        result
128    }
129}
130
131/// A registered source in the catalog.
132pub struct SourceEntry {
133    /// Source name.
134    pub name: String,
135    /// Arrow schema.
136    pub schema: SchemaRef,
137    /// Primary-key columns in declaration order.
138    pub primary_key: Vec<String>,
139    primary_key_indices: Vec<usize>,
140    /// Watermark column name, if configured.
141    pub watermark_column: Option<String>,
142    /// Maximum out-of-orderness for watermark generation.
143    pub max_out_of_orderness: Option<Duration>,
144    /// Whether this source uses `PROCTIME()` watermarks.
145    pub is_processing_time: std::sync::atomic::AtomicBool,
146    pub(crate) source: streaming::Source<ArrowRecord>,
147    pub(crate) sink: streaming::Sink<ArrowRecord>,
148    buffer: SnapshotRing,
149    /// Wakeup handle for `db.insert()` event-driven notification.
150    data_notify: Arc<Notify>,
151}
152
153impl SourceEntry {
154    /// Push a batch to both the channel and the snapshot ring.
155    pub(crate) fn push_and_buffer(
156        &self,
157        batch: RecordBatch,
158    ) -> Result<(), laminar_core::streaming::StreamingError> {
159        validate_source_batch(
160            &self.name,
161            &self.schema,
162            &self.primary_key,
163            &self.primary_key_indices,
164            &batch,
165        )?;
166        self.source.push_arrow(batch.clone())?;
167        self.buffer.push(batch);
168        self.data_notify.notify_one();
169        Ok(())
170    }
171
172    pub(crate) fn snapshot(&self) -> Vec<RecordBatch> {
173        self.buffer.snapshot()
174    }
175
176    pub(crate) fn data_notify(&self) -> Arc<Notify> {
177        Arc::clone(&self.data_notify)
178    }
179}
180
181pub(crate) struct SinkEntry {
182    pub(crate) input: String,
183}
184
185pub(crate) struct QueryEntry {
186    pub(crate) id: u64,
187    pub(crate) sql: String,
188    pub(crate) active: bool,
189}
190
191pub(crate) struct StreamEntry {
192    pub(crate) name: String,
193    emitted_rows: AtomicU64,
194}
195
196impl StreamEntry {
197    pub(crate) fn record_emitted_rows(&self, rows: u64) {
198        self.emitted_rows.fetch_add(rows, Ordering::Relaxed);
199    }
200
201    pub(crate) fn emitted_rows(&self) -> u64 {
202        self.emitted_rows.load(Ordering::Relaxed)
203    }
204}
205
206/// Central registry of sources, sinks, streams, and queries.
207pub struct SourceCatalog {
208    sources: RwLock<HashMap<String, Arc<SourceEntry>>>,
209    sinks: RwLock<HashMap<String, SinkEntry>>,
210    streams: RwLock<HashMap<String, Arc<StreamEntry>>>,
211    queries: RwLock<HashMap<u64, QueryEntry>>,
212    next_query_id: AtomicU64,
213    default_buffer_size: usize,
214    default_backpressure: BackpressureStrategy,
215}
216
217impl SourceCatalog {
218    /// Create a catalog with the given defaults for new sources.
219    #[must_use]
220    pub fn new(buffer_size: usize, backpressure: BackpressureStrategy) -> Self {
221        Self {
222            sources: RwLock::new(HashMap::new()),
223            sinks: RwLock::new(HashMap::new()),
224            streams: RwLock::new(HashMap::new()),
225            queries: RwLock::new(HashMap::new()),
226            next_query_id: AtomicU64::new(1),
227            default_buffer_size: buffer_size,
228            default_backpressure: backpressure,
229        }
230    }
231
232    #[allow(clippy::too_many_arguments)]
233    pub(crate) fn register_source(
234        &self,
235        name: &str,
236        schema: SchemaRef,
237        primary_key: Vec<String>,
238        watermark_column: Option<String>,
239        max_out_of_orderness: Option<Duration>,
240        buffer_size: Option<usize>,
241        backpressure: Option<BackpressureStrategy>,
242    ) -> Result<Arc<SourceEntry>, crate::DbError> {
243        let mut sources = self.sources.write();
244        if sources.contains_key(name) {
245            return Err(crate::DbError::SourceAlreadyExists(name.to_string()));
246        }
247
248        let mut primary_key_indices = Vec::with_capacity(primary_key.len());
249        for column in &primary_key {
250            let index = schema.index_of(column).map_err(|_| {
251                crate::DbError::InvalidOperation(format!(
252                    "source '{name}' primary-key column '{column}' is absent from its schema"
253                ))
254            })?;
255            if primary_key_indices.contains(&index) {
256                return Err(crate::DbError::InvalidOperation(format!(
257                    "source '{name}' primary key repeats column '{column}'"
258                )));
259            }
260            if schema.field(index).is_nullable() {
261                return Err(crate::DbError::InvalidOperation(format!(
262                    "source '{name}' primary-key column '{column}' must be non-nullable"
263                )));
264            }
265            primary_key_indices.push(index);
266        }
267
268        let buf_size = buffer_size.unwrap_or(self.default_buffer_size);
269        let bp = backpressure.unwrap_or(self.default_backpressure);
270
271        // Channel buffer is at least 1024 to avoid blocking on small snapshot rings.
272        let channel_buf = buf_size.max(1024);
273        let config = SourceConfig {
274            channel: streaming::ChannelConfig {
275                buffer_size: channel_buf,
276                backpressure: bp,
277                wait_strategy: WaitStrategy::SpinYield,
278                track_stats: false,
279            },
280            name: Some(name.to_string()),
281        };
282
283        let (source, sink) = streaming::create_with_config::<ArrowRecord>(config);
284
285        let entry = Arc::new(SourceEntry {
286            name: name.to_string(),
287            schema,
288            primary_key,
289            primary_key_indices,
290            watermark_column,
291            max_out_of_orderness,
292            is_processing_time: std::sync::atomic::AtomicBool::new(false),
293            source,
294            sink,
295            buffer: SnapshotRing::new(buf_size),
296            data_notify: Arc::new(Notify::new()),
297        });
298
299        sources.insert(name.to_string(), Arc::clone(&entry));
300        Ok(entry)
301    }
302
303    #[cfg(test)]
304    pub(crate) fn register_source_or_replace(
305        &self,
306        name: &str,
307        schema: SchemaRef,
308        primary_key: Vec<String>,
309        watermark_column: Option<String>,
310        max_out_of_orderness: Option<Duration>,
311        buffer_size: Option<usize>,
312        backpressure: Option<BackpressureStrategy>,
313    ) -> Arc<SourceEntry> {
314        // Remove existing if present
315        self.sources.write().remove(name);
316        // Safe to unwrap since we just removed any conflict
317        self.register_source(
318            name,
319            schema,
320            primary_key,
321            watermark_column,
322            max_out_of_orderness,
323            buffer_size,
324            backpressure,
325        )
326        .unwrap()
327    }
328
329    /// Look up a registered source by name.
330    pub fn get_source(&self, name: &str) -> Option<Arc<SourceEntry>> {
331        self.sources.read().get(name).cloned()
332    }
333
334    /// Returns `true` if the source existed.
335    pub fn drop_source(&self, name: &str) -> bool {
336        self.sources.write().remove(name).is_some()
337    }
338
339    pub(crate) fn register_sink(&self, name: &str, input: &str) -> Result<(), crate::DbError> {
340        let mut sinks = self.sinks.write();
341        if sinks.contains_key(name) {
342            return Err(crate::DbError::SinkAlreadyExists(name.to_string()));
343        }
344        sinks.insert(
345            name.to_string(),
346            SinkEntry {
347                input: input.to_string(),
348            },
349        );
350        Ok(())
351    }
352
353    /// Returns `true` if the sink existed.
354    pub fn drop_sink(&self, name: &str) -> bool {
355        self.sinks.write().remove(name).is_some()
356    }
357
358    pub(crate) fn register_stream(&self, name: &str) -> Result<(), crate::DbError> {
359        let mut streams = self.streams.write();
360        if streams.contains_key(name) {
361            return Err(crate::DbError::StreamAlreadyExists(name.to_string()));
362        }
363
364        streams.insert(
365            name.to_string(),
366            Arc::new(StreamEntry {
367                name: name.to_string(),
368                emitted_rows: AtomicU64::new(0),
369            }),
370        );
371        Ok(())
372    }
373
374    pub(crate) fn get_stream_entry(&self, name: &str) -> Option<Arc<StreamEntry>> {
375        self.streams.read().get(name).cloned()
376    }
377
378    /// Returns `true` if the stream existed.
379    pub fn drop_stream(&self, name: &str) -> bool {
380        self.streams.write().remove(name).is_some()
381    }
382
383    /// All registered stream names.
384    pub fn list_streams(&self) -> Vec<String> {
385        self.streams.read().keys().cloned().collect()
386    }
387
388    /// All registered source names.
389    pub fn list_sources(&self) -> Vec<String> {
390        self.sources.read().keys().cloned().collect()
391    }
392
393    /// All registered sink names.
394    pub fn list_sinks(&self) -> Vec<String> {
395        self.sinks.read().keys().cloned().collect()
396    }
397
398    /// Input source/table name for a sink, if registered.
399    pub fn get_sink_input(&self, name: &str) -> Option<String> {
400        self.sinks.read().get(name).map(|e| e.input.clone())
401    }
402
403    pub(crate) fn register_query(&self, sql: &str) -> u64 {
404        let id = self.next_query_id.fetch_add(1, Ordering::Relaxed);
405        let mut queries = self.queries.write();
406        queries.insert(
407            id,
408            QueryEntry {
409                id,
410                sql: sql.to_string(),
411                active: true,
412            },
413        );
414        id
415    }
416
417    pub(crate) fn deactivate_query(&self, id: u64) -> bool {
418        // Cap retained deactivated queries so finished SELECTs can't accumulate
419        // unboundedly; over the cap, the oldest (lowest id) is dropped.
420        const MAX_INACTIVE_QUERIES: usize = 100;
421        let mut queries = self.queries.write();
422        if let Some(entry) = queries.get_mut(&id) {
423            let was_active = entry.active;
424            entry.active = false;
425            if was_active {
426                let inactive_count = queries.values().filter(|q| !q.active).count();
427                if inactive_count > MAX_INACTIVE_QUERIES {
428                    let oldest_inactive_id =
429                        queries.values().filter(|q| !q.active).map(|q| q.id).min();
430                    if let Some(oldest_id) = oldest_inactive_id {
431                        queries.remove(&oldest_id);
432                    }
433                }
434            }
435            true
436        } else {
437            false
438        }
439    }
440
441    pub(crate) fn list_queries(&self) -> Vec<(u64, String, bool)> {
442        self.queries
443            .read()
444            .values()
445            .map(|q| (q.id, q.sql.clone(), q.active))
446            .collect()
447    }
448
449    /// Schema for DESCRIBE queries.
450    pub fn describe_source(&self, name: &str) -> Option<SchemaRef> {
451        self.sources.read().get(name).map(|e| e.schema.clone())
452    }
453}
454
455#[cfg(test)]
456mod tests;
laminar_db/catalog/mod.rs

laminar_db/catalog/
mod.rs
