Skip to main content

mongreldb_query/
lib.rs

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