Skip to main content

mongreldb_query/
lib.rs

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