1use 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
31pub 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 pub(crate) sparse_state: Arc<SparseVectorState>,
45 pub(in crate::query) spatial: Arc<Mutex<SpatialIndexManager>>,
46 pub(crate) cancellation: CancellationRegistry,
47 pub(crate) csr: Arc<Mutex<HashMap<String, CsrIndex>>>,
49 pub(crate) governor: Arc<MemoryGovernor>,
51 pub(crate) sequences: Arc<LiteSequenceRegistry>,
55 pub(in crate::query) kv_local: Arc<KvLocalState>,
58 pub(crate) indexes: Arc<crate::index::IndexCatalog>,
61 #[cfg(not(target_arch = "wasm32"))]
63 pub(crate) fts_outbound: Option<Arc<crate::sync::FtsOutbound<S>>>,
64 #[cfg(not(target_arch = "wasm32"))]
66 pub(crate) spatial_outbound: Option<Arc<crate::sync::SpatialOutbound<S>>>,
67 #[cfg(not(target_arch = "wasm32"))]
69 pub(crate) kv_outbound: Option<Arc<crate::sync::KvOutbound<S>>>,
70}
71
72pub 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 #[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 #[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 #[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 pub fn sequences(&self) -> &Arc<LiteSequenceRegistry> {
146 &self.sequences
147 }
148
149 pub fn register_collection(&self, _name: &str) {}
151 pub fn register_strict_collection(&self, _name: &str) {}
153 pub fn register_all_collections(&self) {}
155 pub fn register_columnar_collection(&self, _name: &str) {}
157
158 pub async fn execute_sql(&self, sql: &str) -> Result<QueryResult, LiteError> {
160 self.execute_sql_with_params(sql, &[]).await
161 }
162
163 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 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 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
286fn 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#[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#[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#[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#[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}