From 7a852ec5d6bd572573d4e1996e98d67140e97b9e Mon Sep 17 00:00:00 2001
From: Extreme <1403539044@qq.com>
Date: Fri, 19 Jan 2024 14:16:50 +0800
Subject: [PATCH] feat: regex filter pod name

Signed-off-by: Extreme <1403539044@qq.com>
---
 cmd/kail/main.go |  5 +++++
 ds_builder.go    | 22 ++++++++++++++++++++--
 filter.go        | 26 ++++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/cmd/kail/main.go b/cmd/kail/main.go
index cfc94fb..0c30fe1 100644
--- a/cmd/kail/main.go
+++ b/cmd/kail/main.go
@@ -47,6 +47,7 @@ var (
 	flagJob         = kingpin.Flag("job", "job").Short('j').PlaceHolder("NAME").Strings()
 	flagNode        = kingpin.Flag("node", "node").PlaceHolder("NAME").Strings()
 	flagIng         = kingpin.Flag("ing", "ingress").PlaceHolder("NAME").Strings()
+	flagRegex       = kingpin.Flag("regex", "regex to filter pod name").PlaceHolder("REGEX").String()
 
 	flagContext = kingpin.Flag("context", "kubernetes context").PlaceHolder("CONTEXT-NAME").String()
 
@@ -288,6 +289,10 @@ func createDSBuilder() kail.DSBuilder {
 		dsb = dsb.WithIngress(ids...)
 	}
 
+	if flagRegex != nil && *flagRegex != "" {
+		dsb = dsb.WithRegex(*flagRegex)
+	}
+
 	return dsb
 }
 
diff --git a/ds_builder.go b/ds_builder.go
index 49f8b39..1f2ad14 100644
--- a/ds_builder.go
+++ b/ds_builder.go
@@ -3,8 +3,6 @@ package kail
 import (
 	"context"
 
-	"github.com/boz/kcache/types/job"
-
 	logutil "github.com/boz/go-logutil"
 	"github.com/boz/kcache/filter"
 	"github.com/boz/kcache/join"
@@ -12,6 +10,7 @@ import (
 	"github.com/boz/kcache/types/daemonset"
 	"github.com/boz/kcache/types/deployment"
 	"github.com/boz/kcache/types/ingress"
+	"github.com/boz/kcache/types/job"
 	"github.com/boz/kcache/types/pod"
 	"github.com/boz/kcache/types/replicaset"
 	"github.com/boz/kcache/types/replicationcontroller"
@@ -37,6 +36,7 @@ type DSBuilder interface {
 	WithStatefulSet(id ...nsname.NSName) DSBuilder
 	WithJob(id ...nsname.NSName) DSBuilder
 	WithIngress(id ...nsname.NSName) DSBuilder
+	WithRegex(s string) DSBuilder
 
 	Create(ctx context.Context, cs kubernetes.Interface) (DS, error)
 }
@@ -60,6 +60,7 @@ type dsBuilder struct {
 	statefulsets []nsname.NSName
 	jobs         []nsname.NSName
 	ingresses    []nsname.NSName
+	regex        string
 }
 
 func (b *dsBuilder) WithIgnore(selector ...labels.Selector) DSBuilder {
@@ -132,6 +133,11 @@ func (b *dsBuilder) WithIngress(id ...nsname.NSName) DSBuilder {
 	return b
 }
 
+func (b *dsBuilder) WithRegex(s string) DSBuilder {
+	b.regex = s
+	return b
+}
+
 func (b *dsBuilder) Create(ctx context.Context, cs kubernetes.Interface) (DS, error) {
 	log := logutil.FromContextOrDefault(ctx)
 
@@ -178,6 +184,18 @@ func (b *dsBuilder) Create(ctx context.Context, cs kubernetes.Interface) (DS, er
 			return nil, log.Err(err, "labels filter")
 		}
 	}
+	if len(b.regex) != 0 {
+		regexFilter, err := NameRegexFileter(b.regex)
+		if err != nil {
+			ds.closeAll()
+			return nil, log.Err(err, "regex filter")
+		}
+		ds.pods, err = ds.pods.CloneWithFilter(regexFilter)
+		if err != nil {
+			ds.closeAll()
+			return nil, log.Err(err, "regex filter")
+		}
+	}
 
 	if len(b.selectors) != 0 {
 		filters := make([]filter.Filter, 0, len(b.selectors))
diff --git a/filter.go b/filter.go
index 148c27b..967b57e 100644
--- a/filter.go
+++ b/filter.go
@@ -1,10 +1,13 @@
 package kail
 
 import (
+	"regexp"
 	"sort"
 
+	"github.com/boz/kcache/filter"
 	"github.com/boz/kcache/nsname"
 	"k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 type ContainerFilter interface {
@@ -71,3 +74,26 @@ func SourcesForPod(
 
 	return id, sources
 }
+
+func NameRegexFileter(regex string) (filter.Filter, error) {
+	compile, err := regexp.Compile(regex)
+	if err != nil {
+		return nil, err
+	}
+	return nameRegexFilter{compile}, nil
+}
+
+type nameRegexFilter struct {
+	regex *regexp.Regexp
+}
+
+func (f nameRegexFilter) Accept(obj metav1.Object) bool {
+	return f.regex.MatchString(obj.GetName())
+}
+
+func (f nameRegexFilter) Equals(other filter.Filter) bool {
+	if other, ok := other.(nameRegexFilter); ok {
+		return f.regex.String() == other.regex.String()
+	}
+	return false
+}