1pub mod arrow_conv;
25mod commands;
26mod error;
27pub mod extended_sql_functions;
28mod external_modules;
29mod fk_join;
30mod native_agg;
31mod percentile;
32mod scan;
33mod shadow;
34mod udf;
35
36pub use error::{MongrelQueryError, Result};
37pub use external_modules::{
38 ExternalBaseWrite, ExternalModuleDescriptor, ExternalModuleIndex, ExternalModuleRegistry,
39 ExternalPlan, ExternalPlanRequest, ExternalScan, ExternalTable, ExternalTableModule,
40 ExternalTxn, ExternalWriteOp, ExternalWriteResult, ModuleConnectCtx,
41};
42
43use arrow::array::{Array, ArrayRef, Int64Array, StringArray};
44use arrow::datatypes::SchemaRef;
45use arrow::record_batch::RecordBatch;
46use datafusion::catalog::{Session, TableProvider};
47use datafusion::common::{DataFusionError, Result as DFResult};
48use datafusion::logical_expr::{AggregateUDF, Expr, ScalarUDF, TableType, WindowUDF};
49use datafusion::physical_plan::ExecutionPlan;
50use datafusion::prelude::SessionContext;
51use mongreldb_core::{
52 AlterColumn, ColumnFlags, Cursor, Database, Schema as CoreSchema, Table, TypeId,
53};
54use parking_lot::Mutex;
55use std::collections::{HashMap, HashSet};
56use std::sync::Arc;
57
58pub struct MongrelProvider {
61 db: Arc<Mutex<Table>>,
62 schema: SchemaRef,
63}
64
65#[derive(Debug, Clone)]
66pub(crate) struct ViewDef {
67 pub sql: String,
68 pub schema: CoreSchema,
69 pub input_types: HashMap<u16, Option<TypeId>>,
70}
71
72impl MongrelProvider {
73 pub fn new(db: Arc<Mutex<Table>>) -> Result<Self> {
74 let schema = {
75 let db = db.lock();
76 arrow_conv::arrow_schema(db.schema())?
77 };
78 Ok(Self { db, schema })
79 }
80
81 pub fn arrow_schema(&self) -> SchemaRef {
82 self.schema.clone()
83 }
84}
85
86impl std::fmt::Debug for MongrelProvider {
87 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88 f.debug_struct("MongrelProvider").finish_non_exhaustive()
89 }
90}
91
92#[async_trait::async_trait]
93impl TableProvider for MongrelProvider {
94 fn schema(&self) -> SchemaRef {
95 self.schema.clone()
96 }
97
98 fn table_type(&self) -> TableType {
99 TableType::Base
100 }
101
102 fn supports_filters_pushdown(
107 &self,
108 filters: &[&Expr],
109 ) -> DFResult<Vec<datafusion::logical_expr::TableProviderFilterPushDown>> {
110 use datafusion::logical_expr::TableProviderFilterPushDown;
111 let schema_ref = self.db.lock().schema().clone();
112 Ok(filters
113 .iter()
114 .map(|f| match translate_filter(f, &schema_ref) {
115 Some(
116 mongreldb_core::Condition::FmContains { .. }
117 | mongreldb_core::Condition::FmContainsAll { .. },
118 ) => TableProviderFilterPushDown::Inexact,
119 Some(_) => TableProviderFilterPushDown::Exact,
120 None => match translate_ann_search(f, &schema_ref)
121 .or_else(|| translate_sparse_match(f, &schema_ref))
122 {
123 Some(_) => TableProviderFilterPushDown::Exact,
124 None => TableProviderFilterPushDown::Unsupported,
125 },
126 })
127 .collect())
128 }
129
130 async fn scan(
131 &self,
132 _state: &dyn Session,
133 projection: Option<&Vec<usize>>,
134 filters: &[Expr],
135 _limit: Option<usize>,
136 ) -> DFResult<Arc<dyn ExecutionPlan>> {
137 let core_err = |e: mongreldb_core::MongrelError| {
138 DataFusionError::External(Box::new(MongrelQueryError::Core(e)))
139 };
140 let mut db = self.db.lock();
141 let snap = db.snapshot();
142 let schema_ref = db.schema().clone();
143
144 let translated: Vec<mongreldb_core::Condition> = filters
146 .iter()
147 .filter_map(|f| {
148 translate_filter(f, &schema_ref)
149 .or_else(|| translate_ann_search(f, &schema_ref))
150 .or_else(|| translate_sparse_match(f, &schema_ref))
151 })
152 .collect();
153
154 if !translated.is_empty() {
157 db.ensure_indexes_complete().map_err(core_err)?;
158 }
159
160 let empty_proj = projection.map(|p| p.is_empty()).unwrap_or(false);
164 if empty_proj {
165 let total: usize = if translated.is_empty() {
166 mongreldb_core::trace::QueryTrace::record(|t| {
167 t.scan_mode = mongreldb_core::trace::ScanMode::CountMetadata;
168 });
169 db.count() as usize
170 } else if let Some(count) = db.count_conditions(&translated, snap).map_err(core_err)? {
171 count as usize
172 } else {
173 match schema_ref.columns.first() {
174 Some(cdef) => {
175 let one = [cdef.id];
176 let cols = match db
177 .query_columns_native_cached(&translated, Some(&one), snap)
178 .map_err(core_err)?
179 {
180 Some(c) => c,
181 None => db
182 .visible_columns_native(snap, Some(&one))
183 .map_err(core_err)?,
184 };
185 mongreldb_core::trace::QueryTrace::record(|t| {
186 t.scan_mode = mongreldb_core::trace::ScanMode::Materialized;
187 });
188 cols.first().map(|(_, c)| c.len()).unwrap_or(0)
189 }
190 None => 0,
191 }
192 };
193 return Ok(Arc::new(scan::MongrelScanExec::new_row_count(total)));
194 }
195
196 let (col_ids, scan_schema): (Vec<u16>, SchemaRef) = match projection {
201 Some(p) if !p.is_empty() => {
202 let ids = p.iter().map(|&idx| schema_ref.columns[idx].id).collect();
203 let fields: Vec<arrow::datatypes::Field> = p
204 .iter()
205 .map(|&idx| self.schema.field(idx).clone())
206 .collect();
207 (ids, Arc::new(arrow::datatypes::Schema::new(fields)))
208 }
209 _ => (
210 schema_ref.columns.iter().map(|c| c.id).collect(),
211 self.schema.clone(),
212 ),
213 };
214
215 let mut proj_pairs: Vec<(u16, mongreldb_core::schema::TypeId)> =
217 Vec::with_capacity(col_ids.len());
218 let mut types: Vec<mongreldb_core::schema::TypeId> = Vec::with_capacity(col_ids.len());
219 for cid in &col_ids {
220 let ty = schema_ref
221 .columns
222 .iter()
223 .find(|c| c.id == *cid)
224 .map(|c| c.ty)
225 .ok_or_else(|| {
226 DataFusionError::External(Box::new(MongrelQueryError::Arrow(format!(
227 "unknown column {cid}"
228 ))))
229 })?;
230 proj_pairs.push((*cid, ty));
231 types.push(ty);
232 }
233
234 let col_stats_map = if translated.is_empty() {
238 db.exact_column_stats(snap, &col_ids).map_err(core_err)?
239 } else {
240 None
241 };
242 let column_stats: Vec<datafusion::physical_plan::ColumnStatistics> = col_ids
243 .iter()
244 .map(|cid| scan::to_col_statistics(col_stats_map.as_ref().and_then(|m| m.get(cid))))
245 .collect();
246
247 if translated.is_empty()
251 && db.run_count() == 1
252 && db.memtable_is_empty()
253 && db.mutable_run_len() == 0
254 && db.single_run_is_clean()
255 {
256 let shadow = shadow::ArrowShadow::new(db.dir());
257 let run_ids: HashSet<u128> = db.run_ids().into_iter().collect();
258 shadow.sweep(&run_ids);
259 if let Some(&run_id) = run_ids.iter().next() {
260 if let Some(batch) = shadow.try_read(run_id) {
261 if let Some(projected) =
262 project_batch(&batch, &col_ids, &schema_ref, &scan_schema)
263 {
264 mongreldb_core::trace::QueryTrace::record(|t| {
265 t.scan_mode = mongreldb_core::trace::ScanMode::ArrowShadow;
266 });
267 return Ok(Arc::new(scan::MongrelScanExec::new_batch(
268 scan_schema,
269 projected,
270 column_stats,
271 )));
272 }
273 }
274 }
275 }
276
277 let cursor: Option<Box<dyn Cursor>> = db
284 .scan_cursor(snap, proj_pairs, &translated)
285 .map_err(core_err)?;
286 if let Some(cursor) = cursor {
287 let num_rows = cursor.remaining_rows();
288 let residual = extract_residual_filter(filters, &col_ids, &schema_ref);
290 return Ok(Arc::new(scan::MongrelScanExec::new_cursor(
291 scan_schema,
292 types,
293 cursor,
294 num_rows,
295 column_stats,
296 residual,
297 )));
298 }
299
300 let cols = if !translated.is_empty() {
303 match db
304 .query_columns_native_cached(&translated, Some(&col_ids), snap)
305 .map_err(core_err)?
306 {
307 Some(c) => c,
308 None => db
309 .visible_columns_native(snap, Some(&col_ids))
310 .map_err(core_err)?,
311 }
312 } else {
313 db.visible_columns_native(snap, Some(&col_ids))
314 .map_err(core_err)?
315 };
316
317 let mut ordered: Vec<mongreldb_core::columnar::NativeColumn> =
319 Vec::with_capacity(col_ids.len());
320 for cid in &col_ids {
321 let col = cols
322 .iter()
323 .find(|(id, _)| id == cid)
324 .map(|(_, c)| c.clone())
325 .ok_or_else(|| {
326 DataFusionError::External(Box::new(MongrelQueryError::Arrow(format!(
327 "missing column {cid}"
328 ))))
329 })?;
330 ordered.push(col);
331 }
332 let num_rows = ordered.first().map(|c| c.len()).unwrap_or(0);
333
334 let shadow_write: Option<(
336 std::path::PathBuf,
337 u128,
338 Vec<arrow::array::ArrayRef>,
339 SchemaRef,
340 )> = if translated.is_empty()
341 && db.run_count() == 1
342 && db.memtable_is_empty()
343 && db.mutable_run_len() == 0
344 && db.single_run_is_clean()
345 {
346 let all_schema_ids: Vec<u16> = schema_ref.columns.iter().map(|c| c.id).collect();
347 if col_ids == all_schema_ids {
348 let dir = db.dir().to_path_buf();
349 let run_id = db.run_ids().first().copied();
350 run_id.map(|rid| {
351 let arrays = ordered
352 .iter()
353 .zip(types.iter())
354 .map(|(col, &ty)| arrow_conv::native_to_array(ty, col))
355 .collect::<Result<_>>()
356 .unwrap_or_default();
357 (dir, rid, arrays, scan_schema.clone())
358 })
359 } else {
360 None
361 }
362 } else {
363 None
364 };
365
366 drop(db);
367
368 if let Some((dir, run_id, arrays, schema)) = shadow_write {
371 if let Ok(batch) = RecordBatch::try_new(schema, arrays) {
372 shadow::ArrowShadow::new(&dir).write(run_id, &batch);
373 }
374 }
375
376 mongreldb_core::trace::QueryTrace::record(|t| {
377 t.scan_mode = mongreldb_core::trace::ScanMode::Materialized;
378 t.row_materialized = true;
379 });
380 Ok(Arc::new(scan::MongrelScanExec::new(
381 scan_schema,
382 ordered,
383 types,
384 num_rows,
385 column_stats,
386 )))
387 }
388}
389
390fn project_batch(
394 batch: &RecordBatch,
395 col_ids: &[u16],
396 schema_ref: &mongreldb_core::schema::Schema,
397 scan_schema: &arrow::datatypes::SchemaRef,
398) -> Option<RecordBatch> {
399 let arrays: Vec<arrow::array::ArrayRef> = col_ids
401 .iter()
402 .map(|cid| {
403 let name = schema_ref
405 .columns
406 .iter()
407 .find(|c| c.id == *cid)
408 .map(|c| c.name.as_str())?;
409 let idx = batch.schema().index_of(name).ok()?;
411 Some(batch.column(idx).clone())
412 })
413 .collect::<Option<Vec<_>>>()?;
414 RecordBatch::try_new(scan_schema.clone(), arrays).ok()
415}
416
417pub(crate) fn translate_filter(
435 expr: &Expr,
436 schema: &mongreldb_core::Schema,
437) -> Option<mongreldb_core::Condition> {
438 use datafusion::common::ScalarValue;
439 use datafusion::logical_expr::{Between, BinaryExpr, Like, Operator};
440 use mongreldb_core::{ColumnFlags, Condition, IndexKind, TypeId, Value};
441
442 let int_val = |s: &ScalarValue| match s {
446 ScalarValue::Int8(Some(v)) => Some(*v as i64),
447 ScalarValue::Int16(Some(v)) => Some(*v as i64),
448 ScalarValue::Int32(Some(v)) => Some(*v as i64),
449 ScalarValue::Int64(Some(v)) => Some(*v),
450 ScalarValue::UInt8(Some(v)) => Some(*v as i64),
451 ScalarValue::UInt16(Some(v)) => Some(*v as i64),
452 ScalarValue::UInt32(Some(v)) => Some(*v as i64),
453 ScalarValue::UInt64(Some(v)) => Some(*v as i64),
454 ScalarValue::Date32(Some(v)) => Some(*v as i64),
455 ScalarValue::TimestampSecond(Some(v), _) => Some(*v),
456 ScalarValue::TimestampMillisecond(Some(v), _) => Some(*v),
457 ScalarValue::TimestampMicrosecond(Some(v), _) => Some(*v),
458 ScalarValue::TimestampNanosecond(Some(v), _) => Some(*v),
459 _ => None,
460 };
461 let float_val = |s: &ScalarValue| match s {
462 ScalarValue::Float32(Some(f)) => Some(*f as f64),
463 ScalarValue::Float64(Some(f)) => Some(*f),
464 _ => None,
465 };
466 let bytes_val = |s: &ScalarValue| match s {
467 ScalarValue::Utf8(Some(s)) => Some(s.as_bytes().to_vec()),
468 _ => None,
469 };
470 let _ = bytes_val; let val = |s: &ScalarValue| -> Option<Value> {
473 if let Some(i) = int_val(s) {
477 return Some(Value::Int64(i));
478 }
479 match s {
480 ScalarValue::Utf8(Some(s)) => Some(Value::Bytes(s.as_bytes().to_vec())),
481 ScalarValue::Float32(Some(f)) => Some(Value::Float64(*f as f64)),
482 ScalarValue::Float64(Some(f)) => Some(Value::Float64(*f)),
483 ScalarValue::Boolean(Some(b)) => Some(Value::Bool(*b)),
484 _ => None,
485 }
486 };
487
488 let col_def = |name: &str| schema.columns.iter().find(|c| c.name == name);
489 let has_fm = |cid: u16| {
490 schema
491 .indexes
492 .iter()
493 .any(|i| i.column_id == cid && i.kind == IndexKind::FmIndex)
494 };
495 let has_bitmap = |cid: u16| {
496 schema
497 .indexes
498 .iter()
499 .any(|i| i.column_id == cid && i.kind == IndexKind::Bitmap)
500 };
501
502 match expr {
503 Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
506 if *op == Operator::Or {
508 return try_or_as_bitmap_in(expr, schema);
509 }
510 let left = peel_cast(left);
512 let right = peel_cast(right);
513 let (col_name, scalar, flipped) = match (left.as_ref(), right.as_ref()) {
514 (Expr::Column(c), Expr::Literal(s, _)) => (&c.name, s, false),
515 (Expr::Literal(s, _), Expr::Column(c)) => (&c.name, s, true),
516 _ => return None,
517 };
518 let op = if flipped { flip_op(*op)? } else { *op };
519 let cdef = col_def(col_name)?;
520
521 if op == Operator::Eq {
523 let v = val(scalar)?;
524 if has_bitmap(cdef.id) {
525 return Some(Condition::BitmapEq {
526 column_id: cdef.id,
527 value: v.encode_key(),
528 });
529 }
530 if cdef.flags.contains(ColumnFlags::PRIMARY_KEY) {
531 return Some(Condition::Pk(v.encode_key()));
532 }
533 return None;
534 }
535
536 match cdef.ty {
539 TypeId::Int8
540 | TypeId::Int16
541 | TypeId::Int32
542 | TypeId::Int64
543 | TypeId::UInt8
544 | TypeId::UInt16
545 | TypeId::UInt32
546 | TypeId::UInt64
547 | TypeId::TimestampNanos
548 | TypeId::Date32 => {
549 let v = int_val(scalar)?;
550 let (lo, hi) = int_bounds(op, v)?;
551 Some(Condition::Range {
552 column_id: cdef.id,
553 lo,
554 hi,
555 })
556 }
557 TypeId::Float32 | TypeId::Float64 => {
558 let v = float_val(scalar)?;
559 let (lo, lo_inc, hi, hi_inc) = float_bounds(op, v)?;
560 Some(Condition::RangeF64 {
561 column_id: cdef.id,
562 lo,
563 lo_inclusive: lo_inc,
564 hi,
565 hi_inclusive: hi_inc,
566 })
567 }
568 _ => None,
569 }
570 }
571
572 Expr::Between(Between {
574 expr,
575 negated,
576 low,
577 high,
578 }) => {
579 if *negated {
580 return None;
581 }
582 let Expr::Column(c) = expr.as_ref() else {
583 return None;
584 };
585 let cdef = col_def(&c.name)?;
586 let (lo_s, hi_s) = match (low.as_ref(), high.as_ref()) {
587 (Expr::Literal(lo, _), Expr::Literal(hi, _)) => (lo, hi),
588 _ => return None,
589 };
590 match cdef.ty {
591 TypeId::Int8
592 | TypeId::Int16
593 | TypeId::Int32
594 | TypeId::Int64
595 | TypeId::UInt8
596 | TypeId::UInt16
597 | TypeId::UInt32
598 | TypeId::UInt64
599 | TypeId::TimestampNanos
600 | TypeId::Date32 => {
601 let (Some(lo), Some(hi)) = (int_val(lo_s), int_val(hi_s)) else {
602 return None;
603 };
604 Some(Condition::Range {
605 column_id: cdef.id,
606 lo,
607 hi,
608 })
609 }
610 TypeId::Float32 | TypeId::Float64 => {
611 let (Some(lo), Some(hi)) = (float_val(lo_s), float_val(hi_s)) else {
612 return None;
613 };
614 Some(Condition::RangeF64 {
615 column_id: cdef.id,
616 lo,
617 lo_inclusive: true,
618 hi,
619 hi_inclusive: true,
620 })
621 }
622 _ => None,
623 }
624 }
625
626 Expr::Like(Like {
628 negated,
629 expr,
630 pattern,
631 ..
632 }) => {
633 if *negated {
634 return None;
635 }
636 let Expr::Column(c) = expr.as_ref() else {
637 return None;
638 };
639 let Expr::Literal(ScalarValue::Utf8(Some(pat)), _) = pattern.as_ref() else {
640 return None;
641 };
642 let cdef = col_def(&c.name)?;
643 if has_bitmap(cdef.id) {
647 if let Some(prefix) = anchored_like_prefix(pat) {
648 return Some(Condition::BytesPrefix {
649 column_id: cdef.id,
650 prefix: mongreldb_core::Value::Bytes(prefix.as_bytes().to_vec())
651 .encode_key(),
652 });
653 }
654 }
655 if !has_fm(cdef.id) {
656 return None;
657 }
658 let segments: Vec<Vec<u8>> = pat
662 .split(['%', '_'])
663 .filter(|s| s.len() >= 3)
664 .map(|s| s.as_bytes().to_vec())
665 .collect();
666 match segments.len() {
667 0 => longest_like_segment(pat).map(|seg| Condition::FmContains {
668 column_id: cdef.id,
669 pattern: seg,
670 }),
671 1 => Some(Condition::FmContains {
672 column_id: cdef.id,
673 pattern: segments.into_iter().next().unwrap(),
674 }),
675 _ => Some(Condition::FmContainsAll {
676 column_id: cdef.id,
677 patterns: segments,
678 }),
679 }
680 }
681
682 Expr::InList(il) if !il.negated => {
686 let Expr::Column(c) = il.expr.as_ref() else {
687 return None;
688 };
689 let cdef = col_def(&c.name)?;
690 if !has_bitmap(cdef.id) {
691 return None;
692 }
693 let values: Vec<Vec<u8>> = il
694 .list
695 .iter()
696 .filter_map(|e| match e {
697 Expr::Literal(s, _) => val(s).map(|v| v.encode_key()),
698 _ => None,
699 })
700 .collect();
701 if values.is_empty() || values.len() != il.list.len() {
702 return None;
703 }
704 Some(Condition::BitmapIn {
705 column_id: cdef.id,
706 values,
707 })
708 }
709
710 Expr::IsNull(inner) => {
712 let col_name = match inner.as_ref() {
713 Expr::Column(c) => &c.name,
714 _ => return None,
715 };
716 let cdef = col_def(col_name)?;
717 Some(Condition::IsNull { column_id: cdef.id })
718 }
719
720 Expr::IsNotNull(inner) => {
722 let col_name = match inner.as_ref() {
723 Expr::Column(c) => &c.name,
724 _ => return None,
725 };
726 let cdef = col_def(col_name)?;
727 Some(Condition::IsNotNull { column_id: cdef.id })
728 }
729
730 _ => None,
731 }
732}
733
734pub(crate) fn extract_residual_filter(
738 filters: &[Expr],
739 col_ids: &[u16],
740 schema: &mongreldb_core::Schema,
741) -> Option<std::sync::Arc<scan::ResidualFilter>> {
742 use datafusion::common::ScalarValue;
743 use datafusion::logical_expr::Like;
744 for f in filters {
745 if let Expr::Like(Like {
746 negated: false,
747 expr,
748 pattern,
749 ..
750 }) = f
751 {
752 let Expr::Column(c) = expr.as_ref() else {
753 continue;
754 };
755 let Expr::Literal(ScalarValue::Utf8(Some(pat)), _) = pattern.as_ref() else {
756 continue;
757 };
758 let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
759 let col_idx = col_ids.iter().position(|&id| id == cdef.id)?;
760 return Some(std::sync::Arc::new(scan::ResidualFilter::new(
761 col_idx,
762 pat.as_bytes().to_vec(),
763 )));
764 }
765 }
766 None
767}
768
769pub(crate) fn translate_ann_search(
775 expr: &Expr,
776 schema: &mongreldb_core::Schema,
777) -> Option<mongreldb_core::Condition> {
778 use datafusion::common::ScalarValue;
779 use mongreldb_core::Condition;
780
781 let Expr::ScalarFunction(sf) = expr else {
782 return None;
783 };
784 if !sf.func.name().eq_ignore_ascii_case("ann_search") || sf.args.len() != 3 {
785 return None;
786 }
787 let (Expr::Column(c), query_expr, k_expr) = (&sf.args[0], &sf.args[1], &sf.args[2]) else {
788 return None;
789 };
790 let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
791 let json = match query_expr {
792 Expr::Literal(ScalarValue::Utf8(Some(s)), _) => s.as_str(),
793 _ => return None,
794 };
795 let k: i64 = match k_expr {
796 Expr::Literal(scalar, _) => match scalar {
797 ScalarValue::Int64(Some(k)) => *k,
798 ScalarValue::UInt64(Some(k)) => *k as i64,
799 ScalarValue::Int32(Some(k)) => *k as i64,
800 _ => return None,
801 },
802 _ => return None,
803 };
804 let query: Vec<f32> = serde_json::from_str(json).ok()?;
805 Some(Condition::Ann {
806 column_id: cdef.id,
807 query,
808 k: k.max(1) as usize,
809 })
810}
811
812pub(crate) fn translate_sparse_match(
817 expr: &Expr,
818 schema: &mongreldb_core::Schema,
819) -> Option<mongreldb_core::Condition> {
820 use datafusion::common::ScalarValue;
821 use mongreldb_core::Condition;
822
823 let Expr::ScalarFunction(sf) = expr else {
824 return None;
825 };
826 if !sf.func.name().eq_ignore_ascii_case("sparse_match") || sf.args.len() != 3 {
827 return None;
828 }
829 let (Expr::Column(c), query_expr, k_expr) = (&sf.args[0], &sf.args[1], &sf.args[2]) else {
830 return None;
831 };
832 let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
833 let json = match query_expr {
834 Expr::Literal(ScalarValue::Utf8(Some(s)), _) => s.as_str(),
835 _ => return None,
836 };
837 let k: i64 = match k_expr {
838 Expr::Literal(scalar, _) => match scalar {
839 ScalarValue::Int64(Some(k)) => *k,
840 ScalarValue::UInt64(Some(k)) => *k as i64,
841 ScalarValue::Int32(Some(k)) => *k as i64,
842 _ => return None,
843 },
844 _ => return None,
845 };
846 let query: Vec<(u32, f32)> = serde_json::from_str(json).ok()?;
847 Some(Condition::SparseMatch {
848 column_id: cdef.id,
849 query,
850 k: k.max(1) as usize,
851 })
852}
853
854fn flip_op(op: datafusion::logical_expr::Operator) -> Option<datafusion::logical_expr::Operator> {
856 use datafusion::logical_expr::Operator;
857 Some(match op {
858 Operator::Eq => Operator::Eq,
859 Operator::Lt => Operator::Gt,
860 Operator::Gt => Operator::Lt,
861 Operator::LtEq => Operator::GtEq,
862 Operator::GtEq => Operator::LtEq,
863 _ => return None,
864 })
865}
866
867fn int_bounds(op: datafusion::logical_expr::Operator, v: i64) -> Option<(i64, i64)> {
870 use datafusion::logical_expr::Operator;
871 Some(match op {
872 Operator::Gt => (v.saturating_add(1), i64::MAX),
873 Operator::GtEq => (v, i64::MAX),
874 Operator::Lt => (i64::MIN, v.saturating_sub(1)),
875 Operator::LtEq => (i64::MIN, v),
876 _ => return None,
877 })
878}
879
880fn float_bounds(op: datafusion::logical_expr::Operator, v: f64) -> Option<(f64, bool, f64, bool)> {
882 use datafusion::logical_expr::Operator;
883 Some(match op {
884 Operator::Gt => (v, false, f64::INFINITY, false),
885 Operator::GtEq => (v, true, f64::INFINITY, false),
886 Operator::Lt => (f64::NEG_INFINITY, false, v, false),
887 Operator::LtEq => (f64::NEG_INFINITY, false, v, true),
888 _ => return None,
889 })
890}
891
892fn longest_like_segment(pat: &str) -> Option<Vec<u8>> {
899 pat.split(['%', '_'])
900 .map(|s| s.as_bytes())
901 .max_by_key(|s| s.len())
902 .filter(|s| !s.is_empty())
903 .map(|s| s.to_vec())
904}
905
906fn anchored_like_prefix(pat: &str) -> Option<&str> {
911 let rest = pat.strip_suffix('%')?;
912 if rest.is_empty() || rest.contains(['%', '_']) {
913 return None;
914 }
915 Some(rest)
916}
917
918fn peel_cast(expr: &Expr) -> std::borrow::Cow<'_, Expr> {
922 match expr {
923 Expr::Cast(datafusion::logical_expr::Cast { expr, .. }) => std::borrow::Cow::Borrowed(expr),
924 _ => std::borrow::Cow::Borrowed(expr),
925 }
926}
927
928fn try_or_as_bitmap_in(
933 expr: &Expr,
934 schema: &mongreldb_core::Schema,
935) -> Option<mongreldb_core::Condition> {
936 use datafusion::logical_expr::{BinaryExpr, Operator};
937 let mut values: Vec<Vec<u8>> = Vec::new();
938 let mut target_col: Option<u16> = None;
939 let mut stack = vec![expr];
940 while let Some(e) = stack.pop() {
941 match e {
942 Expr::BinaryExpr(BinaryExpr {
943 left,
944 op: Operator::Or,
945 right,
946 }) => {
947 stack.push(left);
948 stack.push(right);
949 }
950 Expr::BinaryExpr(BinaryExpr {
951 left,
952 op: Operator::Eq,
953 right,
954 }) => {
955 let (col_name, scalar) = match (left.as_ref(), right.as_ref()) {
956 (Expr::Column(c), Expr::Literal(s, _)) => (&c.name, s),
957 (Expr::Literal(s, _), Expr::Column(c)) => (&c.name, s),
958 _ => return None,
959 };
960 let cdef = schema.columns.iter().find(|c| &c.name == col_name)?;
961 if !schema
962 .indexes
963 .iter()
964 .any(|i| i.column_id == cdef.id && i.kind == mongreldb_core::IndexKind::Bitmap)
965 {
966 return None;
967 }
968 match target_col {
969 None => target_col = Some(cdef.id),
970 Some(id) if id != cdef.id => return None,
971 _ => {}
972 }
973 let v = match scalar {
974 datafusion::common::ScalarValue::Int64(Some(v)) => {
975 mongreldb_core::Value::Int64(*v)
976 }
977 datafusion::common::ScalarValue::Utf8(Some(s)) => {
978 mongreldb_core::Value::Bytes(s.as_bytes().to_vec())
979 }
980 datafusion::common::ScalarValue::Float64(Some(f)) => {
981 mongreldb_core::Value::Float64(*f)
982 }
983 datafusion::common::ScalarValue::Boolean(Some(b)) => {
984 mongreldb_core::Value::Bool(*b)
985 }
986 _ => return None,
987 };
988 values.push(v.encode_key());
989 }
990 _ => return None,
991 }
992 }
993 let col_id = target_col?;
994 if values.is_empty() {
995 return None;
996 }
997 Some(mongreldb_core::Condition::BitmapIn {
998 column_id: col_id,
999 values,
1000 })
1001}
1002
1003fn sp_ident_name(expr: &sqlparser::ast::Expr) -> Option<&str> {
1010 use sqlparser::ast::Expr;
1011 match expr {
1012 Expr::Identifier(ident) => Some(ident.value.as_str()),
1013 Expr::CompoundIdentifier(idents) => idents.last().map(|i| i.value.as_str()),
1014 _ => None,
1015 }
1016}
1017
1018fn sp_literal(expr: &sqlparser::ast::Expr) -> Option<mongreldb_core::Value> {
1021 use sqlparser::ast::Expr;
1022 let v = match expr {
1023 Expr::Value(v) => v,
1024 _ => return None,
1025 };
1026 use sqlparser::ast::Value as SpValue;
1027 match &v.value {
1028 SpValue::Number(s, _) => s
1029 .parse::<i64>()
1030 .map(mongreldb_core::Value::Int64)
1031 .or_else(|_| s.parse::<f64>().map(mongreldb_core::Value::Float64))
1032 .ok(),
1033 SpValue::SingleQuotedString(s) => Some(mongreldb_core::Value::Bytes(s.as_bytes().to_vec())),
1034 SpValue::Boolean(b) => Some(mongreldb_core::Value::Bool(*b)),
1035 _ => None,
1036 }
1037}
1038
1039fn is_int_ty(ty: mongreldb_core::schema::TypeId) -> bool {
1040 use mongreldb_core::schema::TypeId::*;
1041 matches!(
1042 ty,
1043 Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 | TimestampNanos | Date32
1044 )
1045}
1046
1047fn is_float_ty(ty: mongreldb_core::schema::TypeId) -> bool {
1048 matches!(
1049 ty,
1050 mongreldb_core::schema::TypeId::Float32 | mongreldb_core::schema::TypeId::Float64
1051 )
1052}
1053
1054fn translate_sqlparser_predicate(
1057 expr: &sqlparser::ast::Expr,
1058 schema: &mongreldb_core::Schema,
1059) -> Option<mongreldb_core::Condition> {
1060 use mongreldb_core::{schema::ColumnFlags, Condition, IndexKind, Value};
1061 use sqlparser::ast::{BinaryOperator, Expr};
1062
1063 let col_def = |name: &str| schema.columns.iter().find(|c| c.name == name);
1064 let has_bitmap = |cid: u16| {
1065 schema
1066 .indexes
1067 .iter()
1068 .any(|i| i.column_id == cid && i.kind == IndexKind::Bitmap)
1069 };
1070
1071 match expr {
1072 Expr::BinaryOp {
1074 left,
1075 op: BinaryOperator::Or,
1076 right,
1077 } => {
1078 let mut values: Vec<Vec<u8>> = Vec::new();
1079 let mut target: Option<u16> = None;
1080 let mut stack: Vec<&Expr> = vec![left.as_ref(), right.as_ref()];
1081 while let Some(e) = stack.pop() {
1082 match e {
1083 Expr::BinaryOp {
1084 left,
1085 op: BinaryOperator::Or,
1086 right,
1087 } => {
1088 stack.push(left.as_ref());
1089 stack.push(right.as_ref());
1090 }
1091 Expr::BinaryOp {
1092 left,
1093 op: BinaryOperator::Eq,
1094 right,
1095 } => {
1096 let (name, lit) = match (left.as_ref(), right.as_ref()) {
1097 (l, r) if sp_ident_name(l).is_some() && sp_literal(r).is_some() => {
1098 (l, r)
1099 }
1100 (l, r) if sp_ident_name(r).is_some() && sp_literal(l).is_some() => {
1101 (r, l)
1102 }
1103 _ => return None,
1104 };
1105 let cdef = col_def(sp_ident_name(name)?)?;
1106 if !has_bitmap(cdef.id) {
1107 return None;
1108 }
1109 match target {
1110 None => target = Some(cdef.id),
1111 Some(id) if id != cdef.id => return None,
1112 _ => {}
1113 }
1114 values.push(sp_literal(lit)?.encode_key());
1115 }
1116 _ => return None,
1117 }
1118 }
1119 let cid = target?;
1120 (!values.is_empty()).then_some(Condition::BitmapIn {
1121 column_id: cid,
1122 values,
1123 })
1124 }
1125 Expr::BinaryOp { left, op, right } => {
1127 let flipped;
1128 let (col_expr, lit_expr) = match (
1129 sp_ident_name(left),
1130 sp_literal(right),
1131 sp_ident_name(right),
1132 sp_literal(left),
1133 ) {
1134 (Some(_), Some(_), _, _) => {
1135 flipped = false;
1136 (left.as_ref(), right.as_ref())
1137 }
1138 (_, _, Some(_), Some(_)) => {
1139 flipped = true;
1140 (right.as_ref(), left.as_ref())
1141 }
1142 _ => return None,
1143 };
1144 let name = sp_ident_name(col_expr)?;
1145 let cdef = col_def(name)?;
1146 let v = sp_literal(lit_expr)?;
1147 use sqlparser::ast::BinaryOperator::*;
1148 match op {
1151 Eq => {
1152 if has_bitmap(cdef.id) {
1153 Some(Condition::BitmapEq {
1154 column_id: cdef.id,
1155 value: v.encode_key(),
1156 })
1157 } else if cdef.flags.contains(ColumnFlags::PRIMARY_KEY) {
1158 Some(Condition::Pk(v.encode_key()))
1159 } else {
1160 None
1161 }
1162 }
1163 Lt | LtEq | Gt | GtEq if is_int_ty(cdef.ty) => {
1164 let n = match v {
1165 Value::Int64(n) => n,
1166 _ => return None,
1167 };
1168 let (lo, hi) = match (flipped, op) {
1170 (false, Lt) | (true, Gt) => (i64::MIN, n.saturating_sub(1)),
1171 (false, Gt) | (true, Lt) => (n.saturating_add(1), i64::MAX),
1172 (false, LtEq) | (true, GtEq) => (i64::MIN, n),
1173 (false, GtEq) | (true, LtEq) => (n, i64::MAX),
1174 _ => (i64::MIN, i64::MAX),
1175 };
1176 Some(Condition::Range {
1177 column_id: cdef.id,
1178 lo,
1179 hi,
1180 })
1181 }
1182 Lt | LtEq | Gt | GtEq if is_float_ty(cdef.ty) => {
1183 let f = match v {
1184 Value::Float64(f) => f,
1185 _ => return None,
1186 };
1187 let (lo, li, hi, hi_i) = match (flipped, op) {
1188 (false, Lt) | (true, Gt) => (f64::NEG_INFINITY, true, f, false),
1189 (false, Gt) | (true, Lt) => (f, false, f64::INFINITY, true),
1190 (false, LtEq) | (true, GtEq) => (f64::NEG_INFINITY, true, f, true),
1191 (false, GtEq) | (true, LtEq) => (f, true, f64::INFINITY, true),
1192 _ => (f64::NEG_INFINITY, true, f64::INFINITY, true),
1193 };
1194 Some(Condition::RangeF64 {
1195 column_id: cdef.id,
1196 lo,
1197 lo_inclusive: li,
1198 hi,
1199 hi_inclusive: hi_i,
1200 })
1201 }
1202 _ => None,
1203 }
1204 }
1205 Expr::Between {
1206 expr,
1207 negated,
1208 low,
1209 high,
1210 } if !negated => {
1211 let name = sp_ident_name(expr)?;
1212 let cdef = col_def(name)?;
1213 if is_int_ty(cdef.ty) {
1214 let lo = match sp_literal(low)? {
1215 Value::Int64(n) => n,
1216 _ => return None,
1217 };
1218 let hi = match sp_literal(high)? {
1219 Value::Int64(n) => n,
1220 _ => return None,
1221 };
1222 Some(Condition::Range {
1223 column_id: cdef.id,
1224 lo,
1225 hi,
1226 })
1227 } else if is_float_ty(cdef.ty) {
1228 let lo = match sp_literal(low)? {
1229 Value::Float64(f) => f,
1230 _ => return None,
1231 };
1232 let hi = match sp_literal(high)? {
1233 Value::Float64(f) => f,
1234 _ => return None,
1235 };
1236 Some(Condition::RangeF64 {
1237 column_id: cdef.id,
1238 lo,
1239 lo_inclusive: true,
1240 hi,
1241 hi_inclusive: true,
1242 })
1243 } else {
1244 None
1245 }
1246 }
1247 Expr::InList {
1248 expr,
1249 list,
1250 negated,
1251 } if !negated => {
1252 let name = sp_ident_name(expr)?;
1253 let cdef = col_def(name)?;
1254 if !has_bitmap(cdef.id) {
1255 return None;
1256 }
1257 let values: Vec<Vec<u8>> = list
1258 .iter()
1259 .map(|e| sp_literal(e).map(|v| v.encode_key()))
1260 .collect::<Option<_>>()?;
1261 (!values.is_empty()).then_some(Condition::BitmapIn {
1262 column_id: cdef.id,
1263 values,
1264 })
1265 }
1266 Expr::IsNull(inner) => {
1269 let cid = col_def(sp_ident_name(inner)?)?.id;
1270 Some(Condition::IsNull { column_id: cid })
1271 }
1272 Expr::IsNotNull(inner) => {
1273 let cid = col_def(sp_ident_name(inner)?)?.id;
1274 Some(Condition::IsNotNull { column_id: cid })
1275 }
1276 _ => None,
1277 }
1278}
1279
1280fn translate_sqlparser_filter(
1283 expr: &sqlparser::ast::Expr,
1284 schema: &mongreldb_core::Schema,
1285) -> Option<Vec<mongreldb_core::Condition>> {
1286 use sqlparser::ast::{BinaryOperator, Expr};
1287 let mut out = Vec::new();
1288 let mut stack = vec![expr];
1289 while let Some(e) = stack.pop() {
1290 match e {
1291 Expr::BinaryOp {
1292 left,
1293 op: BinaryOperator::And,
1294 right,
1295 } => {
1296 stack.push(left.as_ref());
1297 stack.push(right.as_ref());
1298 }
1299 other => out.push(translate_sqlparser_predicate(other, schema)?),
1300 }
1301 }
1302 Some(out)
1303}
1304
1305pub struct MongrelSession {
1309 ctx: SessionContext,
1310 db: Option<Arc<Mutex<Table>>>,
1311 database: Option<Arc<Database>>,
1315 cache: ResultCache,
1316 plan_cache: parking_lot::Mutex<HashMap<String, datafusion::logical_expr::LogicalPlan>>,
1318 tables: parking_lot::Mutex<HashMap<String, Arc<Mutex<Table>>>>,
1320 views: parking_lot::Mutex<HashMap<String, ViewDef>>,
1324 attached_databases: parking_lot::Mutex<HashMap<String, Arc<Database>>>,
1328 sql_txn: parking_lot::Mutex<Option<Vec<commands::PendingSqlOp>>>,
1332 savepoints: parking_lot::Mutex<Vec<(String, usize)>>,
1335 sql_fn_state: Arc<extended_sql_functions::ExtendedSqlState>,
1337 external_modules: Arc<ExternalModuleRegistry>,
1340}
1341
1342type CacheKey = (String, u64);
1344type ResultCache = parking_lot::Mutex<std::collections::HashMap<CacheKey, Arc<Vec<RecordBatch>>>>;
1345
1346impl MongrelSession {
1347 pub fn new(db: Table) -> Self {
1351 let db = Arc::new(Mutex::new(db));
1352 let ctx = SessionContext::new();
1353 let sql_fn_state = Arc::new(extended_sql_functions::ExtendedSqlState::default());
1354 register_mongrel_functions(&ctx, Arc::clone(&sql_fn_state));
1355 let external_modules = Arc::new(ExternalModuleRegistry::default());
1356 Self {
1357 ctx,
1358 db: Some(db),
1359 database: None,
1360 cache: parking_lot::Mutex::new(std::collections::HashMap::new()),
1361 plan_cache: parking_lot::Mutex::new(HashMap::new()),
1362 tables: parking_lot::Mutex::new(HashMap::new()),
1363 views: parking_lot::Mutex::new(HashMap::new()),
1364 attached_databases: parking_lot::Mutex::new(HashMap::new()),
1365 savepoints: parking_lot::Mutex::new(Vec::new()),
1366 sql_txn: parking_lot::Mutex::new(None),
1367 sql_fn_state,
1368 external_modules,
1369 }
1370 }
1371
1372 pub fn new_with_external_modules(
1373 db: Table,
1374 modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
1375 ) -> Result<Self> {
1376 let session = Self::new(db);
1377 for module in modules {
1378 session.register_external_module(module)?;
1379 }
1380 Ok(session)
1381 }
1382
1383 pub fn open(database: Arc<Database>) -> Result<Self> {
1388 Self::open_with_external_modules(database, std::iter::empty())
1389 }
1390
1391 pub fn open_with_external_modules(
1392 database: Arc<Database>,
1393 modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
1394 ) -> Result<Self> {
1395 let ctx = SessionContext::new();
1396 let sql_fn_state = Arc::new(extended_sql_functions::ExtendedSqlState::default());
1397 register_mongrel_functions(&ctx, Arc::clone(&sql_fn_state));
1398 let external_modules = Arc::new(ExternalModuleRegistry::default());
1399 for module in modules {
1400 external_modules.register(module)?;
1401 }
1402
1403 let mut tables: HashMap<String, Arc<Mutex<Table>>> = HashMap::new();
1404 for name in database.table_names() {
1405 let handle = database.table(&name)?;
1406 let provider = MongrelProvider::new(handle.clone())?;
1407 ctx.register_table(&name, Arc::new(provider))
1408 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1409 tables.insert(name, handle);
1410 }
1411 for entry in database.external_tables() {
1412 let provider = external_modules.external_table_provider(&database, &entry)?;
1413 ctx.register_table(&entry.name, provider)
1414 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1415 }
1416
1417 let primary = {
1420 let mut names: Vec<&String> = tables.keys().collect();
1421 names.sort();
1422 names.first().and_then(|n| tables.get(*n).cloned())
1423 };
1424
1425 Ok(Self {
1426 ctx,
1427 db: primary,
1428 database: Some(database),
1429 cache: parking_lot::Mutex::new(std::collections::HashMap::new()),
1430 plan_cache: parking_lot::Mutex::new(HashMap::new()),
1431 tables: parking_lot::Mutex::new(tables),
1432 views: parking_lot::Mutex::new(HashMap::new()),
1433 attached_databases: parking_lot::Mutex::new(HashMap::new()),
1434 savepoints: parking_lot::Mutex::new(Vec::new()),
1435 sql_txn: parking_lot::Mutex::new(None),
1436 sql_fn_state,
1437 external_modules,
1438 })
1439 }
1440
1441 pub fn register_external_module(&self, module: Arc<dyn ExternalTableModule>) -> Result<()> {
1442 self.external_modules.register(module)?;
1443 self.clear_cache();
1444 Ok(())
1445 }
1446
1447 pub fn db(&self) -> Option<&Arc<Mutex<Table>>> {
1451 self.db.as_ref()
1452 }
1453
1454 pub fn create_view(&self, name: &str, sql: &str) {
1459 self.create_view_with_schema(name, sql, CoreSchema::default(), HashMap::new());
1460 }
1461
1462 pub(crate) fn create_view_with_schema(
1463 &self,
1464 name: &str,
1465 sql: &str,
1466 schema: CoreSchema,
1467 input_types: HashMap<u16, Option<TypeId>>,
1468 ) {
1469 self.views.lock().insert(
1470 name.to_string(),
1471 ViewDef {
1472 sql: sql.to_string(),
1473 schema,
1474 input_types,
1475 },
1476 );
1477 }
1478
1479 pub fn drop_view(&self, name: &str) {
1481 self.views.lock().remove(name);
1482 }
1483
1484 pub(crate) fn view_schema(&self, name: &str) -> Option<CoreSchema> {
1485 self.views.lock().get(name).map(|view| view.schema.clone())
1486 }
1487
1488 pub(crate) fn view_definition(&self, name: &str) -> Option<ViewDef> {
1489 self.views.lock().get(name).cloned()
1490 }
1491
1492 pub async fn register(&self, name: &str) -> Result<()> {
1494 let db = self.db.clone().ok_or(MongrelQueryError::Core(
1495 mongreldb_core::MongrelError::NotFound("no primary table".into()),
1496 ))?;
1497 let provider = MongrelProvider::new(db.clone())?;
1498 self.tables.lock().insert(name.to_string(), db);
1499 self.ctx
1500 .register_table(name, Arc::new(provider))
1501 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1502 Ok(())
1503 }
1504
1505 pub async fn register_db(&self, name: &str, db: Table) -> Result<()> {
1511 let db_arc = Arc::new(Mutex::new(db));
1512 let provider = MongrelProvider::new(db_arc.clone())?;
1513 self.tables.lock().insert(name.to_string(), db_arc);
1514 self.ctx
1515 .register_table(name, Arc::new(provider))
1516 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1517 Ok(())
1518 }
1519
1520 fn refresh_registered_table(&self, db: &Arc<Database>, name: &str) -> Result<()> {
1521 self.ctx
1522 .deregister_table(name)
1523 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1524 let handle = db.table(name)?;
1525 let provider = MongrelProvider::new(handle.clone())?;
1526 self.ctx
1527 .register_table(name, Arc::new(provider))
1528 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1529 self.tables.lock().insert(name.to_string(), handle);
1530 Ok(())
1531 }
1532
1533 fn try_sqlite_master(&self, sql: &str) -> Result<Option<Vec<RecordBatch>>> {
1544 let lower = sql.to_ascii_lowercase();
1545 if !lower.contains("sqlite_master") && !lower.contains("sqlite_schema") {
1546 return Ok(None);
1547 }
1548 use arrow::array::{ArrayRef, Int64Array, StringArray};
1549 use arrow::datatypes::{DataType, Field, Schema};
1550 use arrow::record_batch::RecordBatch;
1551
1552 let mut types: Vec<String> = Vec::new();
1553 let mut names: Vec<String> = Vec::new();
1554 let mut tbl_names: Vec<String> = Vec::new();
1555
1556 for name in self.tables.lock().keys() {
1558 types.push("table".into());
1559 names.push(name.clone());
1560 tbl_names.push(name.clone());
1561 }
1562 for name in self.views.lock().keys() {
1564 types.push("view".into());
1565 names.push(name.clone());
1566 tbl_names.push(name.clone());
1567 }
1568 if let Some(db) = &self.database {
1570 for t in db.triggers() {
1571 let target_name = match &t.target {
1572 mongreldb_core::trigger::TriggerTarget::Table(n) | mongreldb_core::trigger::TriggerTarget::View(n) => n.clone(),
1573 };
1574 types.push("trigger".into());
1575 names.push(t.name.clone());
1576 tbl_names.push(target_name);
1577 }
1578 }
1579
1580 let schema = Arc::new(Schema::new(vec![
1581 Field::new("type", DataType::Utf8, false),
1582 Field::new("name", DataType::Utf8, false),
1583 Field::new("tbl_name", DataType::Utf8, false),
1584 Field::new("rootpage", DataType::Int64, false),
1585 Field::new("sql", DataType::Utf8, true),
1586 ]));
1587 let n = names.len();
1588 let rootpages: Vec<i64> = vec![0; n];
1589 let sqls: Vec<Option<&str>> = vec![None; n];
1590 let batch = RecordBatch::try_new(
1591 schema,
1592 vec![
1593 Arc::new(StringArray::from(types)) as ArrayRef,
1594 Arc::new(StringArray::from(names)) as ArrayRef,
1595 Arc::new(StringArray::from(tbl_names)) as ArrayRef,
1596 Arc::new(Int64Array::from(rootpages)) as ArrayRef,
1597 Arc::new(StringArray::from(sqls)) as ArrayRef,
1598 ],
1599 )
1600 .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
1601 Ok(Some(vec![batch]))
1602 }
1603
1604 fn try_direct_dispatch(&self, sql: &str) -> Result<Option<Vec<RecordBatch>>> {
1611 use arrow::array::ArrayRef;
1612 use mongreldb_core::Condition;
1613 use sqlparser::ast::{Expr, Query, SelectItem, SetExpr, Statement, TableFactor};
1614 use sqlparser::dialect::PostgreSqlDialect;
1615 use sqlparser::parser::Parser;
1616
1617 let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
1619 return Ok(None);
1620 };
1621 if stmts.len() != 1 {
1622 return Ok(None);
1623 }
1624 let Statement::Query(query) = stmts.into_iter().next().unwrap() else {
1625 return Ok(None);
1626 };
1627 let Query { body, .. } = *query;
1628 let select = match *body {
1629 SetExpr::Select(s) => *s,
1630 _ => return Ok(None),
1631 };
1632 let lower_sql = sql.to_lowercase();
1635 if lower_sql.contains(" limit ") || lower_sql.contains(" offset ") {
1636 return Ok(None);
1637 }
1638 use sqlparser::ast::GroupByExpr;
1640 if select.distinct.is_some()
1641 || !matches!(&select.group_by, GroupByExpr::Expressions(e, _) if e.is_empty())
1642 || select.having.is_some()
1643 || select.from.len() != 1
1644 || !select.from[0].joins.is_empty()
1645 {
1646 return Ok(None);
1647 }
1648 let table_name = match &select.from[0].relation {
1649 TableFactor::Table { name, .. } => Some(name.to_string()),
1650 _ => return Ok(None),
1651 };
1652 let Some(table_name) = table_name else {
1653 return Ok(None);
1654 };
1655
1656 if select.selection.is_none() {
1662 return Ok(None);
1663 }
1664
1665 let mut proj_names: Option<Vec<String>> = None;
1667 for item in &select.projection {
1668 match item {
1669 SelectItem::Wildcard(_) => {}
1670 SelectItem::UnnamedExpr(Expr::Identifier(ident)) => {
1671 proj_names
1672 .get_or_insert_with(Vec::new)
1673 .push(ident.value.clone());
1674 }
1675 SelectItem::UnnamedExpr(Expr::CompoundIdentifier(idents)) => {
1676 if let Some(last) = idents.last() {
1677 proj_names
1678 .get_or_insert_with(Vec::new)
1679 .push(last.value.clone());
1680 }
1681 }
1682 _ => return Ok(None),
1683 }
1684 }
1685
1686 let handle = match self.tables.lock().get(&table_name).cloned() {
1688 Some(h) => h,
1689 None => return Ok(None),
1690 };
1691
1692 let mut db = handle.lock();
1693 let schema = db.schema().clone();
1694 let conditions: Vec<Condition> = match &select.selection {
1697 Some(expr) => match translate_sqlparser_filter(expr, &schema) {
1698 Some(c) => c,
1699 None => return Ok(None),
1700 },
1701 None => Vec::new(),
1702 };
1703 if !conditions.is_empty() && db.ensure_indexes_complete().is_err() {
1704 return Ok(None);
1705 }
1706 let snap = db.snapshot();
1707
1708 let mut col_ids: Vec<u16> = Vec::new();
1710 let mut fields: Vec<arrow::datatypes::Field> = Vec::new();
1711 let resolve_col = |name: &str| -> Option<&mongreldb_core::schema::ColumnDef> {
1712 schema.columns.iter().find(|c| c.name == name)
1713 };
1714 match &proj_names {
1715 None => {
1716 for c in &schema.columns {
1717 col_ids.push(c.id);
1718 fields.push(arrow::datatypes::Field::new(
1719 &c.name,
1720 arrow_conv::arrow_data_type(&c.ty)?,
1721 c.flags.contains(mongreldb_core::ColumnFlags::NULLABLE),
1722 ));
1723 }
1724 }
1725 Some(names) => {
1726 for n in names {
1727 let cdef = match resolve_col(n) {
1728 Some(c) => c,
1729 None => return Ok(None), };
1731 col_ids.push(cdef.id);
1732 fields.push(arrow::datatypes::Field::new(
1733 &cdef.name,
1734 arrow_conv::arrow_data_type(&cdef.ty)?,
1735 cdef.flags.contains(mongreldb_core::ColumnFlags::NULLABLE),
1736 ));
1737 }
1738 }
1739 }
1740
1741 let cols = if !conditions.is_empty() {
1743 match db.query_columns_native_cached(&conditions, Some(&col_ids), snap) {
1744 Ok(Some(c)) => c,
1745 Ok(None) => db
1746 .visible_columns_native(snap, Some(&col_ids))
1747 .map_err(MongrelQueryError::Core)?,
1748 Err(_) => return Ok(None),
1749 }
1750 } else {
1751 db.visible_columns_native(snap, Some(&col_ids))
1752 .map_err(MongrelQueryError::Core)?
1753 };
1754 drop(db);
1755
1756 let mut arrays: Vec<ArrayRef> = Vec::with_capacity(col_ids.len());
1758 for cid in &col_ids {
1759 let col = cols
1760 .iter()
1761 .find(|(id, _)| id == cid)
1762 .map(|(_, c)| c.clone());
1763 let Some(col) = col else { return Ok(None) };
1764 let ty = schema
1765 .columns
1766 .iter()
1767 .find(|c| c.id == *cid)
1768 .map(|c| c.ty)
1769 .unwrap_or(mongreldb_core::schema::TypeId::Int64);
1770 arrays.push(arrow_conv::native_to_array(ty, &col)?);
1771 }
1772 let batch_schema = Arc::new(arrow::datatypes::Schema::new(fields));
1773 let batch = RecordBatch::try_new(batch_schema, arrays)
1774 .map_err(|e| MongrelQueryError::Arrow(format!("direct dispatch batch build: {e}")))?;
1775
1776 mongreldb_core::trace::QueryTrace::record(|t| {
1777 t.scan_mode = mongreldb_core::trace::ScanMode::DirectDispatch;
1778 t.planning_nanos = 0; });
1780 Ok(Some(vec![batch]))
1781 }
1782
1783 pub async fn run(&self, sql: &str) -> Result<Vec<RecordBatch>> {
1788 if let Some(inner) = strip_explain_query_plan(sql) {
1789 return self.explain_query_plan(inner).await;
1790 }
1791 if let Some(batches) = commands::try_run_command(self, sql).await? {
1792 return Ok(batches);
1793 }
1794 let lower = sql.trim_start().to_lowercase();
1796 if lower.starts_with("create table") {
1797 if let Some(db) = &self.database {
1798 let (name, schema) = parse_create_table(sql)?;
1799 db.create_table(&name, schema)?;
1800 let handle = db.table(&name)?;
1801 let provider = MongrelProvider::new(handle.clone())?;
1802 self.ctx
1803 .register_table(&name, Arc::new(provider))
1804 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1805 self.tables.lock().insert(name, handle);
1806 self.clear_cache();
1807 return Ok(Vec::new());
1808 }
1809 }
1810 if lower.starts_with("drop table") {
1811 if let Some(db) = &self.database {
1812 let (name, if_exists) = parse_drop_table(sql)?;
1813 let drop_result = db.drop_table(&name);
1814 if let Err(e) = drop_result {
1815 let is_not_found = matches!(e, mongreldb_core::MongrelError::NotFound(_));
1817 if !(if_exists && is_not_found) {
1818 return Err(e.into());
1819 }
1820 } else {
1821 self.ctx
1822 .deregister_table(&name)
1823 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1824 self.tables.lock().remove(&name);
1825 }
1826 self.clear_cache();
1827 return Ok(Vec::new());
1828 }
1829 }
1830 if lower.starts_with("alter table") {
1831 if let Some(db) = &self.database {
1832 match parse_alter_table(sql)? {
1833 ParsedAlterTable::RenameTable { old_name, new_name } => {
1834 db.rename_table(&old_name, &new_name)?;
1835 self.ctx
1839 .deregister_table(&old_name)
1840 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1841 self.tables.lock().remove(&old_name);
1842 let handle = db.table(&new_name)?;
1843 let provider = MongrelProvider::new(handle.clone())?;
1844 self.ctx
1845 .register_table(&new_name, Arc::new(provider))
1846 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1847 self.tables.lock().insert(new_name, handle);
1848 }
1849 ParsedAlterTable::RenameColumn {
1850 table_name,
1851 column_name,
1852 new_name,
1853 } => {
1854 db.alter_column(&table_name, &column_name, AlterColumn::rename(new_name))?;
1855 self.refresh_registered_table(db, &table_name)?;
1856 }
1857 ParsedAlterTable::AlterColumnType {
1858 table_name,
1859 column_name,
1860 ty,
1861 } => {
1862 db.alter_column(&table_name, &column_name, AlterColumn::set_type(ty))?;
1863 self.refresh_registered_table(db, &table_name)?;
1864 }
1865 ParsedAlterTable::SetNotNull {
1866 table_name,
1867 column_name,
1868 } => {
1869 let flags = current_column_flags(db, &table_name, &column_name)?
1870 .without(ColumnFlags::NULLABLE);
1871 db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
1872 self.refresh_registered_table(db, &table_name)?;
1873 }
1874 ParsedAlterTable::DropNotNull {
1875 table_name,
1876 column_name,
1877 } => {
1878 let flags = current_column_flags(db, &table_name, &column_name)?
1879 .with(ColumnFlags::NULLABLE);
1880 db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
1881 self.refresh_registered_table(db, &table_name)?;
1882 }
1883 }
1884 self.clear_cache();
1885 return Ok(Vec::new());
1886 }
1887 }
1888
1889 let resolved = self.resolve_view_sql(sql);
1892 let resolved = self.rewrite_external_module_compat_sql(&resolved);
1893 let resolved = rewrite_compat_function_calls(&resolved);
1894 let effective_sql = normalize_sql(&resolved);
1898 let sql = effective_sql.as_str();
1899 let epoch = self.cache_epoch();
1903 let key = (sql.to_string(), epoch);
1904 let result_cacheable = !extended_sql_functions::contains_volatile_extended_function(sql);
1905 if result_cacheable {
1906 if let Some(hit) = self.cache.lock().get(&key) {
1907 return Ok((**hit).clone());
1908 }
1909 }
1910 if let Some(batches) = self.try_sqlite_master(sql)? {
1913 if result_cacheable {
1914 self.cache.lock().insert(key, Arc::new(batches.clone()));
1915 }
1916 return Ok(batches);
1917 }
1918 if let Some(batches) = self.try_direct_dispatch(sql)? {
1923 if result_cacheable {
1924 self.cache.lock().insert(key, Arc::new(batches.clone()));
1925 }
1926 return Ok(batches);
1927 }
1928 let plan_start = std::time::Instant::now();
1930 let external_module_scan = self.query_references_external_module(sql);
1931 let df = {
1932 let cached_plan = self.plan_cache.lock().get(sql).cloned();
1933 if let Some(plan) = cached_plan {
1934 datafusion::dataframe::DataFrame::new(self.ctx.state(), plan)
1935 } else {
1936 let df = self
1937 .ctx
1938 .sql(sql)
1939 .await
1940 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1941 self.plan_cache
1942 .lock()
1943 .insert(sql.to_string(), df.logical_plan().clone());
1944 df
1945 }
1946 };
1947 let planning_nanos = plan_start.elapsed().as_nanos() as u64;
1950 mongreldb_core::trace::QueryTrace::record(|t| t.planning_nanos = planning_nanos);
1951
1952 let agg_key = sql_cache_key(sql);
1957 let batches = match self.try_native_aggregate(df.logical_plan(), agg_key) {
1958 Ok(Some(batch)) => vec![batch],
1959 _ => {
1960 match self.try_fk_join(df.logical_plan()) {
1964 Ok(Some(b)) => {
1965 mongreldb_core::trace::QueryTrace::record(|t| {
1967 t.join_mode = mongreldb_core::trace::JoinMode::FkBitmap;
1968 });
1969 b
1970 }
1971 _ => df
1972 .collect()
1973 .await
1974 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?,
1975 }
1976 }
1977 };
1978 if external_module_scan {
1979 mongreldb_core::trace::QueryTrace::record(|t| {
1980 t.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
1981 });
1982 }
1983 if result_cacheable {
1984 self.cache.lock().insert(key, Arc::new(batches.clone()));
1985 }
1986 Ok(batches)
1987 }
1988
1989 pub async fn run_sql_traced(
2001 &self,
2002 sql: &str,
2003 ) -> Result<(Vec<RecordBatch>, mongreldb_core::trace::QueryTrace)> {
2004 mongreldb_core::trace::QueryTrace::push_scope();
2005 let result = self.run(sql).await;
2006 let trace = mongreldb_core::trace::QueryTrace::pop_scope();
2007 Ok((result?, trace))
2008 }
2009
2010 pub fn clear_cache(&self) {
2013 self.cache.lock().clear();
2014 self.plan_cache.lock().clear();
2015 }
2016
2017 async fn explain_query_plan(&self, sql: &str) -> Result<Vec<RecordBatch>> {
2018 let explain_sql = format!("EXPLAIN {}", sql.trim().trim_end_matches(';'));
2019 let batches = self
2020 .ctx
2021 .sql(&explain_sql)
2022 .await
2023 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?
2024 .collect()
2025 .await
2026 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2027 let mut detail = self.mongrel_query_plan_details(sql);
2028 for batch in &batches {
2029 if batch.num_columns() < 2 {
2030 continue;
2031 }
2032 let Some(plan_type) = batch.column(0).as_any().downcast_ref::<StringArray>() else {
2033 continue;
2034 };
2035 let Some(plan) = batch.column(1).as_any().downcast_ref::<StringArray>() else {
2036 continue;
2037 };
2038 for row in 0..batch.num_rows() {
2039 let prefix = plan_type.value(row);
2040 for line in plan.value(row).lines() {
2041 let line = line.trim();
2042 if !line.is_empty() {
2043 detail.push(format!("DATAFUSION {prefix}: {line}"));
2044 }
2045 }
2046 }
2047 }
2048 if detail.is_empty() {
2049 detail.push("plan unavailable".to_string());
2050 }
2051 let ids = (0..detail.len()).map(|i| i as i64).collect::<Vec<_>>();
2052 let parents = vec![0_i64; detail.len()];
2053 let notused = vec![0_i64; detail.len()];
2054 let schema = Arc::new(arrow::datatypes::Schema::new(vec![
2055 arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int64, false),
2056 arrow::datatypes::Field::new("parent", arrow::datatypes::DataType::Int64, false),
2057 arrow::datatypes::Field::new("notused", arrow::datatypes::DataType::Int64, false),
2058 arrow::datatypes::Field::new("detail", arrow::datatypes::DataType::Utf8, false),
2059 ]));
2060 let batch = RecordBatch::try_new(
2061 schema,
2062 vec![
2063 Arc::new(Int64Array::from(ids)) as ArrayRef,
2064 Arc::new(Int64Array::from(parents)),
2065 Arc::new(Int64Array::from(notused)),
2066 Arc::new(StringArray::from(detail)),
2067 ],
2068 )
2069 .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
2070 Ok(vec![batch])
2071 }
2072
2073 fn mongrel_query_plan_details(&self, sql: &str) -> Vec<String> {
2074 use sqlparser::ast::{GroupByExpr, OrderByKind, SetExpr, Statement};
2075 use sqlparser::dialect::PostgreSqlDialect;
2076 use sqlparser::parser::Parser;
2077
2078 let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
2079 return Vec::new();
2080 };
2081 let Some(Statement::Query(query)) = stmts.first() else {
2082 return Vec::new();
2083 };
2084
2085 fn collect(session: &MongrelSession, query: &sqlparser::ast::Query, out: &mut Vec<String>) {
2086 use sqlparser::ast::{SetOperator, TableWithJoins};
2087 match query.body.as_ref() {
2088 SetExpr::Select(select) => {
2089 for TableWithJoins { relation, joins } in &select.from {
2090 session.push_table_plan(relation, select.selection.as_ref(), out);
2091 for join in joins {
2092 session.push_table_plan(&join.relation, None, out);
2093 }
2094 }
2095 if select.distinct.is_some() {
2096 out.push("USE TEMP B-TREE FOR DISTINCT".to_string());
2097 }
2098 let grouped = match &select.group_by {
2099 GroupByExpr::All(_) => true,
2100 GroupByExpr::Expressions(exprs, _) => !exprs.is_empty(),
2101 };
2102 if grouped {
2103 out.push("USE TEMP B-TREE FOR GROUP BY".to_string());
2104 }
2105 let ordered = query.order_by.as_ref().is_some_and(|order_by| {
2106 matches!(order_by.kind, OrderByKind::All(_))
2107 || matches!(&order_by.kind, OrderByKind::Expressions(exprs) if !exprs.is_empty())
2108 });
2109 if ordered {
2110 out.push("USE TEMP B-TREE FOR ORDER BY".to_string());
2111 }
2112 }
2113 SetExpr::Query(query) => collect(session, query, out),
2114 SetExpr::SetOperation {
2115 left, op, right, ..
2116 } => {
2117 let label = match op {
2118 SetOperator::Union => "COMPOUND QUERY UNION",
2119 SetOperator::Except => "COMPOUND QUERY EXCEPT",
2120 SetOperator::Intersect => "COMPOUND QUERY INTERSECT",
2121 _ => "COMPOUND QUERY",
2122 };
2123 out.push(label.to_string());
2124 collect_set_expr(session, left, out);
2125 collect_set_expr(session, right, out);
2126 }
2127 _ => {}
2128 }
2129 }
2130
2131 fn collect_set_expr(session: &MongrelSession, expr: &SetExpr, out: &mut Vec<String>) {
2132 match expr {
2133 SetExpr::Select(select) => {
2134 for table in &select.from {
2135 session.push_table_plan(&table.relation, select.selection.as_ref(), out);
2136 }
2137 }
2138 SetExpr::Query(query) => collect(session, query, out),
2139 SetExpr::SetOperation { left, right, .. } => {
2140 collect_set_expr(session, left, out);
2141 collect_set_expr(session, right, out);
2142 }
2143 _ => {}
2144 }
2145 }
2146
2147 let mut out = Vec::new();
2148 collect(self, query, &mut out);
2149 out
2150 }
2151
2152 fn push_table_plan(
2153 &self,
2154 relation: &sqlparser::ast::TableFactor,
2155 selection: Option<&sqlparser::ast::Expr>,
2156 out: &mut Vec<String>,
2157 ) {
2158 let sqlparser::ast::TableFactor::Table { name, alias, .. } = relation else {
2159 out.push("SCAN SUBQUERY".to_string());
2160 return;
2161 };
2162 let table_name = name.to_string();
2163 let display_name = alias
2164 .as_ref()
2165 .map(|alias| alias.name.value.clone())
2166 .unwrap_or_else(|| table_name.clone());
2167 let Some(handle) = self.tables.lock().get(&table_name).cloned() else {
2168 out.push(format!("SCAN {display_name}"));
2169 return;
2170 };
2171 let schema = handle.lock().schema().clone();
2172 let searchable = selection
2173 .and_then(|expr| translate_sqlparser_filter(expr, &schema))
2174 .is_some_and(|conditions| !conditions.is_empty());
2175 if searchable {
2176 out.push(format!("SEARCH {display_name} USING MONGREL INDEX"));
2177 } else {
2178 out.push(format!("SCAN {display_name}"));
2179 }
2180 }
2181
2182 fn resolve_view_sql(&self, sql: &str) -> String {
2187 let views = self.views.lock();
2188 if views.is_empty() {
2189 return sql.to_string();
2190 }
2191 let mut result = sql.to_string();
2192 for (name, view) in views.iter() {
2193 result = replace_from_view(&result, name, &view.sql);
2194 }
2195 result
2196 }
2197
2198 fn rewrite_external_module_compat_sql(&self, sql: &str) -> String {
2199 let Some(db) = &self.database else {
2200 return sql.to_string();
2201 };
2202 rewrite_fts_match_compat_sql(sql, db)
2203 }
2204
2205 fn query_references_external_module(&self, sql: &str) -> bool {
2206 use sqlparser::ast::Statement;
2207 use sqlparser::dialect::PostgreSqlDialect;
2208 use sqlparser::parser::Parser;
2209
2210 Parser::parse_sql(&PostgreSqlDialect {}, sql)
2211 .ok()
2212 .is_some_and(|statements| {
2213 statements.iter().any(|statement| match statement {
2214 Statement::Query(query) => self.query_uses_external_module(query),
2215 _ => false,
2216 })
2217 })
2218 }
2219
2220 fn query_uses_external_module(&self, query: &sqlparser::ast::Query) -> bool {
2221 self.set_expr_uses_external_module(query.body.as_ref())
2222 }
2223
2224 fn set_expr_uses_external_module(&self, expr: &sqlparser::ast::SetExpr) -> bool {
2225 use sqlparser::ast::SetExpr;
2226
2227 match expr {
2228 SetExpr::Select(select) => select
2229 .from
2230 .iter()
2231 .any(|table| self.table_with_joins_uses_external_module(table)),
2232 SetExpr::Query(query) => self.query_uses_external_module(query),
2233 SetExpr::SetOperation { left, right, .. } => {
2234 self.set_expr_uses_external_module(left)
2235 || self.set_expr_uses_external_module(right)
2236 }
2237 _ => false,
2238 }
2239 }
2240
2241 fn table_with_joins_uses_external_module(
2242 &self,
2243 table: &sqlparser::ast::TableWithJoins,
2244 ) -> bool {
2245 self.table_factor_uses_external_module(&table.relation)
2246 || table
2247 .joins
2248 .iter()
2249 .any(|join| self.table_factor_uses_external_module(&join.relation))
2250 }
2251
2252 fn table_factor_uses_external_module(&self, relation: &sqlparser::ast::TableFactor) -> bool {
2253 use sqlparser::ast::{Expr, TableFactor};
2254
2255 match relation {
2256 TableFactor::Table { name, args, .. } => {
2257 let table_name = name.to_string();
2258 self.database
2259 .as_ref()
2260 .is_some_and(|db| db.external_table(&table_name).is_some())
2261 || (args.is_some() && self.external_modules.contains(&table_name))
2262 }
2263 TableFactor::Function { name, .. } => self.external_modules.contains(&name.to_string()),
2264 TableFactor::TableFunction {
2265 expr: Expr::Function(func),
2266 ..
2267 } => self.external_modules.contains(&func.name.to_string()),
2268 TableFactor::Derived { subquery, .. } => self.query_uses_external_module(subquery),
2269 _ => false,
2270 }
2271 }
2272
2273 fn cache_epoch(&self) -> u64 {
2276 if let Some(db) = &self.database {
2277 db.visible_epoch().0
2278 } else {
2279 self.combined_epoch()
2280 }
2281 }
2282
2283 fn combined_epoch(&self) -> u64 {
2284 let primary = self.db.as_ref().expect("no primary table");
2285 let mut combined = primary.lock().snapshot().epoch.0;
2286 let tables = self.tables.lock();
2287 for arc in tables.values() {
2288 if !Arc::ptr_eq(arc, primary) {
2289 let e = arc.lock().snapshot().epoch.0;
2290 combined = combined.wrapping_mul(31).wrapping_add(e);
2291 }
2292 }
2293 combined
2294 }
2295
2296 fn try_native_aggregate(
2300 &self,
2301 plan: &datafusion::logical_expr::LogicalPlan,
2302 cache_key: u64,
2303 ) -> Result<Option<RecordBatch>> {
2304 let Some(primary) = self.db.as_ref() else {
2305 return Ok(None);
2306 };
2307 let mut db = primary.lock();
2308 let schema = db.schema().clone();
2309 let snap = db.snapshot();
2310 native_agg::try_native_aggregate(&mut db, &schema, snap, plan, cache_key)
2311 }
2312
2313 fn try_fk_join(
2317 &self,
2318 plan: &datafusion::logical_expr::LogicalPlan,
2319 ) -> Result<Option<Vec<RecordBatch>>> {
2320 let tables = self.tables.lock();
2321 fk_join::try_fk_join(&tables, plan)
2322 }
2323
2324 pub fn context(&self) -> &SessionContext {
2325 &self.ctx
2326 }
2327
2328 pub fn register_scalar_udf(&self, f: ScalarUDF) {
2334 self.ctx.register_udf(f);
2335 self.clear_cache();
2336 }
2337
2338 pub fn register_aggregate_udf(&self, f: AggregateUDF) {
2340 self.ctx.register_udaf(f);
2341 self.clear_cache();
2342 }
2343
2344 pub fn register_window_udf(&self, f: WindowUDF) {
2346 self.ctx.register_udwf(f);
2347 self.clear_cache();
2348 }
2349}
2350
2351fn register_mongrel_functions(
2352 ctx: &SessionContext,
2353 sql_fn_state: Arc<extended_sql_functions::ExtendedSqlState>,
2354) {
2355 ctx.register_udf(ScalarUDF::from(udf::AnnSearchUdf::new()));
2356 ctx.register_udf(ScalarUDF::from(udf::SparseMatchUdf::new()));
2357 ctx.register_udf(ScalarUDF::from(udf::RTreeIntersectsUdf::new()));
2358 for udaf in percentile::percentile_udafs() {
2359 ctx.register_udaf(udaf);
2360 }
2361 extended_sql_functions::register_extended_sql_functions_with_state(ctx, sql_fn_state);
2362}
2363
2364fn strip_explain_query_plan(sql: &str) -> Option<&str> {
2365 let trimmed = sql.trim_start();
2366 let lower = trimmed.to_ascii_lowercase();
2367 if !lower.starts_with("explain") {
2368 return None;
2369 }
2370 let after_explain = trimmed.get(7..)?.trim_start();
2371 let after_explain_lower = after_explain.to_ascii_lowercase();
2372 if !after_explain_lower.starts_with("query") {
2373 return None;
2374 }
2375 let after_query = after_explain.get(5..)?.trim_start();
2376 let after_query_lower = after_query.to_ascii_lowercase();
2377 if !after_query_lower.starts_with("plan") {
2378 return None;
2379 }
2380 Some(after_query.get(4..)?.trim_start())
2381}
2382
2383#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2384enum SqlCompatTokenKind {
2385 Ident,
2386 String,
2387 Dot,
2388 LParen,
2389 RParen,
2390 Comma,
2391}
2392
2393#[derive(Debug, Clone)]
2394struct SqlCompatToken {
2395 kind: SqlCompatTokenKind,
2396 raw: String,
2397 normalized: String,
2398 start: usize,
2399 end: usize,
2400}
2401
2402#[derive(Debug, Clone)]
2403struct FtsMatchBinding {
2404 query_ref: String,
2405}
2406
2407#[derive(Debug, Clone)]
2408struct SqlReplacement {
2409 start: usize,
2410 end: usize,
2411 replacement: String,
2412}
2413
2414fn rewrite_fts_match_compat_sql(sql: &str, db: &Database) -> String {
2415 let tokens = sql_compat_tokens(sql);
2416 if tokens.is_empty() {
2417 return sql.to_string();
2418 }
2419 let bindings = fts_match_bindings(sql, db, &tokens);
2420 if bindings.is_empty() {
2421 return sql.to_string();
2422 }
2423 let unique_refs = bindings
2424 .values()
2425 .map(|binding| binding.query_ref.as_str())
2426 .collect::<HashSet<_>>();
2427 let unique_binding = if unique_refs.len() == 1 {
2428 bindings.values().next().cloned()
2429 } else {
2430 None
2431 };
2432 let mut replacements = Vec::new();
2433 for (idx, token) in tokens.iter().enumerate() {
2434 if token.kind != SqlCompatTokenKind::Ident || token.normalized != "match" {
2435 continue;
2436 }
2437 let Some(rhs) = tokens.get(idx + 1) else {
2438 continue;
2439 };
2440 if rhs.kind != SqlCompatTokenKind::String {
2441 continue;
2442 }
2443 let Some((lhs_start, _lhs_end, query_ref)) =
2444 fts_match_lhs_query_ref(&tokens, idx, &bindings, unique_binding.as_ref())
2445 else {
2446 continue;
2447 };
2448 replacements.push(SqlReplacement {
2449 start: lhs_start,
2450 end: rhs.end,
2451 replacement: format!("{query_ref}.query = {}", rhs.raw),
2452 });
2453 }
2454 apply_sql_replacements(sql, &replacements)
2455}
2456
2457fn fts_match_lhs_query_ref(
2458 tokens: &[SqlCompatToken],
2459 match_idx: usize,
2460 bindings: &HashMap<String, FtsMatchBinding>,
2461 unique_binding: Option<&FtsMatchBinding>,
2462) -> Option<(usize, usize, String)> {
2463 if match_idx == 0 {
2464 return None;
2465 }
2466 let lhs = tokens.get(match_idx - 1)?;
2467 if lhs.kind != SqlCompatTokenKind::Ident {
2468 return None;
2469 }
2470
2471 if match_idx >= 3
2472 && tokens.get(match_idx - 2)?.kind == SqlCompatTokenKind::Dot
2473 && tokens.get(match_idx - 3)?.kind == SqlCompatTokenKind::Ident
2474 {
2475 let owner = tokens.get(match_idx - 3)?;
2476 let binding = bindings.get(&owner.normalized)?;
2477 if lhs.normalized == "query" || lhs.normalized == "text" {
2478 return Some((owner.start, lhs.end, binding.query_ref.clone()));
2479 }
2480 return None;
2481 }
2482
2483 if let Some(binding) = bindings.get(&lhs.normalized) {
2484 return Some((lhs.start, lhs.end, binding.query_ref.clone()));
2485 }
2486 if lhs.normalized == "text" {
2487 let binding = unique_binding?;
2488 return Some((lhs.start, lhs.end, binding.query_ref.clone()));
2489 }
2490 None
2491}
2492
2493fn fts_match_bindings(
2494 sql: &str,
2495 db: &Database,
2496 tokens: &[SqlCompatToken],
2497) -> HashMap<String, FtsMatchBinding> {
2498 let mut out = HashMap::new();
2499 let mut i = 0;
2500 while i < tokens.len() {
2501 let token = &tokens[i];
2502 let starts_table_ref = token.kind == SqlCompatTokenKind::Ident
2503 && matches!(token.normalized.as_str(), "from" | "join");
2504 if !starts_table_ref {
2505 i += 1;
2506 continue;
2507 }
2508 let mut table_idx = i + 1;
2509 if tokens
2510 .get(table_idx)
2511 .is_some_and(|token| token.kind == SqlCompatTokenKind::LParen)
2512 {
2513 i += 1;
2514 continue;
2515 }
2516 let Some(table) = tokens.get(table_idx) else {
2517 break;
2518 };
2519 if table.kind != SqlCompatTokenKind::Ident {
2520 i += 1;
2521 continue;
2522 }
2523 let mut table_name = table.normalized.clone();
2524 let mut table_ref = table.raw.clone();
2525 if tokens
2526 .get(table_idx + 1)
2527 .is_some_and(|token| token.kind == SqlCompatTokenKind::Dot)
2528 && tokens
2529 .get(table_idx + 2)
2530 .is_some_and(|token| token.kind == SqlCompatTokenKind::Ident)
2531 {
2532 let qualified = tokens.get(table_idx + 2).unwrap();
2533 table_name = qualified.normalized.clone();
2534 table_ref = sql[table.start..qualified.end].to_string();
2535 table_idx += 2;
2536 }
2537 if !is_fts_docs_table(db, &table_name) {
2538 i = table_idx + 1;
2539 continue;
2540 }
2541 let mut query_ref = table_ref.clone();
2542 let mut alias_key = None;
2543 let mut next = table_idx + 1;
2544 if tokens.get(next).is_some_and(|token| {
2545 token.kind == SqlCompatTokenKind::Ident && token.normalized == "as"
2546 }) {
2547 next += 1;
2548 }
2549 if let Some(alias) = tokens.get(next) {
2550 if alias.kind == SqlCompatTokenKind::Ident && !is_table_ref_boundary(&alias.normalized)
2551 {
2552 alias_key = Some(alias.normalized.clone());
2553 query_ref = alias.raw.clone();
2554 next += 1;
2555 }
2556 }
2557 out.insert(
2558 table_name,
2559 FtsMatchBinding {
2560 query_ref: query_ref.clone(),
2561 },
2562 );
2563 if let Some(alias_key) = alias_key {
2564 out.insert(alias_key, FtsMatchBinding { query_ref });
2565 }
2566 i = next;
2567 }
2568 out
2569}
2570
2571fn is_fts_docs_table(db: &Database, name: &str) -> bool {
2572 db.external_table(name)
2573 .is_some_and(|entry| entry.module == "fts_docs")
2574}
2575
2576fn is_table_ref_boundary(normalized: &str) -> bool {
2577 matches!(
2578 normalized,
2579 "where"
2580 | "join"
2581 | "left"
2582 | "right"
2583 | "inner"
2584 | "outer"
2585 | "full"
2586 | "cross"
2587 | "on"
2588 | "using"
2589 | "group"
2590 | "order"
2591 | "having"
2592 | "limit"
2593 | "offset"
2594 | "union"
2595 | "except"
2596 | "intersect"
2597 )
2598}
2599
2600fn sql_compat_tokens(sql: &str) -> Vec<SqlCompatToken> {
2601 let bytes = sql.as_bytes();
2602 let mut tokens = Vec::new();
2603 let mut i = 0;
2604 while i < bytes.len() {
2605 match bytes[i] {
2606 b if b.is_ascii_whitespace() => i += 1,
2607 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
2608 i += 2;
2609 while i < bytes.len() && bytes[i] != b'\n' {
2610 i += 1;
2611 }
2612 }
2613 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
2614 i = skip_block_comment(bytes, i);
2615 }
2616 b'\'' => {
2617 let end = skip_quoted(bytes, i, b'\'');
2618 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
2619 i = end;
2620 }
2621 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
2622 let end = skip_quoted(bytes, i + 1, b'\'');
2623 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
2624 i = end;
2625 }
2626 b'$' => {
2627 let (end, matched) = skip_dollar_quoted(bytes, i);
2628 if matched {
2629 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
2630 i = end;
2631 } else {
2632 i += 1;
2633 }
2634 }
2635 b'"' => {
2636 let end = skip_quoted(bytes, i, b'"');
2637 let raw = sql[i..end].to_string();
2638 let normalized = unquote_sql_ident(&raw).to_ascii_lowercase();
2639 tokens.push(SqlCompatToken {
2640 kind: SqlCompatTokenKind::Ident,
2641 raw,
2642 normalized,
2643 start: i,
2644 end,
2645 });
2646 i = end;
2647 }
2648 b'.' => {
2649 tokens.push(sql_token(sql, SqlCompatTokenKind::Dot, i, i + 1));
2650 i += 1;
2651 }
2652 b'(' => {
2653 tokens.push(sql_token(sql, SqlCompatTokenKind::LParen, i, i + 1));
2654 i += 1;
2655 }
2656 b')' => {
2657 tokens.push(sql_token(sql, SqlCompatTokenKind::RParen, i, i + 1));
2658 i += 1;
2659 }
2660 b',' => {
2661 tokens.push(sql_token(sql, SqlCompatTokenKind::Comma, i, i + 1));
2662 i += 1;
2663 }
2664 b if is_sql_ident_byte(b) => {
2665 let start = i;
2666 i += 1;
2667 while i < bytes.len() && is_sql_ident_byte(bytes[i]) {
2668 i += 1;
2669 }
2670 tokens.push(sql_token(sql, SqlCompatTokenKind::Ident, start, i));
2671 }
2672 _ => i += 1,
2673 }
2674 }
2675 tokens
2676}
2677
2678fn sql_token(sql: &str, kind: SqlCompatTokenKind, start: usize, end: usize) -> SqlCompatToken {
2679 let raw = sql[start..end].to_string();
2680 SqlCompatToken {
2681 kind,
2682 normalized: raw.to_ascii_lowercase(),
2683 raw,
2684 start,
2685 end,
2686 }
2687}
2688
2689fn unquote_sql_ident(raw: &str) -> String {
2690 if raw.len() >= 2 && raw.starts_with('"') && raw.ends_with('"') {
2691 raw[1..raw.len() - 1].replace("\"\"", "\"")
2692 } else {
2693 raw.to_string()
2694 }
2695}
2696
2697fn apply_sql_replacements(sql: &str, replacements: &[SqlReplacement]) -> String {
2698 if replacements.is_empty() {
2699 return sql.to_string();
2700 }
2701 let mut ordered = replacements.to_vec();
2702 ordered.sort_by_key(|replacement| replacement.start);
2703 let mut out = String::with_capacity(sql.len());
2704 let mut cursor = 0;
2705 for replacement in ordered {
2706 if replacement.start < cursor || replacement.end > sql.len() {
2707 continue;
2708 }
2709 out.push_str(&sql[cursor..replacement.start]);
2710 out.push_str(&replacement.replacement);
2711 cursor = replacement.end;
2712 }
2713 out.push_str(&sql[cursor..]);
2714 out
2715}
2716
2717fn rewrite_compat_function_calls(sql: &str) -> String {
2718 let bytes = sql.as_bytes();
2719 let mut out = String::with_capacity(sql.len());
2720 let mut i = 0;
2721 while i < bytes.len() {
2722 match bytes[i] {
2723 b'\'' => i = copy_quoted_to_string(&mut out, bytes, i, b'\''),
2724 b'"' => i = copy_quoted_to_string(&mut out, bytes, i, b'"'),
2725 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
2726 out.push(bytes[i] as char);
2727 i += 1;
2728 i = copy_quoted_to_string(&mut out, bytes, i, b'\'');
2729 }
2730 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
2731 out.push('-');
2732 out.push('-');
2733 i += 2;
2734 while i < bytes.len() {
2735 let ch = bytes[i] as char;
2736 out.push(ch);
2737 i += 1;
2738 if ch == '\n' {
2739 break;
2740 }
2741 }
2742 }
2743 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
2744 let start = i;
2745 i = skip_block_comment(bytes, i);
2746 out.push_str(&sql[start..i.min(bytes.len())]);
2747 }
2748 b'$' => {
2749 let start_len = out.len();
2750 let (next, matched) = copy_dollar_quoted_to_string(&mut out, bytes, i);
2751 if matched {
2752 i = next;
2753 } else {
2754 out.truncate(start_len);
2755 out.push('$');
2756 i += 1;
2757 }
2758 }
2759 b'g' | b'G' | b'm' | b'M' | b't' | b'T' => {
2760 if let Some((replacement, next)) = compat_function_rewrite_at(sql, i) {
2761 out.push_str(&replacement);
2762 i = next;
2763 } else {
2764 out.push(bytes[i] as char);
2765 i += 1;
2766 }
2767 }
2768 _ => {
2769 out.push(bytes[i] as char);
2770 i += 1;
2771 }
2772 }
2773 }
2774 out
2775}
2776
2777fn compat_function_rewrite_at(sql: &str, start: usize) -> Option<(String, usize)> {
2778 let bytes = sql.as_bytes();
2779 let (name, kind) = if ident_eq_at(bytes, start, b"max") {
2780 ("max", CompatRewriteKind::ScalarMax)
2781 } else if ident_eq_at(bytes, start, b"min") {
2782 ("min", CompatRewriteKind::ScalarMin)
2783 } else if ident_eq_at(bytes, start, b"group_concat") {
2784 ("group_concat", CompatRewriteKind::GroupConcat)
2785 } else if ident_eq_at(bytes, start, b"total") {
2786 ("total", CompatRewriteKind::Total)
2787 } else {
2788 return None;
2789 };
2790 let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
2791 let after_name = start + name.len();
2792 let after_ok = bytes
2793 .get(after_name)
2794 .is_some_and(|b| !is_sql_ident_byte(*b));
2795 if !before_ok || !after_ok {
2796 return None;
2797 }
2798 let mut open = after_name;
2799 while open < bytes.len() && bytes[open].is_ascii_whitespace() {
2800 open += 1;
2801 }
2802 if bytes.get(open) != Some(&b'(') {
2803 return None;
2804 }
2805 let summary = call_arg_summary(sql, open)?;
2806 match kind {
2807 CompatRewriteKind::ScalarMax if summary.top_level_commas > 0 => {
2808 Some(("__mongreldb_scalar_max(".to_string(), open + 1))
2809 }
2810 CompatRewriteKind::ScalarMin if summary.top_level_commas > 0 => {
2811 Some(("__mongreldb_scalar_min(".to_string(), open + 1))
2812 }
2813 CompatRewriteKind::GroupConcat => {
2814 let args = &sql[open + 1..summary.close];
2815 let rewritten = if summary.top_level_commas == 0 {
2816 format!("string_agg({args}, ',')")
2817 } else {
2818 format!("string_agg({args})")
2819 };
2820 Some((rewritten, summary.close + 1))
2821 }
2822 CompatRewriteKind::Total if summary.top_level_commas == 0 => {
2823 let args = &sql[open + 1..summary.close];
2824 let suffix_end = aggregate_suffix_end(sql, summary.close + 1);
2825 let suffix = &sql[summary.close + 1..suffix_end];
2826 Some((
2827 format!("coalesce(cast(sum({args}){suffix} as double), 0.0)"),
2828 suffix_end,
2829 ))
2830 }
2831 _ => None,
2832 }
2833}
2834
2835#[derive(Clone, Copy)]
2836enum CompatRewriteKind {
2837 ScalarMax,
2838 ScalarMin,
2839 GroupConcat,
2840 Total,
2841}
2842
2843fn ident_eq_at(bytes: &[u8], start: usize, ident: &[u8]) -> bool {
2844 bytes
2845 .get(start..start + ident.len())
2846 .is_some_and(|slice| slice.eq_ignore_ascii_case(ident))
2847}
2848
2849fn is_sql_ident_byte(b: u8) -> bool {
2850 b.is_ascii_alphanumeric() || b == b'_' || b == b'$'
2851}
2852
2853fn keyword_at(bytes: &[u8], start: usize, keyword: &[u8]) -> bool {
2854 if !ident_eq_at(bytes, start, keyword) {
2855 return false;
2856 }
2857 let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
2858 let after = start + keyword.len();
2859 let after_ok = after >= bytes.len() || !is_sql_ident_byte(bytes[after]);
2860 before_ok && after_ok
2861}
2862
2863fn skip_sql_whitespace(bytes: &[u8], mut i: usize) -> usize {
2864 while i < bytes.len() && bytes[i].is_ascii_whitespace() {
2865 i += 1;
2866 }
2867 i
2868}
2869
2870fn aggregate_suffix_end(sql: &str, start: usize) -> usize {
2871 let bytes = sql.as_bytes();
2872 let mut suffix_end = start;
2873 let mut i = skip_sql_whitespace(bytes, start);
2874
2875 if keyword_at(bytes, i, b"filter") {
2876 let open = skip_sql_whitespace(bytes, i + b"filter".len());
2877 if bytes.get(open) != Some(&b'(') {
2878 return start;
2879 }
2880 let Some(summary) = call_arg_summary(sql, open) else {
2881 return start;
2882 };
2883 suffix_end = summary.close + 1;
2884 i = skip_sql_whitespace(bytes, suffix_end);
2885 }
2886
2887 if keyword_at(bytes, i, b"over") {
2888 let after_over = skip_sql_whitespace(bytes, i + b"over".len());
2889 if bytes.get(after_over) == Some(&b'(') {
2890 let Some(summary) = call_arg_summary(sql, after_over) else {
2891 return suffix_end;
2892 };
2893 suffix_end = summary.close + 1;
2894 } else {
2895 let mut end = after_over;
2896 while end < bytes.len() && is_sql_ident_byte(bytes[end]) {
2897 end += 1;
2898 }
2899 if end > after_over {
2900 suffix_end = end;
2901 }
2902 }
2903 }
2904
2905 suffix_end
2906}
2907
2908struct CallArgSummary {
2909 close: usize,
2910 top_level_commas: usize,
2911}
2912
2913fn call_arg_summary(sql: &str, open: usize) -> Option<CallArgSummary> {
2914 let bytes = sql.as_bytes();
2915 let mut depth = 1;
2916 let mut i = open + 1;
2917 let mut top_level_commas = 0;
2918 while i < bytes.len() {
2919 match bytes[i] {
2920 b'\'' => i = skip_quoted(bytes, i, b'\''),
2921 b'"' => i = skip_quoted(bytes, i, b'"'),
2922 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
2923 i = skip_quoted(bytes, i + 1, b'\'')
2924 }
2925 b'$' => {
2926 let (next, matched) = skip_dollar_quoted(bytes, i);
2927 i = if matched { next } else { i + 1 };
2928 }
2929 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
2930 i += 2;
2931 while i < bytes.len() && bytes[i] != b'\n' {
2932 i += 1;
2933 }
2934 }
2935 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
2936 i = skip_block_comment(bytes, i);
2937 }
2938 b'(' => {
2939 depth += 1;
2940 i += 1;
2941 }
2942 b')' => {
2943 depth -= 1;
2944 if depth == 0 {
2945 return Some(CallArgSummary {
2946 close: i,
2947 top_level_commas,
2948 });
2949 }
2950 i += 1;
2951 }
2952 b',' if depth == 1 => {
2953 top_level_commas += 1;
2954 i += 1;
2955 }
2956 _ => i += 1,
2957 }
2958 }
2959 None
2960}
2961
2962fn copy_quoted_to_string(out: &mut String, bytes: &[u8], start: usize, delim: u8) -> usize {
2963 let end = skip_quoted(bytes, start, delim);
2964 out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
2965 end
2966}
2967
2968fn skip_quoted(bytes: &[u8], start: usize, delim: u8) -> usize {
2969 let mut i = start;
2970 if i < bytes.len() {
2971 i += 1;
2972 }
2973 while i < bytes.len() {
2974 if bytes[i] == delim {
2975 i += 1;
2976 if i < bytes.len() && bytes[i] == delim {
2977 i += 1;
2978 continue;
2979 }
2980 break;
2981 }
2982 i += 1;
2983 }
2984 i
2985}
2986
2987fn copy_dollar_quoted_to_string(out: &mut String, bytes: &[u8], start: usize) -> (usize, bool) {
2988 let (end, matched) = skip_dollar_quoted(bytes, start);
2989 if matched {
2990 out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
2991 }
2992 (end, matched)
2993}
2994
2995fn skip_dollar_quoted(bytes: &[u8], start: usize) -> (usize, bool) {
2996 if bytes.get(start) != Some(&b'$') {
2997 return (start, false);
2998 }
2999 let mut j = start + 1;
3000 while j < bytes.len() && (bytes[j].is_ascii_alphanumeric() || bytes[j] == b'_') {
3001 j += 1;
3002 }
3003 if bytes.get(j) != Some(&b'$') {
3004 return (start, false);
3005 }
3006 let tag = &bytes[start..=j];
3007 let mut i = j + 1;
3008 while i + tag.len() <= bytes.len() {
3009 if &bytes[i..i + tag.len()] == tag {
3010 return (i + tag.len(), true);
3011 }
3012 i += 1;
3013 }
3014 (start, false)
3015}
3016
3017fn skip_block_comment(bytes: &[u8], start: usize) -> usize {
3018 let mut i = start + 2;
3019 let mut depth = 1;
3020 while i + 1 < bytes.len() && depth > 0 {
3021 if bytes[i] == b'/' && bytes[i + 1] == b'*' {
3022 depth += 1;
3023 i += 2;
3024 } else if bytes[i] == b'*' && bytes[i + 1] == b'/' {
3025 depth -= 1;
3026 i += 2;
3027 } else {
3028 i += 1;
3029 }
3030 }
3031 i
3032}
3033
3034fn sql_cache_key(sql: &str) -> u64 {
3036 use std::hash::{Hash, Hasher};
3037 let mut h = std::collections::hash_map::DefaultHasher::new();
3038 sql.hash(&mut h);
3039 h.finish()
3040}
3041
3042fn replace_from_view(sql: &str, name: &str, view_sql: &str) -> String {
3049 let lower = sql.to_ascii_lowercase();
3050 let bytes = lower.as_bytes();
3051 let name_b = name.as_bytes();
3052 let mut i = 0usize;
3053 while let Some(rel) = lower[i..].find("from") {
3054 let from_start = i + rel;
3055 let after_from = from_start + 4;
3056 i = after_from;
3057 if from_start > 0 && is_ident_byte(bytes[from_start - 1]) {
3059 continue;
3060 }
3061 let mut j = after_from;
3063 while j < bytes.len() && bytes[j].is_ascii_whitespace() {
3064 j += 1;
3065 }
3066 if j == after_from || !bytes[j..].starts_with(name_b) {
3067 continue;
3068 }
3069 let after_name = j + name_b.len();
3070 if after_name < bytes.len() && is_ident_byte(bytes[after_name]) {
3072 continue;
3073 }
3074 let mut out = String::with_capacity(sql.len() + view_sql.len() + name.len() + 8);
3077 out.push_str(&sql[..from_start]);
3078 out.push_str(&sql[from_start..j]);
3079 out.push('(');
3080 out.push_str(view_sql);
3081 out.push_str(") AS ");
3082 out.push_str(name);
3083 out.push_str(&sql[after_name..]);
3084 return out;
3085 }
3086 sql.to_string()
3087}
3088
3089fn is_ident_byte(b: u8) -> bool {
3090 b.is_ascii_alphanumeric() || b == b'_'
3091}
3092
3093fn normalize_sql(sql: &str) -> String {
3103 let b = sql.as_bytes();
3104 let n = b.len();
3105 let mut out: Vec<u8> = Vec::with_capacity(n);
3106 let mut want_space = false;
3109 let mut i = 0usize;
3110 while i < n {
3111 let c = b[i];
3112 if c.is_ascii_whitespace() {
3116 want_space = true;
3117 i += 1;
3118 continue;
3119 }
3120 if c == b'-' && i + 1 < n && b[i + 1] == b'-' {
3121 i += 2;
3123 while i < n && b[i] != b'\n' {
3124 i += 1;
3125 }
3126 want_space = !out.is_empty();
3127 continue;
3128 }
3129 if c == b'/' && i + 1 < n && b[i + 1] == b'*' {
3130 i += 2;
3133 let mut depth = 1usize;
3134 while i + 1 < n && depth > 0 {
3135 if b[i] == b'/' && b[i + 1] == b'*' {
3136 depth += 1;
3137 i += 2;
3138 } else if b[i] == b'*' && b[i + 1] == b'/' {
3139 depth -= 1;
3140 i += 2;
3141 } else {
3142 i += 1;
3143 }
3144 }
3145 want_space = !out.is_empty();
3146 continue;
3147 }
3148 if want_space && !out.is_empty() {
3150 out.push(b' ');
3151 }
3152 want_space = false;
3153 match c {
3154 b'E' | b'e' if i + 1 < n && b[i + 1] == b'\'' => {
3156 out.push(c);
3157 i += 1;
3158 i = copy_quoted(&mut out, b, i, n, b'\'');
3159 continue;
3160 }
3161 b'\'' => {
3163 i = copy_quoted(&mut out, b, i, n, b'\'');
3164 continue;
3165 }
3166 b'"' => {
3168 i = copy_quoted(&mut out, b, i, n, b'"');
3169 continue;
3170 }
3171 b'$' => {
3173 let (consumed, matched) = copy_dollar_quoted(&mut out, b, i, n);
3174 if matched {
3175 i = consumed;
3176 continue;
3177 }
3178 out.push(c);
3179 i += 1;
3180 continue;
3181 }
3182 _ => {
3183 out.push(c);
3184 i += 1;
3185 }
3186 }
3187 }
3188 String::from_utf8(out).unwrap_or_else(|_| sql.to_string())
3189}
3190
3191fn copy_quoted(out: &mut Vec<u8>, b: &[u8], start: usize, n: usize, delim: u8) -> usize {
3195 out.push(b[start]);
3196 let mut i = start + 1;
3197 while i < n {
3198 let c = b[i];
3199 out.push(c);
3200 if c == delim {
3201 if i + 1 < n && b[i + 1] == delim {
3203 out.push(b[i + 1]);
3204 i += 2;
3205 continue;
3206 }
3207 return i + 1;
3208 }
3209 i += 1;
3210 }
3211 i
3212}
3213
3214fn copy_dollar_quoted(out: &mut Vec<u8>, b: &[u8], start: usize, n: usize) -> (usize, bool) {
3218 let mut j = start + 1;
3222 let tag_start = j;
3223 while j < n && b[j] != b'$' && is_dollar_tag_byte(b[j]) {
3224 j += 1;
3225 }
3226 if j >= n || b[j] != b'$' {
3227 return (start + 1, false);
3228 }
3229 if tag_start < j && !(b[tag_start].is_ascii_alphabetic() || b[tag_start] == b'_') {
3230 return (start + 1, false);
3231 }
3232 let close_end = j + 1; let delim = &b[start..close_end];
3234 out.extend_from_slice(delim);
3236 let mut k = close_end;
3238 while k + delim.len() <= n {
3239 if &b[k..k + delim.len()] == delim {
3240 out.extend_from_slice(delim);
3241 return (k + delim.len(), true);
3242 }
3243 out.push(b[k]);
3244 k += 1;
3245 }
3246 out.extend_from_slice(&b[close_end..n]);
3248 (n, true)
3249}
3250
3251fn is_dollar_tag_byte(b: u8) -> bool {
3252 b.is_ascii_alphanumeric() || b == b'_'
3253}
3254
3255fn strip_prefix_ci<'a>(s: &'a str, prefix: &str) -> Option<&'a str> {
3257 let bytes = s.as_bytes();
3258 let pb = prefix.as_bytes();
3259 if bytes.len() >= pb.len() && bytes[..pb.len()].eq_ignore_ascii_case(pb) {
3260 Some(&s[pb.len()..])
3261 } else {
3262 None
3263 }
3264}
3265
3266const COLUMN_CONSTRAINTS: &[(&str, u32)] = &[
3279 ("primary key", ColumnFlags::PRIMARY_KEY),
3280 ("autoincrement", ColumnFlags::AUTO_INCREMENT),
3287 ("auto_increment", ColumnFlags::AUTO_INCREMENT),
3288];
3289
3290fn parse_column_constraints(constraint_text: &str) -> ColumnFlags {
3296 let normalized = constraint_text.to_lowercase();
3297 let mut flags = ColumnFlags::empty();
3298 for (phrase, bit) in COLUMN_CONSTRAINTS {
3299 if normalized.contains(phrase) {
3300 flags = flags.with(*bit);
3301 }
3302 }
3303 flags
3304}
3305
3306fn parse_sql_type(ty_str: &str) -> Result<mongreldb_core::schema::TypeId> {
3307 use mongreldb_core::schema::TypeId;
3308
3309 match ty_str.trim().trim_end_matches(';').to_lowercase().as_str() {
3310 "bigint" | "int8" | "int64" | "integer" | "int" => Ok(TypeId::Int64),
3311 "double" | "float8" | "float64" | "real" | "float" => Ok(TypeId::Float64),
3312 "varchar" | "text" | "string" | "bytes" => Ok(TypeId::Bytes),
3313 "boolean" | "bool" => Ok(TypeId::Bool),
3314 other => Err(MongrelQueryError::Schema(format!(
3315 "unsupported column type: {other}"
3316 ))),
3317 }
3318}
3319
3320fn parse_create_table(sql: &str) -> Result<(String, mongreldb_core::schema::Schema)> {
3326 use mongreldb_core::schema::*;
3327
3328 let open = sql
3329 .find('(')
3330 .ok_or(MongrelQueryError::Schema("CREATE TABLE missing '('".into()))?;
3331 let close = sql
3332 .rfind(')')
3333 .ok_or(MongrelQueryError::Schema("CREATE TABLE missing ')'".into()))?;
3334 let head = sql[..open].trim();
3335 let after_kw = strip_prefix_ci(head, "CREATE TABLE")
3336 .or_else(|| strip_prefix_ci(head, "create table"))
3337 .unwrap_or("")
3338 .trim();
3339 let after_kw = after_kw
3341 .strip_prefix("IF NOT EXISTS")
3342 .or_else(|| after_kw.strip_prefix("if not exists"))
3343 .map(str::trim)
3344 .unwrap_or(after_kw);
3345 let name = after_kw.trim_matches('"').to_string();
3346 if name.is_empty() {
3347 return Err(MongrelQueryError::Schema(
3348 "CREATE TABLE missing table name".into(),
3349 ));
3350 }
3351
3352 let body = &sql[open + 1..close];
3353 let mut columns = Vec::new();
3354 let schema_id: u64 = 0; for (i, raw) in body.split(',').enumerate() {
3356 let part = raw.trim();
3357 if part.is_empty() {
3358 continue;
3359 }
3360 let mut tokens = part.split_whitespace();
3361 let col_name = tokens
3362 .next()
3363 .ok_or(MongrelQueryError::Schema("missing column name".into()))?
3364 .trim_matches('"');
3365 let ty_str = tokens
3366 .next()
3367 .ok_or(MongrelQueryError::Schema("missing column type".into()))?
3368 .to_lowercase();
3369 let ty = parse_sql_type(&ty_str)?;
3370 let constraint_clause: String = tokens.collect::<Vec<_>>().join(" ");
3374 let flags = parse_column_constraints(&constraint_clause);
3375 columns.push(ColumnDef {
3376 id: (i + 1) as u16,
3377 name: col_name.to_string(),
3378 ty,
3379 flags,
3380 });
3381 }
3382
3383 Ok((
3384 name,
3385 Schema {
3386 schema_id,
3387 columns,
3388 indexes: vec![],
3389 colocation: vec![],
3390 constraints: Default::default(),
3391 },
3392 ))
3393}
3394
3395fn parse_drop_table(sql: &str) -> Result<(String, bool)> {
3397 let head = sql.trim();
3398 let after_kw = strip_prefix_ci(head, "DROP TABLE")
3399 .or_else(|| strip_prefix_ci(head, "drop table"))
3400 .unwrap_or("")
3401 .trim();
3402 let (rest, if_exists) = if let Some(r) = after_kw
3404 .strip_prefix("IF EXISTS")
3405 .or_else(|| after_kw.strip_prefix("if exists"))
3406 .map(str::trim)
3407 {
3408 (r, true)
3409 } else {
3410 (after_kw, false)
3411 };
3412 let name = rest.trim_matches(';').trim_matches('"').trim();
3413 if name.is_empty() {
3414 return Err(MongrelQueryError::Schema(
3415 "DROP TABLE missing table name".into(),
3416 ));
3417 }
3418 Ok((name.to_string(), if_exists))
3419}
3420
3421enum ParsedAlterTable {
3422 RenameTable {
3423 old_name: String,
3424 new_name: String,
3425 },
3426 RenameColumn {
3427 table_name: String,
3428 column_name: String,
3429 new_name: String,
3430 },
3431 AlterColumnType {
3432 table_name: String,
3433 column_name: String,
3434 ty: mongreldb_core::schema::TypeId,
3435 },
3436 SetNotNull {
3437 table_name: String,
3438 column_name: String,
3439 },
3440 DropNotNull {
3441 table_name: String,
3442 column_name: String,
3443 },
3444}
3445
3446fn current_column_flags(db: &Arc<Database>, table: &str, column: &str) -> Result<ColumnFlags> {
3447 let handle = db.table(table)?;
3448 let table = handle.lock();
3449 table
3450 .schema()
3451 .column(column)
3452 .map(|c| c.flags)
3453 .ok_or_else(|| MongrelQueryError::Schema(format!("unknown column {column}")))
3454}
3455
3456fn parse_alter_table(sql: &str) -> Result<ParsedAlterTable> {
3457 let trimmed = strip_statement_semicolon(sql.trim());
3458 let after_kw = strip_prefix_ci(trimmed, "ALTER TABLE")
3459 .ok_or_else(|| MongrelQueryError::Schema("not an ALTER TABLE statement".into()))?
3460 .trim();
3461 let (table_name, rest) = take_sql_ident(after_kw, "ALTER TABLE missing table name")?;
3462 let rest = rest.trim();
3463
3464 if let Some(after) = strip_prefix_ci(rest, "RENAME TO") {
3465 let new_name = parse_trailing_identifier(after, "ALTER TABLE missing new table name")?;
3466 return Ok(ParsedAlterTable::RenameTable {
3467 old_name: table_name,
3468 new_name,
3469 });
3470 }
3471
3472 if let Some(after) = strip_prefix_ci(rest, "RENAME COLUMN") {
3473 let (column_name, after_col) =
3474 take_sql_ident(after, "ALTER TABLE RENAME COLUMN missing column name")?;
3475 let after_to = strip_prefix_ci(after_col.trim(), "TO").ok_or_else(|| {
3476 MongrelQueryError::Schema("ALTER TABLE RENAME COLUMN missing TO".into())
3477 })?;
3478 let new_name = parse_trailing_identifier(
3479 after_to,
3480 "ALTER TABLE RENAME COLUMN missing new column name",
3481 )?;
3482 return Ok(ParsedAlterTable::RenameColumn {
3483 table_name,
3484 column_name,
3485 new_name,
3486 });
3487 }
3488
3489 let after_alter = strip_prefix_ci(rest, "ALTER COLUMN")
3490 .or_else(|| strip_prefix_ci(rest, "ALTER"))
3491 .ok_or_else(|| {
3492 MongrelQueryError::Schema(
3493 "ALTER TABLE must be RENAME TO, RENAME COLUMN, or ALTER COLUMN".into(),
3494 )
3495 })?;
3496 let (column_name, action) =
3497 take_sql_ident(after_alter, "ALTER TABLE ALTER COLUMN missing column name")?;
3498 let action = action.trim();
3499
3500 if let Some(after_type) =
3501 strip_prefix_ci(action, "TYPE").or_else(|| strip_prefix_ci(action, "SET DATA TYPE"))
3502 {
3503 let ty = parse_type_tail(after_type)?;
3504 return Ok(ParsedAlterTable::AlterColumnType {
3505 table_name,
3506 column_name,
3507 ty,
3508 });
3509 }
3510 if strip_prefix_ci(action, "SET NOT NULL").is_some() {
3511 return Ok(ParsedAlterTable::SetNotNull {
3512 table_name,
3513 column_name,
3514 });
3515 }
3516 if strip_prefix_ci(action, "DROP NOT NULL").is_some() {
3517 return Ok(ParsedAlterTable::DropNotNull {
3518 table_name,
3519 column_name,
3520 });
3521 }
3522
3523 Err(MongrelQueryError::Schema(
3524 "unsupported ALTER COLUMN action".into(),
3525 ))
3526}
3527
3528fn strip_statement_semicolon(s: &str) -> &str {
3529 s.trim().trim_end_matches(';').trim()
3530}
3531
3532fn take_sql_ident<'a>(s: &'a str, missing: &str) -> Result<(String, &'a str)> {
3533 let s = s.trim();
3534 if s.is_empty() {
3535 return Err(MongrelQueryError::Schema(missing.into()));
3536 }
3537 if let Some(rest) = s.strip_prefix('"') {
3538 let Some(end) = rest.find('"') else {
3539 return Err(MongrelQueryError::Schema(
3540 "unterminated quoted identifier".into(),
3541 ));
3542 };
3543 let ident = rest[..end].to_string();
3544 if ident.is_empty() {
3545 return Err(MongrelQueryError::Schema(missing.into()));
3546 }
3547 return Ok((ident, &rest[end + 1..]));
3548 }
3549 let end = s.find(|c: char| c.is_ascii_whitespace()).unwrap_or(s.len());
3550 let ident = s[..end].trim_matches('"').to_string();
3551 if ident.is_empty() {
3552 return Err(MongrelQueryError::Schema(missing.into()));
3553 }
3554 Ok((ident, &s[end..]))
3555}
3556
3557fn parse_trailing_identifier(s: &str, missing: &str) -> Result<String> {
3558 let (ident, rest) = take_sql_ident(s, missing)?;
3559 if !strip_statement_semicolon(rest).is_empty() {
3560 return Err(MongrelQueryError::Schema(
3561 "unexpected tokens after identifier".into(),
3562 ));
3563 }
3564 Ok(ident)
3565}
3566
3567fn parse_type_tail(s: &str) -> Result<mongreldb_core::schema::TypeId> {
3568 let tail = strip_statement_semicolon(s);
3569 let ty = tail
3570 .split_whitespace()
3571 .next()
3572 .ok_or_else(|| MongrelQueryError::Schema("ALTER COLUMN TYPE missing type".into()))?;
3573 parse_sql_type(ty)
3574}
3575
3576#[cfg(test)]
3577mod tests {
3578 use super::*;
3579
3580 #[test]
3581 fn normalize_collapses_and_trims_whitespace() {
3582 assert_eq!(normalize_sql("SELECT * FROM t"), "SELECT * FROM t");
3583 assert_eq!(normalize_sql(" SELECT * FROM t "), "SELECT * FROM t");
3584 assert_eq!(
3585 normalize_sql("\n\tSELECT\n*\nFROM\n\tt\n"),
3586 "SELECT * FROM t"
3587 );
3588 assert_eq!(
3589 normalize_sql("SELECT a, b FROM t"),
3590 normalize_sql("SELECT a, b FROM t")
3591 );
3592 }
3593
3594 #[test]
3595 fn normalize_preserves_string_literal_whitespace() {
3596 assert_eq!(
3597 normalize_sql("SELECT 'hello world' FROM t"),
3598 "SELECT 'hello world' FROM t"
3599 );
3600 assert_eq!(
3601 normalize_sql("SELECT 'it''s ok' FROM t"),
3602 "SELECT 'it''s ok' FROM t"
3603 );
3604 assert_eq!(
3605 normalize_sql(" SELECT 'a b' FROM t "),
3606 "SELECT 'a b' FROM t"
3607 );
3608 }
3609
3610 #[test]
3611 fn normalize_preserves_quoted_identifier_and_dollar_quote() {
3612 assert_eq!(
3613 normalize_sql(" SELECT \"my col\" FROM t "),
3614 "SELECT \"my col\" FROM t"
3615 );
3616 assert_eq!(
3617 normalize_sql(" SELECT $$a b$$ FROM t "),
3618 "SELECT $$a b$$ FROM t"
3619 );
3620 assert_eq!(
3621 normalize_sql("SELECT $tag$body with spaces$tag$ FROM t"),
3622 "SELECT $tag$body with spaces$tag$ FROM t"
3623 );
3624 }
3625
3626 #[test]
3627 fn normalize_strips_comments() {
3628 assert_eq!(
3629 normalize_sql("SELECT 1 -- trailing comment\nFROM t"),
3630 "SELECT 1 FROM t"
3631 );
3632 assert_eq!(
3633 normalize_sql("SELECT /* block */ 1 FROM t"),
3634 "SELECT 1 FROM t"
3635 );
3636 assert_eq!(
3638 normalize_sql("SELECT /* 'not a string' */ 1 FROM t"),
3639 "SELECT 1 FROM t"
3640 );
3641 assert_eq!(
3643 normalize_sql("SELECT /* outer /* inner */ still outer */ 1 FROM t"),
3644 "SELECT 1 FROM t"
3645 );
3646 }
3647
3648 #[test]
3649 fn normalize_escape_string_preserved() {
3650 assert_eq!(
3651 normalize_sql("SELECT E'line\\nbreak' FROM t"),
3652 "SELECT E'line\\nbreak' FROM t"
3653 );
3654 }
3655
3656 #[test]
3657 fn replace_from_view_matches_whole_word_only() {
3658 let out = replace_from_view("SELECT * FROM logs", "log", "SELECT 1");
3659 assert_eq!(out, "SELECT * FROM logs");
3660
3661 let out = replace_from_view("SELECT * FROM log", "log", "SELECT 1");
3662 assert_eq!(out, "SELECT * FROM (SELECT 1) AS log");
3663
3664 let out = replace_from_view("select * from log where x", "log", "SELECT 1");
3665 assert_eq!(out, "select * from (SELECT 1) AS log where x");
3666
3667 let out = replace_from_view("SELECT * FROM log)", "log", "SELECT 1");
3668 assert_eq!(out, "SELECT * FROM (SELECT 1) AS log)");
3669
3670 let out = replace_from_view("SELECT * xfrom log", "log", "SELECT 1");
3671 assert_eq!(out, "SELECT * xfrom log");
3672 }
3673
3674 #[test]
3675 fn compat_function_rewrite_handles_sqlite_compatibility_calls() {
3676 assert_eq!(
3677 rewrite_compat_function_calls("select max(id), min(id) from t"),
3678 "select max(id), min(id) from t"
3679 );
3680 assert_eq!(
3681 rewrite_compat_function_calls("select max(1, min(2, 3), 'max(4,5)')"),
3682 "select __mongreldb_scalar_max(1, __mongreldb_scalar_min(2, 3), 'max(4,5)')"
3683 );
3684 assert_eq!(
3685 rewrite_compat_function_calls("select /* max(1,2) */ min(1, (2 + 3))"),
3686 "select /* max(1,2) */ __mongreldb_scalar_min(1, (2 + 3))"
3687 );
3688 assert_eq!(
3689 rewrite_compat_function_calls("select max_value, min_value from t"),
3690 "select max_value, min_value from t"
3691 );
3692 assert_eq!(
3693 rewrite_compat_function_calls(
3694 "select group_concat(label), group_concat(label, '|') from t"
3695 ),
3696 "select string_agg(label, ','), string_agg(label, '|') from t"
3697 );
3698 assert_eq!(
3699 rewrite_compat_function_calls("select total(val), total(val) filter (where grp = 2) from t"),
3700 "select coalesce(cast(sum(val) as double), 0.0), coalesce(cast(sum(val) filter (where grp = 2) as double), 0.0) from t"
3701 );
3702 assert_eq!(
3703 rewrite_compat_function_calls(
3704 "select total(val) over (partition by grp order by id) from t"
3705 ),
3706 "select coalesce(cast(sum(val) over (partition by grp order by id) as double), 0.0) from t"
3707 );
3708 }
3709}