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