kubevela/pkg/commands/exec.go

package commands

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/spf13/cobra"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/client-go/kubernetes"
	cmdexec "k8s.io/kubectl/pkg/cmd/exec"
	k8scmdutil "k8s.io/kubectl/pkg/cmd/util"

	"github.com/oam-dev/kubevela/apis/types"
	"github.com/oam-dev/kubevela/pkg/appfile"
	"github.com/oam-dev/kubevela/pkg/appfile/api"
	"github.com/oam-dev/kubevela/pkg/commands/util"
	"github.com/oam-dev/kubevela/pkg/oam"
)

const (
	podRunningTimeoutFlag = "pod-running-timeout"
	defaultPodExecTimeout = 60 * time.Second
	defaultStdin          = true
	defaultTTY            = true
)

// VelaExecOptions creates options for `exec` command
type VelaExecOptions struct {
	Cmd         *cobra.Command
	Args        []string
	Stdin       bool
	TTY         bool
	ServiceName string

	context.Context
	VelaC types.Args
	Env   *types.EnvMeta
	App   *api.Application

	f             k8scmdutil.Factory
	kcExecOptions *cmdexec.ExecOptions
	ClientSet     kubernetes.Interface
}

// NewExecCommand creates `exec` command
func NewExecCommand(c types.Args, ioStreams util.IOStreams) *cobra.Command {
	o := &VelaExecOptions{
		kcExecOptions: &cmdexec.ExecOptions{
			StreamOptions: cmdexec.StreamOptions{
				IOStreams: genericclioptions.IOStreams{
					In:     ioStreams.In,
					Out:    ioStreams.Out,
					ErrOut: ioStreams.ErrOut,
				},
			},
			Executor: &cmdexec.DefaultRemoteExecutor{},
		},
	}
	cmd := &cobra.Command{
		Use:   "exec [flags] APP_NAME -- COMMAND [args...]",
		Short: "Execute command in a container",
		Long:  "Execute command in a container",
		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
			if err := c.SetConfig(); err != nil {
				return err
			}
			o.VelaC = c
			return nil
		},
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) < 1 {
				ioStreams.Error("Please specify an application name.")
				return nil
			}
			argsLenAtDash := cmd.ArgsLenAtDash()
			if argsLenAtDash != 1 {
				ioStreams.Error("Please specify at least one command for the container.")
				return nil
			}

			if err := o.Init(context.Background(), cmd, args); err != nil {
				return err
			}
			if err := o.Complete(); err != nil {
				return err
			}
			if err := o.Run(); err != nil {
				return err
			}
			return nil
		},
		Annotations: map[string]string{
			types.TagCommandType: types.TypeApp,
		},
	}
	cmd.Flags().BoolVarP(&o.Stdin, "stdin", "i", defaultStdin, "Pass stdin to the container")
	cmd.Flags().BoolVarP(&o.TTY, "tty", "t", defaultTTY, "Stdin is a TTY")
	cmd.Flags().Duration(podRunningTimeoutFlag, defaultPodExecTimeout,
		"The length of time (like 5s, 2m, or 3h, higher than zero) to wait until at least one pod is running",
	)
	cmd.Flags().StringVarP(&o.ServiceName, "svc", "s", "", "service name")
	return cmd
}

// Init prepares the arguments accepted by the Exec command
func (o *VelaExecOptions) Init(ctx context.Context, c *cobra.Command, argsIn []string) error {
	o.Context = ctx
	o.Cmd = c
	o.Args = argsIn

	env, err := GetEnv(o.Cmd)
	if err != nil {
		return err
	}
	o.Env = env
	app, err := appfile.LoadApplication(env.Name, o.Args[0])
	if err != nil {
		return err
	}
	o.App = app

	cf := genericclioptions.NewConfigFlags(true)
	cf.Namespace = &o.Env.Namespace
	o.f = k8scmdutil.NewFactory(k8scmdutil.NewMatchVersionFlags(cf))

	if o.ClientSet == nil {
		c, err := kubernetes.NewForConfig(o.VelaC.Config)
		if err != nil {
			return err
		}
		o.ClientSet = c
	}
	return nil
}

// Complete loads data from the command environment
func (o *VelaExecOptions) Complete() error {
	compName, err := o.getComponentName()
	if err != nil {
		return err
	}
	podName, err := o.getPodName(compName)
	if err != nil {
		return err
	}
	o.kcExecOptions.StreamOptions.Stdin = o.Stdin
	o.kcExecOptions.StreamOptions.TTY = o.TTY

	args := make([]string, len(o.Args))
	copy(args, o.Args)
	// args for kcExecOptions MUST be in such formart:
	// [podName, COMMAND...]
	args[0] = podName
	return o.kcExecOptions.Complete(o.f, o.Cmd, args, 1)
}

func (o *VelaExecOptions) getComponentName() (string, error) {
	svcName := o.ServiceName

	if svcName != "" {
		if _, exist := o.App.Services[svcName]; exist {
			return svcName, nil
		}
		o.Cmd.Printf("The service name '%s' is not valid\n", svcName)
	}

	compName, err := util.AskToChooseOneService(appfile.GetComponents(o.App))
	if err != nil {
		return "", err
	}
	return compName, nil
}

func (o *VelaExecOptions) getPodName(compName string) (string, error) {
	podList, err := o.ClientSet.CoreV1().Pods(o.Env.Namespace).List(o.Context, v1.ListOptions{
		LabelSelector: labels.Set(map[string]string{
			// TODO(roywang) except core workloads, not any workloads will pass these label to pod
			// find a rigorous way to get pod by compname
			oam.LabelAppComponent: compName,
		}).String(),
	})
	if err != nil {
		return "", nil
	}
	if podList != nil && len(podList.Items) == 0 {
		return "", fmt.Errorf("cannot get pods")
	}
	for _, p := range podList.Items {
		if strings.HasPrefix(p.Name, compName+"-") {
			return p.Name, nil
		}
	}
	// if no pod with name matched prefix as component name
	// just return the first one
	return podList.Items[0].Name, nil
}

// Run executes a validated remote execution against a pod
func (o *VelaExecOptions) Run() error {
	return o.kcExecOptions.Run()
}