WIP: Allow label/target search for ResultStore sources.

pkg/updater/resultstore/query/query.go

@@ -17,47 +17,76 @@
 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
 }

pkg/updater/resultstore/query/query_test.go

@@ -22,127 +22,173 @@ 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)
 			}
 		})
 	}
 }
 
 func TestTranslateQuery(t *testing.T) {
 	cases := []struct {
-		name      string
-		query     string
-		want      string
-		wantError bool
+		name  string
+		query string
+		want  string
 	}{
 		{
 			name:  "empty",
 			query: "",
 			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)
 			}
 		})
 	}

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)