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