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_catalog_introspection(&self, sql: &str) -> Result<Option<Vec<RecordBatch>>> {
1541 let lower = sql.to_ascii_lowercase();
1542 if !lower.contains("information_schema.tables") {
1543 return Ok(None);
1544 }
1545 use arrow::array::{ArrayRef, Int64Array, StringArray};
1546 use arrow::datatypes::{DataType, Field, Schema};
1547 use arrow::record_batch::RecordBatch;
1548
1549 let mut types: Vec<String> = Vec::new();
1550 let mut names: Vec<String> = Vec::new();
1551 let mut tbl_names: Vec<String> = Vec::new();
1552
1553 for name in self.tables.lock().keys() {
1555 types.push("table".into());
1556 names.push(name.clone());
1557 tbl_names.push(name.clone());
1558 }
1559 for name in self.views.lock().keys() {
1561 types.push("view".into());
1562 names.push(name.clone());
1563 tbl_names.push(name.clone());
1564 }
1565 if let Some(db) = &self.database {
1567 for t in db.triggers() {
1568 let target_name = match &t.target {
1569 mongreldb_core::trigger::TriggerTarget::Table(n)
1570 | mongreldb_core::trigger::TriggerTarget::View(n) => n.clone(),
1571 };
1572 types.push("trigger".into());
1573 names.push(t.name.clone());
1574 tbl_names.push(target_name);
1575 }
1576 }
1577
1578 let schema = Arc::new(Schema::new(vec![
1579 Field::new("type", DataType::Utf8, false),
1580 Field::new("name", DataType::Utf8, false),
1581 Field::new("tbl_name", DataType::Utf8, false),
1582 Field::new("rootpage", DataType::Int64, false),
1583 Field::new("sql", DataType::Utf8, true),
1584 ]));
1585 let n = names.len();
1586 let rootpages: Vec<i64> = vec![0; n];
1587 let sqls: Vec<Option<&str>> = vec![None; n];
1588 let batch = RecordBatch::try_new(
1589 schema,
1590 vec![
1591 Arc::new(StringArray::from(types)) as ArrayRef,
1592 Arc::new(StringArray::from(names)) as ArrayRef,
1593 Arc::new(StringArray::from(tbl_names)) as ArrayRef,
1594 Arc::new(Int64Array::from(rootpages)) as ArrayRef,
1595 Arc::new(StringArray::from(sqls)) as ArrayRef,
1596 ],
1597 )
1598 .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
1599 Ok(Some(vec![batch]))
1600 }
1601
1602 fn try_direct_dispatch(&self, sql: &str) -> Result<Option<Vec<RecordBatch>>> {
1609 use arrow::array::ArrayRef;
1610 use mongreldb_core::Condition;
1611 use sqlparser::ast::{Expr, Query, SelectItem, SetExpr, Statement, TableFactor};
1612 use sqlparser::dialect::PostgreSqlDialect;
1613 use sqlparser::parser::Parser;
1614
1615 let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
1617 return Ok(None);
1618 };
1619 if stmts.len() != 1 {
1620 return Ok(None);
1621 }
1622 let Statement::Query(query) = stmts.into_iter().next().unwrap() else {
1623 return Ok(None);
1624 };
1625 let Query { body, .. } = *query;
1626 let select = match *body {
1627 SetExpr::Select(s) => *s,
1628 _ => return Ok(None),
1629 };
1630 let lower_sql = sql.to_lowercase();
1633 if lower_sql.contains(" limit ") || lower_sql.contains(" offset ") {
1634 return Ok(None);
1635 }
1636 use sqlparser::ast::GroupByExpr;
1638 if select.distinct.is_some()
1639 || !matches!(&select.group_by, GroupByExpr::Expressions(e, _) if e.is_empty())
1640 || select.having.is_some()
1641 || select.from.len() != 1
1642 || !select.from[0].joins.is_empty()
1643 {
1644 return Ok(None);
1645 }
1646 let table_name = match &select.from[0].relation {
1647 TableFactor::Table { name, .. } => Some(name.to_string()),
1648 _ => return Ok(None),
1649 };
1650 let Some(table_name) = table_name else {
1651 return Ok(None);
1652 };
1653
1654 if select.selection.is_none() {
1660 return Ok(None);
1661 }
1662
1663 let mut proj_names: Option<Vec<String>> = None;
1665 for item in &select.projection {
1666 match item {
1667 SelectItem::Wildcard(_) => {}
1668 SelectItem::UnnamedExpr(Expr::Identifier(ident)) => {
1669 proj_names
1670 .get_or_insert_with(Vec::new)
1671 .push(ident.value.clone());
1672 }
1673 SelectItem::UnnamedExpr(Expr::CompoundIdentifier(idents)) => {
1674 if let Some(last) = idents.last() {
1675 proj_names
1676 .get_or_insert_with(Vec::new)
1677 .push(last.value.clone());
1678 }
1679 }
1680 _ => return Ok(None),
1681 }
1682 }
1683
1684 let handle = match self.tables.lock().get(&table_name).cloned() {
1686 Some(h) => h,
1687 None => return Ok(None),
1688 };
1689
1690 let mut db = handle.lock();
1691 let schema = db.schema().clone();
1692 let conditions: Vec<Condition> = match &select.selection {
1695 Some(expr) => match translate_sqlparser_filter(expr, &schema) {
1696 Some(c) => c,
1697 None => return Ok(None),
1698 },
1699 None => Vec::new(),
1700 };
1701 if !conditions.is_empty() && db.ensure_indexes_complete().is_err() {
1702 return Ok(None);
1703 }
1704 let snap = db.snapshot();
1705
1706 let mut col_ids: Vec<u16> = Vec::new();
1708 let mut fields: Vec<arrow::datatypes::Field> = Vec::new();
1709 let resolve_col = |name: &str| -> Option<&mongreldb_core::schema::ColumnDef> {
1710 schema.columns.iter().find(|c| c.name == name)
1711 };
1712 match &proj_names {
1713 None => {
1714 for c in &schema.columns {
1715 col_ids.push(c.id);
1716 fields.push(arrow::datatypes::Field::new(
1717 &c.name,
1718 arrow_conv::arrow_data_type(&c.ty)?,
1719 c.flags.contains(mongreldb_core::ColumnFlags::NULLABLE),
1720 ));
1721 }
1722 }
1723 Some(names) => {
1724 for n in names {
1725 let cdef = match resolve_col(n) {
1726 Some(c) => c,
1727 None => return Ok(None), };
1729 col_ids.push(cdef.id);
1730 fields.push(arrow::datatypes::Field::new(
1731 &cdef.name,
1732 arrow_conv::arrow_data_type(&cdef.ty)?,
1733 cdef.flags.contains(mongreldb_core::ColumnFlags::NULLABLE),
1734 ));
1735 }
1736 }
1737 }
1738
1739 let cols = if !conditions.is_empty() {
1741 match db.query_columns_native_cached(&conditions, Some(&col_ids), snap) {
1742 Ok(Some(c)) => c,
1743 Ok(None) => db
1744 .visible_columns_native(snap, Some(&col_ids))
1745 .map_err(MongrelQueryError::Core)?,
1746 Err(_) => return Ok(None),
1747 }
1748 } else {
1749 db.visible_columns_native(snap, Some(&col_ids))
1750 .map_err(MongrelQueryError::Core)?
1751 };
1752 drop(db);
1753
1754 let mut arrays: Vec<ArrayRef> = Vec::with_capacity(col_ids.len());
1756 for cid in &col_ids {
1757 let col = cols
1758 .iter()
1759 .find(|(id, _)| id == cid)
1760 .map(|(_, c)| c.clone());
1761 let Some(col) = col else { return Ok(None) };
1762 let ty = schema
1763 .columns
1764 .iter()
1765 .find(|c| c.id == *cid)
1766 .map(|c| c.ty)
1767 .unwrap_or(mongreldb_core::schema::TypeId::Int64);
1768 arrays.push(arrow_conv::native_to_array(ty, &col)?);
1769 }
1770 let batch_schema = Arc::new(arrow::datatypes::Schema::new(fields));
1771 let batch = RecordBatch::try_new(batch_schema, arrays)
1772 .map_err(|e| MongrelQueryError::Arrow(format!("direct dispatch batch build: {e}")))?;
1773
1774 mongreldb_core::trace::QueryTrace::record(|t| {
1775 t.scan_mode = mongreldb_core::trace::ScanMode::DirectDispatch;
1776 t.planning_nanos = 0; });
1778 Ok(Some(vec![batch]))
1779 }
1780
1781 pub async fn run(&self, sql: &str) -> Result<Vec<RecordBatch>> {
1786 if let Some(inner) = strip_explain_query_plan(sql) {
1787 return self.explain_query_plan(inner).await;
1788 }
1789 if let Some(batches) = commands::try_run_command(self, sql).await? {
1790 return Ok(batches);
1791 }
1792 let lower = sql.trim_start().to_lowercase();
1794 if lower.starts_with("create table") {
1795 if let Some(db) = &self.database {
1796 let (name, schema) = parse_create_table(sql)?;
1797 db.create_table(&name, schema)?;
1798 let handle = db.table(&name)?;
1799 let provider = MongrelProvider::new(handle.clone())?;
1800 self.ctx
1801 .register_table(&name, Arc::new(provider))
1802 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1803 self.tables.lock().insert(name, handle);
1804 self.clear_cache();
1805 return Ok(Vec::new());
1806 }
1807 }
1808 if lower.starts_with("drop table") {
1809 if let Some(db) = &self.database {
1810 let (name, if_exists) = parse_drop_table(sql)?;
1811 let drop_result = db.drop_table(&name);
1812 if let Err(e) = drop_result {
1813 let is_not_found = matches!(e, mongreldb_core::MongrelError::NotFound(_));
1815 if !(if_exists && is_not_found) {
1816 return Err(e.into());
1817 }
1818 } else {
1819 self.ctx
1820 .deregister_table(&name)
1821 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1822 self.tables.lock().remove(&name);
1823 }
1824 self.clear_cache();
1825 return Ok(Vec::new());
1826 }
1827 }
1828 if lower.starts_with("alter table") {
1829 if let Some(db) = &self.database {
1830 match parse_alter_table(sql)? {
1831 ParsedAlterTable::RenameTable { old_name, new_name } => {
1832 db.rename_table(&old_name, &new_name)?;
1833 self.ctx
1837 .deregister_table(&old_name)
1838 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1839 self.tables.lock().remove(&old_name);
1840 let handle = db.table(&new_name)?;
1841 let provider = MongrelProvider::new(handle.clone())?;
1842 self.ctx
1843 .register_table(&new_name, Arc::new(provider))
1844 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1845 self.tables.lock().insert(new_name, handle);
1846 }
1847 ParsedAlterTable::RenameColumn {
1848 table_name,
1849 column_name,
1850 new_name,
1851 } => {
1852 db.alter_column(&table_name, &column_name, AlterColumn::rename(new_name))?;
1853 self.refresh_registered_table(db, &table_name)?;
1854 }
1855 ParsedAlterTable::AlterColumnType {
1856 table_name,
1857 column_name,
1858 ty,
1859 } => {
1860 db.alter_column(&table_name, &column_name, AlterColumn::set_type(ty))?;
1861 self.refresh_registered_table(db, &table_name)?;
1862 }
1863 ParsedAlterTable::SetNotNull {
1864 table_name,
1865 column_name,
1866 } => {
1867 let flags = current_column_flags(db, &table_name, &column_name)?
1868 .without(ColumnFlags::NULLABLE);
1869 db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
1870 self.refresh_registered_table(db, &table_name)?;
1871 }
1872 ParsedAlterTable::DropNotNull {
1873 table_name,
1874 column_name,
1875 } => {
1876 let flags = current_column_flags(db, &table_name, &column_name)?
1877 .with(ColumnFlags::NULLABLE);
1878 db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
1879 self.refresh_registered_table(db, &table_name)?;
1880 }
1881 }
1882 self.clear_cache();
1883 return Ok(Vec::new());
1884 }
1885 }
1886
1887 let resolved = self.resolve_view_sql(sql);
1890 let resolved = self.rewrite_external_module_compat_sql(&resolved);
1891 let resolved = rewrite_compat_function_calls(&resolved);
1892 let effective_sql = normalize_sql(&resolved);
1896 let sql = effective_sql.as_str();
1897 let epoch = self.cache_epoch();
1901 let key = (sql.to_string(), epoch);
1902 let result_cacheable = !extended_sql_functions::contains_volatile_extended_function(sql);
1903 if result_cacheable {
1904 if let Some(hit) = self.cache.lock().get(&key) {
1905 return Ok((**hit).clone());
1906 }
1907 }
1908 if let Some(batches) = self.try_catalog_introspection(sql)? {
1911 if result_cacheable {
1912 self.cache.lock().insert(key, Arc::new(batches.clone()));
1913 }
1914 return Ok(batches);
1915 }
1916 if let Some(batches) = self.try_direct_dispatch(sql)? {
1921 if result_cacheable {
1922 self.cache.lock().insert(key, Arc::new(batches.clone()));
1923 }
1924 return Ok(batches);
1925 }
1926 let plan_start = std::time::Instant::now();
1928 let external_module_scan = self.query_references_external_module(sql);
1929 let df = {
1930 let cached_plan = self.plan_cache.lock().get(sql).cloned();
1931 if let Some(plan) = cached_plan {
1932 datafusion::dataframe::DataFrame::new(self.ctx.state(), plan)
1933 } else {
1934 let df = self
1935 .ctx
1936 .sql(sql)
1937 .await
1938 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1939 self.plan_cache
1940 .lock()
1941 .insert(sql.to_string(), df.logical_plan().clone());
1942 df
1943 }
1944 };
1945 let planning_nanos = plan_start.elapsed().as_nanos() as u64;
1948 mongreldb_core::trace::QueryTrace::record(|t| t.planning_nanos = planning_nanos);
1949
1950 let agg_key = sql_cache_key(sql);
1955 let batches = match self.try_native_aggregate(df.logical_plan(), agg_key) {
1956 Ok(Some(batch)) => vec![batch],
1957 _ => {
1958 match self.try_fk_join(df.logical_plan()) {
1962 Ok(Some(b)) => {
1963 mongreldb_core::trace::QueryTrace::record(|t| {
1965 t.join_mode = mongreldb_core::trace::JoinMode::FkBitmap;
1966 });
1967 b
1968 }
1969 _ => df
1970 .collect()
1971 .await
1972 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?,
1973 }
1974 }
1975 };
1976 if external_module_scan {
1977 mongreldb_core::trace::QueryTrace::record(|t| {
1978 t.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
1979 });
1980 }
1981 if result_cacheable {
1982 self.cache.lock().insert(key, Arc::new(batches.clone()));
1983 }
1984 Ok(batches)
1985 }
1986
1987 pub async fn run_sql_traced(
1999 &self,
2000 sql: &str,
2001 ) -> Result<(Vec<RecordBatch>, mongreldb_core::trace::QueryTrace)> {
2002 mongreldb_core::trace::QueryTrace::push_scope();
2003 let result = self.run(sql).await;
2004 let trace = mongreldb_core::trace::QueryTrace::pop_scope();
2005 Ok((result?, trace))
2006 }
2007
2008 pub fn clear_cache(&self) {
2011 self.cache.lock().clear();
2012 self.plan_cache.lock().clear();
2013 }
2014
2015 async fn explain_query_plan(&self, sql: &str) -> Result<Vec<RecordBatch>> {
2016 let explain_sql = format!("EXPLAIN {}", sql.trim().trim_end_matches(';'));
2017 let batches = self
2018 .ctx
2019 .sql(&explain_sql)
2020 .await
2021 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?
2022 .collect()
2023 .await
2024 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2025 let mut detail = self.mongrel_query_plan_details(sql);
2026 for batch in &batches {
2027 if batch.num_columns() < 2 {
2028 continue;
2029 }
2030 let Some(plan_type) = batch.column(0).as_any().downcast_ref::<StringArray>() else {
2031 continue;
2032 };
2033 let Some(plan) = batch.column(1).as_any().downcast_ref::<StringArray>() else {
2034 continue;
2035 };
2036 for row in 0..batch.num_rows() {
2037 let prefix = plan_type.value(row);
2038 for line in plan.value(row).lines() {
2039 let line = line.trim();
2040 if !line.is_empty() {
2041 detail.push(format!("DATAFUSION {prefix}: {line}"));
2042 }
2043 }
2044 }
2045 }
2046 if detail.is_empty() {
2047 detail.push("plan unavailable".to_string());
2048 }
2049 let ids = (0..detail.len()).map(|i| i as i64).collect::<Vec<_>>();
2050 let parents = vec![0_i64; detail.len()];
2051 let notused = vec![0_i64; detail.len()];
2052 let schema = Arc::new(arrow::datatypes::Schema::new(vec![
2053 arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int64, false),
2054 arrow::datatypes::Field::new("parent", arrow::datatypes::DataType::Int64, false),
2055 arrow::datatypes::Field::new("notused", arrow::datatypes::DataType::Int64, false),
2056 arrow::datatypes::Field::new("detail", arrow::datatypes::DataType::Utf8, false),
2057 ]));
2058 let batch = RecordBatch::try_new(
2059 schema,
2060 vec![
2061 Arc::new(Int64Array::from(ids)) as ArrayRef,
2062 Arc::new(Int64Array::from(parents)),
2063 Arc::new(Int64Array::from(notused)),
2064 Arc::new(StringArray::from(detail)),
2065 ],
2066 )
2067 .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
2068 Ok(vec![batch])
2069 }
2070
2071 fn mongrel_query_plan_details(&self, sql: &str) -> Vec<String> {
2072 use sqlparser::ast::{GroupByExpr, OrderByKind, SetExpr, Statement};
2073 use sqlparser::dialect::PostgreSqlDialect;
2074 use sqlparser::parser::Parser;
2075
2076 let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
2077 return Vec::new();
2078 };
2079 let Some(Statement::Query(query)) = stmts.first() else {
2080 return Vec::new();
2081 };
2082
2083 fn collect(session: &MongrelSession, query: &sqlparser::ast::Query, out: &mut Vec<String>) {
2084 use sqlparser::ast::{SetOperator, TableWithJoins};
2085 match query.body.as_ref() {
2086 SetExpr::Select(select) => {
2087 for TableWithJoins { relation, joins } in &select.from {
2088 session.push_table_plan(relation, select.selection.as_ref(), out);
2089 for join in joins {
2090 session.push_table_plan(&join.relation, None, out);
2091 }
2092 }
2093 if select.distinct.is_some() {
2094 out.push("USE TEMP B-TREE FOR DISTINCT".to_string());
2095 }
2096 let grouped = match &select.group_by {
2097 GroupByExpr::All(_) => true,
2098 GroupByExpr::Expressions(exprs, _) => !exprs.is_empty(),
2099 };
2100 if grouped {
2101 out.push("USE TEMP B-TREE FOR GROUP BY".to_string());
2102 }
2103 let ordered = query.order_by.as_ref().is_some_and(|order_by| {
2104 matches!(order_by.kind, OrderByKind::All(_))
2105 || matches!(&order_by.kind, OrderByKind::Expressions(exprs) if !exprs.is_empty())
2106 });
2107 if ordered {
2108 out.push("USE TEMP B-TREE FOR ORDER BY".to_string());
2109 }
2110 }
2111 SetExpr::Query(query) => collect(session, query, out),
2112 SetExpr::SetOperation {
2113 left, op, right, ..
2114 } => {
2115 let label = match op {
2116 SetOperator::Union => "COMPOUND QUERY UNION",
2117 SetOperator::Except => "COMPOUND QUERY EXCEPT",
2118 SetOperator::Intersect => "COMPOUND QUERY INTERSECT",
2119 _ => "COMPOUND QUERY",
2120 };
2121 out.push(label.to_string());
2122 collect_set_expr(session, left, out);
2123 collect_set_expr(session, right, out);
2124 }
2125 _ => {}
2126 }
2127 }
2128
2129 fn collect_set_expr(session: &MongrelSession, expr: &SetExpr, out: &mut Vec<String>) {
2130 match expr {
2131 SetExpr::Select(select) => {
2132 for table in &select.from {
2133 session.push_table_plan(&table.relation, select.selection.as_ref(), out);
2134 }
2135 }
2136 SetExpr::Query(query) => collect(session, query, out),
2137 SetExpr::SetOperation { left, right, .. } => {
2138 collect_set_expr(session, left, out);
2139 collect_set_expr(session, right, out);
2140 }
2141 _ => {}
2142 }
2143 }
2144
2145 let mut out = Vec::new();
2146 collect(self, query, &mut out);
2147 out
2148 }
2149
2150 fn push_table_plan(
2151 &self,
2152 relation: &sqlparser::ast::TableFactor,
2153 selection: Option<&sqlparser::ast::Expr>,
2154 out: &mut Vec<String>,
2155 ) {
2156 let sqlparser::ast::TableFactor::Table { name, alias, .. } = relation else {
2157 out.push("SCAN SUBQUERY".to_string());
2158 return;
2159 };
2160 let table_name = name.to_string();
2161 let display_name = alias
2162 .as_ref()
2163 .map(|alias| alias.name.value.clone())
2164 .unwrap_or_else(|| table_name.clone());
2165 let Some(handle) = self.tables.lock().get(&table_name).cloned() else {
2166 out.push(format!("SCAN {display_name}"));
2167 return;
2168 };
2169 let schema = handle.lock().schema().clone();
2170 let searchable = selection
2171 .and_then(|expr| translate_sqlparser_filter(expr, &schema))
2172 .is_some_and(|conditions| !conditions.is_empty());
2173 if searchable {
2174 out.push(format!("SEARCH {display_name} USING MONGREL INDEX"));
2175 } else {
2176 out.push(format!("SCAN {display_name}"));
2177 }
2178 }
2179
2180 fn resolve_view_sql(&self, sql: &str) -> String {
2185 let views = self.views.lock();
2186 if views.is_empty() {
2187 return sql.to_string();
2188 }
2189 let mut result = sql.to_string();
2190 for (name, view) in views.iter() {
2191 result = replace_from_view(&result, name, &view.sql);
2192 }
2193 result
2194 }
2195
2196 fn rewrite_external_module_compat_sql(&self, sql: &str) -> String {
2197 let Some(db) = &self.database else {
2198 return sql.to_string();
2199 };
2200 rewrite_fts_match_compat_sql(sql, db)
2201 }
2202
2203 fn query_references_external_module(&self, sql: &str) -> bool {
2204 use sqlparser::ast::Statement;
2205 use sqlparser::dialect::PostgreSqlDialect;
2206 use sqlparser::parser::Parser;
2207
2208 Parser::parse_sql(&PostgreSqlDialect {}, sql)
2209 .ok()
2210 .is_some_and(|statements| {
2211 statements.iter().any(|statement| match statement {
2212 Statement::Query(query) => self.query_uses_external_module(query),
2213 _ => false,
2214 })
2215 })
2216 }
2217
2218 fn query_uses_external_module(&self, query: &sqlparser::ast::Query) -> bool {
2219 self.set_expr_uses_external_module(query.body.as_ref())
2220 }
2221
2222 fn set_expr_uses_external_module(&self, expr: &sqlparser::ast::SetExpr) -> bool {
2223 use sqlparser::ast::SetExpr;
2224
2225 match expr {
2226 SetExpr::Select(select) => select
2227 .from
2228 .iter()
2229 .any(|table| self.table_with_joins_uses_external_module(table)),
2230 SetExpr::Query(query) => self.query_uses_external_module(query),
2231 SetExpr::SetOperation { left, right, .. } => {
2232 self.set_expr_uses_external_module(left)
2233 || self.set_expr_uses_external_module(right)
2234 }
2235 _ => false,
2236 }
2237 }
2238
2239 fn table_with_joins_uses_external_module(
2240 &self,
2241 table: &sqlparser::ast::TableWithJoins,
2242 ) -> bool {
2243 self.table_factor_uses_external_module(&table.relation)
2244 || table
2245 .joins
2246 .iter()
2247 .any(|join| self.table_factor_uses_external_module(&join.relation))
2248 }
2249
2250 fn table_factor_uses_external_module(&self, relation: &sqlparser::ast::TableFactor) -> bool {
2251 use sqlparser::ast::{Expr, TableFactor};
2252
2253 match relation {
2254 TableFactor::Table { name, args, .. } => {
2255 let table_name = name.to_string();
2256 self.database
2257 .as_ref()
2258 .is_some_and(|db| db.external_table(&table_name).is_some())
2259 || (args.is_some() && self.external_modules.contains(&table_name))
2260 }
2261 TableFactor::Function { name, .. } => self.external_modules.contains(&name.to_string()),
2262 TableFactor::TableFunction {
2263 expr: Expr::Function(func),
2264 ..
2265 } => self.external_modules.contains(&func.name.to_string()),
2266 TableFactor::Derived { subquery, .. } => self.query_uses_external_module(subquery),
2267 _ => false,
2268 }
2269 }
2270
2271 fn cache_epoch(&self) -> u64 {
2274 if let Some(db) = &self.database {
2275 db.visible_epoch().0
2276 } else {
2277 self.combined_epoch()
2278 }
2279 }
2280
2281 fn combined_epoch(&self) -> u64 {
2282 let primary = self.db.as_ref().expect("no primary table");
2283 let mut combined = primary.lock().snapshot().epoch.0;
2284 let tables = self.tables.lock();
2285 for arc in tables.values() {
2286 if !Arc::ptr_eq(arc, primary) {
2287 let e = arc.lock().snapshot().epoch.0;
2288 combined = combined.wrapping_mul(31).wrapping_add(e);
2289 }
2290 }
2291 combined
2292 }
2293
2294 fn try_native_aggregate(
2298 &self,
2299 plan: &datafusion::logical_expr::LogicalPlan,
2300 cache_key: u64,
2301 ) -> Result<Option<RecordBatch>> {
2302 let Some(primary) = self.db.as_ref() else {
2303 return Ok(None);
2304 };
2305 let mut db = primary.lock();
2306 let schema = db.schema().clone();
2307 let snap = db.snapshot();
2308 native_agg::try_native_aggregate(&mut db, &schema, snap, plan, cache_key)
2309 }
2310
2311 fn try_fk_join(
2315 &self,
2316 plan: &datafusion::logical_expr::LogicalPlan,
2317 ) -> Result<Option<Vec<RecordBatch>>> {
2318 let tables = self.tables.lock();
2319 fk_join::try_fk_join(&tables, plan)
2320 }
2321
2322 pub fn context(&self) -> &SessionContext {
2323 &self.ctx
2324 }
2325
2326 pub fn register_scalar_udf(&self, f: ScalarUDF) {
2332 self.ctx.register_udf(f);
2333 self.clear_cache();
2334 }
2335
2336 pub fn register_aggregate_udf(&self, f: AggregateUDF) {
2338 self.ctx.register_udaf(f);
2339 self.clear_cache();
2340 }
2341
2342 pub fn register_window_udf(&self, f: WindowUDF) {
2344 self.ctx.register_udwf(f);
2345 self.clear_cache();
2346 }
2347}
2348
2349fn register_mongrel_functions(
2350 ctx: &SessionContext,
2351 sql_fn_state: Arc<extended_sql_functions::ExtendedSqlState>,
2352) {
2353 ctx.register_udf(ScalarUDF::from(udf::AnnSearchUdf::new()));
2354 ctx.register_udf(ScalarUDF::from(udf::SparseMatchUdf::new()));
2355 ctx.register_udf(ScalarUDF::from(udf::RTreeIntersectsUdf::new()));
2356 for udaf in percentile::percentile_udafs() {
2357 ctx.register_udaf(udaf);
2358 }
2359 extended_sql_functions::register_extended_sql_functions_with_state(ctx, sql_fn_state);
2360}
2361
2362fn strip_explain_query_plan(sql: &str) -> Option<&str> {
2363 let trimmed = sql.trim_start();
2364 let lower = trimmed.to_ascii_lowercase();
2365 if !lower.starts_with("explain") {
2366 return None;
2367 }
2368 let after_explain = trimmed.get(7..)?.trim_start();
2369 let after_explain_lower = after_explain.to_ascii_lowercase();
2370 if !after_explain_lower.starts_with("query") {
2371 return None;
2372 }
2373 let after_query = after_explain.get(5..)?.trim_start();
2374 let after_query_lower = after_query.to_ascii_lowercase();
2375 if !after_query_lower.starts_with("plan") {
2376 return None;
2377 }
2378 Some(after_query.get(4..)?.trim_start())
2379}
2380
2381#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2382enum SqlCompatTokenKind {
2383 Ident,
2384 String,
2385 Dot,
2386 LParen,
2387 RParen,
2388 Comma,
2389}
2390
2391#[derive(Debug, Clone)]
2392struct SqlCompatToken {
2393 kind: SqlCompatTokenKind,
2394 raw: String,
2395 normalized: String,
2396 start: usize,
2397 end: usize,
2398}
2399
2400#[derive(Debug, Clone)]
2401struct FtsMatchBinding {
2402 query_ref: String,
2403}
2404
2405#[derive(Debug, Clone)]
2406struct SqlReplacement {
2407 start: usize,
2408 end: usize,
2409 replacement: String,
2410}
2411
2412fn rewrite_fts_match_compat_sql(sql: &str, db: &Database) -> String {
2413 let tokens = sql_compat_tokens(sql);
2414 if tokens.is_empty() {
2415 return sql.to_string();
2416 }
2417 let bindings = fts_match_bindings(sql, db, &tokens);
2418 if bindings.is_empty() {
2419 return sql.to_string();
2420 }
2421 let unique_refs = bindings
2422 .values()
2423 .map(|binding| binding.query_ref.as_str())
2424 .collect::<HashSet<_>>();
2425 let unique_binding = if unique_refs.len() == 1 {
2426 bindings.values().next().cloned()
2427 } else {
2428 None
2429 };
2430 let mut replacements = Vec::new();
2431 for (idx, token) in tokens.iter().enumerate() {
2432 if token.kind != SqlCompatTokenKind::Ident || token.normalized != "match" {
2433 continue;
2434 }
2435 let Some(rhs) = tokens.get(idx + 1) else {
2436 continue;
2437 };
2438 if rhs.kind != SqlCompatTokenKind::String {
2439 continue;
2440 }
2441 let Some((lhs_start, _lhs_end, query_ref)) =
2442 fts_match_lhs_query_ref(&tokens, idx, &bindings, unique_binding.as_ref())
2443 else {
2444 continue;
2445 };
2446 replacements.push(SqlReplacement {
2447 start: lhs_start,
2448 end: rhs.end,
2449 replacement: format!("{query_ref}.query = {}", rhs.raw),
2450 });
2451 }
2452 apply_sql_replacements(sql, &replacements)
2453}
2454
2455fn fts_match_lhs_query_ref(
2456 tokens: &[SqlCompatToken],
2457 match_idx: usize,
2458 bindings: &HashMap<String, FtsMatchBinding>,
2459 unique_binding: Option<&FtsMatchBinding>,
2460) -> Option<(usize, usize, String)> {
2461 if match_idx == 0 {
2462 return None;
2463 }
2464 let lhs = tokens.get(match_idx - 1)?;
2465 if lhs.kind != SqlCompatTokenKind::Ident {
2466 return None;
2467 }
2468
2469 if match_idx >= 3
2470 && tokens.get(match_idx - 2)?.kind == SqlCompatTokenKind::Dot
2471 && tokens.get(match_idx - 3)?.kind == SqlCompatTokenKind::Ident
2472 {
2473 let owner = tokens.get(match_idx - 3)?;
2474 let binding = bindings.get(&owner.normalized)?;
2475 if lhs.normalized == "query" || lhs.normalized == "text" {
2476 return Some((owner.start, lhs.end, binding.query_ref.clone()));
2477 }
2478 return None;
2479 }
2480
2481 if let Some(binding) = bindings.get(&lhs.normalized) {
2482 return Some((lhs.start, lhs.end, binding.query_ref.clone()));
2483 }
2484 if lhs.normalized == "text" {
2485 let binding = unique_binding?;
2486 return Some((lhs.start, lhs.end, binding.query_ref.clone()));
2487 }
2488 None
2489}
2490
2491fn fts_match_bindings(
2492 sql: &str,
2493 db: &Database,
2494 tokens: &[SqlCompatToken],
2495) -> HashMap<String, FtsMatchBinding> {
2496 let mut out = HashMap::new();
2497 let mut i = 0;
2498 while i < tokens.len() {
2499 let token = &tokens[i];
2500 let starts_table_ref = token.kind == SqlCompatTokenKind::Ident
2501 && matches!(token.normalized.as_str(), "from" | "join");
2502 if !starts_table_ref {
2503 i += 1;
2504 continue;
2505 }
2506 let mut table_idx = i + 1;
2507 if tokens
2508 .get(table_idx)
2509 .is_some_and(|token| token.kind == SqlCompatTokenKind::LParen)
2510 {
2511 i += 1;
2512 continue;
2513 }
2514 let Some(table) = tokens.get(table_idx) else {
2515 break;
2516 };
2517 if table.kind != SqlCompatTokenKind::Ident {
2518 i += 1;
2519 continue;
2520 }
2521 let mut table_name = table.normalized.clone();
2522 let mut table_ref = table.raw.clone();
2523 if tokens
2524 .get(table_idx + 1)
2525 .is_some_and(|token| token.kind == SqlCompatTokenKind::Dot)
2526 && tokens
2527 .get(table_idx + 2)
2528 .is_some_and(|token| token.kind == SqlCompatTokenKind::Ident)
2529 {
2530 let qualified = tokens.get(table_idx + 2).unwrap();
2531 table_name = qualified.normalized.clone();
2532 table_ref = sql[table.start..qualified.end].to_string();
2533 table_idx += 2;
2534 }
2535 if !is_fts_docs_table(db, &table_name) {
2536 i = table_idx + 1;
2537 continue;
2538 }
2539 let mut query_ref = table_ref.clone();
2540 let mut alias_key = None;
2541 let mut next = table_idx + 1;
2542 if tokens.get(next).is_some_and(|token| {
2543 token.kind == SqlCompatTokenKind::Ident && token.normalized == "as"
2544 }) {
2545 next += 1;
2546 }
2547 if let Some(alias) = tokens.get(next) {
2548 if alias.kind == SqlCompatTokenKind::Ident && !is_table_ref_boundary(&alias.normalized)
2549 {
2550 alias_key = Some(alias.normalized.clone());
2551 query_ref = alias.raw.clone();
2552 next += 1;
2553 }
2554 }
2555 out.insert(
2556 table_name,
2557 FtsMatchBinding {
2558 query_ref: query_ref.clone(),
2559 },
2560 );
2561 if let Some(alias_key) = alias_key {
2562 out.insert(alias_key, FtsMatchBinding { query_ref });
2563 }
2564 i = next;
2565 }
2566 out
2567}
2568
2569fn is_fts_docs_table(db: &Database, name: &str) -> bool {
2570 db.external_table(name)
2571 .is_some_and(|entry| entry.module == "fts_docs")
2572}
2573
2574fn is_table_ref_boundary(normalized: &str) -> bool {
2575 matches!(
2576 normalized,
2577 "where"
2578 | "join"
2579 | "left"
2580 | "right"
2581 | "inner"
2582 | "outer"
2583 | "full"
2584 | "cross"
2585 | "on"
2586 | "using"
2587 | "group"
2588 | "order"
2589 | "having"
2590 | "limit"
2591 | "offset"
2592 | "union"
2593 | "except"
2594 | "intersect"
2595 )
2596}
2597
2598fn sql_compat_tokens(sql: &str) -> Vec<SqlCompatToken> {
2599 let bytes = sql.as_bytes();
2600 let mut tokens = Vec::new();
2601 let mut i = 0;
2602 while i < bytes.len() {
2603 match bytes[i] {
2604 b if b.is_ascii_whitespace() => i += 1,
2605 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
2606 i += 2;
2607 while i < bytes.len() && bytes[i] != b'\n' {
2608 i += 1;
2609 }
2610 }
2611 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
2612 i = skip_block_comment(bytes, i);
2613 }
2614 b'\'' => {
2615 let end = skip_quoted(bytes, i, b'\'');
2616 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
2617 i = end;
2618 }
2619 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
2620 let end = skip_quoted(bytes, i + 1, b'\'');
2621 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
2622 i = end;
2623 }
2624 b'$' => {
2625 let (end, matched) = skip_dollar_quoted(bytes, i);
2626 if matched {
2627 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
2628 i = end;
2629 } else {
2630 i += 1;
2631 }
2632 }
2633 b'"' => {
2634 let end = skip_quoted(bytes, i, b'"');
2635 let raw = sql[i..end].to_string();
2636 let normalized = unquote_sql_ident(&raw).to_ascii_lowercase();
2637 tokens.push(SqlCompatToken {
2638 kind: SqlCompatTokenKind::Ident,
2639 raw,
2640 normalized,
2641 start: i,
2642 end,
2643 });
2644 i = end;
2645 }
2646 b'.' => {
2647 tokens.push(sql_token(sql, SqlCompatTokenKind::Dot, i, i + 1));
2648 i += 1;
2649 }
2650 b'(' => {
2651 tokens.push(sql_token(sql, SqlCompatTokenKind::LParen, i, i + 1));
2652 i += 1;
2653 }
2654 b')' => {
2655 tokens.push(sql_token(sql, SqlCompatTokenKind::RParen, i, i + 1));
2656 i += 1;
2657 }
2658 b',' => {
2659 tokens.push(sql_token(sql, SqlCompatTokenKind::Comma, i, i + 1));
2660 i += 1;
2661 }
2662 b if is_sql_ident_byte(b) => {
2663 let start = i;
2664 i += 1;
2665 while i < bytes.len() && is_sql_ident_byte(bytes[i]) {
2666 i += 1;
2667 }
2668 tokens.push(sql_token(sql, SqlCompatTokenKind::Ident, start, i));
2669 }
2670 _ => i += 1,
2671 }
2672 }
2673 tokens
2674}
2675
2676fn sql_token(sql: &str, kind: SqlCompatTokenKind, start: usize, end: usize) -> SqlCompatToken {
2677 let raw = sql[start..end].to_string();
2678 SqlCompatToken {
2679 kind,
2680 normalized: raw.to_ascii_lowercase(),
2681 raw,
2682 start,
2683 end,
2684 }
2685}
2686
2687fn unquote_sql_ident(raw: &str) -> String {
2688 if raw.len() >= 2 && raw.starts_with('"') && raw.ends_with('"') {
2689 raw[1..raw.len() - 1].replace("\"\"", "\"")
2690 } else {
2691 raw.to_string()
2692 }
2693}
2694
2695fn apply_sql_replacements(sql: &str, replacements: &[SqlReplacement]) -> String {
2696 if replacements.is_empty() {
2697 return sql.to_string();
2698 }
2699 let mut ordered = replacements.to_vec();
2700 ordered.sort_by_key(|replacement| replacement.start);
2701 let mut out = String::with_capacity(sql.len());
2702 let mut cursor = 0;
2703 for replacement in ordered {
2704 if replacement.start < cursor || replacement.end > sql.len() {
2705 continue;
2706 }
2707 out.push_str(&sql[cursor..replacement.start]);
2708 out.push_str(&replacement.replacement);
2709 cursor = replacement.end;
2710 }
2711 out.push_str(&sql[cursor..]);
2712 out
2713}
2714
2715fn rewrite_compat_function_calls(sql: &str) -> String {
2716 let bytes = sql.as_bytes();
2717 let mut out = String::with_capacity(sql.len());
2718 let mut i = 0;
2719 while i < bytes.len() {
2720 match bytes[i] {
2721 b'\'' => i = copy_quoted_to_string(&mut out, bytes, i, b'\''),
2722 b'"' => i = copy_quoted_to_string(&mut out, bytes, i, b'"'),
2723 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
2724 out.push(bytes[i] as char);
2725 i += 1;
2726 i = copy_quoted_to_string(&mut out, bytes, i, b'\'');
2727 }
2728 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
2729 out.push('-');
2730 out.push('-');
2731 i += 2;
2732 while i < bytes.len() {
2733 let ch = bytes[i] as char;
2734 out.push(ch);
2735 i += 1;
2736 if ch == '\n' {
2737 break;
2738 }
2739 }
2740 }
2741 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
2742 let start = i;
2743 i = skip_block_comment(bytes, i);
2744 out.push_str(&sql[start..i.min(bytes.len())]);
2745 }
2746 b'$' => {
2747 let start_len = out.len();
2748 let (next, matched) = copy_dollar_quoted_to_string(&mut out, bytes, i);
2749 if matched {
2750 i = next;
2751 } else {
2752 out.truncate(start_len);
2753 out.push('$');
2754 i += 1;
2755 }
2756 }
2757 b'g' | b'G' | b'm' | b'M' | b't' | b'T' => {
2758 if let Some((replacement, next)) = compat_function_rewrite_at(sql, i) {
2759 out.push_str(&replacement);
2760 i = next;
2761 } else {
2762 out.push(bytes[i] as char);
2763 i += 1;
2764 }
2765 }
2766 _ => {
2767 out.push(bytes[i] as char);
2768 i += 1;
2769 }
2770 }
2771 }
2772 out
2773}
2774
2775fn compat_function_rewrite_at(sql: &str, start: usize) -> Option<(String, usize)> {
2776 let bytes = sql.as_bytes();
2777 let (name, kind) = if ident_eq_at(bytes, start, b"max") {
2778 ("max", CompatRewriteKind::ScalarMax)
2779 } else if ident_eq_at(bytes, start, b"min") {
2780 ("min", CompatRewriteKind::ScalarMin)
2781 } else if ident_eq_at(bytes, start, b"group_concat") {
2782 ("group_concat", CompatRewriteKind::GroupConcat)
2783 } else if ident_eq_at(bytes, start, b"total") {
2784 ("total", CompatRewriteKind::Total)
2785 } else {
2786 return None;
2787 };
2788 let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
2789 let after_name = start + name.len();
2790 let after_ok = bytes
2791 .get(after_name)
2792 .is_some_and(|b| !is_sql_ident_byte(*b));
2793 if !before_ok || !after_ok {
2794 return None;
2795 }
2796 let mut open = after_name;
2797 while open < bytes.len() && bytes[open].is_ascii_whitespace() {
2798 open += 1;
2799 }
2800 if bytes.get(open) != Some(&b'(') {
2801 return None;
2802 }
2803 let summary = call_arg_summary(sql, open)?;
2804 match kind {
2805 CompatRewriteKind::ScalarMax if summary.top_level_commas > 0 => {
2806 Some(("__mongreldb_scalar_max(".to_string(), open + 1))
2807 }
2808 CompatRewriteKind::ScalarMin if summary.top_level_commas > 0 => {
2809 Some(("__mongreldb_scalar_min(".to_string(), open + 1))
2810 }
2811 CompatRewriteKind::GroupConcat => {
2812 let args = &sql[open + 1..summary.close];
2813 let rewritten = if summary.top_level_commas == 0 {
2814 format!("string_agg({args}, ',')")
2815 } else {
2816 format!("string_agg({args})")
2817 };
2818 Some((rewritten, summary.close + 1))
2819 }
2820 CompatRewriteKind::Total if summary.top_level_commas == 0 => {
2821 let args = &sql[open + 1..summary.close];
2822 let suffix_end = aggregate_suffix_end(sql, summary.close + 1);
2823 let suffix = &sql[summary.close + 1..suffix_end];
2824 Some((
2825 format!("coalesce(cast(sum({args}){suffix} as double), 0.0)"),
2826 suffix_end,
2827 ))
2828 }
2829 _ => None,
2830 }
2831}
2832
2833#[derive(Clone, Copy)]
2834enum CompatRewriteKind {
2835 ScalarMax,
2836 ScalarMin,
2837 GroupConcat,
2838 Total,
2839}
2840
2841fn ident_eq_at(bytes: &[u8], start: usize, ident: &[u8]) -> bool {
2842 bytes
2843 .get(start..start + ident.len())
2844 .is_some_and(|slice| slice.eq_ignore_ascii_case(ident))
2845}
2846
2847fn is_sql_ident_byte(b: u8) -> bool {
2848 b.is_ascii_alphanumeric() || b == b'_' || b == b'$'
2849}
2850
2851fn keyword_at(bytes: &[u8], start: usize, keyword: &[u8]) -> bool {
2852 if !ident_eq_at(bytes, start, keyword) {
2853 return false;
2854 }
2855 let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
2856 let after = start + keyword.len();
2857 let after_ok = after >= bytes.len() || !is_sql_ident_byte(bytes[after]);
2858 before_ok && after_ok
2859}
2860
2861fn skip_sql_whitespace(bytes: &[u8], mut i: usize) -> usize {
2862 while i < bytes.len() && bytes[i].is_ascii_whitespace() {
2863 i += 1;
2864 }
2865 i
2866}
2867
2868fn aggregate_suffix_end(sql: &str, start: usize) -> usize {
2869 let bytes = sql.as_bytes();
2870 let mut suffix_end = start;
2871 let mut i = skip_sql_whitespace(bytes, start);
2872
2873 if keyword_at(bytes, i, b"filter") {
2874 let open = skip_sql_whitespace(bytes, i + b"filter".len());
2875 if bytes.get(open) != Some(&b'(') {
2876 return start;
2877 }
2878 let Some(summary) = call_arg_summary(sql, open) else {
2879 return start;
2880 };
2881 suffix_end = summary.close + 1;
2882 i = skip_sql_whitespace(bytes, suffix_end);
2883 }
2884
2885 if keyword_at(bytes, i, b"over") {
2886 let after_over = skip_sql_whitespace(bytes, i + b"over".len());
2887 if bytes.get(after_over) == Some(&b'(') {
2888 let Some(summary) = call_arg_summary(sql, after_over) else {
2889 return suffix_end;
2890 };
2891 suffix_end = summary.close + 1;
2892 } else {
2893 let mut end = after_over;
2894 while end < bytes.len() && is_sql_ident_byte(bytes[end]) {
2895 end += 1;
2896 }
2897 if end > after_over {
2898 suffix_end = end;
2899 }
2900 }
2901 }
2902
2903 suffix_end
2904}
2905
2906struct CallArgSummary {
2907 close: usize,
2908 top_level_commas: usize,
2909}
2910
2911fn call_arg_summary(sql: &str, open: usize) -> Option<CallArgSummary> {
2912 let bytes = sql.as_bytes();
2913 let mut depth = 1;
2914 let mut i = open + 1;
2915 let mut top_level_commas = 0;
2916 while i < bytes.len() {
2917 match bytes[i] {
2918 b'\'' => i = skip_quoted(bytes, i, b'\''),
2919 b'"' => i = skip_quoted(bytes, i, b'"'),
2920 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
2921 i = skip_quoted(bytes, i + 1, b'\'')
2922 }
2923 b'$' => {
2924 let (next, matched) = skip_dollar_quoted(bytes, i);
2925 i = if matched { next } else { i + 1 };
2926 }
2927 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
2928 i += 2;
2929 while i < bytes.len() && bytes[i] != b'\n' {
2930 i += 1;
2931 }
2932 }
2933 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
2934 i = skip_block_comment(bytes, i);
2935 }
2936 b'(' => {
2937 depth += 1;
2938 i += 1;
2939 }
2940 b')' => {
2941 depth -= 1;
2942 if depth == 0 {
2943 return Some(CallArgSummary {
2944 close: i,
2945 top_level_commas,
2946 });
2947 }
2948 i += 1;
2949 }
2950 b',' if depth == 1 => {
2951 top_level_commas += 1;
2952 i += 1;
2953 }
2954 _ => i += 1,
2955 }
2956 }
2957 None
2958}
2959
2960fn copy_quoted_to_string(out: &mut String, bytes: &[u8], start: usize, delim: u8) -> usize {
2961 let end = skip_quoted(bytes, start, delim);
2962 out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
2963 end
2964}
2965
2966fn skip_quoted(bytes: &[u8], start: usize, delim: u8) -> usize {
2967 let mut i = start;
2968 if i < bytes.len() {
2969 i += 1;
2970 }
2971 while i < bytes.len() {
2972 if bytes[i] == delim {
2973 i += 1;
2974 if i < bytes.len() && bytes[i] == delim {
2975 i += 1;
2976 continue;
2977 }
2978 break;
2979 }
2980 i += 1;
2981 }
2982 i
2983}
2984
2985fn copy_dollar_quoted_to_string(out: &mut String, bytes: &[u8], start: usize) -> (usize, bool) {
2986 let (end, matched) = skip_dollar_quoted(bytes, start);
2987 if matched {
2988 out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
2989 }
2990 (end, matched)
2991}
2992
2993fn skip_dollar_quoted(bytes: &[u8], start: usize) -> (usize, bool) {
2994 if bytes.get(start) != Some(&b'$') {
2995 return (start, false);
2996 }
2997 let mut j = start + 1;
2998 while j < bytes.len() && (bytes[j].is_ascii_alphanumeric() || bytes[j] == b'_') {
2999 j += 1;
3000 }
3001 if bytes.get(j) != Some(&b'$') {
3002 return (start, false);
3003 }
3004 let tag = &bytes[start..=j];
3005 let mut i = j + 1;
3006 while i + tag.len() <= bytes.len() {
3007 if &bytes[i..i + tag.len()] == tag {
3008 return (i + tag.len(), true);
3009 }
3010 i += 1;
3011 }
3012 (start, false)
3013}
3014
3015fn skip_block_comment(bytes: &[u8], start: usize) -> usize {
3016 let mut i = start + 2;
3017 let mut depth = 1;
3018 while i + 1 < bytes.len() && depth > 0 {
3019 if bytes[i] == b'/' && bytes[i + 1] == b'*' {
3020 depth += 1;
3021 i += 2;
3022 } else if bytes[i] == b'*' && bytes[i + 1] == b'/' {
3023 depth -= 1;
3024 i += 2;
3025 } else {
3026 i += 1;
3027 }
3028 }
3029 i
3030}
3031
3032fn sql_cache_key(sql: &str) -> u64 {
3034 use std::hash::{Hash, Hasher};
3035 let mut h = std::collections::hash_map::DefaultHasher::new();
3036 sql.hash(&mut h);
3037 h.finish()
3038}
3039
3040fn replace_from_view(sql: &str, name: &str, view_sql: &str) -> String {
3047 let lower = sql.to_ascii_lowercase();
3048 let bytes = lower.as_bytes();
3049 let name_b = name.as_bytes();
3050 let mut i = 0usize;
3051 while let Some(rel) = lower[i..].find("from") {
3052 let from_start = i + rel;
3053 let after_from = from_start + 4;
3054 i = after_from;
3055 if from_start > 0 && is_ident_byte(bytes[from_start - 1]) {
3057 continue;
3058 }
3059 let mut j = after_from;
3061 while j < bytes.len() && bytes[j].is_ascii_whitespace() {
3062 j += 1;
3063 }
3064 if j == after_from || !bytes[j..].starts_with(name_b) {
3065 continue;
3066 }
3067 let after_name = j + name_b.len();
3068 if after_name < bytes.len() && is_ident_byte(bytes[after_name]) {
3070 continue;
3071 }
3072 let mut out = String::with_capacity(sql.len() + view_sql.len() + name.len() + 8);
3075 out.push_str(&sql[..from_start]);
3076 out.push_str(&sql[from_start..j]);
3077 out.push('(');
3078 out.push_str(view_sql);
3079 out.push_str(") AS ");
3080 out.push_str(name);
3081 out.push_str(&sql[after_name..]);
3082 return out;
3083 }
3084 sql.to_string()
3085}
3086
3087fn is_ident_byte(b: u8) -> bool {
3088 b.is_ascii_alphanumeric() || b == b'_'
3089}
3090
3091fn normalize_sql(sql: &str) -> String {
3101 let b = sql.as_bytes();
3102 let n = b.len();
3103 let mut out: Vec<u8> = Vec::with_capacity(n);
3104 let mut want_space = false;
3107 let mut i = 0usize;
3108 while i < n {
3109 let c = b[i];
3110 if c.is_ascii_whitespace() {
3114 want_space = true;
3115 i += 1;
3116 continue;
3117 }
3118 if c == b'-' && i + 1 < n && b[i + 1] == b'-' {
3119 i += 2;
3121 while i < n && b[i] != b'\n' {
3122 i += 1;
3123 }
3124 want_space = !out.is_empty();
3125 continue;
3126 }
3127 if c == b'/' && i + 1 < n && b[i + 1] == b'*' {
3128 i += 2;
3131 let mut depth = 1usize;
3132 while i + 1 < n && depth > 0 {
3133 if b[i] == b'/' && b[i + 1] == b'*' {
3134 depth += 1;
3135 i += 2;
3136 } else if b[i] == b'*' && b[i + 1] == b'/' {
3137 depth -= 1;
3138 i += 2;
3139 } else {
3140 i += 1;
3141 }
3142 }
3143 want_space = !out.is_empty();
3144 continue;
3145 }
3146 if want_space && !out.is_empty() {
3148 out.push(b' ');
3149 }
3150 want_space = false;
3151 match c {
3152 b'E' | b'e' if i + 1 < n && b[i + 1] == b'\'' => {
3154 out.push(c);
3155 i += 1;
3156 i = copy_quoted(&mut out, b, i, n, b'\'');
3157 continue;
3158 }
3159 b'\'' => {
3161 i = copy_quoted(&mut out, b, i, n, b'\'');
3162 continue;
3163 }
3164 b'"' => {
3166 i = copy_quoted(&mut out, b, i, n, b'"');
3167 continue;
3168 }
3169 b'$' => {
3171 let (consumed, matched) = copy_dollar_quoted(&mut out, b, i, n);
3172 if matched {
3173 i = consumed;
3174 continue;
3175 }
3176 out.push(c);
3177 i += 1;
3178 continue;
3179 }
3180 _ => {
3181 out.push(c);
3182 i += 1;
3183 }
3184 }
3185 }
3186 String::from_utf8(out).unwrap_or_else(|_| sql.to_string())
3187}
3188
3189fn copy_quoted(out: &mut Vec<u8>, b: &[u8], start: usize, n: usize, delim: u8) -> usize {
3193 out.push(b[start]);
3194 let mut i = start + 1;
3195 while i < n {
3196 let c = b[i];
3197 out.push(c);
3198 if c == delim {
3199 if i + 1 < n && b[i + 1] == delim {
3201 out.push(b[i + 1]);
3202 i += 2;
3203 continue;
3204 }
3205 return i + 1;
3206 }
3207 i += 1;
3208 }
3209 i
3210}
3211
3212fn copy_dollar_quoted(out: &mut Vec<u8>, b: &[u8], start: usize, n: usize) -> (usize, bool) {
3216 let mut j = start + 1;
3220 let tag_start = j;
3221 while j < n && b[j] != b'$' && is_dollar_tag_byte(b[j]) {
3222 j += 1;
3223 }
3224 if j >= n || b[j] != b'$' {
3225 return (start + 1, false);
3226 }
3227 if tag_start < j && !(b[tag_start].is_ascii_alphabetic() || b[tag_start] == b'_') {
3228 return (start + 1, false);
3229 }
3230 let close_end = j + 1; let delim = &b[start..close_end];
3232 out.extend_from_slice(delim);
3234 let mut k = close_end;
3236 while k + delim.len() <= n {
3237 if &b[k..k + delim.len()] == delim {
3238 out.extend_from_slice(delim);
3239 return (k + delim.len(), true);
3240 }
3241 out.push(b[k]);
3242 k += 1;
3243 }
3244 out.extend_from_slice(&b[close_end..n]);
3246 (n, true)
3247}
3248
3249fn is_dollar_tag_byte(b: u8) -> bool {
3250 b.is_ascii_alphanumeric() || b == b'_'
3251}
3252
3253fn strip_prefix_ci<'a>(s: &'a str, prefix: &str) -> Option<&'a str> {
3255 let bytes = s.as_bytes();
3256 let pb = prefix.as_bytes();
3257 if bytes.len() >= pb.len() && bytes[..pb.len()].eq_ignore_ascii_case(pb) {
3258 Some(&s[pb.len()..])
3259 } else {
3260 None
3261 }
3262}
3263
3264const COLUMN_CONSTRAINTS: &[(&str, u32)] = &[
3277 ("primary key", ColumnFlags::PRIMARY_KEY),
3278 ("autoincrement", ColumnFlags::AUTO_INCREMENT),
3285 ("auto_increment", ColumnFlags::AUTO_INCREMENT),
3286];
3287
3288fn parse_column_constraints(constraint_text: &str) -> ColumnFlags {
3294 let normalized = constraint_text.to_lowercase();
3295 let mut flags = ColumnFlags::empty();
3296 for (phrase, bit) in COLUMN_CONSTRAINTS {
3297 if normalized.contains(phrase) {
3298 flags = flags.with(*bit);
3299 }
3300 }
3301 flags
3302}
3303
3304fn parse_sql_type(ty_str: &str) -> Result<mongreldb_core::schema::TypeId> {
3305 use mongreldb_core::schema::TypeId;
3306
3307 match ty_str.trim().trim_end_matches(';').to_lowercase().as_str() {
3308 "bigint" | "int8" | "int64" | "integer" | "int" => Ok(TypeId::Int64),
3309 "double" | "float8" | "float64" | "real" | "float" => Ok(TypeId::Float64),
3310 "varchar" | "text" | "string" | "bytes" => Ok(TypeId::Bytes),
3311 "boolean" | "bool" => Ok(TypeId::Bool),
3312 other => Err(MongrelQueryError::Schema(format!(
3313 "unsupported column type: {other}"
3314 ))),
3315 }
3316}
3317
3318fn parse_create_table(sql: &str) -> Result<(String, mongreldb_core::schema::Schema)> {
3324 use mongreldb_core::schema::*;
3325
3326 let open = sql
3327 .find('(')
3328 .ok_or(MongrelQueryError::Schema("CREATE TABLE missing '('".into()))?;
3329 let close = sql
3330 .rfind(')')
3331 .ok_or(MongrelQueryError::Schema("CREATE TABLE missing ')'".into()))?;
3332 let head = sql[..open].trim();
3333 let after_kw = strip_prefix_ci(head, "CREATE TABLE")
3334 .or_else(|| strip_prefix_ci(head, "create table"))
3335 .unwrap_or("")
3336 .trim();
3337 let after_kw = after_kw
3339 .strip_prefix("IF NOT EXISTS")
3340 .or_else(|| after_kw.strip_prefix("if not exists"))
3341 .map(str::trim)
3342 .unwrap_or(after_kw);
3343 let name = after_kw.trim_matches('"').to_string();
3344 if name.is_empty() {
3345 return Err(MongrelQueryError::Schema(
3346 "CREATE TABLE missing table name".into(),
3347 ));
3348 }
3349
3350 let body = &sql[open + 1..close];
3351 let mut columns = Vec::new();
3352 let schema_id: u64 = 0; for (i, raw) in body.split(',').enumerate() {
3354 let part = raw.trim();
3355 if part.is_empty() {
3356 continue;
3357 }
3358 let mut tokens = part.split_whitespace();
3359 let col_name = tokens
3360 .next()
3361 .ok_or(MongrelQueryError::Schema("missing column name".into()))?
3362 .trim_matches('"');
3363 let ty_str = tokens
3364 .next()
3365 .ok_or(MongrelQueryError::Schema("missing column type".into()))?
3366 .to_lowercase();
3367 let ty = parse_sql_type(&ty_str)?;
3368 let constraint_clause: String = tokens.collect::<Vec<_>>().join(" ");
3372 let flags = parse_column_constraints(&constraint_clause);
3373 columns.push(ColumnDef {
3374 id: (i + 1) as u16,
3375 name: col_name.to_string(),
3376 ty,
3377 flags,
3378 });
3379 }
3380
3381 Ok((
3382 name,
3383 Schema {
3384 schema_id,
3385 columns,
3386 indexes: vec![],
3387 colocation: vec![],
3388 constraints: Default::default(),
3389 clustered: false,
3390 },
3391 ))
3392}
3393
3394fn parse_drop_table(sql: &str) -> Result<(String, bool)> {
3396 let head = sql.trim();
3397 let after_kw = strip_prefix_ci(head, "DROP TABLE")
3398 .or_else(|| strip_prefix_ci(head, "drop table"))
3399 .unwrap_or("")
3400 .trim();
3401 let (rest, if_exists) = if let Some(r) = after_kw
3403 .strip_prefix("IF EXISTS")
3404 .or_else(|| after_kw.strip_prefix("if exists"))
3405 .map(str::trim)
3406 {
3407 (r, true)
3408 } else {
3409 (after_kw, false)
3410 };
3411 let name = rest.trim_matches(';').trim_matches('"').trim();
3412 if name.is_empty() {
3413 return Err(MongrelQueryError::Schema(
3414 "DROP TABLE missing table name".into(),
3415 ));
3416 }
3417 Ok((name.to_string(), if_exists))
3418}
3419
3420enum ParsedAlterTable {
3421 RenameTable {
3422 old_name: String,
3423 new_name: String,
3424 },
3425 RenameColumn {
3426 table_name: String,
3427 column_name: String,
3428 new_name: String,
3429 },
3430 AlterColumnType {
3431 table_name: String,
3432 column_name: String,
3433 ty: mongreldb_core::schema::TypeId,
3434 },
3435 SetNotNull {
3436 table_name: String,
3437 column_name: String,
3438 },
3439 DropNotNull {
3440 table_name: String,
3441 column_name: String,
3442 },
3443}
3444
3445fn current_column_flags(db: &Arc<Database>, table: &str, column: &str) -> Result<ColumnFlags> {
3446 let handle = db.table(table)?;
3447 let table = handle.lock();
3448 table
3449 .schema()
3450 .column(column)
3451 .map(|c| c.flags)
3452 .ok_or_else(|| MongrelQueryError::Schema(format!("unknown column {column}")))
3453}
3454
3455fn parse_alter_table(sql: &str) -> Result<ParsedAlterTable> {
3456 let trimmed = strip_statement_semicolon(sql.trim());
3457 let after_kw = strip_prefix_ci(trimmed, "ALTER TABLE")
3458 .ok_or_else(|| MongrelQueryError::Schema("not an ALTER TABLE statement".into()))?
3459 .trim();
3460 let (table_name, rest) = take_sql_ident(after_kw, "ALTER TABLE missing table name")?;
3461 let rest = rest.trim();
3462
3463 if let Some(after) = strip_prefix_ci(rest, "RENAME TO") {
3464 let new_name = parse_trailing_identifier(after, "ALTER TABLE missing new table name")?;
3465 return Ok(ParsedAlterTable::RenameTable {
3466 old_name: table_name,
3467 new_name,
3468 });
3469 }
3470
3471 if let Some(after) = strip_prefix_ci(rest, "RENAME COLUMN") {
3472 let (column_name, after_col) =
3473 take_sql_ident(after, "ALTER TABLE RENAME COLUMN missing column name")?;
3474 let after_to = strip_prefix_ci(after_col.trim(), "TO").ok_or_else(|| {
3475 MongrelQueryError::Schema("ALTER TABLE RENAME COLUMN missing TO".into())
3476 })?;
3477 let new_name = parse_trailing_identifier(
3478 after_to,
3479 "ALTER TABLE RENAME COLUMN missing new column name",
3480 )?;
3481 return Ok(ParsedAlterTable::RenameColumn {
3482 table_name,
3483 column_name,
3484 new_name,
3485 });
3486 }
3487
3488 let after_alter = strip_prefix_ci(rest, "ALTER COLUMN")
3489 .or_else(|| strip_prefix_ci(rest, "ALTER"))
3490 .ok_or_else(|| {
3491 MongrelQueryError::Schema(
3492 "ALTER TABLE must be RENAME TO, RENAME COLUMN, or ALTER COLUMN".into(),
3493 )
3494 })?;
3495 let (column_name, action) =
3496 take_sql_ident(after_alter, "ALTER TABLE ALTER COLUMN missing column name")?;
3497 let action = action.trim();
3498
3499 if let Some(after_type) =
3500 strip_prefix_ci(action, "TYPE").or_else(|| strip_prefix_ci(action, "SET DATA TYPE"))
3501 {
3502 let ty = parse_type_tail(after_type)?;
3503 return Ok(ParsedAlterTable::AlterColumnType {
3504 table_name,
3505 column_name,
3506 ty,
3507 });
3508 }
3509 if strip_prefix_ci(action, "SET NOT NULL").is_some() {
3510 return Ok(ParsedAlterTable::SetNotNull {
3511 table_name,
3512 column_name,
3513 });
3514 }
3515 if strip_prefix_ci(action, "DROP NOT NULL").is_some() {
3516 return Ok(ParsedAlterTable::DropNotNull {
3517 table_name,
3518 column_name,
3519 });
3520 }
3521
3522 Err(MongrelQueryError::Schema(
3523 "unsupported ALTER COLUMN action".into(),
3524 ))
3525}
3526
3527fn strip_statement_semicolon(s: &str) -> &str {
3528 s.trim().trim_end_matches(';').trim()
3529}
3530
3531fn take_sql_ident<'a>(s: &'a str, missing: &str) -> Result<(String, &'a str)> {
3532 let s = s.trim();
3533 if s.is_empty() {
3534 return Err(MongrelQueryError::Schema(missing.into()));
3535 }
3536 if let Some(rest) = s.strip_prefix('"') {
3537 let Some(end) = rest.find('"') else {
3538 return Err(MongrelQueryError::Schema(
3539 "unterminated quoted identifier".into(),
3540 ));
3541 };
3542 let ident = rest[..end].to_string();
3543 if ident.is_empty() {
3544 return Err(MongrelQueryError::Schema(missing.into()));
3545 }
3546 return Ok((ident, &rest[end + 1..]));
3547 }
3548 let end = s.find(|c: char| c.is_ascii_whitespace()).unwrap_or(s.len());
3549 let ident = s[..end].trim_matches('"').to_string();
3550 if ident.is_empty() {
3551 return Err(MongrelQueryError::Schema(missing.into()));
3552 }
3553 Ok((ident, &s[end..]))
3554}
3555
3556fn parse_trailing_identifier(s: &str, missing: &str) -> Result<String> {
3557 let (ident, rest) = take_sql_ident(s, missing)?;
3558 if !strip_statement_semicolon(rest).is_empty() {
3559 return Err(MongrelQueryError::Schema(
3560 "unexpected tokens after identifier".into(),
3561 ));
3562 }
3563 Ok(ident)
3564}
3565
3566fn parse_type_tail(s: &str) -> Result<mongreldb_core::schema::TypeId> {
3567 let tail = strip_statement_semicolon(s);
3568 let ty = tail
3569 .split_whitespace()
3570 .next()
3571 .ok_or_else(|| MongrelQueryError::Schema("ALTER COLUMN TYPE missing type".into()))?;
3572 parse_sql_type(ty)
3573}
3574
3575#[cfg(test)]
3576mod tests {
3577 use super::*;
3578
3579 #[test]
3580 fn normalize_collapses_and_trims_whitespace() {
3581 assert_eq!(normalize_sql("SELECT * FROM t"), "SELECT * FROM t");
3582 assert_eq!(normalize_sql(" SELECT * FROM t "), "SELECT * FROM t");
3583 assert_eq!(
3584 normalize_sql("\n\tSELECT\n*\nFROM\n\tt\n"),
3585 "SELECT * FROM t"
3586 );
3587 assert_eq!(
3588 normalize_sql("SELECT a, b FROM t"),
3589 normalize_sql("SELECT a, b FROM t")
3590 );
3591 }
3592
3593 #[test]
3594 fn normalize_preserves_string_literal_whitespace() {
3595 assert_eq!(
3596 normalize_sql("SELECT 'hello world' FROM t"),
3597 "SELECT 'hello world' FROM t"
3598 );
3599 assert_eq!(
3600 normalize_sql("SELECT 'it''s ok' FROM t"),
3601 "SELECT 'it''s ok' FROM t"
3602 );
3603 assert_eq!(
3604 normalize_sql(" SELECT 'a b' FROM t "),
3605 "SELECT 'a b' FROM t"
3606 );
3607 }
3608
3609 #[test]
3610 fn normalize_preserves_quoted_identifier_and_dollar_quote() {
3611 assert_eq!(
3612 normalize_sql(" SELECT \"my col\" FROM t "),
3613 "SELECT \"my col\" FROM t"
3614 );
3615 assert_eq!(
3616 normalize_sql(" SELECT $$a b$$ FROM t "),
3617 "SELECT $$a b$$ FROM t"
3618 );
3619 assert_eq!(
3620 normalize_sql("SELECT $tag$body with spaces$tag$ FROM t"),
3621 "SELECT $tag$body with spaces$tag$ FROM t"
3622 );
3623 }
3624
3625 #[test]
3626 fn normalize_strips_comments() {
3627 assert_eq!(
3628 normalize_sql("SELECT 1 -- trailing comment\nFROM t"),
3629 "SELECT 1 FROM t"
3630 );
3631 assert_eq!(
3632 normalize_sql("SELECT /* block */ 1 FROM t"),
3633 "SELECT 1 FROM t"
3634 );
3635 assert_eq!(
3637 normalize_sql("SELECT /* 'not a string' */ 1 FROM t"),
3638 "SELECT 1 FROM t"
3639 );
3640 assert_eq!(
3642 normalize_sql("SELECT /* outer /* inner */ still outer */ 1 FROM t"),
3643 "SELECT 1 FROM t"
3644 );
3645 }
3646
3647 #[test]
3648 fn normalize_escape_string_preserved() {
3649 assert_eq!(
3650 normalize_sql("SELECT E'line\\nbreak' FROM t"),
3651 "SELECT E'line\\nbreak' FROM t"
3652 );
3653 }
3654
3655 #[test]
3656 fn replace_from_view_matches_whole_word_only() {
3657 let out = replace_from_view("SELECT * FROM logs", "log", "SELECT 1");
3658 assert_eq!(out, "SELECT * FROM logs");
3659
3660 let out = replace_from_view("SELECT * FROM log", "log", "SELECT 1");
3661 assert_eq!(out, "SELECT * FROM (SELECT 1) AS log");
3662
3663 let out = replace_from_view("select * from log where x", "log", "SELECT 1");
3664 assert_eq!(out, "select * from (SELECT 1) AS log where x");
3665
3666 let out = replace_from_view("SELECT * FROM log)", "log", "SELECT 1");
3667 assert_eq!(out, "SELECT * FROM (SELECT 1) AS log)");
3668
3669 let out = replace_from_view("SELECT * xfrom log", "log", "SELECT 1");
3670 assert_eq!(out, "SELECT * xfrom log");
3671 }
3672
3673 #[test]
3674 fn compat_function_rewrite_handles_sqlite_compatibility_calls() {
3675 assert_eq!(
3676 rewrite_compat_function_calls("select max(id), min(id) from t"),
3677 "select max(id), min(id) from t"
3678 );
3679 assert_eq!(
3680 rewrite_compat_function_calls("select max(1, min(2, 3), 'max(4,5)')"),
3681 "select __mongreldb_scalar_max(1, __mongreldb_scalar_min(2, 3), 'max(4,5)')"
3682 );
3683 assert_eq!(
3684 rewrite_compat_function_calls("select /* max(1,2) */ min(1, (2 + 3))"),
3685 "select /* max(1,2) */ __mongreldb_scalar_min(1, (2 + 3))"
3686 );
3687 assert_eq!(
3688 rewrite_compat_function_calls("select max_value, min_value from t"),
3689 "select max_value, min_value from t"
3690 );
3691 assert_eq!(
3692 rewrite_compat_function_calls(
3693 "select group_concat(label), group_concat(label, '|') from t"
3694 ),
3695 "select string_agg(label, ','), string_agg(label, '|') from t"
3696 );
3697 assert_eq!(
3698 rewrite_compat_function_calls("select total(val), total(val) filter (where grp = 2) from t"),
3699 "select coalesce(cast(sum(val) as double), 0.0), coalesce(cast(sum(val) filter (where grp = 2) as double), 0.0) from t"
3700 );
3701 assert_eq!(
3702 rewrite_compat_function_calls(
3703 "select total(val) over (partition by grp order by id) from t"
3704 ),
3705 "select coalesce(cast(sum(val) over (partition by grp order by id) as double), 0.0) from t"
3706 );
3707 }
3708}