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