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
3365 .ctx
3366 .sql(&sql)
3367 .await
3368 .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))?;
3369 mongreldb_core::trace::QueryTrace::record(|trace| {
3370 trace.planning_nanos = plan_start.elapsed().as_nanos() as u64;
3371 });
3372 self.collect_dataframe(dataframe, sql_query).await
3373 }
3374
3375 async fn run_as_of_stream(
3376 &self,
3377 query: AsOfQuery,
3378 sql_query: &RegisteredSqlQuery,
3379 ) -> Result<MongrelRecordBatchStream> {
3380 use futures::StreamExt;
3381
3382 let AsOfQuery { sql, registration } = query;
3383 let dataframe = self
3384 .ctx
3385 .sql(&sql)
3386 .await
3387 .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))?;
3388 let stream = self.execute_dataframe_stream(dataframe, sql_query).await?;
3389 let schema = stream.schema();
3390 let guarded = futures::stream::unfold(
3391 (stream, registration),
3392 |(mut stream, registration)| async move {
3393 stream
3394 .next()
3395 .await
3396 .map(|item| (item, (stream, registration)))
3397 },
3398 );
3399 Ok(Box::pin(RecordBatchStreamAdapter::new(schema, guarded)))
3400 }
3401
3402 pub async fn run_stream_with_options(
3403 &self,
3404 sql: &str,
3405 options: SqlQueryOptions,
3406 ) -> Result<MongrelRecordBatchStream> {
3407 let query = self.register_query(options)?;
3408 self.run_stream_with_query(sql, query).await
3409 }
3410
3411 pub async fn run_stream_with_query(
3412 &self,
3413 sql: &str,
3414 query: RegisteredSqlQuery,
3415 ) -> Result<MongrelRecordBatchStream> {
3416 self.run_stream_with_query_mode(sql, query, false)
3417 .await
3418 .map(|(stream, _)| stream)
3419 }
3420
3421 pub async fn run_stream_with_query_for_serialization(
3422 &self,
3423 sql: &str,
3424 query: RegisteredSqlQuery,
3425 ) -> Result<(MongrelRecordBatchStream, SqlStreamCompletion)> {
3426 let (stream, completion) = self.run_stream_with_query_mode(sql, query, true).await?;
3427 let completion = completion.ok_or_else(|| {
3428 MongrelQueryError::InvalidQueryState(
3429 "deferred stream completion was not created".into(),
3430 )
3431 })?;
3432 Ok((stream, completion))
3433 }
3434
3435 async fn run_stream_with_query_mode(
3436 &self,
3437 sql: &str,
3438 query: RegisteredSqlQuery,
3439 defer_completion: bool,
3440 ) -> Result<(MongrelRecordBatchStream, Option<SqlStreamCompletion>)> {
3441 query.set_sql_metadata(sql);
3442 let guard = RegisteredQueryGuard::new(query.clone());
3443 self.fire_test_hook(SqlTestHookPoint::WaitingForSessionPermit);
3444 let permit = tokio::select! {
3445 permit = Arc::clone(&self.execution_gate).acquire_owned() => permit.map_err(|_| {
3446 MongrelQueryError::InvalidQueryState("session execution gate closed".into())
3447 })?,
3448 _ = query.control().cancelled() => {
3449 let error = query.checkpoint().unwrap_err();
3450 guard.fail();
3451 return Err(error);
3452 }
3453 };
3454 query.transition(SqlQueryPhase::Queued, SqlQueryPhase::Planning)?;
3455 self.fire_test_hook(SqlTestHookPoint::Planning);
3456 query.transition(SqlQueryPhase::Planning, SqlQueryPhase::Executing)?;
3457 query.begin_statement(0);
3458 let stream = CURRENT_SQL_QUERY
3459 .scope(query.clone(), async {
3460 tokio::select! {
3461 biased;
3462 result = self.run_stream_internal(sql, &query) => result,
3463 _ = query.control().cancelled() => Err(query.checkpoint().unwrap_err()),
3464 }
3465 })
3466 .await;
3467 let stream = match stream {
3468 Ok(stream) => stream,
3469 Err(error) => {
3470 guard.fail();
3471 return Err(error);
3472 }
3473 };
3474 if query.phase() == SqlQueryPhase::Executing {
3475 query.transition(SqlQueryPhase::Executing, SqlQueryPhase::Streaming)?;
3476 } else if query.phase() != SqlQueryPhase::CommitCritical {
3477 return Err(MongrelQueryError::InvalidQueryState(format!(
3478 "query {} cannot stream from {:?}",
3479 query.id(),
3480 query.phase()
3481 )));
3482 }
3483 let deferred_completion = defer_completion
3484 .then(|| Arc::new(std::sync::atomic::AtomicU8::new(STREAM_COMPLETION_PENDING)));
3485 if defer_completion {
3486 query.begin_serialization()?;
3487 }
3488 let schema = stream.schema();
3489 let completion = deferred_completion
3490 .as_ref()
3491 .map(|state| SqlStreamCompletion {
3492 state: Arc::clone(state),
3493 query: query.clone(),
3494 });
3495 let stream = Box::pin(ManagedQueryStream {
3496 inner: stream,
3497 schema,
3498 query,
3499 guard: Some(guard),
3500 permit: Some(permit),
3501 deferred_completion,
3502 });
3503 Ok((stream, completion))
3504 }
3505
3506 pub async fn run_stream(&self, sql: &str) -> Result<MongrelRecordBatchStream> {
3507 if let Ok(query) = CURRENT_SQL_QUERY.try_with(Clone::clone) {
3508 return Box::pin(self.run_stream_internal(sql, &query)).await;
3509 }
3510 Box::pin(self.run_stream_with_options(sql, SqlQueryOptions::default())).await
3511 }
3512
3513 async fn run_stream_internal(
3518 &self,
3519 sql: &str,
3520 query: &RegisteredSqlQuery,
3521 ) -> Result<MongrelRecordBatchStream> {
3522 let trimmed = sql.trim();
3523 let multiple_statements = trimmed.contains(';')
3524 && !is_trigger_body(trimmed)
3525 && split_sql_statements(trimmed)
3526 .iter()
3527 .filter(|statement| !statement.trim().is_empty())
3528 .count()
3529 > 1;
3530 let rollback = commands::rollback_kind(sql);
3531 let (transaction_open, transaction_aborted) = {
3532 let transaction = self.transaction.lock();
3533 (transaction.staged_ops.is_some(), transaction.aborted)
3534 };
3535 if !multiple_statements && transaction_open && transaction_aborted && rollback.is_none() {
3536 return Err(MongrelQueryError::TransactionAborted);
3537 }
3538 let statement = TransactionStatementGuard::new(
3539 self,
3540 query,
3541 !multiple_statements
3542 && transaction_open
3543 && rollback != Some(commands::RollbackKind::Full),
3544 )?;
3545 self.fire_test_hook(SqlTestHookPoint::BeforeStatementExecution);
3546 let (inner, buffered) = self.run_stream_internal_body(sql, query).await?;
3547 if buffered {
3548 statement.complete();
3549 return Ok(inner);
3550 }
3551 let schema = inner.schema();
3552 Ok(Box::pin(TransactionGuardedStream {
3553 inner,
3554 schema,
3555 statement: Some(statement),
3556 }))
3557 }
3558
3559 async fn run_stream_internal_body(
3560 &self,
3561 sql: &str,
3562 query: &RegisteredSqlQuery,
3563 ) -> Result<(MongrelRecordBatchStream, bool)> {
3564 query.checkpoint()?;
3565 if strip_explain_query_plan(sql).is_some() {
3566 return self
3567 .run_internal(sql, query)
3568 .await
3569 .map(|batches| (buffered_stream(batches), true));
3570 }
3571
3572 let trimmed = sql.trim();
3573 if trimmed.contains(';') && !is_trigger_body(trimmed) {
3574 let stmts = split_sql_statements(trimmed);
3575 let non_empty: Vec<&str> = stmts
3576 .iter()
3577 .map(String::as_str)
3578 .filter(|stmt| !stmt.trim().is_empty() && stmt.trim() != ";")
3579 .collect();
3580 if non_empty.is_empty() {
3581 return Ok((buffered_stream(Vec::new()), true));
3582 }
3583 if stmts.len() > 1 {
3584 for (index, stmt) in non_empty[..non_empty.len() - 1].iter().enumerate() {
3585 query.begin_statement(index);
3586 Box::pin(self.run_internal(stmt, query)).await?;
3587 query.complete_statement(index);
3588 }
3589 query.begin_statement(non_empty.len() - 1);
3590 let stream =
3591 Box::pin(self.run_stream_internal(non_empty[non_empty.len() - 1], query))
3592 .await?;
3593 return Ok((stream, true));
3594 }
3595 }
3596
3597 query.checkpoint()?;
3598 if let Some(batches) = commands::try_run_command(self, sql, query).await? {
3599 return Ok((buffered_stream(batches), true));
3600 }
3601
3602 if let Some(as_of_query) = self.prepare_as_of_query(sql)? {
3603 return self
3604 .run_as_of_stream(as_of_query, query)
3605 .await
3606 .map(|stream| (stream, false));
3607 }
3608
3609 let resolved = self.resolve_view_sql(sql);
3610 let resolved = self.rewrite_external_module_compat_sql(&resolved);
3611 let resolved = rewrite_compat_function_calls(&resolved);
3612 let effective_sql = normalize_sql(&resolved);
3613 let sql = effective_sql.as_str();
3614
3615 if let Some(batches) = self.try_catalog_introspection(sql)? {
3616 return Ok((buffered_stream(batches), true));
3617 }
3618
3619 let plan_start = std::time::Instant::now();
3620 let external_module_scan = self.query_references_external_module(sql);
3621 let df = self.dataframe(sql).await?;
3622 self.track_prepared_name(sql);
3623 mongreldb_core::trace::QueryTrace::record(|trace| {
3624 trace.planning_nanos = plan_start.elapsed().as_nanos() as u64;
3625 });
3626 let stream = self.execute_dataframe_stream(df, query).await?;
3627 if external_module_scan {
3628 mongreldb_core::trace::QueryTrace::record(|trace| {
3629 trace.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
3630 });
3631 }
3632 Ok((stream, false))
3633 }
3634
3635 pub async fn run_with_options(
3636 &self,
3637 sql: &str,
3638 options: SqlQueryOptions,
3639 ) -> Result<Vec<RecordBatch>> {
3640 let query = self.register_query(options)?;
3641 self.run_with_query(sql, query).await
3642 }
3643
3644 pub async fn run_with_query(
3645 &self,
3646 sql: &str,
3647 query: RegisteredSqlQuery,
3648 ) -> Result<Vec<RecordBatch>> {
3649 query.set_sql_metadata(sql);
3650 let guard = RegisteredQueryGuard::new(query.clone());
3651 self.fire_test_hook(SqlTestHookPoint::WaitingForSessionPermit);
3652 let _permit = tokio::select! {
3653 permit = Arc::clone(&self.execution_gate).acquire_owned() => permit.map_err(|_| {
3654 MongrelQueryError::InvalidQueryState("session execution gate closed".into())
3655 })?,
3656 _ = query.control().cancelled() => {
3657 let error = query.checkpoint().unwrap_err();
3658 guard.fail();
3659 return Err(error);
3660 }
3661 };
3662 query.transition(SqlQueryPhase::Queued, SqlQueryPhase::Planning)?;
3663 self.fire_test_hook(SqlTestHookPoint::Planning);
3664 query.transition(SqlQueryPhase::Planning, SqlQueryPhase::Executing)?;
3665 query.begin_statement(0);
3666 let result = CURRENT_COLLECTION_LIMITS
3667 .scope(
3668 SqlCollectionLimits::default(),
3669 CURRENT_SQL_QUERY.scope(query.clone(), async {
3670 tokio::select! {
3671 biased;
3672 result = self.run_internal(sql, &query) => result,
3673 _ = query.control().cancelled() => Err(query.checkpoint().unwrap_err()),
3674 }
3675 }),
3676 )
3677 .await;
3678 match result {
3679 Ok(batches) => {
3680 query.complete_current_statement();
3681 guard.try_complete()?;
3682 Ok(batches)
3683 }
3684 Err(error) => {
3685 if matches!(error, MongrelQueryError::ResultLimitExceeded { .. }) {
3686 guard.fail_result_limit();
3687 } else {
3688 guard.fail();
3689 }
3690 Err(error)
3691 }
3692 }
3693 }
3694
3695 pub async fn run_with_query_for_serialization(
3698 &self,
3699 sql: &str,
3700 query: RegisteredSqlQuery,
3701 ) -> Result<ManagedQueryBatches> {
3702 self.run_with_query_for_serialization_with_limits(
3703 sql,
3704 query,
3705 SqlCollectionLimits::default(),
3706 )
3707 .await
3708 }
3709
3710 pub async fn run_with_query_for_serialization_with_limits(
3712 &self,
3713 sql: &str,
3714 query: RegisteredSqlQuery,
3715 limits: SqlCollectionLimits,
3716 ) -> Result<ManagedQueryBatches> {
3717 query.set_sql_metadata(sql);
3718 let guard = RegisteredQueryGuard::new(query.clone());
3719 self.fire_test_hook(SqlTestHookPoint::WaitingForSessionPermit);
3720 let permit = tokio::select! {
3721 permit = Arc::clone(&self.execution_gate).acquire_owned() => permit.map_err(|_| {
3722 MongrelQueryError::InvalidQueryState("session execution gate closed".into())
3723 })?,
3724 _ = query.control().cancelled() => {
3725 let error = query.checkpoint().unwrap_err();
3726 guard.fail();
3727 return Err(error);
3728 }
3729 };
3730 query.transition(SqlQueryPhase::Queued, SqlQueryPhase::Planning)?;
3731 self.fire_test_hook(SqlTestHookPoint::Planning);
3732 query.transition(SqlQueryPhase::Planning, SqlQueryPhase::Executing)?;
3733 query.begin_statement(0);
3734 let result = CURRENT_COLLECTION_LIMITS
3735 .scope(
3736 limits,
3737 CURRENT_SQL_QUERY.scope(query.clone(), async {
3738 tokio::select! {
3739 biased;
3740 result = self.run_internal(sql, &query) => result,
3741 _ = query.control().cancelled() => Err(query.checkpoint().unwrap_err()),
3742 }
3743 }),
3744 )
3745 .await;
3746 match result {
3747 Ok(batches) => {
3748 query.complete_current_statement();
3749 query.begin_serialization()?;
3750 Ok(ManagedQueryBatches {
3751 batches,
3752 query,
3753 guard: Some(guard),
3754 permit: Some(permit),
3755 })
3756 }
3757 Err(error) => {
3758 if matches!(error, MongrelQueryError::ResultLimitExceeded { .. }) {
3759 guard.fail_result_limit();
3760 } else {
3761 guard.fail();
3762 }
3763 Err(error)
3764 }
3765 }
3766 }
3767
3768 pub async fn run(&self, sql: &str) -> Result<Vec<RecordBatch>> {
3769 if let Ok(query) = CURRENT_SQL_QUERY.try_with(Clone::clone) {
3770 return Box::pin(self.run_internal(sql, &query)).await;
3771 }
3772 Box::pin(self.run_with_options(sql, SqlQueryOptions::default())).await
3773 }
3774
3775 async fn run_internal(
3780 &self,
3781 sql: &str,
3782 query: &RegisteredSqlQuery,
3783 ) -> Result<Vec<RecordBatch>> {
3784 let trimmed = sql.trim();
3785 let multiple_statements = trimmed.contains(';')
3786 && !is_trigger_body(trimmed)
3787 && split_sql_statements(trimmed)
3788 .iter()
3789 .filter(|statement| !statement.trim().is_empty())
3790 .count()
3791 > 1;
3792 let rollback = commands::rollback_kind(sql);
3793 let (transaction_open, transaction_aborted) = {
3794 let transaction = self.transaction.lock();
3795 (transaction.staged_ops.is_some(), transaction.aborted)
3796 };
3797 if !multiple_statements && transaction_open && transaction_aborted && rollback.is_none() {
3798 return Err(MongrelQueryError::TransactionAborted);
3799 }
3800 let statement = TransactionStatementGuard::new(
3801 self,
3802 query,
3803 !multiple_statements
3804 && transaction_open
3805 && rollback != Some(commands::RollbackKind::Full),
3806 )?;
3807 self.fire_test_hook(SqlTestHookPoint::BeforeStatementExecution);
3808 let result = Box::pin(self.run_internal_body(sql, query)).await;
3809 if let Ok(batches) = result.as_ref() {
3810 enforce_collected_limits(batches, query, current_collection_limits())?;
3811 }
3812 if result.is_ok() {
3813 statement.complete();
3814 }
3815 result
3816 }
3817
3818 async fn run_internal_body(
3819 &self,
3820 sql: &str,
3821 query: &RegisteredSqlQuery,
3822 ) -> Result<Vec<RecordBatch>> {
3823 query.checkpoint()?;
3824 if let Some(inner) = strip_explain_query_plan(sql) {
3825 return self.explain_query_plan(inner).await;
3826 }
3827
3828 let trimmed = sql.trim();
3834 if trimmed.contains(';') && !is_trigger_body(trimmed) {
3835 let stmts = split_sql_statements(trimmed);
3836 let non_empty: Vec<&String> = stmts
3838 .iter()
3839 .filter(|s| !s.trim().is_empty() && s.trim() != ";")
3840 .collect();
3841 if non_empty.is_empty() {
3842 return Ok(Vec::new());
3843 }
3844 if stmts.len() > 1 {
3845 let mut last = Vec::new();
3846 let mut statement_index = 0;
3847 for stmt in &stmts {
3848 let stmt = stmt.trim();
3849 if stmt.is_empty() {
3850 continue;
3851 }
3852 query.begin_statement(statement_index);
3853 last = Box::pin(self.run_internal(stmt, query)).await?;
3854 query.complete_statement(statement_index);
3855 statement_index += 1;
3856 }
3857 return Ok(last);
3858 }
3859 }
3860
3861 query.checkpoint()?;
3863 if let Some(batches) = commands::try_run_command(self, sql, query).await? {
3864 return Ok(batches);
3865 }
3866
3867 let trimmed = sql.trim();
3873 if trimmed.contains(';') {
3874 let stmts = split_sql_statements(trimmed);
3875 if stmts.len() > 1 {
3876 let mut last = Vec::new();
3877 let mut statement_index = 0;
3878 for stmt in &stmts {
3879 let stmt = stmt.trim();
3880 if stmt.is_empty() {
3881 continue;
3882 }
3883 query.begin_statement(statement_index);
3884 last = Box::pin(self.run_internal(stmt, query)).await?;
3885 query.complete_statement(statement_index);
3886 statement_index += 1;
3887 }
3888 return Ok(last);
3889 }
3890 }
3891 let lower = sql.trim_start().to_lowercase();
3893 if lower.starts_with("create table") {
3894 if let Some(db) = &self.database {
3895 db.require_for(self.principal().as_ref(), &mongreldb_core::Permission::Ddl)?;
3896 let (name, schema) = parse_create_table(sql)?;
3897 db.create_table(&name, schema)?;
3898 let handle = db.table(&name)?;
3899 let provider = MongrelProvider::new_secured(
3900 handle.clone(),
3901 Arc::clone(db),
3902 name.clone(),
3903 self.principal.clone(),
3904 )?;
3905 self.ctx
3906 .register_table(&name, Arc::new(provider))
3907 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
3908 self.tables.lock().insert(name, handle);
3909 self.invalidate_prepared_statements().await;
3910 self.clear_cache();
3911 return Ok(Vec::new());
3912 }
3913 }
3914 if lower.starts_with("drop table") {
3915 if let Some(db) = &self.database {
3916 db.require_for(self.principal().as_ref(), &mongreldb_core::Permission::Ddl)?;
3917 let (name, if_exists) = parse_drop_table(sql)?;
3918 let drop_result = db.drop_table(&name);
3919 if let Err(e) = drop_result {
3920 let is_not_found = matches!(e, mongreldb_core::MongrelError::NotFound(_));
3922 if !(if_exists && is_not_found) {
3923 return Err(e.into());
3924 }
3925 } else {
3926 self.ctx
3927 .deregister_table(&name)
3928 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
3929 self.tables.lock().remove(&name);
3930 }
3931 self.invalidate_prepared_statements().await;
3932 self.clear_cache();
3933 return Ok(Vec::new());
3934 }
3935 }
3936 if lower.starts_with("alter table") {
3937 if let Some(db) = &self.database {
3938 db.require_for(self.principal().as_ref(), &mongreldb_core::Permission::Ddl)?;
3939 match parse_alter_table(sql)? {
3940 ParsedAlterTable::RenameTable { old_name, new_name } => {
3941 db.rename_table(&old_name, &new_name)?;
3942 self.ctx
3946 .deregister_table(&old_name)
3947 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
3948 self.tables.lock().remove(&old_name);
3949 let handle = db.table(&new_name)?;
3950 let provider = MongrelProvider::new_secured(
3951 handle.clone(),
3952 Arc::clone(db),
3953 new_name.clone(),
3954 self.principal.clone(),
3955 )?;
3956 self.ctx
3957 .register_table(&new_name, Arc::new(provider))
3958 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
3959 self.tables.lock().insert(new_name, handle);
3960 }
3961 ParsedAlterTable::RenameColumn {
3962 table_name,
3963 column_name,
3964 new_name,
3965 } => {
3966 db.alter_column(&table_name, &column_name, AlterColumn::rename(new_name))?;
3967 self.refresh_registered_table(db, &table_name)?;
3968 }
3969 ParsedAlterTable::AlterColumnType {
3970 table_name,
3971 column_name,
3972 ty,
3973 } => {
3974 db.alter_column(&table_name, &column_name, AlterColumn::set_type(ty))?;
3975 self.refresh_registered_table(db, &table_name)?;
3976 }
3977 ParsedAlterTable::SetNotNull {
3978 table_name,
3979 column_name,
3980 } => {
3981 let flags = current_column_flags(db, &table_name, &column_name)?
3982 .without(ColumnFlags::NULLABLE);
3983 db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
3984 self.refresh_registered_table(db, &table_name)?;
3985 }
3986 ParsedAlterTable::DropNotNull {
3987 table_name,
3988 column_name,
3989 } => {
3990 let flags = current_column_flags(db, &table_name, &column_name)?
3991 .with(ColumnFlags::NULLABLE);
3992 db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
3993 self.refresh_registered_table(db, &table_name)?;
3994 }
3995 }
3996 self.invalidate_prepared_statements().await;
3997 self.clear_cache();
3998 return Ok(Vec::new());
3999 }
4000 }
4001
4002 if let Some(as_of_query) = self.prepare_as_of_query(sql)? {
4003 return self.run_as_of(as_of_query, query).await;
4004 }
4005
4006 let resolved = self.resolve_view_sql(sql);
4009 let resolved = self.rewrite_external_module_compat_sql(&resolved);
4010 let resolved = rewrite_compat_function_calls(&resolved);
4011 let effective_sql = normalize_sql(&resolved);
4015 let sql = effective_sql.as_str();
4016 let epoch = self.cache_epoch();
4020 let key = (sql.to_string(), epoch);
4021 let has_ttl_table = self
4022 .tables
4023 .lock()
4024 .values()
4025 .any(|table| table.lock().ttl().is_some());
4026 let result_cacheable = !extended_sql_functions::contains_volatile_extended_function(sql)
4027 && !is_explain_analyze(sql)
4028 && !is_prepared_stmt_sql(sql)
4029 && !has_ttl_table;
4030 if result_cacheable {
4031 if let Some(hit) = self.cache.lock().get(&key) {
4032 return Ok((**hit).clone());
4033 }
4034 }
4035 if let Some(batches) = self.try_catalog_introspection(sql)? {
4038 if result_cacheable {
4039 query.checkpoint()?;
4040 self.cache.lock().insert(key, Arc::new(batches.clone()));
4041 }
4042 return Ok(batches);
4043 }
4044 if let Some(batches) = self.try_direct_dispatch(sql, query)? {
4049 if result_cacheable {
4050 query.checkpoint()?;
4051 self.cache.lock().insert(key, Arc::new(batches.clone()));
4052 }
4053 return Ok(batches);
4054 }
4055 let plan_start = std::time::Instant::now();
4057 let external_module_scan = self.query_references_external_module(sql);
4058 let df = self.dataframe(sql).await?;
4059 self.track_prepared_name(sql);
4062 let planning_nanos = plan_start.elapsed().as_nanos() as u64;
4065 mongreldb_core::trace::QueryTrace::record(|t| t.planning_nanos = planning_nanos);
4066
4067 let agg_key = sql_cache_key(sql);
4072 let batches = match self.try_native_aggregate(df.logical_plan(), agg_key, query) {
4073 Ok(Some(batch)) => vec![batch],
4074 _ => {
4075 match self.try_fk_join(df.logical_plan(), query) {
4079 Ok(Some(b)) => {
4080 mongreldb_core::trace::QueryTrace::record(|t| {
4082 t.join_mode = mongreldb_core::trace::JoinMode::FkBitmap;
4083 });
4084 b
4085 }
4086 _ => self.collect_dataframe(df, query).await?,
4087 }
4088 }
4089 };
4090 if external_module_scan {
4091 mongreldb_core::trace::QueryTrace::record(|t| {
4092 t.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
4093 });
4094 }
4095 if result_cacheable {
4096 query.checkpoint()?;
4097 self.cache.lock().insert(key, Arc::new(batches.clone()));
4098 }
4099 Ok(batches)
4100 }
4101
4102 pub async fn run_sql_traced(
4114 &self,
4115 sql: &str,
4116 ) -> Result<(Vec<RecordBatch>, mongreldb_core::trace::QueryTrace)> {
4117 mongreldb_core::trace::QueryTrace::push_scope();
4118 let result = self.run(sql).await;
4119 let trace = mongreldb_core::trace::QueryTrace::pop_scope();
4120 Ok((result?, trace))
4121 }
4122
4123 pub fn clear_cache(&self) {
4126 self.cache.lock().clear();
4127 self.plan_cache.lock().clear();
4128 }
4129
4130 fn track_prepared_name(&self, sql: &str) {
4134 let lower = sql.trim_start().to_ascii_lowercase();
4135 if let Some(rest) = lower.strip_prefix("prepare ") {
4136 if let Some(name) = parse_stmt_ident(rest) {
4137 self.prepared_names.lock().insert(name);
4138 }
4139 } else if let Some(rest) = lower.strip_prefix("deallocate ") {
4140 if let Some(name) = parse_stmt_ident(rest) {
4141 self.prepared_names.lock().remove(&name);
4142 }
4143 }
4144 }
4145
4146 pub async fn invalidate_prepared_statements(&self) {
4153 let names: Vec<String> = self.prepared_names.lock().drain().collect();
4154 for name in names {
4155 let sql = format!("DEALLOCATE {name}");
4159 if self.ctx.sql(&sql).await.is_err() {
4160 }
4162 }
4163 }
4164
4165 async fn explain_query_plan(&self, sql: &str) -> Result<Vec<RecordBatch>> {
4166 let explain_sql = format!("EXPLAIN {}", sql.trim().trim_end_matches(';'));
4167 let batches = self
4168 .ctx
4169 .sql(&explain_sql)
4170 .await
4171 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?
4172 .collect()
4173 .await
4174 .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
4175 let mut detail = self.mongrel_query_plan_details(sql);
4176 for batch in &batches {
4177 if batch.num_columns() < 2 {
4178 continue;
4179 }
4180 let Some(plan_type) = batch.column(0).as_any().downcast_ref::<StringArray>() else {
4181 continue;
4182 };
4183 let Some(plan) = batch.column(1).as_any().downcast_ref::<StringArray>() else {
4184 continue;
4185 };
4186 for row in 0..batch.num_rows() {
4187 let prefix = plan_type.value(row);
4188 for line in plan.value(row).lines() {
4189 let line = line.trim();
4190 if !line.is_empty() {
4191 detail.push(format!("DATAFUSION {prefix}: {line}"));
4192 }
4193 }
4194 }
4195 }
4196 if detail.is_empty() {
4197 detail.push("plan unavailable".to_string());
4198 }
4199 let ids = (0..detail.len()).map(|i| i as i64).collect::<Vec<_>>();
4200 let parents = vec![0_i64; detail.len()];
4201 let notused = vec![0_i64; detail.len()];
4202 let schema = Arc::new(arrow::datatypes::Schema::new(vec![
4203 arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int64, false),
4204 arrow::datatypes::Field::new("parent", arrow::datatypes::DataType::Int64, false),
4205 arrow::datatypes::Field::new("notused", arrow::datatypes::DataType::Int64, false),
4206 arrow::datatypes::Field::new("detail", arrow::datatypes::DataType::Utf8, false),
4207 ]));
4208 let batch = RecordBatch::try_new(
4209 schema,
4210 vec![
4211 Arc::new(Int64Array::from(ids)) as ArrayRef,
4212 Arc::new(Int64Array::from(parents)),
4213 Arc::new(Int64Array::from(notused)),
4214 Arc::new(StringArray::from(detail)),
4215 ],
4216 )
4217 .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
4218 Ok(vec![batch])
4219 }
4220
4221 fn mongrel_query_plan_details(&self, sql: &str) -> Vec<String> {
4222 use sqlparser::ast::{GroupByExpr, OrderByKind, SetExpr, Statement};
4223 use sqlparser::dialect::PostgreSqlDialect;
4224 use sqlparser::parser::Parser;
4225
4226 let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
4227 return Vec::new();
4228 };
4229 let Some(Statement::Query(query)) = stmts.first() else {
4230 return Vec::new();
4231 };
4232
4233 fn collect(session: &MongrelSession, query: &sqlparser::ast::Query, out: &mut Vec<String>) {
4234 use sqlparser::ast::{SetOperator, TableWithJoins};
4235 match query.body.as_ref() {
4236 SetExpr::Select(select) => {
4237 for TableWithJoins { relation, joins } in &select.from {
4238 session.push_table_plan(relation, select.selection.as_ref(), out);
4239 for join in joins {
4240 session.push_table_plan(&join.relation, None, out);
4241 }
4242 }
4243 if select.distinct.is_some() {
4244 out.push("USE TEMP B-TREE FOR DISTINCT".to_string());
4245 }
4246 let grouped = match &select.group_by {
4247 GroupByExpr::All(_) => true,
4248 GroupByExpr::Expressions(exprs, _) => !exprs.is_empty(),
4249 };
4250 if grouped {
4251 out.push("USE TEMP B-TREE FOR GROUP BY".to_string());
4252 }
4253 let ordered = query.order_by.as_ref().is_some_and(|order_by| {
4254 matches!(order_by.kind, OrderByKind::All(_))
4255 || matches!(&order_by.kind, OrderByKind::Expressions(exprs) if !exprs.is_empty())
4256 });
4257 if ordered {
4258 out.push("USE TEMP B-TREE FOR ORDER BY".to_string());
4259 }
4260 }
4261 SetExpr::Query(query) => collect(session, query, out),
4262 SetExpr::SetOperation {
4263 left, op, right, ..
4264 } => {
4265 let label = match op {
4266 SetOperator::Union => "COMPOUND QUERY UNION",
4267 SetOperator::Except => "COMPOUND QUERY EXCEPT",
4268 SetOperator::Intersect => "COMPOUND QUERY INTERSECT",
4269 _ => "COMPOUND QUERY",
4270 };
4271 out.push(label.to_string());
4272 collect_set_expr(session, left, out);
4273 collect_set_expr(session, right, out);
4274 }
4275 _ => {}
4276 }
4277 }
4278
4279 fn collect_set_expr(session: &MongrelSession, expr: &SetExpr, out: &mut Vec<String>) {
4280 match expr {
4281 SetExpr::Select(select) => {
4282 for table in &select.from {
4283 session.push_table_plan(&table.relation, select.selection.as_ref(), out);
4284 }
4285 }
4286 SetExpr::Query(query) => collect(session, query, out),
4287 SetExpr::SetOperation { left, right, .. } => {
4288 collect_set_expr(session, left, out);
4289 collect_set_expr(session, right, out);
4290 }
4291 _ => {}
4292 }
4293 }
4294
4295 let mut out = Vec::new();
4296 collect(self, query, &mut out);
4297 out
4298 }
4299
4300 fn push_table_plan(
4301 &self,
4302 relation: &sqlparser::ast::TableFactor,
4303 selection: Option<&sqlparser::ast::Expr>,
4304 out: &mut Vec<String>,
4305 ) {
4306 let sqlparser::ast::TableFactor::Table { name, alias, .. } = relation else {
4307 out.push("SCAN SUBQUERY".to_string());
4308 return;
4309 };
4310 let table_name = name.to_string();
4311 let display_name = alias
4312 .as_ref()
4313 .map(|alias| alias.name.value.clone())
4314 .unwrap_or_else(|| table_name.clone());
4315 let Some(handle) = self.tables.lock().get(&table_name).cloned() else {
4316 out.push(format!("SCAN {display_name}"));
4317 return;
4318 };
4319 let schema = handle.lock().schema().clone();
4320 let searchable = selection
4321 .and_then(|expr| translate_sqlparser_filter(expr, &schema))
4322 .is_some_and(|conditions| !conditions.is_empty());
4323 if searchable {
4324 out.push(format!("SEARCH {display_name} USING MONGREL INDEX"));
4325 } else {
4326 out.push(format!("SCAN {display_name}"));
4327 }
4328 }
4329
4330 fn resolve_view_sql(&self, sql: &str) -> String {
4335 let views = self.views.lock();
4336 if views.is_empty() {
4337 return sql.to_string();
4338 }
4339 let mut result = sql.to_string();
4340 for (name, view) in views.iter() {
4341 result = replace_from_view(&result, name, &view.sql);
4342 }
4343 result
4344 }
4345
4346 fn rewrite_external_module_compat_sql(&self, sql: &str) -> String {
4347 let Some(db) = &self.database else {
4348 return sql.to_string();
4349 };
4350 rewrite_fts_match_compat_sql(sql, db)
4351 }
4352
4353 fn query_references_external_module(&self, sql: &str) -> bool {
4354 use sqlparser::ast::Statement;
4355 use sqlparser::dialect::PostgreSqlDialect;
4356 use sqlparser::parser::Parser;
4357
4358 Parser::parse_sql(&PostgreSqlDialect {}, sql)
4359 .ok()
4360 .is_some_and(|statements| {
4361 statements.iter().any(|statement| match statement {
4362 Statement::Query(query) => self.query_uses_external_module(query),
4363 _ => false,
4364 })
4365 })
4366 }
4367
4368 fn query_uses_external_module(&self, query: &sqlparser::ast::Query) -> bool {
4369 self.set_expr_uses_external_module(query.body.as_ref())
4370 }
4371
4372 fn set_expr_uses_external_module(&self, expr: &sqlparser::ast::SetExpr) -> bool {
4373 use sqlparser::ast::SetExpr;
4374
4375 match expr {
4376 SetExpr::Select(select) => select
4377 .from
4378 .iter()
4379 .any(|table| self.table_with_joins_uses_external_module(table)),
4380 SetExpr::Query(query) => self.query_uses_external_module(query),
4381 SetExpr::SetOperation { left, right, .. } => {
4382 self.set_expr_uses_external_module(left)
4383 || self.set_expr_uses_external_module(right)
4384 }
4385 _ => false,
4386 }
4387 }
4388
4389 fn table_with_joins_uses_external_module(
4390 &self,
4391 table: &sqlparser::ast::TableWithJoins,
4392 ) -> bool {
4393 self.table_factor_uses_external_module(&table.relation)
4394 || table
4395 .joins
4396 .iter()
4397 .any(|join| self.table_factor_uses_external_module(&join.relation))
4398 }
4399
4400 fn table_factor_uses_external_module(&self, relation: &sqlparser::ast::TableFactor) -> bool {
4401 use sqlparser::ast::{Expr, TableFactor};
4402
4403 match relation {
4404 TableFactor::Table { name, args, .. } => {
4405 let table_name = name.to_string();
4406 self.database
4407 .as_ref()
4408 .is_some_and(|db| db.external_table(&table_name).is_some())
4409 || (args.is_some() && self.external_modules.contains(&table_name))
4410 }
4411 TableFactor::Function { name, .. } => self.external_modules.contains(&name.to_string()),
4412 TableFactor::TableFunction {
4413 expr: Expr::Function(func),
4414 ..
4415 } => self.external_modules.contains(&func.name.to_string()),
4416 TableFactor::Derived { subquery, .. } => self.query_uses_external_module(subquery),
4417 _ => false,
4418 }
4419 }
4420
4421 fn cache_epoch(&self) -> u64 {
4424 if let Some(db) = &self.database {
4425 db.visible_epoch().0
4426 } else {
4427 self.combined_epoch()
4428 }
4429 }
4430
4431 fn combined_epoch(&self) -> u64 {
4432 let Some(primary) = self.db.as_ref() else {
4433 return 0;
4434 };
4435 let mut combined = primary.lock().snapshot().epoch.0;
4436 let tables = self.tables.lock();
4437 for arc in tables.values() {
4438 if !arc.ptr_eq(primary) {
4439 let e = arc.lock().snapshot().epoch.0;
4440 combined = combined.wrapping_mul(31).wrapping_add(e);
4441 }
4442 }
4443 combined
4444 }
4445
4446 fn try_native_aggregate(
4450 &self,
4451 plan: &datafusion::logical_expr::LogicalPlan,
4452 cache_key: u64,
4453 query: &RegisteredSqlQuery,
4454 ) -> Result<Option<RecordBatch>> {
4455 if self.security_context_active() {
4456 return Ok(None);
4457 }
4458 let Some(primary) = self.db.as_ref() else {
4459 return Ok(None);
4460 };
4461 let mut db = primary.lock();
4462 let schema = db.schema().clone();
4463 let snap = db.snapshot();
4464 native_agg::try_native_aggregate(&mut db, &schema, snap, plan, cache_key, Some(query))
4465 }
4466
4467 fn try_fk_join(
4471 &self,
4472 plan: &datafusion::logical_expr::LogicalPlan,
4473 query: &RegisteredSqlQuery,
4474 ) -> Result<Option<Vec<RecordBatch>>> {
4475 if self.security_context_active() {
4476 return Ok(None);
4477 }
4478 let tables = self.tables.lock();
4479 fk_join::try_fk_join(&tables, plan, Some(query), current_collection_limits())
4480 }
4481
4482 pub fn context(&self) -> &SessionContext {
4483 &self.ctx
4484 }
4485
4486 pub fn register_scalar_udf(&self, f: ScalarUDF) {
4492 self.ctx.register_udf(f);
4493 self.clear_cache();
4494 }
4495
4496 pub fn register_aggregate_udf(&self, f: AggregateUDF) {
4498 self.ctx.register_udaf(f);
4499 self.clear_cache();
4500 }
4501
4502 pub fn register_window_udf(&self, f: WindowUDF) {
4504 self.ctx.register_udwf(f);
4505 self.clear_cache();
4506 }
4507}
4508
4509fn register_mongrel_functions(
4510 ctx: &SessionContext,
4511 sql_fn_state: Arc<extended_sql_functions::ExtendedSqlState>,
4512) {
4513 ctx.register_udf(ScalarUDF::from(udf::AnnSearchUdf::new()));
4514 ctx.register_udf(ScalarUDF::from(udf::SparseMatchUdf::new()));
4515 ctx.register_udf(ScalarUDF::from(udf::RTreeIntersectsUdf::new()));
4516 ctx.register_udf(ScalarUDF::from(udf::FtsRankUdf::new()));
4517 for udaf in percentile::percentile_udafs() {
4518 ctx.register_udaf(udaf);
4519 }
4520 extended_sql_functions::register_extended_sql_functions_with_state(ctx, sql_fn_state);
4521}
4522
4523fn is_trigger_body(sql: &str) -> bool {
4526 let lower = sql.to_ascii_lowercase();
4527 if !lower.contains("create trigger") && !lower.contains("create or replace trigger") {
4528 return false;
4529 }
4530 lower.contains("begin") && lower.contains("end")
4531}
4532
4533fn split_sql_statements(sql: &str) -> Vec<String> {
4538 let b = sql.as_bytes();
4539 let n = b.len();
4540 let mut stmts = Vec::new();
4541 let mut current = String::new();
4542 let mut i = 0;
4543 while i < n {
4544 let c = b[i];
4545 if c == b'-' && i + 1 < n && b[i + 1] == b'-' {
4547 while i < n && b[i] != b'\n' {
4548 current.push(c as char);
4549 i += 1;
4550 }
4551 continue;
4552 }
4553 if c == b'/' && i + 1 < n && b[i + 1] == b'*' {
4555 current.push_str("/*");
4556 i += 2;
4557 let mut depth = 1;
4558 while i + 1 < n && depth > 0 {
4559 if b[i] == b'/' && b[i + 1] == b'*' {
4560 depth += 1;
4561 current.push_str("/*");
4562 i += 2;
4563 } else if b[i] == b'*' && b[i + 1] == b'/' {
4564 depth -= 1;
4565 current.push_str("*/");
4566 i += 2;
4567 } else {
4568 current.push(b[i] as char);
4569 i += 1;
4570 }
4571 }
4572 continue;
4573 }
4574 if c == b'\'' {
4576 current.push('\'');
4577 i += 1;
4578 while i < n {
4579 current.push(b[i] as char);
4580 if b[i] == b'\'' {
4581 if i + 1 < n && b[i + 1] == b'\'' {
4583 current.push('\'');
4584 i += 2;
4585 continue;
4586 }
4587 i += 1;
4588 break;
4589 }
4590 i += 1;
4591 }
4592 continue;
4593 }
4594 if c == b'"' {
4596 current.push('"');
4597 i += 1;
4598 while i < n {
4599 current.push(b[i] as char);
4600 if b[i] == b'"' {
4601 if i + 1 < n && b[i + 1] == b'"' {
4602 current.push('"');
4603 i += 2;
4604 continue;
4605 }
4606 i += 1;
4607 break;
4608 }
4609 i += 1;
4610 }
4611 continue;
4612 }
4613 if c == b'$' {
4615 let start = i;
4617 let mut tag_end = i + 1;
4618 while tag_end < n && b[tag_end] != b'$' && b[tag_end].is_ascii_alphanumeric() {
4619 tag_end += 1;
4620 }
4621 if tag_end < n && b[tag_end] == b'$' {
4622 let tag = &sql[start..=tag_end];
4623 current.push_str(tag);
4624 i = tag_end + 1;
4625 while i < n {
4627 if b[i] == b'$' && sql[i..].starts_with(tag) {
4628 current.push_str(tag);
4629 i += tag.len();
4630 break;
4631 }
4632 current.push(b[i] as char);
4633 i += 1;
4634 }
4635 continue;
4636 }
4637 }
4638 if c == b';' {
4640 current.push(';');
4641 let trimmed = current.trim();
4642 if !trimmed.is_empty() && trimmed != ";" {
4643 stmts.push(current.clone());
4644 }
4645 current.clear();
4646 i += 1;
4647 continue;
4648 }
4649 current.push(c as char);
4650 i += 1;
4651 }
4652 let trimmed = current.trim();
4654 if !trimmed.is_empty() {
4655 stmts.push(current);
4656 }
4657 stmts
4658}
4659
4660fn strip_explain_query_plan(sql: &str) -> Option<&str> {
4661 let trimmed = sql.trim_start();
4662 let lower = trimmed.to_ascii_lowercase();
4663 if !lower.starts_with("explain") {
4664 return None;
4665 }
4666 let after_explain = trimmed.get(7..)?.trim_start();
4667 let after_explain_lower = after_explain.to_ascii_lowercase();
4668 if !after_explain_lower.starts_with("query") {
4669 return None;
4670 }
4671 let after_query = after_explain.get(5..)?.trim_start();
4672 let after_query_lower = after_query.to_ascii_lowercase();
4673 if !after_query_lower.starts_with("plan") {
4674 return None;
4675 }
4676 Some(after_query.get(4..)?.trim_start())
4677}
4678
4679fn is_prepared_stmt_sql(sql: &str) -> bool {
4685 let lower = sql.trim_start().to_ascii_lowercase();
4686 lower.starts_with("prepare ")
4687 || lower.starts_with("execute ")
4688 || lower.starts_with("deallocate ")
4689}
4690
4691fn parse_stmt_ident(s: &str) -> Option<String> {
4695 let tok = s
4696 .trim_start()
4697 .split(|c: char| c.is_whitespace() || c == '(')
4698 .next()?;
4699 let t = tok.trim_matches(|c| c == '"' || c == '`' || c == '\'');
4700 (!t.is_empty()).then(|| t.to_string())
4701}
4702
4703fn is_explain_analyze(sql: &str) -> bool {
4708 let trimmed = sql.trim_start();
4709 let lower = trimmed.to_ascii_lowercase();
4710 if !lower.starts_with("explain") {
4711 return false;
4712 }
4713 let after = trimmed[7..].trim_start();
4714 after.to_ascii_lowercase().starts_with("analyze")
4715}
4716
4717#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4718enum SqlCompatTokenKind {
4719 Ident,
4720 String,
4721 Dot,
4722 LParen,
4723 RParen,
4724 Comma,
4725}
4726
4727#[derive(Debug, Clone)]
4728struct SqlCompatToken {
4729 kind: SqlCompatTokenKind,
4730 raw: String,
4731 normalized: String,
4732 start: usize,
4733 end: usize,
4734}
4735
4736#[derive(Debug, Clone)]
4737struct FtsMatchBinding {
4738 query_ref: String,
4739}
4740
4741#[derive(Debug, Clone)]
4742struct SqlReplacement {
4743 start: usize,
4744 end: usize,
4745 replacement: String,
4746}
4747
4748fn rewrite_fts_match_compat_sql(sql: &str, db: &Database) -> String {
4749 let tokens = sql_compat_tokens(sql);
4750 if tokens.is_empty() {
4751 return sql.to_string();
4752 }
4753 let bindings = fts_match_bindings(sql, db, &tokens);
4754 if bindings.is_empty() {
4755 return sql.to_string();
4756 }
4757 let unique_refs = bindings
4758 .values()
4759 .map(|binding| binding.query_ref.as_str())
4760 .collect::<HashSet<_>>();
4761 let unique_binding = if unique_refs.len() == 1 {
4762 bindings.values().next().cloned()
4763 } else {
4764 None
4765 };
4766 let mut replacements = Vec::new();
4767 for (idx, token) in tokens.iter().enumerate() {
4768 if token.kind != SqlCompatTokenKind::Ident || token.normalized != "match" {
4769 continue;
4770 }
4771 let Some(rhs) = tokens.get(idx + 1) else {
4772 continue;
4773 };
4774 if rhs.kind != SqlCompatTokenKind::String {
4775 continue;
4776 }
4777 let Some((lhs_start, _lhs_end, query_ref)) =
4778 fts_match_lhs_query_ref(&tokens, idx, &bindings, unique_binding.as_ref())
4779 else {
4780 continue;
4781 };
4782 replacements.push(SqlReplacement {
4783 start: lhs_start,
4784 end: rhs.end,
4785 replacement: format!("{query_ref}.query = {}", rhs.raw),
4786 });
4787 }
4788 apply_sql_replacements(sql, &replacements)
4789}
4790
4791fn fts_match_lhs_query_ref(
4792 tokens: &[SqlCompatToken],
4793 match_idx: usize,
4794 bindings: &HashMap<String, FtsMatchBinding>,
4795 unique_binding: Option<&FtsMatchBinding>,
4796) -> Option<(usize, usize, String)> {
4797 if match_idx == 0 {
4798 return None;
4799 }
4800 let lhs = tokens.get(match_idx - 1)?;
4801 if lhs.kind != SqlCompatTokenKind::Ident {
4802 return None;
4803 }
4804
4805 if match_idx >= 3
4806 && tokens.get(match_idx - 2)?.kind == SqlCompatTokenKind::Dot
4807 && tokens.get(match_idx - 3)?.kind == SqlCompatTokenKind::Ident
4808 {
4809 let owner = tokens.get(match_idx - 3)?;
4810 let binding = bindings.get(&owner.normalized)?;
4811 if lhs.normalized == "query" || lhs.normalized == "text" {
4812 return Some((owner.start, lhs.end, binding.query_ref.clone()));
4813 }
4814 return None;
4815 }
4816
4817 if let Some(binding) = bindings.get(&lhs.normalized) {
4818 return Some((lhs.start, lhs.end, binding.query_ref.clone()));
4819 }
4820 if lhs.normalized == "text" {
4821 let binding = unique_binding?;
4822 return Some((lhs.start, lhs.end, binding.query_ref.clone()));
4823 }
4824 None
4825}
4826
4827fn fts_match_bindings(
4828 sql: &str,
4829 db: &Database,
4830 tokens: &[SqlCompatToken],
4831) -> HashMap<String, FtsMatchBinding> {
4832 let mut out = HashMap::new();
4833 let mut i = 0;
4834 while i < tokens.len() {
4835 let token = &tokens[i];
4836 let starts_table_ref = token.kind == SqlCompatTokenKind::Ident
4837 && matches!(token.normalized.as_str(), "from" | "join");
4838 if !starts_table_ref {
4839 i += 1;
4840 continue;
4841 }
4842 let mut table_idx = i + 1;
4843 if tokens
4844 .get(table_idx)
4845 .is_some_and(|token| token.kind == SqlCompatTokenKind::LParen)
4846 {
4847 i += 1;
4848 continue;
4849 }
4850 let Some(table) = tokens.get(table_idx) else {
4851 break;
4852 };
4853 if table.kind != SqlCompatTokenKind::Ident {
4854 i += 1;
4855 continue;
4856 }
4857 let mut table_name = table.normalized.clone();
4858 let mut table_ref = table.raw.clone();
4859 if let Some(qualified) = tokens.get(table_idx + 2).filter(|qualified| {
4860 tokens
4861 .get(table_idx + 1)
4862 .is_some_and(|token| token.kind == SqlCompatTokenKind::Dot)
4863 && qualified.kind == SqlCompatTokenKind::Ident
4864 }) {
4865 table_name = qualified.normalized.clone();
4866 table_ref = sql[table.start..qualified.end].to_string();
4867 table_idx += 2;
4868 }
4869 if !is_fts_docs_table(db, &table_name) {
4870 i = table_idx + 1;
4871 continue;
4872 }
4873 let mut query_ref = table_ref.clone();
4874 let mut alias_key = None;
4875 let mut next = table_idx + 1;
4876 if tokens.get(next).is_some_and(|token| {
4877 token.kind == SqlCompatTokenKind::Ident && token.normalized == "as"
4878 }) {
4879 next += 1;
4880 }
4881 if let Some(alias) = tokens.get(next) {
4882 if alias.kind == SqlCompatTokenKind::Ident && !is_table_ref_boundary(&alias.normalized)
4883 {
4884 alias_key = Some(alias.normalized.clone());
4885 query_ref = alias.raw.clone();
4886 next += 1;
4887 }
4888 }
4889 out.insert(
4890 table_name,
4891 FtsMatchBinding {
4892 query_ref: query_ref.clone(),
4893 },
4894 );
4895 if let Some(alias_key) = alias_key {
4896 out.insert(alias_key, FtsMatchBinding { query_ref });
4897 }
4898 i = next;
4899 }
4900 out
4901}
4902
4903fn is_fts_docs_table(db: &Database, name: &str) -> bool {
4904 db.external_table(name)
4905 .is_some_and(|entry| entry.module == "fts_docs")
4906}
4907
4908fn is_table_ref_boundary(normalized: &str) -> bool {
4909 matches!(
4910 normalized,
4911 "where"
4912 | "join"
4913 | "left"
4914 | "right"
4915 | "inner"
4916 | "outer"
4917 | "full"
4918 | "cross"
4919 | "on"
4920 | "using"
4921 | "group"
4922 | "order"
4923 | "having"
4924 | "limit"
4925 | "offset"
4926 | "union"
4927 | "except"
4928 | "intersect"
4929 )
4930}
4931
4932fn sql_compat_tokens(sql: &str) -> Vec<SqlCompatToken> {
4933 let bytes = sql.as_bytes();
4934 let mut tokens = Vec::new();
4935 let mut i = 0;
4936 while i < bytes.len() {
4937 match bytes[i] {
4938 b if b.is_ascii_whitespace() => i += 1,
4939 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
4940 i += 2;
4941 while i < bytes.len() && bytes[i] != b'\n' {
4942 i += 1;
4943 }
4944 }
4945 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
4946 i = skip_block_comment(bytes, i);
4947 }
4948 b'\'' => {
4949 let end = skip_quoted(bytes, i, b'\'');
4950 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
4951 i = end;
4952 }
4953 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
4954 let end = skip_quoted(bytes, i + 1, b'\'');
4955 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
4956 i = end;
4957 }
4958 b'$' => {
4959 let (end, matched) = skip_dollar_quoted(bytes, i);
4960 if matched {
4961 tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
4962 i = end;
4963 } else {
4964 i += 1;
4965 }
4966 }
4967 b'"' => {
4968 let end = skip_quoted(bytes, i, b'"');
4969 let raw = sql[i..end].to_string();
4970 let normalized = unquote_sql_ident(&raw).to_ascii_lowercase();
4971 tokens.push(SqlCompatToken {
4972 kind: SqlCompatTokenKind::Ident,
4973 raw,
4974 normalized,
4975 start: i,
4976 end,
4977 });
4978 i = end;
4979 }
4980 b'.' => {
4981 tokens.push(sql_token(sql, SqlCompatTokenKind::Dot, i, i + 1));
4982 i += 1;
4983 }
4984 b'(' => {
4985 tokens.push(sql_token(sql, SqlCompatTokenKind::LParen, i, i + 1));
4986 i += 1;
4987 }
4988 b')' => {
4989 tokens.push(sql_token(sql, SqlCompatTokenKind::RParen, i, i + 1));
4990 i += 1;
4991 }
4992 b',' => {
4993 tokens.push(sql_token(sql, SqlCompatTokenKind::Comma, i, i + 1));
4994 i += 1;
4995 }
4996 b if is_sql_ident_byte(b) => {
4997 let start = i;
4998 i += 1;
4999 while i < bytes.len() && is_sql_ident_byte(bytes[i]) {
5000 i += 1;
5001 }
5002 tokens.push(sql_token(sql, SqlCompatTokenKind::Ident, start, i));
5003 }
5004 _ => i += 1,
5005 }
5006 }
5007 tokens
5008}
5009
5010fn sql_token(sql: &str, kind: SqlCompatTokenKind, start: usize, end: usize) -> SqlCompatToken {
5011 let raw = sql[start..end].to_string();
5012 SqlCompatToken {
5013 kind,
5014 normalized: raw.to_ascii_lowercase(),
5015 raw,
5016 start,
5017 end,
5018 }
5019}
5020
5021fn unquote_sql_ident(raw: &str) -> String {
5022 if raw.len() >= 2 && raw.starts_with('"') && raw.ends_with('"') {
5023 raw[1..raw.len() - 1].replace("\"\"", "\"")
5024 } else {
5025 raw.to_string()
5026 }
5027}
5028
5029fn apply_sql_replacements(sql: &str, replacements: &[SqlReplacement]) -> String {
5030 if replacements.is_empty() {
5031 return sql.to_string();
5032 }
5033 let mut ordered = replacements.to_vec();
5034 ordered.sort_by_key(|replacement| replacement.start);
5035 let mut out = String::with_capacity(sql.len());
5036 let mut cursor = 0;
5037 for replacement in ordered {
5038 if replacement.start < cursor || replacement.end > sql.len() {
5039 continue;
5040 }
5041 out.push_str(&sql[cursor..replacement.start]);
5042 out.push_str(&replacement.replacement);
5043 cursor = replacement.end;
5044 }
5045 out.push_str(&sql[cursor..]);
5046 out
5047}
5048
5049fn rewrite_compat_function_calls(sql: &str) -> String {
5050 let bytes = sql.as_bytes();
5051 let mut out = String::with_capacity(sql.len());
5052 let mut i = 0;
5053 while i < bytes.len() {
5054 match bytes[i] {
5055 b'\'' => i = copy_quoted_to_string(&mut out, bytes, i, b'\''),
5056 b'"' => i = copy_quoted_to_string(&mut out, bytes, i, b'"'),
5057 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
5058 out.push(bytes[i] as char);
5059 i += 1;
5060 i = copy_quoted_to_string(&mut out, bytes, i, b'\'');
5061 }
5062 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
5063 out.push('-');
5064 out.push('-');
5065 i += 2;
5066 while i < bytes.len() {
5067 let ch = bytes[i] as char;
5068 out.push(ch);
5069 i += 1;
5070 if ch == '\n' {
5071 break;
5072 }
5073 }
5074 }
5075 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
5076 let start = i;
5077 i = skip_block_comment(bytes, i);
5078 out.push_str(&sql[start..i.min(bytes.len())]);
5079 }
5080 b'$' => {
5081 let start_len = out.len();
5082 let (next, matched) = copy_dollar_quoted_to_string(&mut out, bytes, i);
5083 if matched {
5084 i = next;
5085 } else {
5086 out.truncate(start_len);
5087 out.push('$');
5088 i += 1;
5089 }
5090 }
5091 b'g' | b'G' | b'm' | b'M' | b't' | b'T' => {
5092 if let Some((replacement, next)) = compat_function_rewrite_at(sql, i) {
5093 out.push_str(&replacement);
5094 i = next;
5095 } else {
5096 out.push(bytes[i] as char);
5097 i += 1;
5098 }
5099 }
5100 _ => {
5101 out.push(bytes[i] as char);
5102 i += 1;
5103 }
5104 }
5105 }
5106 out
5107}
5108
5109fn compat_function_rewrite_at(sql: &str, start: usize) -> Option<(String, usize)> {
5110 let bytes = sql.as_bytes();
5111 let (name, kind) = if ident_eq_at(bytes, start, b"max") {
5112 ("max", CompatRewriteKind::ScalarMax)
5113 } else if ident_eq_at(bytes, start, b"min") {
5114 ("min", CompatRewriteKind::ScalarMin)
5115 } else if ident_eq_at(bytes, start, b"group_concat") {
5116 ("group_concat", CompatRewriteKind::GroupConcat)
5117 } else if ident_eq_at(bytes, start, b"total") {
5118 ("total", CompatRewriteKind::Total)
5119 } else {
5120 return None;
5121 };
5122 let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
5123 let after_name = start + name.len();
5124 let after_ok = bytes
5125 .get(after_name)
5126 .is_some_and(|b| !is_sql_ident_byte(*b));
5127 if !before_ok || !after_ok {
5128 return None;
5129 }
5130 let mut open = after_name;
5131 while open < bytes.len() && bytes[open].is_ascii_whitespace() {
5132 open += 1;
5133 }
5134 if bytes.get(open) != Some(&b'(') {
5135 return None;
5136 }
5137 let summary = call_arg_summary(sql, open)?;
5138 match kind {
5139 CompatRewriteKind::ScalarMax if summary.top_level_commas > 0 => {
5140 Some(("__mongreldb_scalar_max(".to_string(), open + 1))
5141 }
5142 CompatRewriteKind::ScalarMin if summary.top_level_commas > 0 => {
5143 Some(("__mongreldb_scalar_min(".to_string(), open + 1))
5144 }
5145 CompatRewriteKind::GroupConcat => {
5146 let args = &sql[open + 1..summary.close];
5147 let rewritten = if summary.top_level_commas == 0 {
5148 format!("string_agg({args}, ',')")
5149 } else {
5150 format!("string_agg({args})")
5151 };
5152 Some((rewritten, summary.close + 1))
5153 }
5154 CompatRewriteKind::Total if summary.top_level_commas == 0 => {
5155 let args = &sql[open + 1..summary.close];
5156 let suffix_end = aggregate_suffix_end(sql, summary.close + 1);
5157 let suffix = &sql[summary.close + 1..suffix_end];
5158 Some((
5159 format!("coalesce(cast(sum({args}){suffix} as double), 0.0)"),
5160 suffix_end,
5161 ))
5162 }
5163 _ => None,
5164 }
5165}
5166
5167#[derive(Clone, Copy)]
5168enum CompatRewriteKind {
5169 ScalarMax,
5170 ScalarMin,
5171 GroupConcat,
5172 Total,
5173}
5174
5175fn ident_eq_at(bytes: &[u8], start: usize, ident: &[u8]) -> bool {
5176 bytes
5177 .get(start..start + ident.len())
5178 .is_some_and(|slice| slice.eq_ignore_ascii_case(ident))
5179}
5180
5181fn is_sql_ident_byte(b: u8) -> bool {
5182 b.is_ascii_alphanumeric() || b == b'_' || b == b'$'
5183}
5184
5185fn keyword_at(bytes: &[u8], start: usize, keyword: &[u8]) -> bool {
5186 if !ident_eq_at(bytes, start, keyword) {
5187 return false;
5188 }
5189 let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
5190 let after = start + keyword.len();
5191 let after_ok = after >= bytes.len() || !is_sql_ident_byte(bytes[after]);
5192 before_ok && after_ok
5193}
5194
5195fn skip_sql_whitespace(bytes: &[u8], mut i: usize) -> usize {
5196 while i < bytes.len() && bytes[i].is_ascii_whitespace() {
5197 i += 1;
5198 }
5199 i
5200}
5201
5202fn aggregate_suffix_end(sql: &str, start: usize) -> usize {
5203 let bytes = sql.as_bytes();
5204 let mut suffix_end = start;
5205 let mut i = skip_sql_whitespace(bytes, start);
5206
5207 if keyword_at(bytes, i, b"filter") {
5208 let open = skip_sql_whitespace(bytes, i + b"filter".len());
5209 if bytes.get(open) != Some(&b'(') {
5210 return start;
5211 }
5212 let Some(summary) = call_arg_summary(sql, open) else {
5213 return start;
5214 };
5215 suffix_end = summary.close + 1;
5216 i = skip_sql_whitespace(bytes, suffix_end);
5217 }
5218
5219 if keyword_at(bytes, i, b"over") {
5220 let after_over = skip_sql_whitespace(bytes, i + b"over".len());
5221 if bytes.get(after_over) == Some(&b'(') {
5222 let Some(summary) = call_arg_summary(sql, after_over) else {
5223 return suffix_end;
5224 };
5225 suffix_end = summary.close + 1;
5226 } else {
5227 let mut end = after_over;
5228 while end < bytes.len() && is_sql_ident_byte(bytes[end]) {
5229 end += 1;
5230 }
5231 if end > after_over {
5232 suffix_end = end;
5233 }
5234 }
5235 }
5236
5237 suffix_end
5238}
5239
5240struct CallArgSummary {
5241 close: usize,
5242 top_level_commas: usize,
5243}
5244
5245fn call_arg_summary(sql: &str, open: usize) -> Option<CallArgSummary> {
5246 let bytes = sql.as_bytes();
5247 let mut depth = 1;
5248 let mut i = open + 1;
5249 let mut top_level_commas = 0;
5250 while i < bytes.len() {
5251 match bytes[i] {
5252 b'\'' => i = skip_quoted(bytes, i, b'\''),
5253 b'"' => i = skip_quoted(bytes, i, b'"'),
5254 b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
5255 i = skip_quoted(bytes, i + 1, b'\'')
5256 }
5257 b'$' => {
5258 let (next, matched) = skip_dollar_quoted(bytes, i);
5259 i = if matched { next } else { i + 1 };
5260 }
5261 b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
5262 i += 2;
5263 while i < bytes.len() && bytes[i] != b'\n' {
5264 i += 1;
5265 }
5266 }
5267 b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
5268 i = skip_block_comment(bytes, i);
5269 }
5270 b'(' => {
5271 depth += 1;
5272 i += 1;
5273 }
5274 b')' => {
5275 depth -= 1;
5276 if depth == 0 {
5277 return Some(CallArgSummary {
5278 close: i,
5279 top_level_commas,
5280 });
5281 }
5282 i += 1;
5283 }
5284 b',' if depth == 1 => {
5285 top_level_commas += 1;
5286 i += 1;
5287 }
5288 _ => i += 1,
5289 }
5290 }
5291 None
5292}
5293
5294fn copy_quoted_to_string(out: &mut String, bytes: &[u8], start: usize, delim: u8) -> usize {
5295 let end = skip_quoted(bytes, start, delim);
5296 out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
5297 end
5298}
5299
5300fn skip_quoted(bytes: &[u8], start: usize, delim: u8) -> usize {
5301 let mut i = start;
5302 if i < bytes.len() {
5303 i += 1;
5304 }
5305 while i < bytes.len() {
5306 if bytes[i] == delim {
5307 i += 1;
5308 if i < bytes.len() && bytes[i] == delim {
5309 i += 1;
5310 continue;
5311 }
5312 break;
5313 }
5314 i += 1;
5315 }
5316 i
5317}
5318
5319fn copy_dollar_quoted_to_string(out: &mut String, bytes: &[u8], start: usize) -> (usize, bool) {
5320 let (end, matched) = skip_dollar_quoted(bytes, start);
5321 if matched {
5322 out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
5323 }
5324 (end, matched)
5325}
5326
5327fn skip_dollar_quoted(bytes: &[u8], start: usize) -> (usize, bool) {
5328 if bytes.get(start) != Some(&b'$') {
5329 return (start, false);
5330 }
5331 let mut j = start + 1;
5332 while j < bytes.len() && (bytes[j].is_ascii_alphanumeric() || bytes[j] == b'_') {
5333 j += 1;
5334 }
5335 if bytes.get(j) != Some(&b'$') {
5336 return (start, false);
5337 }
5338 let tag = &bytes[start..=j];
5339 let mut i = j + 1;
5340 while i + tag.len() <= bytes.len() {
5341 if &bytes[i..i + tag.len()] == tag {
5342 return (i + tag.len(), true);
5343 }
5344 i += 1;
5345 }
5346 (start, false)
5347}
5348
5349fn skip_block_comment(bytes: &[u8], start: usize) -> usize {
5350 let mut i = start + 2;
5351 let mut depth = 1;
5352 while i + 1 < bytes.len() && depth > 0 {
5353 if bytes[i] == b'/' && bytes[i + 1] == b'*' {
5354 depth += 1;
5355 i += 2;
5356 } else if bytes[i] == b'*' && bytes[i + 1] == b'/' {
5357 depth -= 1;
5358 i += 2;
5359 } else {
5360 i += 1;
5361 }
5362 }
5363 i
5364}
5365
5366fn sql_cache_key(sql: &str) -> u64 {
5368 use std::hash::{Hash, Hasher};
5369 let mut h = std::collections::hash_map::DefaultHasher::new();
5370 sql.hash(&mut h);
5371 h.finish()
5372}
5373
5374fn replace_from_view(sql: &str, name: &str, view_sql: &str) -> String {
5381 let lower = sql.to_ascii_lowercase();
5382 let bytes = lower.as_bytes();
5383 let name_b = name.as_bytes();
5384 let mut i = 0usize;
5385 while let Some(rel) = lower[i..].find("from") {
5386 let from_start = i + rel;
5387 let after_from = from_start + 4;
5388 i = after_from;
5389 if from_start > 0 && is_ident_byte(bytes[from_start - 1]) {
5391 continue;
5392 }
5393 let mut j = after_from;
5395 while j < bytes.len() && bytes[j].is_ascii_whitespace() {
5396 j += 1;
5397 }
5398 if j == after_from || !bytes[j..].starts_with(name_b) {
5399 continue;
5400 }
5401 let after_name = j + name_b.len();
5402 if after_name < bytes.len() && is_ident_byte(bytes[after_name]) {
5404 continue;
5405 }
5406 let mut out = String::with_capacity(sql.len() + view_sql.len() + name.len() + 8);
5409 out.push_str(&sql[..from_start]);
5410 out.push_str(&sql[from_start..j]);
5411 out.push('(');
5412 out.push_str(view_sql);
5413 out.push_str(") AS ");
5414 out.push_str(name);
5415 out.push_str(&sql[after_name..]);
5416 return out;
5417 }
5418 sql.to_string()
5419}
5420
5421fn is_ident_byte(b: u8) -> bool {
5422 b.is_ascii_alphanumeric() || b == b'_'
5423}
5424
5425fn normalize_sql(sql: &str) -> String {
5436 let b = sql.as_bytes();
5437 let n = b.len();
5438 let mut out: Vec<u8> = Vec::with_capacity(n);
5439 let mut want_space = false;
5442 let mut i = 0usize;
5443 while i < n {
5444 let c = b[i];
5445 if c.is_ascii_whitespace() {
5449 want_space = true;
5450 i += 1;
5451 continue;
5452 }
5453 if c == b'-' && i + 1 < n && b[i + 1] == b'-' {
5454 i += 2;
5456 while i < n && b[i] != b'\n' {
5457 i += 1;
5458 }
5459 want_space = !out.is_empty();
5460 continue;
5461 }
5462 if c == b'/' && i + 1 < n && b[i + 1] == b'*' {
5463 let comment_start = i;
5466 i += 2;
5467 let mut depth = 1usize;
5468 while i + 1 < n && depth > 0 {
5469 if b[i] == b'/' && b[i + 1] == b'*' {
5470 depth += 1;
5471 i += 2;
5472 } else if b[i] == b'*' && b[i + 1] == b'/' {
5473 depth -= 1;
5474 i += 2;
5475 } else {
5476 i += 1;
5477 }
5478 }
5479 if depth != 0 {
5480 if want_space && !out.is_empty() {
5484 out.push(b' ');
5485 }
5486 out.extend_from_slice(&b[comment_start..]);
5487 break;
5488 }
5489 want_space = !out.is_empty();
5490 continue;
5491 }
5492 if want_space && !out.is_empty() {
5494 out.push(b' ');
5495 }
5496 want_space = false;
5497 match c {
5498 b'E' | b'e' if i + 1 < n && b[i + 1] == b'\'' => {
5500 out.push(c);
5501 i += 1;
5502 i = copy_quoted(&mut out, b, i, n, b'\'', true);
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 i = copy_quoted(&mut out, b, i, n, b'`', false);
5518 continue;
5519 }
5520 b'[' => {
5522 i = copy_quoted(&mut out, b, i, n, b']', false);
5523 continue;
5524 }
5525 b'$' => {
5527 let (consumed, matched) = copy_dollar_quoted(&mut out, b, i, n);
5528 if matched {
5529 i = consumed;
5530 continue;
5531 }
5532 out.push(c);
5533 i += 1;
5534 continue;
5535 }
5536 _ => {
5537 out.push(c);
5538 i += 1;
5539 }
5540 }
5541 }
5542 String::from_utf8(out).unwrap_or_else(|_| sql.to_string())
5543}
5544
5545fn copy_quoted(
5550 out: &mut Vec<u8>,
5551 b: &[u8],
5552 start: usize,
5553 n: usize,
5554 delim: u8,
5555 backslash_escapes: bool,
5556) -> usize {
5557 out.push(b[start]);
5558 let mut i = start + 1;
5559 while i < n {
5560 let c = b[i];
5561 out.push(c);
5562 if backslash_escapes && c == b'\\' && i + 1 < n {
5563 out.push(b[i + 1]);
5564 i += 2;
5565 continue;
5566 }
5567 if c == delim {
5568 if i + 1 < n && b[i + 1] == delim {
5570 out.push(b[i + 1]);
5571 i += 2;
5572 continue;
5573 }
5574 return i + 1;
5575 }
5576 i += 1;
5577 }
5578 i
5579}
5580
5581fn copy_dollar_quoted(out: &mut Vec<u8>, b: &[u8], start: usize, n: usize) -> (usize, bool) {
5585 let mut j = start + 1;
5589 let tag_start = j;
5590 while j < n && b[j] != b'$' && is_dollar_tag_byte(b[j]) {
5591 j += 1;
5592 }
5593 if j >= n || b[j] != b'$' {
5594 return (start + 1, false);
5595 }
5596 if tag_start < j && !(b[tag_start].is_ascii_alphabetic() || b[tag_start] == b'_') {
5597 return (start + 1, false);
5598 }
5599 let close_end = j + 1; let delim = &b[start..close_end];
5601 out.extend_from_slice(delim);
5603 let mut k = close_end;
5605 while k + delim.len() <= n {
5606 if &b[k..k + delim.len()] == delim {
5607 out.extend_from_slice(delim);
5608 return (k + delim.len(), true);
5609 }
5610 out.push(b[k]);
5611 k += 1;
5612 }
5613 out.extend_from_slice(&b[k..n]);
5616 (n, true)
5617}
5618
5619fn is_dollar_tag_byte(b: u8) -> bool {
5620 b.is_ascii_alphanumeric() || b == b'_'
5621}
5622
5623fn strip_prefix_ci<'a>(s: &'a str, prefix: &str) -> Option<&'a str> {
5625 let bytes = s.as_bytes();
5626 let pb = prefix.as_bytes();
5627 if bytes.len() >= pb.len() && bytes[..pb.len()].eq_ignore_ascii_case(pb) {
5628 Some(&s[pb.len()..])
5629 } else {
5630 None
5631 }
5632}
5633
5634const COLUMN_CONSTRAINTS: &[(&str, u32)] = &[
5647 ("primary key", ColumnFlags::PRIMARY_KEY),
5648 ("autoincrement", ColumnFlags::AUTO_INCREMENT),
5655 ("auto_increment", ColumnFlags::AUTO_INCREMENT),
5656];
5657
5658fn parse_column_constraints(constraint_text: &str) -> ColumnFlags {
5664 let normalized = constraint_text.to_lowercase();
5665 let mut flags = ColumnFlags::empty();
5666 for (phrase, bit) in COLUMN_CONSTRAINTS {
5667 if normalized.contains(phrase) {
5668 flags = flags.with(*bit);
5669 }
5670 }
5671 flags
5672}
5673
5674fn parse_sql_type(ty_str: &str) -> Result<mongreldb_core::schema::TypeId> {
5675 use mongreldb_core::schema::TypeId;
5676
5677 match ty_str.trim().trim_end_matches(';').to_lowercase().as_str() {
5678 "bigint" | "int8" | "int64" | "integer" | "int" => Ok(TypeId::Int64),
5679 "double" | "float8" | "float64" | "real" | "float" => Ok(TypeId::Float64),
5680 "varchar" | "text" | "string" | "bytes" => Ok(TypeId::Bytes),
5681 "boolean" | "bool" => Ok(TypeId::Bool),
5682 other => Err(MongrelQueryError::Schema(format!(
5683 "unsupported column type: {other}"
5684 ))),
5685 }
5686}
5687
5688fn parse_create_table(sql: &str) -> Result<(String, mongreldb_core::schema::Schema)> {
5694 use mongreldb_core::schema::*;
5695
5696 let open = sql
5697 .find('(')
5698 .ok_or(MongrelQueryError::Schema("CREATE TABLE missing '('".into()))?;
5699 let close = sql
5700 .rfind(')')
5701 .ok_or(MongrelQueryError::Schema("CREATE TABLE missing ')'".into()))?;
5702 let head = sql[..open].trim();
5703 let after_kw = strip_prefix_ci(head, "CREATE TABLE")
5704 .or_else(|| strip_prefix_ci(head, "create table"))
5705 .unwrap_or("")
5706 .trim();
5707 let after_kw = after_kw
5709 .strip_prefix("IF NOT EXISTS")
5710 .or_else(|| after_kw.strip_prefix("if not exists"))
5711 .map(str::trim)
5712 .unwrap_or(after_kw);
5713 let name = after_kw.trim_matches('"').to_string();
5714 if name.is_empty() {
5715 return Err(MongrelQueryError::Schema(
5716 "CREATE TABLE missing table name".into(),
5717 ));
5718 }
5719
5720 let body = &sql[open + 1..close];
5721 let mut columns = Vec::new();
5722 let schema_id: u64 = 0; for (i, raw) in body.split(',').enumerate() {
5724 let part = raw.trim();
5725 if part.is_empty() {
5726 continue;
5727 }
5728 let mut tokens = part.split_whitespace();
5729 let col_name = tokens
5730 .next()
5731 .ok_or(MongrelQueryError::Schema("missing column name".into()))?
5732 .trim_matches('"');
5733 let ty_str = tokens
5734 .next()
5735 .ok_or(MongrelQueryError::Schema("missing column type".into()))?
5736 .to_lowercase();
5737 let ty = parse_sql_type(&ty_str)?;
5738 let constraint_clause: String = tokens.collect::<Vec<_>>().join(" ");
5742 let flags = parse_column_constraints(&constraint_clause);
5743 columns.push(ColumnDef {
5744 id: (i + 1) as u16,
5745 name: col_name.to_string(),
5746 ty,
5747 flags,
5748 default_value: None,
5749 });
5750 }
5751
5752 Ok((
5753 name,
5754 Schema {
5755 schema_id,
5756 columns,
5757 indexes: vec![],
5758 colocation: vec![],
5759 constraints: Default::default(),
5760 clustered: false,
5761 },
5762 ))
5763}
5764
5765fn parse_drop_table(sql: &str) -> Result<(String, bool)> {
5767 let head = sql.trim();
5768 let after_kw = strip_prefix_ci(head, "DROP TABLE")
5769 .or_else(|| strip_prefix_ci(head, "drop table"))
5770 .unwrap_or("")
5771 .trim();
5772 let (rest, if_exists) = if let Some(r) = after_kw
5774 .strip_prefix("IF EXISTS")
5775 .or_else(|| after_kw.strip_prefix("if exists"))
5776 .map(str::trim)
5777 {
5778 (r, true)
5779 } else {
5780 (after_kw, false)
5781 };
5782 let name = rest.trim_matches(';').trim_matches('"').trim();
5783 if name.is_empty() {
5784 return Err(MongrelQueryError::Schema(
5785 "DROP TABLE missing table name".into(),
5786 ));
5787 }
5788 Ok((name.to_string(), if_exists))
5789}
5790
5791enum ParsedAlterTable {
5792 RenameTable {
5793 old_name: String,
5794 new_name: String,
5795 },
5796 RenameColumn {
5797 table_name: String,
5798 column_name: String,
5799 new_name: String,
5800 },
5801 AlterColumnType {
5802 table_name: String,
5803 column_name: String,
5804 ty: mongreldb_core::schema::TypeId,
5805 },
5806 SetNotNull {
5807 table_name: String,
5808 column_name: String,
5809 },
5810 DropNotNull {
5811 table_name: String,
5812 column_name: String,
5813 },
5814}
5815
5816fn current_column_flags(db: &Arc<Database>, table: &str, column: &str) -> Result<ColumnFlags> {
5817 let handle = db.table(table)?;
5818 let table = handle.lock();
5819 table
5820 .schema()
5821 .column(column)
5822 .map(|c| c.flags)
5823 .ok_or_else(|| MongrelQueryError::Schema(format!("unknown column {column}")))
5824}
5825
5826fn parse_alter_table(sql: &str) -> Result<ParsedAlterTable> {
5827 let trimmed = strip_statement_semicolon(sql.trim());
5828 let after_kw = strip_prefix_ci(trimmed, "ALTER TABLE")
5829 .ok_or_else(|| MongrelQueryError::Schema("not an ALTER TABLE statement".into()))?
5830 .trim();
5831 let (table_name, rest) = take_sql_ident(after_kw, "ALTER TABLE missing table name")?;
5832 let rest = rest.trim();
5833
5834 if let Some(after) = strip_prefix_ci(rest, "RENAME TO") {
5835 let new_name = parse_trailing_identifier(after, "ALTER TABLE missing new table name")?;
5836 return Ok(ParsedAlterTable::RenameTable {
5837 old_name: table_name,
5838 new_name,
5839 });
5840 }
5841
5842 if let Some(after) = strip_prefix_ci(rest, "RENAME COLUMN") {
5843 let (column_name, after_col) =
5844 take_sql_ident(after, "ALTER TABLE RENAME COLUMN missing column name")?;
5845 let after_to = strip_prefix_ci(after_col.trim(), "TO").ok_or_else(|| {
5846 MongrelQueryError::Schema("ALTER TABLE RENAME COLUMN missing TO".into())
5847 })?;
5848 let new_name = parse_trailing_identifier(
5849 after_to,
5850 "ALTER TABLE RENAME COLUMN missing new column name",
5851 )?;
5852 return Ok(ParsedAlterTable::RenameColumn {
5853 table_name,
5854 column_name,
5855 new_name,
5856 });
5857 }
5858
5859 let after_alter = strip_prefix_ci(rest, "ALTER COLUMN")
5860 .or_else(|| strip_prefix_ci(rest, "ALTER"))
5861 .ok_or_else(|| {
5862 MongrelQueryError::Schema(
5863 "ALTER TABLE must be RENAME TO, RENAME COLUMN, or ALTER COLUMN".into(),
5864 )
5865 })?;
5866 let (column_name, action) =
5867 take_sql_ident(after_alter, "ALTER TABLE ALTER COLUMN missing column name")?;
5868 let action = action.trim();
5869
5870 if let Some(after_type) =
5871 strip_prefix_ci(action, "TYPE").or_else(|| strip_prefix_ci(action, "SET DATA TYPE"))
5872 {
5873 let ty = parse_type_tail(after_type)?;
5874 return Ok(ParsedAlterTable::AlterColumnType {
5875 table_name,
5876 column_name,
5877 ty,
5878 });
5879 }
5880 if strip_prefix_ci(action, "SET NOT NULL").is_some() {
5881 return Ok(ParsedAlterTable::SetNotNull {
5882 table_name,
5883 column_name,
5884 });
5885 }
5886 if strip_prefix_ci(action, "DROP NOT NULL").is_some() {
5887 return Ok(ParsedAlterTable::DropNotNull {
5888 table_name,
5889 column_name,
5890 });
5891 }
5892
5893 Err(MongrelQueryError::Schema(
5894 "unsupported ALTER COLUMN action".into(),
5895 ))
5896}
5897
5898fn strip_statement_semicolon(s: &str) -> &str {
5899 s.trim().trim_end_matches(';').trim()
5900}
5901
5902fn take_sql_ident<'a>(s: &'a str, missing: &str) -> Result<(String, &'a str)> {
5903 let s = s.trim();
5904 if s.is_empty() {
5905 return Err(MongrelQueryError::Schema(missing.into()));
5906 }
5907 if let Some(rest) = s.strip_prefix('"') {
5908 let Some(end) = rest.find('"') else {
5909 return Err(MongrelQueryError::Schema(
5910 "unterminated quoted identifier".into(),
5911 ));
5912 };
5913 let ident = rest[..end].to_string();
5914 if ident.is_empty() {
5915 return Err(MongrelQueryError::Schema(missing.into()));
5916 }
5917 return Ok((ident, &rest[end + 1..]));
5918 }
5919 let end = s.find(|c: char| c.is_ascii_whitespace()).unwrap_or(s.len());
5920 let ident = s[..end].trim_matches('"').to_string();
5921 if ident.is_empty() {
5922 return Err(MongrelQueryError::Schema(missing.into()));
5923 }
5924 Ok((ident, &s[end..]))
5925}
5926
5927fn parse_trailing_identifier(s: &str, missing: &str) -> Result<String> {
5928 let (ident, rest) = take_sql_ident(s, missing)?;
5929 if !strip_statement_semicolon(rest).is_empty() {
5930 return Err(MongrelQueryError::Schema(
5931 "unexpected tokens after identifier".into(),
5932 ));
5933 }
5934 Ok(ident)
5935}
5936
5937fn parse_type_tail(s: &str) -> Result<mongreldb_core::schema::TypeId> {
5938 let tail = strip_statement_semicolon(s);
5939 let ty = tail
5940 .split_whitespace()
5941 .next()
5942 .ok_or_else(|| MongrelQueryError::Schema("ALTER COLUMN TYPE missing type".into()))?;
5943 parse_sql_type(ty)
5944}
5945
5946#[cfg(test)]
5947mod tests {
5948 use super::*;
5949
5950 #[test]
5951 fn bounded_lru_evicts_least_recently_used() {
5952 let mut cache = BoundedLru::new(2);
5953 cache.insert("a", 1);
5954 cache.insert("b", 2);
5955 assert_eq!(cache.get("a"), Some(&1));
5956 cache.insert("c", 3);
5957 assert_eq!(cache.entries.len(), 2);
5958 assert_eq!(cache.get("b"), None);
5959 assert_eq!(cache.get("a"), Some(&1));
5960 assert_eq!(cache.get("c"), Some(&3));
5961 }
5962
5963 #[test]
5964 fn result_cache_skips_entry_over_byte_budget() {
5965 let batch = RecordBatch::try_from_iter(vec![(
5966 "value",
5967 Arc::new(Int64Array::from(vec![1_i64])) as ArrayRef,
5968 )])
5969 .unwrap();
5970 let mut cache = ResultCacheStore::new(2, 1);
5971 cache.insert(("SELECT 1".into(), 0), Arc::new(vec![batch]));
5972 assert!(cache.is_empty());
5973 assert_eq!(cache.bytes, 0);
5974 }
5975
5976 #[tokio::test]
5977 async fn collection_limit_records_typed_terminal_failure() {
5978 let dir = tempfile::tempdir().unwrap();
5979 let database = Arc::new(Database::create(dir.path()).unwrap());
5980 let session = MongrelSession::open(database).unwrap();
5981 let query = session.register_query(SqlQueryOptions::default()).unwrap();
5982 let query_id = query.id();
5983 let error = match session
5984 .run_with_query_for_serialization_with_limits(
5985 "SELECT 1 UNION ALL SELECT 2",
5986 query,
5987 SqlCollectionLimits::new(1, 1024),
5988 )
5989 .await
5990 {
5991 Ok(_) => panic!("expected result limit failure"),
5992 Err(error) => error,
5993 };
5994 assert!(matches!(
5995 error,
5996 MongrelQueryError::ResultLimitExceeded {
5997 committed: false,
5998 ..
5999 }
6000 ));
6001 let status = session.query_registry().status(query_id).unwrap();
6002 assert_eq!(
6003 status.terminal_error.unwrap().category,
6004 QueryTerminalErrorCategory::ResultLimit
6005 );
6006 }
6007
6008 #[tokio::test]
6009 async fn collection_limit_error_keeps_prior_statement_commit() {
6010 let dir = tempfile::tempdir().unwrap();
6011 let database = Arc::new(Database::create(dir.path()).unwrap());
6012 let session = MongrelSession::open(database).unwrap();
6013 let query = session.register_query(SqlQueryOptions::default()).unwrap();
6014 let error = match session
6015 .run_with_query_for_serialization_with_limits(
6016 "CREATE TABLE committed_first (id BIGINT PRIMARY KEY); \
6017 SELECT 1 UNION ALL SELECT 2",
6018 query,
6019 SqlCollectionLimits::new(1, 1024),
6020 )
6021 .await
6022 {
6023 Ok(_) => panic!("expected result limit failure"),
6024 Err(error) => error,
6025 };
6026 assert!(matches!(
6027 error,
6028 MongrelQueryError::ResultLimitExceeded {
6029 committed: true,
6030 committed_statements: 1,
6031 last_commit_epoch: Some(_),
6032 first_commit_statement_index: Some(0),
6033 last_commit_statement_index: Some(0),
6034 statement_index: 1,
6035 ..
6036 }
6037 ));
6038 }
6039
6040 #[tokio::test]
6041 async fn streaming_query_bypasses_result_cache() {
6042 use futures::StreamExt;
6043
6044 let dir = tempfile::tempdir().unwrap();
6045 let database = Arc::new(Database::create(dir.path()).unwrap());
6046 let session = MongrelSession::open(database).unwrap();
6047 let mut stream = session.run_stream("SELECT 1 AS value").await.unwrap();
6048 let batch = stream.next().await.unwrap().unwrap();
6049 assert_eq!(batch.num_rows(), 1);
6050 assert!(stream.next().await.is_none());
6051 assert!(session.cache.lock().is_empty());
6052 }
6053
6054 #[tokio::test]
6055 async fn ttl_tables_bypass_epoch_keyed_result_cache() {
6056 let dir = tempfile::tempdir().unwrap();
6057 let database = Arc::new(Database::create(dir.path()).unwrap());
6058 let session = MongrelSession::open(database).unwrap();
6059 session
6060 .run(
6061 "CREATE TABLE events (id BIGINT PRIMARY KEY, ts TIMESTAMP) \
6062 TTL_COLUMN ts TTL '1 day'",
6063 )
6064 .await
6065 .unwrap();
6066 session.run("SELECT * FROM events").await.unwrap();
6067 assert!(session.cache.lock().is_empty());
6068 }
6069
6070 #[test]
6071 fn normalize_collapses_and_trims_whitespace() {
6072 assert_eq!(normalize_sql("SELECT * FROM t"), "SELECT * FROM t");
6073 assert_eq!(normalize_sql(" SELECT * FROM t "), "SELECT * FROM t");
6074 assert_eq!(
6075 normalize_sql("\n\tSELECT\n*\nFROM\n\tt\n"),
6076 "SELECT * FROM t"
6077 );
6078 assert_eq!(
6079 normalize_sql("SELECT a, b FROM t"),
6080 normalize_sql("SELECT a, b FROM t")
6081 );
6082 }
6083
6084 #[test]
6085 fn boolean_ai_predicate_detection_uses_sql_tokens() {
6086 assert!(contains_boolean_ai_predicate(
6087 "SELECT * FROM docs WHERE ANN_SEARCH (embedding, '[1]', 1)"
6088 ));
6089 assert!(contains_boolean_ai_predicate(
6090 "PREPARE q AS SELECT * FROM docs WHERE sparse_match(sparse, '[]', 1)"
6091 ));
6092 assert!(!contains_boolean_ai_predicate(
6093 "SELECT 'ann_search(x)', sparse_search_scored('docs', 'sparse', '[]', 1) /* ann_search(x) */"
6094 ));
6095 }
6096
6097 #[test]
6098 fn normalize_preserves_string_literal_whitespace() {
6099 assert_eq!(
6100 normalize_sql("SELECT 'hello world' FROM t"),
6101 "SELECT 'hello world' FROM t"
6102 );
6103 assert_eq!(
6104 normalize_sql("SELECT 'it''s ok' FROM t"),
6105 "SELECT 'it''s ok' FROM t"
6106 );
6107 assert_eq!(
6108 normalize_sql(" SELECT 'a b' FROM t "),
6109 "SELECT 'a b' FROM t"
6110 );
6111 }
6112
6113 #[test]
6114 fn normalize_preserves_quoted_identifier_and_dollar_quote() {
6115 assert_eq!(
6116 normalize_sql(" SELECT \"my col\" FROM t "),
6117 "SELECT \"my col\" FROM t"
6118 );
6119 assert_eq!(
6120 normalize_sql(" SELECT $$a b$$ FROM t "),
6121 "SELECT $$a b$$ FROM t"
6122 );
6123 assert_eq!(
6124 normalize_sql("SELECT $tag$body with spaces$tag$ FROM t"),
6125 "SELECT $tag$body with spaces$tag$ FROM t"
6126 );
6127 }
6128
6129 #[test]
6130 fn normalize_strips_comments() {
6131 assert_eq!(
6132 normalize_sql("SELECT 1 -- trailing comment\nFROM t"),
6133 "SELECT 1 FROM t"
6134 );
6135 assert_eq!(
6136 normalize_sql("SELECT /* block */ 1 FROM t"),
6137 "SELECT 1 FROM t"
6138 );
6139 assert_eq!(
6141 normalize_sql("SELECT /* 'not a string' */ 1 FROM t"),
6142 "SELECT 1 FROM t"
6143 );
6144 assert_eq!(
6146 normalize_sql("SELECT /* outer /* inner */ still outer */ 1 FROM t"),
6147 "SELECT 1 FROM t"
6148 );
6149 assert_eq!(
6150 normalize_sql("SELECT 1 /* unterminated"),
6151 "SELECT 1 /* unterminated"
6152 );
6153 assert_eq!(normalize_sql("SELECT 1/*"), "SELECT 1/*");
6154 }
6155
6156 #[test]
6157 fn protocol_sql_fingerprint_is_literal_aware() {
6158 assert_eq!(
6159 normalized_sql_fingerprint(" INSERT INTO t VALUES (1) -- retry\n"),
6160 normalized_sql_fingerprint("INSERT INTO t VALUES (1)")
6161 );
6162 assert_ne!(
6163 normalized_sql_fingerprint("INSERT INTO t VALUES ('a b')"),
6164 normalized_sql_fingerprint("INSERT INTO t VALUES ('a b')")
6165 );
6166 }
6167
6168 #[test]
6169 fn idempotency_class_only_allows_durable_single_writes() {
6170 assert_eq!(
6171 classify_sql_idempotency("SELECT 1"),
6172 SqlIdempotencyClass::ReadOnly
6173 );
6174 for sql in [
6175 "INSERT INTO t VALUES (1)",
6176 "UPDATE t SET value = 2",
6177 "DELETE FROM t",
6178 "TRUNCATE TABLE t",
6179 "CREATE TABLE t (id BIGINT)",
6180 "ALTER TABLE t ADD COLUMN value BIGINT",
6181 "CREATE INDEX t_id ON t (id)",
6182 "CREATE MATERIALIZED VIEW visible_t AS SELECT * FROM t",
6183 "DROP TABLE t",
6184 "DROP MATERIALIZED VIEW visible_t",
6185 ] {
6186 assert_eq!(
6187 classify_sql_idempotency(sql),
6188 SqlIdempotencyClass::SingleWrite,
6189 "{sql}"
6190 );
6191 }
6192 assert_eq!(
6193 classify_sql_idempotency("INSERT INTO t VALUES (1); INSERT INTO t VALUES (2)"),
6194 SqlIdempotencyClass::Unsupported
6195 );
6196 for sql in [
6197 "BEGIN",
6198 "NOTIFY jobs, 'ready'",
6199 "LISTEN jobs",
6200 "ATTACH DATABASE 'other.db' AS other",
6201 "DETACH DATABASE other",
6202 "SHOW TABLES",
6203 "EXPLAIN SELECT 1",
6204 "PRAGMA table_info(t)",
6205 "CREATE VIEW visible_t AS SELECT * FROM t",
6206 "DROP VIEW visible_t",
6207 "DROP TABLE one, two",
6208 "DROP INDEX one, two",
6209 ] {
6210 assert_eq!(
6211 classify_sql_idempotency(sql),
6212 SqlIdempotencyClass::Unsupported,
6213 "{sql}"
6214 );
6215 }
6216 }
6217
6218 #[test]
6219 fn normalize_escape_string_preserved() {
6220 assert_eq!(
6221 normalize_sql("SELECT E'line\\nbreak' FROM t"),
6222 "SELECT E'line\\nbreak' FROM t"
6223 );
6224 assert_eq!(
6225 normalize_sql("SELECT E'it\\\'s /* data */' FROM t"),
6226 "SELECT E'it\\\'s /* data */' FROM t"
6227 );
6228 }
6229
6230 #[test]
6231 fn normalize_quoted_identifiers_preserves_comment_markers() {
6232 for sql in [
6233 "SELECT `a/* data */b` FROM t",
6234 "SELECT [a-- data b] FROM t",
6235 "SELECT `a``b` FROM t",
6236 "SELECT [a]]b] FROM t",
6237 ] {
6238 assert_eq!(normalize_sql(sql), sql);
6239 }
6240 assert_eq!(
6241 normalize_sql("SELECT $tag$unterminated"),
6242 "SELECT $tag$unterminated"
6243 );
6244 }
6245
6246 #[test]
6247 fn replace_from_view_matches_whole_word_only() {
6248 let out = replace_from_view("SELECT * FROM logs", "log", "SELECT 1");
6249 assert_eq!(out, "SELECT * FROM logs");
6250
6251 let out = replace_from_view("SELECT * FROM log", "log", "SELECT 1");
6252 assert_eq!(out, "SELECT * FROM (SELECT 1) AS log");
6253
6254 let out = replace_from_view("select * from log where x", "log", "SELECT 1");
6255 assert_eq!(out, "select * from (SELECT 1) AS log where x");
6256
6257 let out = replace_from_view("SELECT * FROM log)", "log", "SELECT 1");
6258 assert_eq!(out, "SELECT * FROM (SELECT 1) AS log)");
6259
6260 let out = replace_from_view("SELECT * xfrom log", "log", "SELECT 1");
6261 assert_eq!(out, "SELECT * xfrom log");
6262 }
6263
6264 #[test]
6265 fn compat_function_rewrite_handles_sqlite_compatibility_calls() {
6266 assert_eq!(
6267 rewrite_compat_function_calls("select max(id), min(id) from t"),
6268 "select max(id), min(id) from t"
6269 );
6270 assert_eq!(
6271 rewrite_compat_function_calls("select max(1, min(2, 3), 'max(4,5)')"),
6272 "select __mongreldb_scalar_max(1, __mongreldb_scalar_min(2, 3), 'max(4,5)')"
6273 );
6274 assert_eq!(
6275 rewrite_compat_function_calls("select /* max(1,2) */ min(1, (2 + 3))"),
6276 "select /* max(1,2) */ __mongreldb_scalar_min(1, (2 + 3))"
6277 );
6278 assert_eq!(
6279 rewrite_compat_function_calls("select max_value, min_value from t"),
6280 "select max_value, min_value from t"
6281 );
6282 assert_eq!(
6283 rewrite_compat_function_calls(
6284 "select group_concat(label), group_concat(label, '|') from t"
6285 ),
6286 "select string_agg(label, ','), string_agg(label, '|') from t"
6287 );
6288 assert_eq!(
6289 rewrite_compat_function_calls("select total(val), total(val) filter (where grp = 2) from t"),
6290 "select coalesce(cast(sum(val) as double), 0.0), coalesce(cast(sum(val) filter (where grp = 2) as double), 0.0) from t"
6291 );
6292 assert_eq!(
6293 rewrite_compat_function_calls(
6294 "select total(val) over (partition by grp order by id) from t"
6295 ),
6296 "select coalesce(cast(sum(val) over (partition by grp order by id) as double), 0.0) from t"
6297 );
6298 }
6299}