@@ -33,6 +33,7 @@ import (
3333
3434type resultStoreClient interface {
3535 SearchInvocations (context.Context , * resultstore.SearchInvocationsRequest , ... grpc.CallOption ) (* resultstore.SearchInvocationsResponse , error )
36+ SearchConfiguredTargets (context.Context , * resultstore.SearchConfiguredTargetsRequest , ... grpc.CallOption ) (* resultstore.SearchConfiguredTargetsResponse , error )
3637 ExportInvocation (context.Context , * resultstore.ExportInvocationRequest , ... grpc.CallOption ) (* resultstore.ExportInvocationResponse , error )
3738}
3839
@@ -81,35 +82,76 @@ func Connect(ctx context.Context, serviceAccountPath string) (*grpc.ClientConn,
8182}
8283
8384// Search finds all the invocations that satisfies the query condition within a project.
84- func (c * DownloadClient ) Search (ctx context.Context , log logrus.FieldLogger , query , projectID string , fields ... string ) ([]string , error ) {
85- var ids []string
85+ func (c * DownloadClient ) Search (ctx context.Context , log logrus.FieldLogger , query , projectID string ) ([]string , error ) {
86+ var invIDs []string
8687 nextPageToken := ""
87- fieldMaskCtx := fieldMask (
88- ctx ,
89- "next_page_token" ,
90- "invocations.id" ,
91- )
88+ searchTargets := strings .Contains (query , "id.target_id=" )
9289 for {
93- req := & resultstore. SearchInvocationsRequest {
94- Query : query ,
95- ProjectId : projectID ,
96- PageStart : & resultstore. SearchInvocationsRequest_PageToken {
97- PageToken : nextPageToken ,
98- },
90+ var ids [] string
91+ var err error
92+ if searchTargets {
93+ ids , nextPageToken , err = c . targetSearch ( ctx , log , query , projectID , nextPageToken )
94+ } else {
95+ ids , nextPageToken , err = c . invocationSearch ( ctx , log , query , projectID , nextPageToken )
9996 }
100- resp , err := c .client .SearchInvocations (fieldMaskCtx , req )
10197 if err != nil {
10298 return nil , err
10399 }
104- for _ , inv := range resp .GetInvocations () {
105- ids = append (ids , inv .Id .GetInvocationId ())
106- }
107- if resp .GetNextPageToken () == "" {
100+ invIDs = append (invIDs , ids ... )
101+ if nextPageToken == "" {
108102 break
109103 }
110- nextPageToken = resp .GetNextPageToken ()
111104 }
112- return ids , nil
105+ return invIDs , nil
106+ }
107+
108+ func (c * DownloadClient ) invocationSearch (ctx context.Context , log logrus.FieldLogger , query , projectID , nextPageToken string ) ([]string , string , error ) {
109+ fieldMaskCtx := fieldMask (
110+ ctx ,
111+ "next_page_token" ,
112+ "invocations.id" ,
113+ )
114+ req := & resultstore.SearchInvocationsRequest {
115+ Query : query ,
116+ ProjectId : projectID ,
117+ PageStart : & resultstore.SearchInvocationsRequest_PageToken {
118+ PageToken : nextPageToken ,
119+ },
120+ }
121+ resp , err := c .client .SearchInvocations (fieldMaskCtx , req )
122+ if err != nil {
123+ return nil , "" , err
124+ }
125+ var ids []string
126+ for _ , inv := range resp .GetInvocations () {
127+ ids = append (ids , inv .GetId ().GetInvocationId ())
128+ }
129+ return ids , resp .GetNextPageToken (), err
130+ }
131+
132+ func (c * DownloadClient ) targetSearch (ctx context.Context , log logrus.FieldLogger , query , projectID , nextPageToken string ) ([]string , string , error ) {
133+ fieldMaskCtx := fieldMask (
134+ ctx ,
135+ "next_page_token" ,
136+ "configured_targets.id" ,
137+ )
138+ req := & resultstore.SearchConfiguredTargetsRequest {
139+ Query : query ,
140+ ProjectId : projectID ,
141+ Parent : "invocations/-/targets/-" ,
142+ PageStart : & resultstore.SearchConfiguredTargetsRequest_PageToken {
143+ PageToken : nextPageToken ,
144+ },
145+ }
146+ resp , err := c .client .SearchConfiguredTargets (fieldMaskCtx , req )
147+ if err != nil {
148+ return nil , "" , err
149+ }
150+ var ids []string
151+ for _ , target := range resp .GetConfiguredTargets () {
152+ ids = append (ids , target .GetId ().GetInvocationId ())
153+ }
154+ return ids , resp .GetNextPageToken (), err
113155}
114156
115157// FetchResult provides a interface to store Resultstore invocation data.
0 commit comments