Skip to main content

nodedb_lite/query/
engine.rs

1//! Lite query engine: SQL via nodedb-sql over local engines.
2//!
3//! Parses SQL with nodedb-sql, then executes against CRDT, strict,
4//! and columnar engines directly — no DataFusion dependency.
5
6use std::collections::HashMap;
7use std::sync::{Arc, Mutex};
8
9use nodedb_mem::MemoryGovernor;
10use nodedb_sql::types::*;
11use nodedb_types::result::QueryResult;
12use nodedb_types::value::Value;
13
14use crate::engine::columnar::ColumnarEngine;
15use crate::engine::crdt::CrdtEngine;
16use crate::engine::fts::FtsState;
17use crate::engine::graph::index::CsrIndex;
18use crate::engine::htap::HtapBridge;
19use crate::engine::sparse_vector::SparseVectorState;
20use crate::engine::spatial::SpatialIndexManager;
21use crate::engine::strict::StrictEngine;
22use crate::engine::vector::VectorState;
23use crate::error::LiteError;
24use crate::nodedb::{KvLocalState, LockExt};
25use crate::sequence::LiteSequenceRegistry;
26use crate::storage::engine::StorageEngine;
27
28use super::catalog::LiteCatalog;
29use super::meta_ops::CancellationRegistry;
30
31/// Lite-side query engine.
32pub struct LiteQueryEngine<S: StorageEngine> {
33    pub(in crate::query) crdt: Arc<Mutex<CrdtEngine>>,
34    pub(in crate::query) strict: Arc<StrictEngine<S>>,
35    pub(in crate::query) columnar: Arc<ColumnarEngine<S>>,
36    pub(in crate::query) htap: Arc<HtapBridge>,
37    pub(in crate::query) storage: Arc<S>,
38    pub(in crate::query) timeseries:
39        Arc<Mutex<crate::engine::timeseries::engine::TimeseriesEngine>>,
40    pub(crate) vector_state: Arc<VectorState<S>>,
41    pub(crate) array_state: Arc<tokio::sync::Mutex<crate::engine::array::engine::ArrayEngineState>>,
42    pub(crate) fts_state: Arc<FtsState>,
43    /// Sparse-vector inverted index state, shared with the owning `NodeDbLite`.
44    pub(crate) sparse_state: Arc<SparseVectorState>,
45    pub(in crate::query) spatial: Arc<Mutex<SpatialIndexManager>>,
46    pub(crate) cancellation: CancellationRegistry,
47    /// Per-collection CSR graph indices shared with the owning NodeDbLite.
48    pub(crate) csr: Arc<Mutex<HashMap<String, CsrIndex>>>,
49    /// Memory budget governor, shared with the owning NodeDbLite.
50    pub(crate) governor: Arc<MemoryGovernor>,
51    /// Sequence registry backing `nextval` / `currval` / `setval` in a
52    /// SELECT list. Definitions are registered through
53    /// `LiteQueryEngine::sequences`; counters live in memory.
54    pub(crate) sequences: Arc<LiteSequenceRegistry>,
55    /// The public KV API's write buffer and read cache. A SQL-path `TRUNCATE`
56    /// forgets what they hold for the cleared collection.
57    pub(in crate::query) kv_local: Arc<KvLocalState>,
58    /// Secondary-index definitions and entries. The CRDT store maintains them
59    /// on every row write; flush persists them with the CRDT state.
60    pub(crate) indexes: Arc<crate::index::IndexCatalog>,
61    /// Durable outbound queue for FTS sync — `None` when sync is disabled.
62    #[cfg(not(target_arch = "wasm32"))]
63    pub(crate) fts_outbound: Option<Arc<crate::sync::FtsOutbound<S>>>,
64    /// Durable outbound queue for spatial sync — `None` when sync is disabled.
65    #[cfg(not(target_arch = "wasm32"))]
66    pub(crate) spatial_outbound: Option<Arc<crate::sync::SpatialOutbound<S>>>,
67    /// Durable outbound queue for KV write sync — `None` when sync is disabled.
68    #[cfg(not(target_arch = "wasm32"))]
69    pub(crate) kv_outbound: Option<Arc<crate::sync::KvOutbound<S>>>,
70}
71
72/// Construction fields for [`LiteQueryEngine::new`].
73pub struct LiteQueryEngineParams<S: StorageEngine> {
74    pub crdt: Arc<Mutex<CrdtEngine>>,
75    pub strict: Arc<StrictEngine<S>>,
76    pub columnar: Arc<ColumnarEngine<S>>,
77    pub htap: Arc<HtapBridge>,
78    pub storage: Arc<S>,
79    pub timeseries: Arc<Mutex<crate::engine::timeseries::engine::TimeseriesEngine>>,
80    pub vector_state: Arc<VectorState<S>>,
81    pub array_state: Arc<tokio::sync::Mutex<crate::engine::array::engine::ArrayEngineState>>,
82    pub fts_state: Arc<FtsState>,
83    pub sparse_state: Arc<SparseVectorState>,
84    pub spatial: Arc<Mutex<SpatialIndexManager>>,
85    pub csr: Arc<Mutex<HashMap<String, CsrIndex>>>,
86    pub governor: Arc<MemoryGovernor>,
87    pub kv_local: Arc<KvLocalState>,
88}
89
90impl<S: StorageEngine> LiteQueryEngine<S> {
91    pub fn new(params: LiteQueryEngineParams<S>) -> Self {
92        let indexes = Arc::new(crate::index::IndexCatalog::new());
93        params
94            .crdt
95            .lock_or_recover()
96            .set_index_catalog(Arc::clone(&indexes));
97        params.strict.set_index_catalog(Arc::clone(&indexes));
98        params.kv_local.set_index_catalog(Arc::clone(&indexes));
99        Self {
100            crdt: params.crdt,
101            strict: params.strict,
102            columnar: params.columnar,
103            htap: params.htap,
104            storage: params.storage,
105            timeseries: params.timeseries,
106            vector_state: params.vector_state,
107            array_state: params.array_state,
108            fts_state: params.fts_state,
109            sparse_state: params.sparse_state,
110            spatial: params.spatial,
111            cancellation: CancellationRegistry::new(),
112            csr: params.csr,
113            governor: params.governor,
114            sequences: Arc::new(LiteSequenceRegistry::new()),
115            kv_local: params.kv_local,
116            indexes,
117            #[cfg(not(target_arch = "wasm32"))]
118            fts_outbound: None,
119            #[cfg(not(target_arch = "wasm32"))]
120            spatial_outbound: None,
121            #[cfg(not(target_arch = "wasm32"))]
122            kv_outbound: None,
123        }
124    }
125
126    /// Wire the durable FTS outbound queue so SQL-path spatial writes are sync-tracked.
127    #[cfg(not(target_arch = "wasm32"))]
128    pub fn set_fts_outbound(&mut self, q: Arc<crate::sync::FtsOutbound<S>>) {
129        self.fts_outbound = Some(q);
130    }
131
132    /// Wire the durable spatial outbound queue so SQL-path spatial writes are sync-tracked.
133    #[cfg(not(target_arch = "wasm32"))]
134    pub fn set_spatial_outbound(&mut self, q: Arc<crate::sync::SpatialOutbound<S>>) {
135        self.spatial_outbound = Some(q);
136    }
137
138    /// Wire the durable KV outbound queue so SQL-path KV writes are sync-tracked.
139    #[cfg(not(target_arch = "wasm32"))]
140    pub fn set_kv_outbound(&mut self, q: Arc<crate::sync::KvOutbound<S>>) {
141        self.kv_outbound = Some(q);
142    }
143
144    /// The sequence registry SELECT-list accessors read and advance.
145    pub fn sequences(&self) -> &Arc<LiteSequenceRegistry> {
146        &self.sequences
147    }
148
149    /// No-op — collections are auto-discovered via catalog.
150    pub fn register_collection(&self, _name: &str) {}
151    /// No-op — collections are auto-discovered via catalog.
152    pub fn register_strict_collection(&self, _name: &str) {}
153    /// No-op — collections are auto-discovered via catalog.
154    pub fn register_all_collections(&self) {}
155    /// No-op — collections are auto-discovered via catalog.
156    pub fn register_columnar_collection(&self, _name: &str) {}
157
158    /// Execute a SQL query and return results.
159    pub async fn execute_sql(&self, sql: &str) -> Result<QueryResult, LiteError> {
160        self.execute_sql_with_params(sql, &[]).await
161    }
162
163    /// Execute a SQL query with bound `$N` parameters and return results.
164    ///
165    /// Each `Value` in `params` is bound to the corresponding `$1`, `$2`, …
166    /// placeholder in `sql` at the AST level before planning. Supported
167    /// `Value` variants: `Null`, `Bool`, `Integer`, `Float`, `String`, `Uuid`.
168    /// Other variants are treated as `Null`.
169    pub async fn execute_sql_with_params(
170        &self,
171        sql: &str,
172        params: &[Value],
173    ) -> Result<QueryResult, LiteError> {
174        if let Some(result) = self.try_handle_ddl(sql).await {
175            return result;
176        }
177        super::returning::refuse_unflagged_returning(sql)?;
178
179        let metas =
180            crate::nodedb::collection::ddl::load_persisted_collection_metas(self.storage.as_ref())
181                .await
182                .unwrap_or_default();
183        let array_names: Vec<String> = self
184            .array_state
185            .lock()
186            .await
187            .arrays
188            .keys()
189            .cloned()
190            .collect();
191        let catalog = LiteCatalog::new(
192            Arc::clone(&self.crdt),
193            Arc::clone(&self.strict),
194            Arc::clone(&self.columnar),
195            metas,
196        )
197        .with_arrays(array_names)
198        .with_indexes(Arc::clone(&self.indexes));
199
200        let sql_params: Vec<nodedb_sql::ParamValue> = params.iter().map(value_to_param).collect();
201
202        let plans = if sql_params.is_empty() {
203            nodedb_sql::plan_sql(sql, &catalog)
204        } else {
205            nodedb_sql::plan_sql_with_params(sql, &sql_params, &catalog)
206        }
207        .map_err(LiteError::from)?;
208
209        // One statement can plan to several units: `TRUNCATE a, b` yields one
210        // plan per collection. Every unit runs; the last result is the answer.
211        let mut result = QueryResult::empty();
212        for plan in &plans {
213            result = self.execute_plan(plan).await?;
214        }
215        Ok(result)
216    }
217
218    pub(in crate::query) async fn execute_plan(
219        &self,
220        plan: &SqlPlan,
221    ) -> Result<QueryResult, LiteError> {
222        if plan_mutates_source(plan) {
223            let guard = self.fts_state.admit_mutation().await;
224            let result = self.execute_plan_admitted(plan, Some(guard.permit())).await;
225            guard.finish(result)
226        } else {
227            self.execute_plan_admitted(plan, None).await
228        }
229    }
230
231    pub(in crate::query) async fn execute_plan_admitted(
232        &self,
233        plan: &SqlPlan,
234        permit: Option<&crate::engine::fts::coordinator::TextMutationPermit>,
235    ) -> Result<QueryResult, LiteError> {
236        let mut visitor = super::visitor::LiteVisitor {
237            engine: self,
238            permit,
239        };
240        let mut result = nodedb_sql::dispatch(&mut visitor, plan)?.await?;
241        // The shared plan dispatcher hands the visitor a point get without its
242        // SELECT list, so the list is applied here, as the scan path applies
243        // it to a scan.
244        if let SqlPlan::PointGet { projection, .. } = plan {
245            super::visitor::project_scan_result(&mut result, projection, &[], &self.sequences)?;
246        }
247        Ok(result)
248    }
249
250    pub(super) async fn execute_constant_result(
251        &self,
252        columns: &[String],
253        values: &[nodedb_sql::types::SqlValue],
254    ) -> Result<QueryResult, LiteError> {
255        let row = values.iter().map(sql_value_to_value).collect();
256        Ok(QueryResult {
257            columns: columns.to_vec(),
258            rows: vec![row],
259            rows_affected: 0,
260            command: None,
261        })
262    }
263}
264
265pub(super) fn sql_value_to_string(v: &SqlValue) -> String {
266    match v {
267        SqlValue::String(s) => s.clone(),
268        SqlValue::Int(i) => i.to_string(),
269        SqlValue::Float(f) => f.to_string(),
270        SqlValue::Bool(b) => b.to_string(),
271        _ => String::new(),
272    }
273}
274
275pub(super) fn sql_value_to_value(v: &nodedb_sql::types::SqlValue) -> Value {
276    match v {
277        nodedb_sql::types::SqlValue::Int(i) => Value::Integer(*i),
278        nodedb_sql::types::SqlValue::Float(f) => Value::Float(*f),
279        nodedb_sql::types::SqlValue::String(s) => Value::String(s.clone()),
280        nodedb_sql::types::SqlValue::Bool(b) => Value::Bool(*b),
281        nodedb_sql::types::SqlValue::Null => Value::Null,
282        _ => Value::Null,
283    }
284}
285
286/// Convert a `nodedb_types::Value` to the `nodedb_sql::ParamValue` type used
287/// for AST-level parameter binding in `plan_sql_with_params`.
288fn value_to_param(v: &Value) -> nodedb_sql::ParamValue {
289    match v {
290        Value::Null => nodedb_sql::ParamValue::Null,
291        Value::Bool(b) => nodedb_sql::ParamValue::Bool(*b),
292        Value::Integer(n) => nodedb_sql::ParamValue::Int64(*n),
293        Value::Float(f) => nodedb_sql::ParamValue::Float64(*f),
294        Value::String(s) => nodedb_sql::ParamValue::Text(s.clone()),
295        Value::Uuid(s) => nodedb_sql::ParamValue::Text(s.clone()),
296        _ => nodedb_sql::ParamValue::Null,
297    }
298}
299
300/// Build a real governor for tests, with a ceiling covering every engine's limit.
301///
302/// `nodedb_mem` requires `global_ceiling >= sum(engine_limits)`, so the
303/// per-engine limit is `usize::MAX / EngineId::ALL.len()` to avoid overflow.
304#[cfg(test)]
305pub(crate) fn test_governor() -> Arc<nodedb_mem::MemoryGovernor> {
306    let per_engine = usize::MAX / nodedb_mem::EngineId::ALL.len();
307    Arc::new(
308        nodedb_mem::MemoryGovernor::new(nodedb_mem::GovernorConfig {
309            global_ceiling: per_engine * nodedb_mem::EngineId::ALL.len(),
310            engine_limits: nodedb_mem::EngineLimits::uniform(per_engine),
311        })
312        .expect("test governor"),
313    )
314}
315
316/// Build a `ScopedMemory` handle from a test governor for `engine`.
317#[cfg(test)]
318pub(crate) fn test_scoped_memory(
319    governor: &Arc<nodedb_mem::MemoryGovernor>,
320    engine: nodedb_mem::EngineId,
321) -> nodedb_mem::ScopedMemory {
322    nodedb_mem::ScopedMemory::new(
323        Arc::clone(governor),
324        nodedb_types::DatabaseId::DEFAULT,
325        nodedb_types::TenantId::new(0),
326        engine,
327    )
328}
329
330/// Build an in-memory `LiteQueryEngine` with every engine state wired up.
331#[cfg(test)]
332pub(crate) async fn test_engine() -> LiteQueryEngine<crate::PagedbStorageMem> {
333    use crate::engine::array::engine::ArrayEngineState;
334    use crate::engine::spatial::SpatialIndexManager;
335
336    let storage = Arc::new(
337        crate::PagedbStorageMem::open_in_memory()
338            .await
339            .expect("in-memory pagedb"),
340    );
341    let crdt = Arc::new(Mutex::new(CrdtEngine::new(1).expect("crdt")));
342    let governor = test_governor();
343    let strict = Arc::new(StrictEngine::new(Arc::clone(&storage)));
344    let columnar = Arc::new(ColumnarEngine::new(
345        Arc::clone(&storage),
346        test_scoped_memory(&governor, nodedb_mem::EngineId::Columnar),
347    ));
348    let htap = Arc::new(HtapBridge::new());
349    let timeseries = Arc::new(Mutex::new(
350        crate::engine::timeseries::engine::TimeseriesEngine::new(),
351    ));
352    let vector_state = Arc::new(VectorState::new(
353        Arc::clone(&storage),
354        100,
355        test_scoped_memory(&governor, nodedb_mem::EngineId::Vector),
356    ));
357    let array_state = Arc::new(tokio::sync::Mutex::new(
358        ArrayEngineState::open(&storage).await.expect("array"),
359    ));
360    let fts_state = Arc::new(FtsState::new(Arc::clone(&governor)));
361    let spatial = Arc::new(Mutex::new(SpatialIndexManager::new(test_scoped_memory(
362        &governor,
363        nodedb_mem::EngineId::Spatial,
364    ))));
365    LiteQueryEngine::new(LiteQueryEngineParams {
366        crdt,
367        strict,
368        columnar,
369        htap,
370        storage,
371        timeseries,
372        vector_state,
373        array_state,
374        fts_state,
375        sparse_state: Arc::new(SparseVectorState::new()),
376        spatial,
377        csr: Arc::new(Mutex::new(HashMap::new())),
378        governor,
379        kv_local: test_kv_local(),
380    })
381}
382
383/// A KV write buffer and cache for engines built outside `NodeDbLite`.
384#[cfg(test)]
385pub(crate) fn test_kv_local() -> Arc<KvLocalState> {
386    Arc::new(KvLocalState::new(
387        std::num::NonZeroUsize::new(64).expect("non-zero"),
388    ))
389}
390
391fn plan_mutates_source(plan: &SqlPlan) -> bool {
392    match plan {
393        SqlPlan::Insert(_)
394        | SqlPlan::KvInsert(_)
395        | SqlPlan::Upsert(_)
396        | SqlPlan::InsertSelect { .. }
397        | SqlPlan::Update { .. }
398        | SqlPlan::UpdateFrom { .. }
399        | SqlPlan::Delete { .. }
400        | SqlPlan::Truncate { .. }
401        | SqlPlan::Merge(_)
402        | SqlPlan::TimeseriesIngest(_)
403        | SqlPlan::VectorPrimaryInsert(_)
404        | SqlPlan::VectorPrimaryDelete(_)
405        | SqlPlan::VectorPrimaryUpdate(_)
406        | SqlPlan::VectorPrimaryTruncate(_)
407        | SqlPlan::CreateArray(_)
408        | SqlPlan::DropArray { .. }
409        | SqlPlan::AlterArray(_)
410        | SqlPlan::InsertArray(_)
411        | SqlPlan::DeleteArray(_)
412        | SqlPlan::ArrayFlush { .. }
413        | SqlPlan::ArrayCompact { .. } => true,
414        SqlPlan::Join { left, right, .. }
415        | SqlPlan::Intersect { left, right, .. }
416        | SqlPlan::Except { left, right, .. } => {
417            plan_mutates_source(left) || plan_mutates_source(right)
418        }
419        SqlPlan::Aggregate { input, .. } | SqlPlan::Subquery { input, .. } => {
420            plan_mutates_source(input)
421        }
422        SqlPlan::Union { inputs, .. } => inputs.iter().any(plan_mutates_source),
423        SqlPlan::LateralTopK(plan) => plan_mutates_source(&plan.outer),
424        SqlPlan::LateralLoop(plan) => {
425            plan_mutates_source(&plan.outer) || plan_mutates_source(&plan.inner)
426        }
427        SqlPlan::Cte(plan) => {
428            plan.definitions
429                .iter()
430                .any(|(_, plan)| plan_mutates_source(plan))
431                || plan_mutates_source(&plan.outer)
432        }
433        _ => false,
434    }
435}