use std::{fmt::Display, str::FromStr};
use anyhow::Result;
use serde::Serialize;
use super::check::load;
use crate::output::CommandResult;
const COMMAND: &str = "federation graph";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum GraphFormat {
Json,
Dot,
Mermaid,
}
impl FromStr for GraphFormat {
type Err = String;
fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"json" => Ok(GraphFormat::Json),
"dot" => Ok(GraphFormat::Dot),
"mermaid" => Ok(GraphFormat::Mermaid),
other => Err(format!("Unknown format: {other}. Use json, dot, or mermaid")),
}
}
}
impl Display for GraphFormat {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
GraphFormat::Json => write!(f, "json"),
GraphFormat::Dot => write!(f, "dot"),
GraphFormat::Mermaid => write!(f, "mermaid"),
}
}
}
#[derive(Debug, Serialize)]
pub struct FederationGraph {
pub subgraphs: Vec<Subgraph>,
pub edges: Vec<Edge>,
pub unresolved: Vec<Unresolved>,
}
#[derive(Debug, Serialize)]
pub struct Subgraph {
pub name: String,
pub owns: Vec<String>,
pub extends: Vec<String>,
}
#[derive(Debug, Serialize)]
pub struct Edge {
pub from: String,
pub to: String,
pub entity: String,
}
#[derive(Debug, Serialize)]
pub struct Unresolved {
pub subgraph: String,
pub entity: String,
}
pub fn run(schema_paths: &[String], format: GraphFormat) -> Result<CommandResult> {
let mut subgraphs = Vec::with_capacity(schema_paths.len());
let mut problems = Vec::new();
for path in schema_paths {
let schema = match load(path)? {
Ok(schema) => schema,
Err(problem) => {
problems.push(problem);
continue;
},
};
let Some(federation) = schema.federation.filter(|f| f.enabled) else {
problems
.push(format!("{path} declares no federation; compile it with federation enabled"));
continue;
};
let Some(name) = federation.service_name else {
problems.push(format!("{path} has no federation service_name"));
continue;
};
let (extended, owned): (Vec<_>, Vec<_>) =
federation.entities.into_iter().partition(|e| e.extends);
subgraphs.push(Subgraph {
name,
owns: owned.into_iter().map(|e| e.name).collect(),
extends: extended.into_iter().map(|e| e.name).collect(),
});
}
if !problems.is_empty() {
return Ok(CommandResult::validation_failed(COMMAND, problems, "INVALID_SCHEMA"));
}
let graph = build(subgraphs);
let output = match format {
GraphFormat::Json => serde_json::to_value(&graph)?,
GraphFormat::Dot => serde_json::Value::String(to_dot(&graph)),
GraphFormat::Mermaid => serde_json::Value::String(to_mermaid(&graph)),
};
Ok(CommandResult::success("federation/graph", output))
}
fn build(subgraphs: Vec<Subgraph>) -> FederationGraph {
let mut edges = Vec::new();
let mut unresolved = Vec::new();
for from in &subgraphs {
for entity in &from.extends {
let owners: Vec<&Subgraph> =
subgraphs.iter().filter(|s| s.owns.contains(entity)).collect();
if owners.is_empty() {
unresolved.push(Unresolved {
subgraph: from.name.clone(),
entity: entity.clone(),
});
}
for to in owners {
edges.push(Edge {
from: from.name.clone(),
to: to.name.clone(),
entity: entity.clone(),
});
}
}
}
FederationGraph {
subgraphs,
edges,
unresolved,
}
}
fn dot_quoted(s: &str) -> String {
serde_json::Value::String(s.to_owned()).to_string()
}
fn mermaid_text(s: &str) -> String {
s.replace('"', "#quot;")
.replace('<', "#lt;")
.replace('>', "#gt;")
.replace('\n', " ")
}
pub(crate) fn to_dot(graph: &FederationGraph) -> String {
let mut dot = String::from("digraph federation {\n");
let id = |name: &str| graph.subgraphs.iter().position(|s| s.name == name).unwrap_or(0);
for (i, subgraph) in graph.subgraphs.iter().enumerate() {
let label = format!("{}\n[{}]", subgraph.name, subgraph.owns.join(", "));
dot.push_str(&format!(" n{i} [label={}];\n", dot_quoted(&label)));
}
for edge in &graph.edges {
dot.push_str(&format!(
" n{} -> n{} [label={}];\n",
id(&edge.from),
id(&edge.to),
dot_quoted(&edge.entity)
));
}
dot.push_str("}\n");
dot
}
pub(crate) fn to_mermaid(graph: &FederationGraph) -> String {
let mut mermaid = String::from("graph LR\n");
let id = |name: &str| graph.subgraphs.iter().position(|s| s.name == name).unwrap_or(0);
for (i, subgraph) in graph.subgraphs.iter().enumerate() {
let owns: Vec<String> = subgraph.owns.iter().map(|e| mermaid_text(e)).collect();
mermaid.push_str(&format!(
" n{i}[\"{}<br/>[{}]\"]\n",
mermaid_text(&subgraph.name),
owns.join("<br/>")
));
}
for edge in &graph.edges {
mermaid.push_str(&format!(
" n{} -->|\"{}\"| n{}\n",
id(&edge.from),
mermaid_text(&edge.entity),
id(&edge.to)
));
}
mermaid
}