Skip to main content

mongreldb_query/
lib.rs

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