use std::collections::HashMap;
use serde_json::{json, Map, Value};
use crate::command::{
build_plan, check_readable_relations, ensure_context, has_pipeline, plan_query_ast_mut,
ERR_PERMISSION, T2Q_MAX_FEDERATION_ROWS,
};
use crate::computes::{merge_depends_into_ast, InjectInfo};
use crate::datasource::DataSourceConfig;
use crate::permission::{can_read_schema, merge_owner_condition, Context};
use crate::pipeline::{flatten_object_fields, parse_gql, Ast};
use crate::rbac::{ensure_read, merge_row_condition};
use crate::schema::{Location, Profile, Registry};
use route::{detect_cross_source_sort, walk};
mod route;
mod select;
pub const FEDERATION_VERSION: u64 = 3;
struct UnitSpec {
key: String,
source: String,
database: Option<String>,
schema: Option<String>,
model: String,
ast: Ast,
depth: usize,
}
struct EdgeSpec {
parent_model: String,
rel: String,
local: String,
foreign: String,
cardinality: String,
path: Vec<String>,
key: String,
}
pub fn plan_federated(
gql: &str,
params: &Map<String, Value>,
registry: &Registry,
ctx: Option<&Context>,
ds_config: &Value,
) -> Result<Value, String> {
let ds_cfg = DataSourceConfig::from_json(ds_config)?;
let max_rows_per_source = if registry.profile() == Profile::Text2Query {
T2Q_MAX_FEDERATION_ROWS
} else {
super::MAX_FEDERATION_ROWS
};
let mut params = params.clone();
let mut ast = parse_gql(gql)?;
ensure_context(registry, ctx)?;
let root_schema = registry.get(&ast.model)?.clone();
if ctx.is_some() && !can_read_schema(registry.role_rules(), &root_schema, ctx) {
return Err(ERR_PERMISSION.to_string());
}
ensure_read(registry, &root_schema, ctx)?;
if ctx.is_some() {
if let Some(r) = ast.params.get("condition").cloned() {
let key = r.get(1..).unwrap_or("").to_string();
let owner_merged = merge_owner_condition(
registry.role_rules(),
&root_schema,
ctx,
params.get(&key).cloned(),
);
match merge_row_condition(registry, &root_schema, ctx, "read", owner_merged) {
Some(v) => {
params.insert(key, v);
}
None => {
params.remove(&key);
}
}
} else {
let static_owner =
merge_owner_condition(registry.role_rules(), &root_schema, ctx, None);
if let Some(owner) =
merge_row_condition(registry, &root_schema, ctx, "read", static_owner)
.filter(|v| v.as_object().map(|o| !o.is_empty()).unwrap_or(false))
{
ast.params
.insert("condition".to_string(), "@__core_owner__".to_string());
params.insert("__core_owner__".to_string(), owner);
}
}
}
let has_pipeline = has_pipeline(&ast, ¶ms);
let inject = if has_pipeline {
InjectInfo::default()
} else {
merge_depends_into_ast(&mut ast.relations, &root_schema)?
};
check_readable_relations(&ast.relations, &root_schema, registry, ctx)?;
let mut post_ast = ast.clone();
if !has_pipeline {
flatten_object_fields(&mut post_ast, &root_schema);
}
let mut degraded: Vec<Value> = Vec::new();
let (involved, graph_edges) =
select::collect_graph(&root_schema.name, &ast.relations, registry)?;
let mut cands: select::Candidates = HashMap::new();
let mut links: Map<String, Value> = Map::new();
for s in &involved {
let locs = select::candidates_of(registry, s)?;
links.insert(
s.clone(),
Value::Array(locs.iter().map(location_to_link_value).collect()),
);
cands.insert(s.clone(), locs);
degraded.extend(registry.resolve_read_links(s, None)?.degraded);
}
select::filter_reachable(&mut cands, None);
let selection = select::select_links(&involved, &graph_edges, &cands, &ds_cfg)?;
if selection.approximate {
degraded.push(select::degraded_approx());
}
let pinned = selection.locs;
let root_loc = pinned
.get(&root_schema.name)
.cloned()
.unwrap_or_else(|| root_schema.location());
let mut units: Vec<UnitSpec> = Vec::new();
let mut edges: Vec<EdgeSpec> = Vec::new();
let mut fetch_ast = ast;
walk(
&fetch_ast.model,
&ds_cfg,
&pinned,
&mut fetch_ast.fields,
&mut fetch_ast.relations,
&[],
¶ms,
registry,
&mut units,
&mut edges,
&mut degraded,
)?;
detect_cross_source_sort(&fetch_ast, ¶ms, &edges, &mut degraded);
if has_pipeline && !edges.is_empty() {
return Err("联邦查询不支持用户 $pipeline(无法跨源下推)".to_string());
}
edges.sort_by_key(|e| e.path.len());
let root_plan = build_plan(fetch_ast, &post_ast, &inject, ¶ms, registry, ctx)?;
let mut sources = vec![json!({
"key": "0",
"source": root_loc.source,
"database": root_loc.database.map(|s| json!(s)).unwrap_or(Value::Null),
"schema": root_loc.schema.map(|s| json!(s)).unwrap_or(Value::Null),
"model": root_schema.name,
"mode": root_plan.mode.as_str(),
"commands": root_plan.commands,
"sort": root_plan.sort.clone().unwrap_or(Value::Null),
})];
let mut unit_order: Vec<usize> = (0..units.len()).collect();
unit_order.sort_by_key(|&i| units[i].depth);
for &i in &unit_order {
let unit = &units[i];
let mut unit_params = params.clone();
let plan = plan_query_ast_mut(unit.ast.clone(), &mut unit_params, registry, ctx)?;
sources.push(json!({
"key": unit.key,
"source": unit.source,
"database": unit.database.as_ref().map(|s| json!(s)).unwrap_or(Value::Null),
"schema": unit.schema.as_ref().map(|s| json!(s)).unwrap_or(Value::Null),
"model": unit.model,
"mode": plan.mode.as_str(),
"commands": plan.commands,
"sort": plan.sort.clone().unwrap_or(Value::Null),
}));
}
let join_edges: Vec<Value> = edges
.iter()
.map(|e| {
json!({
"parent": e.parent_model,
"rel": e.rel,
"local": e.local,
"foreign": e.foreign,
"cardinality": e.cardinality,
"path": e.path,
"key": e.key,
})
})
.collect();
Ok(json!({
"v": FEDERATION_VERSION,
"kind": "federated",
"root": root_schema.name,
"sources": sources,
"join": { "type": "hash", "edges": join_edges },
"postprocess": root_plan.postprocess.clone().unwrap_or(Value::Null),
"degraded": degraded,
"links": links,
"maxRowsPerSource": max_rows_per_source,
}))
}
fn location_to_link_value(l: &Location) -> Value {
json!({
"source": l.source,
"database": l.database,
"schema": l.schema,
"role": l.role.as_ref().map(|r| r.as_str()),
"priority": l.priority,
})
}
pub(crate) fn unit_index(sources: &[Value]) -> HashMap<String, usize> {
sources
.iter()
.enumerate()
.filter_map(|(i, s)| {
s.get("key")
.and_then(|k| k.as_str())
.map(|k| (k.to_string(), i))
})
.collect()
}