add cmd option "only_on" allowing to limit on what host command may r…

…un (#88)

* add cmd option "only_on" allowing to limit on what host

* add the inversion of condition with !

* refactor shouldRun to be easier to read

README.md

@@ -315,14 +315,15 @@ Each command type supports the following options:
 - `no_auto`: if set to `true` the command will not be executed automatically, but can be executed manually using the `--only` flag.
 - `local`: if set to `true` the command will be executed on the local host (the one running the `spot` command) instead of the remote host(s).
 - `sudo`: if set to `true` the command will be executed with `sudo` privileges.
+- `only_on`: optional, allows to set a list of host names or addresses where the command will be executed. If not set, the command will be executed on all hosts. For example, `only_on: [host1, host2]` will execute command on `host1` and `host2` only. This option also supports reversed condition, so if user wants to execute command on all hosts except some, `!` prefix can be used. For example, `only_on: [!host1, !host2]` will execute command on all hosts except `host1` and `host2`. 
 
-example setting `ignore_errors` and `no_auto` options:
+example setting `ignore_errors`, `no_auto` and `only_on` options:
 
 ```yaml
   commands:
       - name: wait
         script: sleep 5s
-        options: {ignore_errors: true, no_auto: true}
+        options: {ignore_errors: true, no_auto: true, only_on: [host1, host2]}
 ```
 
 Please note that the `sudo` option is not supported for the `sync` command type, but all other command types support it.

pkg/config/command.go

@@ -30,11 +30,12 @@ type Cmd struct {
 
 // CmdOptions defines options for a command
 type CmdOptions struct {
-	IgnoreErrors bool     `yaml:"ignore_errors" toml:"ignore_errors"`
-	NoAuto       bool     `yaml:"no_auto" toml:"no_auto"`
-	Local        bool     `yaml:"local" toml:"local"`
-	Sudo         bool     `yaml:"sudo" toml:"sudo"`
-	Secrets      []string `yaml:"secrets" toml:"secrets"`
+	IgnoreErrors bool     `yaml:"ignore_errors" toml:"ignore_errors"` // ignore errors and continue
+	NoAuto       bool     `yaml:"no_auto" toml:"no_auto"`             // don't run command automatically
+	Local        bool     `yaml:"local" toml:"local"`                 // run command on localhost
+	Sudo         bool     `yaml:"sudo" toml:"sudo"`                   // run command with sudo
+	Secrets      []string `yaml:"secrets" toml:"secrets"`             // list of secrets (keys) to load
+	OnlyOn       []string `yaml:"only_on" toml:"only_on"`             // only run on these hosts
 }
 
 // CopyInternal defines copy command, implemented internally

pkg/runner/runner.go

@@ -136,6 +136,10 @@ func (p *Process) runTaskOnHost(ctx context.Context, tsk *config.Task, hostAddr,
 			// skip command if it has NoAuto option and not in Only list
 			continue
 		}
+		if !p.shouldRunCmd(cmd.Options.OnlyOn, hostName, hostAddr) {
+			log.Printf("[DEBUG] skip command %q on host %q (%s)", cmd.Name, hostAddr, hostName)
+			continue
+		}
 
 		infoMsg := fmt.Sprintf("run command %q on host %q (%s)", cmd.Name, hostAddr, hostName)
 		if hostName == "" {
@@ -236,3 +240,29 @@ func (p *Process) execCommand(ctx context.Context, ec execCmd) (details string,
 		return "", nil, fmt.Errorf("unknown command %q", ec.cmd.Name)
 	}
 }
+
+// shouldRunCmd checks if the command should be executed on the host. If the command has no restrictions
+// (onlyOn field), it will be executed on all hosts. If the command has restrictions, it will be executed
+// only on the hosts that match the restrictions.
+// The onlyOn field can contain hostnames or IP addresses. If the hostname starts with "!", it will be
+// excluded from the list of hosts. If the hostname doesn't start with "!", it will be included in the list
+// of hosts. If the onlyOn field is empty, the command will be executed on all hosts.
+func (p *Process) shouldRunCmd(onlyOn []string, hostName, hostAddr string) bool {
+	if len(onlyOn) == 0 {
+		return true
+	}
+
+	for _, host := range onlyOn {
+		if strings.HasPrefix(host, "!") { // exclude host
+			if hostName == host[1:] || hostAddr == host[1:] {
+				return false
+			}
+			continue
+		}
+		if hostName == host || hostAddr == host { // include host
+			return true
+		}
+	}
+
+	return false
+}

pkg/runner/runner_test.go

@@ -61,6 +61,38 @@ func TestProcess_Run(t *testing.T) {
 		assert.Equal(t, 1, res.Hosts)
 	})
 
+	t.Run("simple playbook with only_on skip", func(t *testing.T) {
+		conf, err := config.New("testdata/conf-simple.yml", nil, nil)
+		require.NoError(t, err)
+		conf.Tasks[0].Commands[0].Options.OnlyOn = []string{"not-existing-host"}
+		p := Process{
+			Concurrency: 1,
+			Connector:   connector,
+			Config:      conf,
+			ColorWriter: executor.NewColorizedWriter(os.Stdout, "", "", "", nil),
+		}
+		res, err := p.Run(ctx, "default", testingHostAndPort)
+		require.NoError(t, err)
+		assert.Equal(t, 6, res.Commands, "should skip one command")
+		assert.Equal(t, 1, res.Hosts)
+	})
+
+	t.Run("simple playbook with only_on include", func(t *testing.T) {
+		conf, err := config.New("testdata/conf-simple.yml", nil, nil)
+		require.NoError(t, err)
+		conf.Tasks[0].Commands[0].Options.OnlyOn = []string{testingHostAndPort}
+		p := Process{
+			Concurrency: 1,
+			Connector:   connector,
+			Config:      conf,
+			ColorWriter: executor.NewColorizedWriter(os.Stdout, "", "", "", nil),
+		}
+		res, err := p.Run(ctx, "default", testingHostAndPort)
+		require.NoError(t, err)
+		assert.Equal(t, 7, res.Commands, "should include the only_on command")
+		assert.Equal(t, 1, res.Hosts)
+	})
+
 	t.Run("with runtime vars", func(t *testing.T) {
 		conf, err := config.New("testdata/conf.yml", nil, nil)
 		require.NoError(t, err)
@@ -545,6 +577,34 @@ func TestProcess_RunTaskWithWait(t *testing.T) {
 	assert.Contains(t, buf.String(), "wait done")
 }
 
+func TestProcess_shouldRunCmd(t *testing.T) {
+	p := &Process{}
+	tests := []struct {
+		name, hostName, hostAddr string
+		onlyOn                   []string
+		want                     bool
+	}{
+		{"Empty onlyOn list", "host1", "192.168.0.1", []string{}, true},
+		{"Hostname included", "host1", "192.168.0.1", []string{"host1", "host2"}, true},
+		{"Hostname excluded", "host1", "192.168.0.1", []string{"!host1", "host2"}, false},
+		{"Host address included", "host1", "192.168.0.1", []string{"192.168.0.1", "192.168.0.2"}, true},
+		{"Host address excluded", "host1", "192.168.0.1", []string{"!192.168.0.1", "192.168.0.2"}, false},
+		{"Host not included", "host1", "192.168.0.1", []string{"host2", "host3"}, false},
+		{"All hosts excluded", "host1", "192.168.0.1", []string{"!host1", "!host2"}, false},
+		{"All hosts included but one", "host3", "192.168.0.3", []string{"host1", "host2", "!host3"}, false},
+		{"Empty hostname, host address included", "", "192.168.0.1", []string{"192.168.0.1", "192.168.0.2"}, true},
+		{"Empty hostname, host address excluded", "", "192.168.0.1", []string{"!192.168.0.1", "192.168.0.2"}, false},
+		{"Empty hostname, host not included", "", "192.168.0.1", []string{"192.168.0.2", "192.168.0.3"}, false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := p.shouldRunCmd(tt.onlyOn, tt.hostName, tt.hostAddr)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}
+
 func startTestContainer(t *testing.T) (hostAndPort string, teardown func()) {
 	ctx := context.Background()
 	pubKey, err := os.ReadFile("testdata/test_ssh_key.pub")