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
43pub type MongrelRecordBatchStream = datafusion::physical_plan::SendableRecordBatchStream;
44
45use arrow::array::{Array, ArrayRef, Int64Array, StringArray};
46use arrow::datatypes::SchemaRef;
47use arrow::record_batch::RecordBatch;
48use datafusion::catalog::{Session, TableProvider};
49use datafusion::common::{DataFusionError, Result as DFResult};
50use datafusion::logical_expr::{AggregateUDF, Expr, ScalarUDF, TableType, WindowUDF};
51use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
52use datafusion::physical_plan::ExecutionPlan;
53use datafusion::prelude::SessionContext;
54use mongreldb_core::{
55 AlterColumn, ColumnFlags, Cursor, Database, OwnedSnapshotGuard, Schema as CoreSchema, Snapshot,
56 Table, TypeId,
57};
58use parking_lot::Mutex;
59use std::borrow::Borrow;
60use std::collections::{HashMap, HashSet};
61use std::hash::Hash;
62use std::sync::Arc;
63
64pub struct MongrelProvider {
67 db: Arc<Mutex<Table>>,
68 schema: SchemaRef,
69 core_schema: CoreSchema,
70 snapshot: Option<Snapshot>,
71 _retention: Option<Arc<OwnedSnapshotGuard>>,
72 security: Option<ProviderSecurity>,
73}
74
75#[derive(Clone)]
76struct ProviderSecurity {
77 database: Arc<Database>,
78 table: String,
79 principal: Option<mongreldb_core::Principal>,
80 principal_catalog_bound: bool,
81}
82
83impl ProviderSecurity {
84 fn principal(&self) -> Option<mongreldb_core::Principal> {
85 self.principal.as_ref().and_then(|principal| {
86 if self.principal_catalog_bound {
87 self.database.resolve_principal(&principal.username)
88 } else {
89 Some(principal.clone())
90 }
91 })
92 }
93}
94
95#[derive(Debug, Clone)]
96pub(crate) struct ViewDef {
97 pub sql: String,
98 pub schema: CoreSchema,
99 pub input_types: HashMap<u16, Option<TypeId>>,
100}
101
102impl MongrelProvider {
103 pub fn new(db: Arc<Mutex<Table>>) -> Result<Self> {
104 let (schema, core_schema) = {
105 let db = db.lock();
106 (arrow_conv::arrow_schema(db.schema())?, db.schema().clone())
107 };
108 Ok(Self {
109 db,
110 schema,
111 core_schema,
112 snapshot: None,
113 _retention: None,
114 security: None,
115 })
116 }
117
118 pub(crate) fn new_secured(
119 db: Arc<Mutex<Table>>,
120 database: Arc<Database>,
121 table: String,
122 principal: Option<mongreldb_core::Principal>,
123 ) -> Result<Self> {
124 let core_schema = db.lock().schema().clone();
125 let schema = arrow_conv::arrow_schema(&core_schema)?;
126 Ok(Self {
127 db,
128 schema,
129 core_schema,
130 snapshot: None,
131 _retention: None,
132 security: Some(ProviderSecurity {
133 principal_catalog_bound: principal.as_ref().is_some_and(|principal| {
134 database.resolve_principal(&principal.username).is_some()
135 }),
136 database,
137 table,
138 principal,
139 }),
140 })
141 }
142
143 fn new_historical(
144 db: Arc<Mutex<Table>>,
145 snapshot: Snapshot,
146 retention: OwnedSnapshotGuard,
147 security: Option<ProviderSecurity>,
148 ) -> Result<Self> {
149 let full_schema = {
150 let db = db.lock();
151 db.schema().clone()
152 };
153 let core_schema = full_schema;
154 let schema = arrow_conv::arrow_schema(&core_schema)?;
155 Ok(Self {
156 db,
157 schema,
158 core_schema,
159 snapshot: Some(snapshot),
160 _retention: Some(Arc::new(retention)),
161 security,
162 })
163 }
164
165 pub fn arrow_schema(&self) -> SchemaRef {
166 self.schema.clone()
167 }
168}
169
170impl std::fmt::Debug for MongrelProvider {
171 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
172 f.debug_struct("MongrelProvider").finish_non_exhaustive()
173 }
174}
175
176struct AsOfRegistration {
177 ctx: SessionContext,
178 table_name: String,
179}
180
181impl Drop for AsOfRegistration {
182 fn drop(&mut self) {
183 let _ = self.ctx.deregister_table(&self.table_name);
184 }
185}
186
187struct AsOfQuery {
188 sql: String,
189 registration: AsOfRegistration,
190}
191
192#[async_trait::async_trait]
193impl TableProvider for MongrelProvider {
194 fn schema(&self) -> SchemaRef {
195 self.schema.clone()
196 }
197
198 fn table_type(&self) -> TableType {
199 TableType::Base
200 }
201
202 fn supports_filters_pushdown(
207 &self,
208 filters: &[&Expr],
209 ) -> DFResult<Vec<datafusion::logical_expr::TableProviderFilterPushDown>> {
210 use datafusion::logical_expr::TableProviderFilterPushDown;
211 if self.snapshot.is_some() || self.security.is_some() {
212 return Ok(vec![
213 TableProviderFilterPushDown::Unsupported;
214 filters.len()
215 ]);
216 }
217 let schema_ref = self.db.lock().schema().clone();
218 Ok(filters
219 .iter()
220 .map(|f| match translate_filter(f, &schema_ref) {
221 Some(
222 mongreldb_core::Condition::FmContains { .. }
223 | mongreldb_core::Condition::FmContainsAll { .. },
224 ) => TableProviderFilterPushDown::Inexact,
225 Some(_) => TableProviderFilterPushDown::Exact,
226 None => match translate_ann_search(f, &schema_ref)
227 .or_else(|| translate_sparse_match(f, &schema_ref))
228 {
229 Some(_) => TableProviderFilterPushDown::Exact,
230 None => TableProviderFilterPushDown::Unsupported,
231 },
232 })
233 .collect())
234 }
235
236 async fn scan(
237 &self,
238 _state: &dyn Session,
239 projection: Option<&Vec<usize>>,
240 filters: &[Expr],
241 _limit: Option<usize>,
242 ) -> DFResult<Arc<dyn ExecutionPlan>> {
243 let core_err = |e: mongreldb_core::MongrelError| {
244 DataFusionError::External(Box::new(MongrelQueryError::Core(e)))
245 };
246 if let Some(security) = &self.security {
247 let principal = security.principal();
248 let allowed = security
249 .database
250 .select_column_ids_for(&security.table, principal.as_ref())
251 .map_err(core_err)?;
252 let projected = projection
253 .map(|projection| {
254 projection
255 .iter()
256 .filter_map(|index| self.core_schema.columns.get(*index))
257 .map(|column| column.id)
258 .collect::<Vec<_>>()
259 })
260 .unwrap_or_else(|| {
261 self.core_schema
262 .columns
263 .iter()
264 .map(|column| column.id)
265 .collect()
266 });
267 if projected.iter().any(|column| !allowed.contains(column)) {
268 security
269 .database
270 .require_columns_for(
271 &security.table,
272 mongreldb_core::ColumnOperation::Select,
273 &projected,
274 principal.as_ref(),
275 )
276 .map_err(core_err)?;
277 }
278 let rows = {
279 let table = self.db.lock();
280 let snapshot = self.snapshot.unwrap_or_else(|| table.snapshot());
281 table.visible_rows(snapshot).map_err(core_err)?
282 };
283 let rows = security
284 .database
285 .secure_rows_for(&security.table, rows, principal.as_ref())
286 .map_err(core_err)?;
287 if projection.is_some_and(Vec::is_empty) {
288 return Ok(Arc::new(scan::MongrelScanExec::new_row_count(rows.len())));
289 }
290 let batch = arrow_conv::rows_to_batch(&rows, &self.core_schema)
291 .map_err(|error| DataFusionError::External(Box::new(error)))?;
292 let batch = match projection {
293 Some(projection) => batch.project(projection).map_err(|error| {
294 DataFusionError::External(Box::new(MongrelQueryError::Arrow(error.to_string())))
295 })?,
296 None => batch,
297 };
298 let schema = batch.schema();
299 let statistics = (0..batch.num_columns())
300 .map(|_| scan::to_col_statistics(None))
301 .collect();
302 return Ok(Arc::new(scan::MongrelScanExec::new_batch(
303 schema, batch, statistics,
304 )));
305 }
306 let mut db = self.db.lock();
307 db.require_select().map_err(core_err)?;
311 let historical = self.snapshot.is_some();
312 let snap = self.snapshot.unwrap_or_else(|| db.snapshot());
313 let schema_ref = db.schema().clone();
314
315 let translated: Vec<mongreldb_core::Condition> = if historical {
317 Vec::new()
318 } else {
319 filters
320 .iter()
321 .filter_map(|f| {
322 translate_filter(f, &schema_ref)
323 .or_else(|| translate_ann_search(f, &schema_ref))
324 .or_else(|| translate_sparse_match(f, &schema_ref))
325 })
326 .collect()
327 };
328
329 if !translated.is_empty() {
332 db.ensure_indexes_complete().map_err(core_err)?;
333 }
334
335 let empty_proj = projection.map(|p| p.is_empty()).unwrap_or(false);
339 if empty_proj {
340 let total: usize = if historical {
341 db.visible_rows(snap).map_err(core_err)?.len()
342 } else if translated.is_empty() {
343 mongreldb_core::trace::QueryTrace::record(|t| {
344 t.scan_mode = mongreldb_core::trace::ScanMode::CountMetadata;
345 });
346 db.count() as usize
347 } else if let Some(count) = db.count_conditions(&translated, snap).map_err(core_err)? {
348 count as usize
349 } else {
350 match schema_ref.columns.first() {
351 Some(cdef) => {
352 let one = [cdef.id];
353 let cols = match db
354 .query_columns_native_cached(&translated, Some(&one), snap)
355 .map_err(core_err)?
356 {
357 Some(c) => c,
358 None => db
359 .visible_columns_native(snap, Some(&one))
360 .map_err(core_err)?,
361 };
362 mongreldb_core::trace::QueryTrace::record(|t| {
363 t.scan_mode = mongreldb_core::trace::ScanMode::Materialized;
364 });
365 cols.first().map(|(_, c)| c.len()).unwrap_or(0)
366 }
367 None => 0,
368 }
369 };
370 return Ok(Arc::new(scan::MongrelScanExec::new_row_count(total)));
371 }
372
373 let (col_ids, scan_schema): (Vec<u16>, SchemaRef) = match projection {
378 Some(p) if !p.is_empty() => {
379 let ids = p.iter().map(|&idx| schema_ref.columns[idx].id).collect();
380 let fields: Vec<arrow::datatypes::Field> = p
381 .iter()
382 .map(|&idx| self.schema.field(idx).clone())
383 .collect();
384 (ids, Arc::new(arrow::datatypes::Schema::new(fields)))
385 }
386 _ => (
387 schema_ref.columns.iter().map(|c| c.id).collect(),
388 self.schema.clone(),
389 ),
390 };
391
392 let mut proj_pairs: Vec<(u16, mongreldb_core::schema::TypeId)> =
394 Vec::with_capacity(col_ids.len());
395 let mut types: Vec<mongreldb_core::schema::TypeId> = Vec::with_capacity(col_ids.len());
396 for cid in &col_ids {
397 let ty = schema_ref
398 .columns
399 .iter()
400 .find(|c| c.id == *cid)
401 .map(|c| c.ty.clone())
402 .ok_or_else(|| {
403 DataFusionError::External(Box::new(MongrelQueryError::Arrow(format!(
404 "unknown column {cid}"
405 ))))
406 })?;
407 proj_pairs.push((*cid, ty.clone()));
408 types.push(ty);
409 }
410
411 let col_stats_map = if !historical && translated.is_empty() {
415 db.exact_column_stats(snap, &col_ids).map_err(core_err)?
416 } else {
417 None
418 };
419 let column_stats: Vec<datafusion::physical_plan::ColumnStatistics> = col_ids
420 .iter()
421 .map(|cid| scan::to_col_statistics(col_stats_map.as_ref().and_then(|m| m.get(cid))))
422 .collect();
423
424 if !historical
428 && translated.is_empty()
429 && db.run_count() == 1
430 && db.memtable_is_empty()
431 && db.mutable_run_len() == 0
432 && db.single_run_is_clean()
433 {
434 let shadow = shadow::ArrowShadow::new(db.dir());
435 let run_ids: HashSet<u128> = db.run_ids().into_iter().collect();
436 shadow.sweep(&run_ids);
437 if let Some(&run_id) = run_ids.iter().next() {
438 if let Some(batch) = shadow.try_read(run_id) {
439 if let Some(projected) =
440 project_batch(&batch, &col_ids, &schema_ref, &scan_schema)
441 {
442 mongreldb_core::trace::QueryTrace::record(|t| {
443 t.scan_mode = mongreldb_core::trace::ScanMode::ArrowShadow;
444 });
445 return Ok(Arc::new(scan::MongrelScanExec::new_batch(
446 scan_schema,
447 projected,
448 column_stats,
449 )));
450 }
451 }
452 }
453 }
454
455 let cursor: Option<Box<dyn Cursor>> = db
462 .scan_cursor(snap, proj_pairs, &translated)
463 .map_err(core_err)?;
464 if let Some(cursor) = cursor {
465 let num_rows = cursor.remaining_rows();
466 let residual = extract_residual_filter(filters, &col_ids, &schema_ref);
468 return Ok(Arc::new(scan::MongrelScanExec::new_cursor(
469 scan_schema,
470 types,
471 cursor,
472 num_rows,
473 column_stats,
474 residual,
475 )));
476 }
477
478 let cols = if !translated.is_empty() {
481 match db
482 .query_columns_native_cached(&translated, Some(&col_ids), snap)
483 .map_err(core_err)?
484 {
485 Some(c) => c,
486 None => db
487 .visible_columns_native(snap, Some(&col_ids))
488 .map_err(core_err)?,
489 }
490 } else {
491 db.visible_columns_native(snap, Some(&col_ids))
492 .map_err(core_err)?
493 };
494
495 let mut ordered: Vec<mongreldb_core::columnar::NativeColumn> =
497 Vec::with_capacity(col_ids.len());
498 for cid in &col_ids {
499 let col = cols
500 .iter()
501 .find(|(id, _)| id == cid)
502 .map(|(_, c)| c.clone())
503 .ok_or_else(|| {
504 DataFusionError::External(Box::new(MongrelQueryError::Arrow(format!(
505 "missing column {cid}"
506 ))))
507 })?;
508 ordered.push(col);
509 }
510 let num_rows = ordered.first().map(|c| c.len()).unwrap_or(0);
511
512 let shadow_write: Option<(
514 std::path::PathBuf,
515 u128,
516 Vec<arrow::array::ArrayRef>,
517 SchemaRef,
518 )> = if !historical
519 && translated.is_empty()
520 && db.run_count() == 1
521 && db.memtable_is_empty()
522 && db.mutable_run_len() == 0
523 && db.single_run_is_clean()
524 {
525 let all_schema_ids: Vec<u16> = schema_ref.columns.iter().map(|c| c.id).collect();
526 if col_ids == all_schema_ids {
527 let dir = db.dir().to_path_buf();
528 let run_id = db.run_ids().first().copied();
529 run_id.map(|rid| {
530 let arrays = ordered
531 .iter()
532 .zip(types.iter())
533 .map(|(col, ty)| arrow_conv::native_to_array(ty.clone(), col))
534 .collect::<Result<_>>()
535 .unwrap_or_default();
536 (dir, rid, arrays, scan_schema.clone())
537 })
538 } else {
539 None
540 }
541 } else {
542 None
543 };
544
545 drop(db);
546
547 if let Some((dir, run_id, arrays, schema)) = shadow_write {
550 if let Ok(batch) = RecordBatch::try_new(schema, arrays) {
551 shadow::ArrowShadow::new(&dir).write(run_id, &batch);
552 }
553 }
554
555 mongreldb_core::trace::QueryTrace::record(|t| {
556 t.scan_mode = mongreldb_core::trace::ScanMode::Materialized;
557 t.row_materialized = true;
558 });
559 Ok(Arc::new(scan::MongrelScanExec::new(
560 scan_schema,
561 ordered,
562 types,
563 num_rows,
564 column_stats,
565 )))
566 }
567}
568
569fn project_batch(
573 batch: &RecordBatch,
574 col_ids: &[u16],
575 schema_ref: &mongreldb_core::schema::Schema,
576 scan_schema: &arrow::datatypes::SchemaRef,
577) -> Option<RecordBatch> {
578 let arrays: Vec<arrow::array::ArrayRef> = col_ids
580 .iter()
581 .map(|cid| {
582 let name = schema_ref
584 .columns
585 .iter()
586 .find(|c| c.id == *cid)
587 .map(|c| c.name.as_str())?;
588 let idx = batch.schema().index_of(name).ok()?;
590 Some(batch.column(idx).clone())
591 })
592 .collect::<Option<Vec<_>>>()?;
593 RecordBatch::try_new(scan_schema.clone(), arrays).ok()
594}
595
596pub(crate) fn translate_filter(
614 expr: &Expr,
615 schema: &mongreldb_core::Schema,
616) -> Option<mongreldb_core::Condition> {
617 use datafusion::common::ScalarValue;
618 use datafusion::logical_expr::{Between, BinaryExpr, Like, Operator};
619 use mongreldb_core::{ColumnFlags, Condition, IndexKind, TypeId, Value};
620
621 let int_val = |s: &ScalarValue| match s {
625 ScalarValue::Int8(Some(v)) => Some(*v as i64),
626 ScalarValue::Int16(Some(v)) => Some(*v as i64),
627 ScalarValue::Int32(Some(v)) => Some(*v as i64),
628 ScalarValue::Int64(Some(v)) => Some(*v),
629 ScalarValue::UInt8(Some(v)) => Some(*v as i64),
630 ScalarValue::UInt16(Some(v)) => Some(*v as i64),
631 ScalarValue::UInt32(Some(v)) => Some(*v as i64),
632 ScalarValue::UInt64(Some(v)) => Some(*v as i64),
633 ScalarValue::Date32(Some(v)) => Some(*v as i64),
634 ScalarValue::TimestampSecond(Some(v), _) => Some(*v),
635 ScalarValue::TimestampMillisecond(Some(v), _) => Some(*v),
636 ScalarValue::TimestampMicrosecond(Some(v), _) => Some(*v),
637 ScalarValue::TimestampNanosecond(Some(v), _) => Some(*v),
638 _ => None,
639 };
640 let float_val = |s: &ScalarValue| match s {
641 ScalarValue::Float32(Some(f)) => Some(*f as f64),
642 ScalarValue::Float64(Some(f)) => Some(*f),
643 _ => None,
644 };
645 let bytes_val = |s: &ScalarValue| match s {
646 ScalarValue::Utf8(Some(s)) => Some(s.as_bytes().to_vec()),
647 _ => None,
648 };
649 let _ = bytes_val; let val = |s: &ScalarValue| -> Option<Value> {
652 if let Some(i) = int_val(s) {
656 return Some(Value::Int64(i));
657 }
658 match s {
659 ScalarValue::Utf8(Some(s)) => Some(Value::Bytes(s.as_bytes().to_vec())),
660 ScalarValue::Float32(Some(f)) => Some(Value::Float64(*f as f64)),
661 ScalarValue::Float64(Some(f)) => Some(Value::Float64(*f)),
662 ScalarValue::Boolean(Some(b)) => Some(Value::Bool(*b)),
663 _ => None,
664 }
665 };
666
667 let col_def = |name: &str| schema.columns.iter().find(|c| c.name == name);
668 let has_fm = |cid: u16| {
669 schema
670 .indexes
671 .iter()
672 .any(|i| i.column_id == cid && i.kind == IndexKind::FmIndex)
673 };
674 let has_bitmap = |cid: u16| {
675 schema
676 .indexes
677 .iter()
678 .any(|i| i.column_id == cid && i.kind == IndexKind::Bitmap)
679 };
680
681 match expr {
682 Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
685 if *op == Operator::Or {
687 return try_or_as_bitmap_in(expr, schema);
688 }
689 let left = peel_cast(left);
691 let right = peel_cast(right);
692 let (col_name, scalar, flipped) = match (left.as_ref(), right.as_ref()) {
693 (Expr::Column(c), Expr::Literal(s, _)) => (&c.name, s, false),
694 (Expr::Literal(s, _), Expr::Column(c)) => (&c.name, s, true),
695 _ => return None,
696 };
697 let op = if flipped { flip_op(*op)? } else { *op };
698 let cdef = col_def(col_name)?;
699
700 if op == Operator::Eq {
702 let v = val(scalar)?;
703 if has_bitmap(cdef.id) {
704 return Some(Condition::BitmapEq {
705 column_id: cdef.id,
706 value: v.encode_key(),
707 });
708 }
709 if cdef.flags.contains(ColumnFlags::PRIMARY_KEY) {
710 return Some(Condition::Pk(v.encode_key()));
711 }
712 return None;
713 }
714
715 match cdef.ty {
718 TypeId::Int8
719 | TypeId::Int16
720 | TypeId::Int32
721 | TypeId::Int64
722 | TypeId::UInt8
723 | TypeId::UInt16
724 | TypeId::UInt32
725 | TypeId::UInt64
726 | TypeId::TimestampNanos
727 | TypeId::Date32 => {
728 let v = int_val(scalar)?;
729 let (lo, hi) = int_bounds(op, v)?;
730 Some(Condition::Range {
731 column_id: cdef.id,
732 lo,
733 hi,
734 })
735 }
736 TypeId::Float32 | TypeId::Float64 => {
737 let v = float_val(scalar)?;
738 let (lo, lo_inc, hi, hi_inc) = float_bounds(op, v)?;
739 Some(Condition::RangeF64 {
740 column_id: cdef.id,
741 lo,
742 lo_inclusive: lo_inc,
743 hi,
744 hi_inclusive: hi_inc,
745 })
746 }
747 _ => None,
748 }
749 }
750
751 Expr::Between(Between {
753 expr,
754 negated,
755 low,
756 high,
757 }) => {
758 if *negated {
759 return None;
760 }
761 let Expr::Column(c) = expr.as_ref() else {
762 return None;
763 };
764 let cdef = col_def(&c.name)?;
765 let (lo_s, hi_s) = match (low.as_ref(), high.as_ref()) {
766 (Expr::Literal(lo, _), Expr::Literal(hi, _)) => (lo, hi),
767 _ => return None,
768 };
769 match cdef.ty {
770 TypeId::Int8
771 | TypeId::Int16
772 | TypeId::Int32
773 | TypeId::Int64
774 | TypeId::UInt8
775 | TypeId::UInt16
776 | TypeId::UInt32
777 | TypeId::UInt64
778 | TypeId::TimestampNanos
779 | TypeId::Date32 => {
780 let (Some(lo), Some(hi)) = (int_val(lo_s), int_val(hi_s)) else {
781 return None;
782 };
783 Some(Condition::Range {
784 column_id: cdef.id,
785 lo,
786 hi,
787 })
788 }
789 TypeId::Float32 | TypeId::Float64 => {
790 let (Some(lo), Some(hi)) = (float_val(lo_s), float_val(hi_s)) else {
791 return None;
792 };
793 Some(Condition::RangeF64 {
794 column_id: cdef.id,
795 lo,
796 lo_inclusive: true,
797 hi,
798 hi_inclusive: true,
799 })
800 }
801 _ => None,
802 }
803 }
804
805 Expr::Like(Like {
807 negated,
808 expr,
809 pattern,
810 ..
811 }) => {
812 if *negated {
813 return None;
814 }
815 let Expr::Column(c) = expr.as_ref() else {
816 return None;
817 };
818 let Expr::Literal(ScalarValue::Utf8(Some(pat)), _) = pattern.as_ref() else {
819 return None;
820 };
821 let cdef = col_def(&c.name)?;
822 if has_bitmap(cdef.id) {
826 if let Some(prefix) = anchored_like_prefix(pat) {
827 return Some(Condition::BytesPrefix {
828 column_id: cdef.id,
829 prefix: mongreldb_core::Value::Bytes(prefix.as_bytes().to_vec())
830 .encode_key(),
831 });
832 }
833 }
834 if !has_fm(cdef.id) {
835 return None;
836 }
837 let segments: Vec<Vec<u8>> = pat
841 .split(['%', '_'])
842 .filter(|s| s.len() >= 3)
843 .map(|s| s.as_bytes().to_vec())
844 .collect();
845 match segments.len() {
846 0 => longest_like_segment(pat).map(|seg| Condition::FmContains {
847 column_id: cdef.id,
848 pattern: seg,
849 }),
850 1 => Some(Condition::FmContains {
851 column_id: cdef.id,
852 pattern: segments.into_iter().next().unwrap(),
853 }),
854 _ => Some(Condition::FmContainsAll {
855 column_id: cdef.id,
856 patterns: segments,
857 }),
858 }
859 }
860
861 Expr::InList(il) if !il.negated => {
865 let Expr::Column(c) = il.expr.as_ref() else {
866 return None;
867 };
868 let cdef = col_def(&c.name)?;
869 if !has_bitmap(cdef.id) {
870 return None;
871 }
872 let values: Vec<Vec<u8>> = il
873 .list
874 .iter()
875 .filter_map(|e| match e {
876 Expr::Literal(s, _) => val(s).map(|v| v.encode_key()),
877 _ => None,
878 })
879 .collect();
880 if values.is_empty() || values.len() != il.list.len() {
881 return None;
882 }
883 Some(Condition::BitmapIn {
884 column_id: cdef.id,
885 values,
886 })
887 }
888
889 Expr::IsNull(inner) => {
891 let col_name = match inner.as_ref() {
892 Expr::Column(c) => &c.name,
893 _ => return None,
894 };
895 let cdef = col_def(col_name)?;
896 Some(Condition::IsNull { column_id: cdef.id })
897 }
898
899 Expr::IsNotNull(inner) => {
901 let col_name = match inner.as_ref() {
902 Expr::Column(c) => &c.name,
903 _ => return None,
904 };
905 let cdef = col_def(col_name)?;
906 Some(Condition::IsNotNull { column_id: cdef.id })
907 }
908
909 _ => None,
910 }
911}
912
913pub(crate) fn extract_residual_filter(
917 filters: &[Expr],
918 col_ids: &[u16],
919 schema: &mongreldb_core::Schema,
920) -> Option<std::sync::Arc<scan::ResidualFilter>> {
921 use datafusion::common::ScalarValue;
922 use datafusion::logical_expr::Like;
923 for f in filters {
924 if let Expr::Like(Like {
925 negated: false,
926 expr,
927 pattern,
928 ..
929 }) = f
930 {
931 let Expr::Column(c) = expr.as_ref() else {
932 continue;
933 };
934 let Expr::Literal(ScalarValue::Utf8(Some(pat)), _) = pattern.as_ref() else {
935 continue;
936 };
937 let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
938 let col_idx = col_ids.iter().position(|&id| id == cdef.id)?;
939 return Some(std::sync::Arc::new(scan::ResidualFilter::new(
940 col_idx,
941 pat.as_bytes().to_vec(),
942 )));
943 }
944 }
945 None
946}
947
948pub(crate) fn translate_ann_search(
954 expr: &Expr,
955 schema: &mongreldb_core::Schema,
956) -> Option<mongreldb_core::Condition> {
957 use datafusion::common::ScalarValue;
958 use mongreldb_core::Condition;
959
960 let Expr::ScalarFunction(sf) = expr else {
961 return None;
962 };
963 if !sf.func.name().eq_ignore_ascii_case("ann_search") || sf.args.len() != 3 {
964 return None;
965 }
966 let (Expr::Column(c), query_expr, k_expr) = (&sf.args[0], &sf.args[1], &sf.args[2]) else {
967 return None;
968 };
969 let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
970 let json = match query_expr {
971 Expr::Literal(ScalarValue::Utf8(Some(s)), _) => s.as_str(),
972 _ => return None,
973 };
974 let k: i64 = match k_expr {
975 Expr::Literal(scalar, _) => match scalar {
976 ScalarValue::Int64(Some(k)) => *k,
977 ScalarValue::UInt64(Some(k)) => *k as i64,
978 ScalarValue::Int32(Some(k)) => *k as i64,
979 _ => return None,
980 },
981 _ => return None,
982 };
983 let query: Vec<f32> = serde_json::from_str(json).ok()?;
984 Some(Condition::Ann {
985 column_id: cdef.id,
986 query,
987 k: k.max(1) as usize,
988 })
989}
990
991pub(crate) fn translate_sparse_match(
996 expr: &Expr,
997 schema: &mongreldb_core::Schema,
998) -> Option<mongreldb_core::Condition> {
999 use datafusion::common::ScalarValue;
1000 use mongreldb_core::Condition;
1001
1002 let Expr::ScalarFunction(sf) = expr else {
1003 return None;
1004 };
1005 if !sf.func.name().eq_ignore_ascii_case("sparse_match") || sf.args.len() != 3 {
1006 return None;
1007 }
1008 let (Expr::Column(c), query_expr, k_expr) = (&sf.args[0], &sf.args[1], &sf.args[2]) else {
1009 return None;
1010 };
1011 let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
1012 let json = match query_expr {
1013 Expr::Literal(ScalarValue::Utf8(Some(s)), _) => s.as_str(),
1014 _ => return None,
1015 };
1016 let k: i64 = match k_expr {
1017 Expr::Literal(scalar, _) => match scalar {
1018 ScalarValue::Int64(Some(k)) => *k,
1019 ScalarValue::UInt64(Some(k)) => *k as i64,
1020 ScalarValue::Int32(Some(k)) => *k as i64,
1021 _ => return None,
1022 },
1023 _ => return None,
1024 };
1025 let query: Vec<(u32, f32)> = serde_json::from_str(json).ok()?;
1026 Some(Condition::SparseMatch {
1027 column_id: cdef.id,
1028 query,
1029 k: k.max(1) as usize,
1030 })
1031}
1032
1033fn flip_op(op: datafusion::logical_expr::Operator) -> Option<datafusion::logical_expr::Operator> {
1035 use datafusion::logical_expr::Operator;
1036 Some(match op {
1037 Operator::Eq => Operator::Eq,
1038 Operator::Lt => Operator::Gt,
1039 Operator::Gt => Operator::Lt,
1040 Operator::LtEq => Operator::GtEq,
1041 Operator::GtEq => Operator::LtEq,
1042 _ => return None,
1043 })
1044}
1045
1046fn int_bounds(op: datafusion::logical_expr::Operator, v: i64) -> Option<(i64, i64)> {
1049 use datafusion::logical_expr::Operator;
1050 Some(match op {
1051 Operator::Gt => (v.saturating_add(1), i64::MAX),
1052 Operator::GtEq => (v, i64::MAX),
1053 Operator::Lt => (i64::MIN, v.saturating_sub(1)),
1054 Operator::LtEq => (i64::MIN, v),
1055 _ => return None,
1056 })
1057}
1058
1059fn float_bounds(op: datafusion::logical_expr::Operator, v: f64) -> Option<(f64, bool, f64, bool)> {
1061 use datafusion::logical_expr::Operator;
1062 Some(match op {
1063 Operator::Gt => (v, false, f64::INFINITY, false),
1064 Operator::GtEq => (v, true, f64::INFINITY, false),
1065 Operator::Lt => (f64::NEG_INFINITY, false, v, false),
1066 Operator::LtEq => (f64::NEG_INFINITY, false, v, true),
1067 _ => return None,
1068 })
1069}
1070
1071fn longest_like_segment(pat: &str) -> Option<Vec<u8>> {
1078 pat.split(['%', '_'])
1079 .map(|s| s.as_bytes())
1080 .max_by_key(|s| s.len())
1081 .filter(|s| !s.is_empty())
1082 .map(|s| s.to_vec())
1083}
1084
1085fn anchored_like_prefix(pat: &str) -> Option<&str> {
1090 let rest = pat.strip_suffix('%')?;
1091 if rest.is_empty() || rest.contains(['%', '_']) {
1092 return None;
1093 }
1094 Some(rest)
1095}
1096
1097fn peel_cast(expr: &Expr) -> std::borrow::Cow<'_, Expr> {
1101 match expr {
1102 Expr::Cast(datafusion::logical_expr::Cast { expr, .. }) => std::borrow::Cow::Borrowed(expr),
1103 _ => std::borrow::Cow::Borrowed(expr),
1104 }
1105}
1106
1107fn try_or_as_bitmap_in(
1112 expr: &Expr,
1113 schema: &mongreldb_core::Schema,
1114) -> Option<mongreldb_core::Condition> {
1115 use datafusion::logical_expr::{BinaryExpr, Operator};
1116 let mut values: Vec<Vec<u8>> = Vec::new();
1117 let mut target_col: Option<u16> = None;
1118 let mut stack = vec![expr];
1119 while let Some(e) = stack.pop() {
1120 match e {
1121 Expr::BinaryExpr(BinaryExpr {
1122 left,
1123 op: Operator::Or,
1124 right,
1125 }) => {
1126 stack.push(left);
1127 stack.push(right);
1128 }
1129 Expr::BinaryExpr(BinaryExpr {
1130 left,
1131 op: Operator::Eq,
1132 right,
1133 }) => {
1134 let (col_name, scalar) = match (left.as_ref(), right.as_ref()) {
1135 (Expr::Column(c), Expr::Literal(s, _)) => (&c.name, s),
1136 (Expr::Literal(s, _), Expr::Column(c)) => (&c.name, s),
1137 _ => return None,
1138 };
1139 let cdef = schema.columns.iter().find(|c| &c.name == col_name)?;
1140 if !schema
1141 .indexes
1142 .iter()
1143 .any(|i| i.column_id == cdef.id && i.kind == mongreldb_core::IndexKind::Bitmap)
1144 {
1145 return None;
1146 }
1147 match target_col {
1148 None => target_col = Some(cdef.id),
1149 Some(id) if id != cdef.id => return None,
1150 _ => {}
1151 }
1152 let v = match scalar {
1153 datafusion::common::ScalarValue::Int64(Some(v)) => {
1154 mongreldb_core::Value::Int64(*v)
1155 }
1156 datafusion::common::ScalarValue::Utf8(Some(s)) => {
1157 mongreldb_core::Value::Bytes(s.as_bytes().to_vec())
1158 }
1159 datafusion::common::ScalarValue::Float64(Some(f)) => {
1160 mongreldb_core::Value::Float64(*f)
1161 }
1162 datafusion::common::ScalarValue::Boolean(Some(b)) => {
1163 mongreldb_core::Value::Bool(*b)
1164 }
1165 _ => return None,
1166 };
1167 values.push(v.encode_key());
1168 }
1169 _ => return None,
1170 }
1171 }
1172 let col_id = target_col?;
1173 if values.is_empty() {
1174 return None;
1175 }
1176 Some(mongreldb_core::Condition::BitmapIn {
1177 column_id: col_id,
1178 values,
1179 })
1180}
1181
1182fn sp_ident_name(expr: &sqlparser::ast::Expr) -> Option<&str> {
1189 use sqlparser::ast::Expr;
1190 match expr {
1191 Expr::Identifier(ident) => Some(ident.value.as_str()),
1192 Expr::CompoundIdentifier(idents) => idents.last().map(|i| i.value.as_str()),
1193 _ => None,
1194 }
1195}
1196
1197fn sp_literal(expr: &sqlparser::ast::Expr) -> Option<mongreldb_core::Value> {
1200 use sqlparser::ast::Expr;
1201 let v = match expr {
1202 Expr::Value(v) => v,
1203 _ => return None,
1204 };
1205 use sqlparser::ast::Value as SpValue;
1206 match &v.value {
1207 SpValue::Number(s, _) => s
1208 .parse::<i64>()
1209 .map(mongreldb_core::Value::Int64)
1210 .or_else(|_| s.parse::<f64>().map(mongreldb_core::Value::Float64))
1211 .ok(),
1212 SpValue::SingleQuotedString(s) => Some(mongreldb_core::Value::Bytes(s.as_bytes().to_vec())),
1213 SpValue::Boolean(b) => Some(mongreldb_core::Value::Bool(*b)),
1214 _ => None,
1215 }
1216}
1217
1218fn is_int_ty(ty: mongreldb_core::schema::TypeId) -> bool {
1219 use mongreldb_core::schema::TypeId::*;
1220 matches!(
1221 ty,
1222 Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 | TimestampNanos | Date32
1223 )
1224}
1225
1226fn is_float_ty(ty: mongreldb_core::schema::TypeId) -> bool {
1227 matches!(
1228 ty,
1229 mongreldb_core::schema::TypeId::Float32 | mongreldb_core::schema::TypeId::Float64
1230 )
1231}
1232
1233fn translate_sqlparser_predicate(
1236 expr: &sqlparser::ast::Expr,
1237 schema: &mongreldb_core::Schema,
1238) -> Option<mongreldb_core::Condition> {
1239 use mongreldb_core::{schema::ColumnFlags, Condition, IndexKind, Value};
1240 use sqlparser::ast::{BinaryOperator, Expr};
1241
1242 let col_def = |name: &str| schema.columns.iter().find(|c| c.name == name);
1243 let has_bitmap = |cid: u16| {
1244 schema
1245 .indexes
1246 .iter()
1247 .any(|i| i.column_id == cid && i.kind == IndexKind::Bitmap)
1248 };
1249
1250 match expr {
1251 Expr::BinaryOp {
1253 left,
1254 op: BinaryOperator::Or,
1255 right,
1256 } => {
1257 let mut values: Vec<Vec<u8>> = Vec::new();
1258 let mut target: Option<u16> = None;
1259 let mut stack: Vec<&Expr> = vec![left.as_ref(), right.as_ref()];
1260 while let Some(e) = stack.pop() {
1261 match e {
1262 Expr::BinaryOp {
1263 left,
1264 op: BinaryOperator::Or,
1265 right,
1266 } => {
1267 stack.push(left.as_ref());
1268 stack.push(right.as_ref());
1269 }
1270 Expr::BinaryOp {
1271 left,
1272 op: BinaryOperator::Eq,
1273 right,
1274 } => {
1275 let (name, lit) = match (left.as_ref(), right.as_ref()) {
1276 (l, r) if sp_ident_name(l).is_some() && sp_literal(r).is_some() => {
1277 (l, r)
1278 }
1279 (l, r) if sp_ident_name(r).is_some() && sp_literal(l).is_some() => {
1280 (r, l)
1281 }
1282 _ => return None,
1283 };
1284 let cdef = col_def(sp_ident_name(name)?)?;
1285 if !has_bitmap(cdef.id) {
1286 return None;
1287 }
1288 match target {
1289 None => target = Some(cdef.id),
1290 Some(id) if id != cdef.id => return None,
1291 _ => {}
1292 }
1293 values.push(sp_literal(lit)?.encode_key());
1294 }
1295 _ => return None,
1296 }
1297 }
1298 let cid = target?;
1299 (!values.is_empty()).then_some(Condition::BitmapIn {
1300 column_id: cid,
1301 values,
1302 })
1303 }
1304 Expr::BinaryOp { left, op, right } => {
1306 let flipped;
1307 let (col_expr, lit_expr) = match (
1308 sp_ident_name(left),
1309 sp_literal(right),
1310 sp_ident_name(right),
1311 sp_literal(left),
1312 ) {
1313 (Some(_), Some(_), _, _) => {
1314 flipped = false;
1315 (left.as_ref(), right.as_ref())
1316 }
1317 (_, _, Some(_), Some(_)) => {
1318 flipped = true;
1319 (right.as_ref(), left.as_ref())
1320 }
1321 _ => return None,
1322 };
1323 let name = sp_ident_name(col_expr)?;
1324 let cdef = col_def(name)?;
1325 let v = sp_literal(lit_expr)?;
1326 use sqlparser::ast::BinaryOperator::*;
1327 match op {
1330 Eq => {
1331 if has_bitmap(cdef.id) {
1332 Some(Condition::BitmapEq {
1333 column_id: cdef.id,
1334 value: v.encode_key(),
1335 })
1336 } else if cdef.flags.contains(ColumnFlags::PRIMARY_KEY) {
1337 Some(Condition::Pk(v.encode_key()))
1338 } else {
1339 None
1340 }
1341 }
1342 Lt | LtEq | Gt | GtEq if is_int_ty(cdef.ty.clone()) => {
1343 let n = match v {
1344 Value::Int64(n) => n,
1345 _ => return None,
1346 };
1347 let (lo, hi) = match (flipped, op) {
1349 (false, Lt) | (true, Gt) => (i64::MIN, n.saturating_sub(1)),
1350 (false, Gt) | (true, Lt) => (n.saturating_add(1), i64::MAX),
1351 (false, LtEq) | (true, GtEq) => (i64::MIN, n),
1352 (false, GtEq) | (true, LtEq) => (n, i64::MAX),
1353 _ => (i64::MIN, i64::MAX),
1354 };
1355 Some(Condition::Range {
1356 column_id: cdef.id,
1357 lo,
1358 hi,
1359 })
1360 }
1361 Lt | LtEq | Gt | GtEq if is_float_ty(cdef.ty.clone()) => {
1362 let f = match v {
1363 Value::Float64(f) => f,
1364 _ => return None,
1365 };
1366 let (lo, li, hi, hi_i) = match (flipped, op) {
1367 (false, Lt) | (true, Gt) => (f64::NEG_INFINITY, true, f, false),
1368 (false, Gt) | (true, Lt) => (f, false, f64::INFINITY, true),
1369 (false, LtEq) | (true, GtEq) => (f64::NEG_INFINITY, true, f, true),
1370 (false, GtEq) | (true, LtEq) => (f, true, f64::INFINITY, true),
1371 _ => (f64::NEG_INFINITY, true, f64::INFINITY, true),
1372 };
1373 Some(Condition::RangeF64 {
1374 column_id: cdef.id,
1375 lo,
1376 lo_inclusive: li,
1377 hi,
1378 hi_inclusive: hi_i,
1379 })
1380 }
1381 _ => None,
1382 }
1383 }
1384 Expr::Between {
1385 expr,
1386 negated,
1387 low,
1388 high,
1389 } if !negated => {
1390 let name = sp_ident_name(expr)?;
1391 let cdef = col_def(name)?;
1392 if is_int_ty(cdef.ty.clone()) {
1393 let lo = match sp_literal(low)? {
1394 Value::Int64(n) => n,
1395 _ => return None,
1396 };
1397 let hi = match sp_literal(high)? {
1398 Value::Int64(n) => n,
1399 _ => return None,
1400 };
1401 Some(Condition::Range {
1402 column_id: cdef.id,
1403 lo,
1404 hi,
1405 })
1406 } else if is_float_ty(cdef.ty.clone()) {
1407 let lo = match sp_literal(low)? {
1408 Value::Float64(f) => f,
1409 _ => return None,
1410 };
1411 let hi = match sp_literal(high)? {
1412 Value::Float64(f) => f,
1413 _ => return None,
1414 };
1415 Some(Condition::RangeF64 {
1416 column_id: cdef.id,
1417 lo,
1418 lo_inclusive: true,
1419 hi,
1420 hi_inclusive: true,
1421 })
1422 } else {
1423 None
1424 }
1425 }
1426 Expr::InList {
1427 expr,
1428 list,
1429 negated,
1430 } if !negated => {
1431 let name = sp_ident_name(expr)?;
1432 let cdef = col_def(name)?;
1433 if !has_bitmap(cdef.id) {
1434 return None;
1435 }
1436 let values: Vec<Vec<u8>> = list
1437 .iter()
1438 .map(|e| sp_literal(e).map(|v| v.encode_key()))
1439 .collect::<Option<_>>()?;
1440 (!values.is_empty()).then_some(Condition::BitmapIn {
1441 column_id: cdef.id,
1442 values,
1443 })
1444 }
1445 Expr::IsNull(inner) => {
1448 let cid = col_def(sp_ident_name(inner)?)?.id;
1449 Some(Condition::IsNull { column_id: cid })
1450 }
1451 Expr::IsNotNull(inner) => {
1452 let cid = col_def(sp_ident_name(inner)?)?.id;
1453 Some(Condition::IsNotNull { column_id: cid })
1454 }
1455 _ => None,
1456 }
1457}
1458
1459fn translate_sqlparser_filter(
1462 expr: &sqlparser::ast::Expr,
1463 schema: &mongreldb_core::Schema,
1464) -> Option<Vec<mongreldb_core::Condition>> {
1465 use sqlparser::ast::{BinaryOperator, Expr};
1466 let mut out = Vec::new();
1467 let mut stack = vec![expr];
1468 while let Some(e) = stack.pop() {
1469 match e {
1470 Expr::BinaryOp {
1471 left,
1472 op: BinaryOperator::And,
1473 right,
1474 } => {
1475 stack.push(left.as_ref());
1476 stack.push(right.as_ref());
1477 }
1478 other => out.push(translate_sqlparser_predicate(other, schema)?),
1479 }
1480 }
1481 Some(out)
1482}
1483
1484pub struct MongrelSession {
1488 ctx: SessionContext,
1489 db: Option<Arc<Mutex<Table>>>,
1490 database: Option<Arc<Database>>,
1494 principal: Option<mongreldb_core::Principal>,
1495 principal_catalog_bound: bool,
1496 cache: ResultCache,
1497 plan_cache: parking_lot::Mutex<BoundedLru<String, datafusion::logical_expr::LogicalPlan>>,
1499 tables: parking_lot::Mutex<HashMap<String, Arc<Mutex<Table>>>>,
1501 views: parking_lot::Mutex<HashMap<String, ViewDef>>,
1505 attached_databases: parking_lot::Mutex<HashMap<String, Arc<Database>>>,
1509 sql_txn: parking_lot::Mutex<Option<Vec<commands::PendingSqlOp>>>,
1513 savepoints: parking_lot::Mutex<Vec<(String, usize)>>,
1516 sql_fn_state: Arc<extended_sql_functions::ExtendedSqlState>,
1518 external_modules: Arc<ExternalModuleRegistry>,
1521 prepared_names: parking_lot::Mutex<std::collections::HashSet<String>>,
1525}
1526
1527type CacheKey = (String, u64);
1529type ResultCache = parking_lot::Mutex<BoundedLru<CacheKey, Arc<Vec<RecordBatch>>>>;
1530
1531const SESSION_CACHE_MAX_ENTRIES: usize = 1024;
1532
1533struct BoundedLru<K, V> {
1534 entries: HashMap<K, (V, u64)>,
1535 clock: u64,
1536 capacity: usize,
1537}
1538
1539impl<K: Clone + Eq + Hash, V> BoundedLru<K, V> {
1540 fn new(capacity: usize) -> Self {
1541 Self {
1542 entries: HashMap::new(),
1543 clock: 0,
1544 capacity: capacity.max(1),
1545 }
1546 }
1547
1548 fn get<Q>(&mut self, key: &Q) -> Option<&V>
1549 where
1550 K: Borrow<Q>,
1551 Q: Eq + Hash + ?Sized,
1552 {
1553 self.clock = self.clock.wrapping_add(1);
1554 let (value, last_used) = self.entries.get_mut(key)?;
1555 *last_used = self.clock;
1556 Some(value)
1557 }
1558
1559 fn insert(&mut self, key: K, value: V) {
1560 self.clock = self.clock.wrapping_add(1);
1561 if let Some(entry) = self.entries.get_mut(&key) {
1562 *entry = (value, self.clock);
1563 return;
1564 }
1565 if self.entries.len() >= self.capacity {
1566 if let Some(oldest) = self
1568 .entries
1569 .iter()
1570 .min_by_key(|(_, (_, last_used))| *last_used)
1571 .map(|(key, _)| key.clone())
1572 {
1573 self.entries.remove(&oldest);
1574 }
1575 }
1576 self.entries.insert(key, (value, self.clock));
1577 }
1578
1579 fn clear(&mut self) {
1580 self.entries.clear();
1581 self.clock = 0;
1582 }
1583}
1584
1585fn buffered_stream(batches: Vec<RecordBatch>) -> MongrelRecordBatchStream {
1586 let schema = batches
1587 .first()
1588 .map(RecordBatch::schema)
1589 .unwrap_or_else(|| Arc::new(arrow::datatypes::Schema::empty()));
1590 let stream = futures::stream::iter(batches.into_iter().map(Ok::<RecordBatch, DataFusionError>));
1591 Box::pin(RecordBatchStreamAdapter::new(schema, stream))
1592}
1593
1594impl MongrelSession {
1595 pub fn new(db: Table) -> Self {
1599 let db = Arc::new(Mutex::new(db));
1600 let ctx = SessionContext::new();
1601 let sql_fn_state = Arc::new(extended_sql_functions::ExtendedSqlState::default());
1602 register_mongrel_functions(&ctx, Arc::clone(&sql_fn_state));
1603 let external_modules = Arc::new(ExternalModuleRegistry::default());
1604 Self {
1605 ctx,
1606 db: Some(db),
1607 database: None,
1608 principal: None,
1609 principal_catalog_bound: false,
1610 cache: parking_lot::Mutex::new(BoundedLru::new(SESSION_CACHE_MAX_ENTRIES)),
1611 plan_cache: parking_lot::Mutex::new(BoundedLru::new(SESSION_CACHE_MAX_ENTRIES)),
1612 tables: parking_lot::Mutex::new(HashMap::new()),
1613 views: parking_lot::Mutex::new(HashMap::new()),
1614 attached_databases: parking_lot::Mutex::new(HashMap::new()),
1615 savepoints: parking_lot::Mutex::new(Vec::new()),
1616 sql_txn: parking_lot::Mutex::new(None),
1617 sql_fn_state,
1618 external_modules,
1619 prepared_names: parking_lot::Mutex::new(std::collections::HashSet::new()),
1620 }
1621 }
1622
1623 pub fn new_with_external_modules(
1624 db: Table,
1625 modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
1626 ) -> Result<Self> {
1627 let session = Self::new(db);
1628 for module in modules {
1629 session.register_external_module(module)?;
1630 }
1631 Ok(session)
1632 }
1633
1634 pub fn open(database: Arc<Database>) -> Result<Self> {
1639 Self::open_with_external_modules(database, std::iter::empty())
1640 }
1641
1642 pub fn open_as(database: Arc<Database>, principal: mongreldb_core::Principal) -> Result<Self> {
1643 Self::open_with_external_modules_as(database, std::iter::empty(), Some(principal))
1644 }
1645
1646 pub fn open_with_external_modules(
1647 database: Arc<Database>,
1648 modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
1649 ) -> Result<Self> {
1650 Self::open_with_external_modules_as(database, modules, None)
1651 }
1652
1653 pub fn open_with_external_modules_as(
1654 database: Arc<Database>,
1655 modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
1656 principal: Option<mongreldb_core::Principal>,
1657 ) -> Result<Self> {
1658 let ctx = SessionContext::new();
1659 let sql_fn_state = Arc::new(extended_sql_functions::ExtendedSqlState::default());
1660 register_mongrel_functions(&ctx, Arc::clone(&sql_fn_state));
1661 let external_modules = Arc::new(ExternalModuleRegistry::default());
1662 let principal_catalog_bound = principal
1663 .as_ref()
1664 .is_some_and(|principal| database.resolve_principal(&principal.username).is_some());
1665 for module in modules {
1666 external_modules.register(module)?;
1667 }
1668
1669 let mut tables: HashMap<String, Arc<Mutex<Table>>> = HashMap::new();
1670 for name in database.table_names() {
1671 let handle = database.table(&name)?;
1672 let provider = MongrelProvider::new_secured(
1673 handle.clone(),
1674 Arc::clone(&database),
1675 name.clone(),
1676 principal.clone(),
1677 )?;
1678 ctx.register_table(&name, Arc::new(provider))
1679 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1680 tables.insert(name, handle);
1681 }
1682 for entry in database.external_tables() {
1683 let provider = external_modules.external_table_provider(&database, &entry)?;
1684 ctx.register_table(&entry.name, provider)
1685 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1686 }
1687
1688 let primary = {
1691 let mut names: Vec<&String> = tables.keys().collect();
1692 names.sort();
1693 names.first().and_then(|n| tables.get(*n).cloned())
1694 };
1695
1696 Ok(Self {
1697 ctx,
1698 db: primary,
1699 database: Some(database),
1700 principal,
1701 principal_catalog_bound,
1702 cache: parking_lot::Mutex::new(BoundedLru::new(SESSION_CACHE_MAX_ENTRIES)),
1703 plan_cache: parking_lot::Mutex::new(BoundedLru::new(SESSION_CACHE_MAX_ENTRIES)),
1704 tables: parking_lot::Mutex::new(tables),
1705 views: parking_lot::Mutex::new(HashMap::new()),
1706 attached_databases: parking_lot::Mutex::new(HashMap::new()),
1707 savepoints: parking_lot::Mutex::new(Vec::new()),
1708 sql_txn: parking_lot::Mutex::new(None),
1709 sql_fn_state,
1710 external_modules,
1711 prepared_names: parking_lot::Mutex::new(std::collections::HashSet::new()),
1712 })
1713 }
1714
1715 pub fn principal(&self) -> Option<mongreldb_core::Principal> {
1716 self.principal.as_ref().and_then(|principal| {
1717 if self.principal_catalog_bound {
1718 self.database
1719 .as_ref()
1720 .and_then(|database| database.resolve_principal(&principal.username))
1721 } else {
1722 Some(principal.clone())
1723 }
1724 })
1725 }
1726
1727 fn security_context_active(&self) -> bool {
1728 self.principal.is_some()
1729 || self.database.as_ref().is_some_and(|database| {
1730 database.principal_snapshot().is_some() || {
1731 let security = database.security_catalog();
1732 !security.rls_tables.is_empty()
1733 || !security.policies.is_empty()
1734 || !security.masks.is_empty()
1735 }
1736 })
1737 }
1738
1739 pub fn register_external_module(&self, module: Arc<dyn ExternalTableModule>) -> Result<()> {
1740 self.external_modules.register(module)?;
1741 self.clear_cache();
1742 Ok(())
1743 }
1744
1745 pub fn db(&self) -> Option<&Arc<Mutex<Table>>> {
1749 self.db.as_ref()
1750 }
1751
1752 pub fn create_view(&self, name: &str, sql: &str) {
1757 self.create_view_with_schema(name, sql, CoreSchema::default(), HashMap::new());
1758 }
1759
1760 pub(crate) fn create_view_with_schema(
1761 &self,
1762 name: &str,
1763 sql: &str,
1764 schema: CoreSchema,
1765 input_types: HashMap<u16, Option<TypeId>>,
1766 ) {
1767 self.views.lock().insert(
1768 name.to_string(),
1769 ViewDef {
1770 sql: sql.to_string(),
1771 schema,
1772 input_types,
1773 },
1774 );
1775 }
1776
1777 pub fn drop_view(&self, name: &str) {
1779 self.views.lock().remove(name);
1780 }
1781
1782 pub(crate) fn view_schema(&self, name: &str) -> Option<CoreSchema> {
1783 self.views.lock().get(name).map(|view| view.schema.clone())
1784 }
1785
1786 pub(crate) fn view_definition(&self, name: &str) -> Option<ViewDef> {
1787 self.views.lock().get(name).cloned()
1788 }
1789
1790 pub async fn register(&self, name: &str) -> Result<()> {
1792 let db = self.db.clone().ok_or(MongrelQueryError::Core(
1793 mongreldb_core::MongrelError::NotFound("no primary table".into()),
1794 ))?;
1795 let provider = MongrelProvider::new(db.clone())?;
1796 self.tables.lock().insert(name.to_string(), db);
1797 self.ctx
1798 .register_table(name, Arc::new(provider))
1799 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1800 Ok(())
1801 }
1802
1803 pub async fn register_db(&self, name: &str, db: Table) -> Result<()> {
1809 let db_arc = Arc::new(Mutex::new(db));
1810 let provider = MongrelProvider::new(db_arc.clone())?;
1811 self.tables.lock().insert(name.to_string(), db_arc);
1812 self.ctx
1813 .register_table(name, Arc::new(provider))
1814 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1815 Ok(())
1816 }
1817
1818 fn refresh_registered_table(&self, db: &Arc<Database>, name: &str) -> Result<()> {
1819 self.ctx
1820 .deregister_table(name)
1821 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1822 let handle = db.table(name)?;
1823 let provider = MongrelProvider::new_secured(
1824 handle.clone(),
1825 Arc::clone(db),
1826 name.to_string(),
1827 self.principal.clone(),
1828 )?;
1829 self.ctx
1830 .register_table(name, Arc::new(provider))
1831 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1832 self.tables.lock().insert(name.to_string(), handle);
1833 Ok(())
1834 }
1835
1836 fn try_catalog_introspection(&self, sql: &str) -> Result<Option<Vec<RecordBatch>>> {
1844 let lower = sql.to_ascii_lowercase();
1845 if !lower.contains("information_schema.tables") {
1846 return Ok(None);
1847 }
1848 use arrow::array::{ArrayRef, Int64Array, StringArray};
1849 use arrow::datatypes::{DataType, Field, Schema};
1850 use arrow::record_batch::RecordBatch;
1851
1852 let mut types: Vec<String> = Vec::new();
1853 let mut names: Vec<String> = Vec::new();
1854 let mut tbl_names: Vec<String> = Vec::new();
1855
1856 let principal = self.principal();
1858 let can_see = |name: &str| match &self.database {
1859 Some(database) => database
1860 .select_column_ids_for(name, principal.as_ref())
1861 .is_ok(),
1862 None => true,
1863 };
1864 for name in self.tables.lock().keys().filter(|name| can_see(name)) {
1865 types.push("table".into());
1866 names.push(name.clone());
1867 tbl_names.push(name.clone());
1868 }
1869 for name in self.views.lock().keys() {
1871 types.push("view".into());
1872 names.push(name.clone());
1873 tbl_names.push(name.clone());
1874 }
1875 if let Some(db) = &self.database {
1877 for t in db.triggers() {
1878 let target_name = match &t.target {
1879 mongreldb_core::trigger::TriggerTarget::Table(n)
1880 | mongreldb_core::trigger::TriggerTarget::View(n) => n.clone(),
1881 };
1882 if !can_see(&target_name) {
1883 continue;
1884 }
1885 types.push("trigger".into());
1886 names.push(t.name.clone());
1887 tbl_names.push(target_name);
1888 }
1889 }
1890
1891 let schema = Arc::new(Schema::new(vec![
1892 Field::new("type", DataType::Utf8, false),
1893 Field::new("name", DataType::Utf8, false),
1894 Field::new("tbl_name", DataType::Utf8, false),
1895 Field::new("rootpage", DataType::Int64, false),
1896 Field::new("sql", DataType::Utf8, true),
1897 ]));
1898 let n = names.len();
1899 let rootpages: Vec<i64> = vec![0; n];
1900 let sqls: Vec<Option<&str>> = vec![None; n];
1901 let batch = RecordBatch::try_new(
1902 schema,
1903 vec![
1904 Arc::new(StringArray::from(types)) as ArrayRef,
1905 Arc::new(StringArray::from(names)) as ArrayRef,
1906 Arc::new(StringArray::from(tbl_names)) as ArrayRef,
1907 Arc::new(Int64Array::from(rootpages)) as ArrayRef,
1908 Arc::new(StringArray::from(sqls)) as ArrayRef,
1909 ],
1910 )
1911 .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
1912 Ok(Some(vec![batch]))
1913 }
1914
1915 fn try_direct_dispatch(&self, sql: &str) -> Result<Option<Vec<RecordBatch>>> {
1921 if self.security_context_active() {
1922 return Ok(None);
1923 }
1924
1925 use arrow::array::ArrayRef;
1926 use mongreldb_core::Condition;
1927 use sqlparser::ast::{Expr, Query, SelectItem, SetExpr, Statement, TableFactor};
1928 use sqlparser::dialect::PostgreSqlDialect;
1929 use sqlparser::parser::Parser;
1930
1931 let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
1933 return Ok(None);
1934 };
1935 if stmts.len() != 1 {
1936 return Ok(None);
1937 }
1938 let Statement::Query(query) = stmts.into_iter().next().unwrap() else {
1939 return Ok(None);
1940 };
1941 let Query { body, .. } = *query;
1942 let select = match *body {
1943 SetExpr::Select(s) => *s,
1944 _ => return Ok(None),
1945 };
1946 let lower_sql = sql.to_lowercase();
1949 if lower_sql.contains(" limit ") || lower_sql.contains(" offset ") {
1950 return Ok(None);
1951 }
1952 use sqlparser::ast::GroupByExpr;
1954 if select.distinct.is_some()
1955 || !matches!(&select.group_by, GroupByExpr::Expressions(e, _) if e.is_empty())
1956 || select.having.is_some()
1957 || select.from.len() != 1
1958 || !select.from[0].joins.is_empty()
1959 {
1960 return Ok(None);
1961 }
1962 let table_name = match &select.from[0].relation {
1963 TableFactor::Table { name, .. } => Some(name.to_string()),
1964 _ => return Ok(None),
1965 };
1966 let Some(table_name) = table_name else {
1967 return Ok(None);
1968 };
1969
1970 if select.selection.is_none() {
1976 return Ok(None);
1977 }
1978
1979 let mut proj_names: Option<Vec<String>> = None;
1981 for item in &select.projection {
1982 match item {
1983 SelectItem::Wildcard(_) => {}
1984 SelectItem::UnnamedExpr(Expr::Identifier(ident)) => {
1985 proj_names
1986 .get_or_insert_with(Vec::new)
1987 .push(ident.value.clone());
1988 }
1989 SelectItem::UnnamedExpr(Expr::CompoundIdentifier(idents)) => {
1990 if let Some(last) = idents.last() {
1991 proj_names
1992 .get_or_insert_with(Vec::new)
1993 .push(last.value.clone());
1994 }
1995 }
1996 _ => return Ok(None),
1997 }
1998 }
1999
2000 let handle = match self.tables.lock().get(&table_name).cloned() {
2002 Some(h) => h,
2003 None => return Ok(None),
2004 };
2005
2006 if let Some(db) = &self.database {
2008 db.require_table(
2009 &table_name,
2010 mongreldb_core::auth_state::RequiredPermission::Select,
2011 )?;
2012 }
2013
2014 let mut db = handle.lock();
2015 let schema = db.schema().clone();
2016 let conditions: Vec<Condition> = match &select.selection {
2019 Some(expr) => match translate_sqlparser_filter(expr, &schema) {
2020 Some(c) => c,
2021 None => return Ok(None),
2022 },
2023 None => Vec::new(),
2024 };
2025 if !conditions.is_empty() && db.ensure_indexes_complete().is_err() {
2026 return Ok(None);
2027 }
2028 let snap = db.snapshot();
2029
2030 let mut col_ids: Vec<u16> = Vec::new();
2032 let mut fields: Vec<arrow::datatypes::Field> = Vec::new();
2033 let resolve_col = |name: &str| -> Option<&mongreldb_core::schema::ColumnDef> {
2034 schema.columns.iter().find(|c| c.name == name)
2035 };
2036 match &proj_names {
2037 None => {
2038 for c in &schema.columns {
2039 col_ids.push(c.id);
2040 fields.push(arrow::datatypes::Field::new(
2041 &c.name,
2042 arrow_conv::arrow_data_type(&c.ty)?,
2043 c.flags.contains(mongreldb_core::ColumnFlags::NULLABLE),
2044 ));
2045 }
2046 }
2047 Some(names) => {
2048 for n in names {
2049 let cdef = match resolve_col(n) {
2050 Some(c) => c,
2051 None => return Ok(None), };
2053 col_ids.push(cdef.id);
2054 fields.push(arrow::datatypes::Field::new(
2055 &cdef.name,
2056 arrow_conv::arrow_data_type(&cdef.ty)?,
2057 cdef.flags.contains(mongreldb_core::ColumnFlags::NULLABLE),
2058 ));
2059 }
2060 }
2061 }
2062
2063 let cols = if !conditions.is_empty() {
2065 match db.query_columns_native_cached(&conditions, Some(&col_ids), snap) {
2066 Ok(Some(c)) => c,
2067 Ok(None) => db
2068 .visible_columns_native(snap, Some(&col_ids))
2069 .map_err(MongrelQueryError::Core)?,
2070 Err(_) => return Ok(None),
2071 }
2072 } else {
2073 db.visible_columns_native(snap, Some(&col_ids))
2074 .map_err(MongrelQueryError::Core)?
2075 };
2076 drop(db);
2077
2078 let mut arrays: Vec<ArrayRef> = Vec::with_capacity(col_ids.len());
2080 for cid in &col_ids {
2081 let col = cols
2082 .iter()
2083 .find(|(id, _)| id == cid)
2084 .map(|(_, c)| c.clone());
2085 let Some(col) = col else { return Ok(None) };
2086 let ty = schema
2087 .columns
2088 .iter()
2089 .find(|c| c.id == *cid)
2090 .map(|c| c.ty.clone())
2091 .unwrap_or(mongreldb_core::schema::TypeId::Int64);
2092 arrays.push(arrow_conv::native_to_array(ty, &col)?);
2093 }
2094 let batch_schema = Arc::new(arrow::datatypes::Schema::new(fields));
2095 let batch = RecordBatch::try_new(batch_schema, arrays)
2096 .map_err(|e| MongrelQueryError::Arrow(format!("direct dispatch batch build: {e}")))?;
2097
2098 mongreldb_core::trace::QueryTrace::record(|t| {
2099 t.scan_mode = mongreldb_core::trace::ScanMode::DirectDispatch;
2100 t.planning_nanos = 0; });
2102 Ok(Some(vec![batch]))
2103 }
2104
2105 async fn dataframe(&self, sql: &str) -> Result<datafusion::dataframe::DataFrame> {
2106 let use_plan_cache = !is_prepared_stmt_sql(sql);
2109 if let Some(plan) = use_plan_cache
2110 .then(|| self.plan_cache.lock().get(sql).cloned())
2111 .flatten()
2112 {
2113 return Ok(datafusion::dataframe::DataFrame::new(
2114 self.ctx.state(),
2115 plan,
2116 ));
2117 }
2118
2119 let df = self
2120 .ctx
2121 .sql(sql)
2122 .await
2123 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2124 if use_plan_cache {
2125 self.plan_cache
2126 .lock()
2127 .insert(sql.to_string(), df.logical_plan().clone());
2128 }
2129 Ok(df)
2130 }
2131
2132 fn prepare_as_of_query(&self, sql: &str) -> Result<Option<AsOfQuery>> {
2133 static AS_OF_RE: std::sync::OnceLock<regex::Regex> = std::sync::OnceLock::new();
2134 static NEXT_TEMP: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
2135 let re = AS_OF_RE.get_or_init(|| {
2136 regex::Regex::new(
2137 r#"(?i)\b(FROM|JOIN)\s+("[^"]+"|[A-Za-z_][A-Za-z0-9_]*)\s+AS\s+OF\s+EPOCH\s+([0-9]+)\b"#,
2138 )
2139 .expect("valid AS OF regex")
2140 });
2141 let mut captures = re.captures_iter(sql);
2142 let Some(capture) = captures.next() else {
2143 return Ok(None);
2144 };
2145 if captures.next().is_some() {
2146 return Err(MongrelQueryError::Schema(
2147 "AS OF EPOCH currently supports one historical table per query".into(),
2148 ));
2149 }
2150 let database = self.database.as_ref().ok_or_else(|| {
2151 MongrelQueryError::Schema("AS OF EPOCH requires a Database-backed session".into())
2152 })?;
2153 let whole = capture.get(0).expect("whole AS OF match");
2154 let keyword = capture.get(1).expect("FROM/JOIN capture").as_str();
2155 let table_token = capture.get(2).expect("table capture").as_str();
2156 let table_name = table_token.trim_matches('"');
2157 let epoch = capture
2158 .get(3)
2159 .expect("epoch capture")
2160 .as_str()
2161 .parse::<u64>()
2162 .map_err(|error| MongrelQueryError::Schema(format!("AS OF epoch: {error}")))?;
2163 let temp_id = NEXT_TEMP.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2164 let temp_name = format!("__mongrel_as_of_{temp_id}");
2165
2166 let mut replace_end = whole.end();
2170 let mut alias = table_token.to_string();
2171 let tail = &sql[replace_end..];
2172 let whitespace = tail.len() - tail.trim_start().len();
2173 let trimmed_tail = tail.trim_start();
2174 if trimmed_tail
2175 .get(..3)
2176 .is_some_and(|prefix| prefix.eq_ignore_ascii_case("as "))
2177 {
2178 let alias_text = &trimmed_tail[3..];
2179 let alias_len = alias_text
2180 .find(|character: char| {
2181 character.is_ascii_whitespace() || character == ',' || character == ';'
2182 })
2183 .unwrap_or(alias_text.len());
2184 if alias_len == 0 {
2185 return Err(MongrelQueryError::Schema(
2186 "AS OF EPOCH AS requires an alias".into(),
2187 ));
2188 }
2189 alias = alias_text[..alias_len].to_string();
2190 replace_end += whitespace + 3 + alias_len;
2191 } else {
2192 let alias_len = trimmed_tail
2193 .find(|character: char| {
2194 character.is_ascii_whitespace() || character == ',' || character == ';'
2195 })
2196 .unwrap_or(trimmed_tail.len());
2197 let candidate = &trimmed_tail[..alias_len];
2198 let keyword = candidate.to_ascii_lowercase();
2199 let clause_keyword = matches!(
2200 keyword.as_str(),
2201 "where"
2202 | "join"
2203 | "left"
2204 | "right"
2205 | "inner"
2206 | "outer"
2207 | "cross"
2208 | "full"
2209 | "on"
2210 | "group"
2211 | "order"
2212 | "having"
2213 | "limit"
2214 | "offset"
2215 | "union"
2216 | "intersect"
2217 | "except"
2218 );
2219 let valid_alias = candidate
2220 .as_bytes()
2221 .first()
2222 .is_some_and(|byte| byte.is_ascii_alphabetic() || *byte == b'_' || *byte == b'"');
2223 if alias_len > 0 && valid_alias && !clause_keyword {
2224 alias = candidate.to_string();
2225 replace_end += whitespace + alias_len;
2226 }
2227 }
2228 let replacement = format!("{keyword} {temp_name} AS {alias}");
2229 let mut rewritten = String::with_capacity(sql.len() + replacement.len());
2230 rewritten.push_str(&sql[..whole.start()]);
2231 rewritten.push_str(&replacement);
2232 rewritten.push_str(&sql[replace_end..]);
2233
2234 let (snapshot, retention) = database.snapshot_at_owned(mongreldb_core::Epoch(epoch))?;
2235 let handle = database.table(table_name)?;
2236 let provider = MongrelProvider::new_historical(
2237 handle,
2238 snapshot,
2239 retention,
2240 Some(ProviderSecurity {
2241 database: Arc::clone(database),
2242 table: table_name.to_string(),
2243 principal: self.principal.clone(),
2244 principal_catalog_bound: self.principal_catalog_bound,
2245 }),
2246 )?;
2247 self.ctx
2248 .register_table(&temp_name, Arc::new(provider))
2249 .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))?;
2250 Ok(Some(AsOfQuery {
2251 sql: rewritten,
2252 registration: AsOfRegistration {
2253 ctx: self.ctx.clone(),
2254 table_name: temp_name,
2255 },
2256 }))
2257 }
2258
2259 async fn run_as_of(&self, query: AsOfQuery) -> Result<Vec<RecordBatch>> {
2260 let AsOfQuery { sql, registration } = query;
2261 let _registration = registration;
2262 let plan_start = std::time::Instant::now();
2263 let dataframe = self
2264 .ctx
2265 .sql(&sql)
2266 .await
2267 .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))?;
2268 mongreldb_core::trace::QueryTrace::record(|trace| {
2269 trace.planning_nanos = plan_start.elapsed().as_nanos() as u64;
2270 });
2271 dataframe
2272 .collect()
2273 .await
2274 .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))
2275 }
2276
2277 async fn run_as_of_stream(&self, query: AsOfQuery) -> Result<MongrelRecordBatchStream> {
2278 use futures::StreamExt;
2279
2280 let AsOfQuery { sql, registration } = query;
2281 let dataframe = self
2282 .ctx
2283 .sql(&sql)
2284 .await
2285 .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))?;
2286 let stream = dataframe
2287 .execute_stream()
2288 .await
2289 .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))?;
2290 let schema = stream.schema();
2291 let guarded = futures::stream::unfold(
2292 (stream, registration),
2293 |(mut stream, registration)| async move {
2294 stream
2295 .next()
2296 .await
2297 .map(|item| (item, (stream, registration)))
2298 },
2299 );
2300 Ok(Box::pin(RecordBatchStreamAdapter::new(schema, guarded)))
2301 }
2302
2303 pub async fn run_stream(&self, sql: &str) -> Result<MongrelRecordBatchStream> {
2308 if strip_explain_query_plan(sql).is_some() {
2309 return self.run(sql).await.map(buffered_stream);
2310 }
2311
2312 let trimmed = sql.trim();
2313 if trimmed.contains(';') && !is_trigger_body(trimmed) {
2314 let stmts = split_sql_statements(trimmed);
2315 let non_empty: Vec<&str> = stmts
2316 .iter()
2317 .map(String::as_str)
2318 .filter(|stmt| !stmt.trim().is_empty() && stmt.trim() != ";")
2319 .collect();
2320 if non_empty.is_empty() {
2321 return Ok(buffered_stream(Vec::new()));
2322 }
2323 if stmts.len() > 1 {
2324 for stmt in &non_empty[..non_empty.len() - 1] {
2325 self.run(stmt).await?;
2326 }
2327 return Box::pin(self.run_stream(non_empty[non_empty.len() - 1])).await;
2328 }
2329 }
2330
2331 if let Some(batches) = commands::try_run_command(self, sql).await? {
2332 return Ok(buffered_stream(batches));
2333 }
2334
2335 if let Some(query) = self.prepare_as_of_query(sql)? {
2336 return self.run_as_of_stream(query).await;
2337 }
2338
2339 let resolved = self.resolve_view_sql(sql);
2340 let resolved = self.rewrite_external_module_compat_sql(&resolved);
2341 let resolved = rewrite_compat_function_calls(&resolved);
2342 let effective_sql = normalize_sql(&resolved);
2343 let sql = effective_sql.as_str();
2344
2345 if let Some(batches) = self.try_catalog_introspection(sql)? {
2346 return Ok(buffered_stream(batches));
2347 }
2348
2349 let plan_start = std::time::Instant::now();
2350 let external_module_scan = self.query_references_external_module(sql);
2351 let df = self.dataframe(sql).await?;
2352 self.track_prepared_name(sql);
2353 mongreldb_core::trace::QueryTrace::record(|trace| {
2354 trace.planning_nanos = plan_start.elapsed().as_nanos() as u64;
2355 });
2356 let stream = df
2357 .execute_stream()
2358 .await
2359 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2360 if external_module_scan {
2361 mongreldb_core::trace::QueryTrace::record(|trace| {
2362 trace.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
2363 });
2364 }
2365 Ok(stream)
2366 }
2367
2368 pub async fn run(&self, sql: &str) -> Result<Vec<RecordBatch>> {
2373 if let Some(inner) = strip_explain_query_plan(sql) {
2374 return self.explain_query_plan(inner).await;
2375 }
2376
2377 let trimmed = sql.trim();
2383 if trimmed.contains(';') && !is_trigger_body(trimmed) {
2384 let stmts = split_sql_statements(trimmed);
2385 let non_empty: Vec<&String> = stmts
2387 .iter()
2388 .filter(|s| !s.trim().is_empty() && s.trim() != ";")
2389 .collect();
2390 if non_empty.is_empty() {
2391 return Ok(Vec::new());
2392 }
2393 if stmts.len() > 1 {
2394 let mut last = Vec::new();
2395 for stmt in &stmts {
2396 let stmt = stmt.trim();
2397 if stmt.is_empty() {
2398 continue;
2399 }
2400 last = Box::pin(self.run(stmt)).await?;
2401 }
2402 return Ok(last);
2403 }
2404 }
2405
2406 if let Some(batches) = commands::try_run_command(self, sql).await? {
2408 return Ok(batches);
2409 }
2410
2411 let trimmed = sql.trim();
2417 if trimmed.contains(';') {
2418 let stmts = split_sql_statements(trimmed);
2419 if stmts.len() > 1 {
2420 let mut last = Vec::new();
2421 for stmt in &stmts {
2422 let stmt = stmt.trim();
2423 if stmt.is_empty() {
2424 continue;
2425 }
2426 last = Box::pin(self.run(stmt)).await?;
2427 }
2428 return Ok(last);
2429 }
2430 }
2431 let lower = sql.trim_start().to_lowercase();
2433 if lower.starts_with("create table") {
2434 if let Some(db) = &self.database {
2435 db.require_for(self.principal().as_ref(), &mongreldb_core::Permission::Ddl)?;
2436 let (name, schema) = parse_create_table(sql)?;
2437 db.create_table(&name, schema)?;
2438 let handle = db.table(&name)?;
2439 let provider = MongrelProvider::new_secured(
2440 handle.clone(),
2441 Arc::clone(db),
2442 name.clone(),
2443 self.principal.clone(),
2444 )?;
2445 self.ctx
2446 .register_table(&name, Arc::new(provider))
2447 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2448 self.tables.lock().insert(name, handle);
2449 self.invalidate_prepared_statements().await;
2450 self.clear_cache();
2451 return Ok(Vec::new());
2452 }
2453 }
2454 if lower.starts_with("drop table") {
2455 if let Some(db) = &self.database {
2456 db.require_for(self.principal().as_ref(), &mongreldb_core::Permission::Ddl)?;
2457 let (name, if_exists) = parse_drop_table(sql)?;
2458 let drop_result = db.drop_table(&name);
2459 if let Err(e) = drop_result {
2460 let is_not_found = matches!(e, mongreldb_core::MongrelError::NotFound(_));
2462 if !(if_exists && is_not_found) {
2463 return Err(e.into());
2464 }
2465 } else {
2466 self.ctx
2467 .deregister_table(&name)
2468 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2469 self.tables.lock().remove(&name);
2470 }
2471 self.invalidate_prepared_statements().await;
2472 self.clear_cache();
2473 return Ok(Vec::new());
2474 }
2475 }
2476 if lower.starts_with("alter table") {
2477 if let Some(db) = &self.database {
2478 db.require_for(self.principal().as_ref(), &mongreldb_core::Permission::Ddl)?;
2479 match parse_alter_table(sql)? {
2480 ParsedAlterTable::RenameTable { old_name, new_name } => {
2481 db.rename_table(&old_name, &new_name)?;
2482 self.ctx
2486 .deregister_table(&old_name)
2487 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2488 self.tables.lock().remove(&old_name);
2489 let handle = db.table(&new_name)?;
2490 let provider = MongrelProvider::new_secured(
2491 handle.clone(),
2492 Arc::clone(db),
2493 new_name.clone(),
2494 self.principal.clone(),
2495 )?;
2496 self.ctx
2497 .register_table(&new_name, Arc::new(provider))
2498 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2499 self.tables.lock().insert(new_name, handle);
2500 }
2501 ParsedAlterTable::RenameColumn {
2502 table_name,
2503 column_name,
2504 new_name,
2505 } => {
2506 db.alter_column(&table_name, &column_name, AlterColumn::rename(new_name))?;
2507 self.refresh_registered_table(db, &table_name)?;
2508 }
2509 ParsedAlterTable::AlterColumnType {
2510 table_name,
2511 column_name,
2512 ty,
2513 } => {
2514 db.alter_column(&table_name, &column_name, AlterColumn::set_type(ty))?;
2515 self.refresh_registered_table(db, &table_name)?;
2516 }
2517 ParsedAlterTable::SetNotNull {
2518 table_name,
2519 column_name,
2520 } => {
2521 let flags = current_column_flags(db, &table_name, &column_name)?
2522 .without(ColumnFlags::NULLABLE);
2523 db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
2524 self.refresh_registered_table(db, &table_name)?;
2525 }
2526 ParsedAlterTable::DropNotNull {
2527 table_name,
2528 column_name,
2529 } => {
2530 let flags = current_column_flags(db, &table_name, &column_name)?
2531 .with(ColumnFlags::NULLABLE);
2532 db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
2533 self.refresh_registered_table(db, &table_name)?;
2534 }
2535 }
2536 self.invalidate_prepared_statements().await;
2537 self.clear_cache();
2538 return Ok(Vec::new());
2539 }
2540 }
2541
2542 if let Some(query) = self.prepare_as_of_query(sql)? {
2543 return self.run_as_of(query).await;
2544 }
2545
2546 let resolved = self.resolve_view_sql(sql);
2549 let resolved = self.rewrite_external_module_compat_sql(&resolved);
2550 let resolved = rewrite_compat_function_calls(&resolved);
2551 let effective_sql = normalize_sql(&resolved);
2555 let sql = effective_sql.as_str();
2556 let epoch = self.cache_epoch();
2560 let key = (sql.to_string(), epoch);
2561 let has_ttl_table = self
2562 .tables
2563 .lock()
2564 .values()
2565 .any(|table| table.lock().ttl().is_some());
2566 let result_cacheable = !extended_sql_functions::contains_volatile_extended_function(sql)
2567 && !is_explain_analyze(sql)
2568 && !is_prepared_stmt_sql(sql)
2569 && !has_ttl_table;
2570 if result_cacheable {
2571 if let Some(hit) = self.cache.lock().get(&key) {
2572 return Ok((**hit).clone());
2573 }
2574 }
2575 if let Some(batches) = self.try_catalog_introspection(sql)? {
2578 if result_cacheable {
2579 self.cache.lock().insert(key, Arc::new(batches.clone()));
2580 }
2581 return Ok(batches);
2582 }
2583 if let Some(batches) = self.try_direct_dispatch(sql)? {
2588 if result_cacheable {
2589 self.cache.lock().insert(key, Arc::new(batches.clone()));
2590 }
2591 return Ok(batches);
2592 }
2593 let plan_start = std::time::Instant::now();
2595 let external_module_scan = self.query_references_external_module(sql);
2596 let df = self.dataframe(sql).await?;
2597 self.track_prepared_name(sql);
2600 let planning_nanos = plan_start.elapsed().as_nanos() as u64;
2603 mongreldb_core::trace::QueryTrace::record(|t| t.planning_nanos = planning_nanos);
2604
2605 let agg_key = sql_cache_key(sql);
2610 let batches = match self.try_native_aggregate(df.logical_plan(), agg_key) {
2611 Ok(Some(batch)) => vec![batch],
2612 _ => {
2613 match self.try_fk_join(df.logical_plan()) {
2617 Ok(Some(b)) => {
2618 mongreldb_core::trace::QueryTrace::record(|t| {
2620 t.join_mode = mongreldb_core::trace::JoinMode::FkBitmap;
2621 });
2622 b
2623 }
2624 _ => df
2625 .collect()
2626 .await
2627 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?,
2628 }
2629 }
2630 };
2631 if external_module_scan {
2632 mongreldb_core::trace::QueryTrace::record(|t| {
2633 t.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
2634 });
2635 }
2636 if result_cacheable {
2637 self.cache.lock().insert(key, Arc::new(batches.clone()));
2638 }
2639 Ok(batches)
2640 }
2641
2642 pub async fn run_sql_traced(
2654 &self,
2655 sql: &str,
2656 ) -> Result<(Vec<RecordBatch>, mongreldb_core::trace::QueryTrace)> {
2657 mongreldb_core::trace::QueryTrace::push_scope();
2658 let result = self.run(sql).await;
2659 let trace = mongreldb_core::trace::QueryTrace::pop_scope();
2660 Ok((result?, trace))
2661 }
2662
2663 pub fn clear_cache(&self) {
2666 self.cache.lock().clear();
2667 self.plan_cache.lock().clear();
2668 }
2669
2670 fn track_prepared_name(&self, sql: &str) {
2674 let lower = sql.trim_start().to_ascii_lowercase();
2675 if let Some(rest) = lower.strip_prefix("prepare ") {
2676 if let Some(name) = parse_stmt_ident(rest) {
2677 self.prepared_names.lock().insert(name);
2678 }
2679 } else if let Some(rest) = lower.strip_prefix("deallocate ") {
2680 if let Some(name) = parse_stmt_ident(rest) {
2681 self.prepared_names.lock().remove(&name);
2682 }
2683 }
2684 }
2685
2686 pub async fn invalidate_prepared_statements(&self) {
2693 let names: Vec<String> = self.prepared_names.lock().drain().collect();
2694 for name in names {
2695 let sql = format!("DEALLOCATE {name}");
2699 if self.ctx.sql(&sql).await.is_err() {
2700 }
2702 }
2703 }
2704
2705 async fn explain_query_plan(&self, sql: &str) -> Result<Vec<RecordBatch>> {
2706 let explain_sql = format!("EXPLAIN {}", sql.trim().trim_end_matches(';'));
2707 let batches = self
2708 .ctx
2709 .sql(&explain_sql)
2710 .await
2711 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?
2712 .collect()
2713 .await
2714 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2715 let mut detail = self.mongrel_query_plan_details(sql);
2716 for batch in &batches {
2717 if batch.num_columns() < 2 {
2718 continue;
2719 }
2720 let Some(plan_type) = batch.column(0).as_any().downcast_ref::<StringArray>() else {
2721 continue;
2722 };
2723 let Some(plan) = batch.column(1).as_any().downcast_ref::<StringArray>() else {
2724 continue;
2725 };
2726 for row in 0..batch.num_rows() {
2727 let prefix = plan_type.value(row);
2728 for line in plan.value(row).lines() {
2729 let line = line.trim();
2730 if !line.is_empty() {
2731 detail.push(format!("DATAFUSION {prefix}: {line}"));
2732 }
2733 }
2734 }
2735 }
2736 if detail.is_empty() {
2737 detail.push("plan unavailable".to_string());
2738 }
2739 let ids = (0..detail.len()).map(|i| i as i64).collect::<Vec<_>>();
2740 let parents = vec![0_i64; detail.len()];
2741 let notused = vec![0_i64; detail.len()];
2742 let schema = Arc::new(arrow::datatypes::Schema::new(vec![
2743 arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int64, false),
2744 arrow::datatypes::Field::new("parent", arrow::datatypes::DataType::Int64, false),
2745 arrow::datatypes::Field::new("notused", arrow::datatypes::DataType::Int64, false),
2746 arrow::datatypes::Field::new("detail", arrow::datatypes::DataType::Utf8, false),
2747 ]));
2748 let batch = RecordBatch::try_new(
2749 schema,
2750 vec![
2751 Arc::new(Int64Array::from(ids)) as ArrayRef,
2752 Arc::new(Int64Array::from(parents)),
2753 Arc::new(Int64Array::from(notused)),
2754 Arc::new(StringArray::from(detail)),
2755 ],
2756 )
2757 .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
2758 Ok(vec![batch])
2759 }
2760
2761 fn mongrel_query_plan_details(&self, sql: &str) -> Vec<String> {
2762 use sqlparser::ast::{GroupByExpr, OrderByKind, SetExpr, Statement};
2763 use sqlparser::dialect::PostgreSqlDialect;
2764 use sqlparser::parser::Parser;
2765
2766 let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
2767 return Vec::new();
2768 };
2769 let Some(Statement::Query(query)) = stmts.first() else {
2770 return Vec::new();
2771 };
2772
2773 fn collect(session: &MongrelSession, query: &sqlparser::ast::Query, out: &mut Vec<String>) {
2774 use sqlparser::ast::{SetOperator, TableWithJoins};
2775 match query.body.as_ref() {
2776 SetExpr::Select(select) => {
2777 for TableWithJoins { relation, joins } in &select.from {
2778 session.push_table_plan(relation, select.selection.as_ref(), out);
2779 for join in joins {
2780 session.push_table_plan(&join.relation, None, out);
2781 }
2782 }
2783 if select.distinct.is_some() {
2784 out.push("USE TEMP B-TREE FOR DISTINCT".to_string());
2785 }
2786 let grouped = match &select.group_by {
2787 GroupByExpr::All(_) => true,
2788 GroupByExpr::Expressions(exprs, _) => !exprs.is_empty(),
2789 };
2790 if grouped {
2791 out.push("USE TEMP B-TREE FOR GROUP BY".to_string());
2792 }
2793 let ordered = query.order_by.as_ref().is_some_and(|order_by| {
2794 matches!(order_by.kind, OrderByKind::All(_))
2795 || matches!(&order_by.kind, OrderByKind::Expressions(exprs) if !exprs.is_empty())
2796 });
2797 if ordered {
2798 out.push("USE TEMP B-TREE FOR ORDER BY".to_string());
2799 }
2800 }
2801 SetExpr::Query(query) => collect(session, query, out),
2802 SetExpr::SetOperation {
2803 left, op, right, ..
2804 } => {
2805 let label = match op {
2806 SetOperator::Union => "COMPOUND QUERY UNION",
2807 SetOperator::Except => "COMPOUND QUERY EXCEPT",
2808 SetOperator::Intersect => "COMPOUND QUERY INTERSECT",
2809 _ => "COMPOUND QUERY",
2810 };
2811 out.push(label.to_string());
2812 collect_set_expr(session, left, out);
2813 collect_set_expr(session, right, out);
2814 }
2815 _ => {}
2816 }
2817 }
2818
2819 fn collect_set_expr(session: &MongrelSession, expr: &SetExpr, out: &mut Vec<String>) {
2820 match expr {
2821 SetExpr::Select(select) => {
2822 for table in &select.from {
2823 session.push_table_plan(&table.relation, select.selection.as_ref(), out);
2824 }
2825 }
2826 SetExpr::Query(query) => collect(session, query, out),
2827 SetExpr::SetOperation { left, right, .. } => {
2828 collect_set_expr(session, left, out);
2829 collect_set_expr(session, right, out);
2830 }
2831 _ => {}
2832 }
2833 }
2834
2835 let mut out = Vec::new();
2836 collect(self, query, &mut out);
2837 out
2838 }
2839
2840 fn push_table_plan(
2841 &self,
2842 relation: &sqlparser::ast::TableFactor,
2843 selection: Option<&sqlparser::ast::Expr>,
2844 out: &mut Vec<String>,
2845 ) {
2846 let sqlparser::ast::TableFactor::Table { name, alias, .. } = relation else {
2847 out.push("SCAN SUBQUERY".to_string());
2848 return;
2849 };
2850 let table_name = name.to_string();
2851 let display_name = alias
2852 .as_ref()
2853 .map(|alias| alias.name.value.clone())
2854 .unwrap_or_else(|| table_name.clone());
2855 let Some(handle) = self.tables.lock().get(&table_name).cloned() else {
2856 out.push(format!("SCAN {display_name}"));
2857 return;
2858 };
2859 let schema = handle.lock().schema().clone();
2860 let searchable = selection
2861 .and_then(|expr| translate_sqlparser_filter(expr, &schema))
2862 .is_some_and(|conditions| !conditions.is_empty());
2863 if searchable {
2864 out.push(format!("SEARCH {display_name} USING MONGREL INDEX"));
2865 } else {
2866 out.push(format!("SCAN {display_name}"));
2867 }
2868 }
2869
2870 fn resolve_view_sql(&self, sql: &str) -> String {
2875 let views = self.views.lock();
2876 if views.is_empty() {
2877 return sql.to_string();
2878 }
2879 let mut result = sql.to_string();
2880 for (name, view) in views.iter() {
2881 result = replace_from_view(&result, name, &view.sql);
2882 }
2883 result
2884 }
2885
2886 fn rewrite_external_module_compat_sql(&self, sql: &str) -> String {
2887 let Some(db) = &self.database else {
2888 return sql.to_string();
2889 };
2890 rewrite_fts_match_compat_sql(sql, db)
2891 }
2892
2893 fn query_references_external_module(&self, sql: &str) -> bool {
2894 use sqlparser::ast::Statement;
2895 use sqlparser::dialect::PostgreSqlDialect;
2896 use sqlparser::parser::Parser;
2897
2898 Parser::parse_sql(&PostgreSqlDialect {}, sql)
2899 .ok()
2900 .is_some_and(|statements| {
2901 statements.iter().any(|statement| match statement {
2902 Statement::Query(query) => self.query_uses_external_module(query),
2903 _ => false,
2904 })
2905 })
2906 }
2907
2908 fn query_uses_external_module(&self, query: &sqlparser::ast::Query) -> bool {
2909 self.set_expr_uses_external_module(query.body.as_ref())
2910 }
2911
2912 fn set_expr_uses_external_module(&self, expr: &sqlparser::ast::SetExpr) -> bool {
2913 use sqlparser::ast::SetExpr;
2914
2915 match expr {
2916 SetExpr::Select(select) => select
2917 .from
2918 .iter()
2919 .any(|table| self.table_with_joins_uses_external_module(table)),
2920 SetExpr::Query(query) => self.query_uses_external_module(query),
2921 SetExpr::SetOperation { left, right, .. } => {
2922 self.set_expr_uses_external_module(left)
2923 || self.set_expr_uses_external_module(right)
2924 }
2925 _ => false,
2926 }
2927 }
2928
2929 fn table_with_joins_uses_external_module(
2930 &self,
2931 table: &sqlparser::ast::TableWithJoins,
2932 ) -> bool {
2933 self.table_factor_uses_external_module(&table.relation)
2934 || table
2935 .joins
2936 .iter()
2937 .any(|join| self.table_factor_uses_external_module(&join.relation))
2938 }
2939
2940 fn table_factor_uses_external_module(&self, relation: &sqlparser::ast::TableFactor) -> bool {
2941 use sqlparser::ast::{Expr, TableFactor};
2942
2943 match relation {
2944 TableFactor::Table { name, args, .. } => {
2945 let table_name = name.to_string();
2946 self.database
2947 .as_ref()
2948 .is_some_and(|db| db.external_table(&table_name).is_some())
2949 || (args.is_some() && self.external_modules.contains(&table_name))
2950 }
2951 TableFactor::Function { name, .. } => self.external_modules.contains(&name.to_string()),
2952 TableFactor::TableFunction {
2953 expr: Expr::Function(func),
2954 ..
2955 } => self.external_modules.contains(&func.name.to_string()),
2956 TableFactor::Derived { subquery, .. } => self.query_uses_external_module(subquery),
2957 _ => false,
2958 }
2959 }
2960
2961 fn cache_epoch(&self) -> u64 {
2964 if let Some(db) = &self.database {
2965 db.visible_epoch().0
2966 } else {
2967 self.combined_epoch()
2968 }
2969 }
2970
2971 fn combined_epoch(&self) -> u64 {
2972 let primary = self.db.as_ref().expect("no primary table");
2973 let mut combined = primary.lock().snapshot().epoch.0;
2974 let tables = self.tables.lock();
2975 for arc in tables.values() {
2976 if !Arc::ptr_eq(arc, primary) {
2977 let e = arc.lock().snapshot().epoch.0;
2978 combined = combined.wrapping_mul(31).wrapping_add(e);
2979 }
2980 }
2981 combined
2982 }
2983
2984 fn try_native_aggregate(
2988 &self,
2989 plan: &datafusion::logical_expr::LogicalPlan,
2990 cache_key: u64,
2991 ) -> Result<Option<RecordBatch>> {
2992 if self.security_context_active() {
2993 return Ok(None);
2994 }
2995 let Some(primary) = self.db.as_ref() else {
2996 return Ok(None);
2997 };
2998 let mut db = primary.lock();
2999 let schema = db.schema().clone();
3000 let snap = db.snapshot();
3001 native_agg::try_native_aggregate(&mut db, &schema, snap, plan, cache_key)
3002 }
3003
3004 fn try_fk_join(
3008 &self,
3009 plan: &datafusion::logical_expr::LogicalPlan,
3010 ) -> Result<Option<Vec<RecordBatch>>> {
3011 if self.security_context_active() {
3012 return Ok(None);
3013 }
3014 let tables = self.tables.lock();
3015 fk_join::try_fk_join(&tables, plan)
3016 }
3017
3018 pub fn context(&self) -> &SessionContext {
3019 &self.ctx
3020 }
3021
3022 pub fn register_scalar_udf(&self, f: ScalarUDF) {
3028 self.ctx.register_udf(f);
3029 self.clear_cache();
3030 }
3031
3032 pub fn register_aggregate_udf(&self, f: AggregateUDF) {
3034 self.ctx.register_udaf(f);
3035 self.clear_cache();
3036 }
3037
3038 pub fn register_window_udf(&self, f: WindowUDF) {
3040 self.ctx.register_udwf(f);
3041 self.clear_cache();
3042 }
3043}
3044
3045fn register_mongrel_functions(
3046 ctx: &SessionContext,
3047 sql_fn_state: Arc<extended_sql_functions::ExtendedSqlState>,
3048) {
3049 ctx.register_udf(ScalarUDF::from(udf::AnnSearchUdf::new()));
3050 ctx.register_udf(ScalarUDF::from(udf::SparseMatchUdf::new()));
3051 ctx.register_udf(ScalarUDF::from(udf::RTreeIntersectsUdf::new()));
3052 ctx.register_udf(ScalarUDF::from(udf::FtsRankUdf::new()));
3053 for udaf in percentile::percentile_udafs() {
3054 ctx.register_udaf(udaf);
3055 }
3056 extended_sql_functions::register_extended_sql_functions_with_state(ctx, sql_fn_state);
3057}
3058
3059fn is_trigger_body(sql: &str) -> bool {
3062 let lower = sql.to_ascii_lowercase();
3063 if !lower.contains("create trigger") && !lower.contains("create or replace trigger") {
3064 return false;
3065 }
3066 lower.contains("begin") && lower.contains("end")
3067}
3068
3069fn split_sql_statements(sql: &str) -> Vec<String> {
3074 let b = sql.as_bytes();
3075 let n = b.len();
3076 let mut stmts = Vec::new();
3077 let mut current = String::new();
3078 let mut i = 0;
3079 while i < n {
3080 let c = b[i];
3081 if c == b'-' && i + 1 < n && b[i + 1] == b'-' {
3083 while i < n && b[i] != b'\n' {
3084 current.push(c as char);
3085 i += 1;
3086 }
3087 continue;
3088 }
3089 if c == b'/' && i + 1 < n && b[i + 1] == b'*' {
3091 current.push_str("/*");
3092 i += 2;
3093 let mut depth = 1;
3094 while i + 1 < n && depth > 0 {
3095 if b[i] == b'/' && b[i + 1] == b'*' {
3096 depth += 1;
3097 current.push_str("/*");
3098 i += 2;
3099 } else if b[i] == b'*' && b[i + 1] == b'/' {
3100 depth -= 1;
3101 current.push_str("*/");
3102 i += 2;
3103 } else {
3104 current.push(b[i] as char);
3105 i += 1;
3106 }
3107 }
3108 continue;
3109 }
3110 if c == b'\'' {
3112 current.push('\'');
3113 i += 1;
3114 while i < n {
3115 current.push(b[i] as char);
3116 if b[i] == b'\'' {
3117 if i + 1 < n && b[i + 1] == b'\'' {
3119 current.push('\'');
3120 i += 2;
3121 continue;
3122 }
3123 i += 1;
3124 break;
3125 }
3126 i += 1;
3127 }
3128 continue;
3129 }
3130 if c == b'"' {
3132 current.push('"');
3133 i += 1;
3134 while i < n {
3135 current.push(b[i] as char);
3136 if b[i] == b'"' {
3137 if i + 1 < n && b[i + 1] == b'"' {
3138 current.push('"');
3139 i += 2;
3140 continue;
3141 }
3142 i += 1;
3143 break;
3144 }
3145 i += 1;
3146 }
3147 continue;
3148 }
3149 if c == b'$' {
3151 let start = i;
3153 let mut tag_end = i + 1;
3154 while tag_end < n && b[tag_end] != b'$' && b[tag_end].is_ascii_alphanumeric() {
3155 tag_end += 1;
3156 }
3157 if tag_end < n && b[tag_end] == b'$' {
3158 let tag = &sql[start..=tag_end];
3159 current.push_str(tag);
3160 i = tag_end + 1;
3161 while i < n {
3163 if b[i] == b'$' && sql[i..].starts_with(tag) {
3164 current.push_str(tag);
3165 i += tag.len();
3166 break;
3167 }
3168 current.push(b[i] as char);
3169 i += 1;
3170 }
3171 continue;
3172 }
3173 }
3174 if c == b';' {
3176 current.push(';');
3177 let trimmed = current.trim();
3178 if !trimmed.is_empty() && trimmed != ";" {
3179 stmts.push(current.clone());
3180 }
3181 current.clear();
3182 i += 1;
3183 continue;
3184 }
3185 current.push(c as char);
3186 i += 1;
3187 }
3188 let trimmed = current.trim();
3190 if !trimmed.is_empty() {
3191 stmts.push(current);
3192 }
3193 stmts
3194}
3195
3196fn strip_explain_query_plan(sql: &str) -> Option<&str> {
3197 let trimmed = sql.trim_start();
3198 let lower = trimmed.to_ascii_lowercase();
3199 if !lower.starts_with("explain") {
3200 return None;
3201 }
3202 let after_explain = trimmed.get(7..)?.trim_start();
3203 let after_explain_lower = after_explain.to_ascii_lowercase();
3204 if !after_explain_lower.starts_with("query") {
3205 return None;
3206 }
3207 let after_query = after_explain.get(5..)?.trim_start();
3208 let after_query_lower = after_query.to_ascii_lowercase();
3209 if !after_query_lower.starts_with("plan") {
3210 return None;
3211 }
3212 Some(after_query.get(4..)?.trim_start())
3213}
3214
3215fn is_prepared_stmt_sql(sql: &str) -> bool {
3221 let lower = sql.trim_start().to_ascii_lowercase();
3222 lower.starts_with("prepare ")
3223 || lower.starts_with("execute ")
3224 || lower.starts_with("deallocate ")
3225}
3226
3227fn parse_stmt_ident(s: &str) -> Option<String> {
3231 let tok = s
3232 .trim_start()
3233 .split(|c: char| c.is_whitespace() || c == '(')
3234 .next()?;
3235 let t = tok.trim_matches(|c| c == '"' || c == '`' || c == '\'');
3236 (!t.is_empty()).then(|| t.to_string())
3237}
3238
3239fn is_explain_analyze(sql: &str) -> bool {
3244 let trimmed = sql.trim_start();
3245 let lower = trimmed.to_ascii_lowercase();
3246 if !lower.starts_with("explain") {
3247 return false;
3248 }
3249 let after = trimmed[7..].trim_start();
3250 after.to_ascii_lowercase().starts_with("analyze")
3251}
3252
3253#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3254enum SqlCompatTokenKind {
3255 Ident,
3256 String,
3257 Dot,
3258 LParen,
3259 RParen,
3260 Comma,
3261}
3262
3263#[derive(Debug, Clone)]
3264struct SqlCompatToken {
3265 kind: SqlCompatTokenKind,
3266 raw: String,
3267 normalized: String,
3268 start: usize,
3269 end: usize,
3270}
3271
3272#[derive(Debug, Clone)]
3273struct FtsMatchBinding {
3274 query_ref: String,
3275}
3276
3277#[derive(Debug, Clone)]
3278struct SqlReplacement {
3279 start: usize,
3280 end: usize,
3281 replacement: String,
3282}
3283
3284fn rewrite_fts_match_compat_sql(sql: &str, db: &Database) -> String {
3285 let tokens = sql_compat_tokens(sql);
3286 if tokens.is_empty() {
3287 return sql.to_string();
3288 }
3289 let bindings = fts_match_bindings(sql, db, &tokens);
3290 if bindings.is_empty() {
3291 return sql.to_string();
3292 }
3293 let unique_refs = bindings
3294 .values()
3295 .map(|binding| binding.query_ref.as_str())
3296 .collect::<HashSet<_>>();
3297 let unique_binding = if unique_refs.len() == 1 {
3298 bindings.values().next().cloned()
3299 } else {
3300 None
3301 };
3302 let mut replacements = Vec::new();
3303 for (idx, token) in tokens.iter().enumerate() {
3304 if token.kind != SqlCompatTokenKind::Ident || token.normalized != "match" {
3305 continue;
3306 }
3307 let Some(rhs) = tokens.get(idx + 1) else {
3308 continue;
3309 };
3310 if rhs.kind != SqlCompatTokenKind::String {
3311 continue;
3312 }
3313 let Some((lhs_start, _lhs_end, query_ref)) =
3314 fts_match_lhs_query_ref(&tokens, idx, &bindings, unique_binding.as_ref())
3315 else {
3316 continue;
3317 };
3318 replacements.push(SqlReplacement {
3319 start: lhs_start,
3320 end: rhs.end,
3321 replacement: format!("{query_ref}.query = {}", rhs.raw),
3322 });
3323 }
3324 apply_sql_replacements(sql, &replacements)
3325}
3326
3327fn fts_match_lhs_query_ref(
3328 tokens: &[SqlCompatToken],
3329 match_idx: usize,
3330 bindings: &HashMap<String, FtsMatchBinding>,
3331 unique_binding: Option<&FtsMatchBinding>,
3332) -> Option<(usize, usize, String)> {
3333 if match_idx == 0 {
3334 return None;
3335 }
3336 let lhs = tokens.get(match_idx - 1)?;
3337 if lhs.kind != SqlCompatTokenKind::Ident {
3338 return None;
3339 }
3340
3341 if match_idx >= 3
3342 && tokens.get(match_idx - 2)?.kind == SqlCompatTokenKind::Dot
3343 && tokens.get(match_idx - 3)?.kind == SqlCompatTokenKind::Ident
3344 {
3345 let owner = tokens.get(match_idx - 3)?;
3346 let binding = bindings.get(&owner.normalized)?;
3347 if lhs.normalized == "query" || lhs.normalized == "text" {
3348 return Some((owner.start, lhs.end, binding.query_ref.clone()));
3349 }
3350 return None;
3351 }
3352
3353 if let Some(binding) = bindings.get(&lhs.normalized) {
3354 return Some((lhs.start, lhs.end, binding.query_ref.clone()));
3355 }
3356 if lhs.normalized == "text" {
3357 let binding = unique_binding?;
3358 return Some((lhs.start, lhs.end, binding.query_ref.clone()));
3359 }
3360 None
3361}
3362
3363fn fts_match_bindings(
3364 sql: &str,
3365 db: &Database,
3366 tokens: &[SqlCompatToken],
3367) -> HashMap<String, FtsMatchBinding> {
3368 let mut out = HashMap::new();
3369 let mut i = 0;
3370 while i < tokens.len() {
3371 let token = &tokens[i];
3372 let starts_table_ref = token.kind == SqlCompatTokenKind::Ident
3373 && matches!(token.normalized.as_str(), "from" | "join");
3374 if !starts_table_ref {
3375 i += 1;
3376 continue;
3377 }
3378 let mut table_idx = i + 1;
3379 if tokens
3380 .get(table_idx)
3381 .is_some_and(|token| token.kind == SqlCompatTokenKind::LParen)
3382 {
3383 i += 1;
3384 continue;
3385 }
3386 let Some(table) = tokens.get(table_idx) else {
3387 break;
3388 };
3389 if table.kind != SqlCompatTokenKind::Ident {
3390 i += 1;
3391 continue;
3392 }
3393 let mut table_name = table.normalized.clone();
3394 let mut table_ref = table.raw.clone();
3395 if tokens
3396 .get(table_idx + 1)
3397 .is_some_and(|token| token.kind == SqlCompatTokenKind::Dot)
3398 && tokens
3399 .get(table_idx + 2)
3400 .is_some_and(|token| token.kind == SqlCompatTokenKind::Ident)
3401 {
3402 let qualified = tokens.get(table_idx + 2).unwrap();
3403 table_name = qualified.normalized.clone();
3404 table_ref = sql[table.start..qualified.end].to_string();
3405 table_idx += 2;
3406 }
3407 if !is_fts_docs_table(db, &table_name) {
3408 i = table_idx + 1;
3409 continue;
3410 }
3411 let mut query_ref = table_ref.clone();
3412 let mut alias_key = None;
3413 let mut next = table_idx + 1;
3414 if tokens.get(next).is_some_and(|token| {
3415 token.kind == SqlCompatTokenKind::Ident && token.normalized == "as"
3416 }) {
3417 next += 1;
3418 }
3419 if let Some(alias) = tokens.get(next) {
3420 if alias.kind == SqlCompatTokenKind::Ident && !is_table_ref_boundary(&alias.normalized)
3421 {
3422 alias_key = Some(alias.normalized.clone());
3423 query_ref = alias.raw.clone();
3424 next += 1;
3425 }
3426 }
3427 out.insert(
3428 table_name,
3429 FtsMatchBinding {
3430 query_ref: query_ref.clone(),
3431 },
3432 );
3433 if let Some(alias_key) = alias_key {
3434 out.insert(alias_key, FtsMatchBinding { query_ref });
3435 }
3436 i = next;
3437 }
3438 out
3439}
3440
3441fn is_fts_docs_table(db: &Database, name: &str) -> bool {
3442 db.external_table(name)
3443 .is_some_and(|entry| entry.module == "fts_docs")
3444}
3445
3446fn is_table_ref_boundary(normalized: &str) -> bool {
3447 matches!(
3448 normalized,
3449 "where"
3450 | "join"
3451 | "left"
3452 | "right"
3453 | "inner"
3454 | "outer"
3455 | "full"
3456 | "cross"
3457 | "on"
3458 | "using"
3459 | "group"
3460 | "order"
3461 | "having"
3462 | "limit"
3463 | "offset"
3464 | "union"
3465 | "except"
3466 | "intersect"
3467 )
3468}
3469
3470fn sql_compat_tokens(sql: &str) -> Vec<SqlCompatToken> {
3471 let bytes = sql.as_bytes();
3472 let mut tokens = Vec::new();
3473 let mut i = 0;
3474 while i < bytes.len() {
3475 match bytes[i] {
3476 b if b.is_ascii_whitespace() => i += 1,
3477 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
3478 i += 2;
3479 while i < bytes.len() && bytes[i] != b'\n' {
3480 i += 1;
3481 }
3482 }
3483 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
3484 i = skip_block_comment(bytes, i);
3485 }
3486 b'\'' => {
3487 let end = skip_quoted(bytes, i, b'\'');
3488 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
3489 i = end;
3490 }
3491 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
3492 let end = skip_quoted(bytes, i + 1, b'\'');
3493 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
3494 i = end;
3495 }
3496 b'$' => {
3497 let (end, matched) = skip_dollar_quoted(bytes, i);
3498 if matched {
3499 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
3500 i = end;
3501 } else {
3502 i += 1;
3503 }
3504 }
3505 b'"' => {
3506 let end = skip_quoted(bytes, i, b'"');
3507 let raw = sql[i..end].to_string();
3508 let normalized = unquote_sql_ident(&raw).to_ascii_lowercase();
3509 tokens.push(SqlCompatToken {
3510 kind: SqlCompatTokenKind::Ident,
3511 raw,
3512 normalized,
3513 start: i,
3514 end,
3515 });
3516 i = end;
3517 }
3518 b'.' => {
3519 tokens.push(sql_token(sql, SqlCompatTokenKind::Dot, i, i + 1));
3520 i += 1;
3521 }
3522 b'(' => {
3523 tokens.push(sql_token(sql, SqlCompatTokenKind::LParen, i, i + 1));
3524 i += 1;
3525 }
3526 b')' => {
3527 tokens.push(sql_token(sql, SqlCompatTokenKind::RParen, i, i + 1));
3528 i += 1;
3529 }
3530 b',' => {
3531 tokens.push(sql_token(sql, SqlCompatTokenKind::Comma, i, i + 1));
3532 i += 1;
3533 }
3534 b if is_sql_ident_byte(b) => {
3535 let start = i;
3536 i += 1;
3537 while i < bytes.len() && is_sql_ident_byte(bytes[i]) {
3538 i += 1;
3539 }
3540 tokens.push(sql_token(sql, SqlCompatTokenKind::Ident, start, i));
3541 }
3542 _ => i += 1,
3543 }
3544 }
3545 tokens
3546}
3547
3548fn sql_token(sql: &str, kind: SqlCompatTokenKind, start: usize, end: usize) -> SqlCompatToken {
3549 let raw = sql[start..end].to_string();
3550 SqlCompatToken {
3551 kind,
3552 normalized: raw.to_ascii_lowercase(),
3553 raw,
3554 start,
3555 end,
3556 }
3557}
3558
3559fn unquote_sql_ident(raw: &str) -> String {
3560 if raw.len() >= 2 && raw.starts_with('"') && raw.ends_with('"') {
3561 raw[1..raw.len() - 1].replace("\"\"", "\"")
3562 } else {
3563 raw.to_string()
3564 }
3565}
3566
3567fn apply_sql_replacements(sql: &str, replacements: &[SqlReplacement]) -> String {
3568 if replacements.is_empty() {
3569 return sql.to_string();
3570 }
3571 let mut ordered = replacements.to_vec();
3572 ordered.sort_by_key(|replacement| replacement.start);
3573 let mut out = String::with_capacity(sql.len());
3574 let mut cursor = 0;
3575 for replacement in ordered {
3576 if replacement.start < cursor || replacement.end > sql.len() {
3577 continue;
3578 }
3579 out.push_str(&sql[cursor..replacement.start]);
3580 out.push_str(&replacement.replacement);
3581 cursor = replacement.end;
3582 }
3583 out.push_str(&sql[cursor..]);
3584 out
3585}
3586
3587fn rewrite_compat_function_calls(sql: &str) -> String {
3588 let bytes = sql.as_bytes();
3589 let mut out = String::with_capacity(sql.len());
3590 let mut i = 0;
3591 while i < bytes.len() {
3592 match bytes[i] {
3593 b'\'' => i = copy_quoted_to_string(&mut out, bytes, i, b'\''),
3594 b'"' => i = copy_quoted_to_string(&mut out, bytes, i, b'"'),
3595 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
3596 out.push(bytes[i] as char);
3597 i += 1;
3598 i = copy_quoted_to_string(&mut out, bytes, i, b'\'');
3599 }
3600 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
3601 out.push('-');
3602 out.push('-');
3603 i += 2;
3604 while i < bytes.len() {
3605 let ch = bytes[i] as char;
3606 out.push(ch);
3607 i += 1;
3608 if ch == '\n' {
3609 break;
3610 }
3611 }
3612 }
3613 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
3614 let start = i;
3615 i = skip_block_comment(bytes, i);
3616 out.push_str(&sql[start..i.min(bytes.len())]);
3617 }
3618 b'$' => {
3619 let start_len = out.len();
3620 let (next, matched) = copy_dollar_quoted_to_string(&mut out, bytes, i);
3621 if matched {
3622 i = next;
3623 } else {
3624 out.truncate(start_len);
3625 out.push('$');
3626 i += 1;
3627 }
3628 }
3629 b'g' | b'G' | b'm' | b'M' | b't' | b'T' => {
3630 if let Some((replacement, next)) = compat_function_rewrite_at(sql, i) {
3631 out.push_str(&replacement);
3632 i = next;
3633 } else {
3634 out.push(bytes[i] as char);
3635 i += 1;
3636 }
3637 }
3638 _ => {
3639 out.push(bytes[i] as char);
3640 i += 1;
3641 }
3642 }
3643 }
3644 out
3645}
3646
3647fn compat_function_rewrite_at(sql: &str, start: usize) -> Option<(String, usize)> {
3648 let bytes = sql.as_bytes();
3649 let (name, kind) = if ident_eq_at(bytes, start, b"max") {
3650 ("max", CompatRewriteKind::ScalarMax)
3651 } else if ident_eq_at(bytes, start, b"min") {
3652 ("min", CompatRewriteKind::ScalarMin)
3653 } else if ident_eq_at(bytes, start, b"group_concat") {
3654 ("group_concat", CompatRewriteKind::GroupConcat)
3655 } else if ident_eq_at(bytes, start, b"total") {
3656 ("total", CompatRewriteKind::Total)
3657 } else {
3658 return None;
3659 };
3660 let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
3661 let after_name = start + name.len();
3662 let after_ok = bytes
3663 .get(after_name)
3664 .is_some_and(|b| !is_sql_ident_byte(*b));
3665 if !before_ok || !after_ok {
3666 return None;
3667 }
3668 let mut open = after_name;
3669 while open < bytes.len() && bytes[open].is_ascii_whitespace() {
3670 open += 1;
3671 }
3672 if bytes.get(open) != Some(&b'(') {
3673 return None;
3674 }
3675 let summary = call_arg_summary(sql, open)?;
3676 match kind {
3677 CompatRewriteKind::ScalarMax if summary.top_level_commas > 0 => {
3678 Some(("__mongreldb_scalar_max(".to_string(), open + 1))
3679 }
3680 CompatRewriteKind::ScalarMin if summary.top_level_commas > 0 => {
3681 Some(("__mongreldb_scalar_min(".to_string(), open + 1))
3682 }
3683 CompatRewriteKind::GroupConcat => {
3684 let args = &sql[open + 1..summary.close];
3685 let rewritten = if summary.top_level_commas == 0 {
3686 format!("string_agg({args}, ',')")
3687 } else {
3688 format!("string_agg({args})")
3689 };
3690 Some((rewritten, summary.close + 1))
3691 }
3692 CompatRewriteKind::Total if summary.top_level_commas == 0 => {
3693 let args = &sql[open + 1..summary.close];
3694 let suffix_end = aggregate_suffix_end(sql, summary.close + 1);
3695 let suffix = &sql[summary.close + 1..suffix_end];
3696 Some((
3697 format!("coalesce(cast(sum({args}){suffix} as double), 0.0)"),
3698 suffix_end,
3699 ))
3700 }
3701 _ => None,
3702 }
3703}
3704
3705#[derive(Clone, Copy)]
3706enum CompatRewriteKind {
3707 ScalarMax,
3708 ScalarMin,
3709 GroupConcat,
3710 Total,
3711}
3712
3713fn ident_eq_at(bytes: &[u8], start: usize, ident: &[u8]) -> bool {
3714 bytes
3715 .get(start..start + ident.len())
3716 .is_some_and(|slice| slice.eq_ignore_ascii_case(ident))
3717}
3718
3719fn is_sql_ident_byte(b: u8) -> bool {
3720 b.is_ascii_alphanumeric() || b == b'_' || b == b'$'
3721}
3722
3723fn keyword_at(bytes: &[u8], start: usize, keyword: &[u8]) -> bool {
3724 if !ident_eq_at(bytes, start, keyword) {
3725 return false;
3726 }
3727 let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
3728 let after = start + keyword.len();
3729 let after_ok = after >= bytes.len() || !is_sql_ident_byte(bytes[after]);
3730 before_ok && after_ok
3731}
3732
3733fn skip_sql_whitespace(bytes: &[u8], mut i: usize) -> usize {
3734 while i < bytes.len() && bytes[i].is_ascii_whitespace() {
3735 i += 1;
3736 }
3737 i
3738}
3739
3740fn aggregate_suffix_end(sql: &str, start: usize) -> usize {
3741 let bytes = sql.as_bytes();
3742 let mut suffix_end = start;
3743 let mut i = skip_sql_whitespace(bytes, start);
3744
3745 if keyword_at(bytes, i, b"filter") {
3746 let open = skip_sql_whitespace(bytes, i + b"filter".len());
3747 if bytes.get(open) != Some(&b'(') {
3748 return start;
3749 }
3750 let Some(summary) = call_arg_summary(sql, open) else {
3751 return start;
3752 };
3753 suffix_end = summary.close + 1;
3754 i = skip_sql_whitespace(bytes, suffix_end);
3755 }
3756
3757 if keyword_at(bytes, i, b"over") {
3758 let after_over = skip_sql_whitespace(bytes, i + b"over".len());
3759 if bytes.get(after_over) == Some(&b'(') {
3760 let Some(summary) = call_arg_summary(sql, after_over) else {
3761 return suffix_end;
3762 };
3763 suffix_end = summary.close + 1;
3764 } else {
3765 let mut end = after_over;
3766 while end < bytes.len() && is_sql_ident_byte(bytes[end]) {
3767 end += 1;
3768 }
3769 if end > after_over {
3770 suffix_end = end;
3771 }
3772 }
3773 }
3774
3775 suffix_end
3776}
3777
3778struct CallArgSummary {
3779 close: usize,
3780 top_level_commas: usize,
3781}
3782
3783fn call_arg_summary(sql: &str, open: usize) -> Option<CallArgSummary> {
3784 let bytes = sql.as_bytes();
3785 let mut depth = 1;
3786 let mut i = open + 1;
3787 let mut top_level_commas = 0;
3788 while i < bytes.len() {
3789 match bytes[i] {
3790 b'\'' => i = skip_quoted(bytes, i, b'\''),
3791 b'"' => i = skip_quoted(bytes, i, b'"'),
3792 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
3793 i = skip_quoted(bytes, i + 1, b'\'')
3794 }
3795 b'$' => {
3796 let (next, matched) = skip_dollar_quoted(bytes, i);
3797 i = if matched { next } else { i + 1 };
3798 }
3799 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
3800 i += 2;
3801 while i < bytes.len() && bytes[i] != b'\n' {
3802 i += 1;
3803 }
3804 }
3805 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
3806 i = skip_block_comment(bytes, i);
3807 }
3808 b'(' => {
3809 depth += 1;
3810 i += 1;
3811 }
3812 b')' => {
3813 depth -= 1;
3814 if depth == 0 {
3815 return Some(CallArgSummary {
3816 close: i,
3817 top_level_commas,
3818 });
3819 }
3820 i += 1;
3821 }
3822 b',' if depth == 1 => {
3823 top_level_commas += 1;
3824 i += 1;
3825 }
3826 _ => i += 1,
3827 }
3828 }
3829 None
3830}
3831
3832fn copy_quoted_to_string(out: &mut String, bytes: &[u8], start: usize, delim: u8) -> usize {
3833 let end = skip_quoted(bytes, start, delim);
3834 out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
3835 end
3836}
3837
3838fn skip_quoted(bytes: &[u8], start: usize, delim: u8) -> usize {
3839 let mut i = start;
3840 if i < bytes.len() {
3841 i += 1;
3842 }
3843 while i < bytes.len() {
3844 if bytes[i] == delim {
3845 i += 1;
3846 if i < bytes.len() && bytes[i] == delim {
3847 i += 1;
3848 continue;
3849 }
3850 break;
3851 }
3852 i += 1;
3853 }
3854 i
3855}
3856
3857fn copy_dollar_quoted_to_string(out: &mut String, bytes: &[u8], start: usize) -> (usize, bool) {
3858 let (end, matched) = skip_dollar_quoted(bytes, start);
3859 if matched {
3860 out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
3861 }
3862 (end, matched)
3863}
3864
3865fn skip_dollar_quoted(bytes: &[u8], start: usize) -> (usize, bool) {
3866 if bytes.get(start) != Some(&b'$') {
3867 return (start, false);
3868 }
3869 let mut j = start + 1;
3870 while j < bytes.len() && (bytes[j].is_ascii_alphanumeric() || bytes[j] == b'_') {
3871 j += 1;
3872 }
3873 if bytes.get(j) != Some(&b'$') {
3874 return (start, false);
3875 }
3876 let tag = &bytes[start..=j];
3877 let mut i = j + 1;
3878 while i + tag.len() <= bytes.len() {
3879 if &bytes[i..i + tag.len()] == tag {
3880 return (i + tag.len(), true);
3881 }
3882 i += 1;
3883 }
3884 (start, false)
3885}
3886
3887fn skip_block_comment(bytes: &[u8], start: usize) -> usize {
3888 let mut i = start + 2;
3889 let mut depth = 1;
3890 while i + 1 < bytes.len() && depth > 0 {
3891 if bytes[i] == b'/' && bytes[i + 1] == b'*' {
3892 depth += 1;
3893 i += 2;
3894 } else if bytes[i] == b'*' && bytes[i + 1] == b'/' {
3895 depth -= 1;
3896 i += 2;
3897 } else {
3898 i += 1;
3899 }
3900 }
3901 i
3902}
3903
3904fn sql_cache_key(sql: &str) -> u64 {
3906 use std::hash::{Hash, Hasher};
3907 let mut h = std::collections::hash_map::DefaultHasher::new();
3908 sql.hash(&mut h);
3909 h.finish()
3910}
3911
3912fn replace_from_view(sql: &str, name: &str, view_sql: &str) -> String {
3919 let lower = sql.to_ascii_lowercase();
3920 let bytes = lower.as_bytes();
3921 let name_b = name.as_bytes();
3922 let mut i = 0usize;
3923 while let Some(rel) = lower[i..].find("from") {
3924 let from_start = i + rel;
3925 let after_from = from_start + 4;
3926 i = after_from;
3927 if from_start > 0 && is_ident_byte(bytes[from_start - 1]) {
3929 continue;
3930 }
3931 let mut j = after_from;
3933 while j < bytes.len() && bytes[j].is_ascii_whitespace() {
3934 j += 1;
3935 }
3936 if j == after_from || !bytes[j..].starts_with(name_b) {
3937 continue;
3938 }
3939 let after_name = j + name_b.len();
3940 if after_name < bytes.len() && is_ident_byte(bytes[after_name]) {
3942 continue;
3943 }
3944 let mut out = String::with_capacity(sql.len() + view_sql.len() + name.len() + 8);
3947 out.push_str(&sql[..from_start]);
3948 out.push_str(&sql[from_start..j]);
3949 out.push('(');
3950 out.push_str(view_sql);
3951 out.push_str(") AS ");
3952 out.push_str(name);
3953 out.push_str(&sql[after_name..]);
3954 return out;
3955 }
3956 sql.to_string()
3957}
3958
3959fn is_ident_byte(b: u8) -> bool {
3960 b.is_ascii_alphanumeric() || b == b'_'
3961}
3962
3963fn normalize_sql(sql: &str) -> String {
3973 let b = sql.as_bytes();
3974 let n = b.len();
3975 let mut out: Vec<u8> = Vec::with_capacity(n);
3976 let mut want_space = false;
3979 let mut i = 0usize;
3980 while i < n {
3981 let c = b[i];
3982 if c.is_ascii_whitespace() {
3986 want_space = true;
3987 i += 1;
3988 continue;
3989 }
3990 if c == b'-' && i + 1 < n && b[i + 1] == b'-' {
3991 i += 2;
3993 while i < n && b[i] != b'\n' {
3994 i += 1;
3995 }
3996 want_space = !out.is_empty();
3997 continue;
3998 }
3999 if c == b'/' && i + 1 < n && b[i + 1] == b'*' {
4000 i += 2;
4003 let mut depth = 1usize;
4004 while i + 1 < n && depth > 0 {
4005 if b[i] == b'/' && b[i + 1] == b'*' {
4006 depth += 1;
4007 i += 2;
4008 } else if b[i] == b'*' && b[i + 1] == b'/' {
4009 depth -= 1;
4010 i += 2;
4011 } else {
4012 i += 1;
4013 }
4014 }
4015 want_space = !out.is_empty();
4016 continue;
4017 }
4018 if want_space && !out.is_empty() {
4020 out.push(b' ');
4021 }
4022 want_space = false;
4023 match c {
4024 b'E' | b'e' if i + 1 < n && b[i + 1] == b'\'' => {
4026 out.push(c);
4027 i += 1;
4028 i = copy_quoted(&mut out, b, i, n, b'\'');
4029 continue;
4030 }
4031 b'\'' => {
4033 i = copy_quoted(&mut out, b, i, n, b'\'');
4034 continue;
4035 }
4036 b'"' => {
4038 i = copy_quoted(&mut out, b, i, n, b'"');
4039 continue;
4040 }
4041 b'$' => {
4043 let (consumed, matched) = copy_dollar_quoted(&mut out, b, i, n);
4044 if matched {
4045 i = consumed;
4046 continue;
4047 }
4048 out.push(c);
4049 i += 1;
4050 continue;
4051 }
4052 _ => {
4053 out.push(c);
4054 i += 1;
4055 }
4056 }
4057 }
4058 String::from_utf8(out).unwrap_or_else(|_| sql.to_string())
4059}
4060
4061fn copy_quoted(out: &mut Vec<u8>, b: &[u8], start: usize, n: usize, delim: u8) -> usize {
4065 out.push(b[start]);
4066 let mut i = start + 1;
4067 while i < n {
4068 let c = b[i];
4069 out.push(c);
4070 if c == delim {
4071 if i + 1 < n && b[i + 1] == delim {
4073 out.push(b[i + 1]);
4074 i += 2;
4075 continue;
4076 }
4077 return i + 1;
4078 }
4079 i += 1;
4080 }
4081 i
4082}
4083
4084fn copy_dollar_quoted(out: &mut Vec<u8>, b: &[u8], start: usize, n: usize) -> (usize, bool) {
4088 let mut j = start + 1;
4092 let tag_start = j;
4093 while j < n && b[j] != b'$' && is_dollar_tag_byte(b[j]) {
4094 j += 1;
4095 }
4096 if j >= n || b[j] != b'$' {
4097 return (start + 1, false);
4098 }
4099 if tag_start < j && !(b[tag_start].is_ascii_alphabetic() || b[tag_start] == b'_') {
4100 return (start + 1, false);
4101 }
4102 let close_end = j + 1; let delim = &b[start..close_end];
4104 out.extend_from_slice(delim);
4106 let mut k = close_end;
4108 while k + delim.len() <= n {
4109 if &b[k..k + delim.len()] == delim {
4110 out.extend_from_slice(delim);
4111 return (k + delim.len(), true);
4112 }
4113 out.push(b[k]);
4114 k += 1;
4115 }
4116 out.extend_from_slice(&b[close_end..n]);
4118 (n, true)
4119}
4120
4121fn is_dollar_tag_byte(b: u8) -> bool {
4122 b.is_ascii_alphanumeric() || b == b'_'
4123}
4124
4125fn strip_prefix_ci<'a>(s: &'a str, prefix: &str) -> Option<&'a str> {
4127 let bytes = s.as_bytes();
4128 let pb = prefix.as_bytes();
4129 if bytes.len() >= pb.len() && bytes[..pb.len()].eq_ignore_ascii_case(pb) {
4130 Some(&s[pb.len()..])
4131 } else {
4132 None
4133 }
4134}
4135
4136const COLUMN_CONSTRAINTS: &[(&str, u32)] = &[
4149 ("primary key", ColumnFlags::PRIMARY_KEY),
4150 ("autoincrement", ColumnFlags::AUTO_INCREMENT),
4157 ("auto_increment", ColumnFlags::AUTO_INCREMENT),
4158];
4159
4160fn parse_column_constraints(constraint_text: &str) -> ColumnFlags {
4166 let normalized = constraint_text.to_lowercase();
4167 let mut flags = ColumnFlags::empty();
4168 for (phrase, bit) in COLUMN_CONSTRAINTS {
4169 if normalized.contains(phrase) {
4170 flags = flags.with(*bit);
4171 }
4172 }
4173 flags
4174}
4175
4176fn parse_sql_type(ty_str: &str) -> Result<mongreldb_core::schema::TypeId> {
4177 use mongreldb_core::schema::TypeId;
4178
4179 match ty_str.trim().trim_end_matches(';').to_lowercase().as_str() {
4180 "bigint" | "int8" | "int64" | "integer" | "int" => Ok(TypeId::Int64),
4181 "double" | "float8" | "float64" | "real" | "float" => Ok(TypeId::Float64),
4182 "varchar" | "text" | "string" | "bytes" => Ok(TypeId::Bytes),
4183 "boolean" | "bool" => Ok(TypeId::Bool),
4184 other => Err(MongrelQueryError::Schema(format!(
4185 "unsupported column type: {other}"
4186 ))),
4187 }
4188}
4189
4190fn parse_create_table(sql: &str) -> Result<(String, mongreldb_core::schema::Schema)> {
4196 use mongreldb_core::schema::*;
4197
4198 let open = sql
4199 .find('(')
4200 .ok_or(MongrelQueryError::Schema("CREATE TABLE missing '('".into()))?;
4201 let close = sql
4202 .rfind(')')
4203 .ok_or(MongrelQueryError::Schema("CREATE TABLE missing ')'".into()))?;
4204 let head = sql[..open].trim();
4205 let after_kw = strip_prefix_ci(head, "CREATE TABLE")
4206 .or_else(|| strip_prefix_ci(head, "create table"))
4207 .unwrap_or("")
4208 .trim();
4209 let after_kw = after_kw
4211 .strip_prefix("IF NOT EXISTS")
4212 .or_else(|| after_kw.strip_prefix("if not exists"))
4213 .map(str::trim)
4214 .unwrap_or(after_kw);
4215 let name = after_kw.trim_matches('"').to_string();
4216 if name.is_empty() {
4217 return Err(MongrelQueryError::Schema(
4218 "CREATE TABLE missing table name".into(),
4219 ));
4220 }
4221
4222 let body = &sql[open + 1..close];
4223 let mut columns = Vec::new();
4224 let schema_id: u64 = 0; for (i, raw) in body.split(',').enumerate() {
4226 let part = raw.trim();
4227 if part.is_empty() {
4228 continue;
4229 }
4230 let mut tokens = part.split_whitespace();
4231 let col_name = tokens
4232 .next()
4233 .ok_or(MongrelQueryError::Schema("missing column name".into()))?
4234 .trim_matches('"');
4235 let ty_str = tokens
4236 .next()
4237 .ok_or(MongrelQueryError::Schema("missing column type".into()))?
4238 .to_lowercase();
4239 let ty = parse_sql_type(&ty_str)?;
4240 let constraint_clause: String = tokens.collect::<Vec<_>>().join(" ");
4244 let flags = parse_column_constraints(&constraint_clause);
4245 columns.push(ColumnDef {
4246 id: (i + 1) as u16,
4247 name: col_name.to_string(),
4248 ty,
4249 flags,
4250 default_value: None,
4251 });
4252 }
4253
4254 Ok((
4255 name,
4256 Schema {
4257 schema_id,
4258 columns,
4259 indexes: vec![],
4260 colocation: vec![],
4261 constraints: Default::default(),
4262 clustered: false,
4263 },
4264 ))
4265}
4266
4267fn parse_drop_table(sql: &str) -> Result<(String, bool)> {
4269 let head = sql.trim();
4270 let after_kw = strip_prefix_ci(head, "DROP TABLE")
4271 .or_else(|| strip_prefix_ci(head, "drop table"))
4272 .unwrap_or("")
4273 .trim();
4274 let (rest, if_exists) = if let Some(r) = after_kw
4276 .strip_prefix("IF EXISTS")
4277 .or_else(|| after_kw.strip_prefix("if exists"))
4278 .map(str::trim)
4279 {
4280 (r, true)
4281 } else {
4282 (after_kw, false)
4283 };
4284 let name = rest.trim_matches(';').trim_matches('"').trim();
4285 if name.is_empty() {
4286 return Err(MongrelQueryError::Schema(
4287 "DROP TABLE missing table name".into(),
4288 ));
4289 }
4290 Ok((name.to_string(), if_exists))
4291}
4292
4293enum ParsedAlterTable {
4294 RenameTable {
4295 old_name: String,
4296 new_name: String,
4297 },
4298 RenameColumn {
4299 table_name: String,
4300 column_name: String,
4301 new_name: String,
4302 },
4303 AlterColumnType {
4304 table_name: String,
4305 column_name: String,
4306 ty: mongreldb_core::schema::TypeId,
4307 },
4308 SetNotNull {
4309 table_name: String,
4310 column_name: String,
4311 },
4312 DropNotNull {
4313 table_name: String,
4314 column_name: String,
4315 },
4316}
4317
4318fn current_column_flags(db: &Arc<Database>, table: &str, column: &str) -> Result<ColumnFlags> {
4319 let handle = db.table(table)?;
4320 let table = handle.lock();
4321 table
4322 .schema()
4323 .column(column)
4324 .map(|c| c.flags)
4325 .ok_or_else(|| MongrelQueryError::Schema(format!("unknown column {column}")))
4326}
4327
4328fn parse_alter_table(sql: &str) -> Result<ParsedAlterTable> {
4329 let trimmed = strip_statement_semicolon(sql.trim());
4330 let after_kw = strip_prefix_ci(trimmed, "ALTER TABLE")
4331 .ok_or_else(|| MongrelQueryError::Schema("not an ALTER TABLE statement".into()))?
4332 .trim();
4333 let (table_name, rest) = take_sql_ident(after_kw, "ALTER TABLE missing table name")?;
4334 let rest = rest.trim();
4335
4336 if let Some(after) = strip_prefix_ci(rest, "RENAME TO") {
4337 let new_name = parse_trailing_identifier(after, "ALTER TABLE missing new table name")?;
4338 return Ok(ParsedAlterTable::RenameTable {
4339 old_name: table_name,
4340 new_name,
4341 });
4342 }
4343
4344 if let Some(after) = strip_prefix_ci(rest, "RENAME COLUMN") {
4345 let (column_name, after_col) =
4346 take_sql_ident(after, "ALTER TABLE RENAME COLUMN missing column name")?;
4347 let after_to = strip_prefix_ci(after_col.trim(), "TO").ok_or_else(|| {
4348 MongrelQueryError::Schema("ALTER TABLE RENAME COLUMN missing TO".into())
4349 })?;
4350 let new_name = parse_trailing_identifier(
4351 after_to,
4352 "ALTER TABLE RENAME COLUMN missing new column name",
4353 )?;
4354 return Ok(ParsedAlterTable::RenameColumn {
4355 table_name,
4356 column_name,
4357 new_name,
4358 });
4359 }
4360
4361 let after_alter = strip_prefix_ci(rest, "ALTER COLUMN")
4362 .or_else(|| strip_prefix_ci(rest, "ALTER"))
4363 .ok_or_else(|| {
4364 MongrelQueryError::Schema(
4365 "ALTER TABLE must be RENAME TO, RENAME COLUMN, or ALTER COLUMN".into(),
4366 )
4367 })?;
4368 let (column_name, action) =
4369 take_sql_ident(after_alter, "ALTER TABLE ALTER COLUMN missing column name")?;
4370 let action = action.trim();
4371
4372 if let Some(after_type) =
4373 strip_prefix_ci(action, "TYPE").or_else(|| strip_prefix_ci(action, "SET DATA TYPE"))
4374 {
4375 let ty = parse_type_tail(after_type)?;
4376 return Ok(ParsedAlterTable::AlterColumnType {
4377 table_name,
4378 column_name,
4379 ty,
4380 });
4381 }
4382 if strip_prefix_ci(action, "SET NOT NULL").is_some() {
4383 return Ok(ParsedAlterTable::SetNotNull {
4384 table_name,
4385 column_name,
4386 });
4387 }
4388 if strip_prefix_ci(action, "DROP NOT NULL").is_some() {
4389 return Ok(ParsedAlterTable::DropNotNull {
4390 table_name,
4391 column_name,
4392 });
4393 }
4394
4395 Err(MongrelQueryError::Schema(
4396 "unsupported ALTER COLUMN action".into(),
4397 ))
4398}
4399
4400fn strip_statement_semicolon(s: &str) -> &str {
4401 s.trim().trim_end_matches(';').trim()
4402}
4403
4404fn take_sql_ident<'a>(s: &'a str, missing: &str) -> Result<(String, &'a str)> {
4405 let s = s.trim();
4406 if s.is_empty() {
4407 return Err(MongrelQueryError::Schema(missing.into()));
4408 }
4409 if let Some(rest) = s.strip_prefix('"') {
4410 let Some(end) = rest.find('"') else {
4411 return Err(MongrelQueryError::Schema(
4412 "unterminated quoted identifier".into(),
4413 ));
4414 };
4415 let ident = rest[..end].to_string();
4416 if ident.is_empty() {
4417 return Err(MongrelQueryError::Schema(missing.into()));
4418 }
4419 return Ok((ident, &rest[end + 1..]));
4420 }
4421 let end = s.find(|c: char| c.is_ascii_whitespace()).unwrap_or(s.len());
4422 let ident = s[..end].trim_matches('"').to_string();
4423 if ident.is_empty() {
4424 return Err(MongrelQueryError::Schema(missing.into()));
4425 }
4426 Ok((ident, &s[end..]))
4427}
4428
4429fn parse_trailing_identifier(s: &str, missing: &str) -> Result<String> {
4430 let (ident, rest) = take_sql_ident(s, missing)?;
4431 if !strip_statement_semicolon(rest).is_empty() {
4432 return Err(MongrelQueryError::Schema(
4433 "unexpected tokens after identifier".into(),
4434 ));
4435 }
4436 Ok(ident)
4437}
4438
4439fn parse_type_tail(s: &str) -> Result<mongreldb_core::schema::TypeId> {
4440 let tail = strip_statement_semicolon(s);
4441 let ty = tail
4442 .split_whitespace()
4443 .next()
4444 .ok_or_else(|| MongrelQueryError::Schema("ALTER COLUMN TYPE missing type".into()))?;
4445 parse_sql_type(ty)
4446}
4447
4448#[cfg(test)]
4449mod tests {
4450 use super::*;
4451
4452 #[test]
4453 fn bounded_lru_evicts_least_recently_used() {
4454 let mut cache = BoundedLru::new(2);
4455 cache.insert("a", 1);
4456 cache.insert("b", 2);
4457 assert_eq!(cache.get("a"), Some(&1));
4458 cache.insert("c", 3);
4459 assert_eq!(cache.entries.len(), 2);
4460 assert_eq!(cache.get("b"), None);
4461 assert_eq!(cache.get("a"), Some(&1));
4462 assert_eq!(cache.get("c"), Some(&3));
4463 }
4464
4465 #[tokio::test]
4466 async fn streaming_query_bypasses_result_cache() {
4467 use futures::StreamExt;
4468
4469 let dir = tempfile::tempdir().unwrap();
4470 let database = Arc::new(Database::create(dir.path()).unwrap());
4471 let session = MongrelSession::open(database).unwrap();
4472 let mut stream = session.run_stream("SELECT 1 AS value").await.unwrap();
4473 let batch = stream.next().await.unwrap().unwrap();
4474 assert_eq!(batch.num_rows(), 1);
4475 assert!(stream.next().await.is_none());
4476 assert!(session.cache.lock().entries.is_empty());
4477 }
4478
4479 #[tokio::test]
4480 async fn ttl_tables_bypass_epoch_keyed_result_cache() {
4481 let dir = tempfile::tempdir().unwrap();
4482 let database = Arc::new(Database::create(dir.path()).unwrap());
4483 let session = MongrelSession::open(database).unwrap();
4484 session
4485 .run(
4486 "CREATE TABLE events (id BIGINT PRIMARY KEY, ts TIMESTAMP) \
4487 TTL_COLUMN ts TTL '1 day'",
4488 )
4489 .await
4490 .unwrap();
4491 session.run("SELECT * FROM events").await.unwrap();
4492 assert!(session.cache.lock().entries.is_empty());
4493 }
4494
4495 #[test]
4496 fn normalize_collapses_and_trims_whitespace() {
4497 assert_eq!(normalize_sql("SELECT * FROM t"), "SELECT * FROM t");
4498 assert_eq!(normalize_sql(" SELECT * FROM t "), "SELECT * FROM t");
4499 assert_eq!(
4500 normalize_sql("\n\tSELECT\n*\nFROM\n\tt\n"),
4501 "SELECT * FROM t"
4502 );
4503 assert_eq!(
4504 normalize_sql("SELECT a, b FROM t"),
4505 normalize_sql("SELECT a, b FROM t")
4506 );
4507 }
4508
4509 #[test]
4510 fn normalize_preserves_string_literal_whitespace() {
4511 assert_eq!(
4512 normalize_sql("SELECT 'hello world' FROM t"),
4513 "SELECT 'hello world' FROM t"
4514 );
4515 assert_eq!(
4516 normalize_sql("SELECT 'it''s ok' FROM t"),
4517 "SELECT 'it''s ok' FROM t"
4518 );
4519 assert_eq!(
4520 normalize_sql(" SELECT 'a b' FROM t "),
4521 "SELECT 'a b' FROM t"
4522 );
4523 }
4524
4525 #[test]
4526 fn normalize_preserves_quoted_identifier_and_dollar_quote() {
4527 assert_eq!(
4528 normalize_sql(" SELECT \"my col\" FROM t "),
4529 "SELECT \"my col\" FROM t"
4530 );
4531 assert_eq!(
4532 normalize_sql(" SELECT $$a b$$ FROM t "),
4533 "SELECT $$a b$$ FROM t"
4534 );
4535 assert_eq!(
4536 normalize_sql("SELECT $tag$body with spaces$tag$ FROM t"),
4537 "SELECT $tag$body with spaces$tag$ FROM t"
4538 );
4539 }
4540
4541 #[test]
4542 fn normalize_strips_comments() {
4543 assert_eq!(
4544 normalize_sql("SELECT 1 -- trailing comment\nFROM t"),
4545 "SELECT 1 FROM t"
4546 );
4547 assert_eq!(
4548 normalize_sql("SELECT /* block */ 1 FROM t"),
4549 "SELECT 1 FROM t"
4550 );
4551 assert_eq!(
4553 normalize_sql("SELECT /* 'not a string' */ 1 FROM t"),
4554 "SELECT 1 FROM t"
4555 );
4556 assert_eq!(
4558 normalize_sql("SELECT /* outer /* inner */ still outer */ 1 FROM t"),
4559 "SELECT 1 FROM t"
4560 );
4561 }
4562
4563 #[test]
4564 fn normalize_escape_string_preserved() {
4565 assert_eq!(
4566 normalize_sql("SELECT E'line\\nbreak' FROM t"),
4567 "SELECT E'line\\nbreak' FROM t"
4568 );
4569 }
4570
4571 #[test]
4572 fn replace_from_view_matches_whole_word_only() {
4573 let out = replace_from_view("SELECT * FROM logs", "log", "SELECT 1");
4574 assert_eq!(out, "SELECT * FROM logs");
4575
4576 let out = replace_from_view("SELECT * FROM log", "log", "SELECT 1");
4577 assert_eq!(out, "SELECT * FROM (SELECT 1) AS log");
4578
4579 let out = replace_from_view("select * from log where x", "log", "SELECT 1");
4580 assert_eq!(out, "select * from (SELECT 1) AS log where x");
4581
4582 let out = replace_from_view("SELECT * FROM log)", "log", "SELECT 1");
4583 assert_eq!(out, "SELECT * FROM (SELECT 1) AS log)");
4584
4585 let out = replace_from_view("SELECT * xfrom log", "log", "SELECT 1");
4586 assert_eq!(out, "SELECT * xfrom log");
4587 }
4588
4589 #[test]
4590 fn compat_function_rewrite_handles_sqlite_compatibility_calls() {
4591 assert_eq!(
4592 rewrite_compat_function_calls("select max(id), min(id) from t"),
4593 "select max(id), min(id) from t"
4594 );
4595 assert_eq!(
4596 rewrite_compat_function_calls("select max(1, min(2, 3), 'max(4,5)')"),
4597 "select __mongreldb_scalar_max(1, __mongreldb_scalar_min(2, 3), 'max(4,5)')"
4598 );
4599 assert_eq!(
4600 rewrite_compat_function_calls("select /* max(1,2) */ min(1, (2 + 3))"),
4601 "select /* max(1,2) */ __mongreldb_scalar_min(1, (2 + 3))"
4602 );
4603 assert_eq!(
4604 rewrite_compat_function_calls("select max_value, min_value from t"),
4605 "select max_value, min_value from t"
4606 );
4607 assert_eq!(
4608 rewrite_compat_function_calls(
4609 "select group_concat(label), group_concat(label, '|') from t"
4610 ),
4611 "select string_agg(label, ','), string_agg(label, '|') from t"
4612 );
4613 assert_eq!(
4614 rewrite_compat_function_calls("select total(val), total(val) filter (where grp = 2) from t"),
4615 "select coalesce(cast(sum(val) as double), 0.0), coalesce(cast(sum(val) filter (where grp = 2) as double), 0.0) from t"
4616 );
4617 assert_eq!(
4618 rewrite_compat_function_calls(
4619 "select total(val) over (partition by grp order by id) from t"
4620 ),
4621 "select coalesce(cast(sum(val) over (partition by grp order by id) as double), 0.0) from t"
4622 );
4623 }
4624}