use serde_json::Value;
use crate::dialect::filter::{build_filter, build_filter_raw, WhereClause};
use crate::dialect::ir::{RowCol, RowShape, SqlStmt};
use crate::dialect::{field_is_bool, scalar_column, Backend, ColumnRef};
use crate::schema::Schema;
use super::{col_fn, count_field_pattern, limit_offset_sql, q, tname};
struct GroupKey {
out: String,
field: String,
}
#[derive(Debug, Clone, Copy, PartialEq)]
enum AggOp {
CountRows,
Count,
Sum,
Avg,
Min,
Max,
}
struct GroupAggCol {
alias: String,
op: AggOp,
field: Option<String>,
}
pub(super) struct GroupSpec {
keys: Vec<GroupKey>,
aggs: Vec<GroupAggCol>,
}
pub(super) fn parse_group(g: &Value) -> Result<GroupSpec, String> {
let obj = g.as_object().ok_or("$group 阶段必须是对象")?;
let mut keys: Vec<GroupKey> = Vec::new();
match obj.get("_id") {
None | Some(Value::Null) => {}
Some(v) => flatten_id(v, "", &mut keys)?,
}
let mut aggs: Vec<GroupAggCol> = Vec::new();
for (alias, acc) in obj {
if alias == "_id" {
continue;
}
let (op, field) = parse_acc(acc)?;
aggs.push(GroupAggCol {
alias: alias.clone(),
op,
field,
});
}
Ok(GroupSpec { keys, aggs })
}
fn flatten_id(v: &Value, prefix: &str, out: &mut Vec<GroupKey>) -> Result<(), String> {
match v {
Value::String(s) => {
let field = s
.strip_prefix('$')
.ok_or_else(|| format!("$group._id 仅支持字段引用,收到 {s}"))?;
let out_key = if prefix.is_empty() {
field.to_string()
} else {
prefix.to_string()
};
out.push(GroupKey {
out: out_key,
field: field.to_string(),
});
Ok(())
}
Value::Object(m) => {
for (k, sub) in m {
let p = if prefix.is_empty() {
k.clone()
} else {
format!("{prefix}.{k}")
};
flatten_id(sub, &p, out)?;
}
Ok(())
}
other => Err(format!(
"$group._id 仅支持字段引用 / 字段嵌套对象 / null,收到 {other}"
)),
}
}
fn parse_acc(v: &Value) -> Result<(AggOp, Option<String>), String> {
let o = v.as_object().ok_or("$group 累积器必须恰有一个算子键")?;
let mut it = o.iter();
let (op, arg) = match (it.next(), it.next()) {
(Some(kv), None) => kv,
_ => return Err("$group 累积器必须恰有一个算子键".to_string()),
};
match op.as_str() {
"$sum" => {
if arg.as_i64() == Some(1) {
return Ok((AggOp::CountRows, None));
}
if let Some(f) = arg.as_str() {
return Ok((AggOp::Sum, Some(field_ref(f)?)));
}
if let Some(f) = count_field_pattern(arg) {
return Ok((AggOp::Count, Some(f)));
}
Err(format!("SQL 后端无法翻译 $group 的 $sum 累积器 {arg}"))
}
"$avg" => Ok((AggOp::Avg, Some(field_ref(acc_field(arg, "$avg")?)?))),
"$min" => Ok((AggOp::Min, Some(field_ref(acc_field(arg, "$min")?)?))),
"$max" => Ok((AggOp::Max, Some(field_ref(acc_field(arg, "$max")?)?))),
other => Err(format!(
"SQL 后端不支持的 $group 累积器 {other}(白名单 $count/$sum/$avg/$min/$max)"
)),
}
}
fn acc_field<'a>(arg: &'a Value, op: &str) -> Result<&'a str, String> {
arg.as_str()
.ok_or_else(|| format!("$group 的 {op} 累积器必须带字段引用"))
}
fn field_ref(s: &str) -> Result<String, String> {
s.strip_prefix('$')
.map(|x| x.to_string())
.ok_or_else(|| format!("$group 累积器仅支持字段引用,收到 {s}"))
}
fn agg_expr(backend: Backend, schema: &Schema, a: &GroupAggCol) -> Result<String, String> {
let col = |f: Option<&str>| -> Result<String, String> {
let f = f.ok_or_else(|| format!("$group 聚合列 \"{}\" 缺少字段", a.alias))?;
let c = scalar_column(schema, f).ok_or_else(|| {
format!(
"$group 聚合字段 \"{f}\" 无法映射到本表标量列(\"{}\")",
a.alias
)
})?;
Ok(format!("t.{}", q(backend, &c)))
};
Ok(match a.op {
AggOp::CountRows => "COUNT(*)".to_string(),
AggOp::Count => format!("COUNT({})", col(a.field.as_deref())?),
AggOp::Sum => format!("SUM({})", col(a.field.as_deref())?),
AggOp::Avg => format!(
"AVG(CAST({} AS {}))",
col(a.field.as_deref())?,
backend.double_type()
),
AggOp::Min => format!("MIN({})", col(a.field.as_deref())?),
AggOp::Max => format!("MAX({})", col(a.field.as_deref())?),
})
}
#[allow(clippy::too_many_arguments)]
pub(super) fn translate_group(
backend: Backend,
schema: &Schema,
registry: &crate::schema::Registry,
spec: &GroupSpec,
root_matches: &[Value],
having: Option<&Value>,
order: &[(String, i64)],
skip: i64,
limit: Option<i64>,
project: Option<&Value>,
param_seq: &mut usize,
warnings: &mut Vec<String>,
) -> Result<SqlStmt, String> {
let mut join_sql = String::new();
let mut keys: Vec<(String, String)> = Vec::new(); let mut joined_rels: Vec<(&str, &crate::schema::RelationDef)> = Vec::new();
for k in &spec.keys {
if let Some((head, rest)) = k.field.split_once('.') {
if let Some(rel) = schema.relations.get(head) {
if rel.rel_type != "one" {
return Err(format!(
"SQL 后端不支持 $group 的 by 键 \"{}\"(many 关系路径扇出会破坏聚合语义)",
k.field
));
}
let rel_schema = registry.get(&rel.model)?;
let col = scalar_column(rel_schema, rest).ok_or_else(|| {
format!(
"SQL 后端无法把 $group 的 by 键 \"{}\" 映射到关系 \"{}\" 的标量列",
k.field, head
)
})?;
if !joined_rels.iter().any(|(n, _)| *n == head) {
let g_alias = q(backend, &format!("g_{}", head));
join_sql.push_str(&format!(
" LEFT JOIN {} {} ON {}.{} = t.{}",
tname(backend, rel_schema),
g_alias,
g_alias,
q(backend, &rel.foreign_field),
q(backend, &rel.local_field),
));
joined_rels.push((head, rel));
}
keys.push((
k.out.clone(),
format!(
"{}.{}",
q(backend, &format!("g_{}", head)),
q(backend, &col)
),
));
continue;
}
}
let col = scalar_column(schema, &k.field).ok_or_else(|| {
format!(
"SQL 后端无法把 $group 的 by 键 \"{}\" 映射到本表标量列(object/array 点号路径未映射为列)",
k.field
)
})?;
keys.push((k.out.clone(), format!("t.{}", q(backend, &col))));
}
let mut aggs: Vec<(String, String)> = Vec::new();
for a in &spec.aggs {
aggs.push((a.alias.clone(), agg_expr(backend, schema, a)?));
}
let expr_of = |name: &str| -> Option<String> {
let by_alias: Option<&str> = if name == "_id" {
if keys.len() == 1 {
Some(keys[0].0.as_str())
} else {
None
}
} else {
name.strip_prefix("_id.")
};
if let Some(a) = by_alias {
if let Some((_, e)) = keys.iter().find(|(o, _)| o == a) {
return Some(e.clone());
}
}
if let Some((_, e)) = keys.iter().find(|(o, _)| o == name) {
return Some(e.clone());
}
aggs.iter().find(|(a, _)| a == name).map(|(_, e)| e.clone())
};
let requested: Option<Vec<String>> = project.and_then(|p| p.as_object()).map(|o| {
o.iter()
.filter(|(k, v)| {
!(k.as_str() == "_id" && (v.as_i64() == Some(0) || v.as_bool() == Some(false)))
})
.map(|(k, _)| k.clone())
.collect()
});
let wanted = |name: &str| -> bool {
requested
.as_ref()
.map(|r| r.iter().any(|x| x == name))
.unwrap_or(true)
};
let mut cols_sql: Vec<String> = Vec::new();
let mut columns: Vec<RowCol> = Vec::new();
for (i, k) in spec.keys.iter().enumerate() {
let out = &k.out;
if !wanted(out) {
continue;
}
let alias = format!("b{i}");
cols_sql.push(format!("{} AS {}", keys[i].1, q(backend, &alias)));
columns.push(RowCol {
alias,
json_path: out.split('.').map(String::from).collect(),
is_array: false,
one: false,
ones: Vec::new(),
sub_shape: None,
always: true,
is_bool: field_is_bool(schema, &k.field),
is_json: false,
});
}
for (i, (alias_name, expr)) in aggs.iter().enumerate() {
if !wanted(alias_name) {
continue;
}
let alias = format!("a{i}");
cols_sql.push(format!("{} AS {}", expr, q(backend, &alias)));
columns.push(RowCol {
alias,
json_path: vec![alias_name.clone()],
is_array: false,
one: false,
ones: Vec::new(),
sub_shape: None,
always: true,
is_bool: false,
is_json: false,
});
}
if cols_sql.is_empty() {
return Err("$group 查询没有可输出的分组键 / 聚合列".to_string());
}
let mut wheres: Vec<WhereClause> = Vec::new();
for m in root_matches {
let wh = build_filter(
m,
backend,
"t",
&col_fn(schema),
param_seq,
Some(&mut *warnings),
)?;
if !wh.text.is_empty() {
wheres.push(wh);
}
}
let mut params: Vec<Value> = Vec::new();
let where_sql = if wheres.is_empty() {
String::new()
} else {
let text = wheres
.iter()
.map(|w| w.text.clone())
.collect::<Vec<_>>()
.join(" AND ");
for w in &wheres {
params.extend(w.params.clone());
}
format!(" WHERE {text}")
};
let group_sql = if keys.is_empty() {
String::new()
} else {
format!(
" GROUP BY {}",
keys.iter()
.map(|(_, e)| e.clone())
.collect::<Vec<_>>()
.join(", ")
)
};
let expr_ref = |name: &str| -> Option<ColumnRef> { expr_of(name).map(ColumnRef::Scalar) };
let having_sql = match having {
Some(h) if !h.is_null() => {
let wh = build_filter_raw(h, backend, &expr_ref, param_seq, Some(&mut *warnings))?;
if wh.text.is_empty() {
String::new()
} else {
params.extend(wh.params);
format!(" HAVING {}", wh.text)
}
}
_ => String::new(),
};
let mut order_parts: Vec<String> = Vec::new();
for (k, dir) in order {
let e = expr_of(k)
.ok_or_else(|| format!("$sort 键 \"{k}\" 不在 $group 的 by 键 / agg 别名域内"))?;
order_parts.push(format!("{} {}", e, if *dir >= 0 { "ASC" } else { "DESC" }));
}
let order_sql = if order_parts.is_empty() {
String::new()
} else {
format!(" ORDER BY {}", order_parts.join(", "))
};
let (limit_sql, limit_params) = limit_offset_sql(backend, limit, skip, param_seq);
params.extend(limit_params);
let text = format!(
"SELECT {} FROM {}{}{}{}{}{}{}",
cols_sql.join(", "),
tname(backend, schema),
" t",
&join_sql,
where_sql,
group_sql,
having_sql,
order_sql,
) + &limit_sql;
Ok(SqlStmt::select(
text,
params,
RowShape {
columns,
present_alias: None,
},
))
}