From b3f7ab657e03b3a7454b6d4481521c52e201b2c9 Mon Sep 17 00:00:00 2001
From: Michelle Shepardson <amerai@google.com>
Date: Fri, 8 Nov 2024 03:13:25 +0000
Subject: [PATCH] Allow label/target search for ResultStore sources.

Instead of requiring `label:"prow"` in queries, allow any number of
label or target atoms in a query. (This still limits the queries we
accept, but it should work for most simple uses of ResultStore sources).

Test: Unit tests, and verified with a local instance of TestGrid that a
ResultStore source config will pick up results from ResultStore
correctly.
---
 pkg/updater/resultstore/query/query.go      |  71 +++++---
 pkg/updater/resultstore/query/query_test.go | 176 ++++++++++++--------
 pkg/updater/resultstore/resultstore.go      |  29 +---
 pkg/updater/resultstore/resultstore_test.go |  82 ++-------
 4 files changed, 185 insertions(+), 173 deletions(-)

diff --git a/pkg/updater/resultstore/query/query.go b/pkg/updater/resultstore/query/query.go
index 3bb941bb..db524177 100644
--- a/pkg/updater/resultstore/query/query.go
+++ b/pkg/updater/resultstore/query/query.go
@@ -17,47 +17,76 @@ limitations under the License.
 package query
 
 import (
+	"errors"
 	"fmt"
 	"regexp"
 	"strings"
 )
 
-func translateAtom(simpleAtom string) (string, error) {
-	if simpleAtom == "" {
-		return "", nil
+type keyValue struct {
+	key   string
+	value string
+}
+
+func translateAtom(simpleAtom keyValue, queryTarget bool) (string, error) {
+	if simpleAtom.key == "" {
+		return "", errors.New("missing key")
 	}
-	// For now, we expect an atom with the exact form `target:"<target>"`
-	// Split the `key:value` atom.
-	parts := strings.SplitN(simpleAtom, ":", 2)
-	if len(parts) != 2 {
-		return "", fmt.Errorf("unrecognized atom %q", simpleAtom)
+	if simpleAtom.value == "" {
+		return "", errors.New("missing value")
 	}
-	key := strings.TrimSpace(parts[0])
-	val := strings.Trim(strings.TrimSpace(parts[1]), `"`)
-
 	switch {
-	case key == "target":
-		return fmt.Sprintf(`id.target_id="%s"`, val), nil
+	case simpleAtom.key == "label" && queryTarget:
+		return fmt.Sprintf(`invocation.invocation_attributes.labels:"%s"`, simpleAtom.value), nil
+	case simpleAtom.key == "label":
+		return fmt.Sprintf(`invocation_attributes.labels:"%s"`, simpleAtom.value), nil
+	case simpleAtom.key == "target":
+		return fmt.Sprintf(`id.target_id="%s"`, simpleAtom.value), nil
 	default:
-		return "", fmt.Errorf("unrecognized atom key %q", key)
+		return "", fmt.Errorf("unknown type of atom %q", simpleAtom.key)
 	}
 }
 
 var (
-	queryRe = regexp.MustCompile(`^target:".*"$`)
+	// Captures any atoms of the form `label:"<label>"` or `target:"<target>"`.
+	atomReStr = `(?P<atom>(?P<key>label|target):"(?P<value>.+?)")`
+	atomRe    = regexp.MustCompile(atomReStr)
+	// A query can only have atoms (above), separated by spaces.
+	queryRe = regexp.MustCompile(`^(` + atomReStr + ` *)+$`)
 )
 
+// TranslateQuery translates a simple query (similar to the syntax of searching invocations in the
+// UI) to a query for searching via API.
+// More at https://github.com/googleapis/googleapis/blob/master/google/devtools/resultstore/v2/resultstore_download.proto.
+//
+// This expects a query consisting of any number of space-separated `label:"<label>"` or
+// `target:"<target>"` atoms.
 func TranslateQuery(simpleQuery string) (string, error) {
 	if simpleQuery == "" {
 		return "", nil
 	}
-	// For now, we expect a query with a single atom, with the exact form `target:"<target>"`
 	if !queryRe.MatchString(simpleQuery) {
-		return "", fmt.Errorf("invalid query %q: must match %q", simpleQuery, queryRe.String())
+		return "", fmt.Errorf("query must consist of only space-separated `label:\"<label>\"` or `target:\"<target>\"` atoms")
+	}
+	var simpleAtoms []keyValue
+	var queryTarget bool
+	matches := atomRe.FindAllStringSubmatch(simpleQuery, -1)
+	for _, match := range matches {
+		if len(match) != 4 {
+			return "", fmt.Errorf("atom %v: want 4 submatches (full match, atom, key, value), but got %d", match, len(match))
+		}
+		simpleAtoms = append(simpleAtoms, keyValue{match[2], match[3]})
+		if match[2] == "target" {
+			queryTarget = true
+		}
 	}
-	query, err := translateAtom(simpleQuery)
-	if err != nil {
-		return "", fmt.Errorf("invalid query %q: %v", simpleQuery, err)
+	var atoms []string
+	for _, simpleAtom := range simpleAtoms {
+		atom, err := translateAtom(simpleAtom, queryTarget)
+		if err != nil {
+			return "", fmt.Errorf("atom %v: %v", simpleAtom, err)
+		}
+		atoms = append(atoms, atom)
 	}
-	return query, nil
+	return strings.Join(atoms, " "), nil
 }
diff --git a/pkg/updater/resultstore/query/query_test.go b/pkg/updater/resultstore/query/query_test.go
index f47ee1ae..82cd0cbb 100644
--- a/pkg/updater/resultstore/query/query_test.go
+++ b/pkg/updater/resultstore/query/query_test.go
@@ -22,63 +22,85 @@ import (
 
 func TestTranslateAtom(t *testing.T) {
 	cases := []struct {
-		name      string
-		atom      string
-		want      string
-		wantError bool
+		name       string
+		atom       keyValue
+		want       string
+		wantTarget string
 	}{
 		{
-			name: "empty",
-			atom: "",
-			want: "",
-		},
-		{
-			name: "basic",
-			atom: `target:"//my-target"`,
-			want: `id.target_id="//my-target"`,
+			name:       "label atom",
+			atom:       keyValue{"label", "foo"},
+			want:       `invocation_attributes.labels:"foo"`,
+			wantTarget: `invocation.invocation_attributes.labels:"foo"`,
 		},
 		{
-			name:      "case-sensitive key",
-			atom:      `TARGET:"//MY-TARGET"`,
-			wantError: true,
+			name:       "target atom",
+			atom:       keyValue{"target", "//my-target"},
+			want:       `id.target_id="//my-target"`,
+			wantTarget: `id.target_id="//my-target"`,
 		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Invocation queries (queryTarget = false)
+			got, err := translateAtom(tc.atom, false)
+			if err != nil {
+				t.Fatalf("translateAtom(%q, %t) errored: %v", tc.atom, false, err)
+			}
+			if tc.want != got {
+				t.Errorf("translateAtom(%q, %t) differed; got %q, want %q", tc.atom, false, got, tc.want)
+			}
+			// Configured Target queries (queryTarget = true)
+			gotTarget, err := translateAtom(tc.atom, true)
+			if err != nil {
+				t.Fatalf("translateAtom(%q, %t) errored: %v", tc.atom, true, err)
+			}
+			if tc.wantTarget != gotTarget {
+				t.Errorf("translateAtom(%q, %t) differed; got %q, want %q", tc.atom, true, gotTarget, tc.wantTarget)
+			}
+		})
+	}
+}
+
+func TestTranslateAtom_Error(t *testing.T) {
+	cases := []struct {
+		name string
+		atom keyValue
+	}{
 		{
-			name: "multiple colons",
-			atom: `target:"//path/to:my-target"`,
-			want: `id.target_id="//path/to:my-target"`,
+			name: "empty",
+			atom: keyValue{"", ""},
 		},
 		{
-			name: "unquoted",
-			atom: `target://my-target`,
-			want: `id.target_id="//my-target"`,
+			name: "case-sensitive key",
+			atom: keyValue{"TARGET", "//MY-TARGET"},
 		},
 		{
-			name: "partial quotes",
-			atom: `target://my-target"`,
-			want: `id.target_id="//my-target"`,
+			name: "missing key",
+			atom: keyValue{"", "//path/to:my-target"},
 		},
 		{
-			name:      "not enough parts",
-			atom:      "target",
-			wantError: true,
+			name: "missing value",
+			atom: keyValue{"target", ""},
 		},
 		{
-			name:      "unknown atom",
-			atom:      "label:foo",
-			wantError: true,
+			name: "unknown atom",
+			atom: keyValue{"custom-property", "foo"},
 		},
 	}
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			got, err := translateAtom(tc.atom)
-			if tc.want != got {
-				t.Errorf("translateAtom(%q) differed; got %q, want %q", tc.atom, got, tc.want)
+			// Invocation queries (queryTarget = false)
+			_, err := translateAtom(tc.atom, false)
+			if err == nil {
+				t.Fatalf("translateAtom(%q, %t): want error but got none", tc.atom, false)
 			}
-			if err == nil && tc.wantError {
-				t.Errorf("translateAtom(%q) did not error as expected", tc.atom)
-			} else if err != nil && !tc.wantError {
-				t.Errorf("translateAtom(%q) errored unexpectedly: %v", tc.atom, err)
+			// Configured Target queries (queryTarget = true)
+			_, err = translateAtom(tc.atom, true)
+			if err == nil {
+				t.Fatalf("translateAtom(%q, %t): want error but got none", tc.atom, true)
 			}
 		})
 	}
@@ -86,10 +108,9 @@ func TestTranslateAtom(t *testing.T) {
 
 func TestTranslateQuery(t *testing.T) {
 	cases := []struct {
-		name      string
-		query     string
-		want      string
-		wantError bool
+		name  string
+		query string
+		want  string
 	}{
 		{
 			name:  "empty",
@@ -97,52 +118,77 @@ func TestTranslateQuery(t *testing.T) {
 			want:  "",
 		},
 		{
-			name:  "basic",
+			name:  "label",
+			query: `label:"foo"`,
+			want:  `invocation_attributes.labels:"foo"`,
+		},
+		{
+			name:  "target",
 			query: `target:"//my-target"`,
 			want:  `id.target_id="//my-target"`,
 		},
 		{
-			name:      "case-sensitive key",
-			query:     `TARGET:"//MY-TARGET"`,
-			wantError: true,
+			name:  "multiple labels",
+			query: `label:"foo" label:"bar"`,
+			want:  `invocation_attributes.labels:"foo" invocation_attributes.labels:"bar"`,
+		},
+		{
+			name:  "mixed",
+			query: `label:"foo" target:"//my-target" label:"bar"`,
+			want:  `invocation.invocation_attributes.labels:"foo" id.target_id="//my-target" invocation.invocation_attributes.labels:"bar"`,
 		},
 		{
 			name:  "multiple colons",
 			query: `target:"//path/to:my-target"`,
 			want:  `id.target_id="//path/to:my-target"`,
 		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			got, err := TranslateQuery(tc.query)
+			if err != nil {
+				t.Fatalf("translateQuery(%q) errored: %v", tc.query, err)
+			}
+			if tc.want != got {
+				t.Errorf("translateQuery(%q) differed; got %q, want %q", tc.query, got, tc.want)
+			}
+		})
+	}
+}
+
+func TestTranslateQuery_Error(t *testing.T) {
+	cases := []struct {
+		name  string
+		query string
+	}{
+		{
+			name:  "case-sensitive key",
+			query: `TARGET:"//MY-TARGET"`,
+		},
 		{
-			name:      "unquoted",
-			query:     `target://my-target`,
-			wantError: true,
+			name:  "unquoted",
+			query: `target://my-target`,
 		},
 		{
-			name:      "partial quotes",
-			query:     `target://my-target"`,
-			wantError: true,
+			name:  "partial quotes",
+			query: `target://my-target"`,
 		},
 		{
-			name:      "invalid query",
-			query:     `label:foo`,
-			wantError: true,
+			name:  "invalid query",
+			query: `label:foo`,
 		},
 		{
-			name:      "partial match",
-			query:     `some_target:foo`,
-			wantError: true,
+			name:  "partial match",
+			query: `some_target:foo`,
 		},
 	}
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			got, err := TranslateQuery(tc.query)
-			if tc.want != got {
-				t.Errorf("translateQuery(%q) differed; got %q, want %q", tc.query, got, tc.want)
-			}
-			if tc.wantError && err == nil {
-				t.Errorf("translateQuery(%q) did not error as expected", tc.query)
-			} else if !tc.wantError && err != nil {
-				t.Errorf("translateQuery(%q) errored unexpectedly: %v", tc.query, err)
+			_, err := TranslateQuery(tc.query)
+			if err == nil {
+				t.Fatalf("translateQuery(%q): want error, got none", tc.query)
 			}
 		})
 	}
diff --git a/pkg/updater/resultstore/resultstore.go b/pkg/updater/resultstore/resultstore.go
index 8e2ed2e2..62a37674 100644
--- a/pkg/updater/resultstore/resultstore.go
+++ b/pkg/updater/resultstore/resultstore.go
@@ -887,36 +887,23 @@ func queryAfter(query string, when time.Time) string {
 	if query == "" {
 		return ""
 	}
-	return fmt.Sprintf("%s timing.start_time>=\"%s\"", query, when.UTC().Format(time.RFC3339))
-}
-
-const (
-	// Use this when searching invocations, e.g. if query does not search for a target.
-	prowLabel = `invocation_attributes.labels:"prow"`
-	// Use this when searching for a configured target, e.g. if query contains `target:"<target>"`.
-	prowTargetLabel = `invocation.invocation_attributes.labels:"prow"`
-)
-
-func queryProw(baseQuery string, stop time.Time) (string, error) {
-	// TODO: ResultStore use is assumed to be Prow-only at the moment. Make this more flexible in future.
-	if baseQuery == "" {
-		return queryAfter(prowLabel, stop), nil
-	}
-	query, err := query.TranslateQuery(baseQuery)
-	if err != nil {
-		return "", err
+	// Note: time arguments are different for invocation vs. configured target searches.
+	timeAtom := fmt.Sprintf(`timing.start_time>="%s"`, when.UTC().Format(time.RFC3339))
+	if strings.Contains(query, "id.target_id=") {
+		timeAtom = fmt.Sprintf(`invocation.timing.start_time>="%s"`, when.UTC().Format(time.RFC3339))
 	}
-	return queryAfter(fmt.Sprintf("%s %s", query, prowTargetLabel), stop), nil
+	return fmt.Sprintf(`%s %s`, query, timeAtom)
 }
 
 func search(ctx context.Context, log logrus.FieldLogger, client *DownloadClient, rsConfig *configpb.ResultStoreConfig, stop time.Time) ([]string, error) {
 	if client == nil {
 		return nil, fmt.Errorf("no ResultStore client provided")
 	}
-	query, err := queryProw(rsConfig.GetQuery(), stop)
+	translatedQuery, err := query.TranslateQuery(rsConfig.GetQuery())
 	if err != nil {
-		return nil, fmt.Errorf("queryProw() failed to create query: %v", err)
+		return nil, fmt.Errorf("query.TranslateQuery(%s): %v", rsConfig.GetQuery(), err)
 	}
+	query := queryAfter(translatedQuery, stop)
 	log.WithField("query", query).Debug("Searching ResultStore.")
 	// Quit if search goes over 5 minutes.
 	ctx, cancel := context.WithTimeout(ctx, 5*time.Minute)
diff --git a/pkg/updater/resultstore/resultstore_test.go b/pkg/updater/resultstore/resultstore_test.go
index 066ba56f..57f0bd87 100644
--- a/pkg/updater/resultstore/resultstore_test.go
+++ b/pkg/updater/resultstore/resultstore_test.go
@@ -203,30 +203,32 @@ func TestColumnReader(t *testing.T) {
 	oneMonthConfig := &configpb.TestGroup{
 		Name:          "a-test-group",
 		DaysOfResults: 30,
+		ResultSource: &configpb.TestGroup_ResultSource{
+			ResultSourceConfig: &configpb.TestGroup_ResultSource_ResultstoreConfig{
+				ResultstoreConfig: &configpb.ResultStoreConfig{
+					Query: `label:"foo"`,
+				},
+			},
+		},
 	}
 	now := time.Now()
 	oneDayAgo := now.AddDate(0, 0, -1)
 	twoDaysAgo := now.AddDate(0, 0, -2)
 	threeDaysAgo := now.AddDate(0, 0, -3)
 	oneMonthAgo := now.AddDate(0, 0, -30)
-	testQueryAfter := queryAfter(prowLabel, oneMonthAgo)
+	testQueryAfter := queryAfter(`invocation_attributes.labels:"foo"`, oneMonthAgo)
 	cases := []struct {
 		name    string
 		client  *fakeClient
-		tg      *configpb.TestGroup
 		want    []updater.InflatedColumn
 		wantErr bool
 	}{
 		{
 			name:    "empty",
-			tg:      oneMonthConfig,
 			wantErr: true,
 		},
 		{
 			name: "basic",
-			tg: &configpb.TestGroup{
-				DaysOfResults: 30,
-			},
 			client: &fakeClient{
 				searches: map[string][]string{
 					testQueryAfter: {"id-1", "id-2"},
@@ -351,7 +353,6 @@ func TestColumnReader(t *testing.T) {
 		},
 		{
 			name: "no results from query",
-			tg:   oneMonthConfig,
 			client: &fakeClient{
 				searches: map[string][]string{},
 				invocations: map[string]FetchResult{
@@ -2595,15 +2596,15 @@ func TestQueryAfter(t *testing.T) {
 		},
 		{
 			name:  "zero",
-			query: prowLabel,
+			query: `invocation_attributes.labels:"foo"`,
 			when:  time.Time{},
-			want:  "invocation_attributes.labels:\"prow\" timing.start_time>=\"0001-01-01T00:00:00Z\"",
+			want:  `invocation_attributes.labels:"foo" timing.start_time>="0001-01-01T00:00:00Z"`,
 		},
 		{
 			name:  "basic",
-			query: prowLabel,
+			query: `invocation_attributes.labels:"foo"`,
 			when:  now,
-			want:  fmt.Sprintf("invocation_attributes.labels:\"prow\" timing.start_time>=\"%s\"", now.UTC().Format(time.RFC3339)),
+			want:  fmt.Sprintf(`invocation_attributes.labels:"foo" timing.start_time>="%s"`, now.UTC().Format(time.RFC3339)),
 		},
 	}
 	for _, tc := range cases {
@@ -2616,63 +2617,10 @@ func TestQueryAfter(t *testing.T) {
 	}
 }
 
-func TestQueryProw(t *testing.T) {
-	now := time.Now()
-	cases := []struct {
-		name      string
-		query     string
-		when      time.Time
-		want      string
-		wantError bool
-	}{
-		{
-			name: "empty",
-			want: "invocation_attributes.labels:\"prow\" timing.start_time>=\"0001-01-01T00:00:00Z\"",
-		},
-		{
-			name:  "zero",
-			query: "",
-			when:  time.Time{},
-			want:  "invocation_attributes.labels:\"prow\" timing.start_time>=\"0001-01-01T00:00:00Z\"",
-		},
-		{
-			name:  "basic",
-			query: "",
-			when:  now,
-			want:  fmt.Sprintf("invocation_attributes.labels:\"prow\" timing.start_time>=\"%s\"", now.UTC().Format(time.RFC3339)),
-		},
-		{
-			name:  "target",
-			query: `target:"my-job"`,
-			when:  now,
-			want:  fmt.Sprintf("id.target_id=\"my-job\" invocation.invocation_attributes.labels:\"prow\" timing.start_time>=\"%s\"", now.UTC().Format(time.RFC3339)),
-		},
-		{
-			name:      "invalid query",
-			query:     `label:foo bar`,
-			when:      now,
-			wantError: true,
-		},
-	}
-	for _, tc := range cases {
-		t.Run(tc.name, func(t *testing.T) {
-			got, err := queryProw(tc.query, tc.when)
-			if !tc.wantError && err != nil {
-				t.Errorf("queryProw(%q, %v) errored: %v", tc.query, tc.when, err)
-			} else if tc.wantError && err == nil {
-				t.Errorf("queryProw(%q, %v) did not error as expected", tc.query, tc.when)
-			}
-			if diff := cmp.Diff(tc.want, got); diff != "" {
-				t.Errorf("queryProw(%q, %v) differed (-want, +got): %s", tc.query, tc.when, diff)
-			}
-		})
-	}
-}
-
 func TestSearch(t *testing.T) {
 	twoDaysAgo := time.Now().AddDate(0, 0, -2)
-	testQueryAfter := queryAfter(prowLabel, twoDaysAgo)
-	testTargetQueryAfter, _ := queryProw(`target:"my-job"`, twoDaysAgo)
+	testQueryAfter := queryAfter(`invocation_attributes.labels:"foo"`, twoDaysAgo)
+	testTargetQueryAfter := queryAfter(`id.target_id="my-job"`, twoDaysAgo)
 	cases := []struct {
 		name     string
 		stop     time.Time
@@ -2695,6 +2643,7 @@ func TestSearch(t *testing.T) {
 			name: "no results",
 			rsConfig: &configpb.ResultStoreConfig{
 				Project: "my-project",
+				Query:   `label:foo`,
 			},
 			client:  &fakeClient{},
 			wantErr: true,
@@ -2703,6 +2652,7 @@ func TestSearch(t *testing.T) {
 			name: "basic",
 			rsConfig: &configpb.ResultStoreConfig{
 				Project: "my-project",
+				Query:   `label:"foo"`,
 			},
 			client: &fakeClient{
 				searches: map[string][]string{