extend template substitution to env for all command type

pkg/config/command.go

@@ -71,7 +71,8 @@ func (cmd *Cmd) GetScript() (command string, rdr io.Reader) {
 	}
 
 	elems := strings.Split(cmd.Script, "\n")
-	if len(elems) > 1 {
+	// export should be treated as multiline for env vars to be set
+	if len(elems) > 1 || strings.Contains(cmd.Script, "export") {
 		log.Printf("[DEBUG] command %q is multiline, using script file", cmd.Name)
 		return "", cmd.scriptFile(cmd.Script)
 	}

pkg/config/command_test.go

@@ -116,6 +116,19 @@ echo 'Goodbye, World!'`,
 			expectedScript:   "",
 			expectedContents: nil,
 		},
+		{
+			name: "single line command with export",
+			cmd: &Cmd{
+				Script: "export GREETING='Hello, World!'",
+			},
+			expectedScript: "",
+			expectedContents: []string{
+				"#!/bin/sh",
+				"set -e",
+				"export GREETING='Hello, World!'",
+				"echo setvar GREETING=${GREETING}",
+			},
+		},
 	}
 
 	for _, tc := range testCases {

pkg/runner/commands.go

@@ -78,7 +78,7 @@ func (ec *execCmd) script(ctx context.Context) (details string, vars map[string]
 // if sudo option is set, it will make a temporary directory and upload the files there,
 // then move it to the final destination with sudo script execution.
 func (ec *execCmd) copy(ctx context.Context) (details string, vars map[string]string, err error) {
-	tmpl := templater{hostAddr: ec.hostAddr, hostName: ec.hostName, task: ec.tsk, command: ec.cmd.Name}
+	tmpl := templater{hostAddr: ec.hostAddr, hostName: ec.hostName, task: ec.tsk, command: ec.cmd.Name, env: ec.cmd.Environment}
 
 	src := tmpl.apply(ec.cmd.Copy.Source)
 	dst := tmpl.apply(ec.cmd.Copy.Dest)
@@ -127,7 +127,7 @@ func (ec *execCmd) copy(ctx context.Context) (details string, vars map[string]st
 // mcopy uploads multiple files to a target host. It calls copy function for each file.
 func (ec *execCmd) mcopy(ctx context.Context) (details string, vars map[string]string, err error) {
 	msgs := []string{}
-	tmpl := templater{hostAddr: ec.hostAddr, hostName: ec.hostName, task: ec.tsk, command: ec.cmd.Name}
+	tmpl := templater{hostAddr: ec.hostAddr, hostName: ec.hostName, task: ec.tsk, command: ec.cmd.Name, env: ec.cmd.Environment}
 	for _, c := range ec.cmd.MCopy {
 		src := tmpl.apply(c.Source)
 		dst := tmpl.apply(c.Dest)
@@ -144,7 +144,7 @@ func (ec *execCmd) mcopy(ctx context.Context) (details string, vars map[string]s
 
 // sync synchronizes files from a source to a destination on a target host.
 func (ec *execCmd) sync(ctx context.Context) (details string, vars map[string]string, err error) {
-	tmpl := templater{hostAddr: ec.hostAddr, hostName: ec.hostName, task: ec.tsk, command: ec.cmd.Name}
+	tmpl := templater{hostAddr: ec.hostAddr, hostName: ec.hostName, task: ec.tsk, command: ec.cmd.Name, env: ec.cmd.Environment}
 	src := tmpl.apply(ec.cmd.Sync.Source)
 	dst := tmpl.apply(ec.cmd.Sync.Dest)
 	details = fmt.Sprintf(" {sync: %s -> %s}", src, dst)
@@ -156,7 +156,7 @@ func (ec *execCmd) sync(ctx context.Context) (details string, vars map[string]st
 
 // delete deletes files on a target host. If sudo option is set, it will execute a sudo rm commands.
 func (ec *execCmd) delete(ctx context.Context) (details string, vars map[string]string, err error) {
-	tmpl := templater{hostAddr: ec.hostAddr, hostName: ec.hostName, task: ec.tsk, command: ec.cmd.Name}
+	tmpl := templater{hostAddr: ec.hostAddr, hostName: ec.hostName, task: ec.tsk, command: ec.cmd.Name, env: ec.cmd.Environment}
 	loc := tmpl.apply(ec.cmd.Delete.Location)
 
 	if !ec.cmd.Options.Sudo {
@@ -297,11 +297,13 @@ type templater struct {
 	hostAddr string
 	hostName string
 	command  string
+	env      map[string]string
 	task     *config.Task
 	err      error
 }
 
 // apply applies templates to a string to replace predefined vars placeholders with actual values
+// it also applies the task environment variables to the string
 func (tm *templater) apply(inp string) string {
 	apply := func(inp, from, to string) string {
 		// replace either {SPOT_REMOTE_HOST} ${SPOT_REMOTE_HOST} or $SPOT_REMOTE_HOST format
@@ -323,5 +325,9 @@ func (tm *templater) apply(inp string) string {
 		res = apply(res, "SPOT_ERROR", "")
 	}
 
+	for k, v := range tm.env {
+		res = apply(res, k, v)
+	}
+
 	return res
 }

pkg/runner/commands_test.go

@@ -82,6 +82,17 @@ func Test_templaterApply(t *testing.T) {
 			},
 			expected: "example.com blah task1 user2:ls",
 		},
+		{
+			name: "env variables",
+			inp:  "${FOO} blah $BAR ${SPOT_REMOTE_USER}:${SPOT_COMMAND}",
+			tmpl: templater{
+				hostAddr: "example.com",
+				command:  "ls",
+				task:     &config.Task{Name: "task1", User: "user2"},
+				env:      map[string]string{"FOO": "foo_val", "BAR": "bar_val"},
+			},
+			expected: "foo_val blah bar_val user2:ls",
+		},
 		{
 			name: "with error msg",
 			inp:  "$SPOT_REMOTE_HOST:$SPOT_REMOTE_USER:$SPOT_COMMAND ${SPOT_ERROR}",

pkg/runner/runner_test.go

@@ -160,6 +160,28 @@ func TestProcess_Run(t *testing.T) {
 		assert.NotContains(t, outWriter.String(), "FOO_SECRET")
 		assert.NotContains(t, outWriter.String(), "BAR_SECRET")
 	})
+
+	t.Run("set variables for copy command", func(t *testing.T) {
+		conf, err := config.New("testdata/conf.yml", nil, nil)
+		require.NoError(t, err)
+
+		p := Process{
+			Concurrency: 1,
+			Connector:   connector,
+			Config:      conf,
+			ColorWriter: executor.NewColorizedWriter(os.Stdout, "", "", "", nil),
+			Only:        []string{"set filename for copy to env", "copy filename from env"},
+		}
+
+		outWriter := &bytes.Buffer{}
+		log.SetOutput(io.MultiWriter(outWriter, os.Stderr))
+
+		res, err := p.Run(ctx, "task1", testingHostAndPort)
+		require.NoError(t, err)
+		assert.Equal(t, 2, res.Commands)
+		assert.Contains(t, outWriter.String(), ` > setvar filename=testdata/conf.yml`)
+	})
+
 }
 
 func TestProcess_RunWithSudo(t *testing.T) {

pkg/runner/testdata/conf.yml

@@ -90,6 +90,14 @@ tasks:
         script: ls -la /srv
         options: {no_auto: true, sudo: true}
 
+      - name: set filename for copy to env
+        script: |
+          export filename="testdata/conf.yml"
+        options: {no_auto: true, sudo: true}
+
+      - name: copy filename from env
+        copy: {src: "${filename}", dst: "/srv/conf.yml"}
+        options: {no_auto: true, sudo: true}
 
   - name: failed_task
     commands: