recover.go

Documentation: github.com/chaos-mesh/chaos-mesh/pkg/chaosctl/cmd

     1  // Copyright 2021 Chaos Mesh Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  // http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  
    16  package cmd
    17  
    18  import (
    19  	"context"
    20  	"fmt"
    21  	"strings"
    22  
    23  	"github.com/go-logr/logr"
    24  	"github.com/pkg/errors"
    25  	"github.com/spf13/cobra"
    26  
    27  	"github.com/chaos-mesh/chaos-mesh/api/v1alpha1"
    28  	"github.com/chaos-mesh/chaos-mesh/pkg/chaosctl/common"
    29  	"github.com/chaos-mesh/chaos-mesh/pkg/chaosctl/recover"
    30  	ctrlclient "github.com/chaos-mesh/chaos-mesh/pkg/ctrl/client"
    31  	"github.com/chaos-mesh/chaos-mesh/pkg/label"
    32  )
    33  
    34  type RecoverOptions struct {
    35  	namespace string
    36  	labels    *[]string
    37  }
    38  
    39  func NewRecoverCommand(logger logr.Logger, builders map[string]recover.RecovererBuilder) (*cobra.Command, error) {
    40  	o := &RecoverOptions{namespace: "default"}
    41  
    42  	recoverCmd := &cobra.Command{
    43  		Use:   `recover (CHAOSTYPE) POD[,POD[,POD...]] [-n NAMESPACE]`,
    44  		Short: `Recover certain chaos from certain pods`,
    45  		Long: `Recover certain chaos from certain pods.
    46  Currently unimplemented.
    47  
    48  Examples:
    49    # Recover network chaos from pods in namespace default
    50    chaosctl recover networkchaos
    51  
    52    # Recover network chaos from certain pods in certain namespace
    53    chaosctl recover networkchaos pod1 pod2 pod3 -n NAMESPACE
    54    
    55    # Recover network chaos from pods with label key=value
    56    chaosctl recover networkchaos -l key=value`,
    57  		ValidArgsFunction: noCompletions,
    58  	}
    59  
    60  	for chaosType, builder := range builders {
    61  		recoverCmd.AddCommand(recoverResourceCommand(o, chaosType, builder))
    62  	}
    63  
    64  	recoverCmd.PersistentFlags().StringVarP(&o.namespace, "namespace", "n", "default", "namespace to find pods")
    65  	o.labels = recoverCmd.PersistentFlags().StringSliceP("label", "l", nil, "labels to select pods")
    66  	err := recoverCmd.RegisterFlagCompletionFunc("namespace", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
    67  		client, cancel, err := common.CreateClient(context.TODO(), managerNamespace, managerSvc)
    68  		if err != nil {
    69  			logger.Error(err, "create client")
    70  			return nil, cobra.ShellCompDirectiveNoFileComp
    71  		}
    72  		defer cancel()
    73  
    74  		completion, err := client.ListNamespace(context.TODO())
    75  		if err != nil {
    76  			logger.Error(err, "complete resource")
    77  			return nil, cobra.ShellCompDirectiveNoFileComp
    78  		}
    79  
    80  		return completion, cobra.ShellCompDirectiveNoFileComp | cobra.ShellCompDirectiveNoSpace
    81  	})
    82  	if err != nil {
    83  		return nil, errors.Wrap(err, "register completion func for flag `namespace`")
    84  	}
    85  	err = recoverCmd.RegisterFlagCompletionFunc("label", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
    86  		return nil, cobra.ShellCompDirectiveNoFileComp
    87  	})
    88  	if err != nil {
    89  		return nil, errors.Wrap(err, "register completion func for flag `label`")
    90  	}
    91  	return recoverCmd, nil
    92  }
    93  
    94  func recoverResourceCommand(option *RecoverOptions, chaosType string, builder recover.RecovererBuilder) *cobra.Command {
    95  	return &cobra.Command{
    96  		Use:   fmt.Sprintf(`%s POD[,POD[,POD...]] [-n NAMESPACE]`, chaosType),
    97  		Short: fmt.Sprintf(`Recover %s from certain pods`, chaosType),
    98  		Long:  fmt.Sprintf(`Recover %s from certain pods`, chaosType),
    99  		RunE: func(cmd *cobra.Command, args []string) error {
   100  			client, cancel, err := common.CreateClient(context.TODO(), managerNamespace, managerSvc)
   101  			if err != nil {
   102  				return err
   103  			}
   104  			defer cancel()
   105  			return option.Run(builder(client), client, args)
   106  		},
   107  		SilenceErrors: true,
   108  		SilenceUsage:  true,
   109  		ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
   110  			if len(args) != 0 {
   111  				return []string{}, cobra.ShellCompDirectiveNoFileComp
   112  			}
   113  			client, cancel, err := common.CreateClient(context.TODO(), managerNamespace, managerSvc)
   114  			if err != nil {
   115  				common.PrettyPrint(errors.Wrap(err, "create client").Error(), 0, common.Red)
   116  				return nil, cobra.ShellCompDirectiveNoFileComp
   117  			}
   118  			defer cancel()
   119  			return option.List(client)
   120  		},
   121  	}
   122  }
   123  
   124  // Run recover
   125  func (o *RecoverOptions) Run(recover recover.Recoverer, client *ctrlclient.CtrlClient, args []string) error {
   126  	pods, err := o.selectPods(client, args)
   127  	if err != nil {
   128  		return err
   129  	}
   130  	ctx, cancel := context.WithCancel(context.Background())
   131  	defer cancel()
   132  
   133  	for _, pod := range pods {
   134  		err = recover.Recover(ctx, pod)
   135  		if err != nil {
   136  			return err
   137  		}
   138  	}
   139  	return nil
   140  }
   141  
   142  // List pods to recover
   143  func (o *RecoverOptions) List(client *ctrlclient.CtrlClient) ([]string, cobra.ShellCompDirective) {
   144  	pods, err := o.selectPods(client, []string{})
   145  	if err != nil {
   146  		common.PrettyPrint(errors.Wrap(err, "select pods").Error(), 0, common.Red)
   147  		return nil, cobra.ShellCompDirectiveNoFileComp
   148  	}
   149  
   150  	var names []string
   151  	for _, pod := range pods {
   152  		names = append(names, pod.Name)
   153  	}
   154  
   155  	return names, cobra.ShellCompDirectiveNoFileComp
   156  }
   157  
   158  func (o *RecoverOptions) selectPods(client *ctrlclient.CtrlClient, names []string) ([]*recover.PartialPod, error) {
   159  	ctx, cancel := context.WithCancel(context.Background())
   160  	defer cancel()
   161  
   162  	selector := v1alpha1.PodSelectorSpec{
   163  		GenericSelectorSpec: v1alpha1.GenericSelectorSpec{
   164  			Namespaces: []string{o.namespace},
   165  		},
   166  	}
   167  
   168  	if len(names) != 0 {
   169  		selector.Pods = map[string][]string{o.namespace: names}
   170  	}
   171  
   172  	if o.labels != nil && len(*o.labels) > 0 {
   173  		labels, err := label.ParseLabel(strings.Join(*o.labels, ","))
   174  		if err != nil {
   175  			return nil, errors.Wrap(err, "parse labels")
   176  		}
   177  		if len(labels) != 0 {
   178  			selector.LabelSelectors = labels
   179  		}
   180  	}
   181  
   182  	return recover.SelectPods(ctx, client, selector)
   183  }
   184