Merge pull request #1274 from michelle192837/query

Add ability to query a single job from Prow in ResultStore.

pb/config/config.proto

@@ -364,6 +364,9 @@ message GCSConfig {
 message ResultStoreConfig {
   // Google Cloud Platform project ID where ResultStore results are stored.
   string project = 1;
+  // A simple query to filter for particular results.
+  // Currently, only allows a query in the form of `target:"<target>"`.
+  string query = 2; 
 }
 
 // Options for where to gather linked issues from.

pkg/updater/resultstore/BUILD.bazel

@@ -4,6 +4,7 @@ go_library(
     name = "go_default_library",
     srcs = [
         "client.go",
+        "query.go",
         "resultstore.go",
     ],
     importpath = "github.com/GoogleCloudPlatform/testgrid/pkg/updater/resultstore",
@@ -27,7 +28,10 @@ go_library(
 
 go_test(
     name = "go_default_test",
-    srcs = ["resultstore_test.go"],
+    srcs = [
+        "query_test.go",
+        "resultstore_test.go",
+    ],
     embed = [":go_default_library"],
     deps = [
         "//pb/config:go_default_library",

pkg/updater/resultstore/query.go

@@ -0,0 +1,63 @@
+/*
+Copyright 2024 The TestGrid Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package resultstore
+
+import (
+	"fmt"
+	"regexp"
+	"strings"
+)
+
+func translateAtom(simpleAtom string) (string, error) {
+	if simpleAtom == "" {
+		return "", nil
+	}
+	// 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)
+	}
+	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
+	default:
+		return "", fmt.Errorf("unrecognized atom key %q", key)
+	}
+}
+
+var (
+	queryRe = regexp.MustCompile(`^target:".*"$`)
+)
+
+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())
+	}
+	query, err := translateAtom(simpleQuery)
+	if err != nil {
+		return "", fmt.Errorf("invalid query %q: %v", simpleQuery, err)
+	}
+	return query, nil
+}

pkg/updater/resultstore/query_test.go

@@ -0,0 +1,149 @@
+/*
+Copyright 2024 The TestGrid Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package resultstore
+
+import (
+	"testing"
+)
+
+func TestTranslateAtom(t *testing.T) {
+	cases := []struct {
+		name      string
+		atom      string
+		want      string
+		wantError bool
+	}{
+		{
+			name: "empty",
+			atom: "",
+			want: "",
+		},
+		{
+			name: "basic",
+			atom: `target:"//my-target"`,
+			want: `id.target_id="//my-target"`,
+		},
+		{
+			name:      "case-sensitive key",
+			atom:      `TARGET:"//MY-TARGET"`,
+			wantError: true,
+		},
+		{
+			name: "multiple colons",
+			atom: `target:"//path/to:my-target"`,
+			want: `id.target_id="//path/to:my-target"`,
+		},
+		{
+			name: "unquoted",
+			atom: `target://my-target`,
+			want: `id.target_id="//my-target"`,
+		},
+		{
+			name: "partial quotes",
+			atom: `target://my-target"`,
+			want: `id.target_id="//my-target"`,
+		},
+		{
+			name:      "not enough parts",
+			atom:      "target",
+			wantError: true,
+		},
+		{
+			name:      "unknown atom",
+			atom:      "label:foo",
+			wantError: true,
+		},
+	}
+
+	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)
+			}
+			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)
+			}
+		})
+	}
+}
+
+func TestTranslateQuery(t *testing.T) {
+	cases := []struct {
+		name      string
+		query     string
+		want      string
+		wantError bool
+	}{
+		{
+			name:  "empty",
+			query: "",
+			want:  "",
+		},
+		{
+			name:  "basic",
+			query: `target:"//my-target"`,
+			want:  `id.target_id="//my-target"`,
+		},
+		{
+			name:      "case-sensitive key",
+			query:     `TARGET:"//MY-TARGET"`,
+			wantError: true,
+		},
+		{
+			name:  "multiple colons",
+			query: `target:"//path/to:my-target"`,
+			want:  `id.target_id="//path/to:my-target"`,
+		},
+		{
+			name:      "unquoted",
+			query:     `target://my-target`,
+			wantError: true,
+		},
+		{
+			name:      "partial quotes",
+			query:     `target://my-target"`,
+			wantError: true,
+		},
+		{
+			name:      "invalid query",
+			query:     `label:foo`,
+			wantError: true,
+		},
+		{
+			name:      "partial match",
+			query:     `some_target:foo`,
+			wantError: true,
+		},
+	}
+
+	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)
+			}
+		})
+	}
+}

pkg/updater/resultstore/resultstore.go

@@ -156,7 +156,7 @@ func extractGroupID(tg *configpb.TestGroup, inv *invocation) string {
 func ColumnReader(client *DownloadClient, reprocess time.Duration) updater.ColumnReader {
 	return func(ctx context.Context, log logrus.FieldLogger, tg *configpb.TestGroup, oldCols []updater.InflatedColumn, defaultStop time.Time, receivers chan<- updater.InflatedColumn) error {
 		stop := updateStop(log, tg, time.Now(), oldCols, defaultStop, reprocess)
-		ids, err := search(ctx, log, client, tg.GetResultSource().GetResultstoreConfig().GetProject(), stop)
+		ids, err := search(ctx, log, client, tg.GetResultSource().GetResultstoreConfig(), stop)
 		if err != nil {
 			return fmt.Errorf("error searching invocations: %v", err)
 		}
@@ -889,21 +889,35 @@ func queryAfter(query string, when time.Time) string {
 	return fmt.Sprintf("%s timing.start_time>=\"%s\"", query, when.UTC().Format(time.RFC3339))
 }
 
-// TODO: Replace these hardcoded values with adjustable ones.
 const (
-	queryProw = "invocation_attributes.labels:\"prow\""
+	prowLabel = `invocation_attributes.labels:"prow"`
 )
 
-func search(ctx context.Context, log logrus.FieldLogger, client *DownloadClient, projectID string, stop time.Time) ([]string, error) {
+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 := translateQuery(baseQuery)
+	if err != nil {
+		return "", err
+	}
+	return queryAfter(fmt.Sprintf("%s %s", query, prowLabel), stop), nil
+}
+
+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 := queryAfter(queryProw, stop)
+	query, err := queryProw(rsConfig.GetQuery(), stop)
+	if err != nil {
+		return nil, fmt.Errorf("queryProw() failed to create query: %v", err)
+	}
 	log.WithField("query", query).Debug("Searching ResultStore.")
 	// Quit if search goes over 5 minutes.
 	ctx, cancel := context.WithTimeout(ctx, 5*time.Minute)
 	defer cancel()
-	ids, err := client.Search(ctx, log, query, projectID)
+	ids, err := client.Search(ctx, log, query, rsConfig.GetProject())
 	log.WithField("ids", len(ids)).WithError(err).Debug("Searched ResultStore.")
 	return ids, err
 }