cluster: add subcommand pull and push to transfer files (#1044)

* cluster: add subcommand pull and push to transfer files

* cluster: use correct path order in args

* cluster & dm: mark `exec` as hidden

* cluster: fix host iteration

* cluster: minor optimization of command render

* cluster: add simple CI test for push/pull

Co-authored-by: SIGSEGV <gnu.crazier@gmail.com>

components/cluster/command/exec.go

@@ -21,8 +21,9 @@ import (
 func newExecCmd() *cobra.Command {
 	opt := manager.ExecOptions{}
 	cmd := &cobra.Command{
-		Use:   "exec <cluster-name>",
-		Short: "Run shell command on host in the tidb cluster",
+		Use:    "exec <cluster-name>",
+		Short:  "Run shell command on host in the tidb cluster",
+		Hidden: true,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			if len(args) != 1 {
 				return cmd.Help()

components/cluster/command/root.go

@@ -149,7 +149,6 @@ func init() {
 		newDestroyCmd(),
 		newCleanCmd(),
 		newUpgradeCmd(),
-		newExecCmd(),
 		newDisplayCmd(),
 		newPruneCmd(),
 		newListCmd(),
@@ -161,6 +160,9 @@ func init() {
 		newRenameCmd(),
 		newEnableCmd(),
 		newDisableCmd(),
+		newExecCmd(),
+		newPullCmd(),
+		newPushCmd(),
 		newTestCmd(), // hidden command for test internally
 		newTelemetryCmd(),
 	)

components/cluster/command/transfer.go

@@ -0,0 +1,78 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package command
+
+import (
+	"github.com/pingcap/tiup/pkg/cluster/manager"
+	"github.com/spf13/cobra"
+)
+
+/* Add a pair of adb like commands to transfer files to or from remote
+   servers. Not using `scp` as the real implementation is not necessarily
+   SSH, not using `transfer` all-in-one command to get rid of complex
+   checking of wheather a path is remote or local, as this is supposed
+   to be only a tiny helper utility.
+*/
+
+func newPullCmd() *cobra.Command {
+	opt := manager.TransferOptions{Pull: true}
+	cmd := &cobra.Command{
+		Use:    "pull <cluster-name> <remote-path> <local-path>",
+		Short:  "(EXPERIMENTAL) Transfer files or directories from host in the tidb cluster to local",
+		Hidden: true,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			if len(args) != 3 {
+				return cmd.Help()
+			}
+
+			clusterName := args[0]
+			opt.Remote = args[1]
+			opt.Local = args[2]
+			teleCommand = append(teleCommand, scrubClusterName(clusterName))
+
+			return cm.Transfer(clusterName, opt, gOpt)
+		},
+	}
+
+	cmd.Flags().StringSliceVarP(&gOpt.Roles, "role", "R", nil, "Only exec on host with specified roles")
+	cmd.Flags().StringSliceVarP(&gOpt.Nodes, "node", "N", nil, "Only exec on host with specified nodes")
+
+	return cmd
+}
+
+func newPushCmd() *cobra.Command {
+	opt := manager.TransferOptions{Pull: false}
+	cmd := &cobra.Command{
+		Use:    "push <cluster-name> <local-path> <remote-path>",
+		Short:  "(EXPERIMENTAL) Transfer files or directories from local to host in the tidb cluster",
+		Hidden: true,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			if len(args) != 3 {
+				return cmd.Help()
+			}
+
+			clusterName := args[0]
+			opt.Local = args[1]
+			opt.Remote = args[2]
+			teleCommand = append(teleCommand, scrubClusterName(clusterName))
+
+			return cm.Transfer(clusterName, opt, gOpt)
+		},
+	}
+
+	cmd.Flags().StringSliceVarP(&gOpt.Roles, "role", "R", nil, "Only exec on host with specified roles")
+	cmd.Flags().StringSliceVarP(&gOpt.Nodes, "node", "N", nil, "Only exec on host with specified nodes")
+
+	return cmd
+}

components/dm/command/exec.go

@@ -21,8 +21,9 @@ import (
 func newExecCmd() *cobra.Command {
 	opt := manager.ExecOptions{}
 	cmd := &cobra.Command{
-		Use:   "exec <cluster-name>",
-		Short: "Run shell command on host in the dm cluster",
+		Use:    "exec <cluster-name>",
+		Short:  "Run shell command on host in the dm cluster",
+		Hidden: true,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			if len(args) != 1 {
 				return cmd.Help()

pkg/cluster/manager/transfer.go

@@ -0,0 +1,149 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package manager
+
+import (
+	"bytes"
+	"fmt"
+	"html/template"
+	"strings"
+
+	"github.com/google/uuid"
+	"github.com/joomcode/errorx"
+	perrs "github.com/pingcap/errors"
+	operator "github.com/pingcap/tiup/pkg/cluster/operation"
+	"github.com/pingcap/tiup/pkg/cluster/spec"
+	"github.com/pingcap/tiup/pkg/cluster/task"
+	"github.com/pingcap/tiup/pkg/logger/log"
+	"github.com/pingcap/tiup/pkg/set"
+)
+
+// TransferOptions for exec shell commanm.
+type TransferOptions struct {
+	Local  string
+	Remote string
+	Pull   bool // default to push
+}
+
+// Transfer copies files from or to host in the tidb cluster.
+func (m *Manager) Transfer(name string, opt TransferOptions, gOpt operator.Options) error {
+	metadata, err := m.meta(name)
+	if err != nil {
+		return err
+	}
+
+	topo := metadata.GetTopology()
+	base := metadata.GetBaseMeta()
+
+	filterRoles := set.NewStringSet(gOpt.Roles...)
+	filterNodes := set.NewStringSet(gOpt.Nodes...)
+
+	var shellTasks []task.Task
+
+	uniqueHosts := map[string]set.StringSet{} // host-sshPort-port -> {remote-path}
+	topo.IterInstance(func(inst spec.Instance) {
+		key := fmt.Sprintf("%s-%d-%d", inst.GetHost(), inst.GetSSHPort(), inst.GetPort())
+		if _, found := uniqueHosts[key]; !found {
+			if len(gOpt.Roles) > 0 && !filterRoles.Exist(inst.Role()) {
+				return
+			}
+
+			if len(gOpt.Nodes) > 0 && !filterNodes.Exist(inst.GetHost()) {
+				return
+			}
+
+			// render remote path
+			instPath := opt.Remote
+			paths, err := renderInstanceSpec(instPath, inst)
+			if err != nil {
+				return // skip
+			}
+			pathSet := set.NewStringSet(paths...)
+			if _, ok := uniqueHosts[key]; ok {
+				uniqueHosts[key].Join(pathSet)
+				return
+			}
+			uniqueHosts[key] = pathSet
+		}
+	})
+
+	srcPath := opt.Local
+	for hostKey, i := range uniqueHosts {
+		host := strings.Split(hostKey, "-")[0]
+		for _, p := range i.Slice() {
+			t := task.NewBuilder()
+			if opt.Pull {
+				t.CopyFile(p, srcPath, host, opt.Pull)
+			} else {
+				t.CopyFile(srcPath, p, host, opt.Pull)
+			}
+			shellTasks = append(shellTasks, t.Build())
+		}
+	}
+
+	t := m.sshTaskBuilder(name, topo, base.User, gOpt).
+		Parallel(false, shellTasks...).
+		Build()
+
+	execCtx := task.NewContext()
+	if err := t.Execute(execCtx); err != nil {
+		if errorx.Cast(err) != nil {
+			// FIXME: Map possible task errors and give suggestions.
+			return err
+		}
+		return perrs.Trace(err)
+	}
+
+	return nil
+}
+
+func renderInstanceSpec(t string, inst spec.Instance) ([]string, error) {
+	result := make([]string, 0)
+	switch inst.ComponentName() {
+	case spec.ComponentTiFlash:
+		for _, d := range strings.Split(inst.DataDir(), ",") {
+			tf := inst
+			tfs, ok := tf.(*spec.TiFlashInstance).InstanceSpec.(spec.TiFlashSpec)
+			if !ok {
+				return result, fmt.Errorf("instance type mismatch for %v", inst)
+			}
+			tfs.DataDir = d
+			key := inst.ID() + d + uuid.New().String()
+			if s, err := renderSpec(t, tfs, key); err == nil {
+				result = append(result, s)
+			}
+		}
+	default:
+		s, err := renderSpec(t, inst, inst.ID())
+		if err != nil {
+			return result, fmt.Errorf("error rendering path for instance %v", inst)
+		}
+		result = append(result, s)
+	}
+	return result, nil
+}
+
+func renderSpec(t string, s interface{}, id string) (string, error) {
+	tpl, err := template.New(id).Option("missingkey=error").Parse(t)
+	if err != nil {
+		return "", err
+	}
+
+	result := bytes.NewBufferString("")
+	if err := tpl.Execute(result, s); err != nil {
+		log.Debugf("missing key when parsing: %s", err)
+		return "", err
+	}
+	return result.String(), nil
+}

tests/tiup-cluster/script/cmd_subtest.sh

@@ -101,6 +101,12 @@ function cmd_subtest() {
 
     tiup-cluster $client --yes clean $name --data --all --ignore-node n1:9090
 
+    # Test push and pull
+    echo "test_transfer $name $RANDOM `date`" > test_transfer_1.txt
+    tiup-cluster $client push $name test_transfer_1.txt "{{ .LogDir }}/test_transfer.txt" -R grafana
+    tiup-cluster $client pull $name "{{ .LogDir }}/test_transfer.txt" test_transfer_2.txt -R grafana
+    diff test_transfer_1.txt test_transfer_2.txt
+
     echo "checking cleanup data and log"
     tiup-cluster $client exec $name -N n1 --command "ls /home/tidb/deploy/prometheus-9090/log/prometheus.log"
     ! tiup-cluster $client exec $name -N n1 --command "ls /home/tidb/deploy/tikv-20160/log/tikv.log"