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