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