vector/
graph.rs

1use std::collections::HashMap;
2use std::fmt::Write as _;
3use std::path::PathBuf;
4
5use clap::Parser;
6use itertools::Itertools;
7
8use crate::config;
9
10#[derive(Parser, Debug)]
11#[command(rename_all = "kebab-case")]
12pub struct Opts {
13    /// Read configuration from one or more files. Wildcard paths are supported.
14    /// File format is detected from the file name.
15    /// If zero files are specified the default config path
16    /// `/etc/vector/vector.yaml` will be targeted.
17    #[arg(
18        id = "config",
19        short,
20        long,
21        env = "VECTOR_CONFIG",
22        value_delimiter(',')
23    )]
24    paths: Vec<PathBuf>,
25
26    /// Vector config files in TOML format.
27    #[arg(id = "config-toml", long, value_delimiter(','))]
28    paths_toml: Vec<PathBuf>,
29
30    /// Vector config files in JSON format.
31    #[arg(id = "config-json", long, value_delimiter(','))]
32    paths_json: Vec<PathBuf>,
33
34    /// Vector config files in YAML format.
35    #[arg(id = "config-yaml", long, value_delimiter(','))]
36    paths_yaml: Vec<PathBuf>,
37
38    /// Read configuration from files in one or more directories.
39    /// File format is detected from the file name.
40    ///
41    /// Files not ending in .toml, .json, .yaml, or .yml will be ignored.
42    #[arg(
43        id = "config-dir",
44        short = 'C',
45        long,
46        env = "VECTOR_CONFIG_DIR",
47        value_delimiter(',')
48    )]
49    pub config_dirs: Vec<PathBuf>,
50
51    /// Set the output format
52    ///
53    /// See https://mermaid.js.org/syntax/flowchart.html#styling-and-classes for
54    /// information on the `mermaid` format.
55    #[arg(id = "format", long, default_value = "dot")]
56    pub format: OutputFormat,
57}
58
59#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
60pub enum OutputFormat {
61    Dot,
62    Mermaid,
63}
64
65impl Opts {
66    fn paths_with_formats(&self) -> Vec<config::ConfigPath> {
67        config::merge_path_lists(vec![
68            (&self.paths, None),
69            (&self.paths_toml, Some(config::Format::Toml)),
70            (&self.paths_json, Some(config::Format::Json)),
71            (&self.paths_yaml, Some(config::Format::Yaml)),
72        ])
73        .map(|(path, hint)| config::ConfigPath::File(path, hint))
74        .chain(
75            self.config_dirs
76                .iter()
77                .map(|dir| config::ConfigPath::Dir(dir.to_path_buf())),
78        )
79        .collect()
80    }
81}
82
83fn node_attributes_to_string(attributes: &HashMap<String, String>, default_shape: &str) -> String {
84    let mut attrs = attributes.clone();
85    if !attrs.contains_key("shape") {
86        attrs.insert("shape".to_string(), default_shape.to_string());
87    }
88    attrs.iter().map(|(k, v)| format!("{k}=\"{v}\"")).join(" ")
89}
90
91pub(crate) fn cmd(opts: &Opts) -> exitcode::ExitCode {
92    let paths = opts.paths_with_formats();
93    let paths = match config::process_paths(&paths) {
94        Some(paths) => paths,
95        None => return exitcode::CONFIG,
96    };
97
98    let config = match config::load_from_paths(&paths) {
99        Ok(config) => config,
100        Err(errs) => {
101            #[allow(clippy::print_stderr)]
102            for err in errs {
103                eprintln!("{err}");
104            }
105            return exitcode::CONFIG;
106        }
107    };
108
109    let format = opts.format;
110    match format {
111        OutputFormat::Dot => render_dot(config),
112        OutputFormat::Mermaid => render_mermaid(config),
113    }
114}
115
116fn render_dot(config: config::Config) -> exitcode::ExitCode {
117    let mut dot = String::from("digraph {\n");
118
119    for (id, source) in config.sources() {
120        writeln!(
121            dot,
122            "  \"{}\" [{}]",
123            id,
124            node_attributes_to_string(&source.graph.node_attributes, "trapezium")
125        )
126        .expect("write to String never fails");
127    }
128
129    for (id, transform) in config.transforms() {
130        writeln!(
131            dot,
132            "  \"{}\" [{}]",
133            id,
134            node_attributes_to_string(&transform.graph.node_attributes, "diamond")
135        )
136        .expect("write to String never fails");
137
138        for input in transform.inputs.iter() {
139            if let Some(port) = &input.port {
140                writeln!(
141                    dot,
142                    "  \"{}\" -> \"{}\" [label=\"{}\"]",
143                    input.component, id, port
144                )
145                .expect("write to String never fails");
146            } else {
147                writeln!(dot, "  \"{input}\" -> \"{id}\"").expect("write to String never fails");
148            }
149        }
150    }
151
152    for (id, sink) in config.sinks() {
153        writeln!(
154            dot,
155            "  \"{}\" [{}]",
156            id,
157            node_attributes_to_string(&sink.graph.node_attributes, "invtrapezium")
158        )
159        .expect("write to String never fails");
160
161        for input in &sink.inputs {
162            if let Some(port) = &input.port {
163                writeln!(
164                    dot,
165                    "  \"{}\" -> \"{}\" [label=\"{}\"]",
166                    input.component, id, port
167                )
168                .expect("write to String never fails");
169            } else {
170                writeln!(dot, "  \"{input}\" -> \"{id}\"").expect("write to String never fails");
171            }
172        }
173    }
174
175    dot += "}";
176
177    #[allow(clippy::print_stdout)]
178    {
179        println!("{dot}");
180    }
181
182    exitcode::OK
183}
184
185fn render_mermaid(config: config::Config) -> exitcode::ExitCode {
186    let mut mermaid = String::from("flowchart TD;\n");
187
188    writeln!(mermaid, "\n  %% Sources").unwrap();
189    for (id, _) in config.sources() {
190        writeln!(mermaid, "  {id}[/{id}/]").unwrap();
191    }
192
193    writeln!(mermaid, "\n  %% Transforms").unwrap();
194    for (id, transform) in config.transforms() {
195        writeln!(mermaid, "  {id}{{{id}}}").unwrap();
196
197        for input in transform.inputs.iter() {
198            if let Some(port) = &input.port {
199                writeln!(mermaid, "  {0} -->|{port}| {id}", input.component).unwrap();
200            } else {
201                writeln!(mermaid, "  {0} --> {id}", input.component).unwrap();
202            }
203        }
204    }
205
206    writeln!(mermaid, "\n  %% Sinks").unwrap();
207    for (id, sink) in config.sinks() {
208        writeln!(mermaid, "  {id}[\\{id}\\]").unwrap();
209
210        for input in &sink.inputs {
211            if let Some(port) = &input.port {
212                writeln!(mermaid, "  {0} -->|{port}| {id}", input.component).unwrap();
213            } else {
214                writeln!(mermaid, "  {0} --> {id}", input.component).unwrap();
215            }
216        }
217    }
218
219    #[allow(clippy::print_stdout)]
220    {
221        println!("{mermaid}");
222    }
223
224    exitcode::OK
225}