laminar_db/db/
read_level.rs

1//! Ad-hoc queries that read materialized views at a level the session chose: after the latest
2//! completed cycle, or at one committed checkpoint, possibly the first to hold what the views'
3//! sources had upstream when the read arrived.
4
5use std::collections::{BTreeMap, HashSet};
6use std::time::Instant;
7
8use arrow::datatypes::SchemaRef;
9use laminar_connectors::config::ConnectorConfig;
10use laminar_sql::parser::{parse_streaming_sql, StreamingStatement};
11
12use super::LaminarDB;
13use crate::error::DbError;
14use crate::mv_store::{
15    CommittedRead, CommittedReadError, CommittedViews, PinnedCut, ReadCut, RecencyTargets,
16    VIEW_READ_DEADLINE,
17};
18use crate::query_rows::{execution_error, QueryRows};
19
20/// How an ad-hoc query reads the materialized views it scans.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum ReadLevel {
23    /// Every view after the latest completed cycle, read under one guard. Recovery may roll the
24    /// result back. Embedded and single-node only: a cluster node holds only its part of each
25    /// view.
26    Latest,
27    /// Every view at the one committed checkpoint the argument resolves to; see
28    /// [`LaminarDB::query_committed`].
29    Committed(CommittedRead),
30}
31
32/// The rows of an ad-hoc read and the committed cut it observed.
33#[derive(Debug)]
34pub struct ReadResult {
35    /// The cut every view was read at; `None` for a [`ReadLevel::Latest`] read.
36    pub cut: Option<ReadCut>,
37    /// The result rows.
38    pub rows: QueryRows,
39}
40
41impl LaminarDB {
42    /// The level a read uses when its session chose none: [`ReadLevel::Latest`] in embedded and
43    /// single-node mode, the newest committed checkpoint on a cluster node.
44    #[must_use]
45    pub fn default_read_level(&self) -> ReadLevel {
46        if self.is_cluster_runtime() {
47            ReadLevel::Committed(CommittedRead::Latest)
48        } else {
49            ReadLevel::Latest
50        }
51    }
52
53    /// Run the query `sql` with every materialized view it scans read at `level`.
54    ///
55    /// Each table is read as a snapshot, so joins between views, sources and reference tables are
56    /// bounded. A committed read loads its views before this returns and holds its checkpoint
57    /// against retention until then; the rows then stream to the reader as it reads them.
58    ///
59    /// # Errors
60    /// [`DbError::InvalidOperation`] if `sql` is not one query or reads a view at
61    /// [`ReadLevel::Latest`] on a cluster node; the errors of [`Self::query_committed`] for a
62    /// committed read; any planning or execution error of `sql`.
63    pub async fn read(&self, sql: &str, level: ReadLevel) -> Result<ReadResult, DbError> {
64        if self.shutdown.load(std::sync::atomic::Ordering::Relaxed) {
65            return Err(DbError::Shutdown);
66        }
67        ensure_query(sql)?;
68        let result = match level {
69            ReadLevel::Latest => {
70                let frame = crate::mv_store::plan_statement(&self.ctx, &self.mv_store, sql).await?;
71                let stream = frame.execute_stream().await.map_err(execution_error)?;
72                if let Some(metrics) = self.engine_metrics.lock().as_ref() {
73                    metrics.views.observe_read("latest", None);
74                }
75                ReadResult {
76                    cut: None,
77                    rows: self.spawn_rows(sql, stream),
78                }
79            }
80            ReadLevel::Committed(at) => {
81                let (cut, stream) = self
82                    .committed_stream(sql, at, Instant::now() + VIEW_READ_DEADLINE)
83                    .await?;
84                ReadResult {
85                    cut: Some(cut),
86                    rows: self.spawn_rows(sql, stream),
87                }
88            }
89        };
90        Ok(result)
91    }
92
93    /// The schema of the rows [`Self::read`] returns for `sql` and the parameter count it
94    /// declares, planned without reading any table.
95    ///
96    /// # Errors
97    /// [`DbError::InvalidOperation`] if `sql` is not one query; any planning error of `sql`.
98    pub async fn read_schema(&self, sql: &str) -> Result<(SchemaRef, usize), DbError> {
99        ensure_query(sql)?;
100        let plan = self.ctx.state().create_logical_plan(sql).await?;
101        let parameters = plan.get_parameter_types()?.len();
102        Ok((
103            std::sync::Arc::new(plan.schema().as_arrow().clone()),
104            parameters,
105        ))
106    }
107
108    /// Plan `sql` at the committed cut `at` resolves to and load its views.
109    pub(crate) async fn committed_stream(
110        &self,
111        sql: &str,
112        at: CommittedRead,
113        deadline: Instant,
114    ) -> Result<
115        (
116            ReadCut,
117            datafusion::physical_plan::SendableRecordBatchStream,
118        ),
119        DbError,
120    > {
121        let planned = crate::mv_store::plan_committed_statement(&self.ctx, sql).await?;
122        let views = self
123            .mv_store
124            .read()
125            .committed_views()
126            .ok_or(CommittedReadError::CheckpointingDisabled)?;
127        let (level, cut) = match at {
128            CommittedRead::Recent { at_least } => (
129                "recent",
130                self.pin_recent(&views, planned.views(), at_least, deadline)
131                    .await?,
132            ),
133            at => ("committed", views.pin(at, deadline).await?),
134        };
135        let (frame, cut) = planned.at_cut(&self.mv_store, views, cut);
136        // View scans load their artifacts while the plan is built.
137        let stream = tokio::time::timeout_at(
138            tokio::time::Instant::from_std(deadline),
139            frame.execute_stream(),
140        )
141        .await
142        .map_err(|_| CommittedReadError::DeadlineExceeded { epoch: cut.epoch })?
143        .map_err(execution_error)?;
144        if let Some(metrics) = self.engine_metrics.lock().as_ref() {
145            metrics.views.observe_read(level, Some(cut.taken_at_ms));
146        }
147        Ok((cut, stream))
148    }
149
150    /// Pin the newest committed cut at or after `at_least` that holds what every source behind
151    /// `scanned` had upstream when this was called.
152    async fn pin_recent(
153        &self,
154        views: &CommittedViews,
155        scanned: &[String],
156        at_least: u64,
157        deadline: Instant,
158    ) -> Result<PinnedCut, DbError> {
159        let arrived = Instant::now();
160        let pinned = match self.recency_targets(scanned, deadline).await {
161            Ok(targets) => {
162                let expedite = |epoch| self.expedite_checkpoint(epoch);
163                views
164                    .pin_covering(&targets, at_least, deadline, &expedite)
165                    .await
166            }
167            Err(error) => Err(error),
168        };
169        if let Some(metrics) = self.engine_metrics.lock().as_ref() {
170            match &pinned {
171                Ok(_) => metrics
172                    .views
173                    .recent_read_wait_seconds
174                    .observe(arrived.elapsed().as_secs_f64()),
175                Err(DbError::CommittedRead(CommittedReadError::RecencyProbe { input, .. })) => {
176                    metrics
177                        .views
178                        .recency_probe_failures_total
179                        .with_label_values(&[input])
180                        .inc();
181                }
182                Err(_) => {}
183            }
184        }
185        pinned
186    }
187
188    /// Each source behind `views`, with its upstream's end now.
189    async fn recency_targets(
190        &self,
191        views: &[String],
192        deadline: Instant,
193    ) -> Result<RecencyTargets, DbError> {
194        let probe_failed = |source: &str, reason: String| CommittedReadError::RecencyProbe {
195            input: source.to_owned(),
196            reason,
197        };
198        let mut probes = Vec::new();
199        for (source, config) in self.view_lineage(views)? {
200            let probe = self
201                .connector_registry
202                .source_recency(&config)
203                .map_err(|error| probe_failed(&source, error.to_string()))?
204                .ok_or_else(|| CommittedReadError::RecencyUnsupported {
205                    input: source.clone(),
206                })?;
207            probes.push((source, probe));
208        }
209        let ends = probes.into_iter().map(|(source, probe)| async move {
210            let end = tokio::time::timeout_at(
211                tokio::time::Instant::from_std(deadline),
212                probe.upstream_end(deadline),
213            )
214            .await;
215            match end {
216                Ok(Ok(end)) => Ok((source, end)),
217                Ok(Err(error)) => Err(probe_failed(&source, error.to_string())),
218                Err(_) => Err(probe_failed(
219                    &source,
220                    "no answer before the read's deadline".into(),
221                )),
222            }
223        });
224        Ok(RecencyTargets::new(
225            futures::future::try_join_all(ends).await?,
226        ))
227    }
228
229    /// The connector configuration of every source whose records reach `views`, through the
230    /// streams, stages and views between them. A table is looked up rather than read through an
231    /// input, so it adds no source.
232    fn view_lineage(&self, views: &[String]) -> Result<BTreeMap<String, ConnectorConfig>, DbError> {
233        use laminar_core::catalog::CatalogObjectKind;
234        let unsupported = |input: String| CommittedReadError::RecencyUnsupported { input };
235        let kinds = self.catalog_namespace.lock().clone();
236        let mut sources = BTreeMap::new();
237        {
238            let manager = self.connector_manager.lock();
239            let mut pending = views.to_vec();
240            let mut visited = HashSet::new();
241            // Each catalog name is expanded once.
242            while let Some(name) = pending.pop() {
243                if !visited.insert(name.clone()) {
244                    continue;
245                }
246                match kinds.get(&name) {
247                    Some(CatalogObjectKind::Stream | CatalogObjectKind::MaterializedView) => {
248                        let stream = manager.streams().get(&name);
249                        let stream = stream.ok_or_else(|| unsupported(name.clone()))?;
250                        pending.extend(crate::sql_analysis::extract_table_references(
251                            &stream.query_sql,
252                        ));
253                    }
254                    Some(CatalogObjectKind::Source) => {
255                        let source = manager.sources().get(&name);
256                        let source = source.filter(|source| source.connector_type.is_some());
257                        let source = source.ok_or_else(|| unsupported(name.clone()))?.clone();
258                        sources.insert(name, source);
259                    }
260                    Some(CatalogObjectKind::Table | CatalogObjectKind::LookupTable) => {}
261                    Some(CatalogObjectKind::Sink) | None => return Err(unsupported(name).into()),
262                }
263            }
264        }
265        sources
266            .into_iter()
267            .map(|(name, registration)| {
268                let config = self.build_registered_source_config(&name, &registration)?;
269                Ok((name, config))
270            })
271            .collect()
272    }
273
274    fn spawn_rows(
275        &self,
276        sql: &str,
277        stream: datafusion::physical_plan::SendableRecordBatchStream,
278    ) -> QueryRows {
279        let query_id = self.catalog.register_query(sql);
280        let catalog = std::sync::Arc::clone(&self.catalog);
281        QueryRows::spawn(
282            stream,
283            tokio_util::sync::CancellationToken::new(),
284            move || {
285                catalog.deactivate_query(query_id);
286            },
287        )
288    }
289}
290
291/// A read runs one query; DDL, DML and session statements go through `execute`.
292fn ensure_query(sql: &str) -> Result<(), DbError> {
293    let statements = parse_streaming_sql(sql)?;
294    match statements.as_slice() {
295        [StreamingStatement::Standard(statement)]
296            if matches!(**statement, sqlparser::ast::Statement::Query(_)) =>
297        {
298            Ok(())
299        }
300        [StreamingStatement::TemporalProbeQuery { .. }] => Err(DbError::Unsupported(
301            "temporal queries require the managed vnode runtime and cannot execute through \
302             DataFusion's ordinary join path"
303                .into(),
304        )),
305        _ => Err(DbError::InvalidOperation(
306            "a read must be exactly one query".into(),
307        )),
308    }
309}
310
311#[cfg(all(test, feature = "kafka"))]
312mod tests {
313    use super::*;
314    use crate::config::LaminarConfig;
315
316    #[tokio::test]
317    async fn lineage_follows_streams_stages_and_views_to_their_sources() {
318        let db = LaminarDB::open_with_config(LaminarConfig::default()).unwrap();
319        for sql in [
320            "CREATE SOURCE orders (k BIGINT, v BIGINT) FROM KAFKA (\
321             'bootstrap.servers' = '127.0.0.1:1', 'topic' = 'orders', 'group.id' = 'lineage') \
322             FORMAT JSON",
323            "CREATE SOURCE pushed (k BIGINT)",
324            "CREATE TABLE dims (k BIGINT PRIMARY KEY, name VARCHAR)",
325            "CREATE STREAM doubled AS SELECT k, v * 2 AS v FROM orders",
326            "CREATE MATERIALIZED VIEW sums AS SELECT k, SUM(v) AS total FROM doubled GROUP BY k",
327            "CREATE MATERIALIZED VIEW big AS SELECT k FROM sums WHERE total > 10",
328            "CREATE MATERIALIZED VIEW spread AS SELECT n, COUNT(*) AS keys \
329             FROM (SELECT k, COUNT(*) AS n FROM orders GROUP BY k) GROUP BY n",
330            "CREATE MATERIALIZED VIEW named AS SELECT s.k, s.total, d.name FROM sums s \
331             JOIN dims d ON s.k = d.k",
332            "CREATE MATERIALIZED VIEW pushed_counts AS SELECT k, COUNT(*) AS n FROM pushed GROUP BY k",
333        ] {
334            db.execute(sql)
335                .await
336                .unwrap_or_else(|error| panic!("{sql}: {error}"));
337        }
338        let lineage = |views: &[&str]| {
339            let views: Vec<String> = views.iter().map(|view| (*view).to_owned()).collect();
340            db.view_lineage(&views)
341                .map(|sources| sources.into_keys().collect::<Vec<_>>())
342        };
343        assert_eq!(lineage(&["big"]).unwrap(), ["orders"]);
344        assert_eq!(lineage(&["spread"]).unwrap(), ["orders"]);
345        assert_eq!(lineage(&["named"]).unwrap(), ["orders"]);
346        assert!(lineage(&[]).unwrap().is_empty());
347        let refused = lineage(&["big", "pushed_counts"]);
348        assert!(
349            matches!(
350                &refused,
351                Err(DbError::CommittedRead(CommittedReadError::RecencyUnsupported { input }))
352                    if input == "pushed"
353            ),
354            "{refused:?}"
355        );
356    }
357}
laminar_db/db/read_level.rs

laminar_db/db/
read_level.rs
