1use std::{
2 collections::{HashMap, HashSet},
3 fmt::Write as _,
4 path::PathBuf,
5};
6
7use clap::Parser;
8use itertools::Itertools;
9use vector_lib::{config::OutputId, id::ComponentKey};
10
11use crate::config::{
12 self,
13 dot_graph::{EdgeAttributes, GraphConfig},
14};
15
16#[derive(Parser, Debug)]
17#[command(rename_all = "kebab-case")]
18pub struct Opts {
19 #[arg(
24 id = "config",
25 short,
26 long,
27 env = "VECTOR_CONFIG",
28 value_delimiter(',')
29 )]
30 paths: Vec<PathBuf>,
31
32 #[arg(id = "config-toml", long, value_delimiter(','))]
34 paths_toml: Vec<PathBuf>,
35
36 #[arg(id = "config-json", long, value_delimiter(','))]
38 paths_json: Vec<PathBuf>,
39
40 #[arg(id = "config-yaml", long, value_delimiter(','))]
42 paths_yaml: Vec<PathBuf>,
43
44 #[arg(
49 id = "config-dir",
50 short = 'C',
51 long,
52 env = "VECTOR_CONFIG_DIR",
53 value_delimiter(',')
54 )]
55 pub config_dirs: Vec<PathBuf>,
56
57 #[arg(id = "format", long, default_value = "dot")]
62 pub format: OutputFormat,
63
64 #[arg(
66 long,
67 env = "VECTOR_DISABLE_ENV_VAR_INTERPOLATION",
68 default_value = "false"
69 )]
70 pub disable_env_var_interpolation: bool,
71}
72
73#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
74pub enum OutputFormat {
75 Dot,
76 Mermaid,
77}
78
79impl Opts {
80 fn paths_with_formats(&self) -> Vec<config::ConfigPath> {
81 config::merge_path_lists(vec![
82 (&self.paths, None),
83 (&self.paths_toml, Some(config::Format::Toml)),
84 (&self.paths_json, Some(config::Format::Json)),
85 (&self.paths_yaml, Some(config::Format::Yaml)),
86 ])
87 .map(|(path, hint)| config::ConfigPath::File(path, hint))
88 .chain(
89 self.config_dirs
90 .iter()
91 .map(|dir| config::ConfigPath::Dir(dir.to_path_buf())),
92 )
93 .collect()
94 }
95}
96
97fn node_attributes_to_string(attributes: &HashMap<String, String>, default_shape: &str) -> String {
98 let mut attrs = attributes.clone();
99 if !attrs.contains_key("shape") {
100 attrs.insert("shape".to_string(), default_shape.to_string());
101 }
102 attrs.iter().map(|(k, v)| format!("{k}=\"{v}\"")).join(" ")
103}
104
105fn edge_attributes_to_string(attributes: &EdgeAttributes, default_label: Option<&str>) -> String {
106 let mut attrs = attributes.0.clone();
107 if let Some(default_label) = default_label
108 && !attrs.contains_key("label")
109 {
110 attrs.insert("label".to_string(), default_label.to_string());
111 }
112 attrs.iter().map(|(k, v)| format!("{k}=\"{v}\"")).join(" ")
113}
114
115pub(crate) fn cmd(opts: &Opts) -> exitcode::ExitCode {
116 let paths = opts.paths_with_formats();
117 let paths = match config::process_paths(&paths) {
118 Some(paths) => paths,
119 None => return exitcode::CONFIG,
120 };
121
122 let config = match config::load_from_paths(&paths, !opts.disable_env_var_interpolation) {
123 Ok(config) => config,
124 Err(errs) => {
125 #[allow(clippy::print_stderr)]
126 for err in errs {
127 eprintln!("{err}");
128 }
129 return exitcode::CONFIG;
130 }
131 };
132
133 let format = opts.format;
134 match format {
135 OutputFormat::Dot => render_dot(config),
136 OutputFormat::Mermaid => render_mermaid(config),
137 }
138}
139
140fn render_dot(config: config::Config) -> exitcode::ExitCode {
141 let mut dot = String::from("digraph {\n");
142
143 let mut written_tables = HashSet::<ComponentKey>::new();
144
145 for (id, table) in config
146 .enrichment_tables
147 .iter()
148 .filter_map(|(key, table)| table.as_source(key))
149 {
150 writeln!(
151 dot,
152 " \"{}\" [{}]",
153 id,
154 node_attributes_to_string(&table.graph.node_attributes, "cylinder")
155 )
156 .expect("write to String never fails");
157 written_tables.insert(id);
158 }
159
160 for (id, table) in config
161 .enrichment_tables
162 .iter()
163 .filter_map(|(key, table)| table.as_sink(key))
164 {
165 if !written_tables.contains(&id) {
166 writeln!(
167 dot,
168 " \"{}\" [{}]",
169 id,
170 node_attributes_to_string(&table.graph.node_attributes, "cylinder")
171 )
172 .expect("write to String never fails");
173 }
174
175 for input in table.inputs.iter() {
176 render_dot_edge(&mut dot, &id, input, &table.graph);
177 }
178 }
179
180 for (id, source) in config.sources() {
181 writeln!(
182 dot,
183 " \"{}\" [{}]",
184 id,
185 node_attributes_to_string(&source.graph.node_attributes, "trapezium")
186 )
187 .expect("write to String never fails");
188 }
189
190 for (id, transform) in config.transforms() {
191 writeln!(
192 dot,
193 " \"{}\" [{}]",
194 id,
195 node_attributes_to_string(&transform.graph.node_attributes, "diamond")
196 )
197 .expect("write to String never fails");
198
199 for input in transform.inputs.iter() {
200 render_dot_edge(&mut dot, id, input, &transform.graph);
201 }
202 }
203
204 for (id, sink) in config.sinks() {
205 writeln!(
206 dot,
207 " \"{}\" [{}]",
208 id,
209 node_attributes_to_string(&sink.graph.node_attributes, "invtrapezium")
210 )
211 .expect("write to String never fails");
212
213 for input in &sink.inputs {
214 render_dot_edge(&mut dot, id, input, &sink.graph);
215 }
216 }
217
218 dot += "}";
219
220 #[allow(clippy::print_stdout)]
221 {
222 println!("{dot}");
223 }
224
225 exitcode::OK
226}
227
228fn render_dot_edge(into: &mut String, id: &ComponentKey, input: &OutputId, graph: &GraphConfig) {
229 let edge_attributes = graph
230 .edge_attributes
231 .get(&input.to_string())
232 .or_else(|| graph.edge_attributes.get(&input.component.to_string()));
233 if let Some(port) = &input.port {
234 writeln!(
235 into,
236 " \"{}\" -> \"{id}\" [{}]",
237 input.component,
238 edge_attributes_to_string(
239 edge_attributes.unwrap_or(&EdgeAttributes::default()),
240 Some(port)
241 )
242 )
243 .expect("write to String never fails");
244 } else if let Some(edge_attributes) = edge_attributes {
245 writeln!(
246 into,
247 " \"{input}\" -> \"{id}\" [{}]",
248 edge_attributes_to_string(edge_attributes, None)
249 )
250 .expect("write to String never fails");
251 } else {
252 writeln!(into, " \"{input}\" -> \"{id}\"").expect("write to String never fails");
253 }
254}
255
256fn render_mermaid(config: config::Config) -> exitcode::ExitCode {
257 let mut mermaid = String::from("flowchart TD;\n");
258
259 writeln!(mermaid, "\n %% Enrichment tables").unwrap();
260 let mut written_tables = HashSet::<ComponentKey>::new();
261
262 for (id, _) in config
263 .enrichment_tables
264 .iter()
265 .filter_map(|(key, table)| table.as_source(key))
266 {
267 writeln!(mermaid, " {id}[({id})]").unwrap();
268 written_tables.insert(id);
269 }
270
271 for (id, table) in config
272 .enrichment_tables
273 .iter()
274 .filter_map(|(key, table)| table.as_sink(key))
275 {
276 if !written_tables.contains(&id) {
277 writeln!(mermaid, " {id}[({id})]").unwrap();
278 }
279
280 for input in table.inputs.iter() {
281 if let Some(port) = &input.port {
282 writeln!(mermaid, " {0} -->|{port}| {id}", input.component).unwrap();
283 } else {
284 writeln!(mermaid, " {0} --> {id}", input.component).unwrap();
285 }
286 }
287 }
288
289 writeln!(mermaid, "\n %% Sources").unwrap();
290 for (id, _) in config.sources() {
291 writeln!(mermaid, " {id}[/{id}/]").unwrap();
292 }
293
294 writeln!(mermaid, "\n %% Transforms").unwrap();
295 for (id, transform) in config.transforms() {
296 writeln!(mermaid, " {id}{{{id}}}").unwrap();
297
298 for input in transform.inputs.iter() {
299 if let Some(port) = &input.port {
300 writeln!(mermaid, " {0} -->|{port}| {id}", input.component).unwrap();
301 } else {
302 writeln!(mermaid, " {0} --> {id}", input.component).unwrap();
303 }
304 }
305 }
306
307 writeln!(mermaid, "\n %% Sinks").unwrap();
308 for (id, sink) in config.sinks() {
309 writeln!(mermaid, " {id}[\\{id}\\]").unwrap();
310
311 for input in &sink.inputs {
312 if let Some(port) = &input.port {
313 writeln!(mermaid, " {0} -->|{port}| {id}", input.component).unwrap();
314 } else {
315 writeln!(mermaid, " {0} --> {id}", input.component).unwrap();
316 }
317 }
318 }
319
320 #[allow(clippy::print_stdout)]
321 {
322 println!("{mermaid}");
323 }
324
325 exitcode::OK
326}