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