feat: filter by annotations

close #42 Signed-off-by: Christian Kotzbauer <[email protected]>
ckotzbauer · Jun 23, 2021 · 376a462 · 376a462
1 parent 46dda87
commit 376a462
Show file tree

Hide file tree

Showing 9 changed files with 115 additions and 57 deletions.
diff --git a/cmd/ha.go b/cmd/ha.go
@@ -12,15 +12,17 @@ var haCmd = &cobra.Command{
 	Use:   "ha",
 	Short: "Creates high-availability report of your workload.",
 	Run: func(cmd *cobra.Command, args []string) {
-		selector, _ := cmd.Flags().GetString("selector")
+		labelSelector, _ := cmd.Flags().GetString("selector")
+		annotationSelector, _ := cmd.Flags().GetString("annotation")
 		namespace, _ := cmd.Flags().GetString("namespace")
 
 		r := ha.HighAvailability{
-			KubeOverrides: overrides,
-			KubeClient:    kubernetes.NewClient(cmd, overrides),
-			Pods:          args,
-			Selector:      selector,
-			Namespace:     namespace,
+			KubeOverrides:      overrides,
+			KubeClient:         kubernetes.NewClient(cmd, overrides),
+			Pods:               args,
+			LabelSelector:      labelSelector,
+			AnnotationSelector: annotationSelector,
+			Namespace:          namespace,
 		}
 
 		list := r.Execute()
@@ -36,5 +38,6 @@ var haCmd = &cobra.Command{
 func init() {
 	rootCmd.AddCommand(haCmd)
 	haCmd.Flags().StringP("selector", "l", "", "Label-Selector")
+	haCmd.Flags().StringP("annotation", "a", "", "Annotation-Selector")
 	// Output
 }
diff --git a/cmd/resources.go b/cmd/resources.go
@@ -22,7 +22,8 @@ var resourcesCmd = &cobra.Command{
 		countDays, _ := cmd.Flags().GetInt64("count-days")
 		timeout, _ := cmd.Flags().GetDuration("timeout")
 
-		selector, _ := cmd.Flags().GetString("selector")
+		labelSelector, _ := cmd.Flags().GetString("selector")
+		annotationSelector, _ := cmd.Flags().GetString("annotation")
 		namespace, _ := cmd.Flags().GetString("namespace")
 
 		r := resources.Resource{
@@ -33,11 +34,12 @@ var resourcesCmd = &cobra.Command{
 				CountDays: countDays,
 				Timeout:   timeout,
 			},
-			KubeOverrides: overrides,
-			KubeClient:    kubernetes.NewClient(cmd, overrides),
-			Pods:          args,
-			Selector:      selector,
-			Namespace:     namespace,
+			KubeOverrides:      overrides,
+			KubeClient:         kubernetes.NewClient(cmd, overrides),
+			Pods:               args,
+			LabelSelector:      labelSelector,
+			AnnotationSelector: annotationSelector,
+			Namespace:          namespace,
 		}
 
 		list := r.Execute()
@@ -59,6 +61,7 @@ func init() {
 	resourcesCmd.Flags().DurationP("timeout", "t", time.Duration(30)*time.Second, "Timeout")
 
 	resourcesCmd.Flags().StringP("selector", "l", "", "Label-Selector")
+	resourcesCmd.Flags().StringP("annotation", "a", "", "Annotation-Selector")
 
 	resourcesCmd.MarkFlagRequired("prometheus-url")
 }
diff --git a/pkg/ha/ha.go b/pkg/ha/ha.go
@@ -3,6 +3,7 @@ package ha
 import (
 	"strings"
 
+	"github.com/ckotzbauer/chekr/pkg/kubernetes"
 	"github.com/ckotzbauer/chekr/pkg/printer"
 	"github.com/ckotzbauer/chekr/pkg/util"
 	"github.com/ddelizia/channelify"
@@ -13,13 +14,12 @@ import (
 )
 
 func (h HighAvailability) Execute() printer.PrintableList {
-	var pods []corev1.Pod
-
-	if h.Selector == "" {
-		pods = h.KubeClient.GetNamespacedPods(h.Namespace, h.Pods)
-	} else if h.Selector != "" {
-		pods = h.KubeClient.ListPods(h.Namespace, h.Selector)
-	}
+	pods := h.KubeClient.ListPods(kubernetes.PodQuery{
+		Namespace:          h.Namespace,
+		LabelSelector:      h.LabelSelector,
+		Names:              h.Pods,
+		AnnotationSelector: h.AnnotationSelector,
+	})
 
 	fn1 := func(h HighAvailability, pod corev1.Pod) printer.Printable {
 		return h.analyzePod(pod)

diff --git a/pkg/ha/types.go b/pkg/ha/types.go
@@ -27,11 +27,12 @@ Standalone pod: If the pod is removed or crashes, it disappears without replacem
 `
 
 type HighAvailability struct {
-	KubeOverrides *clientcmd.ConfigOverrides
-	KubeClient    *kubernetes.KubeClient
-	Pods          []string
-	Namespace     string
-	Selector      string
+	KubeOverrides      *clientcmd.ConfigOverrides
+	KubeClient         *kubernetes.KubeClient
+	Pods               []string
+	Namespace          string
+	LabelSelector      string
+	AnnotationSelector string
 }
 
 type PodAvailability struct {

diff --git a/pkg/kubernetes/kubernetes.go b/pkg/kubernetes/kubernetes.go
@@ -2,8 +2,10 @@ package kubernetes
 
 import (
 	"context"
+	"fmt"
 	"strings"
 
+	"github.com/ckotzbauer/chekr/pkg/util"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
@@ -21,6 +23,13 @@ type KubeClient struct {
 	Config *rest.Config
 }
 
+type PodQuery struct {
+	Namespace          string
+	LabelSelector      string
+	AnnotationSelector string
+	Names              []string
+}
+
 func IsInCluster() bool {
 	_, err := rest.InClusterConfig()
 	return err == nil
@@ -61,34 +70,50 @@ func NewClient(cmd *cobra.Command, configOverrides *clientcmd.ConfigOverrides) *
 	return &KubeClient{Client: client, Config: config}
 }
 
-func (kubeClient *KubeClient) ListPods(namespace, labelSelector string) []corev1.Pod {
-	list, err := kubeClient.Client.CoreV1().Pods(namespace).List(context.Background(), metav1.ListOptions{LabelSelector: labelSelector})
-	if err != nil {
-		logrus.WithError(err).WithField("namespace", namespace).WithField("labelSelector", labelSelector).Fatalf("Could list pods!")
-	}
+func (kubeClient *KubeClient) ListPods(query PodQuery) []corev1.Pod {
+	listOptions := metav1.ListOptions{}
 
-	return list.Items
-}
+	if len(query.LabelSelector) > 0 {
+		listOptions.LabelSelector = query.LabelSelector
+	}
 
-func (kubeClient *KubeClient) GetNamespacedPods(namespace string, names []string) []corev1.Pod {
-	var pods []corev1.Pod
+	list, err := kubeClient.Client.CoreV1().Pods(query.Namespace).List(context.Background(), listOptions)
+	pods := []corev1.Pod{}
 
-	if len(names) == 0 {
-		// all pods of the given namespace
-		list, err := kubeClient.Client.CoreV1().Pods(namespace).List(context.Background(), metav1.ListOptions{})
+	if err != nil {
+		logrus.WithError(
+			err).WithField(
+			"namespace", query.Namespace).WithField(
+			"labelSelector", query.LabelSelector).Fatalf(
+			"Could list pods!")
+	}
 
-		if err == nil {
-			pods = list.Items
+	for _, p := range list.Items {
+		if len(query.Names) > 0 && !util.Contains(query.Names, p.Name) {
+			continue
 		}
-	} else {
-		// the given pods in the given namespace
-		for _, p := range names {
-			pod, err := kubeClient.Client.CoreV1().Pods(namespace).Get(context.Background(), p, metav1.GetOptions{})
 
-			if err == nil {
-				pods = append(pods, *pod)
+		if len(query.AnnotationSelector) > 0 {
+			selectors := util.ParseSelector(query.AnnotationSelector)
+			fmt.Println(selectors)
+			matched := true
+			for _, selector := range selectors {
+				potentialValue := p.Annotations[selector.Key]
+				if selector.Operator == "=" {
+					matched = matched && potentialValue == selector.Value
+				} else if selector.Operator == "!=" {
+					matched = matched && potentialValue != selector.Value
+				} else if selector.Operator == "" {
+					matched = matched && potentialValue != ""
+				}
+			}
+
+			if !matched {
+				continue
 			}
 		}
+
+		pods = append(pods, p)
 	}
 
 	return pods

diff --git a/pkg/resources/resources.go b/pkg/resources/resources.go
@@ -5,6 +5,7 @@ import (
 
 	corev1 "k8s.io/api/core/v1"
 
+	"github.com/ckotzbauer/chekr/pkg/kubernetes"
 	"github.com/ckotzbauer/chekr/pkg/printer"
 	"github.com/ckotzbauer/chekr/pkg/prometheus"
 	"github.com/ckotzbauer/chekr/pkg/util"
@@ -22,13 +23,12 @@ func (r Resource) Execute() printer.PrintableList {
 		Step:  time.Minute * 5,
 	}
 
-	var pods []corev1.Pod
-
-	if r.Selector == "" {
-		pods = r.KubeClient.GetNamespacedPods(r.Namespace, r.Pods)
-	} else if r.Selector != "" {
-		pods = r.KubeClient.ListPods(r.Namespace, r.Selector)
-	}
+	pods := r.KubeClient.ListPods(kubernetes.PodQuery{
+		Namespace:          r.Namespace,
+		LabelSelector:      r.LabelSelector,
+		Names:              r.Pods,
+		AnnotationSelector: r.AnnotationSelector,
+	})
 
 	fn1 := func(r Resource, pod corev1.Pod, v1api v1.API, queryRange v1.Range) printer.Printable {
 		return r.analyzePod(pod, v1api, queryRange)

diff --git a/pkg/resources/types.go b/pkg/resources/types.go
@@ -12,12 +12,13 @@ import (
 )
 
 type Resource struct {
-	Prometheus    prometheus.Prometheus
-	KubeOverrides *clientcmd.ConfigOverrides
-	KubeClient    *kubernetes.KubeClient
-	Pods          []string
-	Namespace     string
-	Selector      string
+	Prometheus         prometheus.Prometheus
+	KubeOverrides      *clientcmd.ConfigOverrides
+	KubeClient         *kubernetes.KubeClient
+	Pods               []string
+	Namespace          string
+	LabelSelector      string
+	AnnotationSelector string
 }
 
 type AnalyzedValues struct {

diff --git a/pkg/util/types.go b/pkg/util/types.go
@@ -4,3 +4,9 @@ type ComputedValue struct {
 	Value      float64
 	Percentage float64
 }
+
+type KeyValueSelector struct {
+	Key string
+	Operator string // can be "=" or "!="
+	Value string
+}
diff --git a/pkg/util/util.go b/pkg/util/util.go
@@ -5,6 +5,7 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"regexp"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
@@ -72,3 +73,21 @@ func DownloadFile(filepath string, url string) error {
 	_, err = io.Copy(out, resp.Body)
 	return err
 }
+
+func ParseSelector(selector string) []KeyValueSelector {
+	r := regexp.MustCompile(`((?:([a-z0-9A-Z\/\-\._]+)(=|!=)?([a-z0-9A-Z\-\._]+)?)+)`)
+	matchPairs := r.FindAllStringSubmatch(selector, 10)
+	selectors := []KeyValueSelector{}
+
+	for _, matchPair := range matchPairs {
+		if len(matchPair) == 5 {
+			selectors = append(selectors, KeyValueSelector{
+				Key:      matchPair[2],
+				Operator: matchPair[3],
+				Value:    matchPair[4],
+			})
+		}
+	}
+
+	return selectors
+}