[dagster-pipes-rust] - Adds support for reporting custom messages

libraries/pipes/implementations/rust/CHANGELOG.md

@@ -2,7 +2,7 @@
 
 ## [UNRELEASED]
 
-- 
+- (pull/70) Added `report_custom_message` method to the `PipesContext`
 
 ## 0.1.7
 

libraries/pipes/implementations/rust/pipes.toml

@@ -3,7 +3,7 @@ error_reporting = false
 
 [messages]
 log = false
-report_custom_message = false
+report_custom_message = true
 report_asset_materialization = false
 report_asset_check = false
 log_external_stream = false

libraries/pipes/implementations/rust/src/bin/pipes_tests.rs

@@ -72,7 +72,7 @@ pub fn main() -> Result<(), DagsterPipesError> {
         std::env::set_var(DAGSTER_PIPES_MESSAGES_ENV_VAR, &messages);
     }
 
-    let context = open_dagster_pipes()?;
+    let mut context = open_dagster_pipes()?;
 
     if let Some(job_name) = args.job_name {
         assert_eq!(context.data.job_name, Some(job_name));
@@ -84,5 +84,19 @@ pub fn main() -> Result<(), DagsterPipesError> {
             serde_json::from_reader(file).expect("extras could not be parsed");
         assert_eq!(context.data.extras, Some(json));
     }
+
+    if let Some(custom_payload_path) = args.custom_payload_path {
+        let file =
+            File::open(custom_payload_path).expect("custom_payload_path could not be opened");
+        let payload = serde_json::from_reader::<File, serde_json::Value>(file)
+            .expect("custom_payload_path could not be parsed")
+            .as_object()
+            .expect("custom payload must be an object")
+            .get("payload")
+            .expect("custom payload must have a 'payload' key")
+            .clone();
+        context.report_custom_message(payload)?
+    }
+
     Ok(())
 }

libraries/pipes/implementations/rust/src/lib.rs

@@ -110,6 +110,17 @@ where
         let msg = PipesMessage::new(Method::ReportAssetCheck, Some(params));
         self.message_channel.write_message(msg)
     }
+
+    pub fn report_custom_message(
+        &mut self,
+        payload: serde_json::Value,
+    ) -> Result<(), MessageWriteError> {
+        let params: HashMap<&str, Option<serde_json::Value>> =
+            HashMap::from([("payload", Some(payload))]);
+
+        let msg = PipesMessage::new(Method::ReportCustomMessage, Some(params));
+        self.message_channel.write_message(msg)
+    }
 }
 
 impl<W> Drop for PipesContext<W>
@@ -156,16 +167,32 @@ pub fn open_dagster_pipes() -> Result<PipesContext<PipesDefaultMessageWriter>, D
 
 #[cfg(test)]
 mod tests {
-    use rstest::rstest;
+    use rstest::{fixture, rstest};
     use std::collections::HashMap;
     use std::fs;
     use tempfile::NamedTempFile;
     use writer::message_writer_channel::{DefaultChannel, FileChannel};
 
     use super::*;
 
-    #[test]
-    fn test_write_pipes_metadata() {
+    #[fixture]
+    fn file_and_context() -> (NamedTempFile, PipesContext<DefaultWriter>) {
+        let file = NamedTempFile::new().unwrap();
+        let context: PipesContext<DefaultWriter> = PipesContext {
+            message_channel: DefaultChannel::File(FileChannel::new(file.path().into())),
+            data: PipesContextData {
+                asset_keys: Some(vec!["asset1".to_string()]),
+                run_id: "012345".to_string(),
+                ..Default::default()
+            },
+        };
+        (file, context)
+    }
+
+    #[rstest]
+    fn test_write_pipes_metadata(
+        #[from(file_and_context)] (file, mut context): (NamedTempFile, PipesContext<DefaultWriter>),
+    ) {
         let asset_metadata = HashMap::from([
             ("int", PipesMetadataValue::from(100)),
             ("float", PipesMetadataValue::from(100.0)),
@@ -208,15 +235,6 @@ mod tests {
             ("job", PipesMetadataValue::from_job("some_job".to_string())),
         ]);
 
-        let file = NamedTempFile::new().unwrap();
-        let mut context: PipesContext<DefaultWriter> = PipesContext {
-            message_channel: DefaultChannel::File(FileChannel::new(file.path().into())),
-            data: PipesContextData {
-                asset_keys: Some(vec!["asset1".to_string()]),
-                run_id: "012345".to_string(),
-                ..Default::default()
-            },
-        };
         context
             .report_asset_materialization("asset1", asset_metadata, Some("v1"))
             .expect("Failed to report asset materialization");
@@ -329,18 +347,10 @@ mod tests {
         })
     )]
     fn test_close_pipes_context(
+        #[from(file_and_context)] (file, mut context): (NamedTempFile, PipesContext<DefaultWriter>),
         #[case] exc: Option<PipesException>,
         #[case] expected_message: serde_json::Value,
     ) {
-        let file = NamedTempFile::new().unwrap();
-        let mut context: PipesContext<DefaultWriter> = PipesContext {
-            message_channel: DefaultChannel::File(FileChannel::new(file.path().into())),
-            data: PipesContextData {
-                asset_keys: Some(vec!["asset1".to_string()]),
-                run_id: "012345".to_string(),
-                ..Default::default()
-            },
-        };
         context.close(exc).expect("Failed to close context");
         assert_eq!(
             serde_json::from_str::<serde_json::Value>(&fs::read_to_string(file.path()).unwrap())
@@ -372,4 +382,25 @@ mod tests {
             })
         );
     }
+
+    #[rstest]
+    fn test_report_custom_message(
+        #[from(file_and_context)] (file, mut context): (NamedTempFile, PipesContext<DefaultWriter>),
+    ) {
+        context
+            .report_custom_message(json!({"key": "value"}))
+            .expect("Failed to report custom message");
+
+        assert_eq!(
+            serde_json::from_str::<serde_json::Value>(&fs::read_to_string(file.path()).unwrap())
+                .unwrap(),
+            json!({
+                "__dagster_pipes_version": "0.1",
+                "method": "report_custom_message",
+                "params": {
+                    "payload": {"key": "value"}
+                },
+            })
+        );
+    }
 }