Add new send_stdout_as key for capturing stdout of custom nodes

dora-rs · Feb 12, 2024 · e85ce44 · e85ce44
1 parent 39fbb6b
commit e85ce44
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 6 deletions.
diff --git a/binaries/daemon/src/spawn.rs b/binaries/daemon/src/spawn.rs
@@ -58,8 +58,7 @@ pub async fn spawn_node(
         clock.clone(),
     )
     .await?;
-    let outputs = node_outputs(&node);
-    let log_output = outputs.contains(&DataId::from("op/logs".to_string()));
+    let send_stdout_to = node.send_stdout_as().map(ToOwned::to_owned);
 
     let mut child = match node.kind {
         dora_core::descriptor::CoreNodeKind::Custom(n) => {
@@ -333,19 +332,21 @@ pub async fn spawn_node(
     tokio::spawn(async move {
         while let Some(message) = rx.recv().await {
             // If log is an output, we're sending the logs to the dataflow
-            if log_output {
+            if let Some(stdout_output_name) = &send_stdout_to {
                 // Convert logs to DataMessage
                 let array = message.into_arrow();
 
                 let array: ArrayData = array.into();
                 let total_len = required_data_size(&array);
-                let mut sample: AVec<u8, ConstAlign<128>> =
-                    AVec::__from_elem(128, 0, total_len as usize);
+                let mut sample: AVec<u8, ConstAlign<128>> = AVec::__from_elem(128, 0, total_len);
 
                 let type_info = copy_array_into_sample(&mut sample, &array);
 
                 let metadata = Metadata::new(uhlc.new_timestamp(), type_info);
-                let output_id = OutputId(node_id.clone(), DataId::from("op/logs".to_string()));
+                let output_id = OutputId(
+                    node_id.clone(),
+                    DataId::from(stdout_output_name.to_string()),
+                );
                 let event = DoraEvent::Logs {
                     dataflow_id,
                     output_id,

diff --git a/libraries/core/src/descriptor/mod.rs b/libraries/core/src/descriptor/mod.rs
@@ -164,6 +164,15 @@ pub struct ResolvedNode {
     pub kind: CoreNodeKind,
 }
 
+impl ResolvedNode {
+    pub fn send_stdout_as(&self) -> Option<&str> {
+        match &self.kind {
+            CoreNodeKind::Runtime(_) => None, // todo: add support for operator-level stdout capture
+            CoreNodeKind::Custom(n) => n.send_stdout_as.as_deref(),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Default, Serialize, Deserialize)]
 pub struct ResolvedDeploy {
     pub machine: String,
@@ -275,6 +284,8 @@ pub struct CustomNode {
     pub envs: Option<BTreeMap<String, EnvValue>>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub build: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub send_stdout_as: Option<String>,
 
     #[serde(flatten)]
     pub run_config: NodeRunConfig,