Skip to main content

mongreldb_query/
lib.rs

1//! DataFusion SQL + Arrow frontend for MongrelDB.
2//!
3//! [`MongrelProvider`] implements DataFusion's `TableProvider`: each `scan()`
4//! takes an MVCC snapshot of the table, materializes the visible columns, and
5//! hands DataFusion a streaming `MongrelScanExec` (see `scan.rs`) that emits one
6//! `RecordBatch` per 65 536-row chunk. DataFusion then runs the SQL —
7//! projection, filter, aggregation, limit — with its own vectorized kernels,
8//! pipelined across those small batches so a `LIMIT` short-circuits and peak
9//! memory stays bounded. MongrelDB owns storage/writes/indexes; DataFusion owns
10//! the vectorized execution.
11//!
12//! Example (skipped from doctests; see `tests/sql.rs` for runnable ones):
13//! ```ignore
14//! # use mongreldb_core::Table;
15//! # use mongreldb_query::MongrelSession;
16//! # async fn run() -> anyhow::Result<()> {
17//! let db = Table::create("travel.mongreldb", /* schema */ unimplemented!(), 1)?;
18//! let session = MongrelSession::new(db);
19//! session.register("travel_trips").await?;
20//! let batches = session.run("select * from travel_trips where cost < 300").await?;
21//! # Ok(()) }
22//! ```
23
24pub mod arrow_conv;
25mod commands;
26mod error;
27pub mod extended_sql_functions;
28mod external_modules;
29mod fk_join;
30mod native_agg;
31mod percentile;
32mod scan;
33mod scored_sql;
34mod shadow;
35mod udf;
36
37pub use error::{MongrelQueryError, Result};
38pub use external_modules::{
39    ExternalBaseWrite, ExternalModuleDescriptor, ExternalModuleIndex, ExternalModuleRegistry,
40    ExternalPlan, ExternalPlanRequest, ExternalScan, ExternalTable, ExternalTableModule,
41    ExternalTxn, ExternalWriteOp, ExternalWriteResult, ModuleConnectCtx,
42};
43
44/// True when SQL calls a ranked Boolean AI UDF. Remote servers use this to
45/// require the bounded scored table functions while embedded sessions retain
46/// the trusted Boolean intersection primitives.
47pub fn contains_boolean_ai_predicate(sql: &str) -> bool {
48    use sqlparser::dialect::GenericDialect;
49    use sqlparser::tokenizer::{Token, Tokenizer};
50
51    let Ok(tokens) = Tokenizer::new(&GenericDialect {}, sql).tokenize() else {
52        return false;
53    };
54    let tokens: Vec<_> = tokens
55        .iter()
56        .filter(|token| !matches!(token, Token::Whitespace(_)))
57        .collect();
58    tokens.windows(2).any(|tokens| {
59        matches!(
60            (tokens[0], tokens[1]),
61            (Token::Word(word), Token::LParen)
62                if word.value.eq_ignore_ascii_case("ann_search")
63                    || word.value.eq_ignore_ascii_case("sparse_match")
64        )
65    })
66}
67
68pub type MongrelRecordBatchStream = datafusion::physical_plan::SendableRecordBatchStream;
69
70use arrow::array::{Array, ArrayRef, Int64Array, StringArray};
71use arrow::datatypes::SchemaRef;
72use arrow::record_batch::RecordBatch;
73use datafusion::catalog::{Session, TableProvider};
74use datafusion::common::{DataFusionError, Result as DFResult};
75use datafusion::logical_expr::{AggregateUDF, Expr, ScalarUDF, TableType, WindowUDF};
76use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
77use datafusion::physical_plan::ExecutionPlan;
78use datafusion::prelude::SessionContext;
79use mongreldb_core::{
80    AlterColumn, ColumnFlags, Cursor, Database, OwnedSnapshotGuard, Schema as CoreSchema, Snapshot,
81    Table, TableHandle, TypeId,
82};
83use parking_lot::Mutex;
84use std::borrow::Borrow;
85use std::collections::{HashMap, HashSet};
86use std::hash::Hash;
87use std::sync::Arc;
88
89/// A MongrelDB table exposed to DataFusion. Holds the live `Table` behind a mutex;
90/// each scan takes a fresh MVCC snapshot.
91pub struct MongrelProvider {
92    db: TableHandle,
93    schema: SchemaRef,
94    core_schema: CoreSchema,
95    snapshot: Option<Snapshot>,
96    _retention: Option<Arc<OwnedSnapshotGuard>>,
97    security: Option<ProviderSecurity>,
98}
99
100#[derive(Clone)]
101struct ProviderSecurity {
102    database: Arc<Database>,
103    table: String,
104    principal: Option<mongreldb_core::Principal>,
105    principal_catalog_bound: bool,
106}
107
108impl ProviderSecurity {
109    fn principal(&self) -> mongreldb_core::Result<Option<mongreldb_core::Principal>> {
110        let Some(principal) = &self.principal else {
111            return Ok(None);
112        };
113        if self.principal_catalog_bound {
114            self.database
115                .resolve_principal(&principal.username)
116                .map(Some)
117                .ok_or(mongreldb_core::MongrelError::AuthRequired)
118        } else {
119            Ok(Some(principal.clone()))
120        }
121    }
122}
123
124#[derive(Debug, Clone)]
125pub(crate) struct ViewDef {
126    pub sql: String,
127    pub schema: CoreSchema,
128    pub input_types: HashMap<u16, Option<TypeId>>,
129}
130
131impl MongrelProvider {
132    pub fn new(db: Arc<Mutex<Table>>) -> Result<Self> {
133        Self::new_handle(db.into())
134    }
135
136    pub(crate) fn new_handle(db: TableHandle) -> Result<Self> {
137        let (schema, core_schema) = {
138            let db = db.lock();
139            (arrow_conv::arrow_schema(db.schema())?, db.schema().clone())
140        };
141        Ok(Self {
142            db,
143            schema,
144            core_schema,
145            snapshot: None,
146            _retention: None,
147            security: None,
148        })
149    }
150
151    pub(crate) fn new_secured(
152        db: TableHandle,
153        database: Arc<Database>,
154        table: String,
155        principal: Option<mongreldb_core::Principal>,
156    ) -> Result<Self> {
157        let core_schema = db.lock().schema().clone();
158        let schema = arrow_conv::arrow_schema(&core_schema)?;
159        Ok(Self {
160            db,
161            schema,
162            core_schema,
163            snapshot: None,
164            _retention: None,
165            security: Some(ProviderSecurity {
166                principal_catalog_bound: principal.as_ref().is_some_and(|principal| {
167                    database.resolve_principal(&principal.username).is_some()
168                }),
169                database,
170                table,
171                principal,
172            }),
173        })
174    }
175
176    fn new_historical(
177        db: TableHandle,
178        snapshot: Snapshot,
179        retention: OwnedSnapshotGuard,
180        security: Option<ProviderSecurity>,
181    ) -> Result<Self> {
182        let full_schema = {
183            let db = db.lock();
184            db.schema().clone()
185        };
186        let core_schema = full_schema;
187        let schema = arrow_conv::arrow_schema(&core_schema)?;
188        Ok(Self {
189            db,
190            schema,
191            core_schema,
192            snapshot: Some(snapshot),
193            _retention: Some(Arc::new(retention)),
194            security,
195        })
196    }
197
198    pub fn arrow_schema(&self) -> SchemaRef {
199        self.schema.clone()
200    }
201}
202
203impl std::fmt::Debug for MongrelProvider {
204    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
205        f.debug_struct("MongrelProvider").finish_non_exhaustive()
206    }
207}
208
209struct AsOfRegistration {
210    ctx: SessionContext,
211    table_name: String,
212}
213
214impl Drop for AsOfRegistration {
215    fn drop(&mut self) {
216        let _ = self.ctx.deregister_table(&self.table_name);
217    }
218}
219
220struct AsOfQuery {
221    sql: String,
222    registration: AsOfRegistration,
223}
224
225#[async_trait::async_trait]
226impl TableProvider for MongrelProvider {
227    fn schema(&self) -> SchemaRef {
228        self.schema.clone()
229    }
230
231    fn table_type(&self) -> TableType {
232        TableType::Base
233    }
234
235    /// Tell DataFusion which filters the pushdown serves exactly so it does not
236    /// double-filter (and, for `ann_search`, never evaluates the no-op UDF).
237    /// LIKE/FM is `Inexact`: the FM pushdown is a substring *superset*, so
238    /// DataFusion must still re-apply the real wildcard semantics.
239    fn supports_filters_pushdown(
240        &self,
241        filters: &[&Expr],
242    ) -> DFResult<Vec<datafusion::logical_expr::TableProviderFilterPushDown>> {
243        use datafusion::logical_expr::TableProviderFilterPushDown;
244        if self.snapshot.is_some() {
245            return Ok(vec![
246                TableProviderFilterPushDown::Unsupported;
247                filters.len()
248            ]);
249        }
250        let schema_ref = self.db.lock().schema().clone();
251        if self.security.is_some() {
252            return Ok(filters
253                .iter()
254                .map(|filter| {
255                    if translate_ann_search(filter, &schema_ref).is_some()
256                        || translate_sparse_match(filter, &schema_ref).is_some()
257                    {
258                        TableProviderFilterPushDown::Exact
259                    } else {
260                        TableProviderFilterPushDown::Unsupported
261                    }
262                })
263                .collect());
264        }
265        Ok(filters
266            .iter()
267            .map(|f| match translate_filter(f, &schema_ref) {
268                Some(
269                    mongreldb_core::Condition::FmContains { .. }
270                    | mongreldb_core::Condition::FmContainsAll { .. },
271                ) => TableProviderFilterPushDown::Inexact,
272                Some(_) => TableProviderFilterPushDown::Exact,
273                None => match translate_ann_search(f, &schema_ref)
274                    .or_else(|| translate_sparse_match(f, &schema_ref))
275                {
276                    Some(_) => TableProviderFilterPushDown::Exact,
277                    None => TableProviderFilterPushDown::Unsupported,
278                },
279            })
280            .collect())
281    }
282
283    async fn scan(
284        &self,
285        _state: &dyn Session,
286        projection: Option<&Vec<usize>>,
287        filters: &[Expr],
288        _limit: Option<usize>,
289    ) -> DFResult<Arc<dyn ExecutionPlan>> {
290        let core_err = |e: mongreldb_core::MongrelError| {
291            DataFusionError::External(Box::new(MongrelQueryError::Core(e)))
292        };
293        if let Some(security) = &self.security {
294            let principal = security.principal().map_err(core_err)?;
295            let allowed = security
296                .database
297                .select_column_ids_for(&security.table, principal.as_ref())
298                .map_err(core_err)?;
299            let projected = projection
300                .map(|projection| {
301                    projection
302                        .iter()
303                        .filter_map(|index| self.core_schema.columns.get(*index))
304                        .map(|column| column.id)
305                        .collect::<Vec<_>>()
306                })
307                .unwrap_or_else(|| {
308                    self.core_schema
309                        .columns
310                        .iter()
311                        .map(|column| column.id)
312                        .collect()
313                });
314            if projected.iter().any(|column| !allowed.contains(column)) {
315                security
316                    .database
317                    .require_columns_for(
318                        &security.table,
319                        mongreldb_core::ColumnOperation::Select,
320                        &projected,
321                        principal.as_ref(),
322                    )
323                    .map_err(core_err)?;
324            }
325            let ai_conditions = filters
326                .iter()
327                .filter_map(|filter| {
328                    translate_ann_search(filter, &self.core_schema)
329                        .or_else(|| translate_sparse_match(filter, &self.core_schema))
330                })
331                .collect::<Vec<_>>();
332            let mut required_columns = projected.clone();
333            required_columns.extend(mongreldb_core::query::condition_columns(&ai_conditions));
334            required_columns.sort_unstable();
335            required_columns.dedup();
336            let rows = if let Some(snapshot) = self.snapshot {
337                let rows = self.db.lock().visible_rows(snapshot).map_err(core_err)?;
338                security
339                    .database
340                    .secure_rows_for(&security.table, rows, principal.as_ref())
341                    .map_err(core_err)?
342            } else {
343                security
344                    .database
345                    .with_authorized_read(
346                        &security.table,
347                        principal.as_ref(),
348                        security.principal_catalog_bound,
349                        |table, snapshot, allowed, effective_principal| {
350                            security.database.require_columns_for(
351                                &security.table,
352                                mongreldb_core::ColumnOperation::Select,
353                                &required_columns,
354                                effective_principal,
355                            )?;
356                            let rows = if ai_conditions.is_empty() {
357                                let mut rows = table.visible_rows(snapshot)?;
358                                if let Some(allowed) = allowed {
359                                    rows.retain(|row| allowed.contains(&row.row_id));
360                                }
361                                rows
362                            } else {
363                                table.query_at_with_allowed(
364                                    &mongreldb_core::Query {
365                                        conditions: ai_conditions.clone(),
366                                        limit: Some(mongreldb_core::query::MAX_FINAL_LIMIT),
367                                        offset: 0,
368                                    },
369                                    snapshot,
370                                    allowed,
371                                )?
372                            };
373                            security.database.secure_rows_for(
374                                &security.table,
375                                rows,
376                                effective_principal,
377                            )
378                        },
379                    )
380                    .map_err(core_err)?
381            };
382            if projection.is_some_and(Vec::is_empty) {
383                return Ok(Arc::new(scan::MongrelScanExec::new_row_count(rows.len())));
384            }
385            let batch = arrow_conv::rows_to_batch(&rows, &self.core_schema)
386                .map_err(|error| DataFusionError::External(Box::new(error)))?;
387            let batch = match projection {
388                Some(projection) => batch.project(projection).map_err(|error| {
389                    DataFusionError::External(Box::new(MongrelQueryError::Arrow(error.to_string())))
390                })?,
391                None => batch,
392            };
393            let schema = batch.schema();
394            let statistics = (0..batch.num_columns())
395                .map(|_| scan::to_col_statistics(None))
396                .collect();
397            return Ok(Arc::new(scan::MongrelScanExec::new_batch(
398                schema, batch, statistics,
399            )));
400        }
401        let mut db = self.db.lock();
402        // Enforce Select permission before any read path (count metadata,
403        // count_conditions, or full scan_cursor). On a credentialless database
404        // this is a no-op.
405        db.require_select().map_err(core_err)?;
406        let historical = self.snapshot.is_some();
407        let snap = self.snapshot.unwrap_or_else(|| db.snapshot());
408        let schema_ref = db.schema().clone();
409
410        // Translate WHERE filters into index-backed Conditions.
411        let translated: Vec<mongreldb_core::Condition> = if historical {
412            Vec::new()
413        } else {
414            filters
415                .iter()
416                .filter_map(|f| {
417                    translate_filter(f, &schema_ref)
418                        .or_else(|| translate_ann_search(f, &schema_ref))
419                        .or_else(|| translate_sparse_match(f, &schema_ref))
420                })
421                .collect()
422        };
423
424        // Index-served conditions require complete live indexes; a deferred
425        // bulk load pays its one-time build here (Phase 14.7 lazy contract).
426        if !translated.is_empty() {
427            db.ensure_indexes_complete().map_err(core_err)?;
428        }
429
430        // `COUNT(*)`-style queries (empty projection) need only a row count.
431        // Unfiltered ⇒ O(1) via the maintained `live_count` metadata; a pushed
432        // WHERE ⇒ decode one column through the pushdown path to count survivors.
433        let empty_proj = projection.map(|p| p.is_empty()).unwrap_or(false);
434        if empty_proj {
435            let total: usize = if historical {
436                db.visible_rows(snap).map_err(core_err)?.len()
437            } else if translated.is_empty() {
438                mongreldb_core::trace::QueryTrace::record(|t| {
439                    t.scan_mode = mongreldb_core::trace::ScanMode::CountMetadata;
440                });
441                db.count() as usize
442            } else if let Some(count) = db.count_conditions(&translated, snap).map_err(core_err)? {
443                count as usize
444            } else {
445                match schema_ref.columns.first() {
446                    Some(cdef) => {
447                        let one = [cdef.id];
448                        let cols = match db
449                            .query_columns_native_cached(&translated, Some(&one), snap)
450                            .map_err(core_err)?
451                        {
452                            Some(c) => c,
453                            None => db
454                                .visible_columns_native(snap, Some(&one))
455                                .map_err(core_err)?,
456                        };
457                        mongreldb_core::trace::QueryTrace::record(|t| {
458                            t.scan_mode = mongreldb_core::trace::ScanMode::Materialized;
459                        });
460                        cols.first().map(|(_, c)| c.len()).unwrap_or(0)
461                    }
462                    None => 0,
463                }
464            };
465            return Ok(Arc::new(scan::MongrelScanExec::new_row_count(total)));
466        }
467
468        // Output column ids + Arrow schema for this scan, in scan-field order.
469        // DataFusion's projection already includes every column a retained
470        // (Inexact / Unsupported) filter still needs, so decoding exactly this
471        // set is correct. `None` ⇒ the full schema.
472        let (col_ids, scan_schema): (Vec<u16>, SchemaRef) = match projection {
473            Some(p) if !p.is_empty() => {
474                let ids = p.iter().map(|&idx| schema_ref.columns[idx].id).collect();
475                let fields: Vec<arrow::datatypes::Field> = p
476                    .iter()
477                    .map(|&idx| self.schema.field(idx).clone())
478                    .collect();
479                (ids, Arc::new(arrow::datatypes::Schema::new(fields)))
480            }
481            _ => (
482                schema_ref.columns.iter().map(|c| c.id).collect(),
483                self.schema.clone(),
484            ),
485        };
486
487        // Projection pairs (column id, type) in scan-field order.
488        let mut proj_pairs: Vec<(u16, mongreldb_core::schema::TypeId)> =
489            Vec::with_capacity(col_ids.len());
490        let mut types: Vec<mongreldb_core::schema::TypeId> = Vec::with_capacity(col_ids.len());
491        for cid in &col_ids {
492            let ty = schema_ref
493                .columns
494                .iter()
495                .find(|c| c.id == *cid)
496                .map(|c| c.ty.clone())
497                .ok_or_else(|| {
498                    DataFusionError::External(Box::new(MongrelQueryError::Arrow(format!(
499                        "unknown column {cid}"
500                    ))))
501                })?;
502            proj_pairs.push((*cid, ty.clone()));
503            types.push(ty);
504        }
505
506        // Phase 7.1: exact per-column min/max from page stats, but only for an
507        // unfiltered full scan over an insert-only table (gated in core). A
508        // pushed WHERE or a table with deletes ⇒ all-Absent (DataFusion scans).
509        let col_stats_map = if !historical && translated.is_empty() {
510            db.exact_column_stats(snap, &col_ids).map_err(core_err)?
511        } else {
512            None
513        };
514        let column_stats: Vec<datafusion::physical_plan::ColumnStatistics> = col_ids
515            .iter()
516            .map(|cid| scan::to_col_statistics(col_stats_map.as_ref().and_then(|m| m.get(cid))))
517            .collect();
518
519        // Phase 15.5: Arrow IPC shadow — zero-copy scan for clean single-run
520        // unfiltered tables. The shadow is a derived Arrow IPC file that was
521        // written on a prior scan; reading it avoids per-column decode entirely.
522        if !historical
523            && translated.is_empty()
524            && db.run_count() == 1
525            && db.memtable_is_empty()
526            && db.mutable_run_len() == 0
527            && db.single_run_is_clean()
528        {
529            let shadow = shadow::ArrowShadow::new(db.dir());
530            let run_ids: HashSet<u128> = db.run_ids().into_iter().collect();
531            shadow.sweep(&run_ids);
532            if let Some(&run_id) = run_ids.iter().next() {
533                if let Some(batch) = shadow.try_read(run_id) {
534                    if let Some(projected) =
535                        project_batch(&batch, &col_ids, &schema_ref, &scan_schema)
536                    {
537                        mongreldb_core::trace::QueryTrace::record(|t| {
538                            t.scan_mode = mongreldb_core::trace::ScanMode::ArrowShadow;
539                        });
540                        return Ok(Arc::new(scan::MongrelScanExec::new_batch(
541                            scan_schema,
542                            projected,
543                            column_stats,
544                        )));
545                    }
546                }
547            }
548        }
549
550        // Phase 6.2 / 16.1: drive a lazy streaming cursor that fuses the
551        // predicate, skips pages with no survivors, and decodes only the
552        // projected columns of surviving pages. `scan_cursor` picks the page-plan
553        // fast path for a single run or the k-way-merge cursor for multi-run —
554        // both avoid fully materializing every row. Anything else (e.g. an empty
555        // table with only memtable rows) falls through to materialize-then-chunk.
556        let cursor: Option<Box<dyn Cursor>> = db
557            .scan_cursor(snap, proj_pairs, &translated)
558            .map_err(core_err)?;
559        if let Some(cursor) = cursor {
560            let num_rows = cursor.remaining_rows();
561            // Phase 16.3a: extract the LIKE pattern for residual pre-filtering.
562            let residual = extract_residual_filter(filters, &col_ids, &schema_ref);
563            return Ok(Arc::new(scan::MongrelScanExec::new_cursor(
564                scan_schema,
565                types,
566                cursor,
567                num_rows,
568                column_stats,
569                residual,
570            )));
571        }
572
573        // Pushdown returns exactly `col_ids` when it accepts; the full-scan
574        // fallback returns all columns, of which we keep `col_ids`.
575        let cols = if !translated.is_empty() {
576            match db
577                .query_columns_native_cached(&translated, Some(&col_ids), snap)
578                .map_err(core_err)?
579            {
580                Some(c) => c,
581                None => db
582                    .visible_columns_native(snap, Some(&col_ids))
583                    .map_err(core_err)?,
584            }
585        } else {
586            db.visible_columns_native(snap, Some(&col_ids))
587                .map_err(core_err)?
588        };
589
590        // Order the decoded columns into scan-field order for the streaming exec.
591        let mut ordered: Vec<mongreldb_core::columnar::NativeColumn> =
592            Vec::with_capacity(col_ids.len());
593        for cid in &col_ids {
594            let col = cols
595                .iter()
596                .find(|(id, _)| id == cid)
597                .map(|(_, c)| c.clone())
598                .ok_or_else(|| {
599                    DataFusionError::External(Box::new(MongrelQueryError::Arrow(format!(
600                        "missing column {cid}"
601                    ))))
602                })?;
603            ordered.push(col);
604        }
605        let num_rows = ordered.first().map(|c| c.len()).unwrap_or(0);
606
607        // Collect data needed for the shadow write before releasing the lock.
608        let shadow_write: Option<(
609            std::path::PathBuf,
610            u128,
611            Vec<arrow::array::ArrayRef>,
612            SchemaRef,
613        )> = if !historical
614            && translated.is_empty()
615            && db.run_count() == 1
616            && db.memtable_is_empty()
617            && db.mutable_run_len() == 0
618            && db.single_run_is_clean()
619        {
620            let all_schema_ids: Vec<u16> = schema_ref.columns.iter().map(|c| c.id).collect();
621            if col_ids == all_schema_ids {
622                let dir = db.dir().to_path_buf();
623                let run_id = db.run_ids().first().copied();
624                run_id.map(|rid| {
625                    let arrays = ordered
626                        .iter()
627                        .zip(types.iter())
628                        .map(|(col, ty)| arrow_conv::native_to_array(ty.clone(), col))
629                        .collect::<Result<_>>()
630                        .unwrap_or_default();
631                    (dir, rid, arrays, scan_schema.clone())
632                })
633            } else {
634                None
635            }
636        } else {
637            None
638        };
639
640        drop(db);
641
642        // Phase 15.5: write the Arrow IPC shadow for future scans (best-effort,
643        // outside the Table lock).
644        if let Some((dir, run_id, arrays, schema)) = shadow_write {
645            if let Ok(batch) = RecordBatch::try_new(schema, arrays) {
646                shadow::ArrowShadow::new(&dir).write(run_id, &batch);
647            }
648        }
649
650        mongreldb_core::trace::QueryTrace::record(|t| {
651            t.scan_mode = mongreldb_core::trace::ScanMode::Materialized;
652            t.row_materialized = true;
653        });
654        Ok(Arc::new(scan::MongrelScanExec::new(
655            scan_schema,
656            ordered,
657            types,
658            num_rows,
659            column_stats,
660        )))
661    }
662}
663
664/// Phase 15.5: project columns from a full-schema shadow `RecordBatch` to match
665/// the scan's requested column IDs and Arrow schema. Returns `None` if any
666/// requested column is not present in the shadow batch (schema mismatch → miss).
667fn project_batch(
668    batch: &RecordBatch,
669    col_ids: &[u16],
670    schema_ref: &mongreldb_core::schema::Schema,
671    scan_schema: &arrow::datatypes::SchemaRef,
672) -> Option<RecordBatch> {
673    // Map schema column ids to field names for lookup in the shadow batch.
674    let arrays: Vec<arrow::array::ArrayRef> = col_ids
675        .iter()
676        .map(|cid| {
677            // Find the column name for this id in the live schema.
678            let name = schema_ref
679                .columns
680                .iter()
681                .find(|c| c.id == *cid)
682                .map(|c| c.name.as_str())?;
683            // Look up the column in the shadow batch by name.
684            let idx = batch.schema().index_of(name).ok()?;
685            Some(batch.column(idx).clone())
686        })
687        .collect::<Option<Vec<_>>>()?;
688    RecordBatch::try_new(scan_schema.clone(), arrays).ok()
689}
690
691/// Translate a DataFusion WHERE filter expression into a MongrelDB
692/// index-backed [`Condition`]. Supported translations (all index/scan-served by
693/// `Table::query_columns_native`):
694///
695/// * `col = literal` → [`Condition::BitmapEq`] (bitmap index) or
696///   [`Condition::Pk`] (primary key).
697/// * `col <, >, <=, >= literal` and `col BETWEEN a AND b` →
698///   [`Condition::Range`] (Int64) / [`Condition::RangeF64`] (Float64).
699/// * `col LIKE '%pat%'` → [`Condition::FmContains`] (FM index). Any `%`/`_`
700///   wildcard pattern is mapped to its longest literal segment; DataFusion
701///   re-applies the real LIKE on the returned batch, so correctness is exact
702///   even though the pushdown is a substring superset.
703///
704/// Everything else is left to DataFusion's post-scan filter. Because DataFusion
705/// always re-applies the full WHERE on the returned batch, a pushdown only ever
706/// needs to return a *superset* of the survivors — it is a pure optimization,
707/// never a correctness risk.
708pub(crate) fn translate_filter(
709    expr: &Expr,
710    schema: &mongreldb_core::Schema,
711) -> Option<mongreldb_core::Condition> {
712    use datafusion::common::ScalarValue;
713    use datafusion::logical_expr::{Between, BinaryExpr, Like, Operator};
714    use mongreldb_core::{ColumnFlags, Condition, IndexKind, TypeId, Value};
715
716    // Extended int extraction: handles every integer width (narrow ints are
717    // stored widened to Int64 internally), Date32, and all Timestamp* precision
718    // variants DataFusion emits. The numeric value is the raw i64.
719    let int_val = |s: &ScalarValue| match s {
720        ScalarValue::Int8(Some(v)) => Some(*v as i64),
721        ScalarValue::Int16(Some(v)) => Some(*v as i64),
722        ScalarValue::Int32(Some(v)) => Some(*v as i64),
723        ScalarValue::Int64(Some(v)) => Some(*v),
724        ScalarValue::UInt8(Some(v)) => Some(*v as i64),
725        ScalarValue::UInt16(Some(v)) => Some(*v as i64),
726        ScalarValue::UInt32(Some(v)) => Some(*v as i64),
727        ScalarValue::UInt64(Some(v)) => Some(*v as i64),
728        ScalarValue::Date32(Some(v)) => Some(*v as i64),
729        ScalarValue::TimestampSecond(Some(v), _) => Some(*v),
730        ScalarValue::TimestampMillisecond(Some(v), _) => Some(*v),
731        ScalarValue::TimestampMicrosecond(Some(v), _) => Some(*v),
732        ScalarValue::TimestampNanosecond(Some(v), _) => Some(*v),
733        _ => None,
734    };
735    let float_val = |s: &ScalarValue| match s {
736        ScalarValue::Float32(Some(f)) => Some(*f as f64),
737        ScalarValue::Float64(Some(f)) => Some(*f),
738        _ => None,
739    };
740    let bytes_val = |s: &ScalarValue| match s {
741        ScalarValue::Utf8(Some(s)) => Some(s.as_bytes().to_vec()),
742        _ => None,
743    };
744    let _ = bytes_val; // retained for clarity; equality uses the generic `val` below.
745
746    let val = |s: &ScalarValue| -> Option<Value> {
747        // Integer literals of any width coerce to Int64 (the storage width);
748        // Float32 widens to Float64. This keeps equality pushdown working on
749        // narrow-int / float32 bitmap and primary-key columns.
750        if let Some(i) = int_val(s) {
751            return Some(Value::Int64(i));
752        }
753        match s {
754            ScalarValue::Utf8(Some(s)) => Some(Value::Bytes(s.as_bytes().to_vec())),
755            ScalarValue::Float32(Some(f)) => Some(Value::Float64(*f as f64)),
756            ScalarValue::Float64(Some(f)) => Some(Value::Float64(*f)),
757            ScalarValue::Boolean(Some(b)) => Some(Value::Bool(*b)),
758            _ => None,
759        }
760    };
761
762    let col_def = |name: &str| schema.columns.iter().find(|c| c.name == name);
763    let has_fm = |cid: u16| {
764        schema
765            .indexes
766            .iter()
767            .any(|i| i.column_id == cid && i.kind == IndexKind::FmIndex)
768    };
769    let has_bitmap = |cid: u16| {
770        schema
771            .indexes
772            .iter()
773            .any(|i| i.column_id == cid && i.kind == IndexKind::Bitmap)
774    };
775
776    match expr {
777        // `col OP literal` (and the mirrored `literal OP col`).
778        // Also handles `col = v1 OR col = v2 OR ...` → BitmapIn.
779        Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
780            // OR-of-equalities on the same column → BitmapIn (Priority 6).
781            if *op == Operator::Or {
782                return try_or_as_bitmap_in(expr, schema);
783            }
784            // Unwrap single-layer Cast wrappers (canonicalization).
785            let left = peel_cast(left);
786            let right = peel_cast(right);
787            let (col_name, scalar, flipped) = match (left.as_ref(), right.as_ref()) {
788                (Expr::Column(c), Expr::Literal(s, _)) => (&c.name, s, false),
789                (Expr::Literal(s, _), Expr::Column(c)) => (&c.name, s, true),
790                _ => return None,
791            };
792            let op = if flipped { flip_op(*op)? } else { *op };
793            let cdef = col_def(col_name)?;
794
795            // Equality: bitmap index or primary key.
796            if op == Operator::Eq {
797                let v = val(scalar)?;
798                if has_bitmap(cdef.id) {
799                    return Some(Condition::BitmapEq {
800                        column_id: cdef.id,
801                        value: v.encode_key(),
802                    });
803                }
804                if cdef.flags.contains(ColumnFlags::PRIMARY_KEY) {
805                    return Some(Condition::Pk(v.encode_key()));
806                }
807                return None;
808            }
809
810            // Range on a typed numeric column. Every integer width is stored
811            // widened to Int64, so they all share the integer Range path.
812            match cdef.ty {
813                TypeId::Int8
814                | TypeId::Int16
815                | TypeId::Int32
816                | TypeId::Int64
817                | TypeId::UInt8
818                | TypeId::UInt16
819                | TypeId::UInt32
820                | TypeId::UInt64
821                | TypeId::TimestampNanos
822                | TypeId::Date32 => {
823                    let v = int_val(scalar)?;
824                    let (lo, hi) = int_bounds(op, v)?;
825                    Some(Condition::Range {
826                        column_id: cdef.id,
827                        lo,
828                        hi,
829                    })
830                }
831                TypeId::Float32 | TypeId::Float64 => {
832                    let v = float_val(scalar)?;
833                    let (lo, lo_inc, hi, hi_inc) = float_bounds(op, v)?;
834                    Some(Condition::RangeF64 {
835                        column_id: cdef.id,
836                        lo,
837                        lo_inclusive: lo_inc,
838                        hi,
839                        hi_inclusive: hi_inc,
840                    })
841                }
842                _ => None,
843            }
844        }
845
846        // `col BETWEEN low AND high` (and `col NOT BETWEEN ...` → skip).
847        Expr::Between(Between {
848            expr,
849            negated,
850            low,
851            high,
852        }) => {
853            if *negated {
854                return None;
855            }
856            let Expr::Column(c) = expr.as_ref() else {
857                return None;
858            };
859            let cdef = col_def(&c.name)?;
860            let (lo_s, hi_s) = match (low.as_ref(), high.as_ref()) {
861                (Expr::Literal(lo, _), Expr::Literal(hi, _)) => (lo, hi),
862                _ => return None,
863            };
864            match cdef.ty {
865                TypeId::Int8
866                | TypeId::Int16
867                | TypeId::Int32
868                | TypeId::Int64
869                | TypeId::UInt8
870                | TypeId::UInt16
871                | TypeId::UInt32
872                | TypeId::UInt64
873                | TypeId::TimestampNanos
874                | TypeId::Date32 => {
875                    let (Some(lo), Some(hi)) = (int_val(lo_s), int_val(hi_s)) else {
876                        return None;
877                    };
878                    Some(Condition::Range {
879                        column_id: cdef.id,
880                        lo,
881                        hi,
882                    })
883                }
884                TypeId::Float32 | TypeId::Float64 => {
885                    let (Some(lo), Some(hi)) = (float_val(lo_s), float_val(hi_s)) else {
886                        return None;
887                    };
888                    Some(Condition::RangeF64 {
889                        column_id: cdef.id,
890                        lo,
891                        lo_inclusive: true,
892                        hi,
893                        hi_inclusive: true,
894                    })
895                }
896                _ => None,
897            }
898        }
899
900        // `col LIKE pattern` → FM-index substring on the longest literal segment.
901        Expr::Like(Like {
902            negated,
903            expr,
904            pattern,
905            ..
906        }) => {
907            if *negated {
908                return None;
909            }
910            let Expr::Column(c) = expr.as_ref() else {
911                return None;
912            };
913            let Expr::Literal(ScalarValue::Utf8(Some(pat)), _) = pattern.as_ref() else {
914                return None;
915            };
916            let cdef = col_def(&c.name)?;
917            // §5.6: anchored prefix `LIKE 'literal%'` (no embedded wildcards)
918            // on a bitmap-indexed column → exact BytesPrefix, tighter than the
919            // FM substring superset. Checked before the FM path.
920            if has_bitmap(cdef.id) {
921                if let Some(prefix) = anchored_like_prefix(pat) {
922                    return Some(Condition::BytesPrefix {
923                        column_id: cdef.id,
924                        prefix: mongreldb_core::Value::Bytes(prefix.as_bytes().to_vec())
925                            .encode_key(),
926                    });
927                }
928            }
929            if !has_fm(cdef.id) {
930                return None;
931            }
932            // Priority 12: extract ALL literal segments (≥3 chars) and intersect
933            // their FM results for a much tighter superset than the single
934            // longest segment. Falls back to the longest when only one qualifies.
935            let segments: Vec<Vec<u8>> = pat
936                .split(['%', '_'])
937                .filter(|s| s.len() >= 3)
938                .map(|s| s.as_bytes().to_vec())
939                .collect();
940            match segments.len() {
941                0 => longest_like_segment(pat).map(|seg| Condition::FmContains {
942                    column_id: cdef.id,
943                    pattern: seg,
944                }),
945                1 => Some(Condition::FmContains {
946                    column_id: cdef.id,
947                    pattern: segments.into_iter().next().unwrap(),
948                }),
949                _ => Some(Condition::FmContainsAll {
950                    column_id: cdef.id,
951                    patterns: segments,
952                }),
953            }
954        }
955
956        // `col IN (lit1, lit2, …)` → BitmapIn (bitmap union). Phase 13.5:
957        // runtime-filter pushdown for semi-joins and IN-list filters. Only when
958        // the column has a bitmap index and every list entry is a literal.
959        Expr::InList(il) if !il.negated => {
960            let Expr::Column(c) = il.expr.as_ref() else {
961                return None;
962            };
963            let cdef = col_def(&c.name)?;
964            if !has_bitmap(cdef.id) {
965                return None;
966            }
967            let values: Vec<Vec<u8>> = il
968                .list
969                .iter()
970                .filter_map(|e| match e {
971                    Expr::Literal(s, _) => val(s).map(|v| v.encode_key()),
972                    _ => None,
973                })
974                .collect();
975            if values.is_empty() || values.len() != il.list.len() {
976                return None;
977            }
978            Some(Condition::BitmapIn {
979                column_id: cdef.id,
980                values,
981            })
982        }
983
984        // `col IS NULL` → page-stat-pruned column scan for null validity.
985        Expr::IsNull(inner) => {
986            let col_name = match inner.as_ref() {
987                Expr::Column(c) => &c.name,
988                _ => return None,
989            };
990            let cdef = col_def(col_name)?;
991            Some(Condition::IsNull { column_id: cdef.id })
992        }
993
994        // `col IS NOT NULL` → complement of IS NULL.
995        Expr::IsNotNull(inner) => {
996            let col_name = match inner.as_ref() {
997                Expr::Column(c) => &c.name,
998                _ => return None,
999            };
1000            let cdef = col_def(col_name)?;
1001            Some(Condition::IsNotNull { column_id: cdef.id })
1002        }
1003
1004        _ => None,
1005    }
1006}
1007
1008/// Phase 16.3a: extract the SQL `LIKE` pattern from `filters` for residual
1009/// pre-filtering on `NativeColumn` buffers. Returns a `ResidualFilter` when a
1010/// non-negated LIKE on a Bytes column is found among the filters.
1011pub(crate) fn extract_residual_filter(
1012    filters: &[Expr],
1013    col_ids: &[u16],
1014    schema: &mongreldb_core::Schema,
1015) -> Option<std::sync::Arc<scan::ResidualFilter>> {
1016    use datafusion::common::ScalarValue;
1017    use datafusion::logical_expr::Like;
1018    for f in filters {
1019        if let Expr::Like(Like {
1020            negated: false,
1021            expr,
1022            pattern,
1023            ..
1024        }) = f
1025        {
1026            let Expr::Column(c) = expr.as_ref() else {
1027                continue;
1028            };
1029            let Expr::Literal(ScalarValue::Utf8(Some(pat)), _) = pattern.as_ref() else {
1030                continue;
1031            };
1032            let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
1033            let col_idx = col_ids.iter().position(|&id| id == cdef.id)?;
1034            return Some(std::sync::Arc::new(scan::ResidualFilter::new(
1035                col_idx,
1036                pat.as_bytes().to_vec(),
1037            )));
1038        }
1039    }
1040    None
1041}
1042
1043/// Translate `ann_search(<embedding-col>, '<json f32 array>', k)` — the SQL hook
1044/// for HNSW semantic search — into [`Condition::Ann`]. The `ann_search` UDF is
1045/// registered by [`MongrelSession`] purely so the SQL parses; the provider's
1046/// pushdown serves the real top-k, and `supports_filters_pushdown` marks the
1047/// filter `Exact` so DataFusion never evaluates the (no-op) UDF itself.
1048pub(crate) fn translate_ann_search(
1049    expr: &Expr,
1050    schema: &mongreldb_core::Schema,
1051) -> Option<mongreldb_core::Condition> {
1052    use datafusion::common::ScalarValue;
1053    use mongreldb_core::Condition;
1054
1055    let Expr::ScalarFunction(sf) = expr else {
1056        return None;
1057    };
1058    if !sf.func.name().eq_ignore_ascii_case("ann_search") || sf.args.len() != 3 {
1059        return None;
1060    }
1061    let (Expr::Column(c), query_expr, k_expr) = (&sf.args[0], &sf.args[1], &sf.args[2]) else {
1062        return None;
1063    };
1064    let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
1065    let json = match query_expr {
1066        Expr::Literal(ScalarValue::Utf8(Some(s)), _) => s.as_str(),
1067        _ => return None,
1068    };
1069    let k: usize = match k_expr {
1070        Expr::Literal(scalar, _) => match scalar {
1071            ScalarValue::Int64(Some(k)) => usize::try_from(*k).ok()?,
1072            ScalarValue::UInt64(Some(k)) => usize::try_from(*k).ok()?,
1073            ScalarValue::Int32(Some(k)) => usize::try_from(*k).ok()?,
1074            _ => return None,
1075        },
1076        _ => return None,
1077    };
1078    let query: Vec<f32> = serde_json::from_str(json).ok()?;
1079    Some(Condition::Ann {
1080        column_id: cdef.id,
1081        query,
1082        k,
1083    })
1084}
1085
1086/// Translate `sparse_match(<sparse-col>, '<json [[token, weight], …]>', k)` —
1087/// the SQL hook for SPLADE-style sparse retrieval — into
1088/// [`Condition::SparseMatch`]. The UDF is registered by [`MongrelSession`]
1089/// purely so the SQL parses; the provider's pushdown serves the real top-k.
1090pub(crate) fn translate_sparse_match(
1091    expr: &Expr,
1092    schema: &mongreldb_core::Schema,
1093) -> Option<mongreldb_core::Condition> {
1094    use datafusion::common::ScalarValue;
1095    use mongreldb_core::Condition;
1096
1097    let Expr::ScalarFunction(sf) = expr else {
1098        return None;
1099    };
1100    if !sf.func.name().eq_ignore_ascii_case("sparse_match") || sf.args.len() != 3 {
1101        return None;
1102    }
1103    let (Expr::Column(c), query_expr, k_expr) = (&sf.args[0], &sf.args[1], &sf.args[2]) else {
1104        return None;
1105    };
1106    let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
1107    let json = match query_expr {
1108        Expr::Literal(ScalarValue::Utf8(Some(s)), _) => s.as_str(),
1109        _ => return None,
1110    };
1111    let k: usize = match k_expr {
1112        Expr::Literal(scalar, _) => match scalar {
1113            ScalarValue::Int64(Some(k)) => usize::try_from(*k).ok()?,
1114            ScalarValue::UInt64(Some(k)) => usize::try_from(*k).ok()?,
1115            ScalarValue::Int32(Some(k)) => usize::try_from(*k).ok()?,
1116            _ => return None,
1117        },
1118        _ => return None,
1119    };
1120    let query: Vec<(u32, f32)> = serde_json::from_str(json).ok()?;
1121    Some(Condition::SparseMatch {
1122        column_id: cdef.id,
1123        query,
1124        k,
1125    })
1126}
1127
1128/// Mirror a comparison operator for the `literal OP col` form.
1129fn flip_op(op: datafusion::logical_expr::Operator) -> Option<datafusion::logical_expr::Operator> {
1130    use datafusion::logical_expr::Operator;
1131    Some(match op {
1132        Operator::Eq => Operator::Eq,
1133        Operator::Lt => Operator::Gt,
1134        Operator::Gt => Operator::Lt,
1135        Operator::LtEq => Operator::GtEq,
1136        Operator::GtEq => Operator::LtEq,
1137        _ => return None,
1138    })
1139}
1140
1141/// Convert `col OP v` into inclusive Int64 `[lo, hi]` bounds (exact for all of
1142/// `<`, `>`, `<=`, `>=` via saturating ±1).
1143fn int_bounds(op: datafusion::logical_expr::Operator, v: i64) -> Option<(i64, i64)> {
1144    use datafusion::logical_expr::Operator;
1145    Some(match op {
1146        Operator::Gt => (v.saturating_add(1), i64::MAX),
1147        Operator::GtEq => (v, i64::MAX),
1148        Operator::Lt => (i64::MIN, v.saturating_sub(1)),
1149        Operator::LtEq => (i64::MIN, v),
1150        _ => return None,
1151    })
1152}
1153
1154/// Convert `col OP v` into Float64 bounds with per-bound inclusivity.
1155fn float_bounds(op: datafusion::logical_expr::Operator, v: f64) -> Option<(f64, bool, f64, bool)> {
1156    use datafusion::logical_expr::Operator;
1157    Some(match op {
1158        Operator::Gt => (v, false, f64::INFINITY, false),
1159        Operator::GtEq => (v, true, f64::INFINITY, false),
1160        Operator::Lt => (f64::NEG_INFINITY, false, v, false),
1161        Operator::LtEq => (f64::NEG_INFINITY, false, v, true),
1162        _ => return None,
1163    })
1164}
1165
1166/// Longest contiguous literal (non-`%`, non-`_`) segment of a SQL LIKE pattern;
1167/// `None` if the pattern is all wildcards (matches everything → no pushdown).
1168/// Splitting on BOTH wildcards (not just `%`) keeps the segment a true literal
1169/// substring of every match, so the FM-index search is a correct *superset* —
1170/// e.g. `%City_1%` ⇒ segment `City` (not the literal `City_1`, which no match
1171/// like `City11` actually contains). DataFusion re-applies the real wildcard.
1172fn longest_like_segment(pat: &str) -> Option<Vec<u8>> {
1173    pat.split(['%', '_'])
1174        .map(|s| s.as_bytes())
1175        .max_by_key(|s| s.len())
1176        .filter(|s| !s.is_empty())
1177        .map(|s| s.to_vec())
1178}
1179
1180/// Detect an anchored-prefix LIKE pattern: `literal%` with no `%` or `_` in
1181/// the literal part and a single trailing `%`. Returns the prefix (without the
1182/// `%`). Used to emit an exact `BytesPrefix` condition on bitmap-indexed
1183/// columns — tighter than the FM substring superset. (§5.6)
1184fn anchored_like_prefix(pat: &str) -> Option<&str> {
1185    let rest = pat.strip_suffix('%')?;
1186    if rest.is_empty() || rest.contains(['%', '_']) {
1187        return None;
1188    }
1189    Some(rest)
1190}
1191
1192/// Unwrap a single-layer `Expr::Cast` wrapper to enable pushdown for queries
1193/// like `WHERE CAST(col AS BIGINT) = 5` (canonicalization). Returns the
1194/// original `Box` unchanged for non-cast expressions.
1195fn peel_cast(expr: &Expr) -> std::borrow::Cow<'_, Expr> {
1196    match expr {
1197        Expr::Cast(datafusion::logical_expr::Cast { expr, .. }) => std::borrow::Cow::Borrowed(expr),
1198        _ => std::borrow::Cow::Borrowed(expr),
1199    }
1200}
1201
1202/// Flatten an OR tree of same-column equality comparisons into a `BitmapIn`.
1203/// Handles `col = v1 OR col = v2 OR ...` (and nested OR) that DataFusion's
1204/// optimizer may not have rewritten into `IN`. Returns `None` if the OR spans
1205/// different columns, non-equality comparisons, or a non-bitmap-indexed column.
1206fn try_or_as_bitmap_in(
1207    expr: &Expr,
1208    schema: &mongreldb_core::Schema,
1209) -> Option<mongreldb_core::Condition> {
1210    use datafusion::logical_expr::{BinaryExpr, Operator};
1211    let mut values: Vec<Vec<u8>> = Vec::new();
1212    let mut target_col: Option<u16> = None;
1213    let mut stack = vec![expr];
1214    while let Some(e) = stack.pop() {
1215        match e {
1216            Expr::BinaryExpr(BinaryExpr {
1217                left,
1218                op: Operator::Or,
1219                right,
1220            }) => {
1221                stack.push(left);
1222                stack.push(right);
1223            }
1224            Expr::BinaryExpr(BinaryExpr {
1225                left,
1226                op: Operator::Eq,
1227                right,
1228            }) => {
1229                let (col_name, scalar) = match (left.as_ref(), right.as_ref()) {
1230                    (Expr::Column(c), Expr::Literal(s, _)) => (&c.name, s),
1231                    (Expr::Literal(s, _), Expr::Column(c)) => (&c.name, s),
1232                    _ => return None,
1233                };
1234                let cdef = schema.columns.iter().find(|c| &c.name == col_name)?;
1235                if !schema
1236                    .indexes
1237                    .iter()
1238                    .any(|i| i.column_id == cdef.id && i.kind == mongreldb_core::IndexKind::Bitmap)
1239                {
1240                    return None;
1241                }
1242                match target_col {
1243                    None => target_col = Some(cdef.id),
1244                    Some(id) if id != cdef.id => return None,
1245                    _ => {}
1246                }
1247                let v = match scalar {
1248                    datafusion::common::ScalarValue::Int64(Some(v)) => {
1249                        mongreldb_core::Value::Int64(*v)
1250                    }
1251                    datafusion::common::ScalarValue::Utf8(Some(s)) => {
1252                        mongreldb_core::Value::Bytes(s.as_bytes().to_vec())
1253                    }
1254                    datafusion::common::ScalarValue::Float64(Some(f)) => {
1255                        mongreldb_core::Value::Float64(*f)
1256                    }
1257                    datafusion::common::ScalarValue::Boolean(Some(b)) => {
1258                        mongreldb_core::Value::Bool(*b)
1259                    }
1260                    _ => return None,
1261                };
1262                values.push(v.encode_key());
1263            }
1264            _ => return None,
1265        }
1266    }
1267    let col_id = target_col?;
1268    if values.is_empty() {
1269        return None;
1270    }
1271    Some(mongreldb_core::Condition::BitmapIn {
1272        column_id: col_id,
1273        values,
1274    })
1275}
1276
1277// ──────────────────────────────────────────────────────────────────────────
1278// §5.3 direct SQL dispatch: translate a sqlparser AST WHERE clause into the
1279// engine's exact Condition set (no DataFusion involvement). Only predicates
1280// whose Condition is EXACT are accepted; everything else returns None so the
1281// caller falls through to the DataFusion path (which re-applies residuals).
1282
1283fn sp_ident_name(expr: &sqlparser::ast::Expr) -> Option<&str> {
1284    use sqlparser::ast::Expr;
1285    match expr {
1286        Expr::Identifier(ident) => Some(ident.value.as_str()),
1287        Expr::CompoundIdentifier(idents) => idents.last().map(|i| i.value.as_str()),
1288        _ => None,
1289    }
1290}
1291
1292/// A sqlparser literal → core Value. Numbers widen to Int64 (or Float64 if they
1293/// don't fit i64); single-quoted strings → Bytes; booleans → Bool.
1294fn sp_literal(expr: &sqlparser::ast::Expr) -> Option<mongreldb_core::Value> {
1295    use sqlparser::ast::Expr;
1296    let v = match expr {
1297        Expr::Value(v) => v,
1298        _ => return None,
1299    };
1300    use sqlparser::ast::Value as SpValue;
1301    match &v.value {
1302        SpValue::Number(s, _) => s
1303            .parse::<i64>()
1304            .map(mongreldb_core::Value::Int64)
1305            .or_else(|_| s.parse::<f64>().map(mongreldb_core::Value::Float64))
1306            .ok(),
1307        SpValue::SingleQuotedString(s) => Some(mongreldb_core::Value::Bytes(s.as_bytes().to_vec())),
1308        SpValue::Boolean(b) => Some(mongreldb_core::Value::Bool(*b)),
1309        _ => None,
1310    }
1311}
1312
1313fn is_int_ty(ty: mongreldb_core::schema::TypeId) -> bool {
1314    use mongreldb_core::schema::TypeId::*;
1315    matches!(
1316        ty,
1317        Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 | TimestampNanos | Date32
1318    )
1319}
1320
1321fn is_float_ty(ty: mongreldb_core::schema::TypeId) -> bool {
1322    matches!(
1323        ty,
1324        mongreldb_core::schema::TypeId::Float32 | mongreldb_core::schema::TypeId::Float64
1325    )
1326}
1327
1328/// Translate ONE sqlparser predicate `Expr` into one exact `Condition`.
1329/// Returns `None` for anything inexact or unsupported (→ caller falls through).
1330fn translate_sqlparser_predicate(
1331    expr: &sqlparser::ast::Expr,
1332    schema: &mongreldb_core::Schema,
1333) -> Option<mongreldb_core::Condition> {
1334    use mongreldb_core::{schema::ColumnFlags, Condition, IndexKind, Value};
1335    use sqlparser::ast::{BinaryOperator, Expr};
1336
1337    let col_def = |name: &str| schema.columns.iter().find(|c| c.name == name);
1338    let has_bitmap = |cid: u16| {
1339        schema
1340            .indexes
1341            .iter()
1342            .any(|i| i.column_id == cid && i.kind == IndexKind::Bitmap)
1343    };
1344
1345    match expr {
1346        // `a = b OR a = c …` (one column, all literals) → BitmapIn.
1347        Expr::BinaryOp {
1348            left,
1349            op: BinaryOperator::Or,
1350            right,
1351        } => {
1352            let mut values: Vec<Vec<u8>> = Vec::new();
1353            let mut target: Option<u16> = None;
1354            let mut stack: Vec<&Expr> = vec![left.as_ref(), right.as_ref()];
1355            while let Some(e) = stack.pop() {
1356                match e {
1357                    Expr::BinaryOp {
1358                        left,
1359                        op: BinaryOperator::Or,
1360                        right,
1361                    } => {
1362                        stack.push(left.as_ref());
1363                        stack.push(right.as_ref());
1364                    }
1365                    Expr::BinaryOp {
1366                        left,
1367                        op: BinaryOperator::Eq,
1368                        right,
1369                    } => {
1370                        let (name, lit) = match (left.as_ref(), right.as_ref()) {
1371                            (l, r) if sp_ident_name(l).is_some() && sp_literal(r).is_some() => {
1372                                (l, r)
1373                            }
1374                            (l, r) if sp_ident_name(r).is_some() && sp_literal(l).is_some() => {
1375                                (r, l)
1376                            }
1377                            _ => return None,
1378                        };
1379                        let cdef = col_def(sp_ident_name(name)?)?;
1380                        if !has_bitmap(cdef.id) {
1381                            return None;
1382                        }
1383                        match target {
1384                            None => target = Some(cdef.id),
1385                            Some(id) if id != cdef.id => return None,
1386                            _ => {}
1387                        }
1388                        values.push(sp_literal(lit)?.encode_key());
1389                    }
1390                    _ => return None,
1391                }
1392            }
1393            let cid = target?;
1394            (!values.is_empty()).then_some(Condition::BitmapIn {
1395                column_id: cid,
1396                values,
1397            })
1398        }
1399        // Comparison `col OP literal` (or mirrored).
1400        Expr::BinaryOp { left, op, right } => {
1401            let flipped;
1402            let (col_expr, lit_expr) = match (
1403                sp_ident_name(left),
1404                sp_literal(right),
1405                sp_ident_name(right),
1406                sp_literal(left),
1407            ) {
1408                (Some(_), Some(_), _, _) => {
1409                    flipped = false;
1410                    (left.as_ref(), right.as_ref())
1411                }
1412                (_, _, Some(_), Some(_)) => {
1413                    flipped = true;
1414                    (right.as_ref(), left.as_ref())
1415                }
1416                _ => return None,
1417            };
1418            let name = sp_ident_name(col_expr)?;
1419            let cdef = col_def(name)?;
1420            let v = sp_literal(lit_expr)?;
1421            use sqlparser::ast::BinaryOperator::*;
1422            // Inline the comparison→Range/RangeF64 bounds, fusing the flip
1423            // (BinaryOperator is not Copy, so we match the &op directly).
1424            match op {
1425                Eq => {
1426                    if has_bitmap(cdef.id) {
1427                        Some(Condition::BitmapEq {
1428                            column_id: cdef.id,
1429                            value: v.encode_key(),
1430                        })
1431                    } else if cdef.flags.contains(ColumnFlags::PRIMARY_KEY) {
1432                        Some(Condition::Pk(v.encode_key()))
1433                    } else {
1434                        None
1435                    }
1436                }
1437                Lt | LtEq | Gt | GtEq if is_int_ty(cdef.ty.clone()) => {
1438                    let n = match v {
1439                        Value::Int64(n) => n,
1440                        _ => return None,
1441                    };
1442                    // `col OP v`, or the mirrored `v OP col` with the flipped op.
1443                    let (lo, hi) = match (flipped, op) {
1444                        (false, Lt) | (true, Gt) => (i64::MIN, n.saturating_sub(1)),
1445                        (false, Gt) | (true, Lt) => (n.saturating_add(1), i64::MAX),
1446                        (false, LtEq) | (true, GtEq) => (i64::MIN, n),
1447                        (false, GtEq) | (true, LtEq) => (n, i64::MAX),
1448                        _ => (i64::MIN, i64::MAX),
1449                    };
1450                    Some(Condition::Range {
1451                        column_id: cdef.id,
1452                        lo,
1453                        hi,
1454                    })
1455                }
1456                Lt | LtEq | Gt | GtEq if is_float_ty(cdef.ty.clone()) => {
1457                    let f = match v {
1458                        Value::Float64(f) => f,
1459                        _ => return None,
1460                    };
1461                    let (lo, li, hi, hi_i) = match (flipped, op) {
1462                        (false, Lt) | (true, Gt) => (f64::NEG_INFINITY, true, f, false),
1463                        (false, Gt) | (true, Lt) => (f, false, f64::INFINITY, true),
1464                        (false, LtEq) | (true, GtEq) => (f64::NEG_INFINITY, true, f, true),
1465                        (false, GtEq) | (true, LtEq) => (f, true, f64::INFINITY, true),
1466                        _ => (f64::NEG_INFINITY, true, f64::INFINITY, true),
1467                    };
1468                    Some(Condition::RangeF64 {
1469                        column_id: cdef.id,
1470                        lo,
1471                        lo_inclusive: li,
1472                        hi,
1473                        hi_inclusive: hi_i,
1474                    })
1475                }
1476                _ => None,
1477            }
1478        }
1479        Expr::Between {
1480            expr,
1481            negated,
1482            low,
1483            high,
1484        } if !negated => {
1485            let name = sp_ident_name(expr)?;
1486            let cdef = col_def(name)?;
1487            if is_int_ty(cdef.ty.clone()) {
1488                let lo = match sp_literal(low)? {
1489                    Value::Int64(n) => n,
1490                    _ => return None,
1491                };
1492                let hi = match sp_literal(high)? {
1493                    Value::Int64(n) => n,
1494                    _ => return None,
1495                };
1496                Some(Condition::Range {
1497                    column_id: cdef.id,
1498                    lo,
1499                    hi,
1500                })
1501            } else if is_float_ty(cdef.ty.clone()) {
1502                let lo = match sp_literal(low)? {
1503                    Value::Float64(f) => f,
1504                    _ => return None,
1505                };
1506                let hi = match sp_literal(high)? {
1507                    Value::Float64(f) => f,
1508                    _ => return None,
1509                };
1510                Some(Condition::RangeF64 {
1511                    column_id: cdef.id,
1512                    lo,
1513                    lo_inclusive: true,
1514                    hi,
1515                    hi_inclusive: true,
1516                })
1517            } else {
1518                None
1519            }
1520        }
1521        Expr::InList {
1522            expr,
1523            list,
1524            negated,
1525        } if !negated => {
1526            let name = sp_ident_name(expr)?;
1527            let cdef = col_def(name)?;
1528            if !has_bitmap(cdef.id) {
1529                return None;
1530            }
1531            let values: Vec<Vec<u8>> = list
1532                .iter()
1533                .map(|e| sp_literal(e).map(|v| v.encode_key()))
1534                .collect::<Option<_>>()?;
1535            (!values.is_empty()).then_some(Condition::BitmapIn {
1536                column_id: cdef.id,
1537                values,
1538            })
1539        }
1540        // `col IS NULL` / `col IS NOT NULL`. sqlparser 0.62 represents these as
1541        // `IsNull(expr)` / `IsNotNull(expr)` (and, defensively, `IsBoolean`).
1542        Expr::IsNull(inner) => {
1543            let cid = col_def(sp_ident_name(inner)?)?.id;
1544            Some(Condition::IsNull { column_id: cid })
1545        }
1546        Expr::IsNotNull(inner) => {
1547            let cid = col_def(sp_ident_name(inner)?)?.id;
1548            Some(Condition::IsNotNull { column_id: cid })
1549        }
1550        _ => None,
1551    }
1552}
1553
1554/// Split a top-level AND tree into conjuncts, translating each to an exact
1555/// Condition. Returns `None` if any conjunct is inexact/unsupported.
1556fn translate_sqlparser_filter(
1557    expr: &sqlparser::ast::Expr,
1558    schema: &mongreldb_core::Schema,
1559) -> Option<Vec<mongreldb_core::Condition>> {
1560    use sqlparser::ast::{BinaryOperator, Expr};
1561    let mut out = Vec::new();
1562    let mut stack = vec![expr];
1563    while let Some(e) = stack.pop() {
1564        match e {
1565            Expr::BinaryOp {
1566                left,
1567                op: BinaryOperator::And,
1568                right,
1569            } => {
1570                stack.push(left.as_ref());
1571                stack.push(right.as_ref());
1572            }
1573            other => out.push(translate_sqlparser_predicate(other, schema)?),
1574        }
1575    }
1576    Some(out)
1577}
1578
1579/// Convenience wrapper: a DataFusion `SessionContext` bound to a live MongrelDB,
1580/// with a result cache keyed by `(sql, snapshot_epoch)` that auto-invalidates
1581/// when a commit advances the epoch.
1582pub struct MongrelSession {
1583    ctx: SessionContext,
1584    db: Option<TableHandle>,
1585    /// P4.1: the multi-table `Database` when opened via `open()`. When `Some`,
1586    /// the cache epoch is driven by `Database::visible_epoch()` instead of the
1587    /// legacy `combined_epoch()` fold.
1588    database: Option<Arc<Database>>,
1589    principal: Option<mongreldb_core::Principal>,
1590    principal_catalog_bound: bool,
1591    cache: ResultCache,
1592    /// Phase 16.5: logical-plan cache keyed by SQL string.
1593    plan_cache: parking_lot::Mutex<BoundedLru<String, datafusion::logical_expr::LogicalPlan>>,
1594    /// `table name → owning Table handle` for every registered table.
1595    tables: scored_sql::TableMap,
1596    /// Phase 17.3: named materialized views — `view name → defining SQL`.
1597    /// On `run("SELECT * FROM <view>")`, the defining SQL is executed (or the
1598    /// result-cache is hit). Invalidated automatically on commit (epoch bump).
1599    views: parking_lot::Mutex<HashMap<String, ViewDef>>,
1600    /// Databases attached via `ATTACH 'path' AS alias`, kept alive for the
1601    /// session's lifetime so their tables remain registered on the DataFusion
1602    /// context. Keyed by alias.
1603    attached_databases: parking_lot::Mutex<HashMap<String, Arc<Database>>>,
1604    /// SQL `BEGIN`/`COMMIT` staging for DML statements. Reads remain
1605    /// snapshot-at-scan; this batches SQL writes atomically when a client sends
1606    /// an explicit transaction block.
1607    sql_txn: parking_lot::Mutex<Option<Vec<commands::PendingSqlOp>>>,
1608    /// SAVEPOINT stack: `(name, staged-ops-length-at-savepoint)`. Truncated on
1609    /// `ROLLBACK TO name` and removed on `RELEASE name`.
1610    savepoints: parking_lot::Mutex<Vec<(String, usize)>>,
1611    /// Per-session state for SQL compatibility functions such as changes().
1612    sql_fn_state: Arc<extended_sql_functions::ExtendedSqlState>,
1613    /// Built-in plus app-provided external table modules available to this
1614    /// session.
1615    external_modules: Arc<ExternalModuleRegistry>,
1616    scored_runtime: Arc<scored_sql::ScoredRuntime>,
1617    /// Names of `PREPARE`d statements tracked so they can be `DEALLOCATE`d when
1618    /// DDL invalidates the tables they reference (DataFusion's prepared-plan
1619    /// store is not cleared by the session's own result/plan caches).
1620    prepared_names: parking_lot::Mutex<std::collections::HashSet<String>>,
1621}
1622
1623/// `(sql, snapshot_epoch) → cached result batches`.
1624type CacheKey = (String, u64);
1625type ResultCache = parking_lot::Mutex<BoundedLru<CacheKey, Arc<Vec<RecordBatch>>>>;
1626
1627const SESSION_CACHE_MAX_ENTRIES: usize = 1024;
1628
1629struct BoundedLru<K, V> {
1630    entries: HashMap<K, (V, u64)>,
1631    clock: u64,
1632    capacity: usize,
1633}
1634
1635impl<K: Clone + Eq + Hash, V> BoundedLru<K, V> {
1636    fn new(capacity: usize) -> Self {
1637        Self {
1638            entries: HashMap::new(),
1639            clock: 0,
1640            capacity: capacity.max(1),
1641        }
1642    }
1643
1644    fn get<Q>(&mut self, key: &Q) -> Option<&V>
1645    where
1646        K: Borrow<Q>,
1647        Q: Eq + Hash + ?Sized,
1648    {
1649        self.clock = self.clock.wrapping_add(1);
1650        let (value, last_used) = self.entries.get_mut(key)?;
1651        *last_used = self.clock;
1652        Some(value)
1653    }
1654
1655    fn insert(&mut self, key: K, value: V) {
1656        self.clock = self.clock.wrapping_add(1);
1657        if let Some(entry) = self.entries.get_mut(&key) {
1658            *entry = (value, self.clock);
1659            return;
1660        }
1661        if self.entries.len() >= self.capacity {
1662            // ponytail: scan 1024 entries on eviction; use a linked map if misses get hot.
1663            if let Some(oldest) = self
1664                .entries
1665                .iter()
1666                .min_by_key(|(_, (_, last_used))| *last_used)
1667                .map(|(key, _)| key.clone())
1668            {
1669                self.entries.remove(&oldest);
1670            }
1671        }
1672        self.entries.insert(key, (value, self.clock));
1673    }
1674
1675    fn clear(&mut self) {
1676        self.entries.clear();
1677        self.clock = 0;
1678    }
1679}
1680
1681fn buffered_stream(batches: Vec<RecordBatch>) -> MongrelRecordBatchStream {
1682    let schema = batches
1683        .first()
1684        .map(RecordBatch::schema)
1685        .unwrap_or_else(|| Arc::new(arrow::datatypes::Schema::empty()));
1686    let stream = futures::stream::iter(batches.into_iter().map(Ok::<RecordBatch, DataFusionError>));
1687    Box::pin(RecordBatchStreamAdapter::new(schema, stream))
1688}
1689
1690impl MongrelSession {
1691    /// Create a session over a live `Table`. Takes ownership; wrap in `Arc` if you
1692    /// need to keep a handle for writes after registering the provider. Registers
1693    /// the `ann_search` UDF so SQL semantic-search predicates parse.
1694    pub fn new(db: Table) -> Self {
1695        let db = Arc::new(Mutex::new(db));
1696        let handle = TableHandle::from(Arc::clone(&db));
1697        let ctx = SessionContext::new();
1698        let sql_fn_state = Arc::new(extended_sql_functions::ExtendedSqlState::default());
1699        register_mongrel_functions(&ctx, Arc::clone(&sql_fn_state));
1700        let tables = Arc::new(parking_lot::Mutex::new(HashMap::new()));
1701        let scored_runtime = scored_sql::ScoredRuntime::from_env();
1702        scored_sql::register(
1703            &ctx,
1704            Arc::clone(&tables),
1705            None,
1706            None,
1707            false,
1708            Arc::clone(&scored_runtime),
1709        );
1710        let external_modules = Arc::new(ExternalModuleRegistry::default());
1711        Self {
1712            ctx,
1713            db: Some(handle),
1714            database: None,
1715            principal: None,
1716            principal_catalog_bound: false,
1717            cache: parking_lot::Mutex::new(BoundedLru::new(SESSION_CACHE_MAX_ENTRIES)),
1718            plan_cache: parking_lot::Mutex::new(BoundedLru::new(SESSION_CACHE_MAX_ENTRIES)),
1719            tables,
1720            views: parking_lot::Mutex::new(HashMap::new()),
1721            attached_databases: parking_lot::Mutex::new(HashMap::new()),
1722            savepoints: parking_lot::Mutex::new(Vec::new()),
1723            sql_txn: parking_lot::Mutex::new(None),
1724            sql_fn_state,
1725            external_modules,
1726            scored_runtime,
1727            prepared_names: parking_lot::Mutex::new(std::collections::HashSet::new()),
1728        }
1729    }
1730
1731    pub fn new_with_external_modules(
1732        db: Table,
1733        modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
1734    ) -> Result<Self> {
1735        let session = Self::new(db);
1736        for module in modules {
1737            session.register_external_module(module)?;
1738        }
1739        Ok(session)
1740    }
1741
1742    /// Open a session over a multi-table [`Database`] (spec §12). Auto-registers
1743    /// every live table as a `MongrelProvider`; the cache epoch is driven by
1744    /// `Database::visible_epoch()` so any table's commit invalidates cached
1745    /// results.
1746    pub fn open(database: Arc<Database>) -> Result<Self> {
1747        Self::open_with_external_modules(database, std::iter::empty())
1748    }
1749
1750    pub fn open_as(database: Arc<Database>, principal: mongreldb_core::Principal) -> Result<Self> {
1751        Self::open_with_external_modules_as(database, std::iter::empty(), Some(principal))
1752    }
1753
1754    pub fn open_with_external_modules(
1755        database: Arc<Database>,
1756        modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
1757    ) -> Result<Self> {
1758        Self::open_with_external_modules_as(database, modules, None)
1759    }
1760
1761    pub fn open_with_external_modules_as(
1762        database: Arc<Database>,
1763        modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
1764        principal: Option<mongreldb_core::Principal>,
1765    ) -> Result<Self> {
1766        let ctx = SessionContext::new();
1767        let sql_fn_state = Arc::new(extended_sql_functions::ExtendedSqlState::default());
1768        register_mongrel_functions(&ctx, Arc::clone(&sql_fn_state));
1769        let external_modules = Arc::new(ExternalModuleRegistry::default());
1770        let principal_catalog_bound = principal
1771            .as_ref()
1772            .is_some_and(|principal| database.resolve_principal(&principal.username).is_some());
1773        for module in modules {
1774            external_modules.register(module)?;
1775        }
1776
1777        let mut tables: HashMap<String, TableHandle> = HashMap::new();
1778        for name in database.table_names() {
1779            let handle = database.table(&name)?;
1780            let provider = MongrelProvider::new_secured(
1781                handle.clone(),
1782                Arc::clone(&database),
1783                name.clone(),
1784                principal.clone(),
1785            )?;
1786            ctx.register_table(&name, Arc::new(provider))
1787                .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1788            tables.insert(name, handle);
1789        }
1790        for entry in database.external_tables() {
1791            let provider = external_modules.external_table_provider(&database, &entry)?;
1792            ctx.register_table(&entry.name, provider)
1793                .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1794        }
1795
1796        // Pick a stable "primary" (lexicographically smallest name) for legacy
1797        // `db()` accessors. If the database is empty, `db()` returns `None`.
1798        let primary = {
1799            let mut names: Vec<&String> = tables.keys().collect();
1800            names.sort();
1801            names.first().and_then(|n| tables.get(*n).cloned())
1802        };
1803        let tables = Arc::new(parking_lot::Mutex::new(tables));
1804        let scored_runtime = scored_sql::ScoredRuntime::from_env();
1805        scored_sql::register(
1806            &ctx,
1807            Arc::clone(&tables),
1808            Some(Arc::clone(&database)),
1809            principal.clone(),
1810            principal_catalog_bound,
1811            Arc::clone(&scored_runtime),
1812        );
1813
1814        Ok(Self {
1815            ctx,
1816            db: primary,
1817            database: Some(database),
1818            principal,
1819            principal_catalog_bound,
1820            cache: parking_lot::Mutex::new(BoundedLru::new(SESSION_CACHE_MAX_ENTRIES)),
1821            plan_cache: parking_lot::Mutex::new(BoundedLru::new(SESSION_CACHE_MAX_ENTRIES)),
1822            tables,
1823            views: parking_lot::Mutex::new(HashMap::new()),
1824            attached_databases: parking_lot::Mutex::new(HashMap::new()),
1825            savepoints: parking_lot::Mutex::new(Vec::new()),
1826            sql_txn: parking_lot::Mutex::new(None),
1827            sql_fn_state,
1828            external_modules,
1829            scored_runtime,
1830            prepared_names: parking_lot::Mutex::new(std::collections::HashSet::new()),
1831        })
1832    }
1833
1834    pub fn principal(&self) -> Option<mongreldb_core::Principal> {
1835        self.principal.as_ref().and_then(|principal| {
1836            if self.principal_catalog_bound {
1837                self.database
1838                    .as_ref()
1839                    .and_then(|database| database.resolve_principal(&principal.username))
1840            } else {
1841                Some(principal.clone())
1842            }
1843        })
1844    }
1845
1846    /// Set timeout, work, and fused-union limits for scored SQL functions.
1847    pub fn set_scored_execution_limits(
1848        &self,
1849        statement_timeout: std::time::Duration,
1850        max_work: usize,
1851        max_fused_candidates: usize,
1852    ) -> Result<()> {
1853        self.scored_runtime
1854            .configure(statement_timeout, max_work, max_fused_candidates)
1855            .map_err(Into::into)
1856    }
1857
1858    fn security_context_active(&self) -> bool {
1859        self.principal.is_some()
1860            || self.database.as_ref().is_some_and(|database| {
1861                database.principal_snapshot().is_some() || {
1862                    let security = database.security_catalog();
1863                    !security.rls_tables.is_empty()
1864                        || !security.policies.is_empty()
1865                        || !security.masks.is_empty()
1866                }
1867            })
1868    }
1869
1870    pub fn register_external_module(&self, module: Arc<dyn ExternalTableModule>) -> Result<()> {
1871        self.external_modules.register(module)?;
1872        self.clear_cache();
1873        Ok(())
1874    }
1875
1876    /// The underlying Table handle (Phase 19.3: used by the daemon for direct
1877    /// put/delete/commit/count access). Returns `None` when the session was
1878    /// opened over an empty `Database`.
1879    pub fn db(&self) -> Option<&TableHandle> {
1880        self.db.as_ref()
1881    }
1882
1883    /// Phase 17.3: create a named materialized view backed by a SQL query.
1884    /// `SELECT * FROM <name>` resolves to the view's defining SQL, which is
1885    /// executed (or served from the result cache) transparently. The view is
1886    /// automatically invalidated on commit (via the epoch-keyed result cache).
1887    pub fn create_view(&self, name: &str, sql: &str) {
1888        self.create_view_with_schema(name, sql, CoreSchema::default(), HashMap::new());
1889    }
1890
1891    pub(crate) fn create_view_with_schema(
1892        &self,
1893        name: &str,
1894        sql: &str,
1895        schema: CoreSchema,
1896        input_types: HashMap<u16, Option<TypeId>>,
1897    ) {
1898        self.views.lock().insert(
1899            name.to_string(),
1900            ViewDef {
1901                sql: sql.to_string(),
1902                schema,
1903                input_types,
1904            },
1905        );
1906    }
1907
1908    /// Drop a named materialized view.
1909    pub fn drop_view(&self, name: &str) {
1910        self.views.lock().remove(name);
1911    }
1912
1913    pub(crate) fn view_schema(&self, name: &str) -> Option<CoreSchema> {
1914        self.views.lock().get(name).map(|view| view.schema.clone())
1915    }
1916
1917    pub(crate) fn view_definition(&self, name: &str) -> Option<ViewDef> {
1918        self.views.lock().get(name).cloned()
1919    }
1920
1921    /// Register the table under `name` so `select * from <name>` resolves.
1922    pub async fn register(&self, name: &str) -> Result<()> {
1923        let db = self.db.clone().ok_or(MongrelQueryError::Core(
1924            mongreldb_core::MongrelError::NotFound("no primary table".into()),
1925        ))?;
1926        let provider = MongrelProvider::new_handle(db.clone())?;
1927        self.tables.lock().insert(name.to_string(), db);
1928        self.ctx
1929            .register_table(name, Arc::new(provider))
1930            .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1931        Ok(())
1932    }
1933
1934    /// Register a second (or further) live `Table` as another table on the same
1935    /// session, enabling cross-table SQL joins. The first `Table` (passed to
1936    /// [`Self::new`]) still owns the result-cache epoch: cached results are
1937    /// invalidated on its commits, so mutate the primary table last or call
1938    /// [`Self::clear_cache`] after writing a secondary table.
1939    pub async fn register_db(&self, name: &str, db: Table) -> Result<()> {
1940        let db_arc = Arc::new(Mutex::new(db));
1941        let provider = MongrelProvider::new(db_arc.clone())?;
1942        self.tables.lock().insert(name.to_string(), db_arc.into());
1943        self.ctx
1944            .register_table(name, Arc::new(provider))
1945            .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1946        Ok(())
1947    }
1948
1949    fn refresh_registered_table(&self, db: &Arc<Database>, name: &str) -> Result<()> {
1950        self.ctx
1951            .deregister_table(name)
1952            .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1953        let handle = db.table(name)?;
1954        let provider = MongrelProvider::new_secured(
1955            handle.clone(),
1956            Arc::clone(db),
1957            name.to_string(),
1958            self.principal.clone(),
1959        )?;
1960        self.ctx
1961            .register_table(name, Arc::new(provider))
1962            .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
1963        self.tables.lock().insert(name.to_string(), handle);
1964        Ok(())
1965    }
1966
1967    /// Run a SQL statement and return the result batches. Repeated identical SQL
1968    /// against the same snapshot returns the cached batches without re-executing.
1969    /// Run a SQL statement and return the result batches. DDL statements
1970    /// (`CREATE TABLE`, `DROP TABLE`, `ALTER TABLE`) are intercepted when a
1971    /// Intercept `SELECT ... FROM information_schema.tables` and return
1972    /// a synthesized batch listing tables, views, and triggers. Returns
1973    /// `None` if the SQL doesn't reference that name.
1974    fn try_catalog_introspection(&self, sql: &str) -> Result<Option<Vec<RecordBatch>>> {
1975        let lower = sql.to_ascii_lowercase();
1976        if !lower.contains("information_schema.tables") {
1977            return Ok(None);
1978        }
1979        use arrow::array::{ArrayRef, Int64Array, StringArray};
1980        use arrow::datatypes::{DataType, Field, Schema};
1981        use arrow::record_batch::RecordBatch;
1982
1983        let mut types: Vec<String> = Vec::new();
1984        let mut names: Vec<String> = Vec::new();
1985        let mut tbl_names: Vec<String> = Vec::new();
1986
1987        // Tables.
1988        let principal = self.principal();
1989        let can_see = |name: &str| match &self.database {
1990            Some(database) => database
1991                .select_column_ids_for(name, principal.as_ref())
1992                .is_ok(),
1993            None => true,
1994        };
1995        for name in self.tables.lock().keys().filter(|name| can_see(name)) {
1996            types.push("table".into());
1997            names.push(name.clone());
1998            tbl_names.push(name.clone());
1999        }
2000        // Views (session-scoped).
2001        for name in self.views.lock().keys() {
2002            types.push("view".into());
2003            names.push(name.clone());
2004            tbl_names.push(name.clone());
2005        }
2006        // Triggers (engine-side, if a Database is attached).
2007        if let Some(db) = &self.database {
2008            for t in db.triggers() {
2009                let target_name = match &t.target {
2010                    mongreldb_core::trigger::TriggerTarget::Table(n)
2011                    | mongreldb_core::trigger::TriggerTarget::View(n) => n.clone(),
2012                };
2013                if !can_see(&target_name) {
2014                    continue;
2015                }
2016                types.push("trigger".into());
2017                names.push(t.name.clone());
2018                tbl_names.push(target_name);
2019            }
2020        }
2021
2022        let schema = Arc::new(Schema::new(vec![
2023            Field::new("type", DataType::Utf8, false),
2024            Field::new("name", DataType::Utf8, false),
2025            Field::new("tbl_name", DataType::Utf8, false),
2026            Field::new("rootpage", DataType::Int64, false),
2027            Field::new("sql", DataType::Utf8, true),
2028        ]));
2029        let n = names.len();
2030        let rootpages: Vec<i64> = vec![0; n];
2031        let sqls: Vec<Option<&str>> = vec![None; n];
2032        let batch = RecordBatch::try_new(
2033            schema,
2034            vec![
2035                Arc::new(StringArray::from(types)) as ArrayRef,
2036                Arc::new(StringArray::from(names)) as ArrayRef,
2037                Arc::new(StringArray::from(tbl_names)) as ArrayRef,
2038                Arc::new(Int64Array::from(rootpages)) as ArrayRef,
2039                Arc::new(StringArray::from(sqls)) as ArrayRef,
2040            ],
2041        )
2042        .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
2043        Ok(Some(vec![batch]))
2044    }
2045
2046    /// §5.3 direct SQL dispatch: recognize a simple single-table `SELECT` from
2047    /// the raw SQL via the vendored `sqlparser` AST and serve it straight from
2048    /// the native column cursor, **bypassing DataFusion parse+plan+optimize**.
2049    /// Returns `Ok(None)` (→ fall through to `ctx.sql()`) for any shape it
2050    /// cannot serve *exactly*, or on any parse error.
2051    fn try_direct_dispatch(&self, sql: &str) -> Result<Option<Vec<RecordBatch>>> {
2052        if self.security_context_active() {
2053            return Ok(None);
2054        }
2055
2056        use arrow::array::ArrayRef;
2057        use mongreldb_core::Condition;
2058        use sqlparser::ast::{Expr, Query, SelectItem, SetExpr, Statement, TableFactor};
2059        use sqlparser::dialect::PostgreSqlDialect;
2060        use sqlparser::parser::Parser;
2061
2062        // Any parse error, or more than one statement → fall through.
2063        let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
2064            return Ok(None);
2065        };
2066        if stmts.len() != 1 {
2067            return Ok(None);
2068        }
2069        let Statement::Query(query) = stmts.into_iter().next().unwrap() else {
2070            return Ok(None);
2071        };
2072        let Query { body, .. } = *query;
2073        let select = match *body {
2074            SetExpr::Select(s) => *s,
2075            _ => return Ok(None),
2076        };
2077        // v1: fall through if LIMIT/OFFSET is present (can't read the fields
2078        // portably; a conservative token check keeps correctness safe).
2079        let lower_sql = sql.to_lowercase();
2080        if lower_sql.contains(" limit ") || lower_sql.contains(" offset ") {
2081            return Ok(None);
2082        }
2083        // Reject shapes we don't handle: DISTINCT / GROUP BY / HAVING / multi-FROM / joins.
2084        use sqlparser::ast::GroupByExpr;
2085        if select.distinct.is_some()
2086            || !matches!(&select.group_by, GroupByExpr::Expressions(e, _) if e.is_empty())
2087            || select.having.is_some()
2088            || select.from.len() != 1
2089            || !select.from[0].joins.is_empty()
2090        {
2091            return Ok(None);
2092        }
2093        let table_name = match &select.from[0].relation {
2094            TableFactor::Table { name, .. } => Some(name.to_string()),
2095            _ => return Ok(None),
2096        };
2097        let Some(table_name) = table_name else {
2098            return Ok(None);
2099        };
2100
2101        // v1 only dispatches FILTERED single-table SELECTs. An unfiltered `SELECT
2102        // *`/`SELECT cols` already streams efficiently through the scan path
2103        // (with ≤65 536-row batch chunking + Arrow shadow writes), which the
2104        // direct path's single-shot column decode can't preserve — so leave it
2105        // to DataFusion. The win here is the cold filtered-SELECT planning cost.
2106        if select.selection.is_none() {
2107            return Ok(None);
2108        }
2109
2110        // Projection: only `*` or a list of bare column identifiers.
2111        let mut proj_names: Option<Vec<String>> = None;
2112        for item in &select.projection {
2113            match item {
2114                SelectItem::Wildcard(_) => {}
2115                SelectItem::UnnamedExpr(Expr::Identifier(ident)) => {
2116                    proj_names
2117                        .get_or_insert_with(Vec::new)
2118                        .push(ident.value.clone());
2119                }
2120                SelectItem::UnnamedExpr(Expr::CompoundIdentifier(idents)) => {
2121                    if let Some(last) = idents.last() {
2122                        proj_names
2123                            .get_or_insert_with(Vec::new)
2124                            .push(last.value.clone());
2125                    }
2126                }
2127                _ => return Ok(None),
2128            }
2129        }
2130
2131        // Resolve the table handle.
2132        let handle = match self.tables.lock().get(&table_name).cloned() {
2133            Some(h) => h,
2134            None => return Ok(None),
2135        };
2136
2137        // SQL SELECT → require Select permission on the target table.
2138        if let Some(db) = &self.database {
2139            db.require_table(
2140                &table_name,
2141                mongreldb_core::auth_state::RequiredPermission::Select,
2142            )?;
2143        }
2144
2145        let mut db = handle.lock();
2146        let schema = db.schema().clone();
2147        // Translate WHERE against the live schema; an inexact/unsupported
2148        // predicate → fall through to DataFusion (which re-applies residuals).
2149        let conditions: Vec<Condition> = match &select.selection {
2150            Some(expr) => match translate_sqlparser_filter(expr, &schema) {
2151                Some(c) => c,
2152                None => return Ok(None),
2153            },
2154            None => Vec::new(),
2155        };
2156        if !conditions.is_empty() && db.ensure_indexes_complete().is_err() {
2157            return Ok(None);
2158        }
2159        let snap = db.snapshot();
2160
2161        // Resolve projected column ids + Arrow field list (in projection order).
2162        let mut col_ids: Vec<u16> = Vec::new();
2163        let mut fields: Vec<arrow::datatypes::Field> = Vec::new();
2164        let resolve_col = |name: &str| -> Option<&mongreldb_core::schema::ColumnDef> {
2165            schema.columns.iter().find(|c| c.name == name)
2166        };
2167        match &proj_names {
2168            None => {
2169                for c in &schema.columns {
2170                    col_ids.push(c.id);
2171                    fields.push(arrow::datatypes::Field::new(
2172                        &c.name,
2173                        arrow_conv::arrow_data_type(&c.ty)?,
2174                        c.flags.contains(mongreldb_core::ColumnFlags::NULLABLE),
2175                    ));
2176                }
2177            }
2178            Some(names) => {
2179                for n in names {
2180                    let cdef = match resolve_col(n) {
2181                        Some(c) => c,
2182                        None => return Ok(None), // unknown column → let DataFusion error
2183                    };
2184                    col_ids.push(cdef.id);
2185                    fields.push(arrow::datatypes::Field::new(
2186                        &cdef.name,
2187                        arrow_conv::arrow_data_type(&cdef.ty)?,
2188                        cdef.flags.contains(mongreldb_core::ColumnFlags::NULLABLE),
2189                    ));
2190                }
2191            }
2192        }
2193
2194        // Execute via the same native column path MongrelProvider::scan uses.
2195        let cols = if !conditions.is_empty() {
2196            match db.query_columns_native_cached(&conditions, Some(&col_ids), snap) {
2197                Ok(Some(c)) => c,
2198                Ok(None) => db
2199                    .visible_columns_native(snap, Some(&col_ids))
2200                    .map_err(MongrelQueryError::Core)?,
2201                Err(_) => return Ok(None),
2202            }
2203        } else {
2204            db.visible_columns_native(snap, Some(&col_ids))
2205                .map_err(MongrelQueryError::Core)?
2206        };
2207        drop(db);
2208
2209        // Order decoded columns into projection order, then build one batch.
2210        let mut arrays: Vec<ArrayRef> = Vec::with_capacity(col_ids.len());
2211        for cid in &col_ids {
2212            let col = cols
2213                .iter()
2214                .find(|(id, _)| id == cid)
2215                .map(|(_, c)| c.clone());
2216            let Some(col) = col else { return Ok(None) };
2217            let ty = schema
2218                .columns
2219                .iter()
2220                .find(|c| c.id == *cid)
2221                .map(|c| c.ty.clone())
2222                .unwrap_or(mongreldb_core::schema::TypeId::Int64);
2223            arrays.push(arrow_conv::native_to_array(ty, &col)?);
2224        }
2225        let batch_schema = Arc::new(arrow::datatypes::Schema::new(fields));
2226        let batch = RecordBatch::try_new(batch_schema, arrays)
2227            .map_err(|e| MongrelQueryError::Arrow(format!("direct dispatch batch build: {e}")))?;
2228
2229        mongreldb_core::trace::QueryTrace::record(|t| {
2230            t.scan_mode = mongreldb_core::trace::ScanMode::DirectDispatch;
2231            t.planning_nanos = 0; // we bypassed DataFusion planning
2232        });
2233        Ok(Some(vec![batch]))
2234    }
2235
2236    async fn dataframe(&self, sql: &str) -> Result<datafusion::dataframe::DataFrame> {
2237        // Prepared commands have their own DataFusion store. Caching EXECUTE
2238        // would create one entry per parameter set.
2239        let use_plan_cache = !is_prepared_stmt_sql(sql);
2240        if let Some(plan) = use_plan_cache
2241            .then(|| self.plan_cache.lock().get(sql).cloned())
2242            .flatten()
2243        {
2244            return Ok(datafusion::dataframe::DataFrame::new(
2245                self.ctx.state(),
2246                plan,
2247            ));
2248        }
2249
2250        let df = self
2251            .ctx
2252            .sql(sql)
2253            .await
2254            .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2255        if use_plan_cache {
2256            self.plan_cache
2257                .lock()
2258                .insert(sql.to_string(), df.logical_plan().clone());
2259        }
2260        Ok(df)
2261    }
2262
2263    fn prepare_as_of_query(&self, sql: &str) -> Result<Option<AsOfQuery>> {
2264        static AS_OF_RE: std::sync::OnceLock<regex::Regex> = std::sync::OnceLock::new();
2265        static NEXT_TEMP: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
2266        let re = AS_OF_RE.get_or_init(|| {
2267            regex::Regex::new(
2268                r#"(?i)\b(FROM|JOIN)\s+("[^"]+"|[A-Za-z_][A-Za-z0-9_]*)\s+AS\s+OF\s+EPOCH\s+([0-9]+)\b"#,
2269            )
2270            .expect("valid AS OF regex")
2271        });
2272        let mut captures = re.captures_iter(sql);
2273        let Some(capture) = captures.next() else {
2274            return Ok(None);
2275        };
2276        if captures.next().is_some() {
2277            return Err(MongrelQueryError::Schema(
2278                "AS OF EPOCH currently supports one historical table per query".into(),
2279            ));
2280        }
2281        let database = self.database.as_ref().ok_or_else(|| {
2282            MongrelQueryError::Schema("AS OF EPOCH requires a Database-backed session".into())
2283        })?;
2284        let whole = capture.get(0).expect("whole AS OF match");
2285        let keyword = capture.get(1).expect("FROM/JOIN capture").as_str();
2286        let table_token = capture.get(2).expect("table capture").as_str();
2287        let table_name = table_token.trim_matches('"');
2288        let epoch = capture
2289            .get(3)
2290            .expect("epoch capture")
2291            .as_str()
2292            .parse::<u64>()
2293            .map_err(|error| MongrelQueryError::Schema(format!("AS OF epoch: {error}")))?;
2294        let temp_id = NEXT_TEMP.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2295        let temp_name = format!("__mongrel_as_of_{temp_id}");
2296
2297        // Preserve an explicit alias after the epoch. Without one, alias the
2298        // temporary provider back to the original table name so qualifiers
2299        // such as `events.id` keep working.
2300        let mut replace_end = whole.end();
2301        let mut alias = table_token.to_string();
2302        let tail = &sql[replace_end..];
2303        let whitespace = tail.len() - tail.trim_start().len();
2304        let trimmed_tail = tail.trim_start();
2305        if trimmed_tail
2306            .get(..3)
2307            .is_some_and(|prefix| prefix.eq_ignore_ascii_case("as "))
2308        {
2309            let alias_text = &trimmed_tail[3..];
2310            let alias_len = alias_text
2311                .find(|character: char| {
2312                    character.is_ascii_whitespace() || character == ',' || character == ';'
2313                })
2314                .unwrap_or(alias_text.len());
2315            if alias_len == 0 {
2316                return Err(MongrelQueryError::Schema(
2317                    "AS OF EPOCH AS requires an alias".into(),
2318                ));
2319            }
2320            alias = alias_text[..alias_len].to_string();
2321            replace_end += whitespace + 3 + alias_len;
2322        } else {
2323            let alias_len = trimmed_tail
2324                .find(|character: char| {
2325                    character.is_ascii_whitespace() || character == ',' || character == ';'
2326                })
2327                .unwrap_or(trimmed_tail.len());
2328            let candidate = &trimmed_tail[..alias_len];
2329            let keyword = candidate.to_ascii_lowercase();
2330            let clause_keyword = matches!(
2331                keyword.as_str(),
2332                "where"
2333                    | "join"
2334                    | "left"
2335                    | "right"
2336                    | "inner"
2337                    | "outer"
2338                    | "cross"
2339                    | "full"
2340                    | "on"
2341                    | "group"
2342                    | "order"
2343                    | "having"
2344                    | "limit"
2345                    | "offset"
2346                    | "union"
2347                    | "intersect"
2348                    | "except"
2349            );
2350            let valid_alias = candidate
2351                .as_bytes()
2352                .first()
2353                .is_some_and(|byte| byte.is_ascii_alphabetic() || *byte == b'_' || *byte == b'"');
2354            if alias_len > 0 && valid_alias && !clause_keyword {
2355                alias = candidate.to_string();
2356                replace_end += whitespace + alias_len;
2357            }
2358        }
2359        let replacement = format!("{keyword} {temp_name} AS {alias}");
2360        let mut rewritten = String::with_capacity(sql.len() + replacement.len());
2361        rewritten.push_str(&sql[..whole.start()]);
2362        rewritten.push_str(&replacement);
2363        rewritten.push_str(&sql[replace_end..]);
2364
2365        let (snapshot, retention) = database.snapshot_at_owned(mongreldb_core::Epoch(epoch))?;
2366        let handle = database.table(table_name)?;
2367        let provider = MongrelProvider::new_historical(
2368            handle,
2369            snapshot,
2370            retention,
2371            Some(ProviderSecurity {
2372                database: Arc::clone(database),
2373                table: table_name.to_string(),
2374                principal: self.principal.clone(),
2375                principal_catalog_bound: self.principal_catalog_bound,
2376            }),
2377        )?;
2378        self.ctx
2379            .register_table(&temp_name, Arc::new(provider))
2380            .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))?;
2381        Ok(Some(AsOfQuery {
2382            sql: rewritten,
2383            registration: AsOfRegistration {
2384                ctx: self.ctx.clone(),
2385                table_name: temp_name,
2386            },
2387        }))
2388    }
2389
2390    async fn run_as_of(&self, query: AsOfQuery) -> Result<Vec<RecordBatch>> {
2391        let AsOfQuery { sql, registration } = query;
2392        let _registration = registration;
2393        let plan_start = std::time::Instant::now();
2394        let dataframe = self
2395            .ctx
2396            .sql(&sql)
2397            .await
2398            .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))?;
2399        mongreldb_core::trace::QueryTrace::record(|trace| {
2400            trace.planning_nanos = plan_start.elapsed().as_nanos() as u64;
2401        });
2402        dataframe
2403            .collect()
2404            .await
2405            .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))
2406    }
2407
2408    async fn run_as_of_stream(&self, query: AsOfQuery) -> Result<MongrelRecordBatchStream> {
2409        use futures::StreamExt;
2410
2411        let AsOfQuery { sql, registration } = query;
2412        let dataframe = self
2413            .ctx
2414            .sql(&sql)
2415            .await
2416            .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))?;
2417        let stream = dataframe
2418            .execute_stream()
2419            .await
2420            .map_err(|error| MongrelQueryError::DataFusion(error.to_string()))?;
2421        let schema = stream.schema();
2422        let guarded = futures::stream::unfold(
2423            (stream, registration),
2424            |(mut stream, registration)| async move {
2425                stream
2426                    .next()
2427                    .await
2428                    .map(|item| (item, (stream, registration)))
2429            },
2430        );
2431        Ok(Box::pin(RecordBatchStreamAdapter::new(schema, guarded)))
2432    }
2433
2434    /// Execute SQL as a DataFusion record-batch stream. Query results bypass
2435    /// the result cache and native batch-producing fast paths. Commands and
2436    /// MongrelDB's recursive-CTE compatibility path remain buffered because
2437    /// they do not produce a DataFusion stream.
2438    pub async fn run_stream(&self, sql: &str) -> Result<MongrelRecordBatchStream> {
2439        if strip_explain_query_plan(sql).is_some() {
2440            return self.run(sql).await.map(buffered_stream);
2441        }
2442
2443        let trimmed = sql.trim();
2444        if trimmed.contains(';') && !is_trigger_body(trimmed) {
2445            let stmts = split_sql_statements(trimmed);
2446            let non_empty: Vec<&str> = stmts
2447                .iter()
2448                .map(String::as_str)
2449                .filter(|stmt| !stmt.trim().is_empty() && stmt.trim() != ";")
2450                .collect();
2451            if non_empty.is_empty() {
2452                return Ok(buffered_stream(Vec::new()));
2453            }
2454            if stmts.len() > 1 {
2455                for stmt in &non_empty[..non_empty.len() - 1] {
2456                    self.run(stmt).await?;
2457                }
2458                return Box::pin(self.run_stream(non_empty[non_empty.len() - 1])).await;
2459            }
2460        }
2461
2462        if let Some(batches) = commands::try_run_command(self, sql).await? {
2463            return Ok(buffered_stream(batches));
2464        }
2465
2466        if let Some(query) = self.prepare_as_of_query(sql)? {
2467            return self.run_as_of_stream(query).await;
2468        }
2469
2470        let resolved = self.resolve_view_sql(sql);
2471        let resolved = self.rewrite_external_module_compat_sql(&resolved);
2472        let resolved = rewrite_compat_function_calls(&resolved);
2473        let effective_sql = normalize_sql(&resolved);
2474        let sql = effective_sql.as_str();
2475
2476        if let Some(batches) = self.try_catalog_introspection(sql)? {
2477            return Ok(buffered_stream(batches));
2478        }
2479
2480        let plan_start = std::time::Instant::now();
2481        let external_module_scan = self.query_references_external_module(sql);
2482        let df = self.dataframe(sql).await?;
2483        self.track_prepared_name(sql);
2484        mongreldb_core::trace::QueryTrace::record(|trace| {
2485            trace.planning_nanos = plan_start.elapsed().as_nanos() as u64;
2486        });
2487        let stream = df
2488            .execute_stream()
2489            .await
2490            .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2491        if external_module_scan {
2492            mongreldb_core::trace::QueryTrace::record(|trace| {
2493                trace.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
2494            });
2495        }
2496        Ok(stream)
2497    }
2498
2499    /// Run a SQL statement: DDL/commands are intercepted; otherwise a result
2500    /// cache keyed by `(normalized SQL, snapshot epoch)` memoizes batches.
2501    /// §5.3: simple single-table SELECTs are served by [`try_direct_dispatch`]
2502    /// (no DataFusion planning) before falling back to the full DataFusion path.
2503    pub async fn run(&self, sql: &str) -> Result<Vec<RecordBatch>> {
2504        if let Some(inner) = strip_explain_query_plan(sql) {
2505            return self.explain_query_plan(inner).await;
2506        }
2507
2508        // Multi-statement support: if the SQL contains semicolons, first try
2509        // to split into individual statements. This must happen BEFORE command
2510        // dispatch because `try_run_command` would fail on multi-statement SQL.
2511        // But CREATE TRIGGER bodies contain semicolons inside BEGIN...END, so
2512        // we only split if the first semicolon is NOT inside a BEGIN...END block.
2513        let trimmed = sql.trim();
2514        if trimmed.contains(';') && !is_trigger_body(trimmed) {
2515            let stmts = split_sql_statements(trimmed);
2516            // If all statements are empty (e.g. ";;;"), return empty result.
2517            let non_empty: Vec<&String> = stmts
2518                .iter()
2519                .filter(|s| !s.trim().is_empty() && s.trim() != ";")
2520                .collect();
2521            if non_empty.is_empty() {
2522                return Ok(Vec::new());
2523            }
2524            if stmts.len() > 1 {
2525                let mut last = Vec::new();
2526                for stmt in &stmts {
2527                    let stmt = stmt.trim();
2528                    if stmt.is_empty() {
2529                        continue;
2530                    }
2531                    last = Box::pin(self.run(stmt)).await?;
2532                }
2533                return Ok(last);
2534            }
2535        }
2536
2537        // Try command dispatch (handles DDL/DML/triggers/views/etc.).
2538        if let Some(batches) = commands::try_run_command(self, sql).await? {
2539            return Ok(batches);
2540        }
2541
2542        // Multi-statement support: if the SQL wasn't handled as a single
2543        // command and contains semicolons, split into individual statements
2544        // and execute each sequentially. This catches `SELECT 1; SELECT 2`
2545        // style multi-statement queries. Commands with semicolons in their
2546        // bodies (CREATE TRIGGER ... BEGIN ... END;) are handled above.
2547        let trimmed = sql.trim();
2548        if trimmed.contains(';') {
2549            let stmts = split_sql_statements(trimmed);
2550            if stmts.len() > 1 {
2551                let mut last = Vec::new();
2552                for stmt in &stmts {
2553                    let stmt = stmt.trim();
2554                    if stmt.is_empty() {
2555                        continue;
2556                    }
2557                    last = Box::pin(self.run(stmt)).await?;
2558                }
2559                return Ok(last);
2560            }
2561        }
2562        // P4.2: intercept DDL when a Database is attached.
2563        let lower = sql.trim_start().to_lowercase();
2564        if lower.starts_with("create table") {
2565            if let Some(db) = &self.database {
2566                db.require_for(self.principal().as_ref(), &mongreldb_core::Permission::Ddl)?;
2567                let (name, schema) = parse_create_table(sql)?;
2568                db.create_table(&name, schema)?;
2569                let handle = db.table(&name)?;
2570                let provider = MongrelProvider::new_secured(
2571                    handle.clone(),
2572                    Arc::clone(db),
2573                    name.clone(),
2574                    self.principal.clone(),
2575                )?;
2576                self.ctx
2577                    .register_table(&name, Arc::new(provider))
2578                    .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2579                self.tables.lock().insert(name, handle);
2580                self.invalidate_prepared_statements().await;
2581                self.clear_cache();
2582                return Ok(Vec::new());
2583            }
2584        }
2585        if lower.starts_with("drop table") {
2586            if let Some(db) = &self.database {
2587                db.require_for(self.principal().as_ref(), &mongreldb_core::Permission::Ddl)?;
2588                let (name, if_exists) = parse_drop_table(sql)?;
2589                let drop_result = db.drop_table(&name);
2590                if let Err(e) = drop_result {
2591                    // IF EXISTS tolerates NotFound.
2592                    let is_not_found = matches!(e, mongreldb_core::MongrelError::NotFound(_));
2593                    if !(if_exists && is_not_found) {
2594                        return Err(e.into());
2595                    }
2596                } else {
2597                    self.ctx
2598                        .deregister_table(&name)
2599                        .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2600                    self.tables.lock().remove(&name);
2601                }
2602                self.invalidate_prepared_statements().await;
2603                self.clear_cache();
2604                return Ok(Vec::new());
2605            }
2606        }
2607        if lower.starts_with("alter table") {
2608            if let Some(db) = &self.database {
2609                db.require_for(self.principal().as_ref(), &mongreldb_core::Permission::Ddl)?;
2610                match parse_alter_table(sql)? {
2611                    ParsedAlterTable::RenameTable { old_name, new_name } => {
2612                        db.rename_table(&old_name, &new_name)?;
2613                        // Re-key DataFusion + the session's handle cache under the new
2614                        // name. The table_id and underlying table object are unchanged
2615                        // by a rename, so a fresh handle resolves to the same table.
2616                        self.ctx
2617                            .deregister_table(&old_name)
2618                            .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2619                        self.tables.lock().remove(&old_name);
2620                        let handle = db.table(&new_name)?;
2621                        let provider = MongrelProvider::new_secured(
2622                            handle.clone(),
2623                            Arc::clone(db),
2624                            new_name.clone(),
2625                            self.principal.clone(),
2626                        )?;
2627                        self.ctx
2628                            .register_table(&new_name, Arc::new(provider))
2629                            .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2630                        self.tables.lock().insert(new_name, handle);
2631                    }
2632                    ParsedAlterTable::RenameColumn {
2633                        table_name,
2634                        column_name,
2635                        new_name,
2636                    } => {
2637                        db.alter_column(&table_name, &column_name, AlterColumn::rename(new_name))?;
2638                        self.refresh_registered_table(db, &table_name)?;
2639                    }
2640                    ParsedAlterTable::AlterColumnType {
2641                        table_name,
2642                        column_name,
2643                        ty,
2644                    } => {
2645                        db.alter_column(&table_name, &column_name, AlterColumn::set_type(ty))?;
2646                        self.refresh_registered_table(db, &table_name)?;
2647                    }
2648                    ParsedAlterTable::SetNotNull {
2649                        table_name,
2650                        column_name,
2651                    } => {
2652                        let flags = current_column_flags(db, &table_name, &column_name)?
2653                            .without(ColumnFlags::NULLABLE);
2654                        db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
2655                        self.refresh_registered_table(db, &table_name)?;
2656                    }
2657                    ParsedAlterTable::DropNotNull {
2658                        table_name,
2659                        column_name,
2660                    } => {
2661                        let flags = current_column_flags(db, &table_name, &column_name)?
2662                            .with(ColumnFlags::NULLABLE);
2663                        db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
2664                        self.refresh_registered_table(db, &table_name)?;
2665                    }
2666                }
2667                self.invalidate_prepared_statements().await;
2668                self.clear_cache();
2669                return Ok(Vec::new());
2670            }
2671        }
2672
2673        if let Some(query) = self.prepare_as_of_query(sql)? {
2674            return self.run_as_of(query).await;
2675        }
2676
2677        // Phase 17.3: intercept `SELECT ... FROM <view_name>` and rewrite to
2678        // the view's defining SQL.
2679        let resolved = self.resolve_view_sql(sql);
2680        let resolved = self.rewrite_external_module_compat_sql(&resolved);
2681        let resolved = rewrite_compat_function_calls(&resolved);
2682        // Canonicalize whitespace outside literals/comments so queries that
2683        // differ only in spacing share a cache key (and parse identically — SQL
2684        // is whitespace-insensitive between tokens).
2685        let effective_sql = normalize_sql(&resolved);
2686        let sql = effective_sql.as_str();
2687        // The cache key uses the Database's visible epoch (P4.1) when opened
2688        // via `open()`, or the legacy `combined_epoch()` fold for multi-table
2689        // sessions created via `new()` + `register_db()`.
2690        let epoch = self.cache_epoch();
2691        let key = (sql.to_string(), epoch);
2692        let has_ttl_table = self
2693            .tables
2694            .lock()
2695            .values()
2696            .any(|table| table.lock().ttl().is_some());
2697        let result_cacheable = !extended_sql_functions::contains_volatile_extended_function(sql)
2698            && !is_explain_analyze(sql)
2699            && !is_prepared_stmt_sql(sql)
2700            && !has_ttl_table;
2701        if result_cacheable {
2702            if let Some(hit) = self.cache.lock().get(&key) {
2703                return Ok((**hit).clone());
2704            }
2705        }
2706        // information_schema.tables: intercept catalog-introspection SELECTs
2707        // and synthesize a result batch.
2708        if let Some(batches) = self.try_catalog_introspection(sql)? {
2709            if result_cacheable {
2710                self.cache.lock().insert(key, Arc::new(batches.clone()));
2711            }
2712            return Ok(batches);
2713        }
2714        // §5.3: direct SQL dispatch for simple single-table SELECTs — bypasses
2715        // DataFusion parse+plan+optimize. Served batches are memoized into the
2716        // result cache like the normal path. Returns None (→ fall through) for
2717        // any shape it cannot serve exactly.
2718        if let Some(batches) = self.try_direct_dispatch(sql)? {
2719            if result_cacheable {
2720                self.cache.lock().insert(key, Arc::new(batches.clone()));
2721            }
2722            return Ok(batches);
2723        }
2724        // Phase 16.5: check the logical-plan cache before re-parsing.
2725        let plan_start = std::time::Instant::now();
2726        let external_module_scan = self.query_references_external_module(sql);
2727        let df = self.dataframe(sql).await?;
2728        // Track prepared-statement names so DDL can DEALLOCATE them (prevents a
2729        // prepared plan from querying a detached/old table after DROP+recreate).
2730        self.track_prepared_name(sql);
2731        // Priority 8: record logical-planning time (parse + plan; ~0 on a
2732        // plan-cache hit), separate from execution.
2733        let planning_nanos = plan_start.elapsed().as_nanos() as u64;
2734        mongreldb_core::trace::QueryTrace::record(|t| t.planning_nanos = planning_nanos);
2735
2736        // Phase 7.2/8.3 fast path: serve a simple single aggregate (SUM/MIN/MAX/
2737        // AVG/COUNT) over the primary table from the incremental aggregate
2738        // cache — warm cache ⇒ delta merge on commit; cold ⇒ vectorized scan.
2739        // Falls through to DataFusion for everything it cannot serve exactly.
2740        let agg_key = sql_cache_key(sql);
2741        let batches = match self.try_native_aggregate(df.logical_plan(), agg_key) {
2742            Ok(Some(batch)) => vec![batch],
2743            _ => {
2744                // Phase 8.1 fast path: serve a PK↔FK equi-join over two
2745                // registered tables via roaring-bitmap intersection, with no
2746                // hash-join materialization. Falls through otherwise.
2747                match self.try_fk_join(df.logical_plan()) {
2748                    Ok(Some(b)) => {
2749                        // Priority 13: the native FK-bitmap path served the join.
2750                        mongreldb_core::trace::QueryTrace::record(|t| {
2751                            t.join_mode = mongreldb_core::trace::JoinMode::FkBitmap;
2752                        });
2753                        b
2754                    }
2755                    _ => df
2756                        .collect()
2757                        .await
2758                        .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?,
2759                }
2760            }
2761        };
2762        if external_module_scan {
2763            mongreldb_core::trace::QueryTrace::record(|t| {
2764                t.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
2765            });
2766        }
2767        if result_cacheable {
2768            self.cache.lock().insert(key, Arc::new(batches.clone()));
2769        }
2770        Ok(batches)
2771    }
2772
2773    /// [`Self::run`] with a captured [`mongreldb_core::trace::QueryTrace`].
2774    ///
2775    /// Runs the SQL query inside a trace-capture scope so that path-decision
2776    /// recordings from both the SQL scan layer (`MongrelProvider::scan`) and
2777    /// the core engine (`Table::native_page_cursor`, `query_columns_native`,
2778    /// `count_conditions`, etc.) are collected into a single returned trace.
2779    ///
2780    /// The session-level result cache returns before `scan()` runs on a hit, so
2781    /// a session-cache hit yields `scan_mode = Unknown`. For scan-level
2782    /// result-cache tracing, use
2783    /// [`mongreldb_core::Table::query_columns_native_cached_traced`].
2784    pub async fn run_sql_traced(
2785        &self,
2786        sql: &str,
2787    ) -> Result<(Vec<RecordBatch>, mongreldb_core::trace::QueryTrace)> {
2788        mongreldb_core::trace::QueryTrace::push_scope();
2789        let result = self.run(sql).await;
2790        let trace = mongreldb_core::trace::QueryTrace::pop_scope();
2791        Ok((result?, trace))
2792    }
2793
2794    /// Drop all cached results (e.g. after a manual data change you want
2795    /// reflected immediately).
2796    pub fn clear_cache(&self) {
2797        self.cache.lock().clear();
2798        self.plan_cache.lock().clear();
2799    }
2800
2801    /// Record/remove a prepared-statement name from the tracking set. Called
2802    /// after the DataFusion execution path runs a `PREPARE`/`DEALLOCATE` so the
2803    /// session knows which names to invalidate when DDL changes the schema.
2804    fn track_prepared_name(&self, sql: &str) {
2805        let lower = sql.trim_start().to_ascii_lowercase();
2806        if let Some(rest) = lower.strip_prefix("prepare ") {
2807            if let Some(name) = parse_stmt_ident(rest) {
2808                self.prepared_names.lock().insert(name);
2809            }
2810        } else if let Some(rest) = lower.strip_prefix("deallocate ") {
2811            if let Some(name) = parse_stmt_ident(rest) {
2812                self.prepared_names.lock().remove(&name);
2813            }
2814        }
2815    }
2816
2817    /// `DEALLOCATE` every tracked prepared statement. Called on table DDL so a
2818    /// prepared plan cannot keep querying a dropped/altered/recreated table
2819    /// (DataFusion's prepared-plan store is otherwise untouched by `clear_cache`
2820    /// and would retain a plan bound to the old table provider). Errors from an
2821    /// unknown name are ignored (the plan was never stored, e.g. a failed
2822    /// PREPARE whose name was optimistically tracked).
2823    pub async fn invalidate_prepared_statements(&self) {
2824        let names: Vec<String> = self.prepared_names.lock().drain().collect();
2825        for name in names {
2826            // `DEALLOCATE <name>` — no params, safe (name was validated at
2827            // PREPARE time on the server path; here it came from our own
2828            // tracking so it is a bare identifier).
2829            let sql = format!("DEALLOCATE {name}");
2830            if self.ctx.sql(&sql).await.is_err() {
2831                // Unknown/stale name — ignore; the goal is just to drop it.
2832            }
2833        }
2834    }
2835
2836    async fn explain_query_plan(&self, sql: &str) -> Result<Vec<RecordBatch>> {
2837        let explain_sql = format!("EXPLAIN {}", sql.trim().trim_end_matches(';'));
2838        let batches = self
2839            .ctx
2840            .sql(&explain_sql)
2841            .await
2842            .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?
2843            .collect()
2844            .await
2845            .map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
2846        let mut detail = self.mongrel_query_plan_details(sql);
2847        for batch in &batches {
2848            if batch.num_columns() < 2 {
2849                continue;
2850            }
2851            let Some(plan_type) = batch.column(0).as_any().downcast_ref::<StringArray>() else {
2852                continue;
2853            };
2854            let Some(plan) = batch.column(1).as_any().downcast_ref::<StringArray>() else {
2855                continue;
2856            };
2857            for row in 0..batch.num_rows() {
2858                let prefix = plan_type.value(row);
2859                for line in plan.value(row).lines() {
2860                    let line = line.trim();
2861                    if !line.is_empty() {
2862                        detail.push(format!("DATAFUSION {prefix}: {line}"));
2863                    }
2864                }
2865            }
2866        }
2867        if detail.is_empty() {
2868            detail.push("plan unavailable".to_string());
2869        }
2870        let ids = (0..detail.len()).map(|i| i as i64).collect::<Vec<_>>();
2871        let parents = vec![0_i64; detail.len()];
2872        let notused = vec![0_i64; detail.len()];
2873        let schema = Arc::new(arrow::datatypes::Schema::new(vec![
2874            arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int64, false),
2875            arrow::datatypes::Field::new("parent", arrow::datatypes::DataType::Int64, false),
2876            arrow::datatypes::Field::new("notused", arrow::datatypes::DataType::Int64, false),
2877            arrow::datatypes::Field::new("detail", arrow::datatypes::DataType::Utf8, false),
2878        ]));
2879        let batch = RecordBatch::try_new(
2880            schema,
2881            vec![
2882                Arc::new(Int64Array::from(ids)) as ArrayRef,
2883                Arc::new(Int64Array::from(parents)),
2884                Arc::new(Int64Array::from(notused)),
2885                Arc::new(StringArray::from(detail)),
2886            ],
2887        )
2888        .map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
2889        Ok(vec![batch])
2890    }
2891
2892    fn mongrel_query_plan_details(&self, sql: &str) -> Vec<String> {
2893        use sqlparser::ast::{GroupByExpr, OrderByKind, SetExpr, Statement};
2894        use sqlparser::dialect::PostgreSqlDialect;
2895        use sqlparser::parser::Parser;
2896
2897        let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
2898            return Vec::new();
2899        };
2900        let Some(Statement::Query(query)) = stmts.first() else {
2901            return Vec::new();
2902        };
2903
2904        fn collect(session: &MongrelSession, query: &sqlparser::ast::Query, out: &mut Vec<String>) {
2905            use sqlparser::ast::{SetOperator, TableWithJoins};
2906            match query.body.as_ref() {
2907                SetExpr::Select(select) => {
2908                    for TableWithJoins { relation, joins } in &select.from {
2909                        session.push_table_plan(relation, select.selection.as_ref(), out);
2910                        for join in joins {
2911                            session.push_table_plan(&join.relation, None, out);
2912                        }
2913                    }
2914                    if select.distinct.is_some() {
2915                        out.push("USE TEMP B-TREE FOR DISTINCT".to_string());
2916                    }
2917                    let grouped = match &select.group_by {
2918                        GroupByExpr::All(_) => true,
2919                        GroupByExpr::Expressions(exprs, _) => !exprs.is_empty(),
2920                    };
2921                    if grouped {
2922                        out.push("USE TEMP B-TREE FOR GROUP BY".to_string());
2923                    }
2924                    let ordered = query.order_by.as_ref().is_some_and(|order_by| {
2925                        matches!(order_by.kind, OrderByKind::All(_))
2926                            || matches!(&order_by.kind, OrderByKind::Expressions(exprs) if !exprs.is_empty())
2927                    });
2928                    if ordered {
2929                        out.push("USE TEMP B-TREE FOR ORDER BY".to_string());
2930                    }
2931                }
2932                SetExpr::Query(query) => collect(session, query, out),
2933                SetExpr::SetOperation {
2934                    left, op, right, ..
2935                } => {
2936                    let label = match op {
2937                        SetOperator::Union => "COMPOUND QUERY UNION",
2938                        SetOperator::Except => "COMPOUND QUERY EXCEPT",
2939                        SetOperator::Intersect => "COMPOUND QUERY INTERSECT",
2940                        _ => "COMPOUND QUERY",
2941                    };
2942                    out.push(label.to_string());
2943                    collect_set_expr(session, left, out);
2944                    collect_set_expr(session, right, out);
2945                }
2946                _ => {}
2947            }
2948        }
2949
2950        fn collect_set_expr(session: &MongrelSession, expr: &SetExpr, out: &mut Vec<String>) {
2951            match expr {
2952                SetExpr::Select(select) => {
2953                    for table in &select.from {
2954                        session.push_table_plan(&table.relation, select.selection.as_ref(), out);
2955                    }
2956                }
2957                SetExpr::Query(query) => collect(session, query, out),
2958                SetExpr::SetOperation { left, right, .. } => {
2959                    collect_set_expr(session, left, out);
2960                    collect_set_expr(session, right, out);
2961                }
2962                _ => {}
2963            }
2964        }
2965
2966        let mut out = Vec::new();
2967        collect(self, query, &mut out);
2968        out
2969    }
2970
2971    fn push_table_plan(
2972        &self,
2973        relation: &sqlparser::ast::TableFactor,
2974        selection: Option<&sqlparser::ast::Expr>,
2975        out: &mut Vec<String>,
2976    ) {
2977        let sqlparser::ast::TableFactor::Table { name, alias, .. } = relation else {
2978            out.push("SCAN SUBQUERY".to_string());
2979            return;
2980        };
2981        let table_name = name.to_string();
2982        let display_name = alias
2983            .as_ref()
2984            .map(|alias| alias.name.value.clone())
2985            .unwrap_or_else(|| table_name.clone());
2986        let Some(handle) = self.tables.lock().get(&table_name).cloned() else {
2987            out.push(format!("SCAN {display_name}"));
2988            return;
2989        };
2990        let schema = handle.lock().schema().clone();
2991        let searchable = selection
2992            .and_then(|expr| translate_sqlparser_filter(expr, &schema))
2993            .is_some_and(|conditions| !conditions.is_empty());
2994        if searchable {
2995            out.push(format!("SEARCH {display_name} USING MONGREL INDEX"));
2996        } else {
2997            out.push(format!("SCAN {display_name}"));
2998        }
2999    }
3000
3001    /// A cache key epoch combining the primary table's epoch with every
3002    /// secondary table's, so any registered table's commit invalidates cached
3003    /// results (correctness for multi-table joins).
3004    /// Phase 17.3: rewrite `FROM <view_name>` to `FROM (<view_sql>) AS <view_name>`.
3005    fn resolve_view_sql(&self, sql: &str) -> String {
3006        let views = self.views.lock();
3007        if views.is_empty() {
3008            return sql.to_string();
3009        }
3010        let mut result = sql.to_string();
3011        for (name, view) in views.iter() {
3012            result = replace_from_view(&result, name, &view.sql);
3013        }
3014        result
3015    }
3016
3017    fn rewrite_external_module_compat_sql(&self, sql: &str) -> String {
3018        let Some(db) = &self.database else {
3019            return sql.to_string();
3020        };
3021        rewrite_fts_match_compat_sql(sql, db)
3022    }
3023
3024    fn query_references_external_module(&self, sql: &str) -> bool {
3025        use sqlparser::ast::Statement;
3026        use sqlparser::dialect::PostgreSqlDialect;
3027        use sqlparser::parser::Parser;
3028
3029        Parser::parse_sql(&PostgreSqlDialect {}, sql)
3030            .ok()
3031            .is_some_and(|statements| {
3032                statements.iter().any(|statement| match statement {
3033                    Statement::Query(query) => self.query_uses_external_module(query),
3034                    _ => false,
3035                })
3036            })
3037    }
3038
3039    fn query_uses_external_module(&self, query: &sqlparser::ast::Query) -> bool {
3040        self.set_expr_uses_external_module(query.body.as_ref())
3041    }
3042
3043    fn set_expr_uses_external_module(&self, expr: &sqlparser::ast::SetExpr) -> bool {
3044        use sqlparser::ast::SetExpr;
3045
3046        match expr {
3047            SetExpr::Select(select) => select
3048                .from
3049                .iter()
3050                .any(|table| self.table_with_joins_uses_external_module(table)),
3051            SetExpr::Query(query) => self.query_uses_external_module(query),
3052            SetExpr::SetOperation { left, right, .. } => {
3053                self.set_expr_uses_external_module(left)
3054                    || self.set_expr_uses_external_module(right)
3055            }
3056            _ => false,
3057        }
3058    }
3059
3060    fn table_with_joins_uses_external_module(
3061        &self,
3062        table: &sqlparser::ast::TableWithJoins,
3063    ) -> bool {
3064        self.table_factor_uses_external_module(&table.relation)
3065            || table
3066                .joins
3067                .iter()
3068                .any(|join| self.table_factor_uses_external_module(&join.relation))
3069    }
3070
3071    fn table_factor_uses_external_module(&self, relation: &sqlparser::ast::TableFactor) -> bool {
3072        use sqlparser::ast::{Expr, TableFactor};
3073
3074        match relation {
3075            TableFactor::Table { name, args, .. } => {
3076                let table_name = name.to_string();
3077                self.database
3078                    .as_ref()
3079                    .is_some_and(|db| db.external_table(&table_name).is_some())
3080                    || (args.is_some() && self.external_modules.contains(&table_name))
3081            }
3082            TableFactor::Function { name, .. } => self.external_modules.contains(&name.to_string()),
3083            TableFactor::TableFunction {
3084                expr: Expr::Function(func),
3085                ..
3086            } => self.external_modules.contains(&func.name.to_string()),
3087            TableFactor::Derived { subquery, .. } => self.query_uses_external_module(subquery),
3088            _ => false,
3089        }
3090    }
3091
3092    /// Cache epoch: uses `Database::visible_epoch()` when a Database is
3093    /// attached (P4.1), otherwise falls back to the legacy `combined_epoch()`.
3094    fn cache_epoch(&self) -> u64 {
3095        if let Some(db) = &self.database {
3096            db.visible_epoch().0
3097        } else {
3098            self.combined_epoch()
3099        }
3100    }
3101
3102    fn combined_epoch(&self) -> u64 {
3103        let primary = self.db.as_ref().expect("no primary table");
3104        let mut combined = primary.lock().snapshot().epoch.0;
3105        let tables = self.tables.lock();
3106        for arc in tables.values() {
3107            if !arc.ptr_eq(primary) {
3108                let e = arc.lock().snapshot().epoch.0;
3109                combined = combined.wrapping_mul(31).wrapping_add(e);
3110            }
3111        }
3112        combined
3113    }
3114
3115    /// Attempt the Phase 7.2/8.3 native aggregate fast path against `plan`.
3116    /// Returns `Ok(Some(batch))` when served natively, `Ok(None)` to fall
3117    /// through. `cache_key` ties the result to the incremental cache (Phase 8.3).
3118    fn try_native_aggregate(
3119        &self,
3120        plan: &datafusion::logical_expr::LogicalPlan,
3121        cache_key: u64,
3122    ) -> Result<Option<RecordBatch>> {
3123        if self.security_context_active() {
3124            return Ok(None);
3125        }
3126        let Some(primary) = self.db.as_ref() else {
3127            return Ok(None);
3128        };
3129        let mut db = primary.lock();
3130        let schema = db.schema().clone();
3131        let snap = db.snapshot();
3132        native_agg::try_native_aggregate(&mut db, &schema, snap, plan, cache_key)
3133    }
3134
3135    /// Attempt the Phase 8.1 FK-join (bitmap-intersection) fast path against
3136    /// `plan`. Returns `Ok(Some(batches))` when served natively, `Ok(None)` to
3137    /// fall through to DataFusion.
3138    fn try_fk_join(
3139        &self,
3140        plan: &datafusion::logical_expr::LogicalPlan,
3141    ) -> Result<Option<Vec<RecordBatch>>> {
3142        if self.security_context_active() {
3143            return Ok(None);
3144        }
3145        let tables = self.tables.lock();
3146        fk_join::try_fk_join(&tables, plan)
3147    }
3148
3149    pub fn context(&self) -> &SessionContext {
3150        &self.ctx
3151    }
3152
3153    /// Register a custom scalar SQL function on this session.
3154    ///
3155    /// This is the Rust escape hatch for application-defined SQL functions. The
3156    /// session's plan and result caches are cleared because function resolution
3157    /// can change query output without advancing the storage epoch.
3158    pub fn register_scalar_udf(&self, f: ScalarUDF) {
3159        self.ctx.register_udf(f);
3160        self.clear_cache();
3161    }
3162
3163    /// Register a custom aggregate SQL function on this session.
3164    pub fn register_aggregate_udf(&self, f: AggregateUDF) {
3165        self.ctx.register_udaf(f);
3166        self.clear_cache();
3167    }
3168
3169    /// Register a custom window SQL function on this session.
3170    pub fn register_window_udf(&self, f: WindowUDF) {
3171        self.ctx.register_udwf(f);
3172        self.clear_cache();
3173    }
3174}
3175
3176fn register_mongrel_functions(
3177    ctx: &SessionContext,
3178    sql_fn_state: Arc<extended_sql_functions::ExtendedSqlState>,
3179) {
3180    ctx.register_udf(ScalarUDF::from(udf::AnnSearchUdf::new()));
3181    ctx.register_udf(ScalarUDF::from(udf::SparseMatchUdf::new()));
3182    ctx.register_udf(ScalarUDF::from(udf::RTreeIntersectsUdf::new()));
3183    ctx.register_udf(ScalarUDF::from(udf::FtsRankUdf::new()));
3184    for udaf in percentile::percentile_udafs() {
3185        ctx.register_udaf(udaf);
3186    }
3187    extended_sql_functions::register_extended_sql_functions_with_state(ctx, sql_fn_state);
3188}
3189
3190/// Check if the SQL is a CREATE TRIGGER statement with a BEGIN...END body.
3191/// These contain semicolons inside the body that must NOT be split.
3192fn is_trigger_body(sql: &str) -> bool {
3193    let lower = sql.to_ascii_lowercase();
3194    if !lower.contains("create trigger") && !lower.contains("create or replace trigger") {
3195        return false;
3196    }
3197    lower.contains("begin") && lower.contains("end")
3198}
3199
3200/// Split a SQL string into individual statements on semicolons, respecting
3201/// single-quoted strings (`'...'`), double-quoted identifiers (`"..."`),
3202/// dollar-quoting (`$$...$$` or `$tag$...$tag$`), and line/block comments.
3203/// A trailing semicolon is not an empty statement.
3204fn split_sql_statements(sql: &str) -> Vec<String> {
3205    let b = sql.as_bytes();
3206    let n = b.len();
3207    let mut stmts = Vec::new();
3208    let mut current = String::new();
3209    let mut i = 0;
3210    while i < n {
3211        let c = b[i];
3212        // Line comment: skip to end of line.
3213        if c == b'-' && i + 1 < n && b[i + 1] == b'-' {
3214            while i < n && b[i] != b'\n' {
3215                current.push(c as char);
3216                i += 1;
3217            }
3218            continue;
3219        }
3220        // Block comment: skip to matching close.
3221        if c == b'/' && i + 1 < n && b[i + 1] == b'*' {
3222            current.push_str("/*");
3223            i += 2;
3224            let mut depth = 1;
3225            while i + 1 < n && depth > 0 {
3226                if b[i] == b'/' && b[i + 1] == b'*' {
3227                    depth += 1;
3228                    current.push_str("/*");
3229                    i += 2;
3230                } else if b[i] == b'*' && b[i + 1] == b'/' {
3231                    depth -= 1;
3232                    current.push_str("*/");
3233                    i += 2;
3234                } else {
3235                    current.push(b[i] as char);
3236                    i += 1;
3237                }
3238            }
3239            continue;
3240        }
3241        // Single-quoted string.
3242        if c == b'\'' {
3243            current.push('\'');
3244            i += 1;
3245            while i < n {
3246                current.push(b[i] as char);
3247                if b[i] == b'\'' {
3248                    // Check for doubled quote escape.
3249                    if i + 1 < n && b[i + 1] == b'\'' {
3250                        current.push('\'');
3251                        i += 2;
3252                        continue;
3253                    }
3254                    i += 1;
3255                    break;
3256                }
3257                i += 1;
3258            }
3259            continue;
3260        }
3261        // Double-quoted identifier.
3262        if c == b'"' {
3263            current.push('"');
3264            i += 1;
3265            while i < n {
3266                current.push(b[i] as char);
3267                if b[i] == b'"' {
3268                    if i + 1 < n && b[i + 1] == b'"' {
3269                        current.push('"');
3270                        i += 2;
3271                        continue;
3272                    }
3273                    i += 1;
3274                    break;
3275                }
3276                i += 1;
3277            }
3278            continue;
3279        }
3280        // Dollar-quoting: $tag$ ... $tag$
3281        if c == b'$' {
3282            // Try to read a dollar-quote opener.
3283            let start = i;
3284            let mut tag_end = i + 1;
3285            while tag_end < n && b[tag_end] != b'$' && b[tag_end].is_ascii_alphanumeric() {
3286                tag_end += 1;
3287            }
3288            if tag_end < n && b[tag_end] == b'$' {
3289                let tag = &sql[start..=tag_end];
3290                current.push_str(tag);
3291                i = tag_end + 1;
3292                // Find the closing tag.
3293                while i < n {
3294                    if b[i] == b'$' && sql[i..].starts_with(tag) {
3295                        current.push_str(tag);
3296                        i += tag.len();
3297                        break;
3298                    }
3299                    current.push(b[i] as char);
3300                    i += 1;
3301                }
3302                continue;
3303            }
3304        }
3305        // Semicolon — statement boundary.
3306        if c == b';' {
3307            current.push(';');
3308            let trimmed = current.trim();
3309            if !trimmed.is_empty() && trimmed != ";" {
3310                stmts.push(current.clone());
3311            }
3312            current.clear();
3313            i += 1;
3314            continue;
3315        }
3316        current.push(c as char);
3317        i += 1;
3318    }
3319    // Trailing content without a semicolon.
3320    let trimmed = current.trim();
3321    if !trimmed.is_empty() {
3322        stmts.push(current);
3323    }
3324    stmts
3325}
3326
3327fn strip_explain_query_plan(sql: &str) -> Option<&str> {
3328    let trimmed = sql.trim_start();
3329    let lower = trimmed.to_ascii_lowercase();
3330    if !lower.starts_with("explain") {
3331        return None;
3332    }
3333    let after_explain = trimmed.get(7..)?.trim_start();
3334    let after_explain_lower = after_explain.to_ascii_lowercase();
3335    if !after_explain_lower.starts_with("query") {
3336        return None;
3337    }
3338    let after_query = after_explain.get(5..)?.trim_start();
3339    let after_query_lower = after_query.to_ascii_lowercase();
3340    if !after_query_lower.starts_with("plan") {
3341        return None;
3342    }
3343    Some(after_query.get(4..)?.trim_start())
3344}
3345
3346/// Whether `sql` is a `PREPARE`/`EXECUTE`/`DEALLOCATE` statement. These must
3347/// bypass the session result + logical-plan caches: `EXECUTE name($p)` with
3348/// varying params would otherwise create one cache entry per parameter set
3349/// (unbounded growth), and the prepared plan itself is already held by the
3350/// DataFusion context.
3351fn is_prepared_stmt_sql(sql: &str) -> bool {
3352    let lower = sql.trim_start().to_ascii_lowercase();
3353    lower.starts_with("prepare ")
3354        || lower.starts_with("execute ")
3355        || lower.starts_with("deallocate ")
3356}
3357
3358/// Parse the first identifier from the text following a `PREPARE`/`DEALLOCATE`
3359/// keyword (e.g. `"gt AS SELECT ..."` → `"gt"`, `"p"` → `"p"`), stripping
3360/// surrounding quotes. Used to track prepared-statement names.
3361fn parse_stmt_ident(s: &str) -> Option<String> {
3362    let tok = s
3363        .trim_start()
3364        .split(|c: char| c.is_whitespace() || c == '(')
3365        .next()?;
3366    let t = tok.trim_matches(|c| c == '"' || c == '`' || c == '\'');
3367    (!t.is_empty()).then(|| t.to_string())
3368}
3369
3370/// Whether `sql` is an `EXPLAIN ANALYZE ...` statement. EXPLAIN ANALYZE falls
3371/// through to DataFusion 54 (which executes the plan and reports per-operator
3372/// timing), but its output must never be result-cached: the timing metrics are
3373/// request-specific and would be misleading on a cache hit.
3374fn is_explain_analyze(sql: &str) -> bool {
3375    let trimmed = sql.trim_start();
3376    let lower = trimmed.to_ascii_lowercase();
3377    if !lower.starts_with("explain") {
3378        return false;
3379    }
3380    let after = trimmed[7..].trim_start();
3381    after.to_ascii_lowercase().starts_with("analyze")
3382}
3383
3384#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3385enum SqlCompatTokenKind {
3386    Ident,
3387    String,
3388    Dot,
3389    LParen,
3390    RParen,
3391    Comma,
3392}
3393
3394#[derive(Debug, Clone)]
3395struct SqlCompatToken {
3396    kind: SqlCompatTokenKind,
3397    raw: String,
3398    normalized: String,
3399    start: usize,
3400    end: usize,
3401}
3402
3403#[derive(Debug, Clone)]
3404struct FtsMatchBinding {
3405    query_ref: String,
3406}
3407
3408#[derive(Debug, Clone)]
3409struct SqlReplacement {
3410    start: usize,
3411    end: usize,
3412    replacement: String,
3413}
3414
3415fn rewrite_fts_match_compat_sql(sql: &str, db: &Database) -> String {
3416    let tokens = sql_compat_tokens(sql);
3417    if tokens.is_empty() {
3418        return sql.to_string();
3419    }
3420    let bindings = fts_match_bindings(sql, db, &tokens);
3421    if bindings.is_empty() {
3422        return sql.to_string();
3423    }
3424    let unique_refs = bindings
3425        .values()
3426        .map(|binding| binding.query_ref.as_str())
3427        .collect::<HashSet<_>>();
3428    let unique_binding = if unique_refs.len() == 1 {
3429        bindings.values().next().cloned()
3430    } else {
3431        None
3432    };
3433    let mut replacements = Vec::new();
3434    for (idx, token) in tokens.iter().enumerate() {
3435        if token.kind != SqlCompatTokenKind::Ident || token.normalized != "match" {
3436            continue;
3437        }
3438        let Some(rhs) = tokens.get(idx + 1) else {
3439            continue;
3440        };
3441        if rhs.kind != SqlCompatTokenKind::String {
3442            continue;
3443        }
3444        let Some((lhs_start, _lhs_end, query_ref)) =
3445            fts_match_lhs_query_ref(&tokens, idx, &bindings, unique_binding.as_ref())
3446        else {
3447            continue;
3448        };
3449        replacements.push(SqlReplacement {
3450            start: lhs_start,
3451            end: rhs.end,
3452            replacement: format!("{query_ref}.query = {}", rhs.raw),
3453        });
3454    }
3455    apply_sql_replacements(sql, &replacements)
3456}
3457
3458fn fts_match_lhs_query_ref(
3459    tokens: &[SqlCompatToken],
3460    match_idx: usize,
3461    bindings: &HashMap<String, FtsMatchBinding>,
3462    unique_binding: Option<&FtsMatchBinding>,
3463) -> Option<(usize, usize, String)> {
3464    if match_idx == 0 {
3465        return None;
3466    }
3467    let lhs = tokens.get(match_idx - 1)?;
3468    if lhs.kind != SqlCompatTokenKind::Ident {
3469        return None;
3470    }
3471
3472    if match_idx >= 3
3473        && tokens.get(match_idx - 2)?.kind == SqlCompatTokenKind::Dot
3474        && tokens.get(match_idx - 3)?.kind == SqlCompatTokenKind::Ident
3475    {
3476        let owner = tokens.get(match_idx - 3)?;
3477        let binding = bindings.get(&owner.normalized)?;
3478        if lhs.normalized == "query" || lhs.normalized == "text" {
3479            return Some((owner.start, lhs.end, binding.query_ref.clone()));
3480        }
3481        return None;
3482    }
3483
3484    if let Some(binding) = bindings.get(&lhs.normalized) {
3485        return Some((lhs.start, lhs.end, binding.query_ref.clone()));
3486    }
3487    if lhs.normalized == "text" {
3488        let binding = unique_binding?;
3489        return Some((lhs.start, lhs.end, binding.query_ref.clone()));
3490    }
3491    None
3492}
3493
3494fn fts_match_bindings(
3495    sql: &str,
3496    db: &Database,
3497    tokens: &[SqlCompatToken],
3498) -> HashMap<String, FtsMatchBinding> {
3499    let mut out = HashMap::new();
3500    let mut i = 0;
3501    while i < tokens.len() {
3502        let token = &tokens[i];
3503        let starts_table_ref = token.kind == SqlCompatTokenKind::Ident
3504            && matches!(token.normalized.as_str(), "from" | "join");
3505        if !starts_table_ref {
3506            i += 1;
3507            continue;
3508        }
3509        let mut table_idx = i + 1;
3510        if tokens
3511            .get(table_idx)
3512            .is_some_and(|token| token.kind == SqlCompatTokenKind::LParen)
3513        {
3514            i += 1;
3515            continue;
3516        }
3517        let Some(table) = tokens.get(table_idx) else {
3518            break;
3519        };
3520        if table.kind != SqlCompatTokenKind::Ident {
3521            i += 1;
3522            continue;
3523        }
3524        let mut table_name = table.normalized.clone();
3525        let mut table_ref = table.raw.clone();
3526        if tokens
3527            .get(table_idx + 1)
3528            .is_some_and(|token| token.kind == SqlCompatTokenKind::Dot)
3529            && tokens
3530                .get(table_idx + 2)
3531                .is_some_and(|token| token.kind == SqlCompatTokenKind::Ident)
3532        {
3533            let qualified = tokens.get(table_idx + 2).unwrap();
3534            table_name = qualified.normalized.clone();
3535            table_ref = sql[table.start..qualified.end].to_string();
3536            table_idx += 2;
3537        }
3538        if !is_fts_docs_table(db, &table_name) {
3539            i = table_idx + 1;
3540            continue;
3541        }
3542        let mut query_ref = table_ref.clone();
3543        let mut alias_key = None;
3544        let mut next = table_idx + 1;
3545        if tokens.get(next).is_some_and(|token| {
3546            token.kind == SqlCompatTokenKind::Ident && token.normalized == "as"
3547        }) {
3548            next += 1;
3549        }
3550        if let Some(alias) = tokens.get(next) {
3551            if alias.kind == SqlCompatTokenKind::Ident && !is_table_ref_boundary(&alias.normalized)
3552            {
3553                alias_key = Some(alias.normalized.clone());
3554                query_ref = alias.raw.clone();
3555                next += 1;
3556            }
3557        }
3558        out.insert(
3559            table_name,
3560            FtsMatchBinding {
3561                query_ref: query_ref.clone(),
3562            },
3563        );
3564        if let Some(alias_key) = alias_key {
3565            out.insert(alias_key, FtsMatchBinding { query_ref });
3566        }
3567        i = next;
3568    }
3569    out
3570}
3571
3572fn is_fts_docs_table(db: &Database, name: &str) -> bool {
3573    db.external_table(name)
3574        .is_some_and(|entry| entry.module == "fts_docs")
3575}
3576
3577fn is_table_ref_boundary(normalized: &str) -> bool {
3578    matches!(
3579        normalized,
3580        "where"
3581            | "join"
3582            | "left"
3583            | "right"
3584            | "inner"
3585            | "outer"
3586            | "full"
3587            | "cross"
3588            | "on"
3589            | "using"
3590            | "group"
3591            | "order"
3592            | "having"
3593            | "limit"
3594            | "offset"
3595            | "union"
3596            | "except"
3597            | "intersect"
3598    )
3599}
3600
3601fn sql_compat_tokens(sql: &str) -> Vec<SqlCompatToken> {
3602    let bytes = sql.as_bytes();
3603    let mut tokens = Vec::new();
3604    let mut i = 0;
3605    while i < bytes.len() {
3606        match bytes[i] {
3607            b if b.is_ascii_whitespace() => i += 1,
3608            b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
3609                i += 2;
3610                while i < bytes.len() && bytes[i] != b'\n' {
3611                    i += 1;
3612                }
3613            }
3614            b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
3615                i = skip_block_comment(bytes, i);
3616            }
3617            b'\'' => {
3618                let end = skip_quoted(bytes, i, b'\'');
3619                tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
3620                i = end;
3621            }
3622            b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
3623                let end = skip_quoted(bytes, i + 1, b'\'');
3624                tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
3625                i = end;
3626            }
3627            b'$' => {
3628                let (end, matched) = skip_dollar_quoted(bytes, i);
3629                if matched {
3630                    tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
3631                    i = end;
3632                } else {
3633                    i += 1;
3634                }
3635            }
3636            b'"' => {
3637                let end = skip_quoted(bytes, i, b'"');
3638                let raw = sql[i..end].to_string();
3639                let normalized = unquote_sql_ident(&raw).to_ascii_lowercase();
3640                tokens.push(SqlCompatToken {
3641                    kind: SqlCompatTokenKind::Ident,
3642                    raw,
3643                    normalized,
3644                    start: i,
3645                    end,
3646                });
3647                i = end;
3648            }
3649            b'.' => {
3650                tokens.push(sql_token(sql, SqlCompatTokenKind::Dot, i, i + 1));
3651                i += 1;
3652            }
3653            b'(' => {
3654                tokens.push(sql_token(sql, SqlCompatTokenKind::LParen, i, i + 1));
3655                i += 1;
3656            }
3657            b')' => {
3658                tokens.push(sql_token(sql, SqlCompatTokenKind::RParen, i, i + 1));
3659                i += 1;
3660            }
3661            b',' => {
3662                tokens.push(sql_token(sql, SqlCompatTokenKind::Comma, i, i + 1));
3663                i += 1;
3664            }
3665            b if is_sql_ident_byte(b) => {
3666                let start = i;
3667                i += 1;
3668                while i < bytes.len() && is_sql_ident_byte(bytes[i]) {
3669                    i += 1;
3670                }
3671                tokens.push(sql_token(sql, SqlCompatTokenKind::Ident, start, i));
3672            }
3673            _ => i += 1,
3674        }
3675    }
3676    tokens
3677}
3678
3679fn sql_token(sql: &str, kind: SqlCompatTokenKind, start: usize, end: usize) -> SqlCompatToken {
3680    let raw = sql[start..end].to_string();
3681    SqlCompatToken {
3682        kind,
3683        normalized: raw.to_ascii_lowercase(),
3684        raw,
3685        start,
3686        end,
3687    }
3688}
3689
3690fn unquote_sql_ident(raw: &str) -> String {
3691    if raw.len() >= 2 && raw.starts_with('"') && raw.ends_with('"') {
3692        raw[1..raw.len() - 1].replace("\"\"", "\"")
3693    } else {
3694        raw.to_string()
3695    }
3696}
3697
3698fn apply_sql_replacements(sql: &str, replacements: &[SqlReplacement]) -> String {
3699    if replacements.is_empty() {
3700        return sql.to_string();
3701    }
3702    let mut ordered = replacements.to_vec();
3703    ordered.sort_by_key(|replacement| replacement.start);
3704    let mut out = String::with_capacity(sql.len());
3705    let mut cursor = 0;
3706    for replacement in ordered {
3707        if replacement.start < cursor || replacement.end > sql.len() {
3708            continue;
3709        }
3710        out.push_str(&sql[cursor..replacement.start]);
3711        out.push_str(&replacement.replacement);
3712        cursor = replacement.end;
3713    }
3714    out.push_str(&sql[cursor..]);
3715    out
3716}
3717
3718fn rewrite_compat_function_calls(sql: &str) -> String {
3719    let bytes = sql.as_bytes();
3720    let mut out = String::with_capacity(sql.len());
3721    let mut i = 0;
3722    while i < bytes.len() {
3723        match bytes[i] {
3724            b'\'' => i = copy_quoted_to_string(&mut out, bytes, i, b'\''),
3725            b'"' => i = copy_quoted_to_string(&mut out, bytes, i, b'"'),
3726            b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
3727                out.push(bytes[i] as char);
3728                i += 1;
3729                i = copy_quoted_to_string(&mut out, bytes, i, b'\'');
3730            }
3731            b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
3732                out.push('-');
3733                out.push('-');
3734                i += 2;
3735                while i < bytes.len() {
3736                    let ch = bytes[i] as char;
3737                    out.push(ch);
3738                    i += 1;
3739                    if ch == '\n' {
3740                        break;
3741                    }
3742                }
3743            }
3744            b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
3745                let start = i;
3746                i = skip_block_comment(bytes, i);
3747                out.push_str(&sql[start..i.min(bytes.len())]);
3748            }
3749            b'$' => {
3750                let start_len = out.len();
3751                let (next, matched) = copy_dollar_quoted_to_string(&mut out, bytes, i);
3752                if matched {
3753                    i = next;
3754                } else {
3755                    out.truncate(start_len);
3756                    out.push('$');
3757                    i += 1;
3758                }
3759            }
3760            b'g' | b'G' | b'm' | b'M' | b't' | b'T' => {
3761                if let Some((replacement, next)) = compat_function_rewrite_at(sql, i) {
3762                    out.push_str(&replacement);
3763                    i = next;
3764                } else {
3765                    out.push(bytes[i] as char);
3766                    i += 1;
3767                }
3768            }
3769            _ => {
3770                out.push(bytes[i] as char);
3771                i += 1;
3772            }
3773        }
3774    }
3775    out
3776}
3777
3778fn compat_function_rewrite_at(sql: &str, start: usize) -> Option<(String, usize)> {
3779    let bytes = sql.as_bytes();
3780    let (name, kind) = if ident_eq_at(bytes, start, b"max") {
3781        ("max", CompatRewriteKind::ScalarMax)
3782    } else if ident_eq_at(bytes, start, b"min") {
3783        ("min", CompatRewriteKind::ScalarMin)
3784    } else if ident_eq_at(bytes, start, b"group_concat") {
3785        ("group_concat", CompatRewriteKind::GroupConcat)
3786    } else if ident_eq_at(bytes, start, b"total") {
3787        ("total", CompatRewriteKind::Total)
3788    } else {
3789        return None;
3790    };
3791    let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
3792    let after_name = start + name.len();
3793    let after_ok = bytes
3794        .get(after_name)
3795        .is_some_and(|b| !is_sql_ident_byte(*b));
3796    if !before_ok || !after_ok {
3797        return None;
3798    }
3799    let mut open = after_name;
3800    while open < bytes.len() && bytes[open].is_ascii_whitespace() {
3801        open += 1;
3802    }
3803    if bytes.get(open) != Some(&b'(') {
3804        return None;
3805    }
3806    let summary = call_arg_summary(sql, open)?;
3807    match kind {
3808        CompatRewriteKind::ScalarMax if summary.top_level_commas > 0 => {
3809            Some(("__mongreldb_scalar_max(".to_string(), open + 1))
3810        }
3811        CompatRewriteKind::ScalarMin if summary.top_level_commas > 0 => {
3812            Some(("__mongreldb_scalar_min(".to_string(), open + 1))
3813        }
3814        CompatRewriteKind::GroupConcat => {
3815            let args = &sql[open + 1..summary.close];
3816            let rewritten = if summary.top_level_commas == 0 {
3817                format!("string_agg({args}, ',')")
3818            } else {
3819                format!("string_agg({args})")
3820            };
3821            Some((rewritten, summary.close + 1))
3822        }
3823        CompatRewriteKind::Total if summary.top_level_commas == 0 => {
3824            let args = &sql[open + 1..summary.close];
3825            let suffix_end = aggregate_suffix_end(sql, summary.close + 1);
3826            let suffix = &sql[summary.close + 1..suffix_end];
3827            Some((
3828                format!("coalesce(cast(sum({args}){suffix} as double), 0.0)"),
3829                suffix_end,
3830            ))
3831        }
3832        _ => None,
3833    }
3834}
3835
3836#[derive(Clone, Copy)]
3837enum CompatRewriteKind {
3838    ScalarMax,
3839    ScalarMin,
3840    GroupConcat,
3841    Total,
3842}
3843
3844fn ident_eq_at(bytes: &[u8], start: usize, ident: &[u8]) -> bool {
3845    bytes
3846        .get(start..start + ident.len())
3847        .is_some_and(|slice| slice.eq_ignore_ascii_case(ident))
3848}
3849
3850fn is_sql_ident_byte(b: u8) -> bool {
3851    b.is_ascii_alphanumeric() || b == b'_' || b == b'$'
3852}
3853
3854fn keyword_at(bytes: &[u8], start: usize, keyword: &[u8]) -> bool {
3855    if !ident_eq_at(bytes, start, keyword) {
3856        return false;
3857    }
3858    let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
3859    let after = start + keyword.len();
3860    let after_ok = after >= bytes.len() || !is_sql_ident_byte(bytes[after]);
3861    before_ok && after_ok
3862}
3863
3864fn skip_sql_whitespace(bytes: &[u8], mut i: usize) -> usize {
3865    while i < bytes.len() && bytes[i].is_ascii_whitespace() {
3866        i += 1;
3867    }
3868    i
3869}
3870
3871fn aggregate_suffix_end(sql: &str, start: usize) -> usize {
3872    let bytes = sql.as_bytes();
3873    let mut suffix_end = start;
3874    let mut i = skip_sql_whitespace(bytes, start);
3875
3876    if keyword_at(bytes, i, b"filter") {
3877        let open = skip_sql_whitespace(bytes, i + b"filter".len());
3878        if bytes.get(open) != Some(&b'(') {
3879            return start;
3880        }
3881        let Some(summary) = call_arg_summary(sql, open) else {
3882            return start;
3883        };
3884        suffix_end = summary.close + 1;
3885        i = skip_sql_whitespace(bytes, suffix_end);
3886    }
3887
3888    if keyword_at(bytes, i, b"over") {
3889        let after_over = skip_sql_whitespace(bytes, i + b"over".len());
3890        if bytes.get(after_over) == Some(&b'(') {
3891            let Some(summary) = call_arg_summary(sql, after_over) else {
3892                return suffix_end;
3893            };
3894            suffix_end = summary.close + 1;
3895        } else {
3896            let mut end = after_over;
3897            while end < bytes.len() && is_sql_ident_byte(bytes[end]) {
3898                end += 1;
3899            }
3900            if end > after_over {
3901                suffix_end = end;
3902            }
3903        }
3904    }
3905
3906    suffix_end
3907}
3908
3909struct CallArgSummary {
3910    close: usize,
3911    top_level_commas: usize,
3912}
3913
3914fn call_arg_summary(sql: &str, open: usize) -> Option<CallArgSummary> {
3915    let bytes = sql.as_bytes();
3916    let mut depth = 1;
3917    let mut i = open + 1;
3918    let mut top_level_commas = 0;
3919    while i < bytes.len() {
3920        match bytes[i] {
3921            b'\'' => i = skip_quoted(bytes, i, b'\''),
3922            b'"' => i = skip_quoted(bytes, i, b'"'),
3923            b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
3924                i = skip_quoted(bytes, i + 1, b'\'')
3925            }
3926            b'$' => {
3927                let (next, matched) = skip_dollar_quoted(bytes, i);
3928                i = if matched { next } else { i + 1 };
3929            }
3930            b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
3931                i += 2;
3932                while i < bytes.len() && bytes[i] != b'\n' {
3933                    i += 1;
3934                }
3935            }
3936            b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
3937                i = skip_block_comment(bytes, i);
3938            }
3939            b'(' => {
3940                depth += 1;
3941                i += 1;
3942            }
3943            b')' => {
3944                depth -= 1;
3945                if depth == 0 {
3946                    return Some(CallArgSummary {
3947                        close: i,
3948                        top_level_commas,
3949                    });
3950                }
3951                i += 1;
3952            }
3953            b',' if depth == 1 => {
3954                top_level_commas += 1;
3955                i += 1;
3956            }
3957            _ => i += 1,
3958        }
3959    }
3960    None
3961}
3962
3963fn copy_quoted_to_string(out: &mut String, bytes: &[u8], start: usize, delim: u8) -> usize {
3964    let end = skip_quoted(bytes, start, delim);
3965    out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
3966    end
3967}
3968
3969fn skip_quoted(bytes: &[u8], start: usize, delim: u8) -> usize {
3970    let mut i = start;
3971    if i < bytes.len() {
3972        i += 1;
3973    }
3974    while i < bytes.len() {
3975        if bytes[i] == delim {
3976            i += 1;
3977            if i < bytes.len() && bytes[i] == delim {
3978                i += 1;
3979                continue;
3980            }
3981            break;
3982        }
3983        i += 1;
3984    }
3985    i
3986}
3987
3988fn copy_dollar_quoted_to_string(out: &mut String, bytes: &[u8], start: usize) -> (usize, bool) {
3989    let (end, matched) = skip_dollar_quoted(bytes, start);
3990    if matched {
3991        out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
3992    }
3993    (end, matched)
3994}
3995
3996fn skip_dollar_quoted(bytes: &[u8], start: usize) -> (usize, bool) {
3997    if bytes.get(start) != Some(&b'$') {
3998        return (start, false);
3999    }
4000    let mut j = start + 1;
4001    while j < bytes.len() && (bytes[j].is_ascii_alphanumeric() || bytes[j] == b'_') {
4002        j += 1;
4003    }
4004    if bytes.get(j) != Some(&b'$') {
4005        return (start, false);
4006    }
4007    let tag = &bytes[start..=j];
4008    let mut i = j + 1;
4009    while i + tag.len() <= bytes.len() {
4010        if &bytes[i..i + tag.len()] == tag {
4011            return (i + tag.len(), true);
4012        }
4013        i += 1;
4014    }
4015    (start, false)
4016}
4017
4018fn skip_block_comment(bytes: &[u8], start: usize) -> usize {
4019    let mut i = start + 2;
4020    let mut depth = 1;
4021    while i + 1 < bytes.len() && depth > 0 {
4022        if bytes[i] == b'/' && bytes[i + 1] == b'*' {
4023            depth += 1;
4024            i += 2;
4025        } else if bytes[i] == b'*' && bytes[i + 1] == b'/' {
4026            depth -= 1;
4027            i += 2;
4028        } else {
4029            i += 1;
4030        }
4031    }
4032    i
4033}
4034
4035/// Stable 64-bit cache key for a SQL string (Phase 8.3 incremental cache).
4036fn sql_cache_key(sql: &str) -> u64 {
4037    use std::hash::{Hash, Hasher};
4038    let mut h = std::collections::hash_map::DefaultHasher::new();
4039    sql.hash(&mut h);
4040    h.finish()
4041}
4042
4043/// Replace the first whole-word `FROM <name>` reference (case-insensitive) in
4044/// `sql` with `FROM (<view_sql>) AS <name>`. Unlike a raw substring search this
4045/// requires a word boundary on both sides, so a view named `log` will **not**
4046/// rewrite `FROM logs` (the prior behavior matched the `from log` prefix and
4047/// left a dangling `s`). Original (non-lowercased) casing is preserved outside
4048/// the rewritten span.
4049fn replace_from_view(sql: &str, name: &str, view_sql: &str) -> String {
4050    let lower = sql.to_ascii_lowercase();
4051    let bytes = lower.as_bytes();
4052    let name_b = name.as_bytes();
4053    let mut i = 0usize;
4054    while let Some(rel) = lower[i..].find("from") {
4055        let from_start = i + rel;
4056        let after_from = from_start + 4;
4057        i = after_from;
4058        // Left boundary: "from" must not be a suffix of a longer identifier.
4059        if from_start > 0 && is_ident_byte(bytes[from_start - 1]) {
4060            continue;
4061        }
4062        // Must be followed by whitespace then the name.
4063        let mut j = after_from;
4064        while j < bytes.len() && bytes[j].is_ascii_whitespace() {
4065            j += 1;
4066        }
4067        if j == after_from || !bytes[j..].starts_with(name_b) {
4068            continue;
4069        }
4070        let after_name = j + name_b.len();
4071        // Right boundary: the name must not be a prefix of a longer identifier.
4072        if after_name < bytes.len() && is_ident_byte(bytes[after_name]) {
4073            continue;
4074        }
4075        // Preserve the original `FROM ` casing/whitespace (sql[from_start..j]),
4076        // then wrap the view body as a subquery aliased back to the view name.
4077        let mut out = String::with_capacity(sql.len() + view_sql.len() + name.len() + 8);
4078        out.push_str(&sql[..from_start]);
4079        out.push_str(&sql[from_start..j]);
4080        out.push('(');
4081        out.push_str(view_sql);
4082        out.push_str(") AS ");
4083        out.push_str(name);
4084        out.push_str(&sql[after_name..]);
4085        return out;
4086    }
4087    sql.to_string()
4088}
4089
4090fn is_ident_byte(b: u8) -> bool {
4091    b.is_ascii_alphanumeric() || b == b'_'
4092}
4093
4094/// Canonicalize a SQL string for caching/parsing: collapse runs of ASCII
4095/// whitespace outside of literals/comments to a single space and trim. String
4096/// literals (`'...'`, with `''` escapes), quoted identifiers (`"..."`), escape
4097/// strings (`E'...'`), line comments (`--`), block comments (`/* */`), and
4098/// dollar-quoting (`$tag$...$tag$`) are passed through verbatim so their
4099/// internal whitespace (which IS semantically significant) is never altered.
4100/// SQL parsing is whitespace-insensitive outside literals, so the normalized
4101/// form parses identically while making `SELECT  *  FROM t`, `SELECT * FROM t`,
4102/// and `\n  SELECT * FROM t  \n` share one cache key.
4103fn normalize_sql(sql: &str) -> String {
4104    let b = sql.as_bytes();
4105    let n = b.len();
4106    let mut out: Vec<u8> = Vec::with_capacity(n);
4107    // Whether a single separating space should precede the next emitted token
4108    // (i.e. we're between tokens, not at the very start of the output).
4109    let mut want_space = false;
4110    let mut i = 0usize;
4111    while i < n {
4112        let c = b[i];
4113        // Whitespace and comments both act only as token separators — they set
4114        // the pending-space flag but never emit a byte themselves, so a run of
4115        // "1  -- c\nFROM" collapses to a single separating space.
4116        if c.is_ascii_whitespace() {
4117            want_space = true;
4118            i += 1;
4119            continue;
4120        }
4121        if c == b'-' && i + 1 < n && b[i + 1] == b'-' {
4122            // Line comment: skip to end of line.
4123            i += 2;
4124            while i < n && b[i] != b'\n' {
4125                i += 1;
4126            }
4127            want_space = !out.is_empty();
4128            continue;
4129        }
4130        if c == b'/' && i + 1 < n && b[i + 1] == b'*' {
4131            // Block comment: skip to the matching close `*/`, honoring nesting
4132            // (Postgres/DataFusion allow `/* /* */ */`).
4133            i += 2;
4134            let mut depth = 1usize;
4135            while i + 1 < n && depth > 0 {
4136                if b[i] == b'/' && b[i + 1] == b'*' {
4137                    depth += 1;
4138                    i += 2;
4139                } else if b[i] == b'*' && b[i + 1] == b'/' {
4140                    depth -= 1;
4141                    i += 2;
4142                } else {
4143                    i += 1;
4144                }
4145            }
4146            want_space = !out.is_empty();
4147            continue;
4148        }
4149        // A real token byte (or a literal/quote opener) — emit the separator.
4150        if want_space && !out.is_empty() {
4151            out.push(b' ');
4152        }
4153        want_space = false;
4154        match c {
4155            // Escape string E'...' (backslash escapes; '' is still an escape).
4156            b'E' | b'e' if i + 1 < n && b[i + 1] == b'\'' => {
4157                out.push(c);
4158                i += 1;
4159                i = copy_quoted(&mut out, b, i, n, b'\'');
4160                continue;
4161            }
4162            // Single-quoted string literal ('...' with '' escape).
4163            b'\'' => {
4164                i = copy_quoted(&mut out, b, i, n, b'\'');
4165                continue;
4166            }
4167            // Double-quoted identifier ("..." with "" escape).
4168            b'"' => {
4169                i = copy_quoted(&mut out, b, i, n, b'"');
4170                continue;
4171            }
4172            // Dollar-quoting: $tag$ ... $tag$ (tag optional/empty).
4173            b'$' => {
4174                let (consumed, matched) = copy_dollar_quoted(&mut out, b, i, n);
4175                if matched {
4176                    i = consumed;
4177                    continue;
4178                }
4179                out.push(c);
4180                i += 1;
4181                continue;
4182            }
4183            _ => {
4184                out.push(c);
4185                i += 1;
4186            }
4187        }
4188    }
4189    String::from_utf8(out).unwrap_or_else(|_| sql.to_string())
4190}
4191
4192/// Copy a quote-delimited span starting at `start` (the opening quote byte is
4193/// `delim`), including the opening and closing delimiters and any doubled
4194/// escapes, verbatim into `out`. Returns the index past the closing quote.
4195fn copy_quoted(out: &mut Vec<u8>, b: &[u8], start: usize, n: usize, delim: u8) -> usize {
4196    out.push(b[start]);
4197    let mut i = start + 1;
4198    while i < n {
4199        let c = b[i];
4200        out.push(c);
4201        if c == delim {
4202            // Doubled delimiter (e.g. '' or "") is an escape, not the end.
4203            if i + 1 < n && b[i + 1] == delim {
4204                out.push(b[i + 1]);
4205                i += 2;
4206                continue;
4207            }
4208            return i + 1;
4209        }
4210        i += 1;
4211    }
4212    i
4213}
4214
4215/// Copy a dollar-quoted span starting at the opening `$`. Returns
4216/// `(index_past_close, true)` if a matching close delimiter was found, or
4217/// `(start + 1, false)` if this `$` does not open a dollar-quote.
4218fn copy_dollar_quoted(out: &mut Vec<u8>, b: &[u8], start: usize, n: usize) -> (usize, bool) {
4219    // Parse the opening delimiter: '$' [tag] '$'. An empty tag ($$..$$) is
4220    // allowed; a non-empty tag must be identifier bytes starting with a
4221    // letter/underscore.
4222    let mut j = start + 1;
4223    let tag_start = j;
4224    while j < n && b[j] != b'$' && is_dollar_tag_byte(b[j]) {
4225        j += 1;
4226    }
4227    if j >= n || b[j] != b'$' {
4228        return (start + 1, false);
4229    }
4230    if tag_start < j && !(b[tag_start].is_ascii_alphabetic() || b[tag_start] == b'_') {
4231        return (start + 1, false);
4232    }
4233    let close_end = j + 1; // index just past the opening '$'
4234    let delim = &b[start..close_end];
4235    // Copy the opening delimiter verbatim.
4236    out.extend_from_slice(delim);
4237    // Find the matching close delimiter.
4238    let mut k = close_end;
4239    while k + delim.len() <= n {
4240        if &b[k..k + delim.len()] == delim {
4241            out.extend_from_slice(delim);
4242            return (k + delim.len(), true);
4243        }
4244        out.push(b[k]);
4245        k += 1;
4246    }
4247    // Unterminated: copy the remainder verbatim (don't corrupt).
4248    out.extend_from_slice(&b[close_end..n]);
4249    (n, true)
4250}
4251
4252fn is_dollar_tag_byte(b: u8) -> bool {
4253    b.is_ascii_alphanumeric() || b == b'_'
4254}
4255
4256/// Strip an ASCII case-insensitive prefix from `s`, returning the remainder.
4257fn strip_prefix_ci<'a>(s: &'a str, prefix: &str) -> Option<&'a str> {
4258    let bytes = s.as_bytes();
4259    let pb = prefix.as_bytes();
4260    if bytes.len() >= pb.len() && bytes[..pb.len()].eq_ignore_ascii_case(pb) {
4261        Some(&s[pb.len()..])
4262    } else {
4263        None
4264    }
4265}
4266
4267/// Recognized column constraints in `CREATE TABLE` column definitions. Each
4268/// entry maps a SQL phrase (matched case-insensitively as a substring of the
4269/// whitespace-normalized constraint clause) to the [`ColumnFlags`] bit it sets.
4270///
4271/// Multi-word phrases such as `"primary key"` match regardless of internal
4272/// spacing because the clause is normalized to single spaces before matching.
4273///
4274/// **Adding a new column constraint is a one-line change:** append `(phrase,
4275/// flag)` here. This keeps the DDL shim's grammar in one place rather than
4276/// scattering `contains(...)` checks across the parser. (A full SQL grammar is
4277/// deliberately out of scope — only the DDL shapes handled below are
4278/// intercepted here; all query parsing is delegated to DataFusion.)
4279const COLUMN_CONSTRAINTS: &[(&str, u32)] = &[
4280    ("primary key", ColumnFlags::PRIMARY_KEY),
4281    // Both spellings are accepted: `AUTOINCREMENT` (SQLite) and `AUTO_INCREMENT`
4282    // (MySQL). The engine enforces that the flag is valid only on a single
4283    // non-nullable `Int64` primary key (see `Schema::validate_auto_increment`),
4284    // so recognizing the keyword on any column here is safe — invalid
4285    // placements are rejected at table-creation time, before the schema is
4286    // durably logged.
4287    ("autoincrement", ColumnFlags::AUTO_INCREMENT),
4288    ("auto_increment", ColumnFlags::AUTO_INCREMENT),
4289];
4290
4291/// Translate a column's constraint clause (the text following `<name> <type>`
4292/// in a `CREATE TABLE` column definition) into [`ColumnFlags`]. The clause is
4293/// lowercased and its internal whitespace collapsed to single spaces so
4294/// multi-word phrases match regardless of formatting. See
4295/// [`COLUMN_CONSTRAINTS`] for the recognized phrases; add new ones there.
4296fn parse_column_constraints(constraint_text: &str) -> ColumnFlags {
4297    let normalized = constraint_text.to_lowercase();
4298    let mut flags = ColumnFlags::empty();
4299    for (phrase, bit) in COLUMN_CONSTRAINTS {
4300        if normalized.contains(phrase) {
4301            flags = flags.with(*bit);
4302        }
4303    }
4304    flags
4305}
4306
4307fn parse_sql_type(ty_str: &str) -> Result<mongreldb_core::schema::TypeId> {
4308    use mongreldb_core::schema::TypeId;
4309
4310    match ty_str.trim().trim_end_matches(';').to_lowercase().as_str() {
4311        "bigint" | "int8" | "int64" | "integer" | "int" => Ok(TypeId::Int64),
4312        "double" | "float8" | "float64" | "real" | "float" => Ok(TypeId::Float64),
4313        "varchar" | "text" | "string" | "bytes" => Ok(TypeId::Bytes),
4314        "boolean" | "bool" => Ok(TypeId::Bool),
4315        other => Err(MongrelQueryError::Schema(format!(
4316            "unsupported column type: {other}"
4317        ))),
4318    }
4319}
4320
4321/// Parse `CREATE TABLE [IF NOT EXISTS] <name> (<col> <type> <constraints>, ...)`
4322/// into a MongrelDB table name + schema. Supports BIGINT/INTEGER/INT, DOUBLE,
4323/// VARCHAR/TEXT, BOOLEAN. Recognized column constraints (`PRIMARY KEY`,
4324/// `AUTOINCREMENT` / `AUTO_INCREMENT`) are listed in [`COLUMN_CONSTRAINTS`].
4325/// Table name may be double-quoted.
4326fn parse_create_table(sql: &str) -> Result<(String, mongreldb_core::schema::Schema)> {
4327    use mongreldb_core::schema::*;
4328
4329    let open = sql
4330        .find('(')
4331        .ok_or(MongrelQueryError::Schema("CREATE TABLE missing '('".into()))?;
4332    let close = sql
4333        .rfind(')')
4334        .ok_or(MongrelQueryError::Schema("CREATE TABLE missing ')'".into()))?;
4335    let head = sql[..open].trim();
4336    let after_kw = strip_prefix_ci(head, "CREATE TABLE")
4337        .or_else(|| strip_prefix_ci(head, "create table"))
4338        .unwrap_or("")
4339        .trim();
4340    // Skip optional `IF NOT EXISTS`.
4341    let after_kw = after_kw
4342        .strip_prefix("IF NOT EXISTS")
4343        .or_else(|| after_kw.strip_prefix("if not exists"))
4344        .map(str::trim)
4345        .unwrap_or(after_kw);
4346    let name = after_kw.trim_matches('"').to_string();
4347    if name.is_empty() {
4348        return Err(MongrelQueryError::Schema(
4349            "CREATE TABLE missing table name".into(),
4350        ));
4351    }
4352
4353    let body = &sql[open + 1..close];
4354    let mut columns = Vec::new();
4355    let schema_id: u64 = 0; // Database::create_table overrides with the table_id.
4356    for (i, raw) in body.split(',').enumerate() {
4357        let part = raw.trim();
4358        if part.is_empty() {
4359            continue;
4360        }
4361        let mut tokens = part.split_whitespace();
4362        let col_name = tokens
4363            .next()
4364            .ok_or(MongrelQueryError::Schema("missing column name".into()))?
4365            .trim_matches('"');
4366        let ty_str = tokens
4367            .next()
4368            .ok_or(MongrelQueryError::Schema("missing column type".into()))?
4369            .to_lowercase();
4370        let ty = parse_sql_type(&ty_str)?;
4371        // Everything after `<name> <type>` is the column's constraint clause
4372        // (e.g. `PRIMARY KEY`, `PRIMARY KEY AUTOINCREMENT`). The remaining
4373        // tokens are matched against `COLUMN_CONSTRAINTS`.
4374        let constraint_clause: String = tokens.collect::<Vec<_>>().join(" ");
4375        let flags = parse_column_constraints(&constraint_clause);
4376        columns.push(ColumnDef {
4377            id: (i + 1) as u16,
4378            name: col_name.to_string(),
4379            ty,
4380            flags,
4381            default_value: None,
4382        });
4383    }
4384
4385    Ok((
4386        name,
4387        Schema {
4388            schema_id,
4389            columns,
4390            indexes: vec![],
4391            colocation: vec![],
4392            constraints: Default::default(),
4393            clustered: false,
4394        },
4395    ))
4396}
4397
4398/// Parse `DROP TABLE [IF EXISTS] <name>`. Returns `(name, if_exists)`.
4399fn parse_drop_table(sql: &str) -> Result<(String, bool)> {
4400    let head = sql.trim();
4401    let after_kw = strip_prefix_ci(head, "DROP TABLE")
4402        .or_else(|| strip_prefix_ci(head, "drop table"))
4403        .unwrap_or("")
4404        .trim();
4405    // Detect optional `IF EXISTS`.
4406    let (rest, if_exists) = if let Some(r) = after_kw
4407        .strip_prefix("IF EXISTS")
4408        .or_else(|| after_kw.strip_prefix("if exists"))
4409        .map(str::trim)
4410    {
4411        (r, true)
4412    } else {
4413        (after_kw, false)
4414    };
4415    let name = rest.trim_matches(';').trim_matches('"').trim();
4416    if name.is_empty() {
4417        return Err(MongrelQueryError::Schema(
4418            "DROP TABLE missing table name".into(),
4419        ));
4420    }
4421    Ok((name.to_string(), if_exists))
4422}
4423
4424enum ParsedAlterTable {
4425    RenameTable {
4426        old_name: String,
4427        new_name: String,
4428    },
4429    RenameColumn {
4430        table_name: String,
4431        column_name: String,
4432        new_name: String,
4433    },
4434    AlterColumnType {
4435        table_name: String,
4436        column_name: String,
4437        ty: mongreldb_core::schema::TypeId,
4438    },
4439    SetNotNull {
4440        table_name: String,
4441        column_name: String,
4442    },
4443    DropNotNull {
4444        table_name: String,
4445        column_name: String,
4446    },
4447}
4448
4449fn current_column_flags(db: &Arc<Database>, table: &str, column: &str) -> Result<ColumnFlags> {
4450    let handle = db.table(table)?;
4451    let table = handle.lock();
4452    table
4453        .schema()
4454        .column(column)
4455        .map(|c| c.flags)
4456        .ok_or_else(|| MongrelQueryError::Schema(format!("unknown column {column}")))
4457}
4458
4459fn parse_alter_table(sql: &str) -> Result<ParsedAlterTable> {
4460    let trimmed = strip_statement_semicolon(sql.trim());
4461    let after_kw = strip_prefix_ci(trimmed, "ALTER TABLE")
4462        .ok_or_else(|| MongrelQueryError::Schema("not an ALTER TABLE statement".into()))?
4463        .trim();
4464    let (table_name, rest) = take_sql_ident(after_kw, "ALTER TABLE missing table name")?;
4465    let rest = rest.trim();
4466
4467    if let Some(after) = strip_prefix_ci(rest, "RENAME TO") {
4468        let new_name = parse_trailing_identifier(after, "ALTER TABLE missing new table name")?;
4469        return Ok(ParsedAlterTable::RenameTable {
4470            old_name: table_name,
4471            new_name,
4472        });
4473    }
4474
4475    if let Some(after) = strip_prefix_ci(rest, "RENAME COLUMN") {
4476        let (column_name, after_col) =
4477            take_sql_ident(after, "ALTER TABLE RENAME COLUMN missing column name")?;
4478        let after_to = strip_prefix_ci(after_col.trim(), "TO").ok_or_else(|| {
4479            MongrelQueryError::Schema("ALTER TABLE RENAME COLUMN missing TO".into())
4480        })?;
4481        let new_name = parse_trailing_identifier(
4482            after_to,
4483            "ALTER TABLE RENAME COLUMN missing new column name",
4484        )?;
4485        return Ok(ParsedAlterTable::RenameColumn {
4486            table_name,
4487            column_name,
4488            new_name,
4489        });
4490    }
4491
4492    let after_alter = strip_prefix_ci(rest, "ALTER COLUMN")
4493        .or_else(|| strip_prefix_ci(rest, "ALTER"))
4494        .ok_or_else(|| {
4495            MongrelQueryError::Schema(
4496                "ALTER TABLE must be RENAME TO, RENAME COLUMN, or ALTER COLUMN".into(),
4497            )
4498        })?;
4499    let (column_name, action) =
4500        take_sql_ident(after_alter, "ALTER TABLE ALTER COLUMN missing column name")?;
4501    let action = action.trim();
4502
4503    if let Some(after_type) =
4504        strip_prefix_ci(action, "TYPE").or_else(|| strip_prefix_ci(action, "SET DATA TYPE"))
4505    {
4506        let ty = parse_type_tail(after_type)?;
4507        return Ok(ParsedAlterTable::AlterColumnType {
4508            table_name,
4509            column_name,
4510            ty,
4511        });
4512    }
4513    if strip_prefix_ci(action, "SET NOT NULL").is_some() {
4514        return Ok(ParsedAlterTable::SetNotNull {
4515            table_name,
4516            column_name,
4517        });
4518    }
4519    if strip_prefix_ci(action, "DROP NOT NULL").is_some() {
4520        return Ok(ParsedAlterTable::DropNotNull {
4521            table_name,
4522            column_name,
4523        });
4524    }
4525
4526    Err(MongrelQueryError::Schema(
4527        "unsupported ALTER COLUMN action".into(),
4528    ))
4529}
4530
4531fn strip_statement_semicolon(s: &str) -> &str {
4532    s.trim().trim_end_matches(';').trim()
4533}
4534
4535fn take_sql_ident<'a>(s: &'a str, missing: &str) -> Result<(String, &'a str)> {
4536    let s = s.trim();
4537    if s.is_empty() {
4538        return Err(MongrelQueryError::Schema(missing.into()));
4539    }
4540    if let Some(rest) = s.strip_prefix('"') {
4541        let Some(end) = rest.find('"') else {
4542            return Err(MongrelQueryError::Schema(
4543                "unterminated quoted identifier".into(),
4544            ));
4545        };
4546        let ident = rest[..end].to_string();
4547        if ident.is_empty() {
4548            return Err(MongrelQueryError::Schema(missing.into()));
4549        }
4550        return Ok((ident, &rest[end + 1..]));
4551    }
4552    let end = s.find(|c: char| c.is_ascii_whitespace()).unwrap_or(s.len());
4553    let ident = s[..end].trim_matches('"').to_string();
4554    if ident.is_empty() {
4555        return Err(MongrelQueryError::Schema(missing.into()));
4556    }
4557    Ok((ident, &s[end..]))
4558}
4559
4560fn parse_trailing_identifier(s: &str, missing: &str) -> Result<String> {
4561    let (ident, rest) = take_sql_ident(s, missing)?;
4562    if !strip_statement_semicolon(rest).is_empty() {
4563        return Err(MongrelQueryError::Schema(
4564            "unexpected tokens after identifier".into(),
4565        ));
4566    }
4567    Ok(ident)
4568}
4569
4570fn parse_type_tail(s: &str) -> Result<mongreldb_core::schema::TypeId> {
4571    let tail = strip_statement_semicolon(s);
4572    let ty = tail
4573        .split_whitespace()
4574        .next()
4575        .ok_or_else(|| MongrelQueryError::Schema("ALTER COLUMN TYPE missing type".into()))?;
4576    parse_sql_type(ty)
4577}
4578
4579#[cfg(test)]
4580mod tests {
4581    use super::*;
4582
4583    #[test]
4584    fn bounded_lru_evicts_least_recently_used() {
4585        let mut cache = BoundedLru::new(2);
4586        cache.insert("a", 1);
4587        cache.insert("b", 2);
4588        assert_eq!(cache.get("a"), Some(&1));
4589        cache.insert("c", 3);
4590        assert_eq!(cache.entries.len(), 2);
4591        assert_eq!(cache.get("b"), None);
4592        assert_eq!(cache.get("a"), Some(&1));
4593        assert_eq!(cache.get("c"), Some(&3));
4594    }
4595
4596    #[tokio::test]
4597    async fn streaming_query_bypasses_result_cache() {
4598        use futures::StreamExt;
4599
4600        let dir = tempfile::tempdir().unwrap();
4601        let database = Arc::new(Database::create(dir.path()).unwrap());
4602        let session = MongrelSession::open(database).unwrap();
4603        let mut stream = session.run_stream("SELECT 1 AS value").await.unwrap();
4604        let batch = stream.next().await.unwrap().unwrap();
4605        assert_eq!(batch.num_rows(), 1);
4606        assert!(stream.next().await.is_none());
4607        assert!(session.cache.lock().entries.is_empty());
4608    }
4609
4610    #[tokio::test]
4611    async fn ttl_tables_bypass_epoch_keyed_result_cache() {
4612        let dir = tempfile::tempdir().unwrap();
4613        let database = Arc::new(Database::create(dir.path()).unwrap());
4614        let session = MongrelSession::open(database).unwrap();
4615        session
4616            .run(
4617                "CREATE TABLE events (id BIGINT PRIMARY KEY, ts TIMESTAMP) \
4618                 TTL_COLUMN ts TTL '1 day'",
4619            )
4620            .await
4621            .unwrap();
4622        session.run("SELECT * FROM events").await.unwrap();
4623        assert!(session.cache.lock().entries.is_empty());
4624    }
4625
4626    #[test]
4627    fn normalize_collapses_and_trims_whitespace() {
4628        assert_eq!(normalize_sql("SELECT * FROM t"), "SELECT * FROM t");
4629        assert_eq!(normalize_sql("  SELECT  *   FROM   t  "), "SELECT * FROM t");
4630        assert_eq!(
4631            normalize_sql("\n\tSELECT\n*\nFROM\n\tt\n"),
4632            "SELECT * FROM t"
4633        );
4634        assert_eq!(
4635            normalize_sql("SELECT   a,   b   FROM   t"),
4636            normalize_sql("SELECT a, b FROM t")
4637        );
4638    }
4639
4640    #[test]
4641    fn boolean_ai_predicate_detection_uses_sql_tokens() {
4642        assert!(contains_boolean_ai_predicate(
4643            "SELECT * FROM docs WHERE ANN_SEARCH (embedding, '[1]', 1)"
4644        ));
4645        assert!(contains_boolean_ai_predicate(
4646            "PREPARE q AS SELECT * FROM docs WHERE sparse_match(sparse, '[]', 1)"
4647        ));
4648        assert!(!contains_boolean_ai_predicate(
4649            "SELECT 'ann_search(x)', sparse_search_scored('docs', 'sparse', '[]', 1) /* ann_search(x) */"
4650        ));
4651    }
4652
4653    #[test]
4654    fn normalize_preserves_string_literal_whitespace() {
4655        assert_eq!(
4656            normalize_sql("SELECT 'hello   world' FROM t"),
4657            "SELECT 'hello   world' FROM t"
4658        );
4659        assert_eq!(
4660            normalize_sql("SELECT 'it''s   ok' FROM t"),
4661            "SELECT 'it''s   ok' FROM t"
4662        );
4663        assert_eq!(
4664            normalize_sql("  SELECT  'a  b'  FROM  t  "),
4665            "SELECT 'a  b' FROM t"
4666        );
4667    }
4668
4669    #[test]
4670    fn normalize_preserves_quoted_identifier_and_dollar_quote() {
4671        assert_eq!(
4672            normalize_sql("  SELECT  \"my col\"  FROM  t  "),
4673            "SELECT \"my col\" FROM t"
4674        );
4675        assert_eq!(
4676            normalize_sql("  SELECT  $$a   b$$  FROM  t  "),
4677            "SELECT $$a   b$$ FROM t"
4678        );
4679        assert_eq!(
4680            normalize_sql("SELECT $tag$body   with spaces$tag$ FROM t"),
4681            "SELECT $tag$body   with spaces$tag$ FROM t"
4682        );
4683    }
4684
4685    #[test]
4686    fn normalize_strips_comments() {
4687        assert_eq!(
4688            normalize_sql("SELECT 1 -- trailing comment\nFROM t"),
4689            "SELECT 1 FROM t"
4690        );
4691        assert_eq!(
4692            normalize_sql("SELECT /* block */ 1 FROM t"),
4693            "SELECT 1 FROM t"
4694        );
4695        // Comment with a quote-like body must not confuse the scanner.
4696        assert_eq!(
4697            normalize_sql("SELECT /* 'not a string' */ 1 FROM t"),
4698            "SELECT 1 FROM t"
4699        );
4700        // Nested block comments are honored (Postgres/DataFusion allow nesting).
4701        assert_eq!(
4702            normalize_sql("SELECT /* outer /* inner */ still outer */ 1 FROM t"),
4703            "SELECT 1 FROM t"
4704        );
4705    }
4706
4707    #[test]
4708    fn normalize_escape_string_preserved() {
4709        assert_eq!(
4710            normalize_sql("SELECT E'line\\nbreak' FROM t"),
4711            "SELECT E'line\\nbreak' FROM t"
4712        );
4713    }
4714
4715    #[test]
4716    fn replace_from_view_matches_whole_word_only() {
4717        let out = replace_from_view("SELECT * FROM logs", "log", "SELECT 1");
4718        assert_eq!(out, "SELECT * FROM logs");
4719
4720        let out = replace_from_view("SELECT * FROM log", "log", "SELECT 1");
4721        assert_eq!(out, "SELECT * FROM (SELECT 1) AS log");
4722
4723        let out = replace_from_view("select * from log where x", "log", "SELECT 1");
4724        assert_eq!(out, "select * from (SELECT 1) AS log where x");
4725
4726        let out = replace_from_view("SELECT * FROM log)", "log", "SELECT 1");
4727        assert_eq!(out, "SELECT * FROM (SELECT 1) AS log)");
4728
4729        let out = replace_from_view("SELECT * xfrom log", "log", "SELECT 1");
4730        assert_eq!(out, "SELECT * xfrom log");
4731    }
4732
4733    #[test]
4734    fn compat_function_rewrite_handles_sqlite_compatibility_calls() {
4735        assert_eq!(
4736            rewrite_compat_function_calls("select max(id), min(id) from t"),
4737            "select max(id), min(id) from t"
4738        );
4739        assert_eq!(
4740            rewrite_compat_function_calls("select max(1, min(2, 3), 'max(4,5)')"),
4741            "select __mongreldb_scalar_max(1, __mongreldb_scalar_min(2, 3), 'max(4,5)')"
4742        );
4743        assert_eq!(
4744            rewrite_compat_function_calls("select /* max(1,2) */ min(1, (2 + 3))"),
4745            "select /* max(1,2) */ __mongreldb_scalar_min(1, (2 + 3))"
4746        );
4747        assert_eq!(
4748            rewrite_compat_function_calls("select max_value, min_value from t"),
4749            "select max_value, min_value from t"
4750        );
4751        assert_eq!(
4752            rewrite_compat_function_calls(
4753                "select group_concat(label), group_concat(label, '|') from t"
4754            ),
4755            "select string_agg(label, ','), string_agg(label, '|') from t"
4756        );
4757        assert_eq!(
4758            rewrite_compat_function_calls("select total(val), total(val) filter (where grp = 2) from t"),
4759            "select coalesce(cast(sum(val) as double), 0.0), coalesce(cast(sum(val) filter (where grp = 2) as double), 0.0) from t"
4760        );
4761        assert_eq!(
4762            rewrite_compat_function_calls(
4763                "select total(val) over (partition by grp order by id) from t"
4764            ),
4765            "select coalesce(cast(sum(val) over (partition by grp order by id) as double), 0.0) from t"
4766        );
4767    }
4768}