From 20bb61e95a0582da0431e22dfb6015f588b32f21 Mon Sep 17 00:00:00 2001
From: guoxudong <guoxudong.dev@gmial.com>
Date: Tue, 23 Mar 2021 11:03:10 +0800
Subject: [PATCH] add labelselector

---
 cmd/plugin/cli/root.go               |  4 +++-
 doc/en/more-info.md                  | 10 +++++++--
 doc/zh/{pod-lens.md => more-info.md} | 12 ++++++++---
 pkg/plugin/plugin.go                 | 31 ++++++++++++++++++++--------
 4 files changed, 42 insertions(+), 15 deletions(-)
 rename doc/zh/{pod-lens.md => more-info.md} (57%)

diff --git a/cmd/plugin/cli/root.go b/cmd/plugin/cli/root.go
index ad688b6..9c95766 100644
--- a/cmd/plugin/cli/root.go
+++ b/cmd/plugin/cli/root.go
@@ -20,6 +20,7 @@ import (
 var (
 	KubernetesConfigFlags *genericclioptions.ConfigFlags
 	allNamespacesFlag     bool
+	labelFlag             string
 )
 
 func RootCmd() *cobra.Command {
@@ -47,7 +48,7 @@ $ kubectl pod-lens prometheus-prometheus-operator-prometheus-0
 			argsChannel := make(chan string, 1)
 			argsChannel <- podName
 
-			if err := plugin.RunPlugin(KubernetesConfigFlags, argsChannel, allNamespacesFlag); err != nil {
+			if err := plugin.RunPlugin(KubernetesConfigFlags, argsChannel, allNamespacesFlag, labelFlag); err != nil {
 				return errors.Cause(err)
 			}
 
@@ -60,6 +61,7 @@ $ kubectl pod-lens prometheus-prometheus-operator-prometheus-0
 	KubernetesConfigFlags = genericclioptions.NewConfigFlags(false)
 	KubernetesConfigFlags.AddFlags(cmd.Flags())
 	cmd.Flags().BoolVarP(&allNamespacesFlag, "all-namespaces", "A", false, "query all objects in all API groups, both namespaced and non-namespaced")
+	cmd.Flags().StringVarP(&labelFlag, "selector", "l", "", "Selector (label query) to filter on, only supports '=' and a parameter (e.g. -l key1=value1)")
 
 	klog.InitFlags(nil)
 	cmd.Flags().AddGoFlagSet(flag.CommandLine)
diff --git a/doc/en/more-info.md b/doc/en/more-info.md
index 7804a27..b3dd823 100644
--- a/doc/en/more-info.md
+++ b/doc/en/more-info.md
@@ -10,8 +10,14 @@ title: More Info
 kubectl pod-lens
 ```
 
-##  Show pod-related resources
+## Show pod-related resources
 
 ```shell
 kubectl pod-lens <pod-name>
-```
\ No newline at end of file
+```
+
+## Assign LabelSelector
+
+```shell
+kubectl pod-lens <pod-name> -l app=demo
+```
diff --git a/doc/zh/pod-lens.md b/doc/zh/more-info.md
similarity index 57%
rename from doc/zh/pod-lens.md
rename to doc/zh/more-info.md
index 9832d5c..b94b314 100644
--- a/doc/zh/pod-lens.md
+++ b/doc/zh/more-info.md
@@ -1,5 +1,5 @@
 ---
-title: Pod Lens
+title: 更多信息
 ---
 
 [![asciicast](https://asciinema.org/a/400180.svg)](https://asciinema.org/a/400180)
@@ -10,8 +10,14 @@ title: Pod Lens
 kubectl pod-lens
 ```
 
-##  展示 Pod 相关 K8S 资源
+## 展示 Pod 相关 K8S 资源
 
 ```shell
 kubectl pod-lens <pod-name>
-```
\ No newline at end of file
+```
+
+## 指定 LabelSelector
+
+```shell
+kubectl pod-lens <pod-name> -l app=demo
+```
diff --git a/pkg/plugin/plugin.go b/pkg/plugin/plugin.go
index 8331577..f5f7e37 100644
--- a/pkg/plugin/plugin.go
+++ b/pkg/plugin/plugin.go
@@ -5,6 +5,7 @@ import (
 	select_pod "github.com/sunny0826/kubectl-pod-lens/pkg/select-pod"
 	"k8s.io/klog"
 	"os"
+	"regexp"
 	"strings"
 
 	mapset "github.com/deckarep/golang-set"
@@ -103,13 +104,25 @@ func (sf *SnifferPlugin) findNodeByName() error {
 	return nil
 }
 
-func (sf *SnifferPlugin) getLabelByPod() error {
+func (sf *SnifferPlugin) getLabelByPod(labelFlag string) error {
+	if labelFlag != "" {
+		match, err := regexp.MatchString("[a-z0-9/-]+=([a-z0-9/-]+)( |$)", labelFlag)
+		if err != nil {
+			return err
+		}
+		if match {
+			sf.LabelSelector = labelFlag
+			return nil
+		} else {
+			return errors.New(labelFlag + " is incorrectly formatted.")
+		}
+	}
 	var labelSelector string
 	labels := sf.PodObject.Labels
-	if _, ok := labels["release"]; ok {
-		labelSelector = "release=" + labels["release"]
-	} else if _, ok = labels["app"]; ok {
+	if _, ok := labels["app"]; ok {
 		labelSelector = "app=" + labels["app"]
+	} else if _, ok = labels["release"]; ok {
+		labelSelector = "release=" + labels["release"]
 	} else if _, ok = labels["k8s-app"]; ok {
 		labelSelector = "k8s-app=" + labels["k8s-app"]
 	} else if _, ok = labels["app.kubernetes.io/name"]; ok {
@@ -557,7 +570,7 @@ func (sf *SnifferPlugin) printResource() error {
 	return nil
 }
 
-func RunPlugin(configFlags *genericclioptions.ConfigFlags, outputCh chan string, allNamespacesFlag bool) error {
+func RunPlugin(configFlags *genericclioptions.ConfigFlags, outputCh chan string, allNamespacesFlag bool, labelFlag string) error {
 	klog.V(1).Info("start run plugins")
 	sf, err := NewSnifferPlugin(configFlags)
 	if err != nil {
@@ -570,15 +583,15 @@ func RunPlugin(configFlags *genericclioptions.ConfigFlags, outputCh chan string,
 		namespace = getNamespace(configFlags)
 	}
 
-	if err := sf.findPodByName(podName, namespace); err != nil {
+	if err = sf.findPodByName(podName, namespace); err != nil {
 		return err
 	}
 
-	if err := sf.findNodeByName(); err != nil {
+	if err = sf.findNodeByName(); err != nil {
 		return err
 	}
 
-	if err := sf.getOwnerByPod(); err != nil {
+	if err = sf.getOwnerByPod(); err != nil {
 		return err
 	}
 
@@ -586,7 +599,7 @@ func RunPlugin(configFlags *genericclioptions.ConfigFlags, outputCh chan string,
 		return err
 	}
 
-	if err = sf.getLabelByPod(); err != nil {
+	if err = sf.getLabelByPod(labelFlag); err != nil {
 		return err
 	}