From 65be7ad74176505ffb23a2a7ac8e5c1dda08e841 Mon Sep 17 00:00:00 2001
From: maksimowiczm <maksimowiczmateusz13@gmail.com>
Date: Wed, 1 May 2024 17:27:12 +0200
Subject: [PATCH] Add stdin option to command chains

---
 CHANGELOG.md                                  |  1 +
 .../api/request_collection/chain_source.md    |  1 +
 docs/src/user_guide/chains.md                 |  3 +-
 docs/src/user_guide/filter_query.md           |  3 +-
 src/collection/models.rs                      |  5 +-
 src/template.rs                               | 37 ++++++--
 src/template/render.rs                        | 91 ++++++++++++++-----
 7 files changed, 108 insertions(+), 33 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 31dc81d0..2ebebaa1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 - Add action to save response body to file ([#183](https://github.com/LucasPickering/slumber/issues/183))
 - Add `theme` field to the config, to configure colors ([#193](https://github.com/LucasPickering/slumber/issues/193))
   - [See docs](https://slumber.lucaspickering.me/book/api/configuration/theme.html) for more info
+- Add `stdin` option to command chains ([#190](https://github.com/LucasPickering/slumber/issues/190))
 
 ### Changed
 
diff --git a/docs/src/api/request_collection/chain_source.md b/docs/src/api/request_collection/chain_source.md
index 0e0730e9..2e5ebe16 100644
--- a/docs/src/api/request_collection/chain_source.md
+++ b/docs/src/api/request_collection/chain_source.md
@@ -107,6 +107,7 @@ Execute a command and use its stdout as the rendered value.
 | Field     | Type         | Description                                                 | Default  |
 | --------- | ------------ | ----------------------------------------------------------- | -------- |
 | `command` | `Template[]` | Command to execute, in the format `[program, ...arguments]` | Required |
+| `stdin`   | `Template`   | Standard input which will be piped into the command         | None     |
 
 ### File
 
diff --git a/docs/src/user_guide/chains.md b/docs/src/user_guide/chains.md
index 563d2113..2806996f 100644
--- a/docs/src/user_guide/chains.md
+++ b/docs/src/user_guide/chains.md
@@ -99,7 +99,8 @@ chains:
       recipe: login
   auth_token:
     source: !command
-      command: [sh, -c, "echo '{{chains.auth_token_raw}}' | cut -d':' -f2"]
+      command: [ "cut", "-d':'", "-f2" ]
+      stdin: "{{chains.auth_token_raw}}"
 
 requests:
   login: !request
diff --git a/docs/src/user_guide/filter_query.md b/docs/src/user_guide/filter_query.md
index 09cd6ac5..b823e225 100644
--- a/docs/src/user_guide/filter_query.md
+++ b/docs/src/user_guide/filter_query.md
@@ -77,7 +77,8 @@ chains:
       recipe: login
   auth_token:
     source: !command
-      command: [sh, -c, "echo '{{chains.auth_token_raw}}' | jq .token"]
+      command: [ "jq", ".token" ]
+      stdin: "{{chains.auth_token_raw}}
 
 requests:
   login: !request
diff --git a/src/collection/models.rs b/src/collection/models.rs
index 66f82e5b..6dbdcefb 100644
--- a/src/collection/models.rs
+++ b/src/collection/models.rs
@@ -244,7 +244,10 @@ pub enum ChainSource {
         section: ChainRequestSection,
     },
     /// Run an external command to get a result
-    Command { command: Vec<Template> },
+    Command {
+        command: Vec<Template>,
+        stdin: Option<Template>,
+    },
     /// Load data from a file
     File { path: Template },
     /// Prompt the user for a value
diff --git a/src/template.rs b/src/template.rs
index 6721495e..116998a4 100644
--- a/src/template.rs
+++ b/src/template.rs
@@ -201,6 +201,7 @@ mod tests {
         let profile_data = indexmap! {"field1".into() => "field".into()};
         let source = ChainSource::Command {
             command: vec!["echo".into(), "chain".into()],
+            stdin: None,
         };
         let overrides = indexmap! {
             "field1".into() => "override".into(),
@@ -581,10 +582,20 @@ mod tests {
     }
 
     /// Test success with chained command
+    #[rstest]
+    #[case::with_stdin(&["tail"], Some("hello!"), "hello!")]
+    #[case::raw_command(&["echo", "-n", "hello!"], None, "hello!")]
     #[tokio::test]
-    async fn test_chain_command() {
-        let command = vec!["echo".into(), "-n".into(), "hello!".into()];
-        let chain = create!(Chain, source: ChainSource::Command { command });
+    async fn test_chain_command(
+        #[case] command: &[&str],
+        #[case] stdin: Option<&str>,
+        #[case] expected: &str,
+    ) {
+        let source = ChainSource::Command {
+            command: command.iter().copied().map(Template::from).collect(),
+            stdin: stdin.map(Template::from),
+        };
+        let chain = create!(Chain, source: source);
         let context = create!(
             TemplateContext,
             collection: create!(
@@ -593,21 +604,29 @@ mod tests {
             ),
         );
 
-        assert_eq!(render!("{{chains.chain1}}", context).unwrap(), "hello!");
+        assert_eq!(render!("{{chains.chain1}}", context).unwrap(), expected);
     }
 
     /// Test failure with chained command
     #[rstest]
-    #[case::no_command(&[], "No command given")]
-    #[case::unknown_command(&["totally not a program"], "No such file or directory")]
-    #[case::command_error(&["head", "/dev/random"], "invalid utf-8 sequence")]
+    #[case::no_command(&[], None, "No command given")]
+    #[case::unknown_command(&["totally not a program"], None, "No such file or directory")]
+    #[case::command_error(&["head", "/dev/random"], None, "invalid utf-8 sequence")]
+    #[case::stdin_error(
+        &["tail"],
+        Some("{{chains.stdin}}"),
+        "Resolving chain `chain1`: Rendering nested template for field `stdin`: \
+         Resolving chain `stdin`: Unknown chain: stdin"
+    )]
     #[tokio::test]
     async fn test_chain_command_error(
         #[case] command: &[&str],
+        #[case] stdin: Option<&str>,
         #[case] expected_error: &str,
     ) {
         let source = ChainSource::Command {
             command: command.iter().copied().map(Template::from).collect(),
+            stdin: stdin.map(Template::from),
         };
         let chain = create!(Chain, source: source);
         let context = create!(
@@ -767,7 +786,7 @@ mod tests {
         let command_chain = create!(
             Chain,
             id: "command".into(),
-            source: ChainSource::Command { command },
+            source: ChainSource::Command { command, stdin: None},
         );
 
         let context = create!(
@@ -802,7 +821,7 @@ mod tests {
         let command_chain = create!(
             Chain,
             id: "command".into(),
-            source: ChainSource::Command { command },
+            source: ChainSource::Command { command, stdin: None },
         );
 
         let context = create!(
diff --git a/src/template/render.rs b/src/template/render.rs
index ad49468d..52b07ee1 100644
--- a/src/template/render.rs
+++ b/src/template/render.rs
@@ -19,9 +19,10 @@ use futures::future;
 use std::{
     env,
     path::PathBuf,
+    process::Stdio,
     sync::{atomic::Ordering, Arc},
 };
-use tokio::{fs, process::Command, sync::oneshot};
+use tokio::{fs, io::AsyncWriteExt, process::Command, sync::oneshot};
 use tracing::{debug, debug_span, instrument, trace};
 
 /// Outcome of rendering a single chunk. This allows attaching some metadata to
@@ -284,9 +285,13 @@ impl<'a> TemplateSource<'a> for ChainTemplateSource<'a> {
                 ChainSource::File { path } => {
                     self.render_file(context, path).await?
                 }
-                ChainSource::Command { command } => {
+                ChainSource::Command { command, stdin } => {
                     // No way to guess content type on this
-                    (self.render_command(context, command).await?, None)
+                    (
+                        self.render_command(context, command, stdin.as_ref())
+                            .await?,
+                        None,
+                    )
                 }
                 ChainSource::Prompt { message, default } => (
                     self.render_prompt(
@@ -484,6 +489,7 @@ impl<'a> ChainTemplateSource<'a> {
         &self,
         context: &TemplateContext,
         command: &[Template],
+        stdin: Option<&Template>,
     ) -> Result<Vec<u8>, ChainError> {
         // Render each arg in the command
         let command = future::try_join_all(command.iter().enumerate().map(
@@ -501,26 +507,69 @@ impl<'a> ChainTemplateSource<'a> {
         let [program, args @ ..] = command.as_slice() else {
             return Err(ChainError::CommandMissing);
         };
-        debug_span!("Executing command", ?command)
-            .in_scope(|| async {
-                let output = Command::new(program)
-                    .args(args)
-                    .output()
-                    .await
-                    .map_err(|error| ChainError::Command {
-                        command: command.to_owned(),
-                        error,
-                    })
-                    .traced()?;
-
-                debug!(
-                    stdout = %String::from_utf8_lossy(&output.stdout),
-                    stderr = %String::from_utf8_lossy(&output.stderr),
-                    "Command success"
-                );
-                Ok(output.stdout)
+
+        let _ = debug_span!("Executing command", ?command).entered();
+
+        // Render the stdin template, if present
+        let input = if let Some(template) = stdin {
+            let input =
+                template.render_stitched(context).await.map_err(|error| {
+                    ChainError::Nested {
+                        field: "stdin".into(),
+                        error: error.into(),
+                    }
+                })?;
+
+            Some(input)
+        } else {
+            None
+        };
+
+        // Spawn the command process
+        let mut process = Command::new(program)
+            .args(args)
+            .stdin(Stdio::piped())
+            .stdout(Stdio::piped())
+            .stderr(Stdio::piped())
+            .spawn()
+            .map_err(|error| ChainError::Command {
+                command: command.to_owned(),
+                error,
             })
+            .traced()?;
+
+        // Write the stdin to the process
+        if let Some(input) = input {
+            process
+                .stdin
+                .as_mut()
+                .expect("Process missing stdin")
+                .write_all(input.as_bytes())
+                .await
+                .map_err(|error| ChainError::Command {
+                    command: command.to_owned(),
+                    error,
+                })
+                .traced()?;
+        }
+
+        // Wait for the process to finish
+        let output = process
+            .wait_with_output()
             .await
+            .map_err(|error| ChainError::Command {
+                command: command.to_owned(),
+                error,
+            })
+            .traced()?;
+
+        debug!(
+            stdout = %String::from_utf8_lossy(&output.stdout),
+            stderr = %String::from_utf8_lossy(&output.stderr),
+            "Command success"
+        );
+
+        Ok(output.stdout)
     }
 
     /// Render a value by asking the user to provide it