use crate::{detect_content, detect_extension, CanonicalValue, Format, SqlDialect};
use serde_json::Number;
use sqlparser::ast::{
BinaryOperator, Expr, FunctionArg, FunctionArgExpr, FunctionArguments, GroupByExpr,
JoinConstraint, JoinOperator, LimitClause, ObjectNamePart, OrderByExpr, OrderByKind, Query,
SelectItem, SetExpr, Statement, TableFactor, TableWithJoins, Value, ValueWithSpan,
};
use sqlparser::dialect::GenericDialect;
use sqlparser::parser::Parser;
use std::cmp::Ordering;
use std::collections::{HashMap, HashSet};
use std::error::Error;
use std::path::{Path, PathBuf};
type Row = Vec<(String, CanonicalValue)>;
pub fn run(sql: &str, base_dir: &Path) -> Result<CanonicalValue, Box<dyn Error>> {
let stmts = Parser::parse_sql(&GenericDialect {}, sql)?;
let stmt = stmts.into_iter().next().ok_or("empty SQL")?;
match stmt {
Statement::Query(q) => execute_query(*q, base_dir),
_ => Err("only SELECT statements are supported".into()),
}
}
fn execute_query(query: Query, base_dir: &Path) -> Result<CanonicalValue, Box<dyn Error>> {
let body = (*query.body).clone();
let select = match body {
SetExpr::Select(s) => *s,
_ => return Err("only simple SELECT is supported".into()),
};
let mut rows = load_from(&select.from, base_dir)?;
if let Some(pred) = &select.selection {
let mut kept = Vec::new();
for row in rows {
if matches!(eval(pred.clone(), &row), Ok(CanonicalValue::Bool(true))) {
kept.push(row);
}
}
rows = kept;
}
let has_agg = select.projection.iter().any(|item| match item {
SelectItem::UnnamedExpr(e) | SelectItem::ExprWithAlias { expr: e, .. } => {
contains_aggregate(e)
}
_ => false,
});
let group_by = extract_group_by(&select.group_by);
let mut out_rows = if !group_by.is_empty() || has_agg {
aggregate(rows, &select.projection, &group_by, &select.having)?
} else {
project(rows, &select.projection)?
};
if select.distinct.is_some() {
let mut seen: HashSet<String> = HashSet::new();
out_rows.retain(|r| {
let key = r
.iter()
.map(|(_, v)| format!("{v:?}"))
.collect::<Vec<_>>()
.join("\x1f");
seen.insert(key)
});
}
if let Some(order) = &query.order_by {
if let OrderByKind::Expressions(exprs) = &order.kind {
sort_rows(&mut out_rows, exprs)?;
}
}
if let Some(lc) = &query.limit_clause {
if let Some(limit_expr) = extract_limit_expr(lc) {
if let Ok(CanonicalValue::Number(n)) = eval(limit_expr, &Vec::new()) {
if let Some(limit) = n.as_u64() {
out_rows.truncate(limit as usize);
}
}
}
}
Ok(CanonicalValue::Array(
out_rows.into_iter().map(CanonicalValue::Object).collect(),
))
}
fn extract_group_by(g: &GroupByExpr) -> Vec<Expr> {
match g {
GroupByExpr::Expressions(exprs, _) => exprs.clone(),
_ => Vec::new(),
}
}
fn load_from(from: &[TableWithJoins], base_dir: &Path) -> Result<Vec<Row>, Box<dyn Error>> {
if from.is_empty() {
return Err("FROM is required".into());
}
let mut rows = load_table_with_joins(&from[0], base_dir)?;
for extra in &from[1..] {
let more = load_table_with_joins(extra, base_dir)?;
let mut crossed = Vec::new();
for a in &rows {
for b in &more {
let mut combined = a.clone();
combined.extend(b.clone());
crossed.push(combined);
}
}
rows = crossed;
}
Ok(rows)
}
fn load_table_with_joins(t: &TableWithJoins, base_dir: &Path) -> Result<Vec<Row>, Box<dyn Error>> {
let mut rows = load_factor(&t.relation, base_dir)?;
for join in &t.joins {
let right = load_factor(&join.relation, base_dir)?;
rows = apply_join(rows, right, &join.join_operator)?;
}
Ok(rows)
}
fn load_factor(factor: &TableFactor, base_dir: &Path) -> Result<Vec<Row>, Box<dyn Error>> {
let (name, alias) = match factor {
TableFactor::Table { name, alias, .. } => (name, alias),
_ => return Err("only table references are supported in FROM".into()),
};
let raw = name
.0
.iter()
.filter_map(|p| match p {
ObjectNamePart::Identifier(i) => Some(i.value.clone()),
_ => None,
})
.collect::<Vec<_>>()
.join(".");
let alias_name = alias
.as_ref()
.map(|a| a.name.value.clone())
.unwrap_or_else(|| {
Path::new(&raw)
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or(&raw)
.to_string()
});
let mut all_rows = Vec::new();
if is_glob(&raw) {
let pattern = resolve_path(&raw, base_dir).to_string_lossy().to_string();
let entries = glob::glob(&pattern).map_err(|e| format!("invalid glob {pattern}: {e}"))?;
for entry in entries {
let p = entry?;
let mut r = load_file(&p)?;
all_rows.append(&mut r);
}
} else {
all_rows = load_file(&resolve_path(&raw, base_dir))?;
}
Ok(all_rows
.into_iter()
.map(|row| {
row.into_iter()
.map(|(k, v)| (format!("{alias_name}.{k}"), v))
.collect()
})
.collect())
}
fn is_glob(s: &str) -> bool {
s.contains('*') || s.contains('?') || s.contains('[')
}
fn resolve_path(raw: &str, base_dir: &Path) -> PathBuf {
let candidate = base_dir.join(raw);
if candidate.exists() || is_glob(raw) || Path::new(raw).is_absolute() {
return candidate;
}
let parts: Vec<&str> = raw.split('.').collect();
if parts.len() > 2 {
let prefix = parts[..parts.len() - 2].join("/");
let suffix = parts[parts.len() - 2..].join(".");
let slashed = base_dir.join(format!("{prefix}/{suffix}"));
if slashed.exists() {
return slashed;
}
}
candidate
}
fn load_file(path: &Path) -> Result<Vec<Row>, Box<dyn Error>> {
let bytes = std::fs::read(path)?;
let fmt = match detect_extension(path) {
Some(f) => f,
None => detect_content(&bytes)?,
};
let value = decode(fmt, &bytes)?;
match value {
CanonicalValue::Array(rows) => Ok(rows
.into_iter()
.filter_map(|r| match r {
CanonicalValue::Object(o) => Some(o),
_ => None,
})
.collect()),
_ => Err(format!("{} did not produce a table", path.display()).into()),
}
}
fn decode(fmt: Format, bytes: &[u8]) -> Result<CanonicalValue, Box<dyn Error>> {
let text = || -> Result<&str, Box<dyn Error>> { Ok(std::str::from_utf8(bytes)?) };
match fmt {
Format::Csv => CanonicalValue::from_csv(text()?),
Format::Tsv => CanonicalValue::from_tsv(text()?),
Format::Json => CanonicalValue::from_json(text()?),
Format::Ndjson => CanonicalValue::from_ndjson(text()?),
Format::Yaml => CanonicalValue::from_yaml(text()?),
Format::Xml => CanonicalValue::from_xml(text()?),
Format::Toml => CanonicalValue::from_toml(text()?),
Format::Xlsx => {
let mut sheets = CanonicalValue::from_xlsx_workbook(bytes)?;
if sheets.len() != 1 {
return Err("multi-sheet XLSX requires a single sheet".into());
}
Ok(sheets.remove(0).1)
}
Format::Ods => CanonicalValue::from_ods(bytes),
Format::Sql => {
let mut sheets = CanonicalValue::from_sql_workbook(text()?, SqlDialect::Generic)?;
if sheets.len() != 1 {
return Err("multi-table SQL requires a single table".into());
}
Ok(sheets.remove(0).1)
}
}
}
fn apply_join(
left: Vec<Row>,
right: Vec<Row>,
op: &JoinOperator,
) -> Result<Vec<Row>, Box<dyn Error>> {
let (constraint, is_left) = match op {
JoinOperator::Inner(c) => (c, false),
JoinOperator::Join(c) => (c, false),
JoinOperator::LeftOuter(c) => (c, true),
other => return Err(format!("unsupported join variant: {other:?}").into()),
};
let on = match constraint {
JoinConstraint::On(e) => e.clone(),
_ => return Err("only JOIN ... ON is supported".into()),
};
let mut result = Vec::new();
for l in &left {
let mut matched = false;
for r in &right {
let mut combined = l.clone();
combined.extend(r.clone());
if matches!(eval(on.clone(), &combined), Ok(CanonicalValue::Bool(true))) {
result.push(combined);
matched = true;
}
}
if is_left && !matched {
result.push(l.clone());
}
}
Ok(result)
}
fn eval(expr: Expr, row: &Row) -> Result<CanonicalValue, Box<dyn Error>> {
match expr {
Expr::Identifier(i) => resolve_column(&i.value, row),
Expr::CompoundIdentifier(parts) => {
let name = parts
.iter()
.map(|p| p.value.clone())
.collect::<Vec<_>>()
.join(".");
resolve_column(&name, row)
}
Expr::Value(ValueWithSpan { value, .. }) => Ok(literal(value)),
Expr::BinaryOp { left, op, right } => {
let l = eval(*left, row)?;
let r = eval(*right, row)?;
binary_op(l, &op, r)
}
Expr::Nested(inner) => eval(*inner, row),
other => Err(format!("unsupported expression: {other:?}").into()),
}
}
fn resolve_column(name: &str, row: &Row) -> Result<CanonicalValue, Box<dyn Error>> {
if let Some((_, v)) = row.iter().find(|(k, _)| k == name) {
return Ok(v.clone());
}
let suffix = format!(".{name}");
if let Some((_, v)) = row.iter().find(|(k, _)| k.ends_with(&suffix)) {
return Ok(v.clone());
}
Err(format!("column {name} not found").into())
}
fn literal(v: Value) -> CanonicalValue {
match v {
Value::Number(n, _) => {
if let Ok(i) = n.parse::<i64>() {
CanonicalValue::Number(Number::from(i))
} else if let Ok(f) = n.parse::<f64>() {
Number::from_f64(f)
.map(CanonicalValue::Number)
.unwrap_or(CanonicalValue::Null)
} else {
CanonicalValue::String(n)
}
}
Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => CanonicalValue::String(s),
Value::Boolean(b) => CanonicalValue::Bool(b),
Value::Null => CanonicalValue::Null,
_ => CanonicalValue::Null,
}
}
fn binary_op(
l: CanonicalValue,
op: &BinaryOperator,
r: CanonicalValue,
) -> Result<CanonicalValue, Box<dyn Error>> {
let ord = || compare_values(&l, &r);
match op {
BinaryOperator::Eq => return Ok(CanonicalValue::Bool(ord() == Ordering::Equal)),
BinaryOperator::NotEq => return Ok(CanonicalValue::Bool(ord() != Ordering::Equal)),
BinaryOperator::Lt => return Ok(CanonicalValue::Bool(ord() == Ordering::Less)),
BinaryOperator::LtEq => return Ok(CanonicalValue::Bool(ord() != Ordering::Greater)),
BinaryOperator::Gt => return Ok(CanonicalValue::Bool(ord() == Ordering::Greater)),
BinaryOperator::GtEq => return Ok(CanonicalValue::Bool(ord() != Ordering::Less)),
BinaryOperator::And => {
let a = matches!(l, CanonicalValue::Bool(true));
let b = matches!(r, CanonicalValue::Bool(true));
return Ok(CanonicalValue::Bool(a && b));
}
BinaryOperator::Or => {
let a = matches!(l, CanonicalValue::Bool(true));
let b = matches!(r, CanonicalValue::Bool(true));
return Ok(CanonicalValue::Bool(a || b));
}
_ => {}
}
let a = to_f64(&l);
let b = to_f64(&r);
let result = match (a, b) {
(Some(a), Some(b)) => match op {
BinaryOperator::Plus => Some(a + b),
BinaryOperator::Minus => Some(a - b),
BinaryOperator::Multiply => Some(a * b),
BinaryOperator::Divide => Some(a / b),
_ => None,
},
_ => None,
};
match result {
Some(f) => Ok(num(f)),
None => Err(format!("cannot apply {op:?} to {l:?} and {r:?}").into()),
}
}
fn to_f64(v: &CanonicalValue) -> Option<f64> {
match v {
CanonicalValue::Number(n) => n.as_f64(),
CanonicalValue::String(s) => s.parse().ok(),
CanonicalValue::Bool(b) => Some(if *b { 1.0 } else { 0.0 }),
_ => None,
}
}
fn num(f: f64) -> CanonicalValue {
if f.is_finite() && f.fract() == 0.0 && f.abs() < i64::MAX as f64 {
CanonicalValue::Number(Number::from(f as i64))
} else {
Number::from_f64(f)
.map(CanonicalValue::Number)
.unwrap_or(CanonicalValue::Null)
}
}
fn project(rows: Vec<Row>, projection: &[SelectItem]) -> Result<Vec<Row>, Box<dyn Error>> {
let mut out = Vec::with_capacity(rows.len());
for row in rows {
let mut new_row = Vec::new();
for item in projection {
match item {
SelectItem::UnnamedExpr(expr) => {
let name = expr_name(expr);
new_row.push((name, eval(expr.clone(), &row)?));
}
SelectItem::ExprWithAlias { expr, alias } => {
new_row.push((alias.value.clone(), eval(expr.clone(), &row)?));
}
SelectItem::Wildcard(_) => {
for (k, v) in &row {
let name = k.rsplit('.').next().unwrap_or(k).to_string();
new_row.push((name, v.clone()));
}
}
SelectItem::QualifiedWildcard(name, _) => {
let prefix = format!("{name}.");
for (k, v) in &row {
if let Some(short) = k.strip_prefix(&prefix) {
new_row.push((short.to_string(), v.clone()));
}
}
}
SelectItem::ExprWithAliases { .. } => {
return Err("multi-alias SELECT not supported".into());
}
}
}
out.push(new_row);
}
Ok(out)
}
fn expr_name(expr: &Expr) -> String {
match expr {
Expr::Identifier(i) => i.value.clone(),
Expr::CompoundIdentifier(parts) => {
parts.last().map(|p| p.value.clone()).unwrap_or_default()
}
Expr::Function(f) => {
let fname = f.name.to_string().to_lowercase();
let arg = match &f.args {
FunctionArguments::List(l) => l
.args
.first()
.and_then(|a| match a {
FunctionArg::Unnamed(FunctionArgExpr::Expr(e)) => Some(expr_name(e)),
_ => None,
})
.unwrap_or_default(),
_ => String::new(),
};
format!("{fname}({arg})")
}
Expr::BinaryOp { left, op, right } => {
format!("{} {} {}", expr_name(left), op, expr_name(right))
}
Expr::Nested(inner) => expr_name(inner),
_ => "expr".to_string(),
}
}
fn aggregate(
rows: Vec<Row>,
projection: &[SelectItem],
group_by: &[Expr],
having: &Option<Expr>,
) -> Result<Vec<Row>, Box<dyn Error>> {
let mut groups: HashMap<Vec<String>, Vec<Row>> = HashMap::new();
if group_by.is_empty() {
groups.insert(Vec::new(), rows);
} else {
for row in rows {
let mut key = Vec::new();
for g in group_by {
let v = eval(g.clone(), &row)?;
key.push(format!("{v:?}"));
}
groups.entry(key).or_default().push(row);
}
}
let mut out = Vec::new();
for group in groups.into_values() {
let mut new_row = Vec::new();
for item in projection {
match item {
SelectItem::UnnamedExpr(expr) => {
let name = expr_name(expr);
new_row.push((name, eval_aggregated(expr, &group)?));
}
SelectItem::ExprWithAlias { expr, alias } => {
new_row.push((alias.value.clone(), eval_aggregated(expr, &group)?));
}
_ => return Err("unsupported aggregate SELECT item".into()),
}
}
if let Some(h) = having {
if !matches!(eval(h.clone(), &new_row), Ok(CanonicalValue::Bool(true))) {
continue;
}
}
out.push(new_row);
}
Ok(out)
}
fn eval_aggregated(expr: &Expr, group: &[Row]) -> Result<CanonicalValue, Box<dyn Error>> {
if let Expr::Function(f) = expr {
let name = f.name.to_string().to_uppercase();
if matches!(name.as_str(), "SUM" | "AVG" | "MIN" | "MAX" | "COUNT") {
let arg = match &f.args {
FunctionArguments::List(l) => l.args.first().and_then(|a| match a {
FunctionArg::Unnamed(FunctionArgExpr::Expr(e)) => Some(e.clone()),
_ => None,
}),
_ => None,
};
return compute_aggregate(&name, arg, group);
}
}
if let Some(first) = group.first() {
return eval(expr.clone(), first);
}
Ok(CanonicalValue::Null)
}
fn compute_aggregate(
name: &str,
arg: Option<Expr>,
group: &[Row],
) -> Result<CanonicalValue, Box<dyn Error>> {
let arg = match arg {
Some(a) => a,
None => {
return match name {
"COUNT" => Ok(CanonicalValue::Number(Number::from(group.len() as i64))),
_ => Err(format!("{name} requires an argument").into()),
};
}
};
let mut count = 0usize;
let mut sum = 0.0f64;
let mut best: Option<f64> = None;
for row in group {
if let Ok(v) = eval(arg.clone(), row) {
if let CanonicalValue::Null = v {
continue;
}
if let Some(f) = to_f64(&v) {
sum += f;
count += 1;
best = Some(match best {
None => f,
Some(b) => match name {
"MIN" => b.min(f),
"MAX" => b.max(f),
_ => b,
},
});
}
}
}
Ok(match name {
"SUM" => num(sum),
"AVG" => {
if count > 0 {
num(sum / count as f64)
} else {
CanonicalValue::Null
}
}
"COUNT" => CanonicalValue::Number(Number::from(count as i64)),
"MIN" | "MAX" => best.map(num).unwrap_or(CanonicalValue::Null),
_ => CanonicalValue::Null,
})
}
fn extract_limit_expr(lc: &LimitClause) -> Option<Expr> {
match lc {
LimitClause::LimitOffset { limit, .. } => limit.clone(),
LimitClause::OffsetCommaLimit { limit, .. } => Some(limit.clone()),
}
}
fn sort_rows(rows: &mut [Row], order_by: &[OrderByExpr]) -> Result<(), Box<dyn Error>> {
let keys: Vec<(Expr, bool)> = order_by
.iter()
.map(|o| (o.expr.clone(), o.options.asc.unwrap_or(true)))
.collect();
rows.sort_by(|a, b| {
for (expr, asc) in &keys {
let av = eval(expr.clone(), a).unwrap_or(CanonicalValue::Null);
let bv = eval(expr.clone(), b).unwrap_or(CanonicalValue::Null);
let ord = compare_values(&av, &bv);
let ord = if *asc { ord } else { ord.reverse() };
if ord != Ordering::Equal {
return ord;
}
}
Ordering::Equal
});
Ok(())
}
fn compare_values(a: &CanonicalValue, b: &CanonicalValue) -> Ordering {
match (a, b) {
(CanonicalValue::Null, CanonicalValue::Null) => Ordering::Equal,
(CanonicalValue::Null, _) => Ordering::Less,
(_, CanonicalValue::Null) => Ordering::Greater,
_ => {
if let (Some(x), Some(y)) = (to_f64(a), to_f64(b)) {
return x.partial_cmp(&y).unwrap_or(Ordering::Equal);
}
scalar_string(a).cmp(&scalar_string(b))
}
}
}
fn scalar_string(v: &CanonicalValue) -> String {
match v {
CanonicalValue::Null => String::new(),
CanonicalValue::Bool(b) => b.to_string(),
CanonicalValue::Number(n) => n.to_string(),
CanonicalValue::String(s) => s.clone(),
CanonicalValue::Bytes(b) => b.iter().map(|x| format!("{x:02x}")).collect(),
_ => String::new(),
}
}
fn contains_aggregate(expr: &Expr) -> bool {
match expr {
Expr::Function(f) => {
let name = f.name.to_string().to_uppercase();
matches!(name.as_str(), "SUM" | "AVG" | "MIN" | "MAX" | "COUNT")
}
Expr::BinaryOp { left, right, .. } => contains_aggregate(left) || contains_aggregate(right),
Expr::Nested(inner) => contains_aggregate(inner),
_ => false,
}
}