Highlight matched strings in the log lines with the include option (#231)

superbrothers · web-flow · commit 9fbaa18064da · 2023-02-19T08:46:38.000+09:00
diff --git a/stern/tail.go b/stern/tail.go
@@ -23,6 +23,7 @@ import (
 	"hash/fnv"
 	"io"
 	"regexp"
+	"sort"
 	"strings"
 	"text/template"
 	"time"
@@ -68,6 +69,9 @@ type TailOptions struct {
 	TailLines    *int64
 	Follow       bool
 	OnlyLogLines bool
+
+	// regexp for highlighting the matched string
+	reHightlight *regexp.Regexp
 }
 
 type ResumeRequest struct {
@@ -99,6 +103,34 @@ func (o TailOptions) IsInclude(msg string) bool {
 	return false
 }
 
+var colorHighlight = color.New(color.FgRed, color.Bold).SprintFunc()
+
+func (o TailOptions) HighlightMatchedString(msg string) string {
+	if len(o.Include) == 0 {
+		return msg
+	}
+
+	if o.reHightlight == nil {
+		ss := make([]string, len(o.Include))
+		for i, rin := range o.Include {
+			ss[i] = rin.String()
+		}
+
+		// We expect a longer match
+		sort.Slice(ss, func(i, j int) bool {
+			return len(ss[i]) > len(ss[j])
+		})
+
+		o.reHightlight = regexp.MustCompile("(" + strings.Join(ss, "|") + ")")
+	}
+
+	msg = o.reHightlight.ReplaceAllStringFunc(msg, func(part string) string {
+		return colorHighlight(part)
+	})
+
+	return msg
+}
+
 func (o TailOptions) UpdateTimezone(timestamp string) (string, error) {
 	t, err := time.ParseInLocation(time.RFC3339Nano, timestamp, time.UTC)
 	if err != nil {
@@ -292,7 +324,8 @@ func (t *Tail) consumeLine(line string) {
 		return
 	}
 
-	msg := content
+	msg := t.Options.HighlightMatchedString(content)
+
 	if t.Options.Timestamps {
 		updatedTs, err := t.Options.UpdateTimezone(rfc3339Nano)
 		if err != nil {
@@ -301,6 +334,7 @@ func (t *Tail) consumeLine(line string) {
 		}
 		msg = updatedTs + " " + msg
 	}
+
 	t.Print(msg)
 }
 
diff --git a/stern/tail_test.go b/stern/tail_test.go
@@ -11,6 +11,7 @@ import (
 	"text/template"
 	"time"
 
+	"github.com/fatih/color"
 	"k8s.io/client-go/kubernetes/fake"
 )
 
@@ -365,3 +366,64 @@ func TestRemoveSubsecond(t *testing.T) {
 		}
 	}
 }
+
+func TestHighlightMatchedString(t *testing.T) {
+	tests := []struct {
+		msg      string
+		include  []*regexp.Regexp
+		expected string
+	}{
+		{
+			"test matched",
+			[]*regexp.Regexp{
+				regexp.MustCompile(`test`),
+			},
+			"\x1b[31;1mtest\x1b[0m matched",
+		},
+		{
+			"test not-matched",
+			[]*regexp.Regexp{
+				regexp.MustCompile(`hoge`),
+			},
+			"test not-matched",
+		},
+		{
+			"test matched",
+			[]*regexp.Regexp{
+				regexp.MustCompile(`not-matched`),
+				regexp.MustCompile(`matched`),
+			},
+			"test \x1b[31;1mmatched\x1b[0m",
+		},
+		{
+			"test multiple matched",
+			[]*regexp.Regexp{
+				regexp.MustCompile(`multiple`),
+				regexp.MustCompile(`matched`),
+			},
+			"test \x1b[31;1mmultiple\x1b[0m \x1b[31;1mmatched\x1b[0m",
+		},
+		{
+			"test match on the longer one",
+			[]*regexp.Regexp{
+				regexp.MustCompile(`match`),
+				regexp.MustCompile(`match on the longer one`),
+			},
+			"test \x1b[31;1mmatch on the longer one\x1b[0m",
+		},
+	}
+
+	orig := color.NoColor
+	color.NoColor = false
+	defer func() {
+		color.NoColor = orig
+	}()
+
+	for i, tt := range tests {
+		o := &TailOptions{Include: tt.include}
+		actual := o.HighlightMatchedString(tt.msg)
+		if actual != tt.expected {
+			t.Errorf("%d: expected %q, but actual %q", i, tt.expected, actual)
+		}
+	}
+}
