use serde_json::{json, Value};
use crate::schema::{Registry, Schema};
use crate::dialect::filter::{
build_filter, build_filter_with_relations, RelPredResolver, Warnings, WhereClause,
};
use crate::dialect::ir::{RowCol, RowShape, SqlStmt};
use crate::dialect::{field_is_bool, Backend};
use crate::pipeline::REL_PRED_PREFIX;
use super::group_agg::{self, GroupSpec};
use super::lookup_join::{resolve_join, Join};
use super::relation_agg;
use super::{col_fn, limit_offset_sql, projection_fields, q, tname};
struct PredResolver<'a> {
preds: &'a [relation_agg::SqlRelPredicate],
backend: Backend,
schema: &'a Schema,
registry: &'a Registry,
}
impl RelPredResolver for PredResolver<'_> {
fn resolve(
&self,
field: &str,
cond: &Value,
param_seq: &mut usize,
warnings: Warnings<'_>,
) -> Result<Option<WhereClause>, String> {
if !field.starts_with(REL_PRED_PREFIX) {
return Ok(None);
}
let as_name = field.rsplit_once('.').map(|(a, _)| a).unwrap_or(field);
let Some(p) = self.preds.iter().find(|p| p.as_name == as_name) else {
return Err(format!(
"关系聚合谓词代理字段 \"{field}\" 没有对应的 $lookup"
));
};
let exists = cond
.as_object()
.and_then(|o| o.get("$exists"))
.and_then(|v| v.as_bool())
.ok_or_else(|| {
format!("关系聚合谓词代理字段 \"{field}\" 的条件非法(须为 {{\"$exists\": bool}})")
})?;
Ok(Some(relation_agg::exists_clause(
self.backend,
"t",
self.schema,
self.registry,
p,
!exists,
param_seq,
warnings,
)?))
}
}
pub(super) fn lookup_extra_condition(pipeline: &[Value]) -> Option<Value> {
let mut extras: Vec<Value> = Vec::new();
for stage in pipeline {
let Some(m) = stage.get("$match") else {
continue;
};
let parts: Vec<&Value> = match m.get("$and").and_then(|a| a.as_array()) {
Some(a) => a.iter().collect(),
None => vec![m],
};
for p in parts {
let is_join_key = p.get("$expr").is_some() && p.to_string().contains("$$");
if !is_join_key {
extras.push(p.clone());
}
}
}
match extras.len() {
0 => None,
1 => Some(extras.remove(0)),
_ => Some(json!({ "$and": extras })),
}
}
#[derive(Debug, Clone)]
struct ComputeAgg {
key: String,
rel_name: String,
model: String,
local_col: String,
foreign_col: String,
op: String,
field: Option<String>,
extra: Option<Value>,
}
fn resolve_compute_agg(schema: &Schema, lo: &Value) -> Option<ComputeAgg> {
let as_name = lo.get("as").and_then(|v| v.as_str())?;
let key = as_name.strip_prefix('_')?;
let agg = schema.compute(key)?.agg.clone()?;
let (op, path) = agg
.as_object()
.and_then(|o| o.iter().next())
.map(|(k, v)| (k.clone(), v.as_str().unwrap_or("").to_string()))?;
let (rel_name, field) = match path.split_once('.') {
Some((r, f)) => (r.to_string(), Some(f.to_string())),
None => (path, None),
};
let rel_def = schema.relations.get(&rel_name)?;
let pipeline: Vec<Value> = lo
.get("pipeline")
.and_then(|v| v.as_array())
.cloned()
.unwrap_or_default();
Some(ComputeAgg {
key: key.to_string(),
rel_name,
model: rel_def.model.clone(),
local_col: rel_def.local_field.clone(),
foreign_col: rel_def.foreign_field.clone(),
op,
field,
extra: lookup_extra_condition(&pipeline),
})
}
#[allow(clippy::too_many_arguments)]
fn collect_lookup(
schema: &Schema,
registry: &Registry,
lo: &Value,
parent: Option<usize>,
parent_path: &[String],
parent_ones: &[bool],
joins: &mut Vec<Join>,
warnings: &mut Vec<String>,
) -> Result<(), String> {
let alias = lo
.get("as")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let from = lo
.get("from")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let Some(mut j) = resolve_join(schema, registry, &alias, &from) else {
warnings.push(format!("$lookup 关系 {} 无法匹配 schema,跳过 JOIN", alias));
return Ok(());
};
let child_pipeline: Vec<Value> = lo
.get("pipeline")
.and_then(|v| v.as_array())
.cloned()
.unwrap_or_default();
j.parent = parent;
j.path = {
let mut p = parent_path.to_vec();
p.push(j.rel_name.clone());
p
};
j.ones = {
let mut o = parent_ones.to_vec();
o.push(j.one);
o
};
j.extra = lookup_extra_condition(&child_pipeline);
j.child_sort = child_pipeline.iter().find_map(|s| s.get("$sort").cloned());
j.child_skip = child_pipeline
.iter()
.find_map(|s| s.get("$skip").and_then(|v| v.as_i64()));
j.child_limit = child_pipeline
.iter()
.find_map(|s| s.get("$limit").and_then(|v| v.as_i64()));
let idx = joins.len();
let path = j.path.clone();
let ones = j.ones.clone();
let model = j.model.clone();
joins.push(j);
let rel_schema = registry.get(&model)?.clone();
for stage in &child_pipeline {
if let Some(nested) = stage.get("$lookup") {
collect_lookup(
&rel_schema,
registry,
nested,
Some(idx),
&path,
&ones,
joins,
warnings,
)?;
}
}
Ok(())
}
fn child_order_sql(backend: Backend, rel_schema: &Schema, j: &Join) -> Result<String, String> {
let mut parts: Vec<String> = Vec::new();
if let Some(s) = &j.child_sort {
if let Some(o) = s.as_object() {
for (k, dir) in o {
let Some(c) = crate::dialect::scalar_column(rel_schema, k) else {
return Err(format!(
"$lookup 关系 {} 的子 $sort 字段 {} 无法映射到本层标量列,窗口函数下推不支持",
j.alias, k
));
};
let d = dir.as_i64().unwrap_or(1);
parts.push(format!(
"c.{} {}",
q(backend, &c),
if d >= 0 { "ASC" } else { "DESC" }
));
}
}
}
if parts.is_empty() {
parts.push(format!("c.{} ASC", q(backend, "_id")));
}
Ok(parts.join(", "))
}
pub(super) fn translate_aggregate(
backend: Backend,
schema: &Schema,
pipeline: &[Value],
registry: &Registry,
warnings: &mut Vec<String>,
unsupported: &mut Vec<Value>,
) -> Result<Vec<SqlStmt>, String> {
let mut param_seq = 0usize;
let mut root_matches: Vec<Value> = Vec::new();
let mut root_order: Vec<String> = Vec::new();
let mut root_limit: Option<i64> = None;
let mut root_offset: i64 = 0;
let mut joins: Vec<Join> = Vec::new();
let mut compute_aggs: Vec<ComputeAgg> = Vec::new();
let mut rel_preds: Vec<relation_agg::SqlRelPredicate> = Vec::new();
let mut project_on: Option<Vec<String>> = None;
let mut group: Option<GroupSpec> = None;
let mut group_having: Option<Value> = None;
let mut group_order: Vec<(String, i64)> = Vec::new();
let mut group_skip: i64 = 0;
let mut group_limit: Option<i64> = None;
let mut group_project: Option<Value> = None;
for stage in pipeline {
if let Some(g) = stage.get("$group") {
group = Some(group_agg::parse_group(g)?);
} else if let Some(m) = stage.get("$match") {
if group.is_some() {
group_having = Some(m.clone());
} else {
root_matches.push(m.clone());
}
} else if let Some(lo) = stage.get("$lookup") {
let as_name = lo.get("as").and_then(|v| v.as_str()).unwrap_or("");
if as_name.starts_with(REL_PRED_PREFIX) {
rel_preds.push(relation_agg::parse(lo, schema, registry, backend)?);
} else {
match resolve_compute_agg(schema, lo) {
Some(ca) => compute_aggs.push(ca),
None => {
collect_lookup(schema, registry, lo, None, &[], &[], &mut joins, warnings)?
}
}
}
} else if let Some(s) = stage.get("$sort") {
if let Some(o) = s.as_object() {
for (k, dir) in o {
if group.is_some() {
group_order.push((k.clone(), dir.as_i64().unwrap_or(1)));
continue;
}
let resolved: Option<String> = match k.split_once('.') {
Some((head, rest)) => {
joins.iter().position(|j| j.rel_name == head).and_then(|i| {
registry
.get(&joins[i].model)
.ok()
.and_then(|rel| crate::dialect::scalar_column(rel, rest))
.map(|c| format!("r{}.{}", i, q(backend, &c)))
})
}
None => {
if schema.relations.iter().any(|(n, _)| n.as_str() == k) {
None
} else {
col_fn(schema)(k).map(|c| format!("t.{}", q(backend, &c)))
}
}
};
match resolved {
Some(expr) => {
let d = dir.as_i64().unwrap_or(1);
root_order.push(format!(
"{} {}",
expr,
if d >= 0 { "ASC" } else { "DESC" }
));
}
None => {
warnings.push(format!(
"$sort 字段 {k} 无法映射到有效列(未知字段 / object·array 字段 / 无对应 $lookup),未下推排序"
));
unsupported.push(json!({
"code": "sortField",
"field": k,
"reason": "$sort 字段无法映射到有效列,需 Host 侧兜底排序",
}));
}
}
}
}
} else if let Some(l) = stage.get("$skip") {
if group.is_some() {
group_skip = l.as_i64().unwrap_or(0);
} else {
root_offset = l.as_i64().unwrap_or(0);
}
} else if let Some(l) = stage.get("$limit") {
if group.is_some() {
group_limit = l.as_i64();
} else {
root_limit = l.as_i64();
}
} else if let Some(p) = stage.get("$project") {
if group.is_some() {
group_project = Some(p.clone());
} else {
super::check_projection_supported(schema, Some(p))?;
project_on = Some(projection_fields(schema, Some(p)));
}
} else if stage.get("$facet").is_some() || stage.get("$replaceRoot").is_some() {
if group.is_none() {
let stage_name = if stage.get("$facet").is_some() {
"$facet"
} else {
"$replaceRoot"
};
return Err(format!("SQL 后端暂不支持的聚合阶段 {stage_name}"));
}
} else if let Some(u) = stage.get("$unwind") {
let path = u.get("path").and_then(|v| v.as_str()).unwrap_or("");
let rel = path.strip_prefix('$').unwrap_or(path);
if !joins
.iter()
.any(|j| j.path.len() == 1 && j.path[0] == rel && j.one)
{
return Err(format!(
"SQL 后端不支持的 $unwind 路径 {path}(非 one 关系 $lookup 产物)"
));
}
} else if let Some(af) = stage.get("$addFields") {
let keys: Vec<String> = af
.as_object()
.map(|o| o.keys().cloned().collect())
.unwrap_or_default();
let all_resolved = !keys.is_empty()
&& keys
.iter()
.all(|k| compute_aggs.iter().any(|ca| &ca.key == k));
if !all_resolved {
return Err(
"SQL 后端暂不支持的聚合阶段 $addFields:拒绝静默忽略后返回未聚合的原始行"
.to_string(),
);
}
} else {
let stage_name = stage
.as_object()
.and_then(|m| m.keys().next())
.cloned()
.unwrap_or_else(|| "?".to_string());
return Err(format!(
"SQL 后端暂不支持的聚合阶段 {stage_name}:拒绝静默忽略后返回未聚合的原始行"
));
}
}
if let Some(spec) = group {
let stmt = group_agg::translate_group(
backend,
schema,
&spec,
&root_matches,
group_having.as_ref(),
&group_order,
group_skip,
group_limit,
group_project.as_ref(),
&mut param_seq,
warnings,
)?;
return Ok(vec![stmt]);
}
let selected = project_on.unwrap_or_else(|| schema.fields.keys().cloned().collect());
let mut cols_sql = Vec::new();
let mut columns: Vec<RowCol> = Vec::new();
cols_sql.push(format!("t.{}", q(backend, "_id")));
columns.push(RowCol::scalar("_id", &["_id"]));
let root_selected: Vec<String> = selected
.iter()
.filter(|f| schema.relations.iter().all(|(rn, _)| rn != *f))
.filter(|f| !compute_aggs.iter().any(|ca| &ca.key == *f))
.cloned()
.collect();
for f in &root_selected {
if let Some(c) = col_fn(schema)(f) {
cols_sql.push(format!("t.{}", q(backend, &c)));
columns.push(RowCol::scalar_bool(
&c,
&[f.as_str()],
field_is_bool(schema, f),
));
}
}
if cols_sql.is_empty() && !selected.is_empty() {
cols_sql.push(format!("t.{}", q(backend, "_id")));
columns.push(RowCol::scalar("_id", &["_id"]));
}
let mut from_sql = format!("{} t", tname(backend, schema));
let mut all_params: Vec<Value> = Vec::new();
let mut join_params: Vec<Value> = Vec::new();
for (i, j) in joins.iter().enumerate() {
let r = format!("r{}", i);
let parent_alias = match j.parent {
Some(p) => format!("r{}", p),
None => "t".to_string(),
};
let rel_schema = match registry.get(&j.model) {
Ok(s) => s,
Err(e) => {
warnings.push(format!(
"$lookup 关系 {} 目标不可定位,跳过 JOIN: {}",
j.alias, e
));
continue;
}
};
let has_topn = j.child_limit.is_some() || j.child_skip.is_some();
let mut extra_sql = String::new();
if let Some(extra) = &j.extra {
let alias_for_extra = if has_topn { "c" } else { r.as_str() };
let wh = build_filter(
extra,
backend,
alias_for_extra,
&col_fn(rel_schema),
&mut param_seq,
Some(&mut *warnings),
)?;
if !wh.text.is_empty() {
extra_sql = wh.text;
join_params.extend(wh.params);
}
}
let child_ref = if has_topn {
let order = child_order_sql(backend, rel_schema, j)?;
let mut inner = format!(
"SELECT c.*, ROW_NUMBER() OVER (PARTITION BY c.{} ORDER BY {}) AS {} FROM {} c",
q(backend, &j.foreign_col),
order,
q(backend, "__rn"),
tname(backend, rel_schema),
);
if !extra_sql.is_empty() {
inner.push_str(&format!(" WHERE {}", extra_sql));
extra_sql.clear();
}
let mut conds: Vec<String> = Vec::new();
let skip = j.child_skip.unwrap_or(0);
if skip > 0 {
conds.push(format!("w.{} > {}", q(backend, "__rn"), skip));
}
if let Some(lm) = j.child_limit {
conds.push(format!("w.{} <= {}", q(backend, "__rn"), skip + lm));
}
if conds.is_empty() {
format!("({})", inner)
} else {
format!(
"(SELECT * FROM ({}) w WHERE {})",
inner,
conds.join(" AND ")
)
}
} else {
tname(backend, rel_schema)
};
let on_extra = if extra_sql.is_empty() {
String::new()
} else {
format!(" AND {}", extra_sql)
};
from_sql.push_str(&format!(
" LEFT JOIN {} {} ON {}.{} = {}.{}{}",
child_ref,
r,
r,
q(backend, &j.foreign_col),
parent_alias,
q(backend, &j.local_col),
on_extra,
));
let rel_cols = projection_fields(rel_schema, None);
for rf in rel_cols {
if rel_schema
.fields
.get(&rf)
.map(|f| f.field_type == "object" || f.field_type == "array")
.unwrap_or(false)
{
continue;
}
let alias_col = format!("{}_{}_{}", j.alias, i, rf);
cols_sql.push(format!(
"{}.{} AS {}",
r,
q(backend, &rf),
q(backend, &alias_col)
));
let mut json_path = j.path.clone();
json_path.push(rf.clone());
columns.push(RowCol {
alias: alias_col,
json_path,
is_array: true,
one: *j.ones.first().unwrap_or(&false),
ones: j.ones.clone(),
sub_shape: None,
always: false,
is_bool: field_is_bool(rel_schema, &rf),
});
}
}
for (i, ca) in compute_aggs.iter().enumerate() {
let rel_schema = match registry.get(&ca.model) {
Ok(s) => s,
Err(e) => {
return Err(format!(
"计算列 \"{}\" 的聚合关系 \"{}\" 目标不可定位: {}",
ca.key, ca.rel_name, e
))
}
};
let alias = format!("a{}", i);
let mut extra_sql = String::new();
if let Some(extra) = &ca.extra {
let wh = build_filter(
extra,
backend,
"c",
&col_fn(rel_schema),
&mut param_seq,
Some(&mut *warnings),
)?;
if !wh.text.is_empty() {
extra_sql = wh.text;
join_params.extend(wh.params);
}
}
let fk = q(backend, &ca.foreign_col);
let agg_expr = match ca.op.as_str() {
"$count" => "COUNT(*)".to_string(),
op @ ("$sum" | "$avg" | "$min" | "$max") => {
let f = ca
.field
.as_deref()
.ok_or_else(|| format!("计算列 \"{}\" 的 {op} 必须引用关系字段", ca.key))?;
let col = crate::dialect::scalar_column(rel_schema, f).ok_or_else(|| {
format!(
"计算列 \"{}\" 的聚合字段 {} 无法映射到关系 {} 的标量列",
ca.key, f, ca.rel_name
)
})?;
let fn_name = match op {
"$sum" => "SUM",
"$avg" => "AVG",
"$min" => "MIN",
_ => "MAX",
};
let arg_sql = format!("c.{}", q(backend, &col));
if op == "$avg" {
format!("AVG(CAST({arg_sql} AS {}))", backend.double_type())
} else {
format!("{}({})", fn_name, arg_sql)
}
}
other => {
return Err(format!(
"计算列 \"{}\" 的 agg 算子 {other} 不在白名单($count/$sum/$avg/$min/$max)",
ca.key
))
}
};
let mut inner = format!(
"SELECT c.{} AS {}, {} AS {} FROM {} c",
fk,
q(backend, "fk"),
agg_expr,
q(backend, "v"),
tname(backend, rel_schema),
);
if !extra_sql.is_empty() {
inner.push_str(&format!(" WHERE {}", extra_sql));
}
inner.push_str(&format!(" GROUP BY c.{}", fk));
from_sql.push_str(&format!(
" LEFT JOIN ({}) {} ON {}.{} = t.{}",
inner,
alias,
alias,
q(backend, "fk"),
q(backend, &ca.local_col),
));
let out = if ca.op == "$count" {
format!("COALESCE({}.{}, 0)", alias, q(backend, "v"))
} else {
format!("{}.{}", alias, q(backend, "v"))
};
cols_sql.push(format!("{} AS {}", out, q(backend, &ca.key)));
columns.push(RowCol::computed(&ca.key, &[ca.key.as_str()]));
}
let resolver = PredResolver {
preds: &rel_preds,
backend,
schema,
registry,
};
let mut root_wheres: Vec<WhereClause> = Vec::new();
for m in &root_matches {
let wh = build_filter_with_relations(
m,
backend,
"t",
&col_fn(schema),
Some(&resolver),
&mut param_seq,
Some(&mut *warnings),
)?;
if !wh.text.is_empty() {
root_wheres.push(wh);
}
}
all_params.extend(join_params);
let where_sql = if root_wheres.is_empty() {
String::new()
} else {
let text = root_wheres
.iter()
.map(|w| w.text.clone())
.collect::<Vec<_>>()
.join(" AND ");
for w in &root_wheres {
all_params.extend(w.params.clone());
}
format!(" WHERE {}", text)
};
let order_sql = if root_order.is_empty() {
String::new()
} else {
format!(" ORDER BY {}", root_order.join(", "))
};
let (limit_sql, limit_params) =
limit_offset_sql(backend, root_limit, root_offset, &mut param_seq);
all_params.extend(limit_params);
let select_list = if cols_sql.is_empty() {
q(backend, "_id")
} else if !cols_sql.iter().any(|c| c.contains("__present")) {
cols_sql.push("t.__present AS __present".to_string());
cols_sql.join(", ")
} else {
cols_sql.join(", ")
};
let text = format!(
"SELECT {} FROM {}{}{}{}",
select_list, from_sql, where_sql, order_sql, limit_sql,
);
Ok(vec![SqlStmt::select(
text,
all_params,
RowShape {
columns,
present_alias: Some("__present".to_string()),
},
)])
}