use serde_json::{json, Value};
use crate::schema::{Location, Registry, Schema};
pub const ERR_WRITE_CROSS_SOURCE_PREFIX: &str = "ERR_FEDERATION_WRITE_CROSS_SOURCE:";
pub fn err_write_cross_source(schema: &str, sources: &[String]) -> String {
format!(
"{} {} 的写链路跨连接(sources: {});同连接内跨 db/schema 才可同事务,跨连接无跨源事务",
ERR_WRITE_CROSS_SOURCE_PREFIX,
schema,
sources.join(", ")
)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum WriteLinkPolicy {
#[default]
Reject,
PrimaryOnly,
}
#[derive(Debug, Clone)]
pub struct WriteLinks {
pub source: String,
pub targets: Vec<Location>,
pub degraded: Vec<Value>,
}
pub fn resolve_write_links(
schema: &Schema,
links: &[Location],
policy: WriteLinkPolicy,
) -> Result<WriteLinks, String> {
let mut targets: Vec<Location> = Vec::new();
for l in links {
if !targets.contains(l) {
targets.push(l.clone());
}
}
if targets.is_empty() {
targets.push(schema.location());
}
let mut sources: Vec<String> = Vec::new();
for t in &targets {
if !sources.iter().any(|s| s == &t.source) {
sources.push(t.source.clone());
}
}
if sources.len() <= 1 {
return Ok(WriteLinks {
source: targets[0].source.clone(),
targets,
degraded: Vec::new(),
});
}
match policy {
WriteLinkPolicy::Reject => Err(err_write_cross_source(&schema.name, &sources)),
WriteLinkPolicy::PrimaryOnly => Ok(WriteLinks {
source: targets[0].source.clone(),
targets: vec![targets[0].clone()],
degraded: vec![json!({
"code": "writePrimaryOnly",
"layer": "federation",
"message": format!(
"{} 的写链路跨连接(sources: {});按策略 writePrimaryOnly 仅写主链路,从链路未同步",
schema.name,
sources.join(", ")
),
"hint": "将同名链路收敛到同一连接,或由上游补偿同步",
})],
}),
}
}
pub(crate) fn attach_write_links(
plan: Value,
registry: &Registry,
schemas: &[&str],
policy: WriteLinkPolicy,
) -> Result<Value, String> {
let mut declared: Vec<Location> = Vec::new();
let mut resolved: Vec<WriteLinks> = Vec::new();
for name in schemas {
let schema = registry.get(name)?;
let links = registry.links_of(name)?;
for l in links {
if !declared.contains(l) {
declared.push(l.clone());
}
}
resolved.push(resolve_write_links(schema, links, policy)?);
}
if declared.len() <= 1 {
return Ok(plan);
}
let mut targets: Vec<Location> = Vec::new();
let mut degraded: Vec<Value> = Vec::new();
for wl in &resolved {
for t in &wl.targets {
if !targets.contains(t) {
targets.push(t.clone());
}
}
degraded.extend(wl.degraded.iter().cloned());
}
let mut declared_sources: Vec<String> = Vec::new();
for t in &declared {
if !declared_sources.iter().any(|s| s == &t.source) {
declared_sources.push(t.source.clone());
}
}
let atomic = declared_sources.len() <= 1;
let source = targets
.first()
.map(|t| t.source.clone())
.unwrap_or_default();
let targets_json: Vec<Value> = targets
.iter()
.map(|t| {
json!({
"source": t.source,
"database": t.database,
"schema": t.schema,
})
})
.collect();
let mut plan = plan;
if let Some(obj) = plan.as_object_mut() {
obj.insert(
"writeLinks".to_string(),
json!({
"atomic": atomic,
"source": source,
"targets": targets_json,
"degraded": degraded,
}),
);
if !degraded.is_empty() {
let entry = obj
.entry("degraded".to_string())
.or_insert_with(|| Value::Array(Vec::new()));
match entry.as_array_mut() {
Some(arr) => arr.extend(degraded.iter().cloned()),
None => *entry = Value::Array(degraded.clone()),
}
}
}
Ok(plan)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::schema::Registry;
use serde_json::json;
fn loc(source: &str, database: Option<&str>) -> Location {
Location {
source: source.to_string(),
database: database.map(String::from),
schema: None,
role: None,
priority: None,
}
}
fn order_defn() -> Value {
json!({
"name": "Order", "collection": "orders", "timestamps": false,
"fields": { "code": { "type": "string" } }, "relations": {}
})
}
fn replica_defn() -> Value {
json!({ "name": "Order", "replica": true })
}
#[test]
fn same_source_multi_target_is_atomic() {
let mut r = Registry::new();
r.register_batch(
&[
(order_defn(), loc("mongo_main", Some("orders_db"))),
(replica_defn(), loc("mongo_main", Some("archived_db"))),
],
None,
)
.unwrap();
let schema = r.get("Order").cloned().unwrap();
let links = r.links_of("Order").unwrap().to_vec();
let out = resolve_write_links(&schema, &links, WriteLinkPolicy::Reject).unwrap();
assert_eq!(out.source, "mongo_main");
assert_eq!(out.targets.len(), 2);
assert!(out.degraded.is_empty());
}
#[test]
fn cross_source_reject_errors() {
let mut r = Registry::new();
r.register_batch(
&[
(order_defn(), loc("mongo_main", Some("orders_db"))),
(replica_defn(), loc("pg_main", Some("analytics"))),
],
None,
)
.unwrap();
let schema = r.get("Order").cloned().unwrap();
let links = r.links_of("Order").unwrap().to_vec();
let err = resolve_write_links(&schema, &links, WriteLinkPolicy::Reject)
.expect_err("跨 source 默认应拒绝");
assert!(
err.starts_with(ERR_WRITE_CROSS_SOURCE_PREFIX),
"err = {err}"
);
}
#[test]
fn cross_source_primary_only_degrades() {
let mut r = Registry::new();
r.register_batch(
&[
(order_defn(), loc("mongo_main", Some("orders_db"))),
(replica_defn(), loc("pg_main", Some("analytics"))),
],
None,
)
.unwrap();
let schema = r.get("Order").cloned().unwrap();
let links = r.links_of("Order").unwrap().to_vec();
let out = resolve_write_links(&schema, &links, WriteLinkPolicy::PrimaryOnly).unwrap();
assert_eq!(out.source, "mongo_main");
assert_eq!(out.targets.len(), 1, "PrimaryOnly 仅写主链路");
assert_eq!(out.degraded.len(), 1);
assert_eq!(out.degraded[0]["code"], "writePrimaryOnly");
}
#[test]
fn single_location_schema_plan_unchanged() {
let mut r = Registry::new();
r.register_batch(
&[(order_defn(), loc("mongo_main", Some("orders_db")))],
None,
)
.unwrap();
let plan = json!({ "command": { "kind": "insertOne" }, "returns": {} });
let out = attach_write_links(plan.clone(), &r, &["Order"], r.write_link_policy()).unwrap();
assert_eq!(out, plan, "单落点 schema 计划形状必须零变更");
assert!(out.get("writeLinks").is_none());
}
#[test]
fn multi_location_same_source_attaches_write_links() {
let mut r = Registry::new();
r.register_batch(
&[
(order_defn(), loc("mongo_main", Some("orders_db"))),
(replica_defn(), loc("mongo_main", Some("archived_db"))),
],
None,
)
.unwrap();
let plan = json!({ "command": { "kind": "insertOne" } });
let out = attach_write_links(plan, &r, &["Order"], r.write_link_policy()).unwrap();
assert_eq!(out["writeLinks"]["atomic"], true);
assert_eq!(out["writeLinks"]["source"], "mongo_main");
assert_eq!(out["writeLinks"]["targets"].as_array().unwrap().len(), 2);
}
#[test]
fn cross_source_primary_only_never_atomic() {
let mut r = Registry::new();
r.register_batch(
&[
(order_defn(), loc("mongo_main", Some("orders_db"))),
(replica_defn(), loc("pg_main", Some("analytics"))),
],
None,
)
.unwrap();
r.set_write_link_policy(WriteLinkPolicy::PrimaryOnly);
let plan = json!({ "command": { "kind": "insertOne" } });
let out = attach_write_links(plan, &r, &["Order"], r.write_link_policy()).unwrap();
assert_eq!(out["writeLinks"]["atomic"], false);
let codes: Vec<Value> = out["degraded"]
.as_array()
.unwrap()
.iter()
.map(|d| d["code"].clone())
.collect();
assert!(codes.contains(&json!("writePrimaryOnly")), "out = {out}");
}
}