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