dora_core/descriptor/
mod.rs

1use dora_message::{
2    config::{Input, InputMapping, NodeRunConfig},
3    descriptor::{GitRepoRev, NodeSource},
4    id::{DataId, NodeId, OperatorId},
5};
6use eyre::{bail, Context, OptionExt, Result};
7use std::{
8    collections::{BTreeMap, HashMap},
9    env::consts::EXE_EXTENSION,
10    path::{Path, PathBuf},
11    process::Stdio,
12};
13use tokio::process::Command;
14
15// reexport for compatibility
16pub use dora_message::descriptor::{
17    CoreNodeKind, CustomNode, Descriptor, Node, OperatorConfig, OperatorDefinition, OperatorSource,
18    PythonSource, ResolvedNode, RuntimeNode, SingleOperatorDefinition, DYNAMIC_SOURCE,
19    SHELL_SOURCE,
20};
21pub use validate::ResolvedNodeExt;
22pub use visualize::collect_dora_timers;
23
24mod validate;
25mod visualize;
26
27pub trait DescriptorExt {
28    fn resolve_aliases_and_set_defaults(&self) -> eyre::Result<BTreeMap<NodeId, ResolvedNode>>;
29    fn visualize_as_mermaid(&self) -> eyre::Result<String>;
30    fn blocking_read(path: &Path) -> eyre::Result<Descriptor>;
31    fn parse(buf: Vec<u8>) -> eyre::Result<Descriptor>;
32    fn check(&self, working_dir: &Path) -> eyre::Result<()>;
33    fn check_in_daemon(&self, working_dir: &Path, coordinator_is_remote: bool) -> eyre::Result<()>;
34}
35
36pub const SINGLE_OPERATOR_DEFAULT_ID: &str = "op";
37
38impl DescriptorExt for Descriptor {
39    fn resolve_aliases_and_set_defaults(&self) -> eyre::Result<BTreeMap<NodeId, ResolvedNode>> {
40        let default_op_id = OperatorId::from(SINGLE_OPERATOR_DEFAULT_ID.to_string());
41
42        let single_operator_nodes: HashMap<_, _> = self
43            .nodes
44            .iter()
45            .filter_map(|n| {
46                n.operator
47                    .as_ref()
48                    .map(|op| (&n.id, op.id.as_ref().unwrap_or(&default_op_id)))
49            })
50            .collect();
51
52        let mut resolved = BTreeMap::new();
53        for mut node in self.nodes.clone() {
54            // adjust input mappings
55            let mut node_kind = node_kind_mut(&mut node)?;
56            let input_mappings: Vec<_> = match &mut node_kind {
57                NodeKindMut::Standard { inputs, .. } => inputs.values_mut().collect(),
58                NodeKindMut::Runtime(node) => node
59                    .operators
60                    .iter_mut()
61                    .flat_map(|op| op.config.inputs.values_mut())
62                    .collect(),
63                NodeKindMut::Custom(node) => node.run_config.inputs.values_mut().collect(),
64                NodeKindMut::Operator(operator) => operator.config.inputs.values_mut().collect(),
65            };
66            for mapping in input_mappings
67                .into_iter()
68                .filter_map(|i| match &mut i.mapping {
69                    InputMapping::Timer { .. } => None,
70                    InputMapping::User(m) => Some(m),
71                })
72            {
73                if let Some(op_name) = single_operator_nodes.get(&mapping.source).copied() {
74                    mapping.output = DataId::from(format!("{op_name}/{}", mapping.output));
75                }
76            }
77
78            // resolve nodes
79            let kind = match node_kind {
80                NodeKindMut::Standard {
81                    path,
82                    source,
83                    inputs: _,
84                } => CoreNodeKind::Custom(CustomNode {
85                    path: path.clone(),
86                    source,
87                    args: node.args,
88                    build: node.build,
89                    send_stdout_as: node.send_stdout_as,
90                    run_config: NodeRunConfig {
91                        inputs: node.inputs,
92                        outputs: node.outputs,
93                    },
94                    envs: None,
95                }),
96                NodeKindMut::Custom(node) => CoreNodeKind::Custom(node.clone()),
97                NodeKindMut::Runtime(node) => CoreNodeKind::Runtime(node.clone()),
98                NodeKindMut::Operator(op) => CoreNodeKind::Runtime(RuntimeNode {
99                    operators: vec![OperatorDefinition {
100                        id: op.id.clone().unwrap_or_else(|| default_op_id.clone()),
101                        config: op.config.clone(),
102                    }],
103                }),
104            };
105
106            resolved.insert(
107                node.id.clone(),
108                ResolvedNode {
109                    id: node.id,
110                    name: node.name,
111                    description: node.description,
112                    env: node.env,
113                    deploy: node.deploy,
114                    kind,
115                },
116            );
117        }
118
119        Ok(resolved)
120    }
121
122    fn visualize_as_mermaid(&self) -> eyre::Result<String> {
123        let resolved = self.resolve_aliases_and_set_defaults()?;
124        let flowchart = visualize::visualize_nodes(&resolved);
125
126        Ok(flowchart)
127    }
128
129    fn blocking_read(path: &Path) -> eyre::Result<Descriptor> {
130        let buf = std::fs::read(path).context("failed to open given file")?;
131        Descriptor::parse(buf)
132    }
133
134    fn parse(buf: Vec<u8>) -> eyre::Result<Descriptor> {
135        serde_yaml::from_slice(&buf).context("failed to parse given descriptor")
136    }
137
138    fn check(&self, working_dir: &Path) -> eyre::Result<()> {
139        validate::check_dataflow(self, working_dir, None, false)
140            .wrap_err("Dataflow could not be validated.")
141    }
142
143    fn check_in_daemon(&self, working_dir: &Path, coordinator_is_remote: bool) -> eyre::Result<()> {
144        validate::check_dataflow(self, working_dir, None, coordinator_is_remote)
145            .wrap_err("Dataflow could not be validated.")
146    }
147}
148
149pub async fn read_as_descriptor(path: &Path) -> eyre::Result<Descriptor> {
150    let buf = tokio::fs::read(path)
151        .await
152        .context("failed to open given file")?;
153    Descriptor::parse(buf)
154}
155
156fn node_kind_mut(node: &mut Node) -> eyre::Result<NodeKindMut> {
157    match node.kind()? {
158        NodeKind::Standard(_) => {
159            let source = match (&node.git, &node.branch, &node.tag, &node.rev) {
160                (None, None, None, None) => NodeSource::Local,
161                (Some(repo), branch, tag, rev) => {
162                    let rev = match (branch, tag, rev) {
163                        (None, None, None) => None,
164                        (Some(branch), None, None) => Some(GitRepoRev::Branch(branch.clone())),
165                        (None, Some(tag), None) => Some(GitRepoRev::Tag(tag.clone())),
166                        (None, None, Some(rev)) => Some(GitRepoRev::Rev(rev.clone())),
167                        other @ (_, _, _) => {
168                            eyre::bail!("only one of `branch`, `tag`, and `rev` are allowed (got {other:?})")
169                        }
170                    };
171                    NodeSource::GitBranch {
172                        repo: repo.clone(),
173                        rev,
174                    }
175                }
176                (None, _, _, _) => {
177                    eyre::bail!("`git` source required when using branch, tag, or rev")
178                }
179            };
180
181            Ok(NodeKindMut::Standard {
182                path: node.path.as_ref().ok_or_eyre("missing `path` attribute")?,
183                source,
184                inputs: &mut node.inputs,
185            })
186        }
187        NodeKind::Runtime(_) => node
188            .operators
189            .as_mut()
190            .map(NodeKindMut::Runtime)
191            .ok_or_eyre("no operators"),
192        NodeKind::Custom(_) => node
193            .custom
194            .as_mut()
195            .map(NodeKindMut::Custom)
196            .ok_or_eyre("no custom"),
197        NodeKind::Operator(_) => node
198            .operator
199            .as_mut()
200            .map(NodeKindMut::Operator)
201            .ok_or_eyre("no operator"),
202    }
203}
204
205pub fn source_is_url(source: &str) -> bool {
206    source.contains("://")
207}
208
209pub fn resolve_path(source: &str, working_dir: &Path) -> Result<PathBuf> {
210    let path = Path::new(&source);
211    let path = if path.extension().is_none() {
212        path.with_extension(EXE_EXTENSION)
213    } else {
214        path.to_owned()
215    };
216
217    // Search path within current working directory
218    if let Ok(abs_path) = working_dir.join(&path).canonicalize() {
219        Ok(abs_path)
220    // Search path within $PATH
221    } else if which::which("uv").is_ok() {
222        // spawn: uv run which <path>
223        let which = if cfg!(windows) { "where" } else { "which" };
224        let _output = Command::new("uv")
225            .arg("run")
226            .arg(which)
227            .arg(&path)
228            .stdout(Stdio::null())
229            .spawn()
230            .context("Could not find binary within uv")?;
231        Ok(path)
232    } else if let Ok(abs_path) = which::which(&path) {
233        Ok(abs_path)
234    } else {
235        bail!("Could not find source path {}", path.display())
236    }
237}
238
239pub trait NodeExt {
240    fn kind(&self) -> eyre::Result<NodeKind>;
241}
242
243impl NodeExt for Node {
244    fn kind(&self) -> eyre::Result<NodeKind> {
245        match (&self.path, &self.operators, &self.custom, &self.operator) {
246            (None, None, None, None) => {
247                eyre::bail!(
248                    "node `{}` requires a `path`, `custom`, or `operators` field",
249                    self.id
250                )
251            }
252            (None, None, None, Some(operator)) => Ok(NodeKind::Operator(operator)),
253            (None, None, Some(custom), None) => Ok(NodeKind::Custom(custom)),
254            (None, Some(runtime), None, None) => Ok(NodeKind::Runtime(runtime)),
255            (Some(path), None, None, None) => Ok(NodeKind::Standard(path)),
256            _ => {
257                eyre::bail!(
258                    "node `{}` has multiple exclusive fields set, only one of `path`, `custom`, `operators` and `operator` is allowed",
259                    self.id
260                )
261            }
262        }
263    }
264}
265
266#[derive(Debug)]
267pub enum NodeKind<'a> {
268    Standard(&'a String),
269    /// Dora runtime node
270    Runtime(&'a RuntimeNode),
271    Custom(&'a CustomNode),
272    Operator(&'a SingleOperatorDefinition),
273}
274
275#[derive(Debug)]
276enum NodeKindMut<'a> {
277    Standard {
278        path: &'a String,
279        source: NodeSource,
280        inputs: &'a mut BTreeMap<DataId, Input>,
281    },
282    /// Dora runtime node
283    Runtime(&'a mut RuntimeNode),
284    Custom(&'a mut CustomNode),
285    Operator(&'a mut SingleOperatorDefinition),
286}