vector/
graph.rs

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    /// Read configuration from one or more files. Wildcard paths are supported.
20    /// File format is detected from the file name.
21    /// If zero files are specified the default config path
22    /// `/etc/vector/vector.yaml` will be targeted.
23    #[arg(
24        id = "config",
25        short,
26        long,
27        env = "VECTOR_CONFIG",
28        value_delimiter(',')
29    )]
30    paths: Vec<PathBuf>,
31
32    /// Vector config files in TOML format.
33    #[arg(id = "config-toml", long, value_delimiter(','))]
34    paths_toml: Vec<PathBuf>,
35
36    /// Vector config files in JSON format.
37    #[arg(id = "config-json", long, value_delimiter(','))]
38    paths_json: Vec<PathBuf>,
39
40    /// Vector config files in YAML format.
41    #[arg(id = "config-yaml", long, value_delimiter(','))]
42    paths_yaml: Vec<PathBuf>,
43
44    /// Read configuration from files in one or more directories.
45    /// File format is detected from the file name.
46    ///
47    /// Files not ending in .toml, .json, .yaml, or .yml will be ignored.
48    #[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    /// Set the output format
58    ///
59    /// See https://mermaid.js.org/syntax/flowchart.html#styling-and-classes for
60    /// information on the `mermaid` format.
61    #[arg(id = "format", long, default_value = "dot")]
62    pub format: OutputFormat,
63
64    /// Disable interpolation of environment variables in configuration files.
65    #[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}