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