Use SearchConfiguredTargets when searching targets.

If searching for a target (as opposed to labels or other identifiers), we want to search for those results using a difference method (SearchConfiguredTargets, rather than SearchInvocations). Add the search method and use it if the query contains a target atom (and update the Prow label accordingly for this search). Moved the query code into a separate package so it can be used in config validation without a circular dependency.
GoogleCloudPlatform · Mar 14, 2024 · a6609c4 · a6609c4
1 parent 837ecf5
commit a6609c4
Show file tree

Hide file tree

Showing 10 changed files with 205 additions and 44 deletions.
diff --git a/config/BUILD.bazel b/config/BUILD.bazel
@@ -13,6 +13,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pb/config:go_default_library",
+        "//pkg/updater/resultstore/query:go_default_library",
         "//util/gcs:go_default_library",
         "//util/queue:go_default_library",
         "@com_github_golang_protobuf//proto:go_default_library",

diff --git a/config/config.go b/config/config.go
@@ -28,6 +28,7 @@ import (
 	"github.com/golang/protobuf/proto"
 
 	configpb "github.com/GoogleCloudPlatform/testgrid/pb/config"
+	"github.com/GoogleCloudPlatform/testgrid/pkg/updater/resultstore/query"
 	multierror "github.com/hashicorp/go-multierror"
 )
 
@@ -269,6 +270,9 @@ func validateResultStoreSource(tg *configpb.TestGroup) error {
 		if rs.GetProject() == "" {
 			return errors.New("project ID in resultstore_config cannot be empty")
 		}
+		if _, err := query.TranslateQuery(rs.GetQuery()); err != nil {
+			return fmt.Errorf("invalid ResultStore query %q: %v", rs.GetQuery(), err)
+		}
 	}
 	return nil
 }

diff --git a/config/config_test.go b/config/config_test.go
@@ -498,6 +498,47 @@ func TestValidateResultStoreSource(t *testing.T) {
 			},
 			err: false,
 		},
+		{
+			name: "valid query",
+			tg: &configpb.TestGroup{
+				ResultSource: &configpb.TestGroup_ResultSource{
+					ResultSourceConfig: &configpb.TestGroup_ResultSource_ResultstoreConfig{
+						ResultstoreConfig: &configpb.ResultStoreConfig{
+							Project: "my-project",
+							Query:   `target:"my-job"`,
+						},
+					},
+				},
+			},
+			err: false,
+		},
+		{
+			name: "invalid query",
+			tg: &configpb.TestGroup{
+				ResultSource: &configpb.TestGroup_ResultSource{
+					ResultSourceConfig: &configpb.TestGroup_ResultSource_ResultstoreConfig{
+						ResultstoreConfig: &configpb.ResultStoreConfig{
+							Project: "my-project",
+							Query:   `label:foo bar`,
+						},
+					},
+				},
+			},
+			err: true,
+		},
+		{
+			name: "query without project",
+			tg: &configpb.TestGroup{
+				ResultSource: &configpb.TestGroup_ResultSource{
+					ResultSourceConfig: &configpb.TestGroup_ResultSource_ResultstoreConfig{
+						ResultstoreConfig: &configpb.ResultStoreConfig{
+							Query: `target:"my-job"`,
+						},
+					},
+				},
+			},
+			err: true,
+		},
 		{
 			name: "gcs_prefix and ResultStore defined",
 			tg: &configpb.TestGroup{

diff --git a/pkg/updater/resultstore/BUILD.bazel b/pkg/updater/resultstore/BUILD.bazel
@@ -4,7 +4,6 @@ go_library(
     name = "go_default_library",
     srcs = [
         "client.go",
-        "query.go",
         "resultstore.go",
     ],
     importpath = "github.com/GoogleCloudPlatform/testgrid/pkg/updater/resultstore",
@@ -15,6 +14,7 @@ go_library(
         "//pb/state:go_default_library",
         "//pb/test_status:go_default_library",
         "//pkg/updater:go_default_library",
+        "//pkg/updater/resultstore/query:go_default_library",
         "//util/gcs:go_default_library",
         "@com_github_sirupsen_logrus//:go_default_library",
         "@go_googleapis//google/devtools/resultstore/v2:resultstore_go_proto",
@@ -28,10 +28,7 @@ go_library(
 
 go_test(
     name = "go_default_test",
-    srcs = [
-        "query_test.go",
-        "resultstore_test.go",
-    ],
+    srcs = ["resultstore_test.go"],
     embed = [":go_default_library"],
     deps = [
         "//pb/config:go_default_library",
@@ -59,7 +56,10 @@ filegroup(
 
 filegroup(
     name = "all-srcs",
-    srcs = [":package-srcs"],
+    srcs = [
+        ":package-srcs",
+        "//pkg/updater/resultstore/query:all-srcs",
+    ],
     tags = ["automanaged"],
     visibility = ["//visibility:public"],
 )
diff --git a/pkg/updater/resultstore/client.go b/pkg/updater/resultstore/client.go
@@ -33,6 +33,7 @@ import (
 
 type resultStoreClient interface {
 	SearchInvocations(context.Context, *resultstore.SearchInvocationsRequest, ...grpc.CallOption) (*resultstore.SearchInvocationsResponse, error)
+	SearchConfiguredTargets(context.Context, *resultstore.SearchConfiguredTargetsRequest, ...grpc.CallOption) (*resultstore.SearchConfiguredTargetsResponse, error)
 	ExportInvocation(context.Context, *resultstore.ExportInvocationRequest, ...grpc.CallOption) (*resultstore.ExportInvocationResponse, error)
 }
 
@@ -81,35 +82,76 @@ func Connect(ctx context.Context, serviceAccountPath string) (*grpc.ClientConn,
 }
 
 // Search finds all the invocations that satisfies the query condition within a project.
-func (c *DownloadClient) Search(ctx context.Context, log logrus.FieldLogger, query, projectID string, fields ...string) ([]string, error) {
-	var ids []string
+func (c *DownloadClient) Search(ctx context.Context, log logrus.FieldLogger, query, projectID string) ([]string, error) {
+	var invIDs []string
 	nextPageToken := ""
-	fieldMaskCtx := fieldMask(
-		ctx,
-		"next_page_token",
-		"invocations.id",
-	)
+	searchTargets := strings.Contains(query, "id.target_id=")
 	for {
-		req := &resultstore.SearchInvocationsRequest{
-			Query:     query,
-			ProjectId: projectID,
-			PageStart: &resultstore.SearchInvocationsRequest_PageToken{
-				PageToken: nextPageToken,
-			},
+		var ids []string
+		var err error
+		if searchTargets {
+			ids, nextPageToken, err = c.targetSearch(ctx, log, query, projectID, nextPageToken)
+		} else {
+			ids, nextPageToken, err = c.invocationSearch(ctx, log, query, projectID, nextPageToken)
 		}
-		resp, err := c.client.SearchInvocations(fieldMaskCtx, req)
 		if err != nil {
 			return nil, err
 		}
-		for _, inv := range resp.GetInvocations() {
-			ids = append(ids, inv.Id.GetInvocationId())
-		}
-		if resp.GetNextPageToken() == "" {
+		invIDs = append(invIDs, ids...)
+		if nextPageToken == "" {
 			break
 		}
-		nextPageToken = resp.GetNextPageToken()
 	}
-	return ids, nil
+	return invIDs, nil
+}
+
+func (c *DownloadClient) invocationSearch(ctx context.Context, log logrus.FieldLogger, query, projectID, nextPageToken string) ([]string, string, error) {
+	fieldMaskCtx := fieldMask(
+		ctx,
+		"next_page_token",
+		"invocations.id",
+	)
+	req := &resultstore.SearchInvocationsRequest{
+		Query:     query,
+		ProjectId: projectID,
+		PageStart: &resultstore.SearchInvocationsRequest_PageToken{
+			PageToken: nextPageToken,
+		},
+	}
+	resp, err := c.client.SearchInvocations(fieldMaskCtx, req)
+	if err != nil {
+		return nil, "", err
+	}
+	var ids []string
+	for _, inv := range resp.GetInvocations() {
+		ids = append(ids, inv.GetId().GetInvocationId())
+	}
+	return ids, resp.GetNextPageToken(), err
+}
+
+func (c *DownloadClient) targetSearch(ctx context.Context, log logrus.FieldLogger, query, projectID, nextPageToken string) ([]string, string, error) {
+	fieldMaskCtx := fieldMask(
+		ctx,
+		"next_page_token",
+		"configured_targets.id",
+	)
+	req := &resultstore.SearchConfiguredTargetsRequest{
+		Query:     query,
+		ProjectId: projectID,
+		Parent:    "invocations/-/targets/-",
+		PageStart: &resultstore.SearchConfiguredTargetsRequest_PageToken{
+			PageToken: nextPageToken,
+		},
+	}
+	resp, err := c.client.SearchConfiguredTargets(fieldMaskCtx, req)
+	if err != nil {
+		return nil, "", err
+	}
+	var ids []string
+	for _, target := range resp.GetConfiguredTargets() {
+		ids = append(ids, target.GetId().GetInvocationId())
+	}
+	return ids, resp.GetNextPageToken(), err
 }
 
 // FetchResult provides a interface to store Resultstore invocation data.

diff --git a/pkg/updater/resultstore/query/BUILD.bazel b/pkg/updater/resultstore/query/BUILD.bazel
@@ -0,0 +1,28 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "go_default_library",
+    srcs = ["query.go"],
+    importpath = "github.com/GoogleCloudPlatform/testgrid/pkg/updater/resultstore/query",
+    visibility = ["//visibility:public"],
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = ["query_test.go"],
+    embed = [":go_default_library"],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)
diff --git a/pkg/updater/resultstore/query.go → pkg/updater/resultstore/query/query.go b/pkg/updater/resultstore/query.go → pkg/updater/resultstore/query/query.go
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package resultstore
+package query
 
 import (
 	"fmt"
@@ -47,7 +47,7 @@ var (
 	queryRe = regexp.MustCompile(`^target:".*"$`)
 )
 
-func translateQuery(simpleQuery string) (string, error) {
+func TranslateQuery(simpleQuery string) (string, error) {
 	if simpleQuery == "" {
 		return "", nil
 	}

diff --git a/pkg/updater/resultstore/query_test.go → pkg/updater/resultstore/query/query_test.go b/pkg/updater/resultstore/query_test.go → pkg/updater/resultstore/query/query_test.go
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package resultstore
+package query
 
 import (
 	"testing"
@@ -135,7 +135,7 @@ func TestTranslateQuery(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			got, err := translateQuery(tc.query)
+			got, err := TranslateQuery(tc.query)
 			if tc.want != got {
 				t.Errorf("translateQuery(%q) differed; got %q, want %q", tc.query, got, tc.want)
 			}

diff --git a/pkg/updater/resultstore/resultstore.go b/pkg/updater/resultstore/resultstore.go
@@ -26,6 +26,7 @@ import (
 	"time"
 
 	"github.com/GoogleCloudPlatform/testgrid/pkg/updater"
+	"github.com/GoogleCloudPlatform/testgrid/pkg/updater/resultstore/query"
 	"github.com/GoogleCloudPlatform/testgrid/util/gcs"
 	"github.com/sirupsen/logrus"
 
@@ -890,19 +891,22 @@ func queryAfter(query string, when time.Time) string {
 }
 
 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 := translateQuery(baseQuery)
+	query, err := query.TranslateQuery(baseQuery)
 	if err != nil {
 		return "", err
 	}
-	return queryAfter(fmt.Sprintf("%s %s", query, prowLabel), stop), nil
+	return queryAfter(fmt.Sprintf("%s %s", query, prowTargetLabel), stop), nil
 }
 
 func search(ctx context.Context, log logrus.FieldLogger, client *DownloadClient, rsConfig *configpb.ResultStoreConfig, stop time.Time) ([]string, error) {