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