From ab48d3a3f256145bc2e699b9dc052e32c1e53b8b Mon Sep 17 00:00:00 2001
From: Patrick Meredith <pmeredit@gmail.com>
Date: Tue, 3 Sep 2019 21:50:23 -0400
Subject: [PATCH] Support binary save

---
 src/commands/command.rs   | 15 ++++++++++
 src/commands/save.rs      | 59 ++++++++++++++++++++++++++-------------
 src/commands/to_bson.rs   |  4 +++
 src/commands/to_sqlite.rs |  8 ++++++
 4 files changed, 67 insertions(+), 19 deletions(-)

diff --git a/src/commands/command.rs b/src/commands/command.rs
index 67286e98a1..a79e797fd7 100644
--- a/src/commands/command.rs
+++ b/src/commands/command.rs
@@ -496,6 +496,10 @@ pub trait WholeStreamCommand: Send + Sync {
         args: CommandArgs,
         registry: &registry::CommandRegistry,
     ) -> Result<OutputStream, ShellError>;
+
+    fn is_binary(&self) -> bool {
+        false
+    }
 }
 
 pub trait PerItemCommand: Send + Sync {
@@ -521,6 +525,10 @@ pub trait PerItemCommand: Send + Sync {
         raw_args: &RawCommandArgs,
         input: Tagged<Value>,
     ) -> Result<OutputStream, ShellError>;
+
+    fn is_binary(&self) -> bool {
+        false
+    }
 }
 
 pub enum Command {
@@ -608,6 +616,13 @@ impl Command {
             }
         }
     }
+
+    pub fn is_binary(&self) -> bool {
+        match self {
+            Command::WholeStream(command) => command.is_binary(),
+            Command::PerItem(command) => command.is_binary(),
+        }
+    }
 }
 
 pub struct FnFilterCommand {
diff --git a/src/commands/save.rs b/src/commands/save.rs
index c32c016f4b..cdf41df6e7 100644
--- a/src/commands/save.rs
+++ b/src/commands/save.rs
@@ -96,7 +96,7 @@ fn save(
             }
         }
 
-        let content = if !save_raw {
+        let content : Result<Vec<u8>, ShellError> = if !save_raw {
             if let Some(extension) = full_path.extension() {
                 let command_name = format!("to-{}", extension.to_str().unwrap());
                 if let Some(converter) = registry.get_command(&command_name) {
@@ -116,22 +116,43 @@ fn save(
                     };
                     let mut result = converter.run(new_args.with_input(input), &registry);
                     let result_vec: Vec<Result<ReturnSuccess, ShellError>> = result.drain_vec().await;
-                    let mut result_string = String::new();
-                    for res in result_vec {
-                        match res {
-                            Ok(ReturnSuccess::Value(Tagged { item: Value::Primitive(Primitive::String(s)), .. })) => {
-                                result_string.push_str(&s);
+                    if converter.is_binary() {
+                        let mut result_binary : Vec<u8> = Vec::new();
+                        for res in result_vec {
+                            match res {
+                                Ok(ReturnSuccess::Value(Tagged { item: Value::Binary(b), .. })) => {
+                                    for u in b.into_iter() {
+                                        result_binary.push(u);
+                                    }
+                                }
+                                _ => {
+                                    yield Err(ShellError::labeled_error(
+                                        "Save could not successfully save",
+                                        "unexpected data during binary save",
+                                        name_span,
+                                    ));
+                                },
                             }
-                            _ => {
-                                yield Err(ShellError::labeled_error(
-                                    "Save could not successfully save",
-                                    "unexpected data during save",
-                                    name_span,
-                                ));
-                            },
                         }
+                        Ok(result_binary)
+                    } else {
+                        let mut result_string = String::new();
+                        for res in result_vec {
+                            match res {
+                                Ok(ReturnSuccess::Value(Tagged { item: Value::Primitive(Primitive::String(s)), .. })) => {
+                                    result_string.push_str(&s);
+                                }
+                                _ => {
+                                    yield Err(ShellError::labeled_error(
+                                        "Save could not successfully save",
+                                        "unexpected data during text save",
+                                        name_span,
+                                    ));
+                                },
+                            }
+                        }
+                        Ok(result_string.into_bytes())
                     }
-                    Ok(result_string)
                 } else {
                     let mut result_string = String::new();
                     for res in input {
@@ -148,7 +169,7 @@ fn save(
                             },
                         }
                     }
-                    Ok(result_string)
+                    Ok(result_string.into_bytes())
                 }
             } else {
                 let mut result_string = String::new();
@@ -166,10 +187,10 @@ fn save(
                         },
                     }
                 }
-                Ok(result_string)
+                Ok(result_string.into_bytes())
             }
         } else {
-            string_from(&input)
+            Ok(string_from(&input).into_bytes())
         };
 
         match content {
@@ -185,7 +206,7 @@ fn save(
     Ok(OutputStream::new(stream))
 }
 
-fn string_from(input: &Vec<Tagged<Value>>) -> Result<String, ShellError> {
+fn string_from(input: &Vec<Tagged<Value>>) -> String {
     let mut save_data = String::new();
 
     if input.len() > 0 {
@@ -202,5 +223,5 @@ fn string_from(input: &Vec<Tagged<Value>>) -> Result<String, ShellError> {
         }
     }
 
-    Ok(save_data)
+    save_data
 }
diff --git a/src/commands/to_bson.rs b/src/commands/to_bson.rs
index 35fc9839d4..73232e75b0 100644
--- a/src/commands/to_bson.rs
+++ b/src/commands/to_bson.rs
@@ -26,6 +26,10 @@ impl WholeStreamCommand for ToBSON {
     ) -> Result<OutputStream, ShellError> {
         to_bson(args, registry)
     }
+
+    fn is_binary(&self) -> bool {
+        true
+    }
 }
 
 pub fn value_to_bson_value(v: &Tagged<Value>) -> Result<Bson, ShellError> {
diff --git a/src/commands/to_sqlite.rs b/src/commands/to_sqlite.rs
index dca953c235..e229057637 100644
--- a/src/commands/to_sqlite.rs
+++ b/src/commands/to_sqlite.rs
@@ -27,6 +27,10 @@ impl WholeStreamCommand for ToSQLite {
     ) -> Result<OutputStream, ShellError> {
         to_sqlite(args, registry)
     }
+
+    fn is_binary(&self) -> bool {
+        true
+    }
 }
 
 pub struct ToDB;
@@ -51,6 +55,10 @@ impl WholeStreamCommand for ToDB {
     ) -> Result<OutputStream, ShellError> {
         to_sqlite(args, registry)
     }
+
+    fn is_binary(&self) -> bool {
+        true
+    }
 }
 
 fn comma_concat(acc: String, current: String) -> String {