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