Add --node flag to filter on a specific node (#243)

Takashi Kusumi · web-flow · commit f90f70f8c19a · 2023-03-05T13:08:07.000+09:00
diff --git a/README.md b/README.md
@@ -88,6 +88,7 @@ Supported Kubernetes resources are `pod`, `replicationcontroller`, `service`, `d
  `--max-log-requests`        | `-1`      | Maximum number of concurrent logs to request. Defaults to 50, but 5 when specifying --no-follow
  `--namespace`, `-n`         |           | Kubernetes namespace to use. Default to namespace configured in kubernetes context. To specify multiple namespaces, repeat this or set comma-separated value.
  `--no-follow`               | `false`   | Exit when all logs have been shown.
+ `--node`                    |           | Node name to filter on.
  `--only-log-lines`          | `false`   | Print only log lines
  `--output`, `-o`            | `default` | Specify predefined template. Currently support: [default, raw, json, extjson, ppextjson]
  `--prompt`, `-p`            | `false`   | Toggle interactive prompt for selecting 'app.kubernetes.io/instance' label values.
diff --git a/cmd/cmd.go b/cmd/cmd.go
@@ -73,6 +73,7 @@ type options struct {
 	verbosity           int
 	onlyLogLines        bool
 	maxLogRequests      int
+	node                string
 }
 
 func NewOptions(streams genericclioptions.IOStreams) *options {
@@ -190,12 +191,9 @@ func (o *options) sternConfig() (*stern.Config, error) {
 		}
 	}
 
-	fieldSelector := fields.Everything()
-	if o.fieldSelector != "" {
-		fieldSelector, err = fields.ParseSelector(o.fieldSelector)
-		if err != nil {
-			return nil, errors.Wrap(err, "failed to parse selector as field selector")
-		}
+	fieldSelector, err := o.generateFieldSelector()
+	if err != nil {
+		return nil, err
 	}
 
 	var tailLines *int64
@@ -298,6 +296,7 @@ func (o *options) AddFlags(fs *pflag.FlagSet) {
 	fs.StringVar(&o.kubeConfig, "kube-config", o.kubeConfig, "Path to kubeconfig file to use.")
 	_ = fs.MarkDeprecated("kube-config", "Use --kubeconfig instead.")
 	fs.StringSliceVarP(&o.namespaces, "namespace", "n", o.namespaces, "Kubernetes namespace to use. Default to namespace configured in kubernetes context. To specify multiple namespaces, repeat this or set comma-separated value.")
+	fs.StringVar(&o.node, "node", o.node, "Node name to filter on.")
 	fs.IntVar(&o.maxLogRequests, "max-log-requests", o.maxLogRequests, "Maximum number of concurrent logs to request. Defaults to 50, but 5 when specifying --no-follow")
 	fs.StringVarP(&o.output, "output", "o", o.output, "Specify predefined template. Currently support: [default, raw, json, extjson, ppextjson]")
 	fs.BoolVarP(&o.prompt, "prompt", "p", o.prompt, "Toggle interactive prompt for selecting 'app.kubernetes.io/instance' label values.")
@@ -432,6 +431,25 @@ func (o *options) generateTemplate() (*template.Template, error) {
 	return template, err
 }
 
+func (o *options) generateFieldSelector() (fields.Selector, error) {
+	var queries []string
+	if o.fieldSelector != "" {
+		queries = append(queries, o.fieldSelector)
+	}
+	if o.node != "" {
+		queries = append(queries, fmt.Sprintf("spec.nodeName=%s", o.node))
+	}
+	if len(queries) == 0 {
+		return fields.Everything(), nil
+	}
+
+	fieldSelector, err := fields.ParseSelector(strings.Join(queries, ","))
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to parse selector as field selector")
+	}
+	return fieldSelector, nil
+}
+
 func NewSternCmd(stream genericclioptions.IOStreams) (*cobra.Command, error) {
 	o := NewOptions(stream)
 
diff --git a/cmd/cmd_test.go b/cmd/cmd_test.go
@@ -376,7 +376,7 @@ func TestOptionsSternConfig(t *testing.T) {
 	local, _ := time.LoadLocation("Local")
 	utc, _ := time.LoadLocation("UTC")
 	labelSelector, _ := labels.Parse("l=sel")
-	fieldSelector, _ := fields.ParseSelector("f=field")
+	fieldSelector, _ := fields.ParseSelector("f=field,spec.nodeName=node1")
 
 	re := regexp.MustCompile
 
@@ -451,6 +451,7 @@ func TestOptionsSternConfig(t *testing.T) {
 				o.maxLogRequests = 30
 				o.resource = "res1"
 				o.onlyLogLines = true
+				o.node = "node1"
 
 				return o
 			}(),
@@ -484,6 +485,40 @@ func TestOptionsSternConfig(t *testing.T) {
 			}(),
 			false,
 		},
+		{
+			"fieldSelector without node",
+			func() *options {
+				o := NewOptions(streams)
+				o.fieldSelector = "f=field"
+
+				return o
+			}(),
+			func() *stern.Config {
+				c := defaultConfig()
+				sel, _ := fields.ParseSelector("f=field")
+				c.FieldSelector = sel
+
+				return c
+			}(),
+			false,
+		},
+		{
+			"node without fieldSelector",
+			func() *options {
+				o := NewOptions(streams)
+				o.node = "node1"
+
+				return o
+			}(),
+			func() *stern.Config {
+				c := defaultConfig()
+				sel, _ := fields.ParseSelector("spec.nodeName=node1")
+				c.FieldSelector = sel
+
+				return c
+			}(),
+			false,
+		},
 		{
 			"noFollow has the different default",
 			func() *options {
