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