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
15pub 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 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 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 if let Ok(abs_path) = working_dir.join(&path).canonicalize() {
219 Ok(abs_path)
220 } else if which::which("uv").is_ok() {
222 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 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 Runtime(&'a mut RuntimeNode),
284 Custom(&'a mut CustomNode),
285 Operator(&'a mut SingleOperatorDefinition),
286}