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 query_registry;
33mod scan;
34mod scored_sql;
35mod shadow;
36mod udf;
37
38pub use error::{MongrelQueryError, Result};
39pub use external_modules::{
40 ExternalBaseWrite, ExternalModuleDescriptor, ExternalModuleIndex, ExternalModuleRegistry,
41 ExternalPlan, ExternalPlanRequest, ExternalScan, ExternalTable, ExternalTableModule,
42 ExternalTxn, ExternalWriteOp, ExternalWriteResult, ModuleConnectCtx,
43};
44pub use query_registry::{
45 CancelOutcome, CommitFenceOutcome, DurableOutcome, QueryId, QueryStatus, QueryTerminalError,
46 QueryTerminalErrorCategory, QueryTerminalState, RegisteredQueryGuard, RegisteredSqlQuery,
47 SerializationOutcome, SqlQueryOptions, SqlQueryPhase, SqlQueryRegistry,
48};
49
50pub fn contains_boolean_ai_predicate(sql: &str) -> bool {
54 use sqlparser::dialect::GenericDialect;
55 use sqlparser::tokenizer::{Token, Tokenizer};
56
57 let Ok(tokens) = Tokenizer::new(&GenericDialect {}, sql).tokenize() else {
58 return false;
59 };
60 let tokens: Vec<_> = tokens
61 .iter()
62 .filter(|token| !matches!(token, Token::Whitespace(_)))
63 .collect();
64 tokens.windows(2).any(|tokens| {
65 matches!(
66 (tokens[0], tokens[1]),
67 (Token::Word(word), Token::LParen)
68 if word.value.eq_ignore_ascii_case("ann_search")
69 || word.value.eq_ignore_ascii_case("sparse_match")
70 )
71 })
72}
73
74pub fn normalized_sql_fingerprint(sql: &str) -> [u8; 32] {
78 use sha2::{Digest, Sha256};
79
80 Sha256::digest(normalize_sql(sql).as_bytes()).into()
81}
82
83pub fn is_single_read_only_query(sql: &str) -> bool {
87 use sqlparser::ast::Statement;
88 use sqlparser::dialect::GenericDialect;
89 use sqlparser::parser::Parser;
90
91 Parser::parse_sql(&GenericDialect {}, sql)
92 .is_ok_and(|statements| matches!(statements.as_slice(), [Statement::Query(_)]))
93}
94
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
100pub enum SqlIdempotencyClass {
101 ReadOnly,
102 SingleWrite,
103 Unsupported,
104}
105
106pub fn classify_sql_idempotency(sql: &str) -> SqlIdempotencyClass {
107 use sqlparser::ast::{ObjectType, Statement};
108 use sqlparser::dialect::GenericDialect;
109 use sqlparser::parser::Parser;
110
111 let Ok(statements) = Parser::parse_sql(&GenericDialect {}, sql) else {
112 return SqlIdempotencyClass::Unsupported;
113 };
114 match statements.as_slice() {
115 [Statement::Query(_)] => SqlIdempotencyClass::ReadOnly,
116 [Statement::CreateView(view)] if view.materialized => SqlIdempotencyClass::SingleWrite,
117 [Statement::Insert(_)
118 | Statement::Update(_)
119 | Statement::Delete(_)
120 | Statement::Truncate(_)
121 | Statement::CreateTable(_)
122 | Statement::CreateVirtualTable { .. }
123 | Statement::AlterTable(_)
124 | Statement::CreateIndex(_)
125 | Statement::CreateTrigger(_)
126 | Statement::DropTrigger(_)
127 | Statement::CreatePolicy(_)
128 | Statement::DropPolicy(_)] => SqlIdempotencyClass::SingleWrite,
129 [Statement::Drop {
130 object_type: ObjectType::Table | ObjectType::MaterializedView | ObjectType::Index,
131 names,
132 ..
133 }] if names.len() == 1 => SqlIdempotencyClass::SingleWrite,
134 [_] => SqlIdempotencyClass::Unsupported,
135 _ => SqlIdempotencyClass::Unsupported,
136 }
137}
138
139pub type MongrelRecordBatchStream = datafusion::physical_plan::SendableRecordBatchStream;
140
141#[doc(hidden)]
142#[derive(Debug, Clone, Copy, PartialEq, Eq)]
143pub enum SqlTestHookPoint {
144 AfterRegistration,
145 WaitingForSqlPermit,
146 WaitingForSessionPermit,
147 BeforeServerIdempotencyCheck,
148 Planning,
149 BeforeScanBatch,
150 AfterScanBatch,
151 BeforeCommitFence,
152 InsideCommitCritical,
153 AfterDurableCommit,
154 AfterStatementStaging,
155 BeforeStatementExecution,
156 BeforeAssignmentEvaluation,
157 DuringOrderedDmlPermutation,
158 DuringTriggerExpansion,
159 BeforeSerializationBatch,
160 DuringPaginationDeserialization,
161 AfterSerialization,
162 AfterPageResponseSerialization,
163}
164
165#[doc(hidden)]
166pub type SqlTestHook = Arc<dyn Fn(SqlTestHookPoint) + Send + Sync>;
167
168tokio::task_local! {
169 static CURRENT_SQL_QUERY: RegisteredSqlQuery;
172 static CURRENT_COLLECTION_LIMITS: SqlCollectionLimits;
173}
174
175use arrow::array::{Array, ArrayRef, Int64Array, StringArray};
176use arrow::datatypes::SchemaRef;
177use arrow::record_batch::RecordBatch;
178use datafusion::catalog::{Session, TableProvider};
179use datafusion::common::{DataFusionError, Result as DFResult};
180use datafusion::logical_expr::{AggregateUDF, Expr, ScalarUDF, TableType, WindowUDF};
181use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
182use datafusion::physical_plan::{execute_stream, ExecutionPlan, RecordBatchStream};
183use datafusion::prelude::SessionContext;
184use mongreldb_core::{
185 AlterColumn, ColumnFlags, Cursor, Database, OwnedSnapshotGuard, Schema as CoreSchema, Snapshot,
186 Table, TableHandle, TypeId,
187};
188use parking_lot::Mutex;
189use std::borrow::Borrow;
190use std::collections::{HashMap, HashSet};
191use std::hash::Hash;
192use std::sync::Arc;
193
194use query_registry::SqlTaskContext;
195
196pub struct MongrelProvider {
199 db: TableHandle,
200 schema: SchemaRef,
201 core_schema: CoreSchema,
202 snapshot: Option<Snapshot>,
203 _retention: Option<Arc<OwnedSnapshotGuard>>,
204 security: Option<ProviderSecurity>,
205}
206
207#[derive(Clone)]
208struct ProviderSecurity {
209 database: Arc<Database>,
210 table: String,
211 principal: Option<mongreldb_core::Principal>,
212 principal_catalog_bound: bool,
213}
214
215fn principal_is_catalog_bound(database: &Database, principal: &mongreldb_core::Principal) -> bool {
216 database.require_auth_enabled() || principal.user_id != 0
217}
218
219impl ProviderSecurity {
220 fn principal(&self) -> mongreldb_core::Result<Option<mongreldb_core::Principal>> {
221 let Some(principal) = &self.principal else {
222 return Ok(None);
223 };
224 if self.principal_catalog_bound {
225 self.database
226 .resolve_current_principal(principal)
227 .map(Some)
228 .ok_or(mongreldb_core::MongrelError::AuthRequired)
229 } else {
230 Ok(Some(principal.clone()))
231 }
232 }
233}
234
235#[derive(Debug, Clone)]
236pub(crate) struct ViewDef {
237 pub sql: String,
238 pub schema: CoreSchema,
239 pub input_types: HashMap<u16, Option<TypeId>>,
240}
241
242impl MongrelProvider {
243 pub fn new(db: Arc<Mutex<Table>>) -> Result<Self> {
244 Self::new_handle(db.into())
245 }
246
247 pub(crate) fn new_handle(db: TableHandle) -> Result<Self> {
248 let (schema, core_schema) = {
249 let db = db.lock();
250 (arrow_conv::arrow_schema(db.schema())?, db.schema().clone())
251 };
252 Ok(Self {
253 db,
254 schema,
255 core_schema,
256 snapshot: None,
257 _retention: None,
258 security: None,
259 })
260 }
261
262 pub(crate) fn new_secured(
263 db: TableHandle,
264 database: Arc<Database>,
265 table: String,
266 principal: Option<mongreldb_core::Principal>,
267 ) -> Result<Self> {
268 let core_schema = db.lock().schema().clone();
269 let schema = arrow_conv::arrow_schema(&core_schema)?;
270 Ok(Self {
271 db,
272 schema,
273 core_schema,
274 snapshot: None,
275 _retention: None,
276 security: Some(ProviderSecurity {
277 principal_catalog_bound: principal
278 .as_ref()
279 .is_some_and(|principal| principal_is_catalog_bound(&database, principal)),
280 database,
281 table,
282 principal,
283 }),
284 })
285 }
286
287 fn new_historical(
288 db: TableHandle,
289 snapshot: Snapshot,
290 retention: OwnedSnapshotGuard,
291 security: Option<ProviderSecurity>,
292 ) -> Result<Self> {
293 let full_schema = {
294 let db = db.lock();
295 db.schema().clone()
296 };
297 let core_schema = full_schema;
298 let schema = arrow_conv::arrow_schema(&core_schema)?;
299 Ok(Self {
300 db,
301 schema,
302 core_schema,
303 snapshot: Some(snapshot),
304 _retention: Some(Arc::new(retention)),
305 security,
306 })
307 }
308
309 pub fn arrow_schema(&self) -> SchemaRef {
310 self.schema.clone()
311 }
312}
313
314impl std::fmt::Debug for MongrelProvider {
315 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
316 f.debug_struct("MongrelProvider").finish_non_exhaustive()
317 }
318}
319
320struct AsOfRegistration {
321 ctx: SessionContext,
322 table_name: String,
323}
324
325impl Drop for AsOfRegistration {
326 fn drop(&mut self) {
327 let _ = self.ctx.deregister_table(&self.table_name);
328 }
329}
330
331struct AsOfQuery {
332 sql: String,
333 registration: AsOfRegistration,
334}
335
336#[async_trait::async_trait]
337impl TableProvider for MongrelProvider {
338 fn schema(&self) -> SchemaRef {
339 self.schema.clone()
340 }
341
342 fn table_type(&self) -> TableType {
343 TableType::Base
344 }
345
346 fn supports_filters_pushdown(
351 &self,
352 filters: &[&Expr],
353 ) -> DFResult<Vec<datafusion::logical_expr::TableProviderFilterPushDown>> {
354 use datafusion::logical_expr::TableProviderFilterPushDown;
355 if self.snapshot.is_some() {
356 return Ok(vec![
357 TableProviderFilterPushDown::Unsupported;
358 filters.len()
359 ]);
360 }
361 let schema_ref = self.db.lock().schema().clone();
362 if self.security.is_some() {
363 return Ok(filters
364 .iter()
365 .map(|filter| {
366 if translate_ann_search(filter, &schema_ref).is_some()
367 || translate_sparse_match(filter, &schema_ref).is_some()
368 {
369 TableProviderFilterPushDown::Exact
370 } else {
371 TableProviderFilterPushDown::Unsupported
372 }
373 })
374 .collect());
375 }
376 Ok(filters
377 .iter()
378 .map(|f| match translate_filter(f, &schema_ref) {
379 Some(
380 mongreldb_core::Condition::FmContains { .. }
381 | mongreldb_core::Condition::FmContainsAll { .. },
382 ) => TableProviderFilterPushDown::Inexact,
383 Some(_) => TableProviderFilterPushDown::Exact,
384 None => match translate_ann_search(f, &schema_ref)
385 .or_else(|| translate_sparse_match(f, &schema_ref))
386 {
387 Some(_) => TableProviderFilterPushDown::Exact,
388 None => TableProviderFilterPushDown::Unsupported,
389 },
390 })
391 .collect())
392 }
393
394 async fn scan(
395 &self,
396 _state: &dyn Session,
397 projection: Option<&Vec<usize>>,
398 filters: &[Expr],
399 _limit: Option<usize>,
400 ) -> DFResult<Arc<dyn ExecutionPlan>> {
401 let core_err = |e: mongreldb_core::MongrelError| {
402 DataFusionError::External(Box::new(MongrelQueryError::Core(e)))
403 };
404 let registered_query = CURRENT_SQL_QUERY.try_with(Clone::clone).ok();
405 if let Some(security) = &self.security {
406 let principal = security.principal().map_err(core_err)?;
407 let allowed = security
408 .database
409 .select_column_ids_for(&security.table, principal.as_ref())
410 .map_err(core_err)?;
411 let projected = projection
412 .map(|projection| {
413 projection
414 .iter()
415 .filter_map(|index| self.core_schema.columns.get(*index))
416 .map(|column| column.id)
417 .collect::<Vec<_>>()
418 })
419 .unwrap_or_else(|| {
420 self.core_schema
421 .columns
422 .iter()
423 .map(|column| column.id)
424 .collect()
425 });
426 if projected.iter().any(|column| !allowed.contains(column)) {
427 security
428 .database
429 .require_columns_for(
430 &security.table,
431 mongreldb_core::ColumnOperation::Select,
432 &projected,
433 principal.as_ref(),
434 )
435 .map_err(core_err)?;
436 }
437 let ai_conditions = filters
438 .iter()
439 .filter_map(|filter| {
440 translate_ann_search(filter, &self.core_schema)
441 .or_else(|| translate_sparse_match(filter, &self.core_schema))
442 })
443 .collect::<Vec<_>>();
444 let mut required_columns = projected.clone();
445 required_columns.extend(mongreldb_core::query::condition_columns(&ai_conditions));
446 required_columns.sort_unstable();
447 required_columns.dedup();
448 let rows = if let Some(snapshot) = self.snapshot {
449 let rows = self.db.lock().visible_rows(snapshot).map_err(core_err)?;
450 security
451 .database
452 .secure_rows_for(&security.table, rows, principal.as_ref())
453 .map_err(core_err)?
454 } else {
455 security
456 .database
457 .with_authorized_read(
458 &security.table,
459 principal.as_ref(),
460 security.principal_catalog_bound,
461 |table, snapshot, allowed, effective_principal| {
462 security.database.require_columns_for(
463 &security.table,
464 mongreldb_core::ColumnOperation::Select,
465 &required_columns,
466 effective_principal,
467 )?;
468 let rows = if ai_conditions.is_empty() {
469 let mut rows = table.visible_rows(snapshot)?;
470 if let Some(allowed) = allowed {
471 rows.retain(|row| allowed.contains(&row.row_id));
472 }
473 rows
474 } else {
475 table.query_at_with_allowed(
476 &mongreldb_core::Query {
477 conditions: ai_conditions.clone(),
478 limit: Some(mongreldb_core::query::MAX_FINAL_LIMIT),
479 offset: 0,
480 },
481 snapshot,
482 allowed,
483 )?
484 };
485 security.database.secure_rows_for(
486 &security.table,
487 rows,
488 effective_principal,
489 )
490 },
491 )
492 .map_err(core_err)?
493 };
494 enforce_source_rows_limit(&rows, registered_query.as_ref())?;
495 if projection.is_some_and(Vec::is_empty) {
496 return Ok(Arc::new(scan::MongrelScanExec::new_row_count(rows.len())));
497 }
498 let batch = arrow_conv::rows_to_batch(&rows, &self.core_schema)
499 .map_err(|error| DataFusionError::External(Box::new(error)))?;
500 let batch = match projection {
501 Some(projection) => batch.project(projection).map_err(|error| {
502 DataFusionError::External(Box::new(MongrelQueryError::Arrow(error.to_string())))
503 })?,
504 None => batch,
505 };
506 let schema = batch.schema();
507 let statistics = (0..batch.num_columns())
508 .map(|_| scan::to_col_statistics(None))
509 .collect();
510 return Ok(Arc::new(scan::MongrelScanExec::new_batch(
511 schema, batch, statistics,
512 )));
513 }
514 let mut db = self.db.lock();
515 db.require_select().map_err(core_err)?;
519 let historical = self.snapshot.is_some();
520 let snap = self.snapshot.unwrap_or_else(|| db.snapshot());
521 let schema_ref = db.schema().clone();
522
523 let translated: Vec<mongreldb_core::Condition> = if historical {
525 Vec::new()
526 } else {
527 filters
528 .iter()
529 .filter_map(|f| {
530 translate_filter(f, &schema_ref)
531 .or_else(|| translate_ann_search(f, &schema_ref))
532 .or_else(|| translate_sparse_match(f, &schema_ref))
533 })
534 .collect()
535 };
536
537 if !translated.is_empty() {
540 db.ensure_indexes_complete().map_err(core_err)?;
541 }
542
543 let empty_proj = projection.map(|p| p.is_empty()).unwrap_or(false);
547 if empty_proj {
548 let total: usize = if historical {
549 db.visible_rows(snap).map_err(core_err)?.len()
550 } else if translated.is_empty() {
551 mongreldb_core::trace::QueryTrace::record(|t| {
552 t.scan_mode = mongreldb_core::trace::ScanMode::CountMetadata;
553 });
554 db.count() as usize
555 } else if let Some(count) = db.count_conditions(&translated, snap).map_err(core_err)? {
556 count as usize
557 } else {
558 match schema_ref.columns.first() {
559 Some(cdef) => {
560 let one = [cdef.id];
561 let cols = match db
562 .query_columns_native_cached(&translated, Some(&one), snap)
563 .map_err(core_err)?
564 {
565 Some(c) => c,
566 None => db
567 .visible_columns_native(snap, Some(&one))
568 .map_err(core_err)?,
569 };
570 mongreldb_core::trace::QueryTrace::record(|t| {
571 t.scan_mode = mongreldb_core::trace::ScanMode::Materialized;
572 });
573 cols.first().map(|(_, c)| c.len()).unwrap_or(0)
574 }
575 None => 0,
576 }
577 };
578 return Ok(Arc::new(scan::MongrelScanExec::new_row_count(total)));
579 }
580
581 let (col_ids, scan_schema): (Vec<u16>, SchemaRef) = match projection {
586 Some(p) if !p.is_empty() => {
587 let ids = p.iter().map(|&idx| schema_ref.columns[idx].id).collect();
588 let fields: Vec<arrow::datatypes::Field> = p
589 .iter()
590 .map(|&idx| self.schema.field(idx).clone())
591 .collect();
592 (ids, Arc::new(arrow::datatypes::Schema::new(fields)))
593 }
594 _ => (
595 schema_ref.columns.iter().map(|c| c.id).collect(),
596 self.schema.clone(),
597 ),
598 };
599
600 let mut proj_pairs: Vec<(u16, mongreldb_core::schema::TypeId)> =
602 Vec::with_capacity(col_ids.len());
603 let mut types: Vec<mongreldb_core::schema::TypeId> = Vec::with_capacity(col_ids.len());
604 for cid in &col_ids {
605 let ty = schema_ref
606 .columns
607 .iter()
608 .find(|c| c.id == *cid)
609 .map(|c| c.ty.clone())
610 .ok_or_else(|| {
611 DataFusionError::External(Box::new(MongrelQueryError::Arrow(format!(
612 "unknown column {cid}"
613 ))))
614 })?;
615 proj_pairs.push((*cid, ty.clone()));
616 types.push(ty);
617 }
618
619 let col_stats_map = if !historical && translated.is_empty() {
623 db.exact_column_stats(snap, &col_ids).map_err(core_err)?
624 } else {
625 None
626 };
627 let column_stats: Vec<datafusion::physical_plan::ColumnStatistics> = col_ids
628 .iter()
629 .map(|cid| scan::to_col_statistics(col_stats_map.as_ref().and_then(|m| m.get(cid))))
630 .collect();
631
632 if !historical
636 && translated.is_empty()
637 && db.run_count() == 1
638 && db.memtable_is_empty()
639 && db.mutable_run_len() == 0
640 && db.single_run_is_clean()
641 {
642 let shadow = shadow::ArrowShadow::new(db.dir());
643 let run_ids: HashSet<u128> = db.run_ids().into_iter().collect();
644 shadow.sweep(&run_ids);
645 if let Some(&run_id) = run_ids.iter().next() {
646 if let Some(batch) = shadow.try_read(run_id) {
647 if let Some(projected) =
648 project_batch(&batch, &col_ids, &schema_ref, &scan_schema)
649 {
650 mongreldb_core::trace::QueryTrace::record(|t| {
651 t.scan_mode = mongreldb_core::trace::ScanMode::ArrowShadow;
652 });
653 return Ok(Arc::new(scan::MongrelScanExec::new_batch(
654 scan_schema,
655 projected,
656 column_stats,
657 )));
658 }
659 }
660 }
661 }
662
663 let cursor: Option<Box<dyn Cursor>> = db
670 .scan_cursor(snap, proj_pairs, &translated)
671 .map_err(core_err)?;
672 if let Some(cursor) = cursor {
673 let num_rows = cursor.remaining_rows();
674 let residual = extract_residual_filter(filters, &col_ids, &schema_ref);
676 return Ok(Arc::new(scan::MongrelScanExec::new_cursor(
677 scan_schema,
678 types,
679 cursor,
680 num_rows,
681 column_stats,
682 residual,
683 )));
684 }
685
686 let mut cols = if !translated.is_empty() {
689 match db
690 .query_columns_native_cached(&translated, Some(&col_ids), snap)
691 .map_err(core_err)?
692 {
693 Some(c) => c,
694 None => db
695 .visible_columns_native(snap, Some(&col_ids))
696 .map_err(core_err)?,
697 }
698 } else {
699 db.visible_columns_native(snap, Some(&col_ids))
700 .map_err(core_err)?
701 };
702
703 let mut ordered: Vec<mongreldb_core::columnar::NativeColumn> =
705 Vec::with_capacity(col_ids.len());
706 for cid in &col_ids {
707 let position = cols.iter().position(|(id, _)| id == cid).ok_or_else(|| {
708 DataFusionError::External(Box::new(MongrelQueryError::Arrow(format!(
709 "missing column {cid}"
710 ))))
711 })?;
712 let (_, col) = cols.swap_remove(position);
713 ordered.push(col);
714 }
715 let num_rows = ordered.first().map(|c| c.len()).unwrap_or(0);
716
717 let shadow_write: Option<(
719 std::path::PathBuf,
720 u128,
721 Vec<arrow::array::ArrayRef>,
722 SchemaRef,
723 )> = if !historical
724 && translated.is_empty()
725 && db.run_count() == 1
726 && db.memtable_is_empty()
727 && db.mutable_run_len() == 0
728 && db.single_run_is_clean()
729 {
730 let all_schema_ids: Vec<u16> = schema_ref.columns.iter().map(|c| c.id).collect();
731 if col_ids == all_schema_ids {
732 let dir = db.dir().to_path_buf();
733 let run_id = db.run_ids().first().copied();
734 run_id.map(|rid| {
735 let arrays = ordered
736 .iter()
737 .zip(types.iter())
738 .map(|(col, ty)| arrow_conv::native_to_array(ty.clone(), col))
739 .collect::<Result<_>>()
740 .unwrap_or_default();
741 (dir, rid, arrays, scan_schema.clone())
742 })
743 } else {
744 None
745 }
746 } else {
747 None
748 };
749
750 drop(db);
751
752 if let Some((dir, run_id, arrays, schema)) = shadow_write {
755 if let Ok(batch) = RecordBatch::try_new(schema, arrays) {
756 shadow::ArrowShadow::new(&dir).write(run_id, &batch);
757 }
758 }
759
760 mongreldb_core::trace::QueryTrace::record(|t| {
761 t.scan_mode = mongreldb_core::trace::ScanMode::Materialized;
762 t.row_materialized = true;
763 });
764 Ok(Arc::new(scan::MongrelScanExec::new(
765 scan_schema,
766 ordered,
767 types,
768 num_rows,
769 column_stats,
770 )))
771 }
772}
773
774fn project_batch(
778 batch: &RecordBatch,
779 col_ids: &[u16],
780 schema_ref: &mongreldb_core::schema::Schema,
781 scan_schema: &arrow::datatypes::SchemaRef,
782) -> Option<RecordBatch> {
783 let arrays: Vec<arrow::array::ArrayRef> = col_ids
785 .iter()
786 .map(|cid| {
787 let name = schema_ref
789 .columns
790 .iter()
791 .find(|c| c.id == *cid)
792 .map(|c| c.name.as_str())?;
793 let idx = batch.schema().index_of(name).ok()?;
795 Some(batch.column(idx).clone())
796 })
797 .collect::<Option<Vec<_>>>()?;
798 RecordBatch::try_new(scan_schema.clone(), arrays).ok()
799}
800
801pub(crate) fn translate_filter(
819 expr: &Expr,
820 schema: &mongreldb_core::Schema,
821) -> Option<mongreldb_core::Condition> {
822 use datafusion::common::ScalarValue;
823 use datafusion::logical_expr::{Between, BinaryExpr, Like, Operator};
824 use mongreldb_core::{ColumnFlags, Condition, IndexKind, TypeId, Value};
825
826 let int_val = |s: &ScalarValue| match s {
830 ScalarValue::Int8(Some(v)) => Some(*v as i64),
831 ScalarValue::Int16(Some(v)) => Some(*v as i64),
832 ScalarValue::Int32(Some(v)) => Some(*v as i64),
833 ScalarValue::Int64(Some(v)) => Some(*v),
834 ScalarValue::UInt8(Some(v)) => Some(*v as i64),
835 ScalarValue::UInt16(Some(v)) => Some(*v as i64),
836 ScalarValue::UInt32(Some(v)) => Some(*v as i64),
837 ScalarValue::UInt64(Some(v)) => Some(*v as i64),
838 ScalarValue::Date32(Some(v)) => Some(*v as i64),
839 ScalarValue::TimestampSecond(Some(v), _) => Some(*v),
840 ScalarValue::TimestampMillisecond(Some(v), _) => Some(*v),
841 ScalarValue::TimestampMicrosecond(Some(v), _) => Some(*v),
842 ScalarValue::TimestampNanosecond(Some(v), _) => Some(*v),
843 _ => None,
844 };
845 let float_val = |s: &ScalarValue| match s {
846 ScalarValue::Float32(Some(f)) => Some(*f as f64),
847 ScalarValue::Float64(Some(f)) => Some(*f),
848 _ => None,
849 };
850 let bytes_val = |s: &ScalarValue| match s {
851 ScalarValue::Utf8(Some(s)) => Some(s.as_bytes().to_vec()),
852 _ => None,
853 };
854 let _ = bytes_val; let val = |s: &ScalarValue| -> Option<Value> {
857 if let Some(i) = int_val(s) {
861 return Some(Value::Int64(i));
862 }
863 match s {
864 ScalarValue::Utf8(Some(s)) => Some(Value::Bytes(s.as_bytes().to_vec())),
865 ScalarValue::Float32(Some(f)) => Some(Value::Float64(*f as f64)),
866 ScalarValue::Float64(Some(f)) => Some(Value::Float64(*f)),
867 ScalarValue::Boolean(Some(b)) => Some(Value::Bool(*b)),
868 _ => None,
869 }
870 };
871
872 let col_def = |name: &str| schema.columns.iter().find(|c| c.name == name);
873 let has_fm = |cid: u16| {
874 schema
875 .indexes
876 .iter()
877 .any(|i| i.column_id == cid && i.kind == IndexKind::FmIndex)
878 };
879 let has_bitmap = |cid: u16| {
880 schema
881 .indexes
882 .iter()
883 .any(|i| i.column_id == cid && i.kind == IndexKind::Bitmap)
884 };
885
886 match expr {
887 Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
890 if *op == Operator::Or {
892 return try_or_as_bitmap_in(expr, schema);
893 }
894 let left = peel_cast(left);
896 let right = peel_cast(right);
897 let (col_name, scalar, flipped) = match (left.as_ref(), right.as_ref()) {
898 (Expr::Column(c), Expr::Literal(s, _)) => (&c.name, s, false),
899 (Expr::Literal(s, _), Expr::Column(c)) => (&c.name, s, true),
900 _ => return None,
901 };
902 let op = if flipped { flip_op(*op)? } else { *op };
903 let cdef = col_def(col_name)?;
904
905 if op == Operator::Eq {
907 let v = val(scalar)?;
908 if has_bitmap(cdef.id) {
909 return Some(Condition::BitmapEq {
910 column_id: cdef.id,
911 value: v.encode_key(),
912 });
913 }
914 if cdef.flags.contains(ColumnFlags::PRIMARY_KEY) {
915 return Some(Condition::Pk(v.encode_key()));
916 }
917 return None;
918 }
919
920 match cdef.ty {
923 TypeId::Int8
924 | TypeId::Int16
925 | TypeId::Int32
926 | TypeId::Int64
927 | TypeId::UInt8
928 | TypeId::UInt16
929 | TypeId::UInt32
930 | TypeId::UInt64
931 | TypeId::TimestampNanos
932 | TypeId::Date32 => {
933 let v = int_val(scalar)?;
934 let (lo, hi) = int_bounds(op, v)?;
935 Some(Condition::Range {
936 column_id: cdef.id,
937 lo,
938 hi,
939 })
940 }
941 TypeId::Float32 | TypeId::Float64 => {
942 let v = float_val(scalar)?;
943 let (lo, lo_inc, hi, hi_inc) = float_bounds(op, v)?;
944 Some(Condition::RangeF64 {
945 column_id: cdef.id,
946 lo,
947 lo_inclusive: lo_inc,
948 hi,
949 hi_inclusive: hi_inc,
950 })
951 }
952 _ => None,
953 }
954 }
955
956 Expr::Between(Between {
958 expr,
959 negated,
960 low,
961 high,
962 }) => {
963 if *negated {
964 return None;
965 }
966 let Expr::Column(c) = expr.as_ref() else {
967 return None;
968 };
969 let cdef = col_def(&c.name)?;
970 let (lo_s, hi_s) = match (low.as_ref(), high.as_ref()) {
971 (Expr::Literal(lo, _), Expr::Literal(hi, _)) => (lo, hi),
972 _ => return None,
973 };
974 match cdef.ty {
975 TypeId::Int8
976 | TypeId::Int16
977 | TypeId::Int32
978 | TypeId::Int64
979 | TypeId::UInt8
980 | TypeId::UInt16
981 | TypeId::UInt32
982 | TypeId::UInt64
983 | TypeId::TimestampNanos
984 | TypeId::Date32 => {
985 let (Some(lo), Some(hi)) = (int_val(lo_s), int_val(hi_s)) else {
986 return None;
987 };
988 Some(Condition::Range {
989 column_id: cdef.id,
990 lo,
991 hi,
992 })
993 }
994 TypeId::Float32 | TypeId::Float64 => {
995 let (Some(lo), Some(hi)) = (float_val(lo_s), float_val(hi_s)) else {
996 return None;
997 };
998 Some(Condition::RangeF64 {
999 column_id: cdef.id,
1000 lo,
1001 lo_inclusive: true,
1002 hi,
1003 hi_inclusive: true,
1004 })
1005 }
1006 _ => None,
1007 }
1008 }
1009
1010 Expr::Like(Like {
1012 negated,
1013 expr,
1014 pattern,
1015 ..
1016 }) => {
1017 if *negated {
1018 return None;
1019 }
1020 let Expr::Column(c) = expr.as_ref() else {
1021 return None;
1022 };
1023 let Expr::Literal(ScalarValue::Utf8(Some(pat)), _) = pattern.as_ref() else {
1024 return None;
1025 };
1026 let cdef = col_def(&c.name)?;
1027 if has_bitmap(cdef.id) {
1031 if let Some(prefix) = anchored_like_prefix(pat) {
1032 return Some(Condition::BytesPrefix {
1033 column_id: cdef.id,
1034 prefix: mongreldb_core::Value::Bytes(prefix.as_bytes().to_vec())
1035 .encode_key(),
1036 });
1037 }
1038 }
1039 if !has_fm(cdef.id) {
1040 return None;
1041 }
1042 let segments: Vec<Vec<u8>> = pat
1046 .split(['%', '_'])
1047 .filter(|s| s.len() >= 3)
1048 .map(|s| s.as_bytes().to_vec())
1049 .collect();
1050 match segments.len() {
1051 0 => longest_like_segment(pat).map(|seg| Condition::FmContains {
1052 column_id: cdef.id,
1053 pattern: seg,
1054 }),
1055 1 => segments
1056 .into_iter()
1057 .next()
1058 .map(|pattern| Condition::FmContains {
1059 column_id: cdef.id,
1060 pattern,
1061 }),
1062 _ => Some(Condition::FmContainsAll {
1063 column_id: cdef.id,
1064 patterns: segments,
1065 }),
1066 }
1067 }
1068
1069 Expr::InList(il) if !il.negated => {
1073 let Expr::Column(c) = il.expr.as_ref() else {
1074 return None;
1075 };
1076 let cdef = col_def(&c.name)?;
1077 if !has_bitmap(cdef.id) {
1078 return None;
1079 }
1080 let values: Vec<Vec<u8>> = il
1081 .list
1082 .iter()
1083 .filter_map(|e| match e {
1084 Expr::Literal(s, _) => val(s).map(|v| v.encode_key()),
1085 _ => None,
1086 })
1087 .collect();
1088 if values.is_empty() || values.len() != il.list.len() {
1089 return None;
1090 }
1091 Some(Condition::BitmapIn {
1092 column_id: cdef.id,
1093 values,
1094 })
1095 }
1096
1097 Expr::IsNull(inner) => {
1099 let col_name = match inner.as_ref() {
1100 Expr::Column(c) => &c.name,
1101 _ => return None,
1102 };
1103 let cdef = col_def(col_name)?;
1104 Some(Condition::IsNull { column_id: cdef.id })
1105 }
1106
1107 Expr::IsNotNull(inner) => {
1109 let col_name = match inner.as_ref() {
1110 Expr::Column(c) => &c.name,
1111 _ => return None,
1112 };
1113 let cdef = col_def(col_name)?;
1114 Some(Condition::IsNotNull { column_id: cdef.id })
1115 }
1116
1117 _ => None,
1118 }
1119}
1120
1121pub(crate) fn extract_residual_filter(
1125 filters: &[Expr],
1126 col_ids: &[u16],
1127 schema: &mongreldb_core::Schema,
1128) -> Option<std::sync::Arc<scan::ResidualFilter>> {
1129 use datafusion::common::ScalarValue;
1130 use datafusion::logical_expr::Like;
1131 for f in filters {
1132 if let Expr::Like(Like {
1133 negated: false,
1134 expr,
1135 pattern,
1136 ..
1137 }) = f
1138 {
1139 let Expr::Column(c) = expr.as_ref() else {
1140 continue;
1141 };
1142 let Expr::Literal(ScalarValue::Utf8(Some(pat)), _) = pattern.as_ref() else {
1143 continue;
1144 };
1145 let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
1146 let col_idx = col_ids.iter().position(|&id| id == cdef.id)?;
1147 return Some(std::sync::Arc::new(scan::ResidualFilter::new(
1148 col_idx,
1149 pat.as_bytes().to_vec(),
1150 )));
1151 }
1152 }
1153 None
1154}
1155
1156pub(crate) fn translate_ann_search(
1162 expr: &Expr,
1163 schema: &mongreldb_core::Schema,
1164) -> Option<mongreldb_core::Condition> {
1165 use datafusion::common::ScalarValue;
1166 use mongreldb_core::Condition;
1167
1168 let Expr::ScalarFunction(sf) = expr else {
1169 return None;
1170 };
1171 if !sf.func.name().eq_ignore_ascii_case("ann_search") || sf.args.len() != 3 {
1172 return None;
1173 }
1174 let (Expr::Column(c), query_expr, k_expr) = (&sf.args[0], &sf.args[1], &sf.args[2]) else {
1175 return None;
1176 };
1177 let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
1178 let json = match query_expr {
1179 Expr::Literal(ScalarValue::Utf8(Some(s)), _) => s.as_str(),
1180 _ => return None,
1181 };
1182 let k: usize = match k_expr {
1183 Expr::Literal(scalar, _) => match scalar {
1184 ScalarValue::Int64(Some(k)) => usize::try_from(*k).ok()?,
1185 ScalarValue::UInt64(Some(k)) => usize::try_from(*k).ok()?,
1186 ScalarValue::Int32(Some(k)) => usize::try_from(*k).ok()?,
1187 _ => return None,
1188 },
1189 _ => return None,
1190 };
1191 let query: Vec<f32> = serde_json::from_str(json).ok()?;
1192 Some(Condition::Ann {
1193 column_id: cdef.id,
1194 query,
1195 k,
1196 })
1197}
1198
1199pub(crate) fn translate_sparse_match(
1204 expr: &Expr,
1205 schema: &mongreldb_core::Schema,
1206) -> Option<mongreldb_core::Condition> {
1207 use datafusion::common::ScalarValue;
1208 use mongreldb_core::Condition;
1209
1210 let Expr::ScalarFunction(sf) = expr else {
1211 return None;
1212 };
1213 if !sf.func.name().eq_ignore_ascii_case("sparse_match") || sf.args.len() != 3 {
1214 return None;
1215 }
1216 let (Expr::Column(c), query_expr, k_expr) = (&sf.args[0], &sf.args[1], &sf.args[2]) else {
1217 return None;
1218 };
1219 let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
1220 let json = match query_expr {
1221 Expr::Literal(ScalarValue::Utf8(Some(s)), _) => s.as_str(),
1222 _ => return None,
1223 };
1224 let k: usize = match k_expr {
1225 Expr::Literal(scalar, _) => match scalar {
1226 ScalarValue::Int64(Some(k)) => usize::try_from(*k).ok()?,
1227 ScalarValue::UInt64(Some(k)) => usize::try_from(*k).ok()?,
1228 ScalarValue::Int32(Some(k)) => usize::try_from(*k).ok()?,
1229 _ => return None,
1230 },
1231 _ => return None,
1232 };
1233 let query: Vec<(u32, f32)> = serde_json::from_str(json).ok()?;
1234 Some(Condition::SparseMatch {
1235 column_id: cdef.id,
1236 query,
1237 k,
1238 })
1239}
1240
1241fn flip_op(op: datafusion::logical_expr::Operator) -> Option<datafusion::logical_expr::Operator> {
1243 use datafusion::logical_expr::Operator;
1244 Some(match op {
1245 Operator::Eq => Operator::Eq,
1246 Operator::Lt => Operator::Gt,
1247 Operator::Gt => Operator::Lt,
1248 Operator::LtEq => Operator::GtEq,
1249 Operator::GtEq => Operator::LtEq,
1250 _ => return None,
1251 })
1252}
1253
1254fn int_bounds(op: datafusion::logical_expr::Operator, v: i64) -> Option<(i64, i64)> {
1257 use datafusion::logical_expr::Operator;
1258 Some(match op {
1259 Operator::Gt => (v.saturating_add(1), i64::MAX),
1260 Operator::GtEq => (v, i64::MAX),
1261 Operator::Lt => (i64::MIN, v.saturating_sub(1)),
1262 Operator::LtEq => (i64::MIN, v),
1263 _ => return None,
1264 })
1265}
1266
1267fn float_bounds(op: datafusion::logical_expr::Operator, v: f64) -> Option<(f64, bool, f64, bool)> {
1269 use datafusion::logical_expr::Operator;
1270 Some(match op {
1271 Operator::Gt => (v, false, f64::INFINITY, false),
1272 Operator::GtEq => (v, true, f64::INFINITY, false),
1273 Operator::Lt => (f64::NEG_INFINITY, false, v, false),
1274 Operator::LtEq => (f64::NEG_INFINITY, false, v, true),
1275 _ => return None,
1276 })
1277}
1278
1279fn longest_like_segment(pat: &str) -> Option<Vec<u8>> {
1286 pat.split(['%', '_'])
1287 .map(|s| s.as_bytes())
1288 .max_by_key(|s| s.len())
1289 .filter(|s| !s.is_empty())
1290 .map(|s| s.to_vec())
1291}
1292
1293fn anchored_like_prefix(pat: &str) -> Option<&str> {
1298 let rest = pat.strip_suffix('%')?;
1299 if rest.is_empty() || rest.contains(['%', '_']) {
1300 return None;
1301 }
1302 Some(rest)
1303}
1304
1305fn peel_cast(expr: &Expr) -> std::borrow::Cow<'_, Expr> {
1309 match expr {
1310 Expr::Cast(datafusion::logical_expr::Cast { expr, .. }) => std::borrow::Cow::Borrowed(expr),
1311 _ => std::borrow::Cow::Borrowed(expr),
1312 }
1313}
1314
1315fn try_or_as_bitmap_in(
1320 expr: &Expr,
1321 schema: &mongreldb_core::Schema,
1322) -> Option<mongreldb_core::Condition> {
1323 use datafusion::logical_expr::{BinaryExpr, Operator};
1324 let mut values: Vec<Vec<u8>> = Vec::new();
1325 let mut target_col: Option<u16> = None;
1326 let mut stack = vec![expr];
1327 while let Some(e) = stack.pop() {
1328 match e {
1329 Expr::BinaryExpr(BinaryExpr {
1330 left,
1331 op: Operator::Or,
1332 right,
1333 }) => {
1334 stack.push(left);
1335 stack.push(right);
1336 }
1337 Expr::BinaryExpr(BinaryExpr {
1338 left,
1339 op: Operator::Eq,
1340 right,
1341 }) => {
1342 let (col_name, scalar) = match (left.as_ref(), right.as_ref()) {
1343 (Expr::Column(c), Expr::Literal(s, _)) => (&c.name, s),
1344 (Expr::Literal(s, _), Expr::Column(c)) => (&c.name, s),
1345 _ => return None,
1346 };
1347 let cdef = schema.columns.iter().find(|c| &c.name == col_name)?;
1348 if !schema
1349 .indexes
1350 .iter()
1351 .any(|i| i.column_id == cdef.id && i.kind == mongreldb_core::IndexKind::Bitmap)
1352 {
1353 return None;
1354 }
1355 match target_col {
1356 None => target_col = Some(cdef.id),
1357 Some(id) if id != cdef.id => return None,
1358 _ => {}
1359 }
1360 let v = match scalar {
1361 datafusion::common::ScalarValue::Int64(Some(v)) => {
1362 mongreldb_core::Value::Int64(*v)
1363 }
1364 datafusion::common::ScalarValue::Utf8(Some(s)) => {
1365 mongreldb_core::Value::Bytes(s.as_bytes().to_vec())
1366 }
1367 datafusion::common::ScalarValue::Float64(Some(f)) => {
1368 mongreldb_core::Value::Float64(*f)
1369 }
1370 datafusion::common::ScalarValue::Boolean(Some(b)) => {
1371 mongreldb_core::Value::Bool(*b)
1372 }
1373 _ => return None,
1374 };
1375 values.push(v.encode_key());
1376 }
1377 _ => return None,
1378 }
1379 }
1380 let col_id = target_col?;
1381 if values.is_empty() {
1382 return None;
1383 }
1384 Some(mongreldb_core::Condition::BitmapIn {
1385 column_id: col_id,
1386 values,
1387 })
1388}
1389
1390fn sp_ident_name(expr: &sqlparser::ast::Expr) -> Option<&str> {
1397 use sqlparser::ast::Expr;
1398 match expr {
1399 Expr::Identifier(ident) => Some(ident.value.as_str()),
1400 Expr::CompoundIdentifier(idents) => idents.last().map(|i| i.value.as_str()),
1401 _ => None,
1402 }
1403}
1404
1405fn sp_literal(expr: &sqlparser::ast::Expr) -> Option<mongreldb_core::Value> {
1408 use sqlparser::ast::Expr;
1409 let v = match expr {
1410 Expr::Value(v) => v,
1411 _ => return None,
1412 };
1413 use sqlparser::ast::Value as SpValue;
1414 match &v.value {
1415 SpValue::Number(s, _) => s
1416 .parse::<i64>()
1417 .map(mongreldb_core::Value::Int64)
1418 .or_else(|_| s.parse::<f64>().map(mongreldb_core::Value::Float64))
1419 .ok(),
1420 SpValue::SingleQuotedString(s) => Some(mongreldb_core::Value::Bytes(s.as_bytes().to_vec())),
1421 SpValue::Boolean(b) => Some(mongreldb_core::Value::Bool(*b)),
1422 _ => None,
1423 }
1424}
1425
1426fn is_int_ty(ty: mongreldb_core::schema::TypeId) -> bool {
1427 use mongreldb_core::schema::TypeId::*;
1428 matches!(
1429 ty,
1430 Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 | TimestampNanos | Date32
1431 )
1432}
1433
1434fn is_float_ty(ty: mongreldb_core::schema::TypeId) -> bool {
1435 matches!(
1436 ty,
1437 mongreldb_core::schema::TypeId::Float32 | mongreldb_core::schema::TypeId::Float64
1438 )
1439}
1440
1441fn translate_sqlparser_predicate(
1444 expr: &sqlparser::ast::Expr,
1445 schema: &mongreldb_core::Schema,
1446) -> Option<mongreldb_core::Condition> {
1447 use mongreldb_core::{schema::ColumnFlags, Condition, IndexKind, Value};
1448 use sqlparser::ast::{BinaryOperator, Expr};
1449
1450 let col_def = |name: &str| schema.columns.iter().find(|c| c.name == name);
1451 let has_bitmap = |cid: u16| {
1452 schema
1453 .indexes
1454 .iter()
1455 .any(|i| i.column_id == cid && i.kind == IndexKind::Bitmap)
1456 };
1457
1458 match expr {
1459 Expr::BinaryOp {
1461 left,
1462 op: BinaryOperator::Or,
1463 right,
1464 } => {
1465 let mut values: Vec<Vec<u8>> = Vec::new();
1466 let mut target: Option<u16> = None;
1467 let mut stack: Vec<&Expr> = vec![left.as_ref(), right.as_ref()];
1468 while let Some(e) = stack.pop() {
1469 match e {
1470 Expr::BinaryOp {
1471 left,
1472 op: BinaryOperator::Or,
1473 right,
1474 } => {
1475 stack.push(left.as_ref());
1476 stack.push(right.as_ref());
1477 }
1478 Expr::BinaryOp {
1479 left,
1480 op: BinaryOperator::Eq,
1481 right,
1482 } => {
1483 let (name, lit) = match (left.as_ref(), right.as_ref()) {
1484 (l, r) if sp_ident_name(l).is_some() && sp_literal(r).is_some() => {
1485 (l, r)
1486 }
1487 (l, r) if sp_ident_name(r).is_some() && sp_literal(l).is_some() => {
1488 (r, l)
1489 }
1490 _ => return None,
1491 };
1492 let cdef = col_def(sp_ident_name(name)?)?;
1493 if !has_bitmap(cdef.id) {
1494 return None;
1495 }
1496 match target {
1497 None => target = Some(cdef.id),
1498 Some(id) if id != cdef.id => return None,
1499 _ => {}
1500 }
1501 values.push(sp_literal(lit)?.encode_key());
1502 }
1503 _ => return None,
1504 }
1505 }
1506 let cid = target?;
1507 (!values.is_empty()).then_some(Condition::BitmapIn {
1508 column_id: cid,
1509 values,
1510 })
1511 }
1512 Expr::BinaryOp { left, op, right } => {
1514 let flipped;
1515 let (col_expr, lit_expr) = match (
1516 sp_ident_name(left),
1517 sp_literal(right),
1518 sp_ident_name(right),
1519 sp_literal(left),
1520 ) {
1521 (Some(_), Some(_), _, _) => {
1522 flipped = false;
1523 (left.as_ref(), right.as_ref())
1524 }
1525 (_, _, Some(_), Some(_)) => {
1526 flipped = true;
1527 (right.as_ref(), left.as_ref())
1528 }
1529 _ => return None,
1530 };
1531 let name = sp_ident_name(col_expr)?;
1532 let cdef = col_def(name)?;
1533 let v = sp_literal(lit_expr)?;
1534 use sqlparser::ast::BinaryOperator::*;
1535 match op {
1538 Eq => {
1539 if has_bitmap(cdef.id) {
1540 Some(Condition::BitmapEq {
1541 column_id: cdef.id,
1542 value: v.encode_key(),
1543 })
1544 } else if cdef.flags.contains(ColumnFlags::PRIMARY_KEY) {
1545 Some(Condition::Pk(v.encode_key()))
1546 } else {
1547 None
1548 }
1549 }
1550 Lt | LtEq | Gt | GtEq if is_int_ty(cdef.ty.clone()) => {
1551 let n = match v {
1552 Value::Int64(n) => n,
1553 _ => return None,
1554 };
1555 let (lo, hi) = match (flipped, op) {
1557 (false, Lt) | (true, Gt) => (i64::MIN, n.saturating_sub(1)),
1558 (false, Gt) | (true, Lt) => (n.saturating_add(1), i64::MAX),
1559 (false, LtEq) | (true, GtEq) => (i64::MIN, n),
1560 (false, GtEq) | (true, LtEq) => (n, i64::MAX),
1561 _ => (i64::MIN, i64::MAX),
1562 };
1563 Some(Condition::Range {
1564 column_id: cdef.id,
1565 lo,
1566 hi,
1567 })
1568 }
1569 Lt | LtEq | Gt | GtEq if is_float_ty(cdef.ty.clone()) => {
1570 let f = match v {
1571 Value::Float64(f) => f,
1572 _ => return None,
1573 };
1574 let (lo, li, hi, hi_i) = match (flipped, op) {
1575 (false, Lt) | (true, Gt) => (f64::NEG_INFINITY, true, f, false),
1576 (false, Gt) | (true, Lt) => (f, false, f64::INFINITY, true),
1577 (false, LtEq) | (true, GtEq) => (f64::NEG_INFINITY, true, f, true),
1578 (false, GtEq) | (true, LtEq) => (f, true, f64::INFINITY, true),
1579 _ => (f64::NEG_INFINITY, true, f64::INFINITY, true),
1580 };
1581 Some(Condition::RangeF64 {
1582 column_id: cdef.id,
1583 lo,
1584 lo_inclusive: li,
1585 hi,
1586 hi_inclusive: hi_i,
1587 })
1588 }
1589 _ => None,
1590 }
1591 }
1592 Expr::Between {
1593 expr,
1594 negated,
1595 low,
1596 high,
1597 } if !negated => {
1598 let name = sp_ident_name(expr)?;
1599 let cdef = col_def(name)?;
1600 if is_int_ty(cdef.ty.clone()) {
1601 let lo = match sp_literal(low)? {
1602 Value::Int64(n) => n,
1603 _ => return None,
1604 };
1605 let hi = match sp_literal(high)? {
1606 Value::Int64(n) => n,
1607 _ => return None,
1608 };
1609 Some(Condition::Range {
1610 column_id: cdef.id,
1611 lo,
1612 hi,
1613 })
1614 } else if is_float_ty(cdef.ty.clone()) {
1615 let lo = match sp_literal(low)? {
1616 Value::Float64(f) => f,
1617 _ => return None,
1618 };
1619 let hi = match sp_literal(high)? {
1620 Value::Float64(f) => f,
1621 _ => return None,
1622 };
1623 Some(Condition::RangeF64 {
1624 column_id: cdef.id,
1625 lo,
1626 lo_inclusive: true,
1627 hi,
1628 hi_inclusive: true,
1629 })
1630 } else {
1631 None
1632 }
1633 }
1634 Expr::InList {
1635 expr,
1636 list,
1637 negated,
1638 } if !negated => {
1639 let name = sp_ident_name(expr)?;
1640 let cdef = col_def(name)?;
1641 if !has_bitmap(cdef.id) {
1642 return None;
1643 }
1644 let values: Vec<Vec<u8>> = list
1645 .iter()
1646 .map(|e| sp_literal(e).map(|v| v.encode_key()))
1647 .collect::<Option<_>>()?;
1648 (!values.is_empty()).then_some(Condition::BitmapIn {
1649 column_id: cdef.id,
1650 values,
1651 })
1652 }
1653 Expr::IsNull(inner) => {
1656 let cid = col_def(sp_ident_name(inner)?)?.id;
1657 Some(Condition::IsNull { column_id: cid })
1658 }
1659 Expr::IsNotNull(inner) => {
1660 let cid = col_def(sp_ident_name(inner)?)?.id;
1661 Some(Condition::IsNotNull { column_id: cid })
1662 }
1663 _ => None,
1664 }
1665}
1666
1667fn translate_sqlparser_filter(
1670 expr: &sqlparser::ast::Expr,
1671 schema: &mongreldb_core::Schema,
1672) -> Option<Vec<mongreldb_core::Condition>> {
1673 use sqlparser::ast::{BinaryOperator, Expr};
1674 let mut out = Vec::new();
1675 let mut stack = vec![expr];
1676 while let Some(e) = stack.pop() {
1677 match e {
1678 Expr::BinaryOp {
1679 left,
1680 op: BinaryOperator::And,
1681 right,
1682 } => {
1683 stack.push(left.as_ref());
1684 stack.push(right.as_ref());
1685 }
1686 other => out.push(translate_sqlparser_predicate(other, schema)?),
1687 }
1688 }
1689 Some(out)
1690}
1691
1692pub struct MongrelSession {
1696 ctx: SessionContext,
1697 db: Option<TableHandle>,
1698 database: Option<Arc<Database>>,
1702 principal: Option<mongreldb_core::Principal>,
1703 principal_catalog_bound: bool,
1704 cache: ResultCache,
1705 plan_cache: parking_lot::Mutex<BoundedLru<String, datafusion::logical_expr::LogicalPlan>>,
1707 tables: scored_sql::TableMap,
1709 views: parking_lot::Mutex<HashMap<String, ViewDef>>,
1713 attached_databases: parking_lot::Mutex<HashMap<String, Arc<Database>>>,
1717 transaction: Arc<parking_lot::Mutex<SqlTransactionState>>,
1720 sql_fn_state: Arc<extended_sql_functions::ExtendedSqlState>,
1722 external_modules: Arc<ExternalModuleRegistry>,
1725 scored_runtime: Arc<scored_sql::ScoredRuntime>,
1726 prepared_names: parking_lot::Mutex<std::collections::HashSet<String>>,
1730 query_registry: Arc<SqlQueryRegistry>,
1731 execution_gate: Arc<tokio::sync::Semaphore>,
1732 statement_timeout: parking_lot::RwLock<Option<std::time::Duration>>,
1733 test_hook: parking_lot::Mutex<Option<SqlTestHook>>,
1734}
1735
1736#[derive(Default)]
1737struct SqlTransactionState {
1738 staged_ops: Option<commands::PendingSqlOps>,
1739 aborted: bool,
1740 savepoints: Vec<(String, commands::PendingSqlOpsCheckpoint)>,
1741}
1742
1743type CacheKey = (String, u64);
1745type ResultCache = parking_lot::Mutex<ResultCacheStore>;
1746
1747const SESSION_CACHE_MAX_ENTRIES: usize = 1024;
1748const SESSION_CACHE_MAX_BYTES: usize = 64 * 1024 * 1024;
1749const DEFAULT_COLLECTION_MAX_ROWS: usize = 1_000_000;
1750const DEFAULT_COLLECTION_MAX_BYTES: usize = 64 * 1024 * 1024;
1751
1752fn core_value_memory_bytes(value: &mongreldb_core::Value) -> usize {
1753 use mongreldb_core::Value;
1754 match value {
1755 Value::Null => 0,
1756 Value::Bool(_) => 1,
1757 Value::Int64(_) | Value::Float64(_) => 8,
1758 Value::Bytes(value) | Value::Json(value) => value.len(),
1759 Value::Embedding(value) => value.len().saturating_mul(std::mem::size_of::<f32>()),
1760 Value::Decimal(_) | Value::Uuid(_) => 16,
1761 Value::Interval { .. } => 20,
1762 }
1763}
1764
1765fn core_row_memory_bytes(row: &mongreldb_core::Row) -> usize {
1766 let bucket_bytes = std::mem::size_of::<(u16, mongreldb_core::Value)>()
1767 .saturating_add(2 * std::mem::size_of::<usize>());
1768 row.columns
1769 .capacity()
1770 .saturating_mul(bucket_bytes)
1771 .saturating_add(
1772 row.columns
1773 .values()
1774 .map(core_value_memory_bytes)
1775 .fold(0_usize, usize::saturating_add),
1776 )
1777}
1778
1779fn enforce_source_rows_limit(
1780 rows: &[mongreldb_core::Row],
1781 query: Option<&RegisteredSqlQuery>,
1782) -> DFResult<()> {
1783 let limits = current_collection_limits();
1784 if rows.len() > limits.max_rows {
1785 let message = format!(
1786 "SQL source row limit exceeded ({} > {})",
1787 rows.len(),
1788 limits.max_rows
1789 );
1790 return Err(DataFusionError::External(Box::new(match query {
1791 Some(query) => query.result_limit_error(message),
1792 None => MongrelQueryError::Core(mongreldb_core::MongrelError::ResourceLimitExceeded {
1793 resource: "SQL source rows",
1794 requested: rows.len(),
1795 limit: limits.max_rows,
1796 }),
1797 })));
1798 }
1799 let mut bytes = 0_usize;
1800 for (index, row) in rows.iter().enumerate() {
1801 if index % 256 == 0 {
1802 query
1803 .map(RegisteredSqlQuery::checkpoint)
1804 .transpose()
1805 .map_err(|error| DataFusionError::External(Box::new(error)))?;
1806 }
1807 bytes = bytes.saturating_add(core_row_memory_bytes(row));
1808 if bytes > limits.max_bytes {
1809 let message = format!(
1810 "SQL source row memory limit exceeded ({} > {} bytes)",
1811 bytes, limits.max_bytes
1812 );
1813 return Err(DataFusionError::External(Box::new(match query {
1814 Some(query) => query.result_limit_error(message),
1815 None => {
1816 MongrelQueryError::Core(mongreldb_core::MongrelError::ResourceLimitExceeded {
1817 resource: "SQL source row bytes",
1818 requested: bytes,
1819 limit: limits.max_bytes,
1820 })
1821 }
1822 })));
1823 }
1824 }
1825 Ok(())
1826}
1827
1828#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1830pub struct SqlCollectionLimits {
1831 pub max_rows: usize,
1832 pub max_bytes: usize,
1833}
1834
1835impl SqlCollectionLimits {
1836 pub fn new(max_rows: usize, max_bytes: usize) -> Self {
1837 Self {
1838 max_rows: max_rows.max(1),
1839 max_bytes: max_bytes.max(1),
1840 }
1841 }
1842}
1843
1844impl Default for SqlCollectionLimits {
1845 fn default() -> Self {
1846 Self::new(DEFAULT_COLLECTION_MAX_ROWS, DEFAULT_COLLECTION_MAX_BYTES)
1847 }
1848}
1849
1850fn current_collection_limits() -> SqlCollectionLimits {
1851 CURRENT_COLLECTION_LIMITS
1852 .try_with(|limits| *limits)
1853 .unwrap_or_default()
1854}
1855
1856fn record_batch_memory_bytes(batch: &RecordBatch) -> usize {
1857 batch.columns().iter().fold(0_usize, |total, column| {
1858 total.saturating_add(column.get_array_memory_size())
1859 })
1860}
1861
1862fn take_mongrel_query_error(error: DataFusionError) -> Option<MongrelQueryError> {
1863 match error {
1864 DataFusionError::External(error) => error
1865 .downcast::<MongrelQueryError>()
1866 .ok()
1867 .map(|error| *error),
1868 DataFusionError::Context(_, error) | DataFusionError::Diagnostic(_, error) => {
1869 take_mongrel_query_error(*error)
1870 }
1871 DataFusionError::Collection(errors) => {
1872 errors.into_iter().find_map(take_mongrel_query_error)
1873 }
1874 DataFusionError::Shared(error) => Arc::try_unwrap(error)
1875 .ok()
1876 .and_then(take_mongrel_query_error),
1877 _ => None,
1878 }
1879}
1880
1881fn map_datafusion_error(error: DataFusionError) -> MongrelQueryError {
1882 let message = error.to_string();
1883 take_mongrel_query_error(error).unwrap_or(MongrelQueryError::DataFusion(message))
1884}
1885
1886fn enforce_collected_limits(
1887 batches: &[RecordBatch],
1888 query: &RegisteredSqlQuery,
1889 limits: SqlCollectionLimits,
1890) -> Result<()> {
1891 let mut rows = 0_usize;
1892 let mut bytes = 0_usize;
1893 for batch in batches {
1894 query.checkpoint()?;
1895 rows = rows.saturating_add(batch.num_rows());
1896 bytes = bytes.saturating_add(record_batch_memory_bytes(batch));
1897 if rows > limits.max_rows {
1898 return Err(query.result_limit_error(format!(
1899 "SQL collected row limit exceeded ({} > {})",
1900 rows, limits.max_rows
1901 )));
1902 }
1903 if bytes > limits.max_bytes {
1904 return Err(query.result_limit_error(format!(
1905 "SQL collected Arrow memory limit exceeded ({} > {} bytes)",
1906 bytes, limits.max_bytes
1907 )));
1908 }
1909 }
1910 Ok(())
1911}
1912
1913struct ResultCacheStore {
1914 entries: BoundedLru<CacheKey, CachedResult>,
1915 bytes: usize,
1916 max_bytes: usize,
1917}
1918
1919struct CachedResult {
1920 batches: Arc<Vec<RecordBatch>>,
1921 bytes: usize,
1922}
1923
1924impl ResultCacheStore {
1925 fn new(capacity: usize, max_bytes: usize) -> Self {
1926 Self {
1927 entries: BoundedLru::new(capacity),
1928 bytes: 0,
1929 max_bytes: max_bytes.max(1),
1930 }
1931 }
1932
1933 fn get(&mut self, key: &CacheKey) -> Option<&Arc<Vec<RecordBatch>>> {
1934 self.entries.get(key).map(|entry| &entry.batches)
1935 }
1936
1937 fn insert(&mut self, key: CacheKey, value: Arc<Vec<RecordBatch>>) {
1938 let bytes = value.iter().fold(0_usize, |total, batch| {
1939 total.saturating_add(record_batch_memory_bytes(batch))
1940 });
1941 if bytes > self.max_bytes {
1942 return;
1943 }
1944 if let Some((old, _)) = self.entries.entries.remove(&key) {
1945 self.bytes = self.bytes.saturating_sub(old.bytes);
1946 }
1947 while !self.entries.entries.is_empty()
1948 && (self.bytes.saturating_add(bytes) > self.max_bytes
1949 || self.entries.entries.len() >= self.entries.capacity)
1950 {
1951 let Some(oldest) = self
1952 .entries
1953 .entries
1954 .iter()
1955 .min_by_key(|(_, (_, last_used))| *last_used)
1956 .map(|(key, _)| key.clone())
1957 else {
1958 break;
1959 };
1960 if let Some((old, _)) = self.entries.entries.remove(&oldest) {
1961 self.bytes = self.bytes.saturating_sub(old.bytes);
1962 }
1963 }
1964 self.entries.insert(
1965 key,
1966 CachedResult {
1967 batches: value,
1968 bytes,
1969 },
1970 );
1971 self.bytes = self.bytes.saturating_add(bytes);
1972 }
1973
1974 fn clear(&mut self) {
1975 self.entries.clear();
1976 self.bytes = 0;
1977 }
1978
1979 #[cfg(test)]
1980 fn is_empty(&self) -> bool {
1981 self.entries.entries.is_empty()
1982 }
1983}
1984
1985struct BoundedLru<K, V> {
1986 entries: HashMap<K, (V, u64)>,
1987 clock: u64,
1988 capacity: usize,
1989}
1990
1991impl<K: Clone + Eq + Hash, V> BoundedLru<K, V> {
1992 fn new(capacity: usize) -> Self {
1993 Self {
1994 entries: HashMap::new(),
1995 clock: 0,
1996 capacity: capacity.max(1),
1997 }
1998 }
1999
2000 fn get<Q>(&mut self, key: &Q) -> Option<&V>
2001 where
2002 K: Borrow<Q>,
2003 Q: Eq + Hash + ?Sized,
2004 {
2005 self.clock = self.clock.wrapping_add(1);
2006 let (value, last_used) = self.entries.get_mut(key)?;
2007 *last_used = self.clock;
2008 Some(value)
2009 }
2010
2011 fn insert(&mut self, key: K, value: V) {
2012 self.clock = self.clock.wrapping_add(1);
2013 if let Some(entry) = self.entries.get_mut(&key) {
2014 *entry = (value, self.clock);
2015 return;
2016 }
2017 if self.entries.len() >= self.capacity {
2018 if let Some(oldest) = self
2020 .entries
2021 .iter()
2022 .min_by_key(|(_, (_, last_used))| *last_used)
2023 .map(|(key, _)| key.clone())
2024 {
2025 self.entries.remove(&oldest);
2026 }
2027 }
2028 self.entries.insert(key, (value, self.clock));
2029 }
2030
2031 fn clear(&mut self) {
2032 self.entries.clear();
2033 self.clock = 0;
2034 }
2035}
2036
2037fn buffered_stream(batches: Vec<RecordBatch>) -> MongrelRecordBatchStream {
2038 let schema = batches
2039 .first()
2040 .map(RecordBatch::schema)
2041 .unwrap_or_else(|| Arc::new(arrow::datatypes::Schema::empty()));
2042 let stream = futures::stream::iter(batches.into_iter().map(Ok::<RecordBatch, DataFusionError>));
2043 Box::pin(RecordBatchStreamAdapter::new(schema, stream))
2044}
2045
2046pub struct ManagedQueryBatches {
2049 batches: Vec<RecordBatch>,
2050 query: RegisteredSqlQuery,
2051 guard: Option<RegisteredQueryGuard>,
2052 permit: Option<tokio::sync::OwnedSemaphorePermit>,
2053}
2054
2055impl ManagedQueryBatches {
2056 pub fn batches(&self) -> &[RecordBatch] {
2057 &self.batches
2058 }
2059
2060 pub fn query(&self) -> &RegisteredSqlQuery {
2061 &self.query
2062 }
2063
2064 pub fn complete(self) -> Result<()> {
2065 self.try_complete()
2066 }
2067
2068 pub fn try_complete(mut self) -> Result<()> {
2069 let result = if let Some(guard) = self.guard.take() {
2070 guard.try_complete()
2071 } else {
2072 Ok(())
2073 };
2074 self.permit.take();
2075 result
2076 }
2077
2078 pub fn fail_result_limit(mut self) {
2079 if let Some(guard) = self.guard.take() {
2080 guard.fail_result_limit();
2081 }
2082 self.permit.take();
2083 }
2084
2085 pub fn fail_with_error(
2086 mut self,
2087 code: impl Into<String>,
2088 category: QueryTerminalErrorCategory,
2089 ) {
2090 if let Some(guard) = self.guard.take() {
2091 guard.fail_with_error(code, category);
2092 }
2093 self.permit.take();
2094 }
2095
2096 pub fn fail_serialization(mut self) {
2097 if let Some(guard) = self.guard.take() {
2098 guard.fail_serialization();
2099 }
2100 self.permit.take();
2101 }
2102
2103 pub fn fail(mut self) {
2104 if let Some(guard) = self.guard.take() {
2105 guard.fail();
2106 }
2107 self.permit.take();
2108 }
2109}
2110
2111impl Drop for ManagedQueryBatches {
2112 fn drop(&mut self) {
2113 if self.guard.is_some() {
2114 let _ = self
2115 .query
2116 .request_cancel(mongreldb_core::CancellationReason::ClientDisconnected);
2117 if let Some(guard) = self.guard.take() {
2118 guard.fail();
2119 }
2120 }
2121 }
2122}
2123
2124struct ManagedQueryStream {
2125 inner: MongrelRecordBatchStream,
2126 schema: SchemaRef,
2127 query: RegisteredSqlQuery,
2128 guard: Option<RegisteredQueryGuard>,
2129 permit: Option<tokio::sync::OwnedSemaphorePermit>,
2130 deferred_completion: Option<Arc<std::sync::atomic::AtomicU8>>,
2131}
2132
2133const STREAM_COMPLETION_PENDING: u8 = 0;
2134const STREAM_COMPLETION_SUCCEEDED: u8 = 1;
2135const STREAM_COMPLETION_FAILED: u8 = 2;
2136const STREAM_COMPLETION_FAILING: u8 = 3;
2137const STREAM_COMPLETION_COMPLETING: u8 = 4;
2138
2139#[derive(Clone)]
2140pub struct SqlStreamCompletion {
2141 state: Arc<std::sync::atomic::AtomicU8>,
2142 query: RegisteredSqlQuery,
2143}
2144
2145impl SqlStreamCompletion {
2146 pub fn complete(&self) -> Result<()> {
2147 self.try_complete()
2148 }
2149
2150 pub fn try_complete(&self) -> Result<()> {
2151 if self
2152 .state
2153 .compare_exchange(
2154 STREAM_COMPLETION_PENDING,
2155 STREAM_COMPLETION_COMPLETING,
2156 std::sync::atomic::Ordering::AcqRel,
2157 std::sync::atomic::Ordering::Acquire,
2158 )
2159 .is_err()
2160 {
2161 return match self.state.load(std::sync::atomic::Ordering::Acquire) {
2162 STREAM_COMPLETION_SUCCEEDED => Ok(()),
2163 STREAM_COMPLETION_FAILED => Err(MongrelQueryError::InvalidQueryState(
2164 "SQL stream serialization already failed".into(),
2165 )),
2166 _ => Err(MongrelQueryError::InvalidQueryState(
2167 "SQL stream completion is already in progress".into(),
2168 )),
2169 };
2170 }
2171 let result = self.query.try_complete();
2172 self.state.store(
2173 if result.is_ok() {
2174 STREAM_COMPLETION_SUCCEEDED
2175 } else {
2176 STREAM_COMPLETION_FAILED
2177 },
2178 std::sync::atomic::Ordering::Release,
2179 );
2180 result
2181 }
2182
2183 pub fn fail_result_limit(&self) {
2184 if self
2185 .state
2186 .compare_exchange(
2187 STREAM_COMPLETION_PENDING,
2188 STREAM_COMPLETION_FAILING,
2189 std::sync::atomic::Ordering::AcqRel,
2190 std::sync::atomic::Ordering::Acquire,
2191 )
2192 .is_ok()
2193 {
2194 self.query.fail_result_limit();
2195 self.state.store(
2196 STREAM_COMPLETION_FAILED,
2197 std::sync::atomic::Ordering::Release,
2198 );
2199 }
2200 }
2201
2202 pub fn fail_serialization(&self) {
2203 if self
2204 .state
2205 .compare_exchange(
2206 STREAM_COMPLETION_PENDING,
2207 STREAM_COMPLETION_FAILING,
2208 std::sync::atomic::Ordering::AcqRel,
2209 std::sync::atomic::Ordering::Acquire,
2210 )
2211 .is_ok()
2212 {
2213 self.query.fail_serialization();
2214 self.state.store(
2215 STREAM_COMPLETION_FAILED,
2216 std::sync::atomic::Ordering::Release,
2217 );
2218 }
2219 }
2220}
2221
2222struct TransactionStatementGuard {
2223 transaction: Arc<parking_lot::Mutex<SqlTransactionState>>,
2224 query: RegisteredSqlQuery,
2225 checkpoint: Option<(commands::PendingSqlOpsCheckpoint, usize)>,
2226 completed: bool,
2227}
2228
2229impl TransactionStatementGuard {
2230 fn new(session: &MongrelSession, query: &RegisteredSqlQuery, enabled: bool) -> Result<Self> {
2231 let checkpoint = if enabled {
2232 let mut transaction = session.transaction.lock();
2233 let savepoint_len = transaction.savepoints.len();
2234 match transaction.staged_ops.as_mut() {
2235 Some(ops) => Some((ops.checkpoint()?, savepoint_len)),
2236 None => None,
2237 }
2238 } else {
2239 None
2240 };
2241 Ok(Self {
2242 transaction: Arc::clone(&session.transaction),
2243 query: query.clone(),
2244 checkpoint,
2245 completed: false,
2246 })
2247 }
2248
2249 fn complete(mut self) {
2250 self.completed = true;
2251 }
2252}
2253
2254impl Drop for TransactionStatementGuard {
2255 fn drop(&mut self) {
2256 if self.completed {
2257 return;
2258 }
2259 if self.query.status().durable_outcome.committed {
2260 return;
2261 }
2262 if let Some((checkpoint, savepoint_len)) = self.checkpoint {
2263 let mut transaction = self.transaction.lock();
2264 let restored = transaction
2265 .staged_ops
2266 .as_mut()
2267 .is_some_and(|ops| ops.truncate(checkpoint).is_ok());
2268 if restored {
2269 transaction.savepoints.truncate(savepoint_len);
2270 } else {
2271 transaction.staged_ops = Some(commands::PendingSqlOps::default());
2272 transaction.savepoints.clear();
2273 }
2274 transaction.aborted = true;
2275 }
2276 }
2277}
2278
2279struct TransactionGuardedStream {
2280 inner: MongrelRecordBatchStream,
2281 schema: SchemaRef,
2282 statement: Option<TransactionStatementGuard>,
2283}
2284
2285impl futures::Stream for TransactionGuardedStream {
2286 type Item = datafusion::common::Result<RecordBatch>;
2287
2288 fn poll_next(
2289 mut self: std::pin::Pin<&mut Self>,
2290 context: &mut std::task::Context<'_>,
2291 ) -> std::task::Poll<Option<Self::Item>> {
2292 match self.inner.as_mut().poll_next(context) {
2293 std::task::Poll::Ready(None) => {
2294 if let Some(statement) = self.statement.take() {
2295 statement.complete();
2296 }
2297 std::task::Poll::Ready(None)
2298 }
2299 std::task::Poll::Ready(Some(Err(error))) => {
2300 self.statement.take();
2301 std::task::Poll::Ready(Some(Err(error)))
2302 }
2303 poll => poll,
2304 }
2305 }
2306}
2307
2308impl RecordBatchStream for TransactionGuardedStream {
2309 fn schema(&self) -> SchemaRef {
2310 Arc::clone(&self.schema)
2311 }
2312}
2313
2314impl futures::Stream for ManagedQueryStream {
2315 type Item = datafusion::common::Result<RecordBatch>;
2316
2317 fn poll_next(
2318 mut self: std::pin::Pin<&mut Self>,
2319 context: &mut std::task::Context<'_>,
2320 ) -> std::task::Poll<Option<Self::Item>> {
2321 if let Err(error) = self.query.checkpoint() {
2322 if let Some(guard) = self.guard.take() {
2323 guard.fail();
2324 }
2325 self.permit.take();
2326 self.inner = buffered_stream(Vec::new());
2327 return std::task::Poll::Ready(Some(Err(DataFusionError::External(Box::new(error)))));
2328 }
2329 match self.inner.as_mut().poll_next(context) {
2330 std::task::Poll::Ready(None) => {
2331 self.query.complete_current_statement();
2332 if self.deferred_completion.is_none() {
2333 if let Some(guard) = self.guard.take() {
2334 if let Err(error) = guard.try_complete() {
2335 self.permit.take();
2336 return std::task::Poll::Ready(Some(Err(DataFusionError::External(
2337 Box::new(error),
2338 ))));
2339 }
2340 }
2341 self.permit.take();
2342 }
2343 std::task::Poll::Ready(None)
2344 }
2345 std::task::Poll::Ready(Some(Err(error))) => {
2346 if let Some(guard) = self.guard.take() {
2347 guard.fail();
2348 }
2349 self.permit.take();
2350 std::task::Poll::Ready(Some(Err(error)))
2351 }
2352 item => item,
2353 }
2354 }
2355}
2356
2357impl datafusion::physical_plan::RecordBatchStream for ManagedQueryStream {
2358 fn schema(&self) -> SchemaRef {
2359 Arc::clone(&self.schema)
2360 }
2361}
2362
2363impl Drop for ManagedQueryStream {
2364 fn drop(&mut self) {
2365 if self.guard.is_some() {
2366 let mut completion_state = self
2367 .deferred_completion
2368 .as_ref()
2369 .map(|completion| completion.load(std::sync::atomic::Ordering::Acquire));
2370 while matches!(
2371 completion_state,
2372 Some(STREAM_COMPLETION_FAILING | STREAM_COMPLETION_COMPLETING)
2373 ) {
2374 std::hint::spin_loop();
2375 completion_state = self
2376 .deferred_completion
2377 .as_ref()
2378 .map(|completion| completion.load(std::sync::atomic::Ordering::Acquire));
2379 }
2380 match completion_state {
2381 Some(STREAM_COMPLETION_SUCCEEDED) => {
2382 if let Some(guard) = self.guard.take() {
2383 let result = guard.complete();
2384 debug_assert!(result.is_ok());
2385 }
2386 }
2387 Some(STREAM_COMPLETION_FAILED) => {
2388 if let Some(guard) = self.guard.take() {
2389 guard.fail();
2390 }
2391 }
2392 _ => {
2393 let _ = self
2394 .query
2395 .request_cancel(mongreldb_core::CancellationReason::ClientDisconnected);
2396 if let Some(guard) = self.guard.take() {
2397 guard.fail();
2398 }
2399 }
2400 }
2401 }
2402 }
2403}
2404
2405impl MongrelSession {
2406 pub fn new(db: Table) -> Self {
2410 let db = Arc::new(Mutex::new(db));
2411 let handle = TableHandle::from(Arc::clone(&db));
2412 let ctx = SessionContext::new();
2413 let sql_fn_state = Arc::new(extended_sql_functions::ExtendedSqlState::default());
2414 register_mongrel_functions(&ctx, Arc::clone(&sql_fn_state));
2415 let tables = Arc::new(parking_lot::Mutex::new(HashMap::new()));
2416 let scored_runtime = scored_sql::ScoredRuntime::from_env();
2417 scored_sql::register(
2418 &ctx,
2419 Arc::clone(&tables),
2420 None,
2421 None,
2422 false,
2423 Arc::clone(&scored_runtime),
2424 );
2425 let external_modules = Arc::new(ExternalModuleRegistry::default());
2426 Self {
2427 ctx,
2428 db: Some(handle),
2429 database: None,
2430 principal: None,
2431 principal_catalog_bound: false,
2432 cache: parking_lot::Mutex::new(ResultCacheStore::new(
2433 SESSION_CACHE_MAX_ENTRIES,
2434 SESSION_CACHE_MAX_BYTES,
2435 )),
2436 plan_cache: parking_lot::Mutex::new(BoundedLru::new(SESSION_CACHE_MAX_ENTRIES)),
2437 tables,
2438 views: parking_lot::Mutex::new(HashMap::new()),
2439 attached_databases: parking_lot::Mutex::new(HashMap::new()),
2440 transaction: Arc::new(parking_lot::Mutex::new(SqlTransactionState::default())),
2441 sql_fn_state,
2442 external_modules,
2443 scored_runtime,
2444 prepared_names: parking_lot::Mutex::new(std::collections::HashSet::new()),
2445 query_registry: Arc::new(SqlQueryRegistry::default()),
2446 execution_gate: Arc::new(tokio::sync::Semaphore::new(1)),
2447 statement_timeout: parking_lot::RwLock::new(None),
2448 test_hook: parking_lot::Mutex::new(None),
2449 }
2450 }
2451
2452 pub fn new_with_external_modules(
2453 db: Table,
2454 modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
2455 ) -> Result<Self> {
2456 let session = Self::new(db);
2457 for module in modules {
2458 session.register_external_module(module)?;
2459 }
2460 Ok(session)
2461 }
2462
2463 pub fn open(database: Arc<Database>) -> Result<Self> {
2468 Self::open_with_external_modules(database, std::iter::empty())
2469 }
2470
2471 pub fn open_as(database: Arc<Database>, principal: mongreldb_core::Principal) -> Result<Self> {
2472 Self::open_with_external_modules_as(database, std::iter::empty(), Some(principal))
2473 }
2474
2475 pub fn open_with_external_modules(
2476 database: Arc<Database>,
2477 modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
2478 ) -> Result<Self> {
2479 let principal = database.principal_snapshot();
2480 Self::open_with_external_modules_as(database, modules, principal)
2481 }
2482
2483 pub fn open_with_external_modules_as(
2484 database: Arc<Database>,
2485 modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
2486 principal: Option<mongreldb_core::Principal>,
2487 ) -> Result<Self> {
2488 let ctx = SessionContext::new();
2489 let sql_fn_state = Arc::new(extended_sql_functions::ExtendedSqlState::default());
2490 register_mongrel_functions(&ctx, Arc::clone(&sql_fn_state));
2491 let external_modules = Arc::new(ExternalModuleRegistry::default());
2492 let principal_catalog_bound = principal
2493 .as_ref()
2494 .is_some_and(|principal| principal_is_catalog_bound(&database, principal));
2495 let principal = match principal {
2496 Some(principal) if principal_catalog_bound => Some(
2497 database
2498 .resolve_current_principal(&principal)
2499 .ok_or(mongreldb_core::MongrelError::AuthRequired)?,
2500 ),
2501 principal => principal,
2502 };
2503 for module in modules {
2504 external_modules.register(module)?;
2505 }
2506
2507 let mut tables: HashMap<String, TableHandle> = HashMap::new();
2508 for name in database.table_names() {
2509 let handle = database.table(&name)?;
2510 let provider = MongrelProvider::new_secured(
2511 handle.clone(),
2512 Arc::clone(&database),
2513 name.clone(),
2514 principal.clone(),
2515 )?;
2516 ctx.register_table(&name, Arc::new(provider))
2517 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2518 tables.insert(name, handle);
2519 }
2520 for entry in database.external_tables() {
2521 let provider = external_modules.external_table_provider(&database, &entry)?;
2522 ctx.register_table(&entry.name, provider)
2523 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2524 }
2525
2526 let primary = {
2529 let mut names: Vec<&String> = tables.keys().collect();
2530 names.sort();
2531 names.first().and_then(|n| tables.get(*n).cloned())
2532 };
2533 let tables = Arc::new(parking_lot::Mutex::new(tables));
2534 let scored_runtime = scored_sql::ScoredRuntime::from_env();
2535 scored_sql::register(
2536 &ctx,
2537 Arc::clone(&tables),
2538 Some(Arc::clone(&database)),
2539 principal.clone(),
2540 principal_catalog_bound,
2541 Arc::clone(&scored_runtime),
2542 );
2543
2544 Ok(Self {
2545 ctx,
2546 db: primary,
2547 database: Some(database),
2548 principal,
2549 principal_catalog_bound,
2550 cache: parking_lot::Mutex::new(ResultCacheStore::new(
2551 SESSION_CACHE_MAX_ENTRIES,
2552 SESSION_CACHE_MAX_BYTES,
2553 )),
2554 plan_cache: parking_lot::Mutex::new(BoundedLru::new(SESSION_CACHE_MAX_ENTRIES)),
2555 tables,
2556 views: parking_lot::Mutex::new(HashMap::new()),
2557 attached_databases: parking_lot::Mutex::new(HashMap::new()),
2558 transaction: Arc::new(parking_lot::Mutex::new(SqlTransactionState::default())),
2559 sql_fn_state,
2560 external_modules,
2561 scored_runtime,
2562 prepared_names: parking_lot::Mutex::new(std::collections::HashSet::new()),
2563 query_registry: Arc::new(SqlQueryRegistry::default()),
2564 execution_gate: Arc::new(tokio::sync::Semaphore::new(1)),
2565 statement_timeout: parking_lot::RwLock::new(None),
2566 test_hook: parking_lot::Mutex::new(None),
2567 })
2568 }
2569
2570 pub fn query_registry(&self) -> &Arc<SqlQueryRegistry> {
2571 &self.query_registry
2572 }
2573
2574 pub fn with_query_registry(mut self, registry: Arc<SqlQueryRegistry>) -> Self {
2578 self.query_registry = registry;
2579 self
2580 }
2581
2582 pub fn set_statement_timeout(&self, timeout: Option<std::time::Duration>) {
2583 *self.statement_timeout.write() = timeout;
2584 }
2585
2586 pub fn register_query(&self, mut options: SqlQueryOptions) -> Result<RegisteredSqlQuery> {
2587 if let Some(statement_timeout) = *self.statement_timeout.read() {
2588 options.timeout = Some(
2589 options
2590 .timeout
2591 .map_or(statement_timeout, |timeout| timeout.min(statement_timeout)),
2592 );
2593 }
2594 let query = self.query_registry.register(options)?;
2595 self.fire_test_hook(SqlTestHookPoint::AfterRegistration);
2596 Ok(query)
2597 }
2598
2599 pub fn cancel_query(&self, query_id: QueryId) -> CancelOutcome {
2600 self.query_registry.cancel(query_id)
2601 }
2602
2603 #[doc(hidden)]
2604 pub fn set_test_hook(&self, hook: Option<SqlTestHook>) {
2605 *self.test_hook.lock() = hook;
2606 }
2607
2608 #[doc(hidden)]
2609 pub fn sql_test_hook(&self) -> Option<SqlTestHook> {
2610 self.test_hook.lock().clone()
2611 }
2612
2613 #[doc(hidden)]
2614 pub fn staged_sql_operation_count(&self) -> Option<usize> {
2615 self.transaction
2616 .lock()
2617 .staged_ops
2618 .as_ref()
2619 .map(|ops| ops.len())
2620 }
2621
2622 pub fn fire_test_hook(&self, point: SqlTestHookPoint) {
2623 let hook = self.test_hook.lock().clone();
2624 if let Some(hook) = hook {
2625 hook(point);
2626 }
2627 }
2628
2629 pub(crate) fn current_query(&self) -> Result<RegisteredSqlQuery> {
2630 CURRENT_SQL_QUERY.try_with(Clone::clone).map_err(|_| {
2631 MongrelQueryError::InvalidQueryState("SQL command has no registered query".into())
2632 })
2633 }
2634
2635 pub fn principal(&self) -> Option<mongreldb_core::Principal> {
2636 self.principal.clone()
2637 }
2638
2639 pub fn set_scored_execution_limits(
2641 &self,
2642 statement_timeout: std::time::Duration,
2643 max_work: usize,
2644 max_fused_candidates: usize,
2645 ) -> Result<()> {
2646 self.scored_runtime
2647 .configure(statement_timeout, max_work, max_fused_candidates)
2648 .map_err(Into::into)
2649 }
2650
2651 fn security_context_active(&self) -> bool {
2652 self.principal.is_some()
2653 || self.database.as_ref().is_some_and(|database| {
2654 database.principal_snapshot().is_some() || {
2655 let security = database.security_catalog();
2656 !security.rls_tables.is_empty()
2657 || !security.policies.is_empty()
2658 || !security.masks.is_empty()
2659 }
2660 })
2661 }
2662
2663 pub fn register_external_module(&self, module: Arc<dyn ExternalTableModule>) -> Result<()> {
2664 self.external_modules.register(module)?;
2665 self.clear_cache();
2666 Ok(())
2667 }
2668
2669 pub fn db(&self) -> Option<&TableHandle> {
2673 self.db.as_ref()
2674 }
2675
2676 pub fn create_view(&self, name: &str, sql: &str) {
2681 self.create_view_with_schema(name, sql, CoreSchema::default(), HashMap::new());
2682 }
2683
2684 pub(crate) fn create_view_with_schema(
2685 &self,
2686 name: &str,
2687 sql: &str,
2688 schema: CoreSchema,
2689 input_types: HashMap<u16, Option<TypeId>>,
2690 ) {
2691 self.views.lock().insert(
2692 name.to_string(),
2693 ViewDef {
2694 sql: sql.to_string(),
2695 schema,
2696 input_types,
2697 },
2698 );
2699 }
2700
2701 pub fn drop_view(&self, name: &str) {
2703 self.views.lock().remove(name);
2704 }
2705
2706 pub(crate) fn view_schema(&self, name: &str) -> Option<CoreSchema> {
2707 self.views.lock().get(name).map(|view| view.schema.clone())
2708 }
2709
2710 pub(crate) fn view_definition(&self, name: &str) -> Option<ViewDef> {
2711 self.views.lock().get(name).cloned()
2712 }
2713
2714 pub async fn register(&self, name: &str) -> Result<()> {
2716 let db = self.db.clone().ok_or(MongrelQueryError::Core(
2717 mongreldb_core::MongrelError::NotFound("no primary table".into()),
2718 ))?;
2719 let provider = MongrelProvider::new_handle(db.clone())?;
2720 self.tables.lock().insert(name.to_string(), db);
2721 self.ctx
2722 .register_table(name, Arc::new(provider))
2723 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2724 Ok(())
2725 }
2726
2727 pub async fn register_db(&self, name: &str, db: Table) -> Result<()> {
2733 let db_arc = Arc::new(Mutex::new(db));
2734 let provider = MongrelProvider::new(db_arc.clone())?;
2735 self.tables.lock().insert(name.to_string(), db_arc.into());
2736 self.ctx
2737 .register_table(name, Arc::new(provider))
2738 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2739 Ok(())
2740 }
2741
2742 fn refresh_registered_table(&self, db: &Arc<Database>, name: &str) -> Result<()> {
2743 self.ctx
2744 .deregister_table(name)
2745 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2746 let handle = db.table(name)?;
2747 let provider = MongrelProvider::new_secured(
2748 handle.clone(),
2749 Arc::clone(db),
2750 name.to_string(),
2751 self.principal.clone(),
2752 )?;
2753 self.ctx
2754 .register_table(name, Arc::new(provider))
2755 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2756 self.tables.lock().insert(name.to_string(), handle);
2757 Ok(())
2758 }
2759
2760 fn try_catalog_introspection(&self, sql: &str) -> Result<Option<Vec<RecordBatch>>> {
2768 let lower = sql.to_ascii_lowercase();
2769 if !lower.contains("information_schema.tables") {
2770 return Ok(None);
2771 }
2772 use arrow::array::{ArrayRef, Int64Array, StringArray};
2773 use arrow::datatypes::{DataType, Field, Schema};
2774 use arrow::record_batch::RecordBatch;
2775
2776 let mut types: Vec<String> = Vec::new();
2777 let mut names: Vec<String> = Vec::new();
2778 let mut tbl_names: Vec<String> = Vec::new();
2779
2780 let principal = self.principal();
2782 let can_see = |name: &str| match &self.database {
2783 Some(database) => database
2784 .select_column_ids_for(name, principal.as_ref())
2785 .is_ok(),
2786 None => true,
2787 };
2788 for name in self.tables.lock().keys().filter(|name| can_see(name)) {
2789 types.push("table".into());
2790 names.push(name.clone());
2791 tbl_names.push(name.clone());
2792 }
2793 for name in self.views.lock().keys() {
2795 types.push("view".into());
2796 names.push(name.clone());
2797 tbl_names.push(name.clone());
2798 }
2799 if let Some(db) = &self.database {
2801 for t in db.triggers() {
2802 let target_name = match &t.target {
2803 mongreldb_core::trigger::TriggerTarget::Table(n)
2804 | mongreldb_core::trigger::TriggerTarget::View(n) => n.clone(),
2805 };
2806 if !can_see(&target_name) {
2807 continue;
2808 }
2809 types.push("trigger".into());
2810 names.push(t.name.clone());
2811 tbl_names.push(target_name);
2812 }
2813 }
2814
2815 let schema = Arc::new(Schema::new(vec![
2816 Field::new("type", DataType::Utf8, false),
2817 Field::new("name", DataType::Utf8, false),
2818 Field::new("tbl_name", DataType::Utf8, false),
2819 Field::new("rootpage", DataType::Int64, false),
2820 Field::new("sql", DataType::Utf8, true),
2821 ]));
2822 let n = names.len();
2823 let rootpages: Vec<i64> = vec![0; n];
2824 let sqls: Vec<Option<&str>> = vec![None; n];
2825 let batch = RecordBatch::try_new(
2826 schema,
2827 vec![
2828 Arc::new(StringArray::from(types)) as ArrayRef,
2829 Arc::new(StringArray::from(names)) as ArrayRef,
2830 Arc::new(StringArray::from(tbl_names)) as ArrayRef,
2831 Arc::new(Int64Array::from(rootpages)) as ArrayRef,
2832 Arc::new(StringArray::from(sqls)) as ArrayRef,
2833 ],
2834 )
2835 .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
2836 Ok(Some(vec![batch]))
2837 }
2838
2839 fn try_direct_dispatch(
2845 &self,
2846 sql: &str,
2847 query: &RegisteredSqlQuery,
2848 ) -> Result<Option<Vec<RecordBatch>>> {
2849 query.checkpoint()?;
2850 if self.security_context_active() {
2851 return Ok(None);
2852 }
2853
2854 use arrow::array::ArrayRef;
2855 use mongreldb_core::Condition;
2856 use sqlparser::ast::{Expr, Query, SelectItem, SetExpr, Statement, TableFactor};
2857 use sqlparser::dialect::PostgreSqlDialect;
2858 use sqlparser::parser::Parser;
2859
2860 let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
2862 return Ok(None);
2863 };
2864 query.checkpoint()?;
2865 if stmts.len() != 1 {
2866 return Ok(None);
2867 }
2868 let Some(statement) = stmts.into_iter().next() else {
2869 return Ok(None);
2870 };
2871 let Statement::Query(ast_query) = statement else {
2872 return Ok(None);
2873 };
2874 let Query { body, .. } = *ast_query;
2875 let select = match *body {
2876 SetExpr::Select(s) => *s,
2877 _ => return Ok(None),
2878 };
2879 let lower_sql = sql.to_lowercase();
2882 if lower_sql.contains(" limit ") || lower_sql.contains(" offset ") {
2883 return Ok(None);
2884 }
2885 use sqlparser::ast::GroupByExpr;
2887 if select.distinct.is_some()
2888 || !matches!(&select.group_by, GroupByExpr::Expressions(e, _) if e.is_empty())
2889 || select.having.is_some()
2890 || select.from.len() != 1
2891 || !select.from[0].joins.is_empty()
2892 {
2893 return Ok(None);
2894 }
2895 let table_name = match &select.from[0].relation {
2896 TableFactor::Table { name, .. } => Some(name.to_string()),
2897 _ => return Ok(None),
2898 };
2899 let Some(table_name) = table_name else {
2900 return Ok(None);
2901 };
2902
2903 if select.selection.is_none() {
2909 return Ok(None);
2910 }
2911
2912 let mut proj_names: Option<Vec<String>> = None;
2914 for item in &select.projection {
2915 match item {
2916 SelectItem::Wildcard(_) => {}
2917 SelectItem::UnnamedExpr(Expr::Identifier(ident)) => {
2918 proj_names
2919 .get_or_insert_with(Vec::new)
2920 .push(ident.value.clone());
2921 }
2922 SelectItem::UnnamedExpr(Expr::CompoundIdentifier(idents)) => {
2923 if let Some(last) = idents.last() {
2924 proj_names
2925 .get_or_insert_with(Vec::new)
2926 .push(last.value.clone());
2927 }
2928 }
2929 _ => return Ok(None),
2930 }
2931 }
2932
2933 let handle = match self.tables.lock().get(&table_name).cloned() {
2935 Some(h) => h,
2936 None => return Ok(None),
2937 };
2938
2939 if let Some(db) = &self.database {
2941 db.require_table(
2942 &table_name,
2943 mongreldb_core::auth_state::RequiredPermission::Select,
2944 )?;
2945 }
2946
2947 let mut db = handle.lock();
2948 let schema = db.schema().clone();
2949 let conditions: Vec<Condition> = match &select.selection {
2952 Some(expr) => match translate_sqlparser_filter(expr, &schema) {
2953 Some(c) => c,
2954 None => return Ok(None),
2955 },
2956 None => Vec::new(),
2957 };
2958 if !conditions.is_empty() && db.ensure_indexes_complete().is_err() {
2959 return Ok(None);
2960 }
2961 let snap = db.snapshot();
2962 if db.ttl().is_some()
2963 || db
2964 .count_conditions(&conditions, snap)
2965 .map_err(MongrelQueryError::Core)?
2966 .is_none_or(|rows| rows > current_collection_limits().max_rows as u64)
2967 {
2968 return Ok(None);
2969 }
2970
2971 let mut col_ids: Vec<u16> = Vec::new();
2973 let mut fields: Vec<arrow::datatypes::Field> = Vec::new();
2974 let resolve_col = |name: &str| -> Option<&mongreldb_core::schema::ColumnDef> {
2975 schema.columns.iter().find(|c| c.name == name)
2976 };
2977 match &proj_names {
2978 None => {
2979 for c in &schema.columns {
2980 col_ids.push(c.id);
2981 fields.push(arrow::datatypes::Field::new(
2982 &c.name,
2983 arrow_conv::arrow_data_type(&c.ty)?,
2984 c.flags.contains(mongreldb_core::ColumnFlags::NULLABLE),
2985 ));
2986 }
2987 }
2988 Some(names) => {
2989 for n in names {
2990 let cdef = match resolve_col(n) {
2991 Some(c) => c,
2992 None => return Ok(None), };
2994 col_ids.push(cdef.id);
2995 fields.push(arrow::datatypes::Field::new(
2996 &cdef.name,
2997 arrow_conv::arrow_data_type(&cdef.ty)?,
2998 cdef.flags.contains(mongreldb_core::ColumnFlags::NULLABLE),
2999 ));
3000 }
3001 }
3002 }
3003
3004 let mut cols = if !conditions.is_empty() {
3006 match db.query_columns_native_cached_with_control(
3007 &conditions,
3008 Some(&col_ids),
3009 snap,
3010 query.control(),
3011 ) {
3012 Ok(Some(c)) => c,
3013 Ok(None) => db
3014 .visible_columns_native_with_control(snap, Some(&col_ids), query.control())
3015 .map_err(MongrelQueryError::Core)?,
3016 Err(_) => return Ok(None),
3017 }
3018 } else {
3019 db.visible_columns_native_with_control(snap, Some(&col_ids), query.control())
3020 .map_err(MongrelQueryError::Core)?
3021 };
3022 query.checkpoint()?;
3023 let native_bytes = cols.iter().fold(0_u64, |total, (_, column)| {
3024 total.saturating_add(column.approx_bytes())
3025 });
3026 if native_bytes > current_collection_limits().max_bytes as u64 {
3027 return Ok(None);
3028 }
3029 drop(db);
3030
3031 let mut arrays: Vec<ArrayRef> = Vec::with_capacity(col_ids.len());
3033 for cid in &col_ids {
3034 query.checkpoint()?;
3035 let Some(position) = cols.iter().position(|(id, _)| id == cid) else {
3036 return Ok(None);
3037 };
3038 let (_, col) = cols.swap_remove(position);
3039 let ty = schema
3040 .columns
3041 .iter()
3042 .find(|c| c.id == *cid)
3043 .map(|c| c.ty.clone())
3044 .unwrap_or(mongreldb_core::schema::TypeId::Int64);
3045 arrays.push(arrow_conv::native_to_array_with_query(
3046 ty,
3047 &col,
3048 Some(query),
3049 )?);
3050 }
3051 let batch_schema = Arc::new(arrow::datatypes::Schema::new(fields));
3052 let batch = RecordBatch::try_new(batch_schema, arrays)
3053 .map_err(|e| MongrelQueryError::Arrow(format!("direct dispatch batch build: {e}")))?;
3054
3055 mongreldb_core::trace::QueryTrace::record(|t| {
3056 t.scan_mode = mongreldb_core::trace::ScanMode::DirectDispatch;
3057 t.planning_nanos = 0; });
3059 Ok(Some(vec![batch]))
3060 }
3061
3062 async fn dataframe(&self, sql: &str) -> Result<datafusion::dataframe::DataFrame> {
3063 let use_plan_cache = !is_prepared_stmt_sql(sql);
3066 if let Some(plan) = use_plan_cache
3067 .then(|| self.plan_cache.lock().get(sql).cloned())
3068 .flatten()
3069 {
3070 return Ok(datafusion::dataframe::DataFrame::new(
3071 self.ctx.state(),
3072 plan,
3073 ));
3074 }
3075
3076 let df = self
3077 .ctx
3078 .sql(sql)
3079 .await
3080 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
3081 if use_plan_cache {
3082 self.plan_cache
3083 .lock()
3084 .insert(sql.to_string(), df.logical_plan().clone());
3085 }
3086 Ok(df)
3087 }
3088
3089 async fn dataframe_with_query(
3090 &self,
3091 sql: &str,
3092 query: &RegisteredSqlQuery,
3093 ) -> Result<datafusion::dataframe::DataFrame> {
3094 query.checkpoint()?;
3095 tokio::select! {
3096 biased;
3097 dataframe = self.dataframe(sql) => dataframe,
3098 _ = query.control().cancelled() => match query.checkpoint() {
3099 Err(error) => Err(error),
3100 Ok(()) => Err(MongrelQueryError::InvalidQueryState(
3101 "query cancellation signal resolved without cancellation".into(),
3102 )),
3103 },
3104 }
3105 }
3106
3107 fn prepare_as_of_query(&self, sql: &str) -> Result<Option<AsOfQuery>> {
3108 static AS_OF_RE: std::sync::OnceLock<std::result::Result<regex::Regex, regex::Error>> =
3109 std::sync::OnceLock::new();
3110 static NEXT_TEMP: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
3111 let re = AS_OF_RE
3112 .get_or_init(|| {
3113 regex::Regex::new(
3114 r#"(?i)\b(FROM|JOIN)\s+("[^"]+"|[A-Za-z_][A-Za-z0-9_]*)\s+AS\s+OF\s+EPOCH\s+([0-9]+)\b"#,
3115 )
3116 })
3117 .as_ref()
3118 .map_err(|error| {
3119 MongrelQueryError::InvalidQueryState(format!(
3120 "failed to initialize AS OF parser: {error}"
3121 ))
3122 })?;
3123 let mut captures = re.captures_iter(sql);
3124 let Some(capture) = captures.next() else {
3125 return Ok(None);
3126 };
3127 if captures.next().is_some() {
3128 return Err(MongrelQueryError::Schema(
3129 "AS OF EPOCH currently supports one historical table per query".into(),
3130 ));
3131 }
3132 let database = self.database.as_ref().ok_or_else(|| {
3133 MongrelQueryError::Schema("AS OF EPOCH requires a Database-backed session".into())
3134 })?;
3135 let whole = capture.get(0).ok_or_else(|| {
3136 MongrelQueryError::InvalidQueryState("AS OF parser omitted the whole match".into())
3137 })?;
3138 let keyword = capture
3139 .get(1)
3140 .ok_or_else(|| {
3141 MongrelQueryError::InvalidQueryState("AS OF parser omitted FROM/JOIN".into())
3142 })?
3143 .as_str();
3144 let table_token = capture
3145 .get(2)
3146 .ok_or_else(|| {
3147 MongrelQueryError::InvalidQueryState("AS OF parser omitted the table".into())
3148 })?
3149 .as_str();
3150 let table_name = table_token.trim_matches('"');
3151 let epoch = capture
3152 .get(3)
3153 .ok_or_else(|| {
3154 MongrelQueryError::InvalidQueryState("AS OF parser omitted the epoch".into())
3155 })?
3156 .as_str()
3157 .parse::<u64>()
3158 .map_err(|error| MongrelQueryError::Schema(format!("AS OF epoch: {error}")))?;
3159 let temp_id = NEXT_TEMP.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3160 let temp_name = format!("__mongrel_as_of_{temp_id}");
3161
3162 let mut replace_end = whole.end();
3166 let mut alias = table_token.to_string();
3167 let tail = &sql[replace_end..];
3168 let whitespace = tail.len() - tail.trim_start().len();
3169 let trimmed_tail = tail.trim_start();
3170 if trimmed_tail
3171 .get(..3)
3172 .is_some_and(|prefix| prefix.eq_ignore_ascii_case("as "))
3173 {
3174 let alias_text = &trimmed_tail[3..];
3175 let alias_len = alias_text
3176 .find(|character: char| {
3177 character.is_ascii_whitespace() || character == ',' || character == ';'
3178 })
3179 .unwrap_or(alias_text.len());
3180 if alias_len == 0 {
3181 return Err(MongrelQueryError::Schema(
3182 "AS OF EPOCH AS requires an alias".into(),
3183 ));
3184 }
3185 alias = alias_text[..alias_len].to_string();
3186 replace_end += whitespace + 3 + alias_len;
3187 } else {
3188 let alias_len = trimmed_tail
3189 .find(|character: char| {
3190 character.is_ascii_whitespace() || character == ',' || character == ';'
3191 })
3192 .unwrap_or(trimmed_tail.len());
3193 let candidate = &trimmed_tail[..alias_len];
3194 let keyword = candidate.to_ascii_lowercase();
3195 let clause_keyword = matches!(
3196 keyword.as_str(),
3197 "where"
3198 | "join"
3199 | "left"
3200 | "right"
3201 | "inner"
3202 | "outer"
3203 | "cross"
3204 | "full"
3205 | "on"
3206 | "group"
3207 | "order"
3208 | "having"
3209 | "limit"
3210 | "offset"
3211 | "union"
3212 | "intersect"
3213 | "except"
3214 );
3215 let valid_alias = candidate
3216 .as_bytes()
3217 .first()
3218 .is_some_and(|byte| byte.is_ascii_alphabetic() || *byte == b'_' || *byte == b'"');
3219 if alias_len > 0 && valid_alias && !clause_keyword {
3220 alias = candidate.to_string();
3221 replace_end += whitespace + alias_len;
3222 }
3223 }
3224 let replacement = format!("{keyword} {temp_name} AS {alias}");
3225 let mut rewritten = String::with_capacity(sql.len() + replacement.len());
3226 rewritten.push_str(&sql[..whole.start()]);
3227 rewritten.push_str(&replacement);
3228 rewritten.push_str(&sql[replace_end..]);
3229
3230 let (snapshot, retention) = database.snapshot_at_owned(mongreldb_core::Epoch(epoch))?;
3231 let handle = database.table(table_name)?;
3232 let provider = MongrelProvider::new_historical(
3233 handle,
3234 snapshot,
3235 retention,
3236 Some(ProviderSecurity {
3237 database: Arc::clone(database),
3238 table: table_name.to_string(),
3239 principal: self.principal.clone(),
3240 principal_catalog_bound: self.principal_catalog_bound,
3241 }),
3242 )?;
3243 self.ctx
3244 .register_table(&temp_name, Arc::new(provider))
3245 .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))?;
3246 Ok(Some(AsOfQuery {
3247 sql: rewritten,
3248 registration: AsOfRegistration {
3249 ctx: self.ctx.clone(),
3250 table_name: temp_name,
3251 },
3252 }))
3253 }
3254
3255 async fn execute_dataframe_stream(
3256 &self,
3257 dataframe: datafusion::dataframe::DataFrame,
3258 query: &RegisteredSqlQuery,
3259 ) -> Result<MongrelRecordBatchStream> {
3260 query.checkpoint()?;
3261 let task_context = dataframe.task_ctx();
3262 let session_config = task_context
3263 .session_config()
3264 .clone()
3265 .with_extension(Arc::new(SqlTaskContext::new(
3266 query.clone(),
3267 self.sql_test_hook(),
3268 )));
3269 let task_context = Arc::new(task_context.with_session_config(session_config));
3270 let physical_plan = tokio::select! {
3271 biased;
3272 plan = dataframe.create_physical_plan() => plan
3273 .map_err(map_datafusion_error)?,
3274 _ = query.control().cancelled() => return Err(query.checkpoint().unwrap_err()),
3275 };
3276 query.checkpoint()?;
3277 execute_stream(physical_plan, task_context).map_err(map_datafusion_error)
3278 }
3279
3280 pub(crate) async fn execute_command_source_stream(
3284 &self,
3285 sql: &str,
3286 query: &RegisteredSqlQuery,
3287 ) -> Result<MongrelRecordBatchStream> {
3288 query.checkpoint()?;
3289 if let Some(as_of_query) = self.prepare_as_of_query(sql)? {
3290 return self.run_as_of_stream(as_of_query, query).await;
3291 }
3292
3293 let resolved = self.resolve_view_sql(sql);
3294 let resolved = self.rewrite_external_module_compat_sql(&resolved);
3295 let resolved = rewrite_compat_function_calls(&resolved);
3296 let effective_sql = normalize_sql(&resolved);
3297 let sql = effective_sql.as_str();
3298
3299 if let Some(batches) = self.try_catalog_introspection(sql)? {
3300 return Ok(buffered_stream(batches));
3301 }
3302
3303 let external_module_scan = self.query_references_external_module(sql);
3304 let dataframe = self.dataframe_with_query(sql, query).await?;
3305 let stream = self.execute_dataframe_stream(dataframe, query).await?;
3306 if external_module_scan {
3307 mongreldb_core::trace::QueryTrace::record(|trace| {
3308 trace.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
3309 });
3310 }
3311 Ok(stream)
3312 }
3313
3314 async fn collect_dataframe(
3315 &self,
3316 dataframe: datafusion::dataframe::DataFrame,
3317 query: &RegisteredSqlQuery,
3318 ) -> Result<Vec<RecordBatch>> {
3319 use futures::StreamExt;
3320
3321 let mut stream = self.execute_dataframe_stream(dataframe, query).await?;
3322 let mut batches = Vec::new();
3323 let limits = current_collection_limits();
3324 let mut rows = 0_usize;
3325 let mut bytes = 0_usize;
3326 loop {
3327 let item = tokio::select! {
3328 biased;
3329 _ = query.control().cancelled() => return Err(query.checkpoint().unwrap_err()),
3330 item = stream.next() => item,
3331 };
3332 let Some(batch) = item else {
3333 break;
3334 };
3335 let batch = batch.map_err(map_datafusion_error)?;
3336 rows = rows.saturating_add(batch.num_rows());
3337 bytes = bytes.saturating_add(record_batch_memory_bytes(&batch));
3338 if rows > limits.max_rows {
3339 return Err(query.result_limit_error(format!(
3340 "SQL collected row limit exceeded ({} > {})",
3341 rows, limits.max_rows
3342 )));
3343 }
3344 if bytes > limits.max_bytes {
3345 return Err(query.result_limit_error(format!(
3346 "SQL collected Arrow memory limit exceeded ({} > {} bytes)",
3347 bytes, limits.max_bytes
3348 )));
3349 }
3350 batches.push(batch);
3351 query.checkpoint()?;
3352 }
3353 Ok(batches)
3354 }
3355
3356 async fn run_as_of(
3357 &self,
3358 query: AsOfQuery,
3359 sql_query: &RegisteredSqlQuery,
3360 ) -> Result<Vec<RecordBatch>> {
3361 let AsOfQuery { sql, registration } = query;
3362 let _registration = registration;
3363 let plan_start = std::time::Instant::now();
3364 let dataframe = self.dataframe_with_query(&sql, sql_query).await?;
3365 mongreldb_core::trace::QueryTrace::record(|trace| {
3366 trace.planning_nanos = plan_start.elapsed().as_nanos() as u64;
3367 });
3368 self.collect_dataframe(dataframe, sql_query).await
3369 }
3370
3371 async fn run_as_of_stream(
3372 &self,
3373 query: AsOfQuery,
3374 sql_query: &RegisteredSqlQuery,
3375 ) -> Result<MongrelRecordBatchStream> {
3376 use futures::StreamExt;
3377
3378 let AsOfQuery { sql, registration } = query;
3379 let dataframe = self.dataframe_with_query(&sql, sql_query).await?;
3380 let stream = self.execute_dataframe_stream(dataframe, sql_query).await?;
3381 let schema = stream.schema();
3382 let guarded = futures::stream::unfold(
3383 (stream, registration),
3384 |(mut stream, registration)| async move {
3385 stream
3386 .next()
3387 .await
3388 .map(|item| (item, (stream, registration)))
3389 },
3390 );
3391 Ok(Box::pin(RecordBatchStreamAdapter::new(schema, guarded)))
3392 }
3393
3394 pub async fn run_stream_with_options(
3395 &self,
3396 sql: &str,
3397 options: SqlQueryOptions,
3398 ) -> Result<MongrelRecordBatchStream> {
3399 let query = self.register_query(options)?;
3400 self.run_stream_with_query(sql, query).await
3401 }
3402
3403 pub async fn run_stream_with_query(
3404 &self,
3405 sql: &str,
3406 query: RegisteredSqlQuery,
3407 ) -> Result<MongrelRecordBatchStream> {
3408 self.run_stream_with_query_mode(sql, query, false)
3409 .await
3410 .map(|(stream, _)| stream)
3411 }
3412
3413 pub async fn run_stream_with_query_for_serialization(
3414 &self,
3415 sql: &str,
3416 query: RegisteredSqlQuery,
3417 ) -> Result<(MongrelRecordBatchStream, SqlStreamCompletion)> {
3418 let (stream, completion) = self.run_stream_with_query_mode(sql, query, true).await?;
3419 let completion = completion.ok_or_else(|| {
3420 MongrelQueryError::InvalidQueryState(
3421 "deferred stream completion was not created".into(),
3422 )
3423 })?;
3424 Ok((stream, completion))
3425 }
3426
3427 async fn run_stream_with_query_mode(
3428 &self,
3429 sql: &str,
3430 query: RegisteredSqlQuery,
3431 defer_completion: bool,
3432 ) -> Result<(MongrelRecordBatchStream, Option<SqlStreamCompletion>)> {
3433 query.set_sql_metadata(sql);
3434 let guard = RegisteredQueryGuard::new(query.clone());
3435 self.fire_test_hook(SqlTestHookPoint::WaitingForSessionPermit);
3436 let permit = tokio::select! {
3437 permit = Arc::clone(&self.execution_gate).acquire_owned() => permit.map_err(|_| {
3438 MongrelQueryError::InvalidQueryState("session execution gate closed".into())
3439 })?,
3440 _ = query.control().cancelled() => {
3441 let error = query.checkpoint().unwrap_err();
3442 guard.fail();
3443 return Err(error);
3444 }
3445 };
3446 query.transition(SqlQueryPhase::Queued, SqlQueryPhase::Planning)?;
3447 self.fire_test_hook(SqlTestHookPoint::Planning);
3448 query.transition(SqlQueryPhase::Planning, SqlQueryPhase::Executing)?;
3449 query.begin_statement(0);
3450 let stream = CURRENT_SQL_QUERY
3451 .scope(query.clone(), async {
3452 tokio::select! {
3453 biased;
3454 result = self.run_stream_internal(sql, &query) => result,
3455 _ = query.control().cancelled() => Err(query.checkpoint().unwrap_err()),
3456 }
3457 })
3458 .await;
3459 let stream = match stream {
3460 Ok(stream) => stream,
3461 Err(error) => {
3462 guard.fail();
3463 return Err(error);
3464 }
3465 };
3466 if query.phase() == SqlQueryPhase::Executing {
3467 query.transition(SqlQueryPhase::Executing, SqlQueryPhase::Streaming)?;
3468 } else if query.phase() != SqlQueryPhase::CommitCritical {
3469 return Err(MongrelQueryError::InvalidQueryState(format!(
3470 "query {} cannot stream from {:?}",
3471 query.id(),
3472 query.phase()
3473 )));
3474 }
3475 let deferred_completion = defer_completion
3476 .then(|| Arc::new(std::sync::atomic::AtomicU8::new(STREAM_COMPLETION_PENDING)));
3477 if defer_completion {
3478 query.begin_serialization()?;
3479 }
3480 let schema = stream.schema();
3481 let completion = deferred_completion
3482 .as_ref()
3483 .map(|state| SqlStreamCompletion {
3484 state: Arc::clone(state),
3485 query: query.clone(),
3486 });
3487 let stream = Box::pin(ManagedQueryStream {
3488 inner: stream,
3489 schema,
3490 query,
3491 guard: Some(guard),
3492 permit: Some(permit),
3493 deferred_completion,
3494 });
3495 Ok((stream, completion))
3496 }
3497
3498 pub async fn run_stream(&self, sql: &str) -> Result<MongrelRecordBatchStream> {
3499 if let Ok(query) = CURRENT_SQL_QUERY.try_with(Clone::clone) {
3500 return Box::pin(self.run_stream_internal(sql, &query)).await;
3501 }
3502 Box::pin(self.run_stream_with_options(sql, SqlQueryOptions::default())).await
3503 }
3504
3505 async fn run_stream_internal(
3510 &self,
3511 sql: &str,
3512 query: &RegisteredSqlQuery,
3513 ) -> Result<MongrelRecordBatchStream> {
3514 let trimmed = sql.trim();
3515 let multiple_statements = trimmed.contains(';')
3516 && !is_trigger_body(trimmed)
3517 && split_sql_statements(trimmed)
3518 .iter()
3519 .filter(|statement| !statement.trim().is_empty())
3520 .count()
3521 > 1;
3522 let rollback = commands::rollback_kind(sql);
3523 let (transaction_open, transaction_aborted) = {
3524 let transaction = self.transaction.lock();
3525 (transaction.staged_ops.is_some(), transaction.aborted)
3526 };
3527 if !multiple_statements && transaction_open && transaction_aborted && rollback.is_none() {
3528 return Err(MongrelQueryError::TransactionAborted);
3529 }
3530 let statement = TransactionStatementGuard::new(
3531 self,
3532 query,
3533 !multiple_statements
3534 && transaction_open
3535 && rollback != Some(commands::RollbackKind::Full),
3536 )?;
3537 self.fire_test_hook(SqlTestHookPoint::BeforeStatementExecution);
3538 let (inner, buffered) = self.run_stream_internal_body(sql, query).await?;
3539 if buffered {
3540 statement.complete();
3541 return Ok(inner);
3542 }
3543 let schema = inner.schema();
3544 Ok(Box::pin(TransactionGuardedStream {
3545 inner,
3546 schema,
3547 statement: Some(statement),
3548 }))
3549 }
3550
3551 async fn run_stream_internal_body(
3552 &self,
3553 sql: &str,
3554 query: &RegisteredSqlQuery,
3555 ) -> Result<(MongrelRecordBatchStream, bool)> {
3556 query.checkpoint()?;
3557 if strip_explain_query_plan(sql).is_some() {
3558 return self
3559 .run_internal(sql, query)
3560 .await
3561 .map(|batches| (buffered_stream(batches), true));
3562 }
3563
3564 let trimmed = sql.trim();
3565 if trimmed.contains(';') && !is_trigger_body(trimmed) {
3566 let stmts = split_sql_statements(trimmed);
3567 let non_empty: Vec<&str> = stmts
3568 .iter()
3569 .map(String::as_str)
3570 .filter(|stmt| !stmt.trim().is_empty() && stmt.trim() != ";")
3571 .collect();
3572 if non_empty.is_empty() {
3573 return Ok((buffered_stream(Vec::new()), true));
3574 }
3575 if stmts.len() > 1 {
3576 for (index, stmt) in non_empty[..non_empty.len() - 1].iter().enumerate() {
3577 query.begin_statement(index);
3578 Box::pin(self.run_internal(stmt, query)).await?;
3579 query.complete_statement(index);
3580 }
3581 query.begin_statement(non_empty.len() - 1);
3582 let stream =
3583 Box::pin(self.run_stream_internal(non_empty[non_empty.len() - 1], query))
3584 .await?;
3585 return Ok((stream, true));
3586 }
3587 }
3588
3589 query.checkpoint()?;
3590 if let Some(batches) = commands::try_run_command(self, sql, query).await? {
3591 return Ok((buffered_stream(batches), true));
3592 }
3593
3594 if let Some(as_of_query) = self.prepare_as_of_query(sql)? {
3595 return self
3596 .run_as_of_stream(as_of_query, query)
3597 .await
3598 .map(|stream| (stream, false));
3599 }
3600
3601 let resolved = self.resolve_view_sql(sql);
3602 let resolved = self.rewrite_external_module_compat_sql(&resolved);
3603 let resolved = rewrite_compat_function_calls(&resolved);
3604 let effective_sql = normalize_sql(&resolved);
3605 let sql = effective_sql.as_str();
3606
3607 if let Some(batches) = self.try_catalog_introspection(sql)? {
3608 return Ok((buffered_stream(batches), true));
3609 }
3610
3611 let plan_start = std::time::Instant::now();
3612 let external_module_scan = self.query_references_external_module(sql);
3613 let df = self.dataframe_with_query(sql, query).await?;
3614 self.track_prepared_name(sql);
3615 mongreldb_core::trace::QueryTrace::record(|trace| {
3616 trace.planning_nanos = plan_start.elapsed().as_nanos() as u64;
3617 });
3618 let stream = self.execute_dataframe_stream(df, query).await?;
3619 if external_module_scan {
3620 mongreldb_core::trace::QueryTrace::record(|trace| {
3621 trace.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
3622 });
3623 }
3624 Ok((stream, false))
3625 }
3626
3627 pub async fn run_with_options(
3628 &self,
3629 sql: &str,
3630 options: SqlQueryOptions,
3631 ) -> Result<Vec<RecordBatch>> {
3632 let query = self.register_query(options)?;
3633 self.run_with_query(sql, query).await
3634 }
3635
3636 pub async fn run_with_query(
3637 &self,
3638 sql: &str,
3639 query: RegisteredSqlQuery,
3640 ) -> Result<Vec<RecordBatch>> {
3641 query.set_sql_metadata(sql);
3642 let guard = RegisteredQueryGuard::new(query.clone());
3643 self.fire_test_hook(SqlTestHookPoint::WaitingForSessionPermit);
3644 let _permit = tokio::select! {
3645 permit = Arc::clone(&self.execution_gate).acquire_owned() => permit.map_err(|_| {
3646 MongrelQueryError::InvalidQueryState("session execution gate closed".into())
3647 })?,
3648 _ = query.control().cancelled() => {
3649 let error = query.checkpoint().unwrap_err();
3650 guard.fail();
3651 return Err(error);
3652 }
3653 };
3654 query.transition(SqlQueryPhase::Queued, SqlQueryPhase::Planning)?;
3655 self.fire_test_hook(SqlTestHookPoint::Planning);
3656 query.transition(SqlQueryPhase::Planning, SqlQueryPhase::Executing)?;
3657 query.begin_statement(0);
3658 let result = CURRENT_COLLECTION_LIMITS
3659 .scope(
3660 SqlCollectionLimits::default(),
3661 CURRENT_SQL_QUERY.scope(query.clone(), async {
3662 tokio::select! {
3663 biased;
3664 result = self.run_internal(sql, &query) => result,
3665 _ = query.control().cancelled() => Err(query.checkpoint().unwrap_err()),
3666 }
3667 }),
3668 )
3669 .await;
3670 match result {
3671 Ok(batches) => {
3672 query.complete_current_statement();
3673 guard.try_complete()?;
3674 Ok(batches)
3675 }
3676 Err(error) => {
3677 if matches!(error, MongrelQueryError::ResultLimitExceeded { .. }) {
3678 guard.fail_result_limit();
3679 } else {
3680 guard.fail();
3681 }
3682 Err(error)
3683 }
3684 }
3685 }
3686
3687 pub async fn run_with_query_for_serialization(
3690 &self,
3691 sql: &str,
3692 query: RegisteredSqlQuery,
3693 ) -> Result<ManagedQueryBatches> {
3694 self.run_with_query_for_serialization_with_limits(
3695 sql,
3696 query,
3697 SqlCollectionLimits::default(),
3698 )
3699 .await
3700 }
3701
3702 pub async fn run_with_query_for_serialization_with_limits(
3704 &self,
3705 sql: &str,
3706 query: RegisteredSqlQuery,
3707 limits: SqlCollectionLimits,
3708 ) -> Result<ManagedQueryBatches> {
3709 query.set_sql_metadata(sql);
3710 let guard = RegisteredQueryGuard::new(query.clone());
3711 self.fire_test_hook(SqlTestHookPoint::WaitingForSessionPermit);
3712 let permit = tokio::select! {
3713 permit = Arc::clone(&self.execution_gate).acquire_owned() => permit.map_err(|_| {
3714 MongrelQueryError::InvalidQueryState("session execution gate closed".into())
3715 })?,
3716 _ = query.control().cancelled() => {
3717 let error = query.checkpoint().unwrap_err();
3718 guard.fail();
3719 return Err(error);
3720 }
3721 };
3722 query.transition(SqlQueryPhase::Queued, SqlQueryPhase::Planning)?;
3723 self.fire_test_hook(SqlTestHookPoint::Planning);
3724 query.transition(SqlQueryPhase::Planning, SqlQueryPhase::Executing)?;
3725 query.begin_statement(0);
3726 let result = CURRENT_COLLECTION_LIMITS
3727 .scope(
3728 limits,
3729 CURRENT_SQL_QUERY.scope(query.clone(), async {
3730 tokio::select! {
3731 biased;
3732 result = self.run_internal(sql, &query) => result,
3733 _ = query.control().cancelled() => Err(query.checkpoint().unwrap_err()),
3734 }
3735 }),
3736 )
3737 .await;
3738 match result {
3739 Ok(batches) => {
3740 query.complete_current_statement();
3741 query.begin_serialization()?;
3742 Ok(ManagedQueryBatches {
3743 batches,
3744 query,
3745 guard: Some(guard),
3746 permit: Some(permit),
3747 })
3748 }
3749 Err(error) => {
3750 if matches!(error, MongrelQueryError::ResultLimitExceeded { .. }) {
3751 guard.fail_result_limit();
3752 } else {
3753 guard.fail();
3754 }
3755 Err(error)
3756 }
3757 }
3758 }
3759
3760 pub async fn run(&self, sql: &str) -> Result<Vec<RecordBatch>> {
3761 if let Ok(query) = CURRENT_SQL_QUERY.try_with(Clone::clone) {
3762 return Box::pin(self.run_internal(sql, &query)).await;
3763 }
3764 Box::pin(self.run_with_options(sql, SqlQueryOptions::default())).await
3765 }
3766
3767 async fn run_internal(
3772 &self,
3773 sql: &str,
3774 query: &RegisteredSqlQuery,
3775 ) -> Result<Vec<RecordBatch>> {
3776 let trimmed = sql.trim();
3777 let multiple_statements = trimmed.contains(';')
3778 && !is_trigger_body(trimmed)
3779 && split_sql_statements(trimmed)
3780 .iter()
3781 .filter(|statement| !statement.trim().is_empty())
3782 .count()
3783 > 1;
3784 let rollback = commands::rollback_kind(sql);
3785 let (transaction_open, transaction_aborted) = {
3786 let transaction = self.transaction.lock();
3787 (transaction.staged_ops.is_some(), transaction.aborted)
3788 };
3789 if !multiple_statements && transaction_open && transaction_aborted && rollback.is_none() {
3790 return Err(MongrelQueryError::TransactionAborted);
3791 }
3792 let statement = TransactionStatementGuard::new(
3793 self,
3794 query,
3795 !multiple_statements
3796 && transaction_open
3797 && rollback != Some(commands::RollbackKind::Full),
3798 )?;
3799 self.fire_test_hook(SqlTestHookPoint::BeforeStatementExecution);
3800 let result = Box::pin(self.run_internal_body(sql, query)).await;
3801 if let Ok(batches) = result.as_ref() {
3802 enforce_collected_limits(batches, query, current_collection_limits())?;
3803 }
3804 if result.is_ok() {
3805 statement.complete();
3806 }
3807 result
3808 }
3809
3810 async fn run_internal_body(
3811 &self,
3812 sql: &str,
3813 query: &RegisteredSqlQuery,
3814 ) -> Result<Vec<RecordBatch>> {
3815 query.checkpoint()?;
3816 if let Some(inner) = strip_explain_query_plan(sql) {
3817 return self.explain_query_plan(inner, query).await;
3818 }
3819
3820 let trimmed = sql.trim();
3826 if trimmed.contains(';') && !is_trigger_body(trimmed) {
3827 let stmts = split_sql_statements(trimmed);
3828 let non_empty: Vec<&String> = stmts
3830 .iter()
3831 .filter(|s| !s.trim().is_empty() && s.trim() != ";")
3832 .collect();
3833 if non_empty.is_empty() {
3834 return Ok(Vec::new());
3835 }
3836 if stmts.len() > 1 {
3837 let mut last = Vec::new();
3838 let mut statement_index = 0;
3839 for stmt in &stmts {
3840 let stmt = stmt.trim();
3841 if stmt.is_empty() {
3842 continue;
3843 }
3844 query.begin_statement(statement_index);
3845 last = Box::pin(self.run_internal(stmt, query)).await?;
3846 query.complete_statement(statement_index);
3847 statement_index += 1;
3848 }
3849 return Ok(last);
3850 }
3851 }
3852
3853 query.checkpoint()?;
3855 if let Some(batches) = commands::try_run_command(self, sql, query).await? {
3856 return Ok(batches);
3857 }
3858
3859 let trimmed = sql.trim();
3865 if trimmed.contains(';') {
3866 let stmts = split_sql_statements(trimmed);
3867 if stmts.len() > 1 {
3868 let mut last = Vec::new();
3869 let mut statement_index = 0;
3870 for stmt in &stmts {
3871 let stmt = stmt.trim();
3872 if stmt.is_empty() {
3873 continue;
3874 }
3875 query.begin_statement(statement_index);
3876 last = Box::pin(self.run_internal(stmt, query)).await?;
3877 query.complete_statement(statement_index);
3878 statement_index += 1;
3879 }
3880 return Ok(last);
3881 }
3882 }
3883 let lower = sql.trim_start().to_lowercase();
3885 if lower.starts_with("create table") {
3886 if let Some(db) = &self.database {
3887 db.require_for(self.principal().as_ref(), &mongreldb_core::Permission::Ddl)?;
3888 let (name, schema) = parse_create_table(sql)?;
3889 db.create_table(&name, schema)?;
3890 let handle = db.table(&name)?;
3891 let provider = MongrelProvider::new_secured(
3892 handle.clone(),
3893 Arc::clone(db),
3894 name.clone(),
3895 self.principal.clone(),
3896 )?;
3897 self.ctx
3898 .register_table(&name, Arc::new(provider))
3899 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
3900 self.tables.lock().insert(name, handle);
3901 self.invalidate_prepared_statements(query).await;
3902 self.clear_cache();
3903 return Ok(Vec::new());
3904 }
3905 }
3906 if lower.starts_with("drop table") {
3907 if let Some(db) = &self.database {
3908 db.require_for(self.principal().as_ref(), &mongreldb_core::Permission::Ddl)?;
3909 let (name, if_exists) = parse_drop_table(sql)?;
3910 let drop_result = db.drop_table(&name);
3911 if let Err(e) = drop_result {
3912 let is_not_found = matches!(e, mongreldb_core::MongrelError::NotFound(_));
3914 if !(if_exists && is_not_found) {
3915 return Err(e.into());
3916 }
3917 } else {
3918 self.ctx
3919 .deregister_table(&name)
3920 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
3921 self.tables.lock().remove(&name);
3922 }
3923 self.invalidate_prepared_statements(query).await;
3924 self.clear_cache();
3925 return Ok(Vec::new());
3926 }
3927 }
3928 if lower.starts_with("alter table") {
3929 if let Some(db) = &self.database {
3930 db.require_for(self.principal().as_ref(), &mongreldb_core::Permission::Ddl)?;
3931 match parse_alter_table(sql)? {
3932 ParsedAlterTable::RenameTable { old_name, new_name } => {
3933 db.rename_table(&old_name, &new_name)?;
3934 self.ctx
3938 .deregister_table(&old_name)
3939 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
3940 self.tables.lock().remove(&old_name);
3941 let handle = db.table(&new_name)?;
3942 let provider = MongrelProvider::new_secured(
3943 handle.clone(),
3944 Arc::clone(db),
3945 new_name.clone(),
3946 self.principal.clone(),
3947 )?;
3948 self.ctx
3949 .register_table(&new_name, Arc::new(provider))
3950 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
3951 self.tables.lock().insert(new_name, handle);
3952 }
3953 ParsedAlterTable::RenameColumn {
3954 table_name,
3955 column_name,
3956 new_name,
3957 } => {
3958 db.alter_column(&table_name, &column_name, AlterColumn::rename(new_name))?;
3959 self.refresh_registered_table(db, &table_name)?;
3960 }
3961 ParsedAlterTable::AlterColumnType {
3962 table_name,
3963 column_name,
3964 ty,
3965 } => {
3966 db.alter_column(&table_name, &column_name, AlterColumn::set_type(ty))?;
3967 self.refresh_registered_table(db, &table_name)?;
3968 }
3969 ParsedAlterTable::SetNotNull {
3970 table_name,
3971 column_name,
3972 } => {
3973 let flags = current_column_flags(db, &table_name, &column_name)?
3974 .without(ColumnFlags::NULLABLE);
3975 db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
3976 self.refresh_registered_table(db, &table_name)?;
3977 }
3978 ParsedAlterTable::DropNotNull {
3979 table_name,
3980 column_name,
3981 } => {
3982 let flags = current_column_flags(db, &table_name, &column_name)?
3983 .with(ColumnFlags::NULLABLE);
3984 db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
3985 self.refresh_registered_table(db, &table_name)?;
3986 }
3987 }
3988 self.invalidate_prepared_statements(query).await;
3989 self.clear_cache();
3990 return Ok(Vec::new());
3991 }
3992 }
3993
3994 if let Some(as_of_query) = self.prepare_as_of_query(sql)? {
3995 return self.run_as_of(as_of_query, query).await;
3996 }
3997
3998 let resolved = self.resolve_view_sql(sql);
4001 let resolved = self.rewrite_external_module_compat_sql(&resolved);
4002 let resolved = rewrite_compat_function_calls(&resolved);
4003 let effective_sql = normalize_sql(&resolved);
4007 let sql = effective_sql.as_str();
4008 let epoch = self.cache_epoch();
4012 let key = (sql.to_string(), epoch);
4013 let has_ttl_table = self
4014 .tables
4015 .lock()
4016 .values()
4017 .any(|table| table.lock().ttl().is_some());
4018 let result_cacheable = !extended_sql_functions::contains_volatile_extended_function(sql)
4019 && !is_explain_analyze(sql)
4020 && !is_prepared_stmt_sql(sql)
4021 && !has_ttl_table;
4022 if result_cacheable {
4023 if let Some(hit) = self.cache.lock().get(&key) {
4024 return Ok((**hit).clone());
4025 }
4026 }
4027 if let Some(batches) = self.try_catalog_introspection(sql)? {
4030 if result_cacheable {
4031 query.checkpoint()?;
4032 self.cache.lock().insert(key, Arc::new(batches.clone()));
4033 }
4034 return Ok(batches);
4035 }
4036 if let Some(batches) = self.try_direct_dispatch(sql, query)? {
4041 if result_cacheable {
4042 query.checkpoint()?;
4043 self.cache.lock().insert(key, Arc::new(batches.clone()));
4044 }
4045 return Ok(batches);
4046 }
4047 let plan_start = std::time::Instant::now();
4049 let external_module_scan = self.query_references_external_module(sql);
4050 let df = self.dataframe_with_query(sql, query).await?;
4051 self.track_prepared_name(sql);
4054 let planning_nanos = plan_start.elapsed().as_nanos() as u64;
4057 mongreldb_core::trace::QueryTrace::record(|t| t.planning_nanos = planning_nanos);
4058
4059 let agg_key = sql_cache_key(sql);
4064 let batches = match self.try_native_aggregate(df.logical_plan(), agg_key, query) {
4065 Ok(Some(batch)) => vec![batch],
4066 _ => {
4067 match self.try_fk_join(df.logical_plan(), query) {
4071 Ok(Some(b)) => {
4072 mongreldb_core::trace::QueryTrace::record(|t| {
4074 t.join_mode = mongreldb_core::trace::JoinMode::FkBitmap;
4075 });
4076 b
4077 }
4078 _ => self.collect_dataframe(df, query).await?,
4079 }
4080 }
4081 };
4082 if external_module_scan {
4083 mongreldb_core::trace::QueryTrace::record(|t| {
4084 t.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
4085 });
4086 }
4087 if result_cacheable {
4088 query.checkpoint()?;
4089 self.cache.lock().insert(key, Arc::new(batches.clone()));
4090 }
4091 Ok(batches)
4092 }
4093
4094 pub async fn run_sql_traced(
4106 &self,
4107 sql: &str,
4108 ) -> Result<(Vec<RecordBatch>, mongreldb_core::trace::QueryTrace)> {
4109 mongreldb_core::trace::QueryTrace::push_scope();
4110 let result = self.run(sql).await;
4111 let trace = mongreldb_core::trace::QueryTrace::pop_scope();
4112 Ok((result?, trace))
4113 }
4114
4115 pub fn clear_cache(&self) {
4118 self.cache.lock().clear();
4119 self.plan_cache.lock().clear();
4120 }
4121
4122 fn track_prepared_name(&self, sql: &str) {
4126 let lower = sql.trim_start().to_ascii_lowercase();
4127 if let Some(rest) = lower.strip_prefix("prepare ") {
4128 if let Some(name) = parse_stmt_ident(rest) {
4129 self.prepared_names.lock().insert(name);
4130 }
4131 } else if let Some(rest) = lower.strip_prefix("deallocate ") {
4132 if let Some(name) = parse_stmt_ident(rest) {
4133 self.prepared_names.lock().remove(&name);
4134 }
4135 }
4136 }
4137
4138 async fn invalidate_prepared_statements(&self, query: &RegisteredSqlQuery) {
4145 let names: Vec<String> = self.prepared_names.lock().drain().collect();
4146 for name in names {
4147 let sql = format!("DEALLOCATE {name}");
4151 if self.dataframe_with_query(&sql, query).await.is_err() {
4152 }
4154 }
4155 }
4156
4157 async fn explain_query_plan(
4158 &self,
4159 sql: &str,
4160 query: &RegisteredSqlQuery,
4161 ) -> Result<Vec<RecordBatch>> {
4162 let explain_sql = format!("EXPLAIN {}", sql.trim().trim_end_matches(';'));
4163 let dataframe = self.dataframe_with_query(&explain_sql, query).await?;
4164 let batches = self.collect_dataframe(dataframe, query).await?;
4165 let mut detail = self.mongrel_query_plan_details(sql);
4166 for batch in &batches {
4167 if batch.num_columns() < 2 {
4168 continue;
4169 }
4170 let Some(plan_type) = batch.column(0).as_any().downcast_ref::<StringArray>() else {
4171 continue;
4172 };
4173 let Some(plan) = batch.column(1).as_any().downcast_ref::<StringArray>() else {
4174 continue;
4175 };
4176 for row in 0..batch.num_rows() {
4177 let prefix = plan_type.value(row);
4178 for line in plan.value(row).lines() {
4179 let line = line.trim();
4180 if !line.is_empty() {
4181 detail.push(format!("DATAFUSION {prefix}: {line}"));
4182 }
4183 }
4184 }
4185 }
4186 if detail.is_empty() {
4187 detail.push("plan unavailable".to_string());
4188 }
4189 let ids = (0..detail.len()).map(|i| i as i64).collect::<Vec<_>>();
4190 let parents = vec![0_i64; detail.len()];
4191 let notused = vec![0_i64; detail.len()];
4192 let schema = Arc::new(arrow::datatypes::Schema::new(vec![
4193 arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int64, false),
4194 arrow::datatypes::Field::new("parent", arrow::datatypes::DataType::Int64, false),
4195 arrow::datatypes::Field::new("notused", arrow::datatypes::DataType::Int64, false),
4196 arrow::datatypes::Field::new("detail", arrow::datatypes::DataType::Utf8, false),
4197 ]));
4198 let batch = RecordBatch::try_new(
4199 schema,
4200 vec![
4201 Arc::new(Int64Array::from(ids)) as ArrayRef,
4202 Arc::new(Int64Array::from(parents)),
4203 Arc::new(Int64Array::from(notused)),
4204 Arc::new(StringArray::from(detail)),
4205 ],
4206 )
4207 .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
4208 Ok(vec![batch])
4209 }
4210
4211 fn mongrel_query_plan_details(&self, sql: &str) -> Vec<String> {
4212 use sqlparser::ast::{GroupByExpr, OrderByKind, SetExpr, Statement};
4213 use sqlparser::dialect::PostgreSqlDialect;
4214 use sqlparser::parser::Parser;
4215
4216 let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
4217 return Vec::new();
4218 };
4219 let Some(Statement::Query(query)) = stmts.first() else {
4220 return Vec::new();
4221 };
4222
4223 fn collect(session: &MongrelSession, query: &sqlparser::ast::Query, out: &mut Vec<String>) {
4224 use sqlparser::ast::{SetOperator, TableWithJoins};
4225 match query.body.as_ref() {
4226 SetExpr::Select(select) => {
4227 for TableWithJoins { relation, joins } in &select.from {
4228 session.push_table_plan(relation, select.selection.as_ref(), out);
4229 for join in joins {
4230 session.push_table_plan(&join.relation, None, out);
4231 }
4232 }
4233 if select.distinct.is_some() {
4234 out.push("USE TEMP B-TREE FOR DISTINCT".to_string());
4235 }
4236 let grouped = match &select.group_by {
4237 GroupByExpr::All(_) => true,
4238 GroupByExpr::Expressions(exprs, _) => !exprs.is_empty(),
4239 };
4240 if grouped {
4241 out.push("USE TEMP B-TREE FOR GROUP BY".to_string());
4242 }
4243 let ordered = query.order_by.as_ref().is_some_and(|order_by| {
4244 matches!(order_by.kind, OrderByKind::All(_))
4245 || matches!(&order_by.kind, OrderByKind::Expressions(exprs) if !exprs.is_empty())
4246 });
4247 if ordered {
4248 out.push("USE TEMP B-TREE FOR ORDER BY".to_string());
4249 }
4250 }
4251 SetExpr::Query(query) => collect(session, query, out),
4252 SetExpr::SetOperation {
4253 left, op, right, ..
4254 } => {
4255 let label = match op {
4256 SetOperator::Union => "COMPOUND QUERY UNION",
4257 SetOperator::Except => "COMPOUND QUERY EXCEPT",
4258 SetOperator::Intersect => "COMPOUND QUERY INTERSECT",
4259 _ => "COMPOUND QUERY",
4260 };
4261 out.push(label.to_string());
4262 collect_set_expr(session, left, out);
4263 collect_set_expr(session, right, out);
4264 }
4265 _ => {}
4266 }
4267 }
4268
4269 fn collect_set_expr(session: &MongrelSession, expr: &SetExpr, out: &mut Vec<String>) {
4270 match expr {
4271 SetExpr::Select(select) => {
4272 for table in &select.from {
4273 session.push_table_plan(&table.relation, select.selection.as_ref(), out);
4274 }
4275 }
4276 SetExpr::Query(query) => collect(session, query, out),
4277 SetExpr::SetOperation { left, right, .. } => {
4278 collect_set_expr(session, left, out);
4279 collect_set_expr(session, right, out);
4280 }
4281 _ => {}
4282 }
4283 }
4284
4285 let mut out = Vec::new();
4286 collect(self, query, &mut out);
4287 out
4288 }
4289
4290 fn push_table_plan(
4291 &self,
4292 relation: &sqlparser::ast::TableFactor,
4293 selection: Option<&sqlparser::ast::Expr>,
4294 out: &mut Vec<String>,
4295 ) {
4296 let sqlparser::ast::TableFactor::Table { name, alias, .. } = relation else {
4297 out.push("SCAN SUBQUERY".to_string());
4298 return;
4299 };
4300 let table_name = name.to_string();
4301 let display_name = alias
4302 .as_ref()
4303 .map(|alias| alias.name.value.clone())
4304 .unwrap_or_else(|| table_name.clone());
4305 let Some(handle) = self.tables.lock().get(&table_name).cloned() else {
4306 out.push(format!("SCAN {display_name}"));
4307 return;
4308 };
4309 let schema = handle.lock().schema().clone();
4310 let searchable = selection
4311 .and_then(|expr| translate_sqlparser_filter(expr, &schema))
4312 .is_some_and(|conditions| !conditions.is_empty());
4313 if searchable {
4314 out.push(format!("SEARCH {display_name} USING MONGREL INDEX"));
4315 } else {
4316 out.push(format!("SCAN {display_name}"));
4317 }
4318 }
4319
4320 fn resolve_view_sql(&self, sql: &str) -> String {
4325 let views = self.views.lock();
4326 if views.is_empty() {
4327 return sql.to_string();
4328 }
4329 let mut result = sql.to_string();
4330 for (name, view) in views.iter() {
4331 result = replace_from_view(&result, name, &view.sql);
4332 }
4333 result
4334 }
4335
4336 fn rewrite_external_module_compat_sql(&self, sql: &str) -> String {
4337 let Some(db) = &self.database else {
4338 return sql.to_string();
4339 };
4340 rewrite_fts_match_compat_sql(sql, db)
4341 }
4342
4343 fn query_references_external_module(&self, sql: &str) -> bool {
4344 use sqlparser::ast::Statement;
4345 use sqlparser::dialect::PostgreSqlDialect;
4346 use sqlparser::parser::Parser;
4347
4348 Parser::parse_sql(&PostgreSqlDialect {}, sql)
4349 .ok()
4350 .is_some_and(|statements| {
4351 statements.iter().any(|statement| match statement {
4352 Statement::Query(query) => self.query_uses_external_module(query),
4353 _ => false,
4354 })
4355 })
4356 }
4357
4358 fn query_uses_external_module(&self, query: &sqlparser::ast::Query) -> bool {
4359 self.set_expr_uses_external_module(query.body.as_ref())
4360 }
4361
4362 fn set_expr_uses_external_module(&self, expr: &sqlparser::ast::SetExpr) -> bool {
4363 use sqlparser::ast::SetExpr;
4364
4365 match expr {
4366 SetExpr::Select(select) => select
4367 .from
4368 .iter()
4369 .any(|table| self.table_with_joins_uses_external_module(table)),
4370 SetExpr::Query(query) => self.query_uses_external_module(query),
4371 SetExpr::SetOperation { left, right, .. } => {
4372 self.set_expr_uses_external_module(left)
4373 || self.set_expr_uses_external_module(right)
4374 }
4375 _ => false,
4376 }
4377 }
4378
4379 fn table_with_joins_uses_external_module(
4380 &self,
4381 table: &sqlparser::ast::TableWithJoins,
4382 ) -> bool {
4383 self.table_factor_uses_external_module(&table.relation)
4384 || table
4385 .joins
4386 .iter()
4387 .any(|join| self.table_factor_uses_external_module(&join.relation))
4388 }
4389
4390 fn table_factor_uses_external_module(&self, relation: &sqlparser::ast::TableFactor) -> bool {
4391 use sqlparser::ast::{Expr, TableFactor};
4392
4393 match relation {
4394 TableFactor::Table { name, args, .. } => {
4395 let table_name = name.to_string();
4396 self.database
4397 .as_ref()
4398 .is_some_and(|db| db.external_table(&table_name).is_some())
4399 || (args.is_some() && self.external_modules.contains(&table_name))
4400 }
4401 TableFactor::Function { name, .. } => self.external_modules.contains(&name.to_string()),
4402 TableFactor::TableFunction {
4403 expr: Expr::Function(func),
4404 ..
4405 } => self.external_modules.contains(&func.name.to_string()),
4406 TableFactor::Derived { subquery, .. } => self.query_uses_external_module(subquery),
4407 _ => false,
4408 }
4409 }
4410
4411 fn cache_epoch(&self) -> u64 {
4414 if let Some(db) = &self.database {
4415 db.visible_epoch().0
4416 } else {
4417 self.combined_epoch()
4418 }
4419 }
4420
4421 fn combined_epoch(&self) -> u64 {
4422 let Some(primary) = self.db.as_ref() else {
4423 return 0;
4424 };
4425 let mut combined = primary.lock().snapshot().epoch.0;
4426 let tables = self.tables.lock();
4427 for arc in tables.values() {
4428 if !arc.ptr_eq(primary) {
4429 let e = arc.lock().snapshot().epoch.0;
4430 combined = combined.wrapping_mul(31).wrapping_add(e);
4431 }
4432 }
4433 combined
4434 }
4435
4436 fn try_native_aggregate(
4440 &self,
4441 plan: &datafusion::logical_expr::LogicalPlan,
4442 cache_key: u64,
4443 query: &RegisteredSqlQuery,
4444 ) -> Result<Option<RecordBatch>> {
4445 if self.security_context_active() {
4446 return Ok(None);
4447 }
4448 let Some(primary) = self.db.as_ref() else {
4449 return Ok(None);
4450 };
4451 let mut db = primary.lock();
4452 let schema = db.schema().clone();
4453 let snap = db.snapshot();
4454 native_agg::try_native_aggregate(&mut db, &schema, snap, plan, cache_key, Some(query))
4455 }
4456
4457 fn try_fk_join(
4461 &self,
4462 plan: &datafusion::logical_expr::LogicalPlan,
4463 query: &RegisteredSqlQuery,
4464 ) -> Result<Option<Vec<RecordBatch>>> {
4465 if self.security_context_active() {
4466 return Ok(None);
4467 }
4468 let tables = self.tables.lock();
4469 fk_join::try_fk_join(&tables, plan, Some(query), current_collection_limits())
4470 }
4471
4472 pub fn context(&self) -> &SessionContext {
4473 &self.ctx
4474 }
4475
4476 pub fn register_scalar_udf(&self, f: ScalarUDF) {
4482 self.ctx.register_udf(f);
4483 self.clear_cache();
4484 }
4485
4486 pub fn register_aggregate_udf(&self, f: AggregateUDF) {
4488 self.ctx.register_udaf(f);
4489 self.clear_cache();
4490 }
4491
4492 pub fn register_window_udf(&self, f: WindowUDF) {
4494 self.ctx.register_udwf(f);
4495 self.clear_cache();
4496 }
4497}
4498
4499fn register_mongrel_functions(
4500 ctx: &SessionContext,
4501 sql_fn_state: Arc<extended_sql_functions::ExtendedSqlState>,
4502) {
4503 ctx.register_udf(ScalarUDF::from(udf::AnnSearchUdf::new()));
4504 ctx.register_udf(ScalarUDF::from(udf::SparseMatchUdf::new()));
4505 ctx.register_udf(ScalarUDF::from(udf::RTreeIntersectsUdf::new()));
4506 ctx.register_udf(ScalarUDF::from(udf::FtsRankUdf::new()));
4507 for udaf in percentile::percentile_udafs() {
4508 ctx.register_udaf(udaf);
4509 }
4510 extended_sql_functions::register_extended_sql_functions_with_state(ctx, sql_fn_state);
4511}
4512
4513fn is_trigger_body(sql: &str) -> bool {
4516 let lower = sql.to_ascii_lowercase();
4517 if !lower.contains("create trigger") && !lower.contains("create or replace trigger") {
4518 return false;
4519 }
4520 lower.contains("begin") && lower.contains("end")
4521}
4522
4523fn split_sql_statements(sql: &str) -> Vec<String> {
4528 let b = sql.as_bytes();
4529 let n = b.len();
4530 let mut stmts = Vec::new();
4531 let mut current = String::new();
4532 let mut i = 0;
4533 while i < n {
4534 let c = b[i];
4535 if c == b'-' && i + 1 < n && b[i + 1] == b'-' {
4537 while i < n && b[i] != b'\n' {
4538 current.push(c as char);
4539 i += 1;
4540 }
4541 continue;
4542 }
4543 if c == b'/' && i + 1 < n && b[i + 1] == b'*' {
4545 current.push_str("/*");
4546 i += 2;
4547 let mut depth = 1;
4548 while i + 1 < n && depth > 0 {
4549 if b[i] == b'/' && b[i + 1] == b'*' {
4550 depth += 1;
4551 current.push_str("/*");
4552 i += 2;
4553 } else if b[i] == b'*' && b[i + 1] == b'/' {
4554 depth -= 1;
4555 current.push_str("*/");
4556 i += 2;
4557 } else {
4558 current.push(b[i] as char);
4559 i += 1;
4560 }
4561 }
4562 continue;
4563 }
4564 if c == b'\'' {
4566 current.push('\'');
4567 i += 1;
4568 while i < n {
4569 current.push(b[i] as char);
4570 if b[i] == b'\'' {
4571 if i + 1 < n && b[i + 1] == b'\'' {
4573 current.push('\'');
4574 i += 2;
4575 continue;
4576 }
4577 i += 1;
4578 break;
4579 }
4580 i += 1;
4581 }
4582 continue;
4583 }
4584 if c == b'"' {
4586 current.push('"');
4587 i += 1;
4588 while i < n {
4589 current.push(b[i] as char);
4590 if b[i] == b'"' {
4591 if i + 1 < n && b[i + 1] == b'"' {
4592 current.push('"');
4593 i += 2;
4594 continue;
4595 }
4596 i += 1;
4597 break;
4598 }
4599 i += 1;
4600 }
4601 continue;
4602 }
4603 if c == b'$' {
4605 let start = i;
4607 let mut tag_end = i + 1;
4608 while tag_end < n && b[tag_end] != b'$' && b[tag_end].is_ascii_alphanumeric() {
4609 tag_end += 1;
4610 }
4611 if tag_end < n && b[tag_end] == b'$' {
4612 let tag = &sql[start..=tag_end];
4613 current.push_str(tag);
4614 i = tag_end + 1;
4615 while i < n {
4617 if b[i] == b'$' && sql[i..].starts_with(tag) {
4618 current.push_str(tag);
4619 i += tag.len();
4620 break;
4621 }
4622 current.push(b[i] as char);
4623 i += 1;
4624 }
4625 continue;
4626 }
4627 }
4628 if c == b';' {
4630 current.push(';');
4631 let trimmed = current.trim();
4632 if !trimmed.is_empty() && trimmed != ";" {
4633 stmts.push(current.clone());
4634 }
4635 current.clear();
4636 i += 1;
4637 continue;
4638 }
4639 current.push(c as char);
4640 i += 1;
4641 }
4642 let trimmed = current.trim();
4644 if !trimmed.is_empty() {
4645 stmts.push(current);
4646 }
4647 stmts
4648}
4649
4650fn strip_explain_query_plan(sql: &str) -> Option<&str> {
4651 let trimmed = sql.trim_start();
4652 let lower = trimmed.to_ascii_lowercase();
4653 if !lower.starts_with("explain") {
4654 return None;
4655 }
4656 let after_explain = trimmed.get(7..)?.trim_start();
4657 let after_explain_lower = after_explain.to_ascii_lowercase();
4658 if !after_explain_lower.starts_with("query") {
4659 return None;
4660 }
4661 let after_query = after_explain.get(5..)?.trim_start();
4662 let after_query_lower = after_query.to_ascii_lowercase();
4663 if !after_query_lower.starts_with("plan") {
4664 return None;
4665 }
4666 Some(after_query.get(4..)?.trim_start())
4667}
4668
4669fn is_prepared_stmt_sql(sql: &str) -> bool {
4675 let lower = sql.trim_start().to_ascii_lowercase();
4676 lower.starts_with("prepare ")
4677 || lower.starts_with("execute ")
4678 || lower.starts_with("deallocate ")
4679}
4680
4681fn parse_stmt_ident(s: &str) -> Option<String> {
4685 let tok = s
4686 .trim_start()
4687 .split(|c: char| c.is_whitespace() || c == '(')
4688 .next()?;
4689 let t = tok.trim_matches(|c| c == '"' || c == '`' || c == '\'');
4690 (!t.is_empty()).then(|| t.to_string())
4691}
4692
4693fn is_explain_analyze(sql: &str) -> bool {
4698 let trimmed = sql.trim_start();
4699 let lower = trimmed.to_ascii_lowercase();
4700 if !lower.starts_with("explain") {
4701 return false;
4702 }
4703 let after = trimmed[7..].trim_start();
4704 after.to_ascii_lowercase().starts_with("analyze")
4705}
4706
4707#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4708enum SqlCompatTokenKind {
4709 Ident,
4710 String,
4711 Dot,
4712 LParen,
4713 RParen,
4714 Comma,
4715}
4716
4717#[derive(Debug, Clone)]
4718struct SqlCompatToken {
4719 kind: SqlCompatTokenKind,
4720 raw: String,
4721 normalized: String,
4722 start: usize,
4723 end: usize,
4724}
4725
4726#[derive(Debug, Clone)]
4727struct FtsMatchBinding {
4728 query_ref: String,
4729}
4730
4731#[derive(Debug, Clone)]
4732struct SqlReplacement {
4733 start: usize,
4734 end: usize,
4735 replacement: String,
4736}
4737
4738fn rewrite_fts_match_compat_sql(sql: &str, db: &Database) -> String {
4739 let tokens = sql_compat_tokens(sql);
4740 if tokens.is_empty() {
4741 return sql.to_string();
4742 }
4743 let bindings = fts_match_bindings(sql, db, &tokens);
4744 if bindings.is_empty() {
4745 return sql.to_string();
4746 }
4747 let unique_refs = bindings
4748 .values()
4749 .map(|binding| binding.query_ref.as_str())
4750 .collect::<HashSet<_>>();
4751 let unique_binding = if unique_refs.len() == 1 {
4752 bindings.values().next().cloned()
4753 } else {
4754 None
4755 };
4756 let mut replacements = Vec::new();
4757 for (idx, token) in tokens.iter().enumerate() {
4758 if token.kind != SqlCompatTokenKind::Ident || token.normalized != "match" {
4759 continue;
4760 }
4761 let Some(rhs) = tokens.get(idx + 1) else {
4762 continue;
4763 };
4764 if rhs.kind != SqlCompatTokenKind::String {
4765 continue;
4766 }
4767 let Some((lhs_start, _lhs_end, query_ref)) =
4768 fts_match_lhs_query_ref(&tokens, idx, &bindings, unique_binding.as_ref())
4769 else {
4770 continue;
4771 };
4772 replacements.push(SqlReplacement {
4773 start: lhs_start,
4774 end: rhs.end,
4775 replacement: format!("{query_ref}.query = {}", rhs.raw),
4776 });
4777 }
4778 apply_sql_replacements(sql, &replacements)
4779}
4780
4781fn fts_match_lhs_query_ref(
4782 tokens: &[SqlCompatToken],
4783 match_idx: usize,
4784 bindings: &HashMap<String, FtsMatchBinding>,
4785 unique_binding: Option<&FtsMatchBinding>,
4786) -> Option<(usize, usize, String)> {
4787 if match_idx == 0 {
4788 return None;
4789 }
4790 let lhs = tokens.get(match_idx - 1)?;
4791 if lhs.kind != SqlCompatTokenKind::Ident {
4792 return None;
4793 }
4794
4795 if match_idx >= 3
4796 && tokens.get(match_idx - 2)?.kind == SqlCompatTokenKind::Dot
4797 && tokens.get(match_idx - 3)?.kind == SqlCompatTokenKind::Ident
4798 {
4799 let owner = tokens.get(match_idx - 3)?;
4800 let binding = bindings.get(&owner.normalized)?;
4801 if lhs.normalized == "query" || lhs.normalized == "text" {
4802 return Some((owner.start, lhs.end, binding.query_ref.clone()));
4803 }
4804 return None;
4805 }
4806
4807 if let Some(binding) = bindings.get(&lhs.normalized) {
4808 return Some((lhs.start, lhs.end, binding.query_ref.clone()));
4809 }
4810 if lhs.normalized == "text" {
4811 let binding = unique_binding?;
4812 return Some((lhs.start, lhs.end, binding.query_ref.clone()));
4813 }
4814 None
4815}
4816
4817fn fts_match_bindings(
4818 sql: &str,
4819 db: &Database,
4820 tokens: &[SqlCompatToken],
4821) -> HashMap<String, FtsMatchBinding> {
4822 let mut out = HashMap::new();
4823 let mut i = 0;
4824 while i < tokens.len() {
4825 let token = &tokens[i];
4826 let starts_table_ref = token.kind == SqlCompatTokenKind::Ident
4827 && matches!(token.normalized.as_str(), "from" | "join");
4828 if !starts_table_ref {
4829 i += 1;
4830 continue;
4831 }
4832 let mut table_idx = i + 1;
4833 if tokens
4834 .get(table_idx)
4835 .is_some_and(|token| token.kind == SqlCompatTokenKind::LParen)
4836 {
4837 i += 1;
4838 continue;
4839 }
4840 let Some(table) = tokens.get(table_idx) else {
4841 break;
4842 };
4843 if table.kind != SqlCompatTokenKind::Ident {
4844 i += 1;
4845 continue;
4846 }
4847 let mut table_name = table.normalized.clone();
4848 let mut table_ref = table.raw.clone();
4849 if let Some(qualified) = tokens.get(table_idx + 2).filter(|qualified| {
4850 tokens
4851 .get(table_idx + 1)
4852 .is_some_and(|token| token.kind == SqlCompatTokenKind::Dot)
4853 && qualified.kind == SqlCompatTokenKind::Ident
4854 }) {
4855 table_name = qualified.normalized.clone();
4856 table_ref = sql[table.start..qualified.end].to_string();
4857 table_idx += 2;
4858 }
4859 if !is_fts_docs_table(db, &table_name) {
4860 i = table_idx + 1;
4861 continue;
4862 }
4863 let mut query_ref = table_ref.clone();
4864 let mut alias_key = None;
4865 let mut next = table_idx + 1;
4866 if tokens.get(next).is_some_and(|token| {
4867 token.kind == SqlCompatTokenKind::Ident && token.normalized == "as"
4868 }) {
4869 next += 1;
4870 }
4871 if let Some(alias) = tokens.get(next) {
4872 if alias.kind == SqlCompatTokenKind::Ident && !is_table_ref_boundary(&alias.normalized)
4873 {
4874 alias_key = Some(alias.normalized.clone());
4875 query_ref = alias.raw.clone();
4876 next += 1;
4877 }
4878 }
4879 out.insert(
4880 table_name,
4881 FtsMatchBinding {
4882 query_ref: query_ref.clone(),
4883 },
4884 );
4885 if let Some(alias_key) = alias_key {
4886 out.insert(alias_key, FtsMatchBinding { query_ref });
4887 }
4888 i = next;
4889 }
4890 out
4891}
4892
4893fn is_fts_docs_table(db: &Database, name: &str) -> bool {
4894 db.external_table(name)
4895 .is_some_and(|entry| entry.module == "fts_docs")
4896}
4897
4898fn is_table_ref_boundary(normalized: &str) -> bool {
4899 matches!(
4900 normalized,
4901 "where"
4902 | "join"
4903 | "left"
4904 | "right"
4905 | "inner"
4906 | "outer"
4907 | "full"
4908 | "cross"
4909 | "on"
4910 | "using"
4911 | "group"
4912 | "order"
4913 | "having"
4914 | "limit"
4915 | "offset"
4916 | "union"
4917 | "except"
4918 | "intersect"
4919 )
4920}
4921
4922fn sql_compat_tokens(sql: &str) -> Vec<SqlCompatToken> {
4923 let bytes = sql.as_bytes();
4924 let mut tokens = Vec::new();
4925 let mut i = 0;
4926 while i < bytes.len() {
4927 match bytes[i] {
4928 b if b.is_ascii_whitespace() => i += 1,
4929 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
4930 i += 2;
4931 while i < bytes.len() && bytes[i] != b'\n' {
4932 i += 1;
4933 }
4934 }
4935 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
4936 i = skip_block_comment(bytes, i);
4937 }
4938 b'\'' => {
4939 let end = skip_quoted(bytes, i, b'\'');
4940 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
4941 i = end;
4942 }
4943 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
4944 let end = skip_quoted(bytes, i + 1, b'\'');
4945 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
4946 i = end;
4947 }
4948 b'$' => {
4949 let (end, matched) = skip_dollar_quoted(bytes, i);
4950 if matched {
4951 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
4952 i = end;
4953 } else {
4954 i += 1;
4955 }
4956 }
4957 b'"' => {
4958 let end = skip_quoted(bytes, i, b'"');
4959 let raw = sql[i..end].to_string();
4960 let normalized = unquote_sql_ident(&raw).to_ascii_lowercase();
4961 tokens.push(SqlCompatToken {
4962 kind: SqlCompatTokenKind::Ident,
4963 raw,
4964 normalized,
4965 start: i,
4966 end,
4967 });
4968 i = end;
4969 }
4970 b'.' => {
4971 tokens.push(sql_token(sql, SqlCompatTokenKind::Dot, i, i + 1));
4972 i += 1;
4973 }
4974 b'(' => {
4975 tokens.push(sql_token(sql, SqlCompatTokenKind::LParen, i, i + 1));
4976 i += 1;
4977 }
4978 b')' => {
4979 tokens.push(sql_token(sql, SqlCompatTokenKind::RParen, i, i + 1));
4980 i += 1;
4981 }
4982 b',' => {
4983 tokens.push(sql_token(sql, SqlCompatTokenKind::Comma, i, i + 1));
4984 i += 1;
4985 }
4986 b if is_sql_ident_byte(b) => {
4987 let start = i;
4988 i += 1;
4989 while i < bytes.len() && is_sql_ident_byte(bytes[i]) {
4990 i += 1;
4991 }
4992 tokens.push(sql_token(sql, SqlCompatTokenKind::Ident, start, i));
4993 }
4994 _ => i += 1,
4995 }
4996 }
4997 tokens
4998}
4999
5000fn sql_token(sql: &str, kind: SqlCompatTokenKind, start: usize, end: usize) -> SqlCompatToken {
5001 let raw = sql[start..end].to_string();
5002 SqlCompatToken {
5003 kind,
5004 normalized: raw.to_ascii_lowercase(),
5005 raw,
5006 start,
5007 end,
5008 }
5009}
5010
5011fn unquote_sql_ident(raw: &str) -> String {
5012 if raw.len() >= 2 && raw.starts_with('"') && raw.ends_with('"') {
5013 raw[1..raw.len() - 1].replace("\"\"", "\"")
5014 } else {
5015 raw.to_string()
5016 }
5017}
5018
5019fn apply_sql_replacements(sql: &str, replacements: &[SqlReplacement]) -> String {
5020 if replacements.is_empty() {
5021 return sql.to_string();
5022 }
5023 let mut ordered = replacements.to_vec();
5024 ordered.sort_by_key(|replacement| replacement.start);
5025 let mut out = String::with_capacity(sql.len());
5026 let mut cursor = 0;
5027 for replacement in ordered {
5028 if replacement.start < cursor || replacement.end > sql.len() {
5029 continue;
5030 }
5031 out.push_str(&sql[cursor..replacement.start]);
5032 out.push_str(&replacement.replacement);
5033 cursor = replacement.end;
5034 }
5035 out.push_str(&sql[cursor..]);
5036 out
5037}
5038
5039fn rewrite_compat_function_calls(sql: &str) -> String {
5040 let bytes = sql.as_bytes();
5041 let mut out = String::with_capacity(sql.len());
5042 let mut i = 0;
5043 while i < bytes.len() {
5044 match bytes[i] {
5045 b'\'' => i = copy_quoted_to_string(&mut out, bytes, i, b'\''),
5046 b'"' => i = copy_quoted_to_string(&mut out, bytes, i, b'"'),
5047 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
5048 out.push(bytes[i] as char);
5049 i += 1;
5050 i = copy_quoted_to_string(&mut out, bytes, i, b'\'');
5051 }
5052 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
5053 out.push('-');
5054 out.push('-');
5055 i += 2;
5056 while i < bytes.len() {
5057 let ch = bytes[i] as char;
5058 out.push(ch);
5059 i += 1;
5060 if ch == '\n' {
5061 break;
5062 }
5063 }
5064 }
5065 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
5066 let start = i;
5067 i = skip_block_comment(bytes, i);
5068 out.push_str(&sql[start..i.min(bytes.len())]);
5069 }
5070 b'$' => {
5071 let start_len = out.len();
5072 let (next, matched) = copy_dollar_quoted_to_string(&mut out, bytes, i);
5073 if matched {
5074 i = next;
5075 } else {
5076 out.truncate(start_len);
5077 out.push('$');
5078 i += 1;
5079 }
5080 }
5081 b'g' | b'G' | b'm' | b'M' | b't' | b'T' => {
5082 if let Some((replacement, next)) = compat_function_rewrite_at(sql, i) {
5083 out.push_str(&replacement);
5084 i = next;
5085 } else {
5086 out.push(bytes[i] as char);
5087 i += 1;
5088 }
5089 }
5090 _ => {
5091 out.push(bytes[i] as char);
5092 i += 1;
5093 }
5094 }
5095 }
5096 out
5097}
5098
5099fn compat_function_rewrite_at(sql: &str, start: usize) -> Option<(String, usize)> {
5100 let bytes = sql.as_bytes();
5101 let (name, kind) = if ident_eq_at(bytes, start, b"max") {
5102 ("max", CompatRewriteKind::ScalarMax)
5103 } else if ident_eq_at(bytes, start, b"min") {
5104 ("min", CompatRewriteKind::ScalarMin)
5105 } else if ident_eq_at(bytes, start, b"group_concat") {
5106 ("group_concat", CompatRewriteKind::GroupConcat)
5107 } else if ident_eq_at(bytes, start, b"total") {
5108 ("total", CompatRewriteKind::Total)
5109 } else {
5110 return None;
5111 };
5112 let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
5113 let after_name = start + name.len();
5114 let after_ok = bytes
5115 .get(after_name)
5116 .is_some_and(|b| !is_sql_ident_byte(*b));
5117 if !before_ok || !after_ok {
5118 return None;
5119 }
5120 let mut open = after_name;
5121 while open < bytes.len() && bytes[open].is_ascii_whitespace() {
5122 open += 1;
5123 }
5124 if bytes.get(open) != Some(&b'(') {
5125 return None;
5126 }
5127 let summary = call_arg_summary(sql, open)?;
5128 match kind {
5129 CompatRewriteKind::ScalarMax if summary.top_level_commas > 0 => {
5130 Some(("__mongreldb_scalar_max(".to_string(), open + 1))
5131 }
5132 CompatRewriteKind::ScalarMin if summary.top_level_commas > 0 => {
5133 Some(("__mongreldb_scalar_min(".to_string(), open + 1))
5134 }
5135 CompatRewriteKind::GroupConcat => {
5136 let args = &sql[open + 1..summary.close];
5137 let rewritten = if summary.top_level_commas == 0 {
5138 format!("string_agg({args}, ',')")
5139 } else {
5140 format!("string_agg({args})")
5141 };
5142 Some((rewritten, summary.close + 1))
5143 }
5144 CompatRewriteKind::Total if summary.top_level_commas == 0 => {
5145 let args = &sql[open + 1..summary.close];
5146 let suffix_end = aggregate_suffix_end(sql, summary.close + 1);
5147 let suffix = &sql[summary.close + 1..suffix_end];
5148 Some((
5149 format!("coalesce(cast(sum({args}){suffix} as double), 0.0)"),
5150 suffix_end,
5151 ))
5152 }
5153 _ => None,
5154 }
5155}
5156
5157#[derive(Clone, Copy)]
5158enum CompatRewriteKind {
5159 ScalarMax,
5160 ScalarMin,
5161 GroupConcat,
5162 Total,
5163}
5164
5165fn ident_eq_at(bytes: &[u8], start: usize, ident: &[u8]) -> bool {
5166 bytes
5167 .get(start..start + ident.len())
5168 .is_some_and(|slice| slice.eq_ignore_ascii_case(ident))
5169}
5170
5171fn is_sql_ident_byte(b: u8) -> bool {
5172 b.is_ascii_alphanumeric() || b == b'_' || b == b'$'
5173}
5174
5175fn keyword_at(bytes: &[u8], start: usize, keyword: &[u8]) -> bool {
5176 if !ident_eq_at(bytes, start, keyword) {
5177 return false;
5178 }
5179 let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
5180 let after = start + keyword.len();
5181 let after_ok = after >= bytes.len() || !is_sql_ident_byte(bytes[after]);
5182 before_ok && after_ok
5183}
5184
5185fn skip_sql_whitespace(bytes: &[u8], mut i: usize) -> usize {
5186 while i < bytes.len() && bytes[i].is_ascii_whitespace() {
5187 i += 1;
5188 }
5189 i
5190}
5191
5192fn aggregate_suffix_end(sql: &str, start: usize) -> usize {
5193 let bytes = sql.as_bytes();
5194 let mut suffix_end = start;
5195 let mut i = skip_sql_whitespace(bytes, start);
5196
5197 if keyword_at(bytes, i, b"filter") {
5198 let open = skip_sql_whitespace(bytes, i + b"filter".len());
5199 if bytes.get(open) != Some(&b'(') {
5200 return start;
5201 }
5202 let Some(summary) = call_arg_summary(sql, open) else {
5203 return start;
5204 };
5205 suffix_end = summary.close + 1;
5206 i = skip_sql_whitespace(bytes, suffix_end);
5207 }
5208
5209 if keyword_at(bytes, i, b"over") {
5210 let after_over = skip_sql_whitespace(bytes, i + b"over".len());
5211 if bytes.get(after_over) == Some(&b'(') {
5212 let Some(summary) = call_arg_summary(sql, after_over) else {
5213 return suffix_end;
5214 };
5215 suffix_end = summary.close + 1;
5216 } else {
5217 let mut end = after_over;
5218 while end < bytes.len() && is_sql_ident_byte(bytes[end]) {
5219 end += 1;
5220 }
5221 if end > after_over {
5222 suffix_end = end;
5223 }
5224 }
5225 }
5226
5227 suffix_end
5228}
5229
5230struct CallArgSummary {
5231 close: usize,
5232 top_level_commas: usize,
5233}
5234
5235fn call_arg_summary(sql: &str, open: usize) -> Option<CallArgSummary> {
5236 let bytes = sql.as_bytes();
5237 let mut depth = 1;
5238 let mut i = open + 1;
5239 let mut top_level_commas = 0;
5240 while i < bytes.len() {
5241 match bytes[i] {
5242 b'\'' => i = skip_quoted(bytes, i, b'\''),
5243 b'"' => i = skip_quoted(bytes, i, b'"'),
5244 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
5245 i = skip_quoted(bytes, i + 1, b'\'')
5246 }
5247 b'$' => {
5248 let (next, matched) = skip_dollar_quoted(bytes, i);
5249 i = if matched { next } else { i + 1 };
5250 }
5251 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
5252 i += 2;
5253 while i < bytes.len() && bytes[i] != b'\n' {
5254 i += 1;
5255 }
5256 }
5257 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
5258 i = skip_block_comment(bytes, i);
5259 }
5260 b'(' => {
5261 depth += 1;
5262 i += 1;
5263 }
5264 b')' => {
5265 depth -= 1;
5266 if depth == 0 {
5267 return Some(CallArgSummary {
5268 close: i,
5269 top_level_commas,
5270 });
5271 }
5272 i += 1;
5273 }
5274 b',' if depth == 1 => {
5275 top_level_commas += 1;
5276 i += 1;
5277 }
5278 _ => i += 1,
5279 }
5280 }
5281 None
5282}
5283
5284fn copy_quoted_to_string(out: &mut String, bytes: &[u8], start: usize, delim: u8) -> usize {
5285 let end = skip_quoted(bytes, start, delim);
5286 out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
5287 end
5288}
5289
5290fn skip_quoted(bytes: &[u8], start: usize, delim: u8) -> usize {
5291 let mut i = start;
5292 if i < bytes.len() {
5293 i += 1;
5294 }
5295 while i < bytes.len() {
5296 if bytes[i] == delim {
5297 i += 1;
5298 if i < bytes.len() && bytes[i] == delim {
5299 i += 1;
5300 continue;
5301 }
5302 break;
5303 }
5304 i += 1;
5305 }
5306 i
5307}
5308
5309fn copy_dollar_quoted_to_string(out: &mut String, bytes: &[u8], start: usize) -> (usize, bool) {
5310 let (end, matched) = skip_dollar_quoted(bytes, start);
5311 if matched {
5312 out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
5313 }
5314 (end, matched)
5315}
5316
5317fn skip_dollar_quoted(bytes: &[u8], start: usize) -> (usize, bool) {
5318 if bytes.get(start) != Some(&b'$') {
5319 return (start, false);
5320 }
5321 let mut j = start + 1;
5322 while j < bytes.len() && (bytes[j].is_ascii_alphanumeric() || bytes[j] == b'_') {
5323 j += 1;
5324 }
5325 if bytes.get(j) != Some(&b'$') {
5326 return (start, false);
5327 }
5328 let tag = &bytes[start..=j];
5329 let mut i = j + 1;
5330 while i + tag.len() <= bytes.len() {
5331 if &bytes[i..i + tag.len()] == tag {
5332 return (i + tag.len(), true);
5333 }
5334 i += 1;
5335 }
5336 (start, false)
5337}
5338
5339fn skip_block_comment(bytes: &[u8], start: usize) -> usize {
5340 let mut i = start + 2;
5341 let mut depth = 1;
5342 while i + 1 < bytes.len() && depth > 0 {
5343 if bytes[i] == b'/' && bytes[i + 1] == b'*' {
5344 depth += 1;
5345 i += 2;
5346 } else if bytes[i] == b'*' && bytes[i + 1] == b'/' {
5347 depth -= 1;
5348 i += 2;
5349 } else {
5350 i += 1;
5351 }
5352 }
5353 i
5354}
5355
5356fn sql_cache_key(sql: &str) -> u64 {
5358 use std::hash::{Hash, Hasher};
5359 let mut h = std::collections::hash_map::DefaultHasher::new();
5360 sql.hash(&mut h);
5361 h.finish()
5362}
5363
5364fn replace_from_view(sql: &str, name: &str, view_sql: &str) -> String {
5371 let lower = sql.to_ascii_lowercase();
5372 let bytes = lower.as_bytes();
5373 let name_b = name.as_bytes();
5374 let mut i = 0usize;
5375 while let Some(rel) = lower[i..].find("from") {
5376 let from_start = i + rel;
5377 let after_from = from_start + 4;
5378 i = after_from;
5379 if from_start > 0 && is_ident_byte(bytes[from_start - 1]) {
5381 continue;
5382 }
5383 let mut j = after_from;
5385 while j < bytes.len() && bytes[j].is_ascii_whitespace() {
5386 j += 1;
5387 }
5388 if j == after_from || !bytes[j..].starts_with(name_b) {
5389 continue;
5390 }
5391 let after_name = j + name_b.len();
5392 if after_name < bytes.len() && is_ident_byte(bytes[after_name]) {
5394 continue;
5395 }
5396 let mut out = String::with_capacity(sql.len() + view_sql.len() + name.len() + 8);
5399 out.push_str(&sql[..from_start]);
5400 out.push_str(&sql[from_start..j]);
5401 out.push('(');
5402 out.push_str(view_sql);
5403 out.push_str(") AS ");
5404 out.push_str(name);
5405 out.push_str(&sql[after_name..]);
5406 return out;
5407 }
5408 sql.to_string()
5409}
5410
5411fn is_ident_byte(b: u8) -> bool {
5412 b.is_ascii_alphanumeric() || b == b'_'
5413}
5414
5415fn normalize_sql(sql: &str) -> String {
5426 let b = sql.as_bytes();
5427 let n = b.len();
5428 let mut out: Vec<u8> = Vec::with_capacity(n);
5429 let mut want_space = false;
5432 let mut i = 0usize;
5433 while i < n {
5434 let c = b[i];
5435 if c.is_ascii_whitespace() {
5439 want_space = true;
5440 i += 1;
5441 continue;
5442 }
5443 if c == b'-' && i + 1 < n && b[i + 1] == b'-' {
5444 i += 2;
5446 while i < n && b[i] != b'\n' {
5447 i += 1;
5448 }
5449 want_space = !out.is_empty();
5450 continue;
5451 }
5452 if c == b'/' && i + 1 < n && b[i + 1] == b'*' {
5453 let comment_start = i;
5456 i += 2;
5457 let mut depth = 1usize;
5458 while i + 1 < n && depth > 0 {
5459 if b[i] == b'/' && b[i + 1] == b'*' {
5460 depth += 1;
5461 i += 2;
5462 } else if b[i] == b'*' && b[i + 1] == b'/' {
5463 depth -= 1;
5464 i += 2;
5465 } else {
5466 i += 1;
5467 }
5468 }
5469 if depth != 0 {
5470 if want_space && !out.is_empty() {
5474 out.push(b' ');
5475 }
5476 out.extend_from_slice(&b[comment_start..]);
5477 break;
5478 }
5479 want_space = !out.is_empty();
5480 continue;
5481 }
5482 if want_space && !out.is_empty() {
5484 out.push(b' ');
5485 }
5486 want_space = false;
5487 match c {
5488 b'E' | b'e' if i + 1 < n && b[i + 1] == b'\'' => {
5490 out.push(c);
5491 i += 1;
5492 i = copy_quoted(&mut out, b, i, n, b'\'', true);
5493 continue;
5494 }
5495 b'\'' => {
5497 i = copy_quoted(&mut out, b, i, n, b'\'', false);
5498 continue;
5499 }
5500 b'"' => {
5502 i = copy_quoted(&mut out, b, i, n, b'"', false);
5503 continue;
5504 }
5505 b'`' => {
5507 i = copy_quoted(&mut out, b, i, n, b'`', false);
5508 continue;
5509 }
5510 b'[' => {
5512 i = copy_quoted(&mut out, b, i, n, b']', false);
5513 continue;
5514 }
5515 b'$' => {
5517 let (consumed, matched) = copy_dollar_quoted(&mut out, b, i, n);
5518 if matched {
5519 i = consumed;
5520 continue;
5521 }
5522 out.push(c);
5523 i += 1;
5524 continue;
5525 }
5526 _ => {
5527 out.push(c);
5528 i += 1;
5529 }
5530 }
5531 }
5532 String::from_utf8(out).unwrap_or_else(|_| sql.to_string())
5533}
5534
5535fn copy_quoted(
5540 out: &mut Vec<u8>,
5541 b: &[u8],
5542 start: usize,
5543 n: usize,
5544 delim: u8,
5545 backslash_escapes: bool,
5546) -> usize {
5547 out.push(b[start]);
5548 let mut i = start + 1;
5549 while i < n {
5550 let c = b[i];
5551 out.push(c);
5552 if backslash_escapes && c == b'\\' && i + 1 < n {
5553 out.push(b[i + 1]);
5554 i += 2;
5555 continue;
5556 }
5557 if c == delim {
5558 if i + 1 < n && b[i + 1] == delim {
5560 out.push(b[i + 1]);
5561 i += 2;
5562 continue;
5563 }
5564 return i + 1;
5565 }
5566 i += 1;
5567 }
5568 i
5569}
5570
5571fn copy_dollar_quoted(out: &mut Vec<u8>, b: &[u8], start: usize, n: usize) -> (usize, bool) {
5575 let mut j = start + 1;
5579 let tag_start = j;
5580 while j < n && b[j] != b'$' && is_dollar_tag_byte(b[j]) {
5581 j += 1;
5582 }
5583 if j >= n || b[j] != b'$' {
5584 return (start + 1, false);
5585 }
5586 if tag_start < j && !(b[tag_start].is_ascii_alphabetic() || b[tag_start] == b'_') {
5587 return (start + 1, false);
5588 }
5589 let close_end = j + 1; let delim = &b[start..close_end];
5591 out.extend_from_slice(delim);
5593 let mut k = close_end;
5595 while k + delim.len() <= n {
5596 if &b[k..k + delim.len()] == delim {
5597 out.extend_from_slice(delim);
5598 return (k + delim.len(), true);
5599 }
5600 out.push(b[k]);
5601 k += 1;
5602 }
5603 out.extend_from_slice(&b[k..n]);
5606 (n, true)
5607}
5608
5609fn is_dollar_tag_byte(b: u8) -> bool {
5610 b.is_ascii_alphanumeric() || b == b'_'
5611}
5612
5613fn strip_prefix_ci<'a>(s: &'a str, prefix: &str) -> Option<&'a str> {
5615 let bytes = s.as_bytes();
5616 let pb = prefix.as_bytes();
5617 if bytes.len() >= pb.len() && bytes[..pb.len()].eq_ignore_ascii_case(pb) {
5618 Some(&s[pb.len()..])
5619 } else {
5620 None
5621 }
5622}
5623
5624const COLUMN_CONSTRAINTS: &[(&str, u32)] = &[
5637 ("primary key", ColumnFlags::PRIMARY_KEY),
5638 ("autoincrement", ColumnFlags::AUTO_INCREMENT),
5645 ("auto_increment", ColumnFlags::AUTO_INCREMENT),
5646];
5647
5648fn parse_column_constraints(constraint_text: &str) -> ColumnFlags {
5654 let normalized = constraint_text.to_lowercase();
5655 let mut flags = ColumnFlags::empty();
5656 for (phrase, bit) in COLUMN_CONSTRAINTS {
5657 if normalized.contains(phrase) {
5658 flags = flags.with(*bit);
5659 }
5660 }
5661 flags
5662}
5663
5664fn parse_sql_type(ty_str: &str) -> Result<mongreldb_core::schema::TypeId> {
5665 use mongreldb_core::schema::TypeId;
5666
5667 match ty_str.trim().trim_end_matches(';').to_lowercase().as_str() {
5668 "bigint" | "int8" | "int64" | "integer" | "int" => Ok(TypeId::Int64),
5669 "double" | "float8" | "float64" | "real" | "float" => Ok(TypeId::Float64),
5670 "varchar" | "text" | "string" | "bytes" => Ok(TypeId::Bytes),
5671 "boolean" | "bool" => Ok(TypeId::Bool),
5672 other => Err(MongrelQueryError::Schema(format!(
5673 "unsupported column type: {other}"
5674 ))),
5675 }
5676}
5677
5678fn parse_create_table(sql: &str) -> Result<(String, mongreldb_core::schema::Schema)> {
5684 use mongreldb_core::schema::*;
5685
5686 let open = sql
5687 .find('(')
5688 .ok_or(MongrelQueryError::Schema("CREATE TABLE missing '('".into()))?;
5689 let close = sql
5690 .rfind(')')
5691 .ok_or(MongrelQueryError::Schema("CREATE TABLE missing ')'".into()))?;
5692 let head = sql[..open].trim();
5693 let after_kw = strip_prefix_ci(head, "CREATE TABLE")
5694 .or_else(|| strip_prefix_ci(head, "create table"))
5695 .unwrap_or("")
5696 .trim();
5697 let after_kw = after_kw
5699 .strip_prefix("IF NOT EXISTS")
5700 .or_else(|| after_kw.strip_prefix("if not exists"))
5701 .map(str::trim)
5702 .unwrap_or(after_kw);
5703 let name = after_kw.trim_matches('"').to_string();
5704 if name.is_empty() {
5705 return Err(MongrelQueryError::Schema(
5706 "CREATE TABLE missing table name".into(),
5707 ));
5708 }
5709
5710 let body = &sql[open + 1..close];
5711 let mut columns = Vec::new();
5712 let schema_id: u64 = 0; for (i, raw) in body.split(',').enumerate() {
5714 let part = raw.trim();
5715 if part.is_empty() {
5716 continue;
5717 }
5718 let mut tokens = part.split_whitespace();
5719 let col_name = tokens
5720 .next()
5721 .ok_or(MongrelQueryError::Schema("missing column name".into()))?
5722 .trim_matches('"');
5723 let ty_str = tokens
5724 .next()
5725 .ok_or(MongrelQueryError::Schema("missing column type".into()))?
5726 .to_lowercase();
5727 let ty = parse_sql_type(&ty_str)?;
5728 let constraint_clause: String = tokens.collect::<Vec<_>>().join(" ");
5732 let flags = parse_column_constraints(&constraint_clause);
5733 columns.push(ColumnDef {
5734 id: (i + 1) as u16,
5735 name: col_name.to_string(),
5736 ty,
5737 flags,
5738 default_value: None,
5739 });
5740 }
5741
5742 Ok((
5743 name,
5744 Schema {
5745 schema_id,
5746 columns,
5747 indexes: vec![],
5748 colocation: vec![],
5749 constraints: Default::default(),
5750 clustered: false,
5751 },
5752 ))
5753}
5754
5755fn parse_drop_table(sql: &str) -> Result<(String, bool)> {
5757 let head = sql.trim();
5758 let after_kw = strip_prefix_ci(head, "DROP TABLE")
5759 .or_else(|| strip_prefix_ci(head, "drop table"))
5760 .unwrap_or("")
5761 .trim();
5762 let (rest, if_exists) = if let Some(r) = after_kw
5764 .strip_prefix("IF EXISTS")
5765 .or_else(|| after_kw.strip_prefix("if exists"))
5766 .map(str::trim)
5767 {
5768 (r, true)
5769 } else {
5770 (after_kw, false)
5771 };
5772 let name = rest.trim_matches(';').trim_matches('"').trim();
5773 if name.is_empty() {
5774 return Err(MongrelQueryError::Schema(
5775 "DROP TABLE missing table name".into(),
5776 ));
5777 }
5778 Ok((name.to_string(), if_exists))
5779}
5780
5781enum ParsedAlterTable {
5782 RenameTable {
5783 old_name: String,
5784 new_name: String,
5785 },
5786 RenameColumn {
5787 table_name: String,
5788 column_name: String,
5789 new_name: String,
5790 },
5791 AlterColumnType {
5792 table_name: String,
5793 column_name: String,
5794 ty: mongreldb_core::schema::TypeId,
5795 },
5796 SetNotNull {
5797 table_name: String,
5798 column_name: String,
5799 },
5800 DropNotNull {
5801 table_name: String,
5802 column_name: String,
5803 },
5804}
5805
5806fn current_column_flags(db: &Arc<Database>, table: &str, column: &str) -> Result<ColumnFlags> {
5807 let handle = db.table(table)?;
5808 let table = handle.lock();
5809 table
5810 .schema()
5811 .column(column)
5812 .map(|c| c.flags)
5813 .ok_or_else(|| MongrelQueryError::Schema(format!("unknown column {column}")))
5814}
5815
5816fn parse_alter_table(sql: &str) -> Result<ParsedAlterTable> {
5817 let trimmed = strip_statement_semicolon(sql.trim());
5818 let after_kw = strip_prefix_ci(trimmed, "ALTER TABLE")
5819 .ok_or_else(|| MongrelQueryError::Schema("not an ALTER TABLE statement".into()))?
5820 .trim();
5821 let (table_name, rest) = take_sql_ident(after_kw, "ALTER TABLE missing table name")?;
5822 let rest = rest.trim();
5823
5824 if let Some(after) = strip_prefix_ci(rest, "RENAME TO") {
5825 let new_name = parse_trailing_identifier(after, "ALTER TABLE missing new table name")?;
5826 return Ok(ParsedAlterTable::RenameTable {
5827 old_name: table_name,
5828 new_name,
5829 });
5830 }
5831
5832 if let Some(after) = strip_prefix_ci(rest, "RENAME COLUMN") {
5833 let (column_name, after_col) =
5834 take_sql_ident(after, "ALTER TABLE RENAME COLUMN missing column name")?;
5835 let after_to = strip_prefix_ci(after_col.trim(), "TO").ok_or_else(|| {
5836 MongrelQueryError::Schema("ALTER TABLE RENAME COLUMN missing TO".into())
5837 })?;
5838 let new_name = parse_trailing_identifier(
5839 after_to,
5840 "ALTER TABLE RENAME COLUMN missing new column name",
5841 )?;
5842 return Ok(ParsedAlterTable::RenameColumn {
5843 table_name,
5844 column_name,
5845 new_name,
5846 });
5847 }
5848
5849 let after_alter = strip_prefix_ci(rest, "ALTER COLUMN")
5850 .or_else(|| strip_prefix_ci(rest, "ALTER"))
5851 .ok_or_else(|| {
5852 MongrelQueryError::Schema(
5853 "ALTER TABLE must be RENAME TO, RENAME COLUMN, or ALTER COLUMN".into(),
5854 )
5855 })?;
5856 let (column_name, action) =
5857 take_sql_ident(after_alter, "ALTER TABLE ALTER COLUMN missing column name")?;
5858 let action = action.trim();
5859
5860 if let Some(after_type) =
5861 strip_prefix_ci(action, "TYPE").or_else(|| strip_prefix_ci(action, "SET DATA TYPE"))
5862 {
5863 let ty = parse_type_tail(after_type)?;
5864 return Ok(ParsedAlterTable::AlterColumnType {
5865 table_name,
5866 column_name,
5867 ty,
5868 });
5869 }
5870 if strip_prefix_ci(action, "SET NOT NULL").is_some() {
5871 return Ok(ParsedAlterTable::SetNotNull {
5872 table_name,
5873 column_name,
5874 });
5875 }
5876 if strip_prefix_ci(action, "DROP NOT NULL").is_some() {
5877 return Ok(ParsedAlterTable::DropNotNull {
5878 table_name,
5879 column_name,
5880 });
5881 }
5882
5883 Err(MongrelQueryError::Schema(
5884 "unsupported ALTER COLUMN action".into(),
5885 ))
5886}
5887
5888fn strip_statement_semicolon(s: &str) -> &str {
5889 s.trim().trim_end_matches(';').trim()
5890}
5891
5892fn take_sql_ident<'a>(s: &'a str, missing: &str) -> Result<(String, &'a str)> {
5893 let s = s.trim();
5894 if s.is_empty() {
5895 return Err(MongrelQueryError::Schema(missing.into()));
5896 }
5897 if let Some(rest) = s.strip_prefix('"') {
5898 let Some(end) = rest.find('"') else {
5899 return Err(MongrelQueryError::Schema(
5900 "unterminated quoted identifier".into(),
5901 ));
5902 };
5903 let ident = rest[..end].to_string();
5904 if ident.is_empty() {
5905 return Err(MongrelQueryError::Schema(missing.into()));
5906 }
5907 return Ok((ident, &rest[end + 1..]));
5908 }
5909 let end = s.find(|c: char| c.is_ascii_whitespace()).unwrap_or(s.len());
5910 let ident = s[..end].trim_matches('"').to_string();
5911 if ident.is_empty() {
5912 return Err(MongrelQueryError::Schema(missing.into()));
5913 }
5914 Ok((ident, &s[end..]))
5915}
5916
5917fn parse_trailing_identifier(s: &str, missing: &str) -> Result<String> {
5918 let (ident, rest) = take_sql_ident(s, missing)?;
5919 if !strip_statement_semicolon(rest).is_empty() {
5920 return Err(MongrelQueryError::Schema(
5921 "unexpected tokens after identifier".into(),
5922 ));
5923 }
5924 Ok(ident)
5925}
5926
5927fn parse_type_tail(s: &str) -> Result<mongreldb_core::schema::TypeId> {
5928 let tail = strip_statement_semicolon(s);
5929 let ty = tail
5930 .split_whitespace()
5931 .next()
5932 .ok_or_else(|| MongrelQueryError::Schema("ALTER COLUMN TYPE missing type".into()))?;
5933 parse_sql_type(ty)
5934}
5935
5936#[cfg(test)]
5937mod tests {
5938 use super::*;
5939
5940 #[test]
5941 fn bounded_lru_evicts_least_recently_used() {
5942 let mut cache = BoundedLru::new(2);
5943 cache.insert("a", 1);
5944 cache.insert("b", 2);
5945 assert_eq!(cache.get("a"), Some(&1));
5946 cache.insert("c", 3);
5947 assert_eq!(cache.entries.len(), 2);
5948 assert_eq!(cache.get("b"), None);
5949 assert_eq!(cache.get("a"), Some(&1));
5950 assert_eq!(cache.get("c"), Some(&3));
5951 }
5952
5953 #[test]
5954 fn result_cache_skips_entry_over_byte_budget() {
5955 let batch = RecordBatch::try_from_iter(vec![(
5956 "value",
5957 Arc::new(Int64Array::from(vec![1_i64])) as ArrayRef,
5958 )])
5959 .unwrap();
5960 let mut cache = ResultCacheStore::new(2, 1);
5961 cache.insert(("SELECT 1".into(), 0), Arc::new(vec![batch]));
5962 assert!(cache.is_empty());
5963 assert_eq!(cache.bytes, 0);
5964 }
5965
5966 #[tokio::test]
5967 async fn collection_limit_records_typed_terminal_failure() {
5968 let dir = tempfile::tempdir().unwrap();
5969 let database = Arc::new(Database::create(dir.path()).unwrap());
5970 let session = MongrelSession::open(database).unwrap();
5971 let query = session.register_query(SqlQueryOptions::default()).unwrap();
5972 let query_id = query.id();
5973 let error = match session
5974 .run_with_query_for_serialization_with_limits(
5975 "SELECT 1 UNION ALL SELECT 2",
5976 query,
5977 SqlCollectionLimits::new(1, 1024),
5978 )
5979 .await
5980 {
5981 Ok(_) => panic!("expected result limit failure"),
5982 Err(error) => error,
5983 };
5984 assert!(matches!(
5985 error,
5986 MongrelQueryError::ResultLimitExceeded {
5987 committed: false,
5988 ..
5989 }
5990 ));
5991 let status = session.query_registry().status(query_id).unwrap();
5992 assert_eq!(
5993 status.terminal_error.unwrap().category,
5994 QueryTerminalErrorCategory::ResultLimit
5995 );
5996 }
5997
5998 #[tokio::test]
5999 async fn collection_limit_error_keeps_prior_statement_commit() {
6000 let dir = tempfile::tempdir().unwrap();
6001 let database = Arc::new(Database::create(dir.path()).unwrap());
6002 let session = MongrelSession::open(database).unwrap();
6003 let query = session.register_query(SqlQueryOptions::default()).unwrap();
6004 let error = match session
6005 .run_with_query_for_serialization_with_limits(
6006 "CREATE TABLE committed_first (id BIGINT PRIMARY KEY); \
6007 SELECT 1 UNION ALL SELECT 2",
6008 query,
6009 SqlCollectionLimits::new(1, 1024),
6010 )
6011 .await
6012 {
6013 Ok(_) => panic!("expected result limit failure"),
6014 Err(error) => error,
6015 };
6016 assert!(matches!(
6017 error,
6018 MongrelQueryError::ResultLimitExceeded {
6019 committed: true,
6020 committed_statements: 1,
6021 last_commit_epoch: Some(_),
6022 first_commit_statement_index: Some(0),
6023 last_commit_statement_index: Some(0),
6024 statement_index: 1,
6025 ..
6026 }
6027 ));
6028 }
6029
6030 #[tokio::test]
6031 async fn streaming_query_bypasses_result_cache() {
6032 use futures::StreamExt;
6033
6034 let dir = tempfile::tempdir().unwrap();
6035 let database = Arc::new(Database::create(dir.path()).unwrap());
6036 let session = MongrelSession::open(database).unwrap();
6037 let mut stream = session.run_stream("SELECT 1 AS value").await.unwrap();
6038 let batch = stream.next().await.unwrap().unwrap();
6039 assert_eq!(batch.num_rows(), 1);
6040 assert!(stream.next().await.is_none());
6041 assert!(session.cache.lock().is_empty());
6042 }
6043
6044 #[tokio::test]
6045 async fn ttl_tables_bypass_epoch_keyed_result_cache() {
6046 let dir = tempfile::tempdir().unwrap();
6047 let database = Arc::new(Database::create(dir.path()).unwrap());
6048 let session = MongrelSession::open(database).unwrap();
6049 session
6050 .run(
6051 "CREATE TABLE events (id BIGINT PRIMARY KEY, ts TIMESTAMP) \
6052 TTL_COLUMN ts TTL '1 day'",
6053 )
6054 .await
6055 .unwrap();
6056 session.run("SELECT * FROM events").await.unwrap();
6057 assert!(session.cache.lock().is_empty());
6058 }
6059
6060 #[test]
6061 fn normalize_collapses_and_trims_whitespace() {
6062 assert_eq!(normalize_sql("SELECT * FROM t"), "SELECT * FROM t");
6063 assert_eq!(normalize_sql(" SELECT * FROM t "), "SELECT * FROM t");
6064 assert_eq!(
6065 normalize_sql("\n\tSELECT\n*\nFROM\n\tt\n"),
6066 "SELECT * FROM t"
6067 );
6068 assert_eq!(
6069 normalize_sql("SELECT a, b FROM t"),
6070 normalize_sql("SELECT a, b FROM t")
6071 );
6072 }
6073
6074 #[test]
6075 fn boolean_ai_predicate_detection_uses_sql_tokens() {
6076 assert!(contains_boolean_ai_predicate(
6077 "SELECT * FROM docs WHERE ANN_SEARCH (embedding, '[1]', 1)"
6078 ));
6079 assert!(contains_boolean_ai_predicate(
6080 "PREPARE q AS SELECT * FROM docs WHERE sparse_match(sparse, '[]', 1)"
6081 ));
6082 assert!(!contains_boolean_ai_predicate(
6083 "SELECT 'ann_search(x)', sparse_search_scored('docs', 'sparse', '[]', 1) /* ann_search(x) */"
6084 ));
6085 }
6086
6087 #[test]
6088 fn normalize_preserves_string_literal_whitespace() {
6089 assert_eq!(
6090 normalize_sql("SELECT 'hello world' FROM t"),
6091 "SELECT 'hello world' FROM t"
6092 );
6093 assert_eq!(
6094 normalize_sql("SELECT 'it''s ok' FROM t"),
6095 "SELECT 'it''s ok' FROM t"
6096 );
6097 assert_eq!(
6098 normalize_sql(" SELECT 'a b' FROM t "),
6099 "SELECT 'a b' FROM t"
6100 );
6101 }
6102
6103 #[test]
6104 fn normalize_preserves_quoted_identifier_and_dollar_quote() {
6105 assert_eq!(
6106 normalize_sql(" SELECT \"my col\" FROM t "),
6107 "SELECT \"my col\" FROM t"
6108 );
6109 assert_eq!(
6110 normalize_sql(" SELECT $$a b$$ FROM t "),
6111 "SELECT $$a b$$ FROM t"
6112 );
6113 assert_eq!(
6114 normalize_sql("SELECT $tag$body with spaces$tag$ FROM t"),
6115 "SELECT $tag$body with spaces$tag$ FROM t"
6116 );
6117 }
6118
6119 #[test]
6120 fn normalize_strips_comments() {
6121 assert_eq!(
6122 normalize_sql("SELECT 1 -- trailing comment\nFROM t"),
6123 "SELECT 1 FROM t"
6124 );
6125 assert_eq!(
6126 normalize_sql("SELECT /* block */ 1 FROM t"),
6127 "SELECT 1 FROM t"
6128 );
6129 assert_eq!(
6131 normalize_sql("SELECT /* 'not a string' */ 1 FROM t"),
6132 "SELECT 1 FROM t"
6133 );
6134 assert_eq!(
6136 normalize_sql("SELECT /* outer /* inner */ still outer */ 1 FROM t"),
6137 "SELECT 1 FROM t"
6138 );
6139 assert_eq!(
6140 normalize_sql("SELECT 1 /* unterminated"),
6141 "SELECT 1 /* unterminated"
6142 );
6143 assert_eq!(normalize_sql("SELECT 1/*"), "SELECT 1/*");
6144 }
6145
6146 #[test]
6147 fn protocol_sql_fingerprint_is_literal_aware() {
6148 assert_eq!(
6149 normalized_sql_fingerprint(" INSERT INTO t VALUES (1) -- retry\n"),
6150 normalized_sql_fingerprint("INSERT INTO t VALUES (1)")
6151 );
6152 assert_ne!(
6153 normalized_sql_fingerprint("INSERT INTO t VALUES ('a b')"),
6154 normalized_sql_fingerprint("INSERT INTO t VALUES ('a b')")
6155 );
6156 }
6157
6158 #[test]
6159 fn idempotency_class_only_allows_durable_single_writes() {
6160 assert_eq!(
6161 classify_sql_idempotency("SELECT 1"),
6162 SqlIdempotencyClass::ReadOnly
6163 );
6164 for sql in [
6165 "INSERT INTO t VALUES (1)",
6166 "UPDATE t SET value = 2",
6167 "DELETE FROM t",
6168 "TRUNCATE TABLE t",
6169 "CREATE TABLE t (id BIGINT)",
6170 "ALTER TABLE t ADD COLUMN value BIGINT",
6171 "CREATE INDEX t_id ON t (id)",
6172 "CREATE MATERIALIZED VIEW visible_t AS SELECT * FROM t",
6173 "DROP TABLE t",
6174 "DROP MATERIALIZED VIEW visible_t",
6175 ] {
6176 assert_eq!(
6177 classify_sql_idempotency(sql),
6178 SqlIdempotencyClass::SingleWrite,
6179 "{sql}"
6180 );
6181 }
6182 assert_eq!(
6183 classify_sql_idempotency("INSERT INTO t VALUES (1); INSERT INTO t VALUES (2)"),
6184 SqlIdempotencyClass::Unsupported
6185 );
6186 for sql in [
6187 "BEGIN",
6188 "NOTIFY jobs, 'ready'",
6189 "LISTEN jobs",
6190 "ATTACH DATABASE 'other.db' AS other",
6191 "DETACH DATABASE other",
6192 "SHOW TABLES",
6193 "EXPLAIN SELECT 1",
6194 "PRAGMA table_info(t)",
6195 "CREATE VIEW visible_t AS SELECT * FROM t",
6196 "DROP VIEW visible_t",
6197 "DROP TABLE one, two",
6198 "DROP INDEX one, two",
6199 ] {
6200 assert_eq!(
6201 classify_sql_idempotency(sql),
6202 SqlIdempotencyClass::Unsupported,
6203 "{sql}"
6204 );
6205 }
6206 }
6207
6208 #[test]
6209 fn normalize_escape_string_preserved() {
6210 assert_eq!(
6211 normalize_sql("SELECT E'line\\nbreak' FROM t"),
6212 "SELECT E'line\\nbreak' FROM t"
6213 );
6214 assert_eq!(
6215 normalize_sql("SELECT E'it\\\'s /* data */' FROM t"),
6216 "SELECT E'it\\\'s /* data */' FROM t"
6217 );
6218 }
6219
6220 #[test]
6221 fn normalize_quoted_identifiers_preserves_comment_markers() {
6222 for sql in [
6223 "SELECT `a/* data */b` FROM t",
6224 "SELECT [a-- data b] FROM t",
6225 "SELECT `a``b` FROM t",
6226 "SELECT [a]]b] FROM t",
6227 ] {
6228 assert_eq!(normalize_sql(sql), sql);
6229 }
6230 assert_eq!(
6231 normalize_sql("SELECT $tag$unterminated"),
6232 "SELECT $tag$unterminated"
6233 );
6234 }
6235
6236 #[test]
6237 fn replace_from_view_matches_whole_word_only() {
6238 let out = replace_from_view("SELECT * FROM logs", "log", "SELECT 1");
6239 assert_eq!(out, "SELECT * FROM logs");
6240
6241 let out = replace_from_view("SELECT * FROM log", "log", "SELECT 1");
6242 assert_eq!(out, "SELECT * FROM (SELECT 1) AS log");
6243
6244 let out = replace_from_view("select * from log where x", "log", "SELECT 1");
6245 assert_eq!(out, "select * from (SELECT 1) AS log where x");
6246
6247 let out = replace_from_view("SELECT * FROM log)", "log", "SELECT 1");
6248 assert_eq!(out, "SELECT * FROM (SELECT 1) AS log)");
6249
6250 let out = replace_from_view("SELECT * xfrom log", "log", "SELECT 1");
6251 assert_eq!(out, "SELECT * xfrom log");
6252 }
6253
6254 #[test]
6255 fn compat_function_rewrite_handles_sqlite_compatibility_calls() {
6256 assert_eq!(
6257 rewrite_compat_function_calls("select max(id), min(id) from t"),
6258 "select max(id), min(id) from t"
6259 );
6260 assert_eq!(
6261 rewrite_compat_function_calls("select max(1, min(2, 3), 'max(4,5)')"),
6262 "select __mongreldb_scalar_max(1, __mongreldb_scalar_min(2, 3), 'max(4,5)')"
6263 );
6264 assert_eq!(
6265 rewrite_compat_function_calls("select /* max(1,2) */ min(1, (2 + 3))"),
6266 "select /* max(1,2) */ __mongreldb_scalar_min(1, (2 + 3))"
6267 );
6268 assert_eq!(
6269 rewrite_compat_function_calls("select max_value, min_value from t"),
6270 "select max_value, min_value from t"
6271 );
6272 assert_eq!(
6273 rewrite_compat_function_calls(
6274 "select group_concat(label), group_concat(label, '|') from t"
6275 ),
6276 "select string_agg(label, ','), string_agg(label, '|') from t"
6277 );
6278 assert_eq!(
6279 rewrite_compat_function_calls("select total(val), total(val) filter (where grp = 2) from t"),
6280 "select coalesce(cast(sum(val) as double), 0.0), coalesce(cast(sum(val) filter (where grp = 2) as double), 0.0) from t"
6281 );
6282 assert_eq!(
6283 rewrite_compat_function_calls(
6284 "select total(val) over (partition by grp order by id) from t"
6285 ),
6286 "select coalesce(cast(sum(val) over (partition by grp order by id) as double), 0.0) from t"
6287 );
6288 }
6289}