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