diff --git a/README.md b/README.md index e9a72e06..cf2fcf1b 100644 --- a/README.md +++ b/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. diff --git a/pkg/config/command.go b/pkg/config/command.go index 199d4517..be2013fa 100644 --- a/pkg/config/command.go +++ b/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 diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go index a53ecf66..ff77f66f 100644 --- a/pkg/runner/runner.go +++ b/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 +} diff --git a/pkg/runner/runner_test.go b/pkg/runner/runner_test.go index 56b378dd..5789a763 100644 --- a/pkg/runner/runner_test.go +++ b/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")