pub mod arrow_conv;
mod commands;
mod error;
pub mod extended_sql_functions;
mod external_modules;
mod fk_join;
mod native_agg;
mod percentile;
mod scan;
mod shadow;
mod udf;
pub use error::{MongrelQueryError, Result};
pub use external_modules::{
ExternalBaseWrite, ExternalModuleDescriptor, ExternalModuleIndex, ExternalModuleRegistry,
ExternalPlan, ExternalPlanRequest, ExternalScan, ExternalTable, ExternalTableModule,
ExternalTxn, ExternalWriteOp, ExternalWriteResult, ModuleConnectCtx,
};
use arrow::array::{Array, ArrayRef, Int64Array, StringArray};
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use datafusion::catalog::{Session, TableProvider};
use datafusion::common::{DataFusionError, Result as DFResult};
use datafusion::logical_expr::{AggregateUDF, Expr, ScalarUDF, TableType, WindowUDF};
use datafusion::physical_plan::ExecutionPlan;
use datafusion::prelude::SessionContext;
use mongreldb_core::{
AlterColumn, ColumnFlags, Cursor, Database, Schema as CoreSchema, Table, TypeId,
};
use parking_lot::Mutex;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
pub struct MongrelProvider {
db: Arc<Mutex<Table>>,
schema: SchemaRef,
}
#[derive(Debug, Clone)]
pub(crate) struct ViewDef {
pub sql: String,
pub schema: CoreSchema,
pub input_types: HashMap<u16, Option<TypeId>>,
}
impl MongrelProvider {
pub fn new(db: Arc<Mutex<Table>>) -> Result<Self> {
let schema = {
let db = db.lock();
arrow_conv::arrow_schema(db.schema())?
};
Ok(Self { db, schema })
}
pub fn arrow_schema(&self) -> SchemaRef {
self.schema.clone()
}
}
impl std::fmt::Debug for MongrelProvider {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MongrelProvider").finish_non_exhaustive()
}
}
#[async_trait::async_trait]
impl TableProvider for MongrelProvider {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
fn table_type(&self) -> TableType {
TableType::Base
}
fn supports_filters_pushdown(
&self,
filters: &[&Expr],
) -> DFResult<Vec<datafusion::logical_expr::TableProviderFilterPushDown>> {
use datafusion::logical_expr::TableProviderFilterPushDown;
let schema_ref = self.db.lock().schema().clone();
Ok(filters
.iter()
.map(|f| match translate_filter(f, &schema_ref) {
Some(
mongreldb_core::Condition::FmContains { .. }
| mongreldb_core::Condition::FmContainsAll { .. },
) => TableProviderFilterPushDown::Inexact,
Some(_) => TableProviderFilterPushDown::Exact,
None => match translate_ann_search(f, &schema_ref)
.or_else(|| translate_sparse_match(f, &schema_ref))
{
Some(_) => TableProviderFilterPushDown::Exact,
None => TableProviderFilterPushDown::Unsupported,
},
})
.collect())
}
async fn scan(
&self,
_state: &dyn Session,
projection: Option<&Vec<usize>>,
filters: &[Expr],
_limit: Option<usize>,
) -> DFResult<Arc<dyn ExecutionPlan>> {
let core_err = |e: mongreldb_core::MongrelError| {
DataFusionError::External(Box::new(MongrelQueryError::Core(e)))
};
let mut db = self.db.lock();
let snap = db.snapshot();
let schema_ref = db.schema().clone();
let translated: Vec<mongreldb_core::Condition> = filters
.iter()
.filter_map(|f| {
translate_filter(f, &schema_ref)
.or_else(|| translate_ann_search(f, &schema_ref))
.or_else(|| translate_sparse_match(f, &schema_ref))
})
.collect();
if !translated.is_empty() {
db.ensure_indexes_complete().map_err(core_err)?;
}
let empty_proj = projection.map(|p| p.is_empty()).unwrap_or(false);
if empty_proj {
let total: usize = if translated.is_empty() {
mongreldb_core::trace::QueryTrace::record(|t| {
t.scan_mode = mongreldb_core::trace::ScanMode::CountMetadata;
});
db.count() as usize
} else if let Some(count) = db.count_conditions(&translated, snap).map_err(core_err)? {
count as usize
} else {
match schema_ref.columns.first() {
Some(cdef) => {
let one = [cdef.id];
let cols = match db
.query_columns_native_cached(&translated, Some(&one), snap)
.map_err(core_err)?
{
Some(c) => c,
None => db
.visible_columns_native(snap, Some(&one))
.map_err(core_err)?,
};
mongreldb_core::trace::QueryTrace::record(|t| {
t.scan_mode = mongreldb_core::trace::ScanMode::Materialized;
});
cols.first().map(|(_, c)| c.len()).unwrap_or(0)
}
None => 0,
}
};
return Ok(Arc::new(scan::MongrelScanExec::new_row_count(total)));
}
let (col_ids, scan_schema): (Vec<u16>, SchemaRef) = match projection {
Some(p) if !p.is_empty() => {
let ids = p.iter().map(|&idx| schema_ref.columns[idx].id).collect();
let fields: Vec<arrow::datatypes::Field> = p
.iter()
.map(|&idx| self.schema.field(idx).clone())
.collect();
(ids, Arc::new(arrow::datatypes::Schema::new(fields)))
}
_ => (
schema_ref.columns.iter().map(|c| c.id).collect(),
self.schema.clone(),
),
};
let mut proj_pairs: Vec<(u16, mongreldb_core::schema::TypeId)> =
Vec::with_capacity(col_ids.len());
let mut types: Vec<mongreldb_core::schema::TypeId> = Vec::with_capacity(col_ids.len());
for cid in &col_ids {
let ty = schema_ref
.columns
.iter()
.find(|c| c.id == *cid)
.map(|c| c.ty)
.ok_or_else(|| {
DataFusionError::External(Box::new(MongrelQueryError::Arrow(format!(
"unknown column {cid}"
))))
})?;
proj_pairs.push((*cid, ty));
types.push(ty);
}
let col_stats_map = if translated.is_empty() {
db.exact_column_stats(snap, &col_ids).map_err(core_err)?
} else {
None
};
let column_stats: Vec<datafusion::physical_plan::ColumnStatistics> = col_ids
.iter()
.map(|cid| scan::to_col_statistics(col_stats_map.as_ref().and_then(|m| m.get(cid))))
.collect();
if translated.is_empty()
&& db.run_count() == 1
&& db.memtable_is_empty()
&& db.mutable_run_len() == 0
&& db.single_run_is_clean()
{
let shadow = shadow::ArrowShadow::new(db.dir());
let run_ids: HashSet<u128> = db.run_ids().into_iter().collect();
shadow.sweep(&run_ids);
if let Some(&run_id) = run_ids.iter().next() {
if let Some(batch) = shadow.try_read(run_id) {
if let Some(projected) =
project_batch(&batch, &col_ids, &schema_ref, &scan_schema)
{
mongreldb_core::trace::QueryTrace::record(|t| {
t.scan_mode = mongreldb_core::trace::ScanMode::ArrowShadow;
});
return Ok(Arc::new(scan::MongrelScanExec::new_batch(
scan_schema,
projected,
column_stats,
)));
}
}
}
}
let cursor: Option<Box<dyn Cursor>> = db
.scan_cursor(snap, proj_pairs, &translated)
.map_err(core_err)?;
if let Some(cursor) = cursor {
let num_rows = cursor.remaining_rows();
let residual = extract_residual_filter(filters, &col_ids, &schema_ref);
return Ok(Arc::new(scan::MongrelScanExec::new_cursor(
scan_schema,
types,
cursor,
num_rows,
column_stats,
residual,
)));
}
let cols = if !translated.is_empty() {
match db
.query_columns_native_cached(&translated, Some(&col_ids), snap)
.map_err(core_err)?
{
Some(c) => c,
None => db
.visible_columns_native(snap, Some(&col_ids))
.map_err(core_err)?,
}
} else {
db.visible_columns_native(snap, Some(&col_ids))
.map_err(core_err)?
};
let mut ordered: Vec<mongreldb_core::columnar::NativeColumn> =
Vec::with_capacity(col_ids.len());
for cid in &col_ids {
let col = cols
.iter()
.find(|(id, _)| id == cid)
.map(|(_, c)| c.clone())
.ok_or_else(|| {
DataFusionError::External(Box::new(MongrelQueryError::Arrow(format!(
"missing column {cid}"
))))
})?;
ordered.push(col);
}
let num_rows = ordered.first().map(|c| c.len()).unwrap_or(0);
let shadow_write: Option<(
std::path::PathBuf,
u128,
Vec<arrow::array::ArrayRef>,
SchemaRef,
)> = if translated.is_empty()
&& db.run_count() == 1
&& db.memtable_is_empty()
&& db.mutable_run_len() == 0
&& db.single_run_is_clean()
{
let all_schema_ids: Vec<u16> = schema_ref.columns.iter().map(|c| c.id).collect();
if col_ids == all_schema_ids {
let dir = db.dir().to_path_buf();
let run_id = db.run_ids().first().copied();
run_id.map(|rid| {
let arrays = ordered
.iter()
.zip(types.iter())
.map(|(col, &ty)| arrow_conv::native_to_array(ty, col))
.collect::<Result<_>>()
.unwrap_or_default();
(dir, rid, arrays, scan_schema.clone())
})
} else {
None
}
} else {
None
};
drop(db);
if let Some((dir, run_id, arrays, schema)) = shadow_write {
if let Ok(batch) = RecordBatch::try_new(schema, arrays) {
shadow::ArrowShadow::new(&dir).write(run_id, &batch);
}
}
mongreldb_core::trace::QueryTrace::record(|t| {
t.scan_mode = mongreldb_core::trace::ScanMode::Materialized;
t.row_materialized = true;
});
Ok(Arc::new(scan::MongrelScanExec::new(
scan_schema,
ordered,
types,
num_rows,
column_stats,
)))
}
}
fn project_batch(
batch: &RecordBatch,
col_ids: &[u16],
schema_ref: &mongreldb_core::schema::Schema,
scan_schema: &arrow::datatypes::SchemaRef,
) -> Option<RecordBatch> {
let arrays: Vec<arrow::array::ArrayRef> = col_ids
.iter()
.map(|cid| {
let name = schema_ref
.columns
.iter()
.find(|c| c.id == *cid)
.map(|c| c.name.as_str())?;
let idx = batch.schema().index_of(name).ok()?;
Some(batch.column(idx).clone())
})
.collect::<Option<Vec<_>>>()?;
RecordBatch::try_new(scan_schema.clone(), arrays).ok()
}
pub(crate) fn translate_filter(
expr: &Expr,
schema: &mongreldb_core::Schema,
) -> Option<mongreldb_core::Condition> {
use datafusion::common::ScalarValue;
use datafusion::logical_expr::{Between, BinaryExpr, Like, Operator};
use mongreldb_core::{ColumnFlags, Condition, IndexKind, TypeId, Value};
let int_val = |s: &ScalarValue| match s {
ScalarValue::Int8(Some(v)) => Some(*v as i64),
ScalarValue::Int16(Some(v)) => Some(*v as i64),
ScalarValue::Int32(Some(v)) => Some(*v as i64),
ScalarValue::Int64(Some(v)) => Some(*v),
ScalarValue::UInt8(Some(v)) => Some(*v as i64),
ScalarValue::UInt16(Some(v)) => Some(*v as i64),
ScalarValue::UInt32(Some(v)) => Some(*v as i64),
ScalarValue::UInt64(Some(v)) => Some(*v as i64),
ScalarValue::Date32(Some(v)) => Some(*v as i64),
ScalarValue::TimestampSecond(Some(v), _) => Some(*v),
ScalarValue::TimestampMillisecond(Some(v), _) => Some(*v),
ScalarValue::TimestampMicrosecond(Some(v), _) => Some(*v),
ScalarValue::TimestampNanosecond(Some(v), _) => Some(*v),
_ => None,
};
let float_val = |s: &ScalarValue| match s {
ScalarValue::Float32(Some(f)) => Some(*f as f64),
ScalarValue::Float64(Some(f)) => Some(*f),
_ => None,
};
let bytes_val = |s: &ScalarValue| match s {
ScalarValue::Utf8(Some(s)) => Some(s.as_bytes().to_vec()),
_ => None,
};
let _ = bytes_val;
let val = |s: &ScalarValue| -> Option<Value> {
if let Some(i) = int_val(s) {
return Some(Value::Int64(i));
}
match s {
ScalarValue::Utf8(Some(s)) => Some(Value::Bytes(s.as_bytes().to_vec())),
ScalarValue::Float32(Some(f)) => Some(Value::Float64(*f as f64)),
ScalarValue::Float64(Some(f)) => Some(Value::Float64(*f)),
ScalarValue::Boolean(Some(b)) => Some(Value::Bool(*b)),
_ => None,
}
};
let col_def = |name: &str| schema.columns.iter().find(|c| c.name == name);
let has_fm = |cid: u16| {
schema
.indexes
.iter()
.any(|i| i.column_id == cid && i.kind == IndexKind::FmIndex)
};
let has_bitmap = |cid: u16| {
schema
.indexes
.iter()
.any(|i| i.column_id == cid && i.kind == IndexKind::Bitmap)
};
match expr {
Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
if *op == Operator::Or {
return try_or_as_bitmap_in(expr, schema);
}
let left = peel_cast(left);
let right = peel_cast(right);
let (col_name, scalar, flipped) = match (left.as_ref(), right.as_ref()) {
(Expr::Column(c), Expr::Literal(s, _)) => (&c.name, s, false),
(Expr::Literal(s, _), Expr::Column(c)) => (&c.name, s, true),
_ => return None,
};
let op = if flipped { flip_op(*op)? } else { *op };
let cdef = col_def(col_name)?;
if op == Operator::Eq {
let v = val(scalar)?;
if has_bitmap(cdef.id) {
return Some(Condition::BitmapEq {
column_id: cdef.id,
value: v.encode_key(),
});
}
if cdef.flags.contains(ColumnFlags::PRIMARY_KEY) {
return Some(Condition::Pk(v.encode_key()));
}
return None;
}
match cdef.ty {
TypeId::Int8
| TypeId::Int16
| TypeId::Int32
| TypeId::Int64
| TypeId::UInt8
| TypeId::UInt16
| TypeId::UInt32
| TypeId::UInt64
| TypeId::TimestampNanos
| TypeId::Date32 => {
let v = int_val(scalar)?;
let (lo, hi) = int_bounds(op, v)?;
Some(Condition::Range {
column_id: cdef.id,
lo,
hi,
})
}
TypeId::Float32 | TypeId::Float64 => {
let v = float_val(scalar)?;
let (lo, lo_inc, hi, hi_inc) = float_bounds(op, v)?;
Some(Condition::RangeF64 {
column_id: cdef.id,
lo,
lo_inclusive: lo_inc,
hi,
hi_inclusive: hi_inc,
})
}
_ => None,
}
}
Expr::Between(Between {
expr,
negated,
low,
high,
}) => {
if *negated {
return None;
}
let Expr::Column(c) = expr.as_ref() else {
return None;
};
let cdef = col_def(&c.name)?;
let (lo_s, hi_s) = match (low.as_ref(), high.as_ref()) {
(Expr::Literal(lo, _), Expr::Literal(hi, _)) => (lo, hi),
_ => return None,
};
match cdef.ty {
TypeId::Int8
| TypeId::Int16
| TypeId::Int32
| TypeId::Int64
| TypeId::UInt8
| TypeId::UInt16
| TypeId::UInt32
| TypeId::UInt64
| TypeId::TimestampNanos
| TypeId::Date32 => {
let (Some(lo), Some(hi)) = (int_val(lo_s), int_val(hi_s)) else {
return None;
};
Some(Condition::Range {
column_id: cdef.id,
lo,
hi,
})
}
TypeId::Float32 | TypeId::Float64 => {
let (Some(lo), Some(hi)) = (float_val(lo_s), float_val(hi_s)) else {
return None;
};
Some(Condition::RangeF64 {
column_id: cdef.id,
lo,
lo_inclusive: true,
hi,
hi_inclusive: true,
})
}
_ => None,
}
}
Expr::Like(Like {
negated,
expr,
pattern,
..
}) => {
if *negated {
return None;
}
let Expr::Column(c) = expr.as_ref() else {
return None;
};
let Expr::Literal(ScalarValue::Utf8(Some(pat)), _) = pattern.as_ref() else {
return None;
};
let cdef = col_def(&c.name)?;
if has_bitmap(cdef.id) {
if let Some(prefix) = anchored_like_prefix(pat) {
return Some(Condition::BytesPrefix {
column_id: cdef.id,
prefix: mongreldb_core::Value::Bytes(prefix.as_bytes().to_vec())
.encode_key(),
});
}
}
if !has_fm(cdef.id) {
return None;
}
let segments: Vec<Vec<u8>> = pat
.split(['%', '_'])
.filter(|s| s.len() >= 3)
.map(|s| s.as_bytes().to_vec())
.collect();
match segments.len() {
0 => longest_like_segment(pat).map(|seg| Condition::FmContains {
column_id: cdef.id,
pattern: seg,
}),
1 => Some(Condition::FmContains {
column_id: cdef.id,
pattern: segments.into_iter().next().unwrap(),
}),
_ => Some(Condition::FmContainsAll {
column_id: cdef.id,
patterns: segments,
}),
}
}
Expr::InList(il) if !il.negated => {
let Expr::Column(c) = il.expr.as_ref() else {
return None;
};
let cdef = col_def(&c.name)?;
if !has_bitmap(cdef.id) {
return None;
}
let values: Vec<Vec<u8>> = il
.list
.iter()
.filter_map(|e| match e {
Expr::Literal(s, _) => val(s).map(|v| v.encode_key()),
_ => None,
})
.collect();
if values.is_empty() || values.len() != il.list.len() {
return None;
}
Some(Condition::BitmapIn {
column_id: cdef.id,
values,
})
}
Expr::IsNull(inner) => {
let col_name = match inner.as_ref() {
Expr::Column(c) => &c.name,
_ => return None,
};
let cdef = col_def(col_name)?;
Some(Condition::IsNull { column_id: cdef.id })
}
Expr::IsNotNull(inner) => {
let col_name = match inner.as_ref() {
Expr::Column(c) => &c.name,
_ => return None,
};
let cdef = col_def(col_name)?;
Some(Condition::IsNotNull { column_id: cdef.id })
}
_ => None,
}
}
pub(crate) fn extract_residual_filter(
filters: &[Expr],
col_ids: &[u16],
schema: &mongreldb_core::Schema,
) -> Option<std::sync::Arc<scan::ResidualFilter>> {
use datafusion::common::ScalarValue;
use datafusion::logical_expr::Like;
for f in filters {
if let Expr::Like(Like {
negated: false,
expr,
pattern,
..
}) = f
{
let Expr::Column(c) = expr.as_ref() else {
continue;
};
let Expr::Literal(ScalarValue::Utf8(Some(pat)), _) = pattern.as_ref() else {
continue;
};
let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
let col_idx = col_ids.iter().position(|&id| id == cdef.id)?;
return Some(std::sync::Arc::new(scan::ResidualFilter::new(
col_idx,
pat.as_bytes().to_vec(),
)));
}
}
None
}
pub(crate) fn translate_ann_search(
expr: &Expr,
schema: &mongreldb_core::Schema,
) -> Option<mongreldb_core::Condition> {
use datafusion::common::ScalarValue;
use mongreldb_core::Condition;
let Expr::ScalarFunction(sf) = expr else {
return None;
};
if !sf.func.name().eq_ignore_ascii_case("ann_search") || sf.args.len() != 3 {
return None;
}
let (Expr::Column(c), query_expr, k_expr) = (&sf.args[0], &sf.args[1], &sf.args[2]) else {
return None;
};
let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
let json = match query_expr {
Expr::Literal(ScalarValue::Utf8(Some(s)), _) => s.as_str(),
_ => return None,
};
let k: i64 = match k_expr {
Expr::Literal(scalar, _) => match scalar {
ScalarValue::Int64(Some(k)) => *k,
ScalarValue::UInt64(Some(k)) => *k as i64,
ScalarValue::Int32(Some(k)) => *k as i64,
_ => return None,
},
_ => return None,
};
let query: Vec<f32> = serde_json::from_str(json).ok()?;
Some(Condition::Ann {
column_id: cdef.id,
query,
k: k.max(1) as usize,
})
}
pub(crate) fn translate_sparse_match(
expr: &Expr,
schema: &mongreldb_core::Schema,
) -> Option<mongreldb_core::Condition> {
use datafusion::common::ScalarValue;
use mongreldb_core::Condition;
let Expr::ScalarFunction(sf) = expr else {
return None;
};
if !sf.func.name().eq_ignore_ascii_case("sparse_match") || sf.args.len() != 3 {
return None;
}
let (Expr::Column(c), query_expr, k_expr) = (&sf.args[0], &sf.args[1], &sf.args[2]) else {
return None;
};
let cdef = schema.columns.iter().find(|col| col.name == c.name)?;
let json = match query_expr {
Expr::Literal(ScalarValue::Utf8(Some(s)), _) => s.as_str(),
_ => return None,
};
let k: i64 = match k_expr {
Expr::Literal(scalar, _) => match scalar {
ScalarValue::Int64(Some(k)) => *k,
ScalarValue::UInt64(Some(k)) => *k as i64,
ScalarValue::Int32(Some(k)) => *k as i64,
_ => return None,
},
_ => return None,
};
let query: Vec<(u32, f32)> = serde_json::from_str(json).ok()?;
Some(Condition::SparseMatch {
column_id: cdef.id,
query,
k: k.max(1) as usize,
})
}
fn flip_op(op: datafusion::logical_expr::Operator) -> Option<datafusion::logical_expr::Operator> {
use datafusion::logical_expr::Operator;
Some(match op {
Operator::Eq => Operator::Eq,
Operator::Lt => Operator::Gt,
Operator::Gt => Operator::Lt,
Operator::LtEq => Operator::GtEq,
Operator::GtEq => Operator::LtEq,
_ => return None,
})
}
fn int_bounds(op: datafusion::logical_expr::Operator, v: i64) -> Option<(i64, i64)> {
use datafusion::logical_expr::Operator;
Some(match op {
Operator::Gt => (v.saturating_add(1), i64::MAX),
Operator::GtEq => (v, i64::MAX),
Operator::Lt => (i64::MIN, v.saturating_sub(1)),
Operator::LtEq => (i64::MIN, v),
_ => return None,
})
}
fn float_bounds(op: datafusion::logical_expr::Operator, v: f64) -> Option<(f64, bool, f64, bool)> {
use datafusion::logical_expr::Operator;
Some(match op {
Operator::Gt => (v, false, f64::INFINITY, false),
Operator::GtEq => (v, true, f64::INFINITY, false),
Operator::Lt => (f64::NEG_INFINITY, false, v, false),
Operator::LtEq => (f64::NEG_INFINITY, false, v, true),
_ => return None,
})
}
fn longest_like_segment(pat: &str) -> Option<Vec<u8>> {
pat.split(['%', '_'])
.map(|s| s.as_bytes())
.max_by_key(|s| s.len())
.filter(|s| !s.is_empty())
.map(|s| s.to_vec())
}
fn anchored_like_prefix(pat: &str) -> Option<&str> {
let rest = pat.strip_suffix('%')?;
if rest.is_empty() || rest.contains(['%', '_']) {
return None;
}
Some(rest)
}
fn peel_cast(expr: &Expr) -> std::borrow::Cow<'_, Expr> {
match expr {
Expr::Cast(datafusion::logical_expr::Cast { expr, .. }) => std::borrow::Cow::Borrowed(expr),
_ => std::borrow::Cow::Borrowed(expr),
}
}
fn try_or_as_bitmap_in(
expr: &Expr,
schema: &mongreldb_core::Schema,
) -> Option<mongreldb_core::Condition> {
use datafusion::logical_expr::{BinaryExpr, Operator};
let mut values: Vec<Vec<u8>> = Vec::new();
let mut target_col: Option<u16> = None;
let mut stack = vec![expr];
while let Some(e) = stack.pop() {
match e {
Expr::BinaryExpr(BinaryExpr {
left,
op: Operator::Or,
right,
}) => {
stack.push(left);
stack.push(right);
}
Expr::BinaryExpr(BinaryExpr {
left,
op: Operator::Eq,
right,
}) => {
let (col_name, scalar) = match (left.as_ref(), right.as_ref()) {
(Expr::Column(c), Expr::Literal(s, _)) => (&c.name, s),
(Expr::Literal(s, _), Expr::Column(c)) => (&c.name, s),
_ => return None,
};
let cdef = schema.columns.iter().find(|c| &c.name == col_name)?;
if !schema
.indexes
.iter()
.any(|i| i.column_id == cdef.id && i.kind == mongreldb_core::IndexKind::Bitmap)
{
return None;
}
match target_col {
None => target_col = Some(cdef.id),
Some(id) if id != cdef.id => return None,
_ => {}
}
let v = match scalar {
datafusion::common::ScalarValue::Int64(Some(v)) => {
mongreldb_core::Value::Int64(*v)
}
datafusion::common::ScalarValue::Utf8(Some(s)) => {
mongreldb_core::Value::Bytes(s.as_bytes().to_vec())
}
datafusion::common::ScalarValue::Float64(Some(f)) => {
mongreldb_core::Value::Float64(*f)
}
datafusion::common::ScalarValue::Boolean(Some(b)) => {
mongreldb_core::Value::Bool(*b)
}
_ => return None,
};
values.push(v.encode_key());
}
_ => return None,
}
}
let col_id = target_col?;
if values.is_empty() {
return None;
}
Some(mongreldb_core::Condition::BitmapIn {
column_id: col_id,
values,
})
}
fn sp_ident_name(expr: &sqlparser::ast::Expr) -> Option<&str> {
use sqlparser::ast::Expr;
match expr {
Expr::Identifier(ident) => Some(ident.value.as_str()),
Expr::CompoundIdentifier(idents) => idents.last().map(|i| i.value.as_str()),
_ => None,
}
}
fn sp_literal(expr: &sqlparser::ast::Expr) -> Option<mongreldb_core::Value> {
use sqlparser::ast::Expr;
let v = match expr {
Expr::Value(v) => v,
_ => return None,
};
use sqlparser::ast::Value as SpValue;
match &v.value {
SpValue::Number(s, _) => s
.parse::<i64>()
.map(mongreldb_core::Value::Int64)
.or_else(|_| s.parse::<f64>().map(mongreldb_core::Value::Float64))
.ok(),
SpValue::SingleQuotedString(s) => Some(mongreldb_core::Value::Bytes(s.as_bytes().to_vec())),
SpValue::Boolean(b) => Some(mongreldb_core::Value::Bool(*b)),
_ => None,
}
}
fn is_int_ty(ty: mongreldb_core::schema::TypeId) -> bool {
use mongreldb_core::schema::TypeId::*;
matches!(
ty,
Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 | TimestampNanos | Date32
)
}
fn is_float_ty(ty: mongreldb_core::schema::TypeId) -> bool {
matches!(
ty,
mongreldb_core::schema::TypeId::Float32 | mongreldb_core::schema::TypeId::Float64
)
}
fn translate_sqlparser_predicate(
expr: &sqlparser::ast::Expr,
schema: &mongreldb_core::Schema,
) -> Option<mongreldb_core::Condition> {
use mongreldb_core::{schema::ColumnFlags, Condition, IndexKind, Value};
use sqlparser::ast::{BinaryOperator, Expr};
let col_def = |name: &str| schema.columns.iter().find(|c| c.name == name);
let has_bitmap = |cid: u16| {
schema
.indexes
.iter()
.any(|i| i.column_id == cid && i.kind == IndexKind::Bitmap)
};
match expr {
Expr::BinaryOp {
left,
op: BinaryOperator::Or,
right,
} => {
let mut values: Vec<Vec<u8>> = Vec::new();
let mut target: Option<u16> = None;
let mut stack: Vec<&Expr> = vec![left.as_ref(), right.as_ref()];
while let Some(e) = stack.pop() {
match e {
Expr::BinaryOp {
left,
op: BinaryOperator::Or,
right,
} => {
stack.push(left.as_ref());
stack.push(right.as_ref());
}
Expr::BinaryOp {
left,
op: BinaryOperator::Eq,
right,
} => {
let (name, lit) = match (left.as_ref(), right.as_ref()) {
(l, r) if sp_ident_name(l).is_some() && sp_literal(r).is_some() => {
(l, r)
}
(l, r) if sp_ident_name(r).is_some() && sp_literal(l).is_some() => {
(r, l)
}
_ => return None,
};
let cdef = col_def(sp_ident_name(name)?)?;
if !has_bitmap(cdef.id) {
return None;
}
match target {
None => target = Some(cdef.id),
Some(id) if id != cdef.id => return None,
_ => {}
}
values.push(sp_literal(lit)?.encode_key());
}
_ => return None,
}
}
let cid = target?;
(!values.is_empty()).then_some(Condition::BitmapIn {
column_id: cid,
values,
})
}
Expr::BinaryOp { left, op, right } => {
let flipped;
let (col_expr, lit_expr) = match (
sp_ident_name(left),
sp_literal(right),
sp_ident_name(right),
sp_literal(left),
) {
(Some(_), Some(_), _, _) => {
flipped = false;
(left.as_ref(), right.as_ref())
}
(_, _, Some(_), Some(_)) => {
flipped = true;
(right.as_ref(), left.as_ref())
}
_ => return None,
};
let name = sp_ident_name(col_expr)?;
let cdef = col_def(name)?;
let v = sp_literal(lit_expr)?;
use sqlparser::ast::BinaryOperator::*;
match op {
Eq => {
if has_bitmap(cdef.id) {
Some(Condition::BitmapEq {
column_id: cdef.id,
value: v.encode_key(),
})
} else if cdef.flags.contains(ColumnFlags::PRIMARY_KEY) {
Some(Condition::Pk(v.encode_key()))
} else {
None
}
}
Lt | LtEq | Gt | GtEq if is_int_ty(cdef.ty) => {
let n = match v {
Value::Int64(n) => n,
_ => return None,
};
let (lo, hi) = match (flipped, op) {
(false, Lt) | (true, Gt) => (i64::MIN, n.saturating_sub(1)),
(false, Gt) | (true, Lt) => (n.saturating_add(1), i64::MAX),
(false, LtEq) | (true, GtEq) => (i64::MIN, n),
(false, GtEq) | (true, LtEq) => (n, i64::MAX),
_ => (i64::MIN, i64::MAX),
};
Some(Condition::Range {
column_id: cdef.id,
lo,
hi,
})
}
Lt | LtEq | Gt | GtEq if is_float_ty(cdef.ty) => {
let f = match v {
Value::Float64(f) => f,
_ => return None,
};
let (lo, li, hi, hi_i) = match (flipped, op) {
(false, Lt) | (true, Gt) => (f64::NEG_INFINITY, true, f, false),
(false, Gt) | (true, Lt) => (f, false, f64::INFINITY, true),
(false, LtEq) | (true, GtEq) => (f64::NEG_INFINITY, true, f, true),
(false, GtEq) | (true, LtEq) => (f, true, f64::INFINITY, true),
_ => (f64::NEG_INFINITY, true, f64::INFINITY, true),
};
Some(Condition::RangeF64 {
column_id: cdef.id,
lo,
lo_inclusive: li,
hi,
hi_inclusive: hi_i,
})
}
_ => None,
}
}
Expr::Between {
expr,
negated,
low,
high,
} if !negated => {
let name = sp_ident_name(expr)?;
let cdef = col_def(name)?;
if is_int_ty(cdef.ty) {
let lo = match sp_literal(low)? {
Value::Int64(n) => n,
_ => return None,
};
let hi = match sp_literal(high)? {
Value::Int64(n) => n,
_ => return None,
};
Some(Condition::Range {
column_id: cdef.id,
lo,
hi,
})
} else if is_float_ty(cdef.ty) {
let lo = match sp_literal(low)? {
Value::Float64(f) => f,
_ => return None,
};
let hi = match sp_literal(high)? {
Value::Float64(f) => f,
_ => return None,
};
Some(Condition::RangeF64 {
column_id: cdef.id,
lo,
lo_inclusive: true,
hi,
hi_inclusive: true,
})
} else {
None
}
}
Expr::InList {
expr,
list,
negated,
} if !negated => {
let name = sp_ident_name(expr)?;
let cdef = col_def(name)?;
if !has_bitmap(cdef.id) {
return None;
}
let values: Vec<Vec<u8>> = list
.iter()
.map(|e| sp_literal(e).map(|v| v.encode_key()))
.collect::<Option<_>>()?;
(!values.is_empty()).then_some(Condition::BitmapIn {
column_id: cdef.id,
values,
})
}
Expr::IsNull(inner) => {
let cid = col_def(sp_ident_name(inner)?)?.id;
Some(Condition::IsNull { column_id: cid })
}
Expr::IsNotNull(inner) => {
let cid = col_def(sp_ident_name(inner)?)?.id;
Some(Condition::IsNotNull { column_id: cid })
}
_ => None,
}
}
fn translate_sqlparser_filter(
expr: &sqlparser::ast::Expr,
schema: &mongreldb_core::Schema,
) -> Option<Vec<mongreldb_core::Condition>> {
use sqlparser::ast::{BinaryOperator, Expr};
let mut out = Vec::new();
let mut stack = vec![expr];
while let Some(e) = stack.pop() {
match e {
Expr::BinaryOp {
left,
op: BinaryOperator::And,
right,
} => {
stack.push(left.as_ref());
stack.push(right.as_ref());
}
other => out.push(translate_sqlparser_predicate(other, schema)?),
}
}
Some(out)
}
pub struct MongrelSession {
ctx: SessionContext,
db: Option<Arc<Mutex<Table>>>,
database: Option<Arc<Database>>,
cache: ResultCache,
plan_cache: parking_lot::Mutex<HashMap<String, datafusion::logical_expr::LogicalPlan>>,
tables: parking_lot::Mutex<HashMap<String, Arc<Mutex<Table>>>>,
views: parking_lot::Mutex<HashMap<String, ViewDef>>,
attached_databases: parking_lot::Mutex<HashMap<String, Arc<Database>>>,
sql_txn: parking_lot::Mutex<Option<Vec<commands::PendingSqlOp>>>,
savepoints: parking_lot::Mutex<Vec<(String, usize)>>,
sql_fn_state: Arc<extended_sql_functions::ExtendedSqlState>,
external_modules: Arc<ExternalModuleRegistry>,
}
type CacheKey = (String, u64);
type ResultCache = parking_lot::Mutex<std::collections::HashMap<CacheKey, Arc<Vec<RecordBatch>>>>;
impl MongrelSession {
pub fn new(db: Table) -> Self {
let db = Arc::new(Mutex::new(db));
let ctx = SessionContext::new();
let sql_fn_state = Arc::new(extended_sql_functions::ExtendedSqlState::default());
register_mongrel_functions(&ctx, Arc::clone(&sql_fn_state));
let external_modules = Arc::new(ExternalModuleRegistry::default());
Self {
ctx,
db: Some(db),
database: None,
cache: parking_lot::Mutex::new(std::collections::HashMap::new()),
plan_cache: parking_lot::Mutex::new(HashMap::new()),
tables: parking_lot::Mutex::new(HashMap::new()),
views: parking_lot::Mutex::new(HashMap::new()),
attached_databases: parking_lot::Mutex::new(HashMap::new()),
savepoints: parking_lot::Mutex::new(Vec::new()),
sql_txn: parking_lot::Mutex::new(None),
sql_fn_state,
external_modules,
}
}
pub fn new_with_external_modules(
db: Table,
modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
) -> Result<Self> {
let session = Self::new(db);
for module in modules {
session.register_external_module(module)?;
}
Ok(session)
}
pub fn open(database: Arc<Database>) -> Result<Self> {
Self::open_with_external_modules(database, std::iter::empty())
}
pub fn open_with_external_modules(
database: Arc<Database>,
modules: impl IntoIterator<Item = Arc<dyn ExternalTableModule>>,
) -> Result<Self> {
let ctx = SessionContext::new();
let sql_fn_state = Arc::new(extended_sql_functions::ExtendedSqlState::default());
register_mongrel_functions(&ctx, Arc::clone(&sql_fn_state));
let external_modules = Arc::new(ExternalModuleRegistry::default());
for module in modules {
external_modules.register(module)?;
}
let mut tables: HashMap<String, Arc<Mutex<Table>>> = HashMap::new();
for name in database.table_names() {
let handle = database.table(&name)?;
let provider = MongrelProvider::new(handle.clone())?;
ctx.register_table(&name, Arc::new(provider))
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
tables.insert(name, handle);
}
for entry in database.external_tables() {
let provider = external_modules.external_table_provider(&database, &entry)?;
ctx.register_table(&entry.name, provider)
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
}
let primary = {
let mut names: Vec<&String> = tables.keys().collect();
names.sort();
names.first().and_then(|n| tables.get(*n).cloned())
};
Ok(Self {
ctx,
db: primary,
database: Some(database),
cache: parking_lot::Mutex::new(std::collections::HashMap::new()),
plan_cache: parking_lot::Mutex::new(HashMap::new()),
tables: parking_lot::Mutex::new(tables),
views: parking_lot::Mutex::new(HashMap::new()),
attached_databases: parking_lot::Mutex::new(HashMap::new()),
savepoints: parking_lot::Mutex::new(Vec::new()),
sql_txn: parking_lot::Mutex::new(None),
sql_fn_state,
external_modules,
})
}
pub fn register_external_module(&self, module: Arc<dyn ExternalTableModule>) -> Result<()> {
self.external_modules.register(module)?;
self.clear_cache();
Ok(())
}
pub fn db(&self) -> Option<&Arc<Mutex<Table>>> {
self.db.as_ref()
}
pub fn create_view(&self, name: &str, sql: &str) {
self.create_view_with_schema(name, sql, CoreSchema::default(), HashMap::new());
}
pub(crate) fn create_view_with_schema(
&self,
name: &str,
sql: &str,
schema: CoreSchema,
input_types: HashMap<u16, Option<TypeId>>,
) {
self.views.lock().insert(
name.to_string(),
ViewDef {
sql: sql.to_string(),
schema,
input_types,
},
);
}
pub fn drop_view(&self, name: &str) {
self.views.lock().remove(name);
}
pub(crate) fn view_schema(&self, name: &str) -> Option<CoreSchema> {
self.views.lock().get(name).map(|view| view.schema.clone())
}
pub(crate) fn view_definition(&self, name: &str) -> Option<ViewDef> {
self.views.lock().get(name).cloned()
}
pub async fn register(&self, name: &str) -> Result<()> {
let db = self.db.clone().ok_or(MongrelQueryError::Core(
mongreldb_core::MongrelError::NotFound("no primary table".into()),
))?;
let provider = MongrelProvider::new(db.clone())?;
self.tables.lock().insert(name.to_string(), db);
self.ctx
.register_table(name, Arc::new(provider))
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
Ok(())
}
pub async fn register_db(&self, name: &str, db: Table) -> Result<()> {
let db_arc = Arc::new(Mutex::new(db));
let provider = MongrelProvider::new(db_arc.clone())?;
self.tables.lock().insert(name.to_string(), db_arc);
self.ctx
.register_table(name, Arc::new(provider))
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
Ok(())
}
fn refresh_registered_table(&self, db: &Arc<Database>, name: &str) -> Result<()> {
self.ctx
.deregister_table(name)
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
let handle = db.table(name)?;
let provider = MongrelProvider::new(handle.clone())?;
self.ctx
.register_table(name, Arc::new(provider))
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
self.tables.lock().insert(name.to_string(), handle);
Ok(())
}
fn try_sqlite_master(&self, sql: &str) -> Result<Option<Vec<RecordBatch>>> {
let lower = sql.to_ascii_lowercase();
if !lower.contains("sqlite_master") && !lower.contains("sqlite_schema") {
return Ok(None);
}
use arrow::array::{ArrayRef, Int64Array, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
let mut types: Vec<String> = Vec::new();
let mut names: Vec<String> = Vec::new();
let mut tbl_names: Vec<String> = Vec::new();
for name in self.tables.lock().keys() {
types.push("table".into());
names.push(name.clone());
tbl_names.push(name.clone());
}
for name in self.views.lock().keys() {
types.push("view".into());
names.push(name.clone());
tbl_names.push(name.clone());
}
if let Some(db) = &self.database {
for t in db.triggers() {
let target_name = match &t.target {
mongreldb_core::trigger::TriggerTarget::Table(n) | mongreldb_core::trigger::TriggerTarget::View(n) => n.clone(),
};
types.push("trigger".into());
names.push(t.name.clone());
tbl_names.push(target_name);
}
}
let schema = Arc::new(Schema::new(vec![
Field::new("type", DataType::Utf8, false),
Field::new("name", DataType::Utf8, false),
Field::new("tbl_name", DataType::Utf8, false),
Field::new("rootpage", DataType::Int64, false),
Field::new("sql", DataType::Utf8, true),
]));
let n = names.len();
let rootpages: Vec<i64> = vec![0; n];
let sqls: Vec<Option<&str>> = vec![None; n];
let batch = RecordBatch::try_new(
schema,
vec![
Arc::new(StringArray::from(types)) as ArrayRef,
Arc::new(StringArray::from(names)) as ArrayRef,
Arc::new(StringArray::from(tbl_names)) as ArrayRef,
Arc::new(Int64Array::from(rootpages)) as ArrayRef,
Arc::new(StringArray::from(sqls)) as ArrayRef,
],
)
.map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
Ok(Some(vec![batch]))
}
fn try_direct_dispatch(&self, sql: &str) -> Result<Option<Vec<RecordBatch>>> {
use arrow::array::ArrayRef;
use mongreldb_core::Condition;
use sqlparser::ast::{Expr, Query, SelectItem, SetExpr, Statement, TableFactor};
use sqlparser::dialect::PostgreSqlDialect;
use sqlparser::parser::Parser;
let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
return Ok(None);
};
if stmts.len() != 1 {
return Ok(None);
}
let Statement::Query(query) = stmts.into_iter().next().unwrap() else {
return Ok(None);
};
let Query { body, .. } = *query;
let select = match *body {
SetExpr::Select(s) => *s,
_ => return Ok(None),
};
let lower_sql = sql.to_lowercase();
if lower_sql.contains(" limit ") || lower_sql.contains(" offset ") {
return Ok(None);
}
use sqlparser::ast::GroupByExpr;
if select.distinct.is_some()
|| !matches!(&select.group_by, GroupByExpr::Expressions(e, _) if e.is_empty())
|| select.having.is_some()
|| select.from.len() != 1
|| !select.from[0].joins.is_empty()
{
return Ok(None);
}
let table_name = match &select.from[0].relation {
TableFactor::Table { name, .. } => Some(name.to_string()),
_ => return Ok(None),
};
let Some(table_name) = table_name else {
return Ok(None);
};
if select.selection.is_none() {
return Ok(None);
}
let mut proj_names: Option<Vec<String>> = None;
for item in &select.projection {
match item {
SelectItem::Wildcard(_) => {}
SelectItem::UnnamedExpr(Expr::Identifier(ident)) => {
proj_names
.get_or_insert_with(Vec::new)
.push(ident.value.clone());
}
SelectItem::UnnamedExpr(Expr::CompoundIdentifier(idents)) => {
if let Some(last) = idents.last() {
proj_names
.get_or_insert_with(Vec::new)
.push(last.value.clone());
}
}
_ => return Ok(None),
}
}
let handle = match self.tables.lock().get(&table_name).cloned() {
Some(h) => h,
None => return Ok(None),
};
let mut db = handle.lock();
let schema = db.schema().clone();
let conditions: Vec<Condition> = match &select.selection {
Some(expr) => match translate_sqlparser_filter(expr, &schema) {
Some(c) => c,
None => return Ok(None),
},
None => Vec::new(),
};
if !conditions.is_empty() && db.ensure_indexes_complete().is_err() {
return Ok(None);
}
let snap = db.snapshot();
let mut col_ids: Vec<u16> = Vec::new();
let mut fields: Vec<arrow::datatypes::Field> = Vec::new();
let resolve_col = |name: &str| -> Option<&mongreldb_core::schema::ColumnDef> {
schema.columns.iter().find(|c| c.name == name)
};
match &proj_names {
None => {
for c in &schema.columns {
col_ids.push(c.id);
fields.push(arrow::datatypes::Field::new(
&c.name,
arrow_conv::arrow_data_type(&c.ty)?,
c.flags.contains(mongreldb_core::ColumnFlags::NULLABLE),
));
}
}
Some(names) => {
for n in names {
let cdef = match resolve_col(n) {
Some(c) => c,
None => return Ok(None), };
col_ids.push(cdef.id);
fields.push(arrow::datatypes::Field::new(
&cdef.name,
arrow_conv::arrow_data_type(&cdef.ty)?,
cdef.flags.contains(mongreldb_core::ColumnFlags::NULLABLE),
));
}
}
}
let cols = if !conditions.is_empty() {
match db.query_columns_native_cached(&conditions, Some(&col_ids), snap) {
Ok(Some(c)) => c,
Ok(None) => db
.visible_columns_native(snap, Some(&col_ids))
.map_err(MongrelQueryError::Core)?,
Err(_) => return Ok(None),
}
} else {
db.visible_columns_native(snap, Some(&col_ids))
.map_err(MongrelQueryError::Core)?
};
drop(db);
let mut arrays: Vec<ArrayRef> = Vec::with_capacity(col_ids.len());
for cid in &col_ids {
let col = cols
.iter()
.find(|(id, _)| id == cid)
.map(|(_, c)| c.clone());
let Some(col) = col else { return Ok(None) };
let ty = schema
.columns
.iter()
.find(|c| c.id == *cid)
.map(|c| c.ty)
.unwrap_or(mongreldb_core::schema::TypeId::Int64);
arrays.push(arrow_conv::native_to_array(ty, &col)?);
}
let batch_schema = Arc::new(arrow::datatypes::Schema::new(fields));
let batch = RecordBatch::try_new(batch_schema, arrays)
.map_err(|e| MongrelQueryError::Arrow(format!("direct dispatch batch build: {e}")))?;
mongreldb_core::trace::QueryTrace::record(|t| {
t.scan_mode = mongreldb_core::trace::ScanMode::DirectDispatch;
t.planning_nanos = 0; });
Ok(Some(vec![batch]))
}
pub async fn run(&self, sql: &str) -> Result<Vec<RecordBatch>> {
if let Some(inner) = strip_explain_query_plan(sql) {
return self.explain_query_plan(inner).await;
}
if let Some(batches) = commands::try_run_command(self, sql).await? {
return Ok(batches);
}
let lower = sql.trim_start().to_lowercase();
if lower.starts_with("create table") {
if let Some(db) = &self.database {
let (name, schema) = parse_create_table(sql)?;
db.create_table(&name, schema)?;
let handle = db.table(&name)?;
let provider = MongrelProvider::new(handle.clone())?;
self.ctx
.register_table(&name, Arc::new(provider))
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
self.tables.lock().insert(name, handle);
self.clear_cache();
return Ok(Vec::new());
}
}
if lower.starts_with("drop table") {
if let Some(db) = &self.database {
let (name, if_exists) = parse_drop_table(sql)?;
let drop_result = db.drop_table(&name);
if let Err(e) = drop_result {
let is_not_found = matches!(e, mongreldb_core::MongrelError::NotFound(_));
if !(if_exists && is_not_found) {
return Err(e.into());
}
} else {
self.ctx
.deregister_table(&name)
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
self.tables.lock().remove(&name);
}
self.clear_cache();
return Ok(Vec::new());
}
}
if lower.starts_with("alter table") {
if let Some(db) = &self.database {
match parse_alter_table(sql)? {
ParsedAlterTable::RenameTable { old_name, new_name } => {
db.rename_table(&old_name, &new_name)?;
self.ctx
.deregister_table(&old_name)
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
self.tables.lock().remove(&old_name);
let handle = db.table(&new_name)?;
let provider = MongrelProvider::new(handle.clone())?;
self.ctx
.register_table(&new_name, Arc::new(provider))
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
self.tables.lock().insert(new_name, handle);
}
ParsedAlterTable::RenameColumn {
table_name,
column_name,
new_name,
} => {
db.alter_column(&table_name, &column_name, AlterColumn::rename(new_name))?;
self.refresh_registered_table(db, &table_name)?;
}
ParsedAlterTable::AlterColumnType {
table_name,
column_name,
ty,
} => {
db.alter_column(&table_name, &column_name, AlterColumn::set_type(ty))?;
self.refresh_registered_table(db, &table_name)?;
}
ParsedAlterTable::SetNotNull {
table_name,
column_name,
} => {
let flags = current_column_flags(db, &table_name, &column_name)?
.without(ColumnFlags::NULLABLE);
db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
self.refresh_registered_table(db, &table_name)?;
}
ParsedAlterTable::DropNotNull {
table_name,
column_name,
} => {
let flags = current_column_flags(db, &table_name, &column_name)?
.with(ColumnFlags::NULLABLE);
db.alter_column(&table_name, &column_name, AlterColumn::set_flags(flags))?;
self.refresh_registered_table(db, &table_name)?;
}
}
self.clear_cache();
return Ok(Vec::new());
}
}
let resolved = self.resolve_view_sql(sql);
let resolved = self.rewrite_external_module_compat_sql(&resolved);
let resolved = rewrite_compat_function_calls(&resolved);
let effective_sql = normalize_sql(&resolved);
let sql = effective_sql.as_str();
let epoch = self.cache_epoch();
let key = (sql.to_string(), epoch);
let result_cacheable = !extended_sql_functions::contains_volatile_extended_function(sql);
if result_cacheable {
if let Some(hit) = self.cache.lock().get(&key) {
return Ok((**hit).clone());
}
}
if let Some(batches) = self.try_sqlite_master(sql)? {
if result_cacheable {
self.cache.lock().insert(key, Arc::new(batches.clone()));
}
return Ok(batches);
}
if let Some(batches) = self.try_direct_dispatch(sql)? {
if result_cacheable {
self.cache.lock().insert(key, Arc::new(batches.clone()));
}
return Ok(batches);
}
let plan_start = std::time::Instant::now();
let external_module_scan = self.query_references_external_module(sql);
let df = {
let cached_plan = self.plan_cache.lock().get(sql).cloned();
if let Some(plan) = cached_plan {
datafusion::dataframe::DataFrame::new(self.ctx.state(), plan)
} else {
let df = self
.ctx
.sql(sql)
.await
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
self.plan_cache
.lock()
.insert(sql.to_string(), df.logical_plan().clone());
df
}
};
let planning_nanos = plan_start.elapsed().as_nanos() as u64;
mongreldb_core::trace::QueryTrace::record(|t| t.planning_nanos = planning_nanos);
let agg_key = sql_cache_key(sql);
let batches = match self.try_native_aggregate(df.logical_plan(), agg_key) {
Ok(Some(batch)) => vec![batch],
_ => {
match self.try_fk_join(df.logical_plan()) {
Ok(Some(b)) => {
mongreldb_core::trace::QueryTrace::record(|t| {
t.join_mode = mongreldb_core::trace::JoinMode::FkBitmap;
});
b
}
_ => df
.collect()
.await
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?,
}
}
};
if external_module_scan {
mongreldb_core::trace::QueryTrace::record(|t| {
t.scan_mode = mongreldb_core::trace::ScanMode::ExternalModule;
});
}
if result_cacheable {
self.cache.lock().insert(key, Arc::new(batches.clone()));
}
Ok(batches)
}
pub async fn run_sql_traced(
&self,
sql: &str,
) -> Result<(Vec<RecordBatch>, mongreldb_core::trace::QueryTrace)> {
mongreldb_core::trace::QueryTrace::push_scope();
let result = self.run(sql).await;
let trace = mongreldb_core::trace::QueryTrace::pop_scope();
Ok((result?, trace))
}
pub fn clear_cache(&self) {
self.cache.lock().clear();
self.plan_cache.lock().clear();
}
async fn explain_query_plan(&self, sql: &str) -> Result<Vec<RecordBatch>> {
let explain_sql = format!("EXPLAIN {}", sql.trim().trim_end_matches(';'));
let batches = self
.ctx
.sql(&explain_sql)
.await
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?
.collect()
.await
.map_err(|e| MongrelQueryError::DataFusion(e.to_string()))?;
let mut detail = self.mongrel_query_plan_details(sql);
for batch in &batches {
if batch.num_columns() < 2 {
continue;
}
let Some(plan_type) = batch.column(0).as_any().downcast_ref::<StringArray>() else {
continue;
};
let Some(plan) = batch.column(1).as_any().downcast_ref::<StringArray>() else {
continue;
};
for row in 0..batch.num_rows() {
let prefix = plan_type.value(row);
for line in plan.value(row).lines() {
let line = line.trim();
if !line.is_empty() {
detail.push(format!("DATAFUSION {prefix}: {line}"));
}
}
}
}
if detail.is_empty() {
detail.push("plan unavailable".to_string());
}
let ids = (0..detail.len()).map(|i| i as i64).collect::<Vec<_>>();
let parents = vec![0_i64; detail.len()];
let notused = vec![0_i64; detail.len()];
let schema = Arc::new(arrow::datatypes::Schema::new(vec![
arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int64, false),
arrow::datatypes::Field::new("parent", arrow::datatypes::DataType::Int64, false),
arrow::datatypes::Field::new("notused", arrow::datatypes::DataType::Int64, false),
arrow::datatypes::Field::new("detail", arrow::datatypes::DataType::Utf8, false),
]));
let batch = RecordBatch::try_new(
schema,
vec![
Arc::new(Int64Array::from(ids)) as ArrayRef,
Arc::new(Int64Array::from(parents)),
Arc::new(Int64Array::from(notused)),
Arc::new(StringArray::from(detail)),
],
)
.map_err(|e| MongrelQueryError::Arrow(e.to_string()))?;
Ok(vec![batch])
}
fn mongrel_query_plan_details(&self, sql: &str) -> Vec<String> {
use sqlparser::ast::{GroupByExpr, OrderByKind, SetExpr, Statement};
use sqlparser::dialect::PostgreSqlDialect;
use sqlparser::parser::Parser;
let Ok(stmts) = Parser::parse_sql(&PostgreSqlDialect {}, sql) else {
return Vec::new();
};
let Some(Statement::Query(query)) = stmts.first() else {
return Vec::new();
};
fn collect(session: &MongrelSession, query: &sqlparser::ast::Query, out: &mut Vec<String>) {
use sqlparser::ast::{SetOperator, TableWithJoins};
match query.body.as_ref() {
SetExpr::Select(select) => {
for TableWithJoins { relation, joins } in &select.from {
session.push_table_plan(relation, select.selection.as_ref(), out);
for join in joins {
session.push_table_plan(&join.relation, None, out);
}
}
if select.distinct.is_some() {
out.push("USE TEMP B-TREE FOR DISTINCT".to_string());
}
let grouped = match &select.group_by {
GroupByExpr::All(_) => true,
GroupByExpr::Expressions(exprs, _) => !exprs.is_empty(),
};
if grouped {
out.push("USE TEMP B-TREE FOR GROUP BY".to_string());
}
let ordered = query.order_by.as_ref().is_some_and(|order_by| {
matches!(order_by.kind, OrderByKind::All(_))
|| matches!(&order_by.kind, OrderByKind::Expressions(exprs) if !exprs.is_empty())
});
if ordered {
out.push("USE TEMP B-TREE FOR ORDER BY".to_string());
}
}
SetExpr::Query(query) => collect(session, query, out),
SetExpr::SetOperation {
left, op, right, ..
} => {
let label = match op {
SetOperator::Union => "COMPOUND QUERY UNION",
SetOperator::Except => "COMPOUND QUERY EXCEPT",
SetOperator::Intersect => "COMPOUND QUERY INTERSECT",
_ => "COMPOUND QUERY",
};
out.push(label.to_string());
collect_set_expr(session, left, out);
collect_set_expr(session, right, out);
}
_ => {}
}
}
fn collect_set_expr(session: &MongrelSession, expr: &SetExpr, out: &mut Vec<String>) {
match expr {
SetExpr::Select(select) => {
for table in &select.from {
session.push_table_plan(&table.relation, select.selection.as_ref(), out);
}
}
SetExpr::Query(query) => collect(session, query, out),
SetExpr::SetOperation { left, right, .. } => {
collect_set_expr(session, left, out);
collect_set_expr(session, right, out);
}
_ => {}
}
}
let mut out = Vec::new();
collect(self, query, &mut out);
out
}
fn push_table_plan(
&self,
relation: &sqlparser::ast::TableFactor,
selection: Option<&sqlparser::ast::Expr>,
out: &mut Vec<String>,
) {
let sqlparser::ast::TableFactor::Table { name, alias, .. } = relation else {
out.push("SCAN SUBQUERY".to_string());
return;
};
let table_name = name.to_string();
let display_name = alias
.as_ref()
.map(|alias| alias.name.value.clone())
.unwrap_or_else(|| table_name.clone());
let Some(handle) = self.tables.lock().get(&table_name).cloned() else {
out.push(format!("SCAN {display_name}"));
return;
};
let schema = handle.lock().schema().clone();
let searchable = selection
.and_then(|expr| translate_sqlparser_filter(expr, &schema))
.is_some_and(|conditions| !conditions.is_empty());
if searchable {
out.push(format!("SEARCH {display_name} USING MONGREL INDEX"));
} else {
out.push(format!("SCAN {display_name}"));
}
}
fn resolve_view_sql(&self, sql: &str) -> String {
let views = self.views.lock();
if views.is_empty() {
return sql.to_string();
}
let mut result = sql.to_string();
for (name, view) in views.iter() {
result = replace_from_view(&result, name, &view.sql);
}
result
}
fn rewrite_external_module_compat_sql(&self, sql: &str) -> String {
let Some(db) = &self.database else {
return sql.to_string();
};
rewrite_fts_match_compat_sql(sql, db)
}
fn query_references_external_module(&self, sql: &str) -> bool {
use sqlparser::ast::Statement;
use sqlparser::dialect::PostgreSqlDialect;
use sqlparser::parser::Parser;
Parser::parse_sql(&PostgreSqlDialect {}, sql)
.ok()
.is_some_and(|statements| {
statements.iter().any(|statement| match statement {
Statement::Query(query) => self.query_uses_external_module(query),
_ => false,
})
})
}
fn query_uses_external_module(&self, query: &sqlparser::ast::Query) -> bool {
self.set_expr_uses_external_module(query.body.as_ref())
}
fn set_expr_uses_external_module(&self, expr: &sqlparser::ast::SetExpr) -> bool {
use sqlparser::ast::SetExpr;
match expr {
SetExpr::Select(select) => select
.from
.iter()
.any(|table| self.table_with_joins_uses_external_module(table)),
SetExpr::Query(query) => self.query_uses_external_module(query),
SetExpr::SetOperation { left, right, .. } => {
self.set_expr_uses_external_module(left)
|| self.set_expr_uses_external_module(right)
}
_ => false,
}
}
fn table_with_joins_uses_external_module(
&self,
table: &sqlparser::ast::TableWithJoins,
) -> bool {
self.table_factor_uses_external_module(&table.relation)
|| table
.joins
.iter()
.any(|join| self.table_factor_uses_external_module(&join.relation))
}
fn table_factor_uses_external_module(&self, relation: &sqlparser::ast::TableFactor) -> bool {
use sqlparser::ast::{Expr, TableFactor};
match relation {
TableFactor::Table { name, args, .. } => {
let table_name = name.to_string();
self.database
.as_ref()
.is_some_and(|db| db.external_table(&table_name).is_some())
|| (args.is_some() && self.external_modules.contains(&table_name))
}
TableFactor::Function { name, .. } => self.external_modules.contains(&name.to_string()),
TableFactor::TableFunction {
expr: Expr::Function(func),
..
} => self.external_modules.contains(&func.name.to_string()),
TableFactor::Derived { subquery, .. } => self.query_uses_external_module(subquery),
_ => false,
}
}
fn cache_epoch(&self) -> u64 {
if let Some(db) = &self.database {
db.visible_epoch().0
} else {
self.combined_epoch()
}
}
fn combined_epoch(&self) -> u64 {
let primary = self.db.as_ref().expect("no primary table");
let mut combined = primary.lock().snapshot().epoch.0;
let tables = self.tables.lock();
for arc in tables.values() {
if !Arc::ptr_eq(arc, primary) {
let e = arc.lock().snapshot().epoch.0;
combined = combined.wrapping_mul(31).wrapping_add(e);
}
}
combined
}
fn try_native_aggregate(
&self,
plan: &datafusion::logical_expr::LogicalPlan,
cache_key: u64,
) -> Result<Option<RecordBatch>> {
let Some(primary) = self.db.as_ref() else {
return Ok(None);
};
let mut db = primary.lock();
let schema = db.schema().clone();
let snap = db.snapshot();
native_agg::try_native_aggregate(&mut db, &schema, snap, plan, cache_key)
}
fn try_fk_join(
&self,
plan: &datafusion::logical_expr::LogicalPlan,
) -> Result<Option<Vec<RecordBatch>>> {
let tables = self.tables.lock();
fk_join::try_fk_join(&tables, plan)
}
pub fn context(&self) -> &SessionContext {
&self.ctx
}
pub fn register_scalar_udf(&self, f: ScalarUDF) {
self.ctx.register_udf(f);
self.clear_cache();
}
pub fn register_aggregate_udf(&self, f: AggregateUDF) {
self.ctx.register_udaf(f);
self.clear_cache();
}
pub fn register_window_udf(&self, f: WindowUDF) {
self.ctx.register_udwf(f);
self.clear_cache();
}
}
fn register_mongrel_functions(
ctx: &SessionContext,
sql_fn_state: Arc<extended_sql_functions::ExtendedSqlState>,
) {
ctx.register_udf(ScalarUDF::from(udf::AnnSearchUdf::new()));
ctx.register_udf(ScalarUDF::from(udf::SparseMatchUdf::new()));
ctx.register_udf(ScalarUDF::from(udf::RTreeIntersectsUdf::new()));
for udaf in percentile::percentile_udafs() {
ctx.register_udaf(udaf);
}
extended_sql_functions::register_extended_sql_functions_with_state(ctx, sql_fn_state);
}
fn strip_explain_query_plan(sql: &str) -> Option<&str> {
let trimmed = sql.trim_start();
let lower = trimmed.to_ascii_lowercase();
if !lower.starts_with("explain") {
return None;
}
let after_explain = trimmed.get(7..)?.trim_start();
let after_explain_lower = after_explain.to_ascii_lowercase();
if !after_explain_lower.starts_with("query") {
return None;
}
let after_query = after_explain.get(5..)?.trim_start();
let after_query_lower = after_query.to_ascii_lowercase();
if !after_query_lower.starts_with("plan") {
return None;
}
Some(after_query.get(4..)?.trim_start())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SqlCompatTokenKind {
Ident,
String,
Dot,
LParen,
RParen,
Comma,
}
#[derive(Debug, Clone)]
struct SqlCompatToken {
kind: SqlCompatTokenKind,
raw: String,
normalized: String,
start: usize,
end: usize,
}
#[derive(Debug, Clone)]
struct FtsMatchBinding {
query_ref: String,
}
#[derive(Debug, Clone)]
struct SqlReplacement {
start: usize,
end: usize,
replacement: String,
}
fn rewrite_fts_match_compat_sql(sql: &str, db: &Database) -> String {
let tokens = sql_compat_tokens(sql);
if tokens.is_empty() {
return sql.to_string();
}
let bindings = fts_match_bindings(sql, db, &tokens);
if bindings.is_empty() {
return sql.to_string();
}
let unique_refs = bindings
.values()
.map(|binding| binding.query_ref.as_str())
.collect::<HashSet<_>>();
let unique_binding = if unique_refs.len() == 1 {
bindings.values().next().cloned()
} else {
None
};
let mut replacements = Vec::new();
for (idx, token) in tokens.iter().enumerate() {
if token.kind != SqlCompatTokenKind::Ident || token.normalized != "match" {
continue;
}
let Some(rhs) = tokens.get(idx + 1) else {
continue;
};
if rhs.kind != SqlCompatTokenKind::String {
continue;
}
let Some((lhs_start, _lhs_end, query_ref)) =
fts_match_lhs_query_ref(&tokens, idx, &bindings, unique_binding.as_ref())
else {
continue;
};
replacements.push(SqlReplacement {
start: lhs_start,
end: rhs.end,
replacement: format!("{query_ref}.query = {}", rhs.raw),
});
}
apply_sql_replacements(sql, &replacements)
}
fn fts_match_lhs_query_ref(
tokens: &[SqlCompatToken],
match_idx: usize,
bindings: &HashMap<String, FtsMatchBinding>,
unique_binding: Option<&FtsMatchBinding>,
) -> Option<(usize, usize, String)> {
if match_idx == 0 {
return None;
}
let lhs = tokens.get(match_idx - 1)?;
if lhs.kind != SqlCompatTokenKind::Ident {
return None;
}
if match_idx >= 3
&& tokens.get(match_idx - 2)?.kind == SqlCompatTokenKind::Dot
&& tokens.get(match_idx - 3)?.kind == SqlCompatTokenKind::Ident
{
let owner = tokens.get(match_idx - 3)?;
let binding = bindings.get(&owner.normalized)?;
if lhs.normalized == "query" || lhs.normalized == "text" {
return Some((owner.start, lhs.end, binding.query_ref.clone()));
}
return None;
}
if let Some(binding) = bindings.get(&lhs.normalized) {
return Some((lhs.start, lhs.end, binding.query_ref.clone()));
}
if lhs.normalized == "text" {
let binding = unique_binding?;
return Some((lhs.start, lhs.end, binding.query_ref.clone()));
}
None
}
fn fts_match_bindings(
sql: &str,
db: &Database,
tokens: &[SqlCompatToken],
) -> HashMap<String, FtsMatchBinding> {
let mut out = HashMap::new();
let mut i = 0;
while i < tokens.len() {
let token = &tokens[i];
let starts_table_ref = token.kind == SqlCompatTokenKind::Ident
&& matches!(token.normalized.as_str(), "from" | "join");
if !starts_table_ref {
i += 1;
continue;
}
let mut table_idx = i + 1;
if tokens
.get(table_idx)
.is_some_and(|token| token.kind == SqlCompatTokenKind::LParen)
{
i += 1;
continue;
}
let Some(table) = tokens.get(table_idx) else {
break;
};
if table.kind != SqlCompatTokenKind::Ident {
i += 1;
continue;
}
let mut table_name = table.normalized.clone();
let mut table_ref = table.raw.clone();
if tokens
.get(table_idx + 1)
.is_some_and(|token| token.kind == SqlCompatTokenKind::Dot)
&& tokens
.get(table_idx + 2)
.is_some_and(|token| token.kind == SqlCompatTokenKind::Ident)
{
let qualified = tokens.get(table_idx + 2).unwrap();
table_name = qualified.normalized.clone();
table_ref = sql[table.start..qualified.end].to_string();
table_idx += 2;
}
if !is_fts_docs_table(db, &table_name) {
i = table_idx + 1;
continue;
}
let mut query_ref = table_ref.clone();
let mut alias_key = None;
let mut next = table_idx + 1;
if tokens.get(next).is_some_and(|token| {
token.kind == SqlCompatTokenKind::Ident && token.normalized == "as"
}) {
next += 1;
}
if let Some(alias) = tokens.get(next) {
if alias.kind == SqlCompatTokenKind::Ident && !is_table_ref_boundary(&alias.normalized)
{
alias_key = Some(alias.normalized.clone());
query_ref = alias.raw.clone();
next += 1;
}
}
out.insert(
table_name,
FtsMatchBinding {
query_ref: query_ref.clone(),
},
);
if let Some(alias_key) = alias_key {
out.insert(alias_key, FtsMatchBinding { query_ref });
}
i = next;
}
out
}
fn is_fts_docs_table(db: &Database, name: &str) -> bool {
db.external_table(name)
.is_some_and(|entry| entry.module == "fts_docs")
}
fn is_table_ref_boundary(normalized: &str) -> bool {
matches!(
normalized,
"where"
| "join"
| "left"
| "right"
| "inner"
| "outer"
| "full"
| "cross"
| "on"
| "using"
| "group"
| "order"
| "having"
| "limit"
| "offset"
| "union"
| "except"
| "intersect"
)
}
fn sql_compat_tokens(sql: &str) -> Vec<SqlCompatToken> {
let bytes = sql.as_bytes();
let mut tokens = Vec::new();
let mut i = 0;
while i < bytes.len() {
match bytes[i] {
b if b.is_ascii_whitespace() => i += 1,
b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
i += 2;
while i < bytes.len() && bytes[i] != b'\n' {
i += 1;
}
}
b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
i = skip_block_comment(bytes, i);
}
b'\'' => {
let end = skip_quoted(bytes, i, b'\'');
tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
i = end;
}
b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
let end = skip_quoted(bytes, i + 1, b'\'');
tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
i = end;
}
b'$' => {
let (end, matched) = skip_dollar_quoted(bytes, i);
if matched {
tokens.push(sql_token(sql, SqlCompatTokenKind::String, i, end));
i = end;
} else {
i += 1;
}
}
b'"' => {
let end = skip_quoted(bytes, i, b'"');
let raw = sql[i..end].to_string();
let normalized = unquote_sql_ident(&raw).to_ascii_lowercase();
tokens.push(SqlCompatToken {
kind: SqlCompatTokenKind::Ident,
raw,
normalized,
start: i,
end,
});
i = end;
}
b'.' => {
tokens.push(sql_token(sql, SqlCompatTokenKind::Dot, i, i + 1));
i += 1;
}
b'(' => {
tokens.push(sql_token(sql, SqlCompatTokenKind::LParen, i, i + 1));
i += 1;
}
b')' => {
tokens.push(sql_token(sql, SqlCompatTokenKind::RParen, i, i + 1));
i += 1;
}
b',' => {
tokens.push(sql_token(sql, SqlCompatTokenKind::Comma, i, i + 1));
i += 1;
}
b if is_sql_ident_byte(b) => {
let start = i;
i += 1;
while i < bytes.len() && is_sql_ident_byte(bytes[i]) {
i += 1;
}
tokens.push(sql_token(sql, SqlCompatTokenKind::Ident, start, i));
}
_ => i += 1,
}
}
tokens
}
fn sql_token(sql: &str, kind: SqlCompatTokenKind, start: usize, end: usize) -> SqlCompatToken {
let raw = sql[start..end].to_string();
SqlCompatToken {
kind,
normalized: raw.to_ascii_lowercase(),
raw,
start,
end,
}
}
fn unquote_sql_ident(raw: &str) -> String {
if raw.len() >= 2 && raw.starts_with('"') && raw.ends_with('"') {
raw[1..raw.len() - 1].replace("\"\"", "\"")
} else {
raw.to_string()
}
}
fn apply_sql_replacements(sql: &str, replacements: &[SqlReplacement]) -> String {
if replacements.is_empty() {
return sql.to_string();
}
let mut ordered = replacements.to_vec();
ordered.sort_by_key(|replacement| replacement.start);
let mut out = String::with_capacity(sql.len());
let mut cursor = 0;
for replacement in ordered {
if replacement.start < cursor || replacement.end > sql.len() {
continue;
}
out.push_str(&sql[cursor..replacement.start]);
out.push_str(&replacement.replacement);
cursor = replacement.end;
}
out.push_str(&sql[cursor..]);
out
}
fn rewrite_compat_function_calls(sql: &str) -> String {
let bytes = sql.as_bytes();
let mut out = String::with_capacity(sql.len());
let mut i = 0;
while i < bytes.len() {
match bytes[i] {
b'\'' => i = copy_quoted_to_string(&mut out, bytes, i, b'\''),
b'"' => i = copy_quoted_to_string(&mut out, bytes, i, b'"'),
b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
out.push(bytes[i] as char);
i += 1;
i = copy_quoted_to_string(&mut out, bytes, i, b'\'');
}
b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
out.push('-');
out.push('-');
i += 2;
while i < bytes.len() {
let ch = bytes[i] as char;
out.push(ch);
i += 1;
if ch == '\n' {
break;
}
}
}
b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
let start = i;
i = skip_block_comment(bytes, i);
out.push_str(&sql[start..i.min(bytes.len())]);
}
b'$' => {
let start_len = out.len();
let (next, matched) = copy_dollar_quoted_to_string(&mut out, bytes, i);
if matched {
i = next;
} else {
out.truncate(start_len);
out.push('$');
i += 1;
}
}
b'g' | b'G' | b'm' | b'M' | b't' | b'T' => {
if let Some((replacement, next)) = compat_function_rewrite_at(sql, i) {
out.push_str(&replacement);
i = next;
} else {
out.push(bytes[i] as char);
i += 1;
}
}
_ => {
out.push(bytes[i] as char);
i += 1;
}
}
}
out
}
fn compat_function_rewrite_at(sql: &str, start: usize) -> Option<(String, usize)> {
let bytes = sql.as_bytes();
let (name, kind) = if ident_eq_at(bytes, start, b"max") {
("max", CompatRewriteKind::ScalarMax)
} else if ident_eq_at(bytes, start, b"min") {
("min", CompatRewriteKind::ScalarMin)
} else if ident_eq_at(bytes, start, b"group_concat") {
("group_concat", CompatRewriteKind::GroupConcat)
} else if ident_eq_at(bytes, start, b"total") {
("total", CompatRewriteKind::Total)
} else {
return None;
};
let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
let after_name = start + name.len();
let after_ok = bytes
.get(after_name)
.is_some_and(|b| !is_sql_ident_byte(*b));
if !before_ok || !after_ok {
return None;
}
let mut open = after_name;
while open < bytes.len() && bytes[open].is_ascii_whitespace() {
open += 1;
}
if bytes.get(open) != Some(&b'(') {
return None;
}
let summary = call_arg_summary(sql, open)?;
match kind {
CompatRewriteKind::ScalarMax if summary.top_level_commas > 0 => {
Some(("__mongreldb_scalar_max(".to_string(), open + 1))
}
CompatRewriteKind::ScalarMin if summary.top_level_commas > 0 => {
Some(("__mongreldb_scalar_min(".to_string(), open + 1))
}
CompatRewriteKind::GroupConcat => {
let args = &sql[open + 1..summary.close];
let rewritten = if summary.top_level_commas == 0 {
format!("string_agg({args}, ',')")
} else {
format!("string_agg({args})")
};
Some((rewritten, summary.close + 1))
}
CompatRewriteKind::Total if summary.top_level_commas == 0 => {
let args = &sql[open + 1..summary.close];
let suffix_end = aggregate_suffix_end(sql, summary.close + 1);
let suffix = &sql[summary.close + 1..suffix_end];
Some((
format!("coalesce(cast(sum({args}){suffix} as double), 0.0)"),
suffix_end,
))
}
_ => None,
}
}
#[derive(Clone, Copy)]
enum CompatRewriteKind {
ScalarMax,
ScalarMin,
GroupConcat,
Total,
}
fn ident_eq_at(bytes: &[u8], start: usize, ident: &[u8]) -> bool {
bytes
.get(start..start + ident.len())
.is_some_and(|slice| slice.eq_ignore_ascii_case(ident))
}
fn is_sql_ident_byte(b: u8) -> bool {
b.is_ascii_alphanumeric() || b == b'_' || b == b'$'
}
fn keyword_at(bytes: &[u8], start: usize, keyword: &[u8]) -> bool {
if !ident_eq_at(bytes, start, keyword) {
return false;
}
let before_ok = start == 0 || !is_sql_ident_byte(bytes[start - 1]);
let after = start + keyword.len();
let after_ok = after >= bytes.len() || !is_sql_ident_byte(bytes[after]);
before_ok && after_ok
}
fn skip_sql_whitespace(bytes: &[u8], mut i: usize) -> usize {
while i < bytes.len() && bytes[i].is_ascii_whitespace() {
i += 1;
}
i
}
fn aggregate_suffix_end(sql: &str, start: usize) -> usize {
let bytes = sql.as_bytes();
let mut suffix_end = start;
let mut i = skip_sql_whitespace(bytes, start);
if keyword_at(bytes, i, b"filter") {
let open = skip_sql_whitespace(bytes, i + b"filter".len());
if bytes.get(open) != Some(&b'(') {
return start;
}
let Some(summary) = call_arg_summary(sql, open) else {
return start;
};
suffix_end = summary.close + 1;
i = skip_sql_whitespace(bytes, suffix_end);
}
if keyword_at(bytes, i, b"over") {
let after_over = skip_sql_whitespace(bytes, i + b"over".len());
if bytes.get(after_over) == Some(&b'(') {
let Some(summary) = call_arg_summary(sql, after_over) else {
return suffix_end;
};
suffix_end = summary.close + 1;
} else {
let mut end = after_over;
while end < bytes.len() && is_sql_ident_byte(bytes[end]) {
end += 1;
}
if end > after_over {
suffix_end = end;
}
}
}
suffix_end
}
struct CallArgSummary {
close: usize,
top_level_commas: usize,
}
fn call_arg_summary(sql: &str, open: usize) -> Option<CallArgSummary> {
let bytes = sql.as_bytes();
let mut depth = 1;
let mut i = open + 1;
let mut top_level_commas = 0;
while i < bytes.len() {
match bytes[i] {
b'\'' => i = skip_quoted(bytes, i, b'\''),
b'"' => i = skip_quoted(bytes, i, b'"'),
b'E' | b'e' if i + 1 < bytes.len() && bytes[i + 1] == b'\'' => {
i = skip_quoted(bytes, i + 1, b'\'')
}
b'$' => {
let (next, matched) = skip_dollar_quoted(bytes, i);
i = if matched { next } else { i + 1 };
}
b'-' if i + 1 < bytes.len() && bytes[i + 1] == b'-' => {
i += 2;
while i < bytes.len() && bytes[i] != b'\n' {
i += 1;
}
}
b'/' if i + 1 < bytes.len() && bytes[i + 1] == b'*' => {
i = skip_block_comment(bytes, i);
}
b'(' => {
depth += 1;
i += 1;
}
b')' => {
depth -= 1;
if depth == 0 {
return Some(CallArgSummary {
close: i,
top_level_commas,
});
}
i += 1;
}
b',' if depth == 1 => {
top_level_commas += 1;
i += 1;
}
_ => i += 1,
}
}
None
}
fn copy_quoted_to_string(out: &mut String, bytes: &[u8], start: usize, delim: u8) -> usize {
let end = skip_quoted(bytes, start, delim);
out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
end
}
fn skip_quoted(bytes: &[u8], start: usize, delim: u8) -> usize {
let mut i = start;
if i < bytes.len() {
i += 1;
}
while i < bytes.len() {
if bytes[i] == delim {
i += 1;
if i < bytes.len() && bytes[i] == delim {
i += 1;
continue;
}
break;
}
i += 1;
}
i
}
fn copy_dollar_quoted_to_string(out: &mut String, bytes: &[u8], start: usize) -> (usize, bool) {
let (end, matched) = skip_dollar_quoted(bytes, start);
if matched {
out.push_str(std::str::from_utf8(&bytes[start..end]).unwrap_or_default());
}
(end, matched)
}
fn skip_dollar_quoted(bytes: &[u8], start: usize) -> (usize, bool) {
if bytes.get(start) != Some(&b'$') {
return (start, false);
}
let mut j = start + 1;
while j < bytes.len() && (bytes[j].is_ascii_alphanumeric() || bytes[j] == b'_') {
j += 1;
}
if bytes.get(j) != Some(&b'$') {
return (start, false);
}
let tag = &bytes[start..=j];
let mut i = j + 1;
while i + tag.len() <= bytes.len() {
if &bytes[i..i + tag.len()] == tag {
return (i + tag.len(), true);
}
i += 1;
}
(start, false)
}
fn skip_block_comment(bytes: &[u8], start: usize) -> usize {
let mut i = start + 2;
let mut depth = 1;
while i + 1 < bytes.len() && depth > 0 {
if bytes[i] == b'/' && bytes[i + 1] == b'*' {
depth += 1;
i += 2;
} else if bytes[i] == b'*' && bytes[i + 1] == b'/' {
depth -= 1;
i += 2;
} else {
i += 1;
}
}
i
}
fn sql_cache_key(sql: &str) -> u64 {
use std::hash::{Hash, Hasher};
let mut h = std::collections::hash_map::DefaultHasher::new();
sql.hash(&mut h);
h.finish()
}
fn replace_from_view(sql: &str, name: &str, view_sql: &str) -> String {
let lower = sql.to_ascii_lowercase();
let bytes = lower.as_bytes();
let name_b = name.as_bytes();
let mut i = 0usize;
while let Some(rel) = lower[i..].find("from") {
let from_start = i + rel;
let after_from = from_start + 4;
i = after_from;
if from_start > 0 && is_ident_byte(bytes[from_start - 1]) {
continue;
}
let mut j = after_from;
while j < bytes.len() && bytes[j].is_ascii_whitespace() {
j += 1;
}
if j == after_from || !bytes[j..].starts_with(name_b) {
continue;
}
let after_name = j + name_b.len();
if after_name < bytes.len() && is_ident_byte(bytes[after_name]) {
continue;
}
let mut out = String::with_capacity(sql.len() + view_sql.len() + name.len() + 8);
out.push_str(&sql[..from_start]);
out.push_str(&sql[from_start..j]);
out.push('(');
out.push_str(view_sql);
out.push_str(") AS ");
out.push_str(name);
out.push_str(&sql[after_name..]);
return out;
}
sql.to_string()
}
fn is_ident_byte(b: u8) -> bool {
b.is_ascii_alphanumeric() || b == b'_'
}
fn normalize_sql(sql: &str) -> String {
let b = sql.as_bytes();
let n = b.len();
let mut out: Vec<u8> = Vec::with_capacity(n);
let mut want_space = false;
let mut i = 0usize;
while i < n {
let c = b[i];
if c.is_ascii_whitespace() {
want_space = true;
i += 1;
continue;
}
if c == b'-' && i + 1 < n && b[i + 1] == b'-' {
i += 2;
while i < n && b[i] != b'\n' {
i += 1;
}
want_space = !out.is_empty();
continue;
}
if c == b'/' && i + 1 < n && b[i + 1] == b'*' {
i += 2;
let mut depth = 1usize;
while i + 1 < n && depth > 0 {
if b[i] == b'/' && b[i + 1] == b'*' {
depth += 1;
i += 2;
} else if b[i] == b'*' && b[i + 1] == b'/' {
depth -= 1;
i += 2;
} else {
i += 1;
}
}
want_space = !out.is_empty();
continue;
}
if want_space && !out.is_empty() {
out.push(b' ');
}
want_space = false;
match c {
b'E' | b'e' if i + 1 < n && b[i + 1] == b'\'' => {
out.push(c);
i += 1;
i = copy_quoted(&mut out, b, i, n, b'\'');
continue;
}
b'\'' => {
i = copy_quoted(&mut out, b, i, n, b'\'');
continue;
}
b'"' => {
i = copy_quoted(&mut out, b, i, n, b'"');
continue;
}
b'$' => {
let (consumed, matched) = copy_dollar_quoted(&mut out, b, i, n);
if matched {
i = consumed;
continue;
}
out.push(c);
i += 1;
continue;
}
_ => {
out.push(c);
i += 1;
}
}
}
String::from_utf8(out).unwrap_or_else(|_| sql.to_string())
}
fn copy_quoted(out: &mut Vec<u8>, b: &[u8], start: usize, n: usize, delim: u8) -> usize {
out.push(b[start]);
let mut i = start + 1;
while i < n {
let c = b[i];
out.push(c);
if c == delim {
if i + 1 < n && b[i + 1] == delim {
out.push(b[i + 1]);
i += 2;
continue;
}
return i + 1;
}
i += 1;
}
i
}
fn copy_dollar_quoted(out: &mut Vec<u8>, b: &[u8], start: usize, n: usize) -> (usize, bool) {
let mut j = start + 1;
let tag_start = j;
while j < n && b[j] != b'$' && is_dollar_tag_byte(b[j]) {
j += 1;
}
if j >= n || b[j] != b'$' {
return (start + 1, false);
}
if tag_start < j && !(b[tag_start].is_ascii_alphabetic() || b[tag_start] == b'_') {
return (start + 1, false);
}
let close_end = j + 1; let delim = &b[start..close_end];
out.extend_from_slice(delim);
let mut k = close_end;
while k + delim.len() <= n {
if &b[k..k + delim.len()] == delim {
out.extend_from_slice(delim);
return (k + delim.len(), true);
}
out.push(b[k]);
k += 1;
}
out.extend_from_slice(&b[close_end..n]);
(n, true)
}
fn is_dollar_tag_byte(b: u8) -> bool {
b.is_ascii_alphanumeric() || b == b'_'
}
fn strip_prefix_ci<'a>(s: &'a str, prefix: &str) -> Option<&'a str> {
let bytes = s.as_bytes();
let pb = prefix.as_bytes();
if bytes.len() >= pb.len() && bytes[..pb.len()].eq_ignore_ascii_case(pb) {
Some(&s[pb.len()..])
} else {
None
}
}
const COLUMN_CONSTRAINTS: &[(&str, u32)] = &[
("primary key", ColumnFlags::PRIMARY_KEY),
("autoincrement", ColumnFlags::AUTO_INCREMENT),
("auto_increment", ColumnFlags::AUTO_INCREMENT),
];
fn parse_column_constraints(constraint_text: &str) -> ColumnFlags {
let normalized = constraint_text.to_lowercase();
let mut flags = ColumnFlags::empty();
for (phrase, bit) in COLUMN_CONSTRAINTS {
if normalized.contains(phrase) {
flags = flags.with(*bit);
}
}
flags
}
fn parse_sql_type(ty_str: &str) -> Result<mongreldb_core::schema::TypeId> {
use mongreldb_core::schema::TypeId;
match ty_str.trim().trim_end_matches(';').to_lowercase().as_str() {
"bigint" | "int8" | "int64" | "integer" | "int" => Ok(TypeId::Int64),
"double" | "float8" | "float64" | "real" | "float" => Ok(TypeId::Float64),
"varchar" | "text" | "string" | "bytes" => Ok(TypeId::Bytes),
"boolean" | "bool" => Ok(TypeId::Bool),
other => Err(MongrelQueryError::Schema(format!(
"unsupported column type: {other}"
))),
}
}
fn parse_create_table(sql: &str) -> Result<(String, mongreldb_core::schema::Schema)> {
use mongreldb_core::schema::*;
let open = sql
.find('(')
.ok_or(MongrelQueryError::Schema("CREATE TABLE missing '('".into()))?;
let close = sql
.rfind(')')
.ok_or(MongrelQueryError::Schema("CREATE TABLE missing ')'".into()))?;
let head = sql[..open].trim();
let after_kw = strip_prefix_ci(head, "CREATE TABLE")
.or_else(|| strip_prefix_ci(head, "create table"))
.unwrap_or("")
.trim();
let after_kw = after_kw
.strip_prefix("IF NOT EXISTS")
.or_else(|| after_kw.strip_prefix("if not exists"))
.map(str::trim)
.unwrap_or(after_kw);
let name = after_kw.trim_matches('"').to_string();
if name.is_empty() {
return Err(MongrelQueryError::Schema(
"CREATE TABLE missing table name".into(),
));
}
let body = &sql[open + 1..close];
let mut columns = Vec::new();
let schema_id: u64 = 0; for (i, raw) in body.split(',').enumerate() {
let part = raw.trim();
if part.is_empty() {
continue;
}
let mut tokens = part.split_whitespace();
let col_name = tokens
.next()
.ok_or(MongrelQueryError::Schema("missing column name".into()))?
.trim_matches('"');
let ty_str = tokens
.next()
.ok_or(MongrelQueryError::Schema("missing column type".into()))?
.to_lowercase();
let ty = parse_sql_type(&ty_str)?;
let constraint_clause: String = tokens.collect::<Vec<_>>().join(" ");
let flags = parse_column_constraints(&constraint_clause);
columns.push(ColumnDef {
id: (i + 1) as u16,
name: col_name.to_string(),
ty,
flags,
});
}
Ok((
name,
Schema {
schema_id,
columns,
indexes: vec![],
colocation: vec![],
constraints: Default::default(),
},
))
}
fn parse_drop_table(sql: &str) -> Result<(String, bool)> {
let head = sql.trim();
let after_kw = strip_prefix_ci(head, "DROP TABLE")
.or_else(|| strip_prefix_ci(head, "drop table"))
.unwrap_or("")
.trim();
let (rest, if_exists) = if let Some(r) = after_kw
.strip_prefix("IF EXISTS")
.or_else(|| after_kw.strip_prefix("if exists"))
.map(str::trim)
{
(r, true)
} else {
(after_kw, false)
};
let name = rest.trim_matches(';').trim_matches('"').trim();
if name.is_empty() {
return Err(MongrelQueryError::Schema(
"DROP TABLE missing table name".into(),
));
}
Ok((name.to_string(), if_exists))
}
enum ParsedAlterTable {
RenameTable {
old_name: String,
new_name: String,
},
RenameColumn {
table_name: String,
column_name: String,
new_name: String,
},
AlterColumnType {
table_name: String,
column_name: String,
ty: mongreldb_core::schema::TypeId,
},
SetNotNull {
table_name: String,
column_name: String,
},
DropNotNull {
table_name: String,
column_name: String,
},
}
fn current_column_flags(db: &Arc<Database>, table: &str, column: &str) -> Result<ColumnFlags> {
let handle = db.table(table)?;
let table = handle.lock();
table
.schema()
.column(column)
.map(|c| c.flags)
.ok_or_else(|| MongrelQueryError::Schema(format!("unknown column {column}")))
}
fn parse_alter_table(sql: &str) -> Result<ParsedAlterTable> {
let trimmed = strip_statement_semicolon(sql.trim());
let after_kw = strip_prefix_ci(trimmed, "ALTER TABLE")
.ok_or_else(|| MongrelQueryError::Schema("not an ALTER TABLE statement".into()))?
.trim();
let (table_name, rest) = take_sql_ident(after_kw, "ALTER TABLE missing table name")?;
let rest = rest.trim();
if let Some(after) = strip_prefix_ci(rest, "RENAME TO") {
let new_name = parse_trailing_identifier(after, "ALTER TABLE missing new table name")?;
return Ok(ParsedAlterTable::RenameTable {
old_name: table_name,
new_name,
});
}
if let Some(after) = strip_prefix_ci(rest, "RENAME COLUMN") {
let (column_name, after_col) =
take_sql_ident(after, "ALTER TABLE RENAME COLUMN missing column name")?;
let after_to = strip_prefix_ci(after_col.trim(), "TO").ok_or_else(|| {
MongrelQueryError::Schema("ALTER TABLE RENAME COLUMN missing TO".into())
})?;
let new_name = parse_trailing_identifier(
after_to,
"ALTER TABLE RENAME COLUMN missing new column name",
)?;
return Ok(ParsedAlterTable::RenameColumn {
table_name,
column_name,
new_name,
});
}
let after_alter = strip_prefix_ci(rest, "ALTER COLUMN")
.or_else(|| strip_prefix_ci(rest, "ALTER"))
.ok_or_else(|| {
MongrelQueryError::Schema(
"ALTER TABLE must be RENAME TO, RENAME COLUMN, or ALTER COLUMN".into(),
)
})?;
let (column_name, action) =
take_sql_ident(after_alter, "ALTER TABLE ALTER COLUMN missing column name")?;
let action = action.trim();
if let Some(after_type) =
strip_prefix_ci(action, "TYPE").or_else(|| strip_prefix_ci(action, "SET DATA TYPE"))
{
let ty = parse_type_tail(after_type)?;
return Ok(ParsedAlterTable::AlterColumnType {
table_name,
column_name,
ty,
});
}
if strip_prefix_ci(action, "SET NOT NULL").is_some() {
return Ok(ParsedAlterTable::SetNotNull {
table_name,
column_name,
});
}
if strip_prefix_ci(action, "DROP NOT NULL").is_some() {
return Ok(ParsedAlterTable::DropNotNull {
table_name,
column_name,
});
}
Err(MongrelQueryError::Schema(
"unsupported ALTER COLUMN action".into(),
))
}
fn strip_statement_semicolon(s: &str) -> &str {
s.trim().trim_end_matches(';').trim()
}
fn take_sql_ident<'a>(s: &'a str, missing: &str) -> Result<(String, &'a str)> {
let s = s.trim();
if s.is_empty() {
return Err(MongrelQueryError::Schema(missing.into()));
}
if let Some(rest) = s.strip_prefix('"') {
let Some(end) = rest.find('"') else {
return Err(MongrelQueryError::Schema(
"unterminated quoted identifier".into(),
));
};
let ident = rest[..end].to_string();
if ident.is_empty() {
return Err(MongrelQueryError::Schema(missing.into()));
}
return Ok((ident, &rest[end + 1..]));
}
let end = s.find(|c: char| c.is_ascii_whitespace()).unwrap_or(s.len());
let ident = s[..end].trim_matches('"').to_string();
if ident.is_empty() {
return Err(MongrelQueryError::Schema(missing.into()));
}
Ok((ident, &s[end..]))
}
fn parse_trailing_identifier(s: &str, missing: &str) -> Result<String> {
let (ident, rest) = take_sql_ident(s, missing)?;
if !strip_statement_semicolon(rest).is_empty() {
return Err(MongrelQueryError::Schema(
"unexpected tokens after identifier".into(),
));
}
Ok(ident)
}
fn parse_type_tail(s: &str) -> Result<mongreldb_core::schema::TypeId> {
let tail = strip_statement_semicolon(s);
let ty = tail
.split_whitespace()
.next()
.ok_or_else(|| MongrelQueryError::Schema("ALTER COLUMN TYPE missing type".into()))?;
parse_sql_type(ty)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn normalize_collapses_and_trims_whitespace() {
assert_eq!(normalize_sql("SELECT * FROM t"), "SELECT * FROM t");
assert_eq!(normalize_sql(" SELECT * FROM t "), "SELECT * FROM t");
assert_eq!(
normalize_sql("\n\tSELECT\n*\nFROM\n\tt\n"),
"SELECT * FROM t"
);
assert_eq!(
normalize_sql("SELECT a, b FROM t"),
normalize_sql("SELECT a, b FROM t")
);
}
#[test]
fn normalize_preserves_string_literal_whitespace() {
assert_eq!(
normalize_sql("SELECT 'hello world' FROM t"),
"SELECT 'hello world' FROM t"
);
assert_eq!(
normalize_sql("SELECT 'it''s ok' FROM t"),
"SELECT 'it''s ok' FROM t"
);
assert_eq!(
normalize_sql(" SELECT 'a b' FROM t "),
"SELECT 'a b' FROM t"
);
}
#[test]
fn normalize_preserves_quoted_identifier_and_dollar_quote() {
assert_eq!(
normalize_sql(" SELECT \"my col\" FROM t "),
"SELECT \"my col\" FROM t"
);
assert_eq!(
normalize_sql(" SELECT $$a b$$ FROM t "),
"SELECT $$a b$$ FROM t"
);
assert_eq!(
normalize_sql("SELECT $tag$body with spaces$tag$ FROM t"),
"SELECT $tag$body with spaces$tag$ FROM t"
);
}
#[test]
fn normalize_strips_comments() {
assert_eq!(
normalize_sql("SELECT 1 -- trailing comment\nFROM t"),
"SELECT 1 FROM t"
);
assert_eq!(
normalize_sql("SELECT /* block */ 1 FROM t"),
"SELECT 1 FROM t"
);
assert_eq!(
normalize_sql("SELECT /* 'not a string' */ 1 FROM t"),
"SELECT 1 FROM t"
);
assert_eq!(
normalize_sql("SELECT /* outer /* inner */ still outer */ 1 FROM t"),
"SELECT 1 FROM t"
);
}
#[test]
fn normalize_escape_string_preserved() {
assert_eq!(
normalize_sql("SELECT E'line\\nbreak' FROM t"),
"SELECT E'line\\nbreak' FROM t"
);
}
#[test]
fn replace_from_view_matches_whole_word_only() {
let out = replace_from_view("SELECT * FROM logs", "log", "SELECT 1");
assert_eq!(out, "SELECT * FROM logs");
let out = replace_from_view("SELECT * FROM log", "log", "SELECT 1");
assert_eq!(out, "SELECT * FROM (SELECT 1) AS log");
let out = replace_from_view("select * from log where x", "log", "SELECT 1");
assert_eq!(out, "select * from (SELECT 1) AS log where x");
let out = replace_from_view("SELECT * FROM log)", "log", "SELECT 1");
assert_eq!(out, "SELECT * FROM (SELECT 1) AS log)");
let out = replace_from_view("SELECT * xfrom log", "log", "SELECT 1");
assert_eq!(out, "SELECT * xfrom log");
}
#[test]
fn compat_function_rewrite_handles_sqlite_compatibility_calls() {
assert_eq!(
rewrite_compat_function_calls("select max(id), min(id) from t"),
"select max(id), min(id) from t"
);
assert_eq!(
rewrite_compat_function_calls("select max(1, min(2, 3), 'max(4,5)')"),
"select __mongreldb_scalar_max(1, __mongreldb_scalar_min(2, 3), 'max(4,5)')"
);
assert_eq!(
rewrite_compat_function_calls("select /* max(1,2) */ min(1, (2 + 3))"),
"select /* max(1,2) */ __mongreldb_scalar_min(1, (2 + 3))"
);
assert_eq!(
rewrite_compat_function_calls("select max_value, min_value from t"),
"select max_value, min_value from t"
);
assert_eq!(
rewrite_compat_function_calls(
"select group_concat(label), group_concat(label, '|') from t"
),
"select string_agg(label, ','), string_agg(label, '|') from t"
);
assert_eq!(
rewrite_compat_function_calls("select total(val), total(val) filter (where grp = 2) from t"),
"select coalesce(cast(sum(val) as double), 0.0), coalesce(cast(sum(val) filter (where grp = 2) as double), 0.0) from t"
);
assert_eq!(
rewrite_compat_function_calls(
"select total(val) over (partition by grp order by id) from t"
),
"select coalesce(cast(sum(val) over (partition by grp order by id) as double), 0.0) from t"
);
}
}