Register vars (umputun#151)

* add an explicit register option

umputun#135 (reply in thread)

* better text in docs

* update man

* add validation allowing register with script only

README.md

@@ -524,6 +524,27 @@ commands:
     copy: {src: $FILE_NAME, dest: /tmp/file2}
 ```
 
+Sometimes, exporting variables is not possible or not desired. For such cases, Spot allows to register variables explicitly using the `register` option.
+
+For example:
+
+```yaml
+commands:
+  - name: first command
+    script: |
+      FILE_NAME=/tmp/file1
+      touch $FILE_NAME
+      ANOTHER_VAR=foo
+    register: [FILE_NAME, ANOTHER_VAR]}
+      
+  - name: second command
+    script: |
+      echo "File name is $FILE_NAME, var is $ANOTHER_VAR"
+      
+  - name: third command
+    copy: {src: $FILE_NAME, dest: /tmp/file2}
+```
+
 ### Setting environment variables
 
 Environment variables can be set with `--env` / `-e` cli option. For example: `-e VAR1:VALUE1 -e VAR2:VALUE2`. Environment variables can also be set in the environment file (default `env.yml` can be changed with `--env-file` / `-E` cli flag). For example:

pkg/config/command.go

@@ -29,6 +29,7 @@ type Cmd struct {
 	Environment map[string]string `yaml:"env" toml:"env"`
 	Options     CmdOptions        `yaml:"options" toml:"options,omitempty"`
 	Condition   string            `yaml:"cond" toml:"cond,omitempty"`
+	Register    []string          `yaml:"register" toml:"register"` // register variables from command
 
 	Secrets  map[string]string `yaml:"-" toml:"-"` // loaded secrets, filled by playbook
 	SSHShell string            `yaml:"-" toml:"-"` // shell to use for ssh commands, filled by playbook
@@ -86,7 +87,7 @@ func (cmd *Cmd) GetScript() (command string, rdr io.Reader) {
 	// 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)
+		return "", cmd.scriptFile(cmd.Script, cmd.Register)
 	}
 
 	log.Printf("[DEBUG] command %q is single line, using script string", cmd.Name)
@@ -102,7 +103,7 @@ func (cmd *Cmd) GetWait() (command string, rdr io.Reader) {
 	elems := strings.Split(cmd.Wait.Command, "\n")
 	if len(elems) > 1 {
 		log.Printf("[DEBUG] wait command %q is multiline, using script file", cmd.Name)
-		return "", cmd.scriptFile(cmd.Wait.Command)
+		return "", cmd.scriptFile(cmd.Wait.Command, nil)
 	}
 
 	log.Printf("[DEBUG] wait command %q is single line, using command string", cmd.Name)
@@ -122,7 +123,7 @@ func (cmd *Cmd) GetCondition() (command string, rdr io.Reader, inverted bool) {
 	elems := strings.Split(cond, "\n")
 	if len(elems) > 1 {
 		log.Printf("[DEBUG] condition %q is multiline, using script file", cmd.Name)
-		return "", cmd.scriptFile(cond), inverted
+		return "", cmd.scriptFile(cond, nil), inverted
 	}
 
 	log.Printf("[DEBUG] condition %q is single line, using condition string", cmd.Name)
@@ -167,7 +168,7 @@ func (cmd *Cmd) scriptCommand(inp string) string {
 
 // scriptFile returns a reader for script file. All the lines in the command used as a script, with hashbang,
 // set -e and environment variables.
-func (cmd *Cmd) scriptFile(inp string) (r io.Reader) {
+func (cmd *Cmd) scriptFile(inp string, register []string) (r io.Reader) {
 	var buf bytes.Buffer
 	inp = strings.TrimPrefix(inp, "\n") // trim leading newline if present; can be due to multiline yaml format
 	if !cmd.hasShebang(inp) {
@@ -208,8 +209,15 @@ func (cmd *Cmd) scriptFile(inp string) (r io.Reader) {
 
 	// each exported variable is printed as a setvar command to be captured by the caller
 	if len(exports) > 0 {
-		for i := range exports {
-			buf.WriteString(fmt.Sprintf("echo setvar %s=${%s}\n", exports[i], exports[i]))
+		for _, v := range exports {
+			buf.WriteString(fmt.Sprintf("echo setvar %s=${%s}\n", v, v))
+		}
+	}
+
+	// each register variable is printed as a setvar command to be captured by the caller
+	if len(register) > 0 {
+		for _, v := range register {
+			buf.WriteString(fmt.Sprintf("echo setvar %s=${%s}\n", v, v))
 		}
 	}
 
@@ -377,6 +385,11 @@ func (cmd *Cmd) validate() error {
 	if len(setCmds) == 0 {
 		return fmt.Errorf("one of [%s] must be set", strings.Join(names, ", "))
 	}
+
+	// make sure what register used with script and not with other commands
+	if cmd.Script == "" && len(cmd.Register) > 0 {
+		return fmt.Errorf("register is only allowed with script command")
+	}
 	return nil
 }
 

pkg/config/command_test.go

@@ -371,11 +371,26 @@ func TestCmd_getScriptFile(t *testing.T) {
 			},
 			expected: "#!/bin/sh\nset -e\nexport SEC1=\"secret1\"\nexport SEC2=\"secret2\"\necho 'Hello, World!'\n",
 		},
+		{
+			name: "with exports",
+			cmd: &Cmd{
+				Script: "echo 'Hello, World!'\nexport var1=blah\n export var2=baz",
+			},
+			expected: "#!/bin/sh\nset -e\necho 'Hello, World!'\nexport var1=blah\n export var2=baz\necho setvar var1=${var1}\necho setvar var2=${var2}\n",
+		},
+		{
+			name: "with exports and register",
+			cmd: &Cmd{
+				Script:   "echo 'Hello, World!'\nexport var1=blah\n export var2=baz",
+				Register: []string{"var21", "var22", "var23"},
+			},
+			expected: "#!/bin/sh\nset -e\necho 'Hello, World!'\nexport var1=blah\n export var2=baz\necho setvar var1=${var1}\necho setvar var2=${var2}\necho setvar var21=${var21}\necho setvar var22=${var22}\necho setvar var23=${var23}\n",
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			reader := tt.cmd.scriptFile(tt.cmd.Script)
+			reader := tt.cmd.scriptFile(tt.cmd.Script, tt.cmd.Register)
 			scriptContentBytes, err := io.ReadAll(reader)
 			assert.NoError(t, err)
 			scriptContent := string(scriptContentBytes)
@@ -571,6 +586,9 @@ func TestCmd_validate(t *testing.T) {
 		{"multiple fields set", Cmd{Script: "example_script", Copy: CopyInternal{Source: "source", Dest: "dest"}},
 			"only one of [script, copy] is allowed"},
 		{"nothing set", Cmd{}, "one of [script, copy, mcopy, delete, mdelete, sync, msync, wait, echo] must be set"},
+		{"script with register", Cmd{Script: "example_script", Register: []string{"a", "b"}}, ""},
+		{"unexpected register", Cmd{Copy: CopyInternal{Source: "source", Dest: "dest"}, Register: []string{"a", "b"}},
+			"register is only allowed with script command"},
 	}
 
 	for _, tt := range tbl {

pkg/runner/runner_test.go

@@ -45,7 +45,7 @@ func TestProcess_Run(t *testing.T) {
 		require.NoError(t, err)
 		assert.Equal(t, 8, res.Commands)
 		assert.Equal(t, 1, res.Hosts)
-		assert.EqualValues(t, map[string]string{"bar": "9", "baz": "zzzzz", "foo": "6"}, res.Vars)
+		assert.EqualValues(t, map[string]string{"bar": "9", "bar2": "10", "baz": "zzzzz", "foo": "6", "foo2": "7"}, res.Vars)
 	})
 
 	t.Run("simple playbook", func(t *testing.T) {
@@ -170,7 +170,7 @@ func TestProcess_Run(t *testing.T) {
 		assert.Contains(t, outWriter.String(), `> var foo: 6`)
 		assert.Contains(t, outWriter.String(), `> var bar: 9`)
 		assert.Contains(t, outWriter.String(), `> var baz: qux`, "was not overwritten")
-		assert.EqualValues(t, map[string]string{"bar": "9", "baz": "zzzzz", "foo": "6"}, res.Vars)
+		assert.EqualValues(t, map[string]string{"bar": "9", "bar2": "10", "baz": "zzzzz", "foo": "6", "foo2": "7"}, res.Vars)
 	})
 
 	t.Run("with secrets", func(t *testing.T) {

pkg/runner/testdata/conf.yml

@@ -27,10 +27,13 @@ tasks:
           export foo=$((1 + 2 + 3))
           export bar=$((4 + 5))
           export baz=zzzzz
+          foo2=$((foo + 1))
+          bar2=$((bar + 1))
           ls -laR /tmp
           du -hcs /srv
           cat /tmp/conf.yml
           echo all good, 123
+        register: [foo2, bar2]
 
       - name: runtime variables
         script: echo host:"{SPOT_REMOTE_HOST}", name:"{SPOT_REMOTE_NAME}", cmd:"{SPOT_COMMAND}", user:"{SPOT_REMOTE_USER}", task:"{SPOT_TASK}"

spot.1

@@ -1,4 +1,4 @@
-.TH "SPOT" 1  20230804T082658 spot manual
+.TH "SPOT" 1  20230808T034208 spot manual
 .\" Automatically generated by Pandoc 3.1.6
 .\"
 .\" Define V font for inline verbatim, using C font in formats
@@ -788,6 +788,31 @@ commands:
     copy: {src: $FILE_NAME, dest: /tmp/file2}
 \f[R]
 .fi
+.PP
+Sometimes, exporting variables is not possible or not desired.
+For such cases, Spot allows to register variables explicitly using the
+\f[V]register\f[R] option.
+.PP
+For example:
+.IP
+.nf
+\f[C]
+commands:
+  - name: first command
+    script: |
+      FILE_NAME=/tmp/file1
+      touch $FILE_NAME
+      ANOTHER_VAR=foo
+    register: [FILE_NAME, ANOTHER_VAR]}
+
+  - name: second command
+    script: |
+      echo \[dq]File name is $FILE_NAME, var is $ANOTHER_VAR\[dq]
+
+  - name: third command
+    copy: {src: $FILE_NAME, dest: /tmp/file2}
+\f[R]
+.fi
 .SS Setting environment variables
 .PP
 Environment variables can be set with \f[V]--env\f[R] / \f[V]-e\f[R] cli