use serde_json::{json, Map, Value};
use crate::permission::{merge_owner_condition, Context};
use crate::pipeline::{param, parse_gql};
use crate::rbac::merge_row_condition;
use crate::schema::Registry;
use crate::types::{has_relation_predicate, is_truthy};
use super::cmd::{cmd_count_documents, num_value};
use super::query::{plan_query_ast_mut, resolve_page, Page, QueryPlan};
#[derive(Debug, Clone)]
pub struct CountQueryPlan {
pub query: QueryPlan,
pub page: Page,
pub count_command: Value,
}
impl CountQueryPlan {
pub fn to_value(&self) -> Value {
let mut m = self
.query
.to_value()
.as_object()
.cloned()
.unwrap_or_default();
m.insert("countCommand".to_string(), self.count_command.clone());
if let Some(o) = self.page.to_value().as_object() {
for (k, v) in o {
m.insert(k.clone(), v.clone());
}
}
Value::Object(m)
}
pub fn has_more(&self, total: f64) -> bool {
self.page.has_more(total)
}
}
pub fn plan_query_with_count(
gql: &str,
params: &Map<String, Value>,
registry: &Registry,
ctx: Option<&Context>,
) -> Result<CountQueryPlan, String> {
let mut ast = parse_gql(gql)?;
let schema = registry.get(&ast.model)?;
let mut params = params.clone();
let page = resolve_page(&ast, ¶ms);
let wants_pagination = params.contains_key("page") || params.contains_key("pageSize");
if wants_pagination && !ast.params.contains_key("skip") {
ast.params
.insert("skip".to_string(), "@__core_page_skip__".to_string());
params.insert(
"__core_page_skip__".to_string(),
num_value(page.page * page.page_size),
);
}
if wants_pagination && !ast.params.contains_key("limit") {
ast.params
.insert("limit".to_string(), "@__core_page_limit__".to_string());
params.insert("__core_page_limit__".to_string(), num_value(page.page_size));
}
if let Some(r) = ast.params.get("skip").cloned() {
let key = r.get(1..).unwrap_or("").to_string();
params.insert(key, num_value(page.page * page.page_size));
}
if let Some(r) = ast.params.get("limit").cloned() {
let key = r.get(1..).unwrap_or("").to_string();
params.insert(key, num_value(page.page_size));
}
let count_filter = ast
.params
.get("condition")
.and_then(|r| param(¶ms, Some(r)).cloned())
.filter(is_truthy)
.unwrap_or_else(|| json!({}));
let owner_merged =
merge_owner_condition(registry.role_rules(), schema, ctx, Some(count_filter));
let count_filter = merge_row_condition(registry, schema, ctx, "read", owner_merged)
.unwrap_or_else(|| json!({}));
if has_relation_predicate(schema, &count_filter) {
return Err(
"query_with_count 暂不支持关系聚合谓词条件(§9.6):total 无法与 items 保持一致,\
请改用 query 并在调用方自行统计"
.to_string(),
);
}
let query = plan_query_ast_mut(ast, &mut params, registry, ctx)?;
let count_command = cmd_count_documents(schema, &count_filter);
Ok(CountQueryPlan {
query,
page,
count_command,
})
}