use std::collections::HashSet;
use serde_json::{json, Map, Value};
use crate::command::forbid_t2q;
use crate::permission::Context;
use crate::schema::{Registry, Schema};
use crate::types::{is_truthy, validate_condition, validate_condition_shape, validate_sort_shape};
use super::ast::{Ast, RelAst};
use super::group;
use super::lookup::{build_agg_stages, build_lookup};
use super::relation_filter;
use super::util::{append_order, find_stage_idx, force_t2q_limit, non_nullish, param};
pub fn flatten_object_fields(ast: &mut Ast, schema: &Schema) {
flatten_object_fields_impl(&mut ast.fields, &mut ast.relations, schema);
}
pub(crate) fn flatten_object_fields_impl(
fields: &mut Vec<String>,
relations: &mut Vec<(String, RelAst)>,
schema: &Schema,
) {
let mut removed: Vec<String> = Vec::new();
for (rel_name, rel_ast) in relations.iter() {
let is_object_field = schema
.fields
.get(rel_name)
.map(|f| f.field_type == "object" && f.fields.as_ref().map(is_truthy).unwrap_or(false))
.unwrap_or(false);
if !is_object_field {
continue;
}
fields.extend(
rel_ast
.fields
.iter()
.map(|sub| format!("{}.{}", rel_name, sub)),
);
removed.push(rel_name.clone());
}
relations.retain(|(n, _)| !removed.contains(n));
}
fn override_or_append(stages: &mut Vec<Value>, stage_key: &str, value: Value) {
let stage = Value::Object(Map::from_iter([(stage_key.to_string(), value)]));
match find_stage_idx(stages, stage_key) {
Some(idx) => stages[idx] = stage,
None => stages.push(stage),
}
}
fn custom_pipeline_branch(
stages: &mut Vec<Value>,
root_pipeline: &Value,
root_condition: Option<&Value>,
root_sort: Option<&Value>,
root_skip: Option<&Value>,
root_limit: Option<&Value>,
) -> Result<Value, String> {
if let Some(arr) = root_pipeline.as_array() {
stages.extend(arr.iter().cloned());
}
if let Some(v) = non_nullish(root_condition) {
validate_condition(v)?;
override_or_append(stages, "$match", v.clone());
}
if let Some(v) = non_nullish(root_sort) {
override_or_append(stages, "$sort", v.clone());
}
if let Some(v) = non_nullish(root_skip) {
override_or_append(stages, "$skip", v.clone());
}
if let Some(v) = non_nullish(root_limit) {
override_or_append(stages, "$limit", v.clone());
}
Ok(Value::Array(std::mem::take(stages)))
}
fn root_lookup_stages(
ast: &Ast,
schema: &Schema,
params: &Map<String, Value>,
stages: &mut Vec<Value>,
registry: &Registry,
ctx: Option<&Context>,
) -> Result<(), String> {
for (rel_name, rel_ast) in &ast.relations {
let rel_def = schema.relations.get(rel_name).ok_or_else(|| {
format!(
"关系 \"{}\" 未在 schema \"{}\" 中定义",
rel_name, schema.name
)
})?;
let rel_schema = registry.get(&rel_def.model)?;
stages.push(build_lookup(
rel_name, rel_ast, params, rel_def, rel_schema, schema, 0, 0, registry, ctx,
)?);
if rel_def.rel_type == "one" {
stages.push(json!({
"$unwind": { "path": format!("${}", rel_name), "preserveNullAndEmptyArrays": true }
}));
}
}
Ok(())
}
pub fn build_pipeline(
ast: &mut Ast,
params: &Map<String, Value>,
registry: &Registry,
ctx: Option<&Context>,
) -> Result<Value, String> {
let schema = registry.get(&ast.model)?;
let mut stages: Vec<Value> = Vec::new();
let root_condition = param(params, ast.params.get("condition")).cloned();
let root_group = param(params, ast.params.get("group")).cloned();
let root_having = param(params, ast.params.get("having")).cloned();
let root_sort = param(params, ast.params.get("sort")).cloned();
let root_skip = param(params, ast.params.get("skip")).cloned();
let root_limit = force_t2q_limit(registry, param(params, ast.params.get("limit")));
let root_pipeline = param(params, ast.params.get("pipeline")).cloned();
if let Some(pipe) = non_nullish(root_pipeline.as_ref()).filter(|v| v.is_array()) {
forbid_t2q(registry, "$pipeline 直通")?;
flatten_object_fields(ast, schema);
if let Some((rel_name, _)) = ast
.relations
.iter()
.find(|(n, _)| schema.relations.contains_key(n))
{
return Err(format!(
"`$pipeline` 直通模式不支持关系字段 [{rel_name}](自定义 pipeline 与关系下推不能并用)"
));
}
return custom_pipeline_branch(
&mut stages,
pipe,
root_condition.as_ref(),
root_sort.as_ref(),
root_skip.as_ref(),
root_limit.as_ref(),
);
}
if let Some(cond) = non_nullish(root_condition.as_ref()) {
validate_condition_shape(schema, cond, registry.profile())?;
validate_condition(cond)?;
}
let rel_plan = match non_nullish(root_condition.as_ref()) {
Some(cond) => relation_filter::plan(schema, registry, ctx, cond)?,
None => None,
};
if let Some(group_v) = non_nullish(root_group.as_ref()) {
if rel_plan.is_some() {
return Err(
"$group 查询暂不支持关系聚合谓词(§9.6 属 $condition 段,与分组终结路径暂未互通)"
.to_string(),
);
}
let spec = group::parse(schema, group_v, registry)?;
group::validate_read_permission(schema, ctx, &spec)?;
let stages = group::build_stages(
&spec,
&ast.fields,
ast.relations.is_empty(),
non_nullish(root_condition.as_ref()),
non_nullish(root_having.as_ref()),
non_nullish(root_sort.as_ref()),
non_nullish(root_skip.as_ref()),
non_nullish(root_limit.as_ref()),
)?;
return Ok(Value::Array(stages));
}
if non_nullish(root_having.as_ref()).is_some() {
return Err("$having 需要与 $group 同时使用(无 $group 时没有分组结果可过滤)".to_string());
}
if let Some(srt) = non_nullish(root_sort.as_ref()) {
validate_sort_shape(schema, srt, registry.profile())?;
}
flatten_object_fields(ast, schema);
if let Some(plan) = &rel_plan {
stages.extend(plan.lookups.iter().cloned());
stages.push(json!({ "$match": plan.condition.clone() }));
} else if let Some(cond) = non_nullish(root_condition.as_ref()) {
stages.push(json!({ "$match": cond }));
}
root_lookup_stages(ast, schema, params, &mut stages, registry, ctx)?;
append_order(
&mut stages,
root_sort.as_ref(),
root_skip.as_ref(),
root_limit.as_ref(),
);
let requested_computes: HashSet<String> = schema
.computes
.iter()
.map(|(k, _)| k.clone())
.filter(|k| {
ast.fields
.iter()
.any(|f| f == k || f.starts_with(&format!("{}.", k)))
})
.collect();
stages.extend(build_agg_stages(
schema,
registry,
ctx,
&requested_computes,
)?);
Ok(Value::Array(stages))
}