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