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