From 730d947dc86f17d8d8867c0ba203e3d07191c8f2 Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Mon, 18 Dec 2023 12:53:26 -0800 Subject: [PATCH 01/14] feat(datastore): Query profiling --- datastore/integration_test.go | 260 ++++++++++++++++++++++++++++++++++ datastore/query.go | 178 ++++++++++++++++++++--- datastore/query_test.go | 58 ++++++++ 3 files changed, 474 insertions(+), 22 deletions(-) diff --git a/datastore/integration_test.go b/datastore/integration_test.go index ea50102898af..51575194e09f 100644 --- a/datastore/integration_test.go +++ b/datastore/integration_test.go @@ -24,6 +24,7 @@ import ( "net/url" "os" "reflect" + "slices" "sort" "strings" "testing" @@ -32,6 +33,7 @@ import ( "cloud.google.com/go/internal/testutil" "cloud.google.com/go/internal/uid" "cloud.google.com/go/rpcreplay" + "github.com/google/go-cmp/cmp" "google.golang.org/api/iterator" "google.golang.org/api/option" pb "google.golang.org/genproto/googleapis/datastore/v1" @@ -990,6 +992,106 @@ func TestIntegration_AggregationQueries(t *testing.T) { } +func TestIntegration_RunAggregationQueryWithOptions(t *testing.T) { + ctx := context.Background() + client := newTestClient(ctx, t) + defer client.Close() + + _, _, now, parent, cleanup := createTestEntities(t, ctx, client, "RunAggregationQueryWithOptions", 3) + defer cleanup() + + aggQuery := NewQuery("SQChild").Ancestor(parent).Filter("T=", now).NewAggregationQuery(). + WithSum("I", "i_sum").WithAvg("I", "i_avg").WithCount("count") + wantAggResult := map[string]interface{}{ + "i_sum": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: 6}}, + "i_avg": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: 2}}, + "count": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: 3}}, + } + testCases := []struct { + desc string + wantFailure bool + wantErrMsg string + wantRes AggregationWithOptionsResult + opts []RunOption + }{ + { + desc: "No mode", + wantRes: AggregationWithOptionsResult{ + Result: wantAggResult, + }, + }, + { + desc: "Normal mode", + wantRes: AggregationWithOptionsResult{ + Result: wantAggResult, + }, + opts: []RunOption{QueryModeNormal}, + }, + { + desc: "Explain mode", + wantRes: AggregationWithOptionsResult{ + Stats: &ResultSetStats{ + QueryPlan: &QueryPlan{ + PlanInfo: map[string]interface{}{ + "indexes_used": []interface{}{ + map[string]interface{}{ + "properties": "(T ASC, I ASC, __name__ ASC)", + "query_scope": "Includes Ancestors", + }, + }, + }, + }, + }, + }, + opts: []RunOption{QueryModeExplain}, + }, + { + desc: "ExplainAnalyze mode", + wantRes: AggregationWithOptionsResult{ + Result: wantAggResult, + Stats: &ResultSetStats{ + QueryPlan: &QueryPlan{ + PlanInfo: map[string]interface{}{ + "indexes_used": []interface{}{ + map[string]interface{}{ + "properties": "(T ASC, I ASC, __name__ ASC)", + "query_scope": "Includes Ancestors", + }, + }, + }, + }, + QueryStats: map[string]interface{}{ + "bytes_returned": "74", + "documents_scanned": "0", + "index_entries_scanned": "3", + "results_returned": "1", + "total_execution_time": "14.29 msecs", + }, + }, + }, + opts: []RunOption{QueryModeExplainAnalyze}, + }, + } + + for _, testcase := range testCases { + testutil.Retry(t, 10, time.Second, func(r *testutil.R) { + gotRes, gotErr := client.RunAggregationQueryWithOptions(ctx, aggQuery, testcase.opts...) + if gotErr != nil { + r.Errorf("err: got %v, want: nil", gotErr) + } + + if gotErr == nil && !reflect.DeepEqual(gotRes.Result, testcase.wantRes.Result) { + r.Errorf("%q: Mismatch in aggregation result got: %v, want: %v", testcase.desc, gotRes, testcase.wantRes) + return + } + + if err := isEqualResultSetStats(gotRes.Stats, testcase.wantRes.Stats); err != nil { + r.Errorf("%q: Mismatch in stats %+v", testcase.desc, err) + } + }) + } +} + type ckey struct{} func TestIntegration_LargeQuery(t *testing.T) { @@ -1291,6 +1393,164 @@ func TestIntegration_GetAllWithFieldMismatch(t *testing.T) { } } +func createTestEntities(t *testing.T, ctx context.Context, client *Client, partialNameKey string, count int) ([]*Key, []SQChild, int64, *Key, func()) { + parent := NameKey("SQParent", keyPrefix+partialNameKey+suffix, nil) + now := timeNow.Truncate(time.Millisecond).Unix() + + entities := []SQChild{} + for i := 0; i < count; i++ { + entities = append(entities, SQChild{I: i + 1, T: now, U: now, V: 1.5, W: "str"}) + } + + keys := make([]*Key, len(entities)) + for i := range keys { + keys[i] = IncompleteKey("SQChild", parent) + } + + // Create entities + keys, err := client.PutMulti(ctx, keys, entities) + if err != nil { + t.Fatalf("client.PutMulti: %v", err) + } + return keys, entities, now, parent, func() { + err := client.DeleteMulti(ctx, keys) + if err != nil { + t.Errorf("client.DeleteMulti: %v", err) + } + } +} + +func TestIntegration_RunAndGetAllWithOptions(t *testing.T) { + ctx := context.Background() + client := newTestClient(ctx, t) + defer client.Close() + + keys, entities, now, parent, cleanup := createTestEntities(t, ctx, client, "GetAllWithOptions", 3) + defer cleanup() + query := NewQuery("SQChild").Ancestor(parent).Filter("T=", now).Order("I") + for _, testcase := range []struct { + desc string + wantKeys []*Key + wantStats *ResultSetStats + wantEntities []SQChild + opts []RunOption + }{ + { + desc: "No mode", + wantKeys: keys, + wantEntities: entities, + }, + { + desc: "Normal query mode", + opts: []RunOption{QueryModeNormal}, + wantKeys: keys, + wantEntities: entities, + }, + { + desc: "Explain query mode", + opts: []RunOption{QueryModeExplain}, + wantStats: &ResultSetStats{ + QueryPlan: &QueryPlan{ + PlanInfo: map[string]interface{}{ + "indexes_used": []interface{}{ + map[string]interface{}{ + "properties": "(T ASC, I ASC, __name__ ASC)", + "query_scope": "Includes Ancestors", + }, + }, + }, + }, + }, + }, + { + desc: "ExplainAnalyze query mode", + opts: []RunOption{QueryModeExplainAnalyze}, + wantKeys: keys, + wantStats: &ResultSetStats{ + QueryPlan: &QueryPlan{ + PlanInfo: map[string]interface{}{ + "indexes_used": []interface{}{ + map[string]interface{}{ + "properties": "(T ASC, I ASC, __name__ ASC)", + "query_scope": "Includes Ancestors", + }, + }, + }, + }, + QueryStats: map[string]interface{}{ + "bytes_returned": "552", + "documents_scanned": "3", + "index_entries_scanned": "3", + "results_returned": "3", + "total_execution_time": "14.29 msecs", + }, + }, + wantEntities: entities, + }, + } { + // Test GetAllWithOptions + var gotSQChildsFromGetAll []SQChild + gotRes, gotErr := client.GetAllWithOptions(ctx, query, &gotSQChildsFromGetAll, testcase.opts...) + if gotErr != nil { + t.Errorf("%v err: got: %+v, want: nil", testcase.desc, gotErr) + } + if !testutil.Equal(gotSQChildsFromGetAll, testcase.wantEntities) { + t.Errorf("%v entities: got: %+v, want: %+v", testcase.desc, gotSQChildsFromGetAll, testcase.wantEntities) + } + if !testutil.Equal(gotRes.Keys, testcase.wantKeys) { + t.Errorf("%v keys: got: %+v, want: %+v", testcase.desc, gotRes.Keys, testcase.wantKeys) + } + if err := isEqualResultSetStats(gotRes.Stats, testcase.wantStats); err != nil { + t.Errorf("%v %+v", testcase.desc, err) + } + + // Test Run() + var gotSQChildsFromRun []SQChild + iter := client.Run(ctx, query, testcase.opts...) + for { + var gotSQChild SQChild + _, err := iter.Next(&gotSQChild) + if err == iterator.Done { + break + } + if err != nil { + t.Errorf("%v iter.Next: %v", testcase.desc, err) + } + gotSQChildsFromRun = append(gotSQChildsFromRun, gotSQChild) + } + if !testutil.Equal(gotSQChildsFromRun, testcase.wantEntities) { + t.Errorf("%v entities: got: %+v, want: %+v", testcase.desc, gotSQChildsFromRun, testcase.wantEntities) + } + if err := isEqualResultSetStats(iter.Stats, testcase.wantStats); err != nil { + t.Errorf("%v %+v", testcase.desc, err) + } + } +} + +func ignoreStatsFields(p cmp.Path) bool { + step, ok := p[len(p)-1].(cmp.MapIndex) + // Ignore these fields while comparing stats as these can differ per run + fields := []string{"bytes_returned", "total_execution_time"} + return ok && slices.Contains(fields, step.Key().String()) +} + +func isEqualResultSetStats(got *ResultSetStats, want *ResultSetStats) error { + if (got != nil && want == nil) || (got == nil && want != nil) { + return fmt.Errorf("Stats: got: %+v, want: %+v", got, want) + } + if got != nil { + if !testutil.Equal(got.QueryPlan, want.QueryPlan) { + return fmt.Errorf("Stats.QueryPlan.PlanInfo: got: %+v, want: %+v", got.QueryPlan, want.QueryPlan) + } + + // Compare query stats maps except 'total_execution_time' key + if !testutil.Equal(got.QueryStats, want.QueryStats, cmp.FilterPath(ignoreStatsFields, cmp.Ignore())) { + return fmt.Errorf("Stats.QueryPlan.QueryStats: got: %+v, want: %+v", got.QueryStats, want.QueryStats) + } + } + return nil +} + func TestIntegration_KindlessQueries(t *testing.T) { ctx := context.Background() client := newTestClient(ctx, t) diff --git a/datastore/query.go b/datastore/query.go index e833ff8d8ff4..55bf68eff456 100644 --- a/datastore/query.go +++ b/datastore/query.go @@ -24,6 +24,7 @@ import ( "strconv" "strings" + "cloud.google.com/go/internal/protostruct" "cloud.google.com/go/internal/trace" wrapperspb "github.com/golang/protobuf/ptypes/wrappers" "google.golang.org/api/iterator" @@ -625,6 +626,70 @@ func (c *Client) Count(ctx context.Context, q *Query) (n int, err error) { } } +const ( + // QueryModeNormal is the default mode. Only the query results are returned. + QueryModeNormal QueryMode = iota + // QueryModeExplain returns only the query plan, without any results or execution + // statistics information. + QueryModeExplain + // QueryModeExplainAnalyze returns both the query plan and the execution statistics along + // with the results. + QueryModeExplainAnalyze +) + +// RunOption lets the user provide options while running a query +type RunOption interface { + apply(*runQuerySettings) error +} + +// QueryMode is the mode in which the query request must be processed. +type QueryMode int32 + +func (m QueryMode) apply(s *runQuerySettings) error { + if s.mode != nil { + return errors.New("datastore: only one mode can be specified") + } + pbQueryMode := pb.QueryMode(m) + s.mode = &pbQueryMode + return nil +} + +type runQuerySettings struct { + mode *pb.QueryMode +} + +// newRunQuerySettings creates a runQuerySettings with a given RunOption slice. +func newRunQuerySettings(opts []RunOption) (*runQuerySettings, error) { + s := &runQuerySettings{} + for _, o := range opts { + if o == nil { + return nil, errors.New("datastore: RunOption cannot be nil") + } + err := o.apply(s) + if err != nil { + return nil, err + } + } + return s, nil +} + +// ResultSetStats is planning and execution statistics for the query. +type ResultSetStats struct { + QueryPlan *QueryPlan + QueryStats map[string]interface{} +} + +// QueryPlan is plan for the query. +type QueryPlan struct { + PlanInfo map[string]interface{} +} + +// GetAllWithOptionsResult is the result of call to GetAllWithOptions method +type GetAllWithOptionsResult struct { + Keys []*Key + Stats *ResultSetStats +} + // GetAll runs the provided query in the given context and returns all keys // that match that query, as well as appending the values to dst. // @@ -649,6 +714,15 @@ func (c *Client) GetAll(ctx context.Context, q *Query, dst interface{}) (keys [] ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.Query.GetAll") defer func() { trace.EndSpan(ctx, err) }() + res, err := c.GetAllWithOptions(ctx, q, dst) + return res.Keys, err +} + +// GetAllWithOptions is similar to GetAll but runs the query with provided options +func (c *Client) GetAllWithOptions(ctx context.Context, q *Query, dst interface{}, opts ...RunOption) (res GetAllWithOptionsResult, err error) { + ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.Query.GetAllWithOptions") + defer func() { trace.EndSpan(ctx, err) }() + var ( dv reflect.Value mat multiArgType @@ -658,22 +732,23 @@ func (c *Client) GetAll(ctx context.Context, q *Query, dst interface{}) (keys [] if !q.keysOnly { dv = reflect.ValueOf(dst) if dv.Kind() != reflect.Ptr || dv.IsNil() { - return nil, ErrInvalidEntityType + return res, ErrInvalidEntityType } dv = dv.Elem() mat, elemType = checkMultiArg(dv) if mat == multiArgTypeInvalid || mat == multiArgTypeInterface { - return nil, ErrInvalidEntityType + return res, ErrInvalidEntityType } } - for t := c.Run(ctx, q); ; { + for t := c.Run(ctx, q, opts...); ; { k, e, err := t.next() + res.Stats = t.Stats if err == iterator.Done { break } if err != nil { - return keys, err + return res, err } if !q.keysOnly { ev := reflect.New(elemType) @@ -700,7 +775,7 @@ func (c *Client) GetAll(ctx context.Context, q *Query, dst interface{}) (keys [] // an ErrFieldMismatch is returned. errFieldMismatch = err } else { - return keys, err + return res, err } } if mat != multiArgTypeStructPtr { @@ -708,13 +783,13 @@ func (c *Client) GetAll(ctx context.Context, q *Query, dst interface{}) (keys [] } dv.Set(reflect.Append(dv, ev)) } - keys = append(keys, k) + res.Keys = append(res.Keys, k) } - return keys, errFieldMismatch + return res, errFieldMismatch } -// Run runs the given query in the given context. -func (c *Client) Run(ctx context.Context, q *Query) *Iterator { +// Run runs the given query in the given context with the provided options +func (c *Client) Run(ctx context.Context, q *Query, opts ...RunOption) *Iterator { if q.err != nil { return &Iterator{err: q.err} } @@ -738,6 +813,16 @@ func (c *Client) Run(ctx context.Context, q *Query) *Iterator { } } + runSettings, err := newRunQuerySettings(opts) + if err != nil { + t.err = err + return t + } + + if runSettings.mode != nil { + t.req.Mode = *runSettings.mode + } + if err := q.toRunQueryRequest(t.req); err != nil { t.err = err } @@ -748,22 +833,30 @@ func (c *Client) Run(ctx context.Context, q *Query) *Iterator { func (c *Client) RunAggregationQuery(ctx context.Context, aq *AggregationQuery) (ar AggregationResult, err error) { ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.Query.RunAggregationQuery") defer func() { trace.EndSpan(ctx, err) }() + aro, err := c.RunAggregationQueryWithOptions(ctx, aq, nil) + return aro.Result, err +} + +// RunAggregationQueryWithOptions runs aggregation query (e.g. COUNT) with provided options and returns results from the service. +func (c *Client) RunAggregationQueryWithOptions(ctx context.Context, aq *AggregationQuery, opts ...RunOption) (ar AggregationWithOptionsResult, err error) { + ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.Query.RunAggregationQueryWithOptions") + defer func() { trace.EndSpan(ctx, err) }() if aq == nil { - return nil, errors.New("datastore: aggregation query cannot be nil") + return ar, errors.New("datastore: aggregation query cannot be nil") } if aq.query == nil { - return nil, errors.New("datastore: aggregation query must include nested query") + return ar, errors.New("datastore: aggregation query must include nested query") } if len(aq.aggregationQueries) == 0 { - return nil, errors.New("datastore: aggregation query must contain one or more operators (e.g. count)") + return ar, errors.New("datastore: aggregation query must contain one or more operators (e.g. count)") } q, err := aq.query.toProto() if err != nil { - return nil, err + return ar, err } req := &pb.RunAggregationQueryRequest{ @@ -785,25 +878,35 @@ func (c *Client) RunAggregationQuery(ctx context.Context, aq *AggregationQuery) } } + runSettings, err := newRunQuerySettings(opts) + if err != nil { + return ar, err + } + if runSettings.mode != nil { + req.Mode = *runSettings.mode + } + // Parse the read options. req.ReadOptions, err = parseReadOptions(aq.query) if err != nil { - return nil, err + return ar, err } - res, err := c.client.RunAggregationQuery(ctx, req) + resp, err := c.client.RunAggregationQuery(ctx, req) if err != nil { - return nil, err + return ar, err } - ar = make(AggregationResult) - - // TODO(developer): change batch parsing logic if other aggregations are supported. - for _, a := range res.Batch.AggregationResults { - for k, v := range a.AggregateProperties { - ar[k] = v + if req.Mode != pb.QueryMode_PLAN { + ar.Result = make(AggregationResult) + // TODO(developer): change batch parsing logic if other aggregations are supported. + for _, a := range resp.Batch.AggregationResults { + for k, v := range a.AggregateProperties { + ar.Result[k] = v + } } } + ar.Stats = fromPbResultSetStats(resp.Stats) return ar, nil } @@ -858,6 +961,9 @@ type Iterator struct { pageCursor []byte // entityCursor is the compiled cursor of the next result. entityCursor []byte + + // Stats contains query plan and execution statistics. + Stats *ResultSetStats } // Next returns the key of the next result. When there are no more results, @@ -930,6 +1036,13 @@ func (t *Iterator) nextBatch() error { return err } + if t.req.Mode == pb.QueryMode_PLAN { + // No results to process + t.limit = 0 + t.Stats = fromPbResultSetStats(resp.Stats) + return nil + } + // Adjust any offset from skipped results. skip := resp.Batch.SkippedResults if skip < 0 { @@ -969,9 +1082,24 @@ func (t *Iterator) nextBatch() error { t.pageCursor = resp.Batch.EndCursor t.results = resp.Batch.EntityResults + t.Stats = fromPbResultSetStats(resp.Stats) return nil } +func fromPbResultSetStats(pbstats *pb.ResultSetStats) *ResultSetStats { + if pbstats == nil { + return nil + } + planInfo := protostruct.DecodeToMap(pbstats.QueryPlan.PlanInfo) + queryStats := protostruct.DecodeToMap(pbstats.QueryStats) + return &ResultSetStats{ + QueryPlan: &QueryPlan{ + PlanInfo: planInfo, + }, + QueryStats: queryStats, + } +} + // Cursor returns a cursor for the iterator's current location. func (t *Iterator) Cursor() (c Cursor, err error) { t.ctx = trace.StartSpan(t.ctx, "cloud.google.com/go/datastore.Query.Cursor") @@ -1102,3 +1230,9 @@ func (aq *AggregationQuery) WithAvg(fieldName string, alias string) *Aggregation // AggregationResult contains the results of an aggregation query. type AggregationResult map[string]interface{} + +// AggregationWithOptionsResult contains the results of an aggregation query run with options. +type AggregationWithOptionsResult struct { + Result AggregationResult + Stats *ResultSetStats +} diff --git a/datastore/query_test.go b/datastore/query_test.go index 40fc8f327795..b5e1a608d52f 100644 --- a/datastore/query_test.go +++ b/datastore/query_test.go @@ -20,6 +20,7 @@ import ( "fmt" "reflect" "sort" + "strings" "testing" "cloud.google.com/go/internal/testutil" @@ -858,3 +859,60 @@ func TestAggregationQueryIsNil(t *testing.T) { t.Fatal(err) } } + +func TestQueryModeApply(t *testing.T) { + pbNormal := pb.QueryMode(QueryModeNormal) + for _, testcase := range []struct { + desc string + existingMode *pb.QueryMode + newMode QueryMode + wantErrMsg string + }{ + { + desc: "Multiple modes", + existingMode: &pbNormal, + newMode: QueryModeExplain, + wantErrMsg: "only one mode can be specified", + }, + { + desc: "Single mode", + existingMode: nil, + newMode: QueryModeExplain, + }, + } { + gotErr := testcase.newMode.apply(&runQuerySettings{mode: testcase.existingMode}) + if (gotErr == nil && testcase.wantErrMsg != "") || + (gotErr != nil && !strings.Contains(gotErr.Error(), testcase.wantErrMsg)) { + t.Errorf("%v: apply got: %v want: %v", testcase.desc, gotErr, testcase.wantErrMsg) + } + } +} + +func TestNewRunQuerySettings(t *testing.T) { + for _, testcase := range []struct { + desc string + opts []RunOption + wantErrMsg string + }{ + { + desc: "nil RunOption", + opts: []RunOption{QueryModeNormal, nil}, + wantErrMsg: "cannot be nil", + }, + { + desc: "success RunOption", + opts: []RunOption{QueryModeNormal}, + }, + { + desc: "multiple modes", + opts: []RunOption{QueryModeNormal, QueryModeExplainAnalyze}, + wantErrMsg: "only one mode can be specified", + }, + } { + _, gotErr := newRunQuerySettings(testcase.opts) + if (gotErr == nil && testcase.wantErrMsg != "") || + (gotErr != nil && !strings.Contains(gotErr.Error(), testcase.wantErrMsg)) { + t.Errorf("%v: newRunQuerySettings got: %v want: %v", testcase.desc, gotErr, testcase.wantErrMsg) + } + } +} From 5825267718913afb9be40fa06429f59b907a8c8b Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Mon, 18 Dec 2023 13:02:10 -0800 Subject: [PATCH 02/14] feat(datastore): remove comment --- datastore/integration_test.go | 1 - 1 file changed, 1 deletion(-) diff --git a/datastore/integration_test.go b/datastore/integration_test.go index 51575194e09f..9dc7d4e0531e 100644 --- a/datastore/integration_test.go +++ b/datastore/integration_test.go @@ -1543,7 +1543,6 @@ func isEqualResultSetStats(got *ResultSetStats, want *ResultSetStats) error { return fmt.Errorf("Stats.QueryPlan.PlanInfo: got: %+v, want: %+v", got.QueryPlan, want.QueryPlan) } - // Compare query stats maps except 'total_execution_time' key if !testutil.Equal(got.QueryStats, want.QueryStats, cmp.FilterPath(ignoreStatsFields, cmp.Ignore())) { return fmt.Errorf("Stats.QueryPlan.QueryStats: got: %+v, want: %+v", got.QueryStats, want.QueryStats) } From 4d8a1f17955fa254827f6b960f6a2db38c5afa4f Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Tue, 26 Dec 2023 20:54:15 -0800 Subject: [PATCH 03/14] feat(datastore): Refactor test --- datastore/integration_test.go | 27 +++++++++------------------ datastore/query.go | 2 +- 2 files changed, 10 insertions(+), 19 deletions(-) diff --git a/datastore/integration_test.go b/datastore/integration_test.go index 9dc7d4e0531e..18a4fc6e72f9 100644 --- a/datastore/integration_test.go +++ b/datastore/integration_test.go @@ -24,7 +24,6 @@ import ( "net/url" "os" "reflect" - "slices" "sort" "strings" "testing" @@ -33,7 +32,6 @@ import ( "cloud.google.com/go/internal/testutil" "cloud.google.com/go/internal/uid" "cloud.google.com/go/rpcreplay" - "github.com/google/go-cmp/cmp" "google.golang.org/api/iterator" "google.golang.org/api/option" pb "google.golang.org/genproto/googleapis/datastore/v1" @@ -997,7 +995,7 @@ func TestIntegration_RunAggregationQueryWithOptions(t *testing.T) { client := newTestClient(ctx, t) defer client.Close() - _, _, now, parent, cleanup := createTestEntities(t, ctx, client, "RunAggregationQueryWithOptions", 3) + _, _, now, parent, cleanup := createTestEntities(ctx, t, client, "RunAggregationQueryWithOptions", 3) defer cleanup() aggQuery := NewQuery("SQChild").Ancestor(parent).Filter("T=", now).NewAggregationQuery(). @@ -1061,11 +1059,9 @@ func TestIntegration_RunAggregationQueryWithOptions(t *testing.T) { }, }, QueryStats: map[string]interface{}{ - "bytes_returned": "74", "documents_scanned": "0", "index_entries_scanned": "3", "results_returned": "1", - "total_execution_time": "14.29 msecs", }, }, }, @@ -1393,7 +1389,7 @@ func TestIntegration_GetAllWithFieldMismatch(t *testing.T) { } } -func createTestEntities(t *testing.T, ctx context.Context, client *Client, partialNameKey string, count int) ([]*Key, []SQChild, int64, *Key, func()) { +func createTestEntities(ctx context.Context, t *testing.T, client *Client, partialNameKey string, count int) ([]*Key, []SQChild, int64, *Key, func()) { parent := NameKey("SQParent", keyPrefix+partialNameKey+suffix, nil) now := timeNow.Truncate(time.Millisecond).Unix() @@ -1425,7 +1421,7 @@ func TestIntegration_RunAndGetAllWithOptions(t *testing.T) { client := newTestClient(ctx, t) defer client.Close() - keys, entities, now, parent, cleanup := createTestEntities(t, ctx, client, "GetAllWithOptions", 3) + keys, entities, now, parent, cleanup := createTestEntities(ctx, t, client, "GetAllWithOptions", 3) defer cleanup() query := NewQuery("SQChild").Ancestor(parent).Filter("T=", now).Order("I") for _, testcase := range []struct { @@ -1478,11 +1474,9 @@ func TestIntegration_RunAndGetAllWithOptions(t *testing.T) { }, }, QueryStats: map[string]interface{}{ - "bytes_returned": "552", "documents_scanned": "3", "index_entries_scanned": "3", "results_returned": "3", - "total_execution_time": "14.29 msecs", }, }, wantEntities: entities, @@ -1527,13 +1521,6 @@ func TestIntegration_RunAndGetAllWithOptions(t *testing.T) { } } -func ignoreStatsFields(p cmp.Path) bool { - step, ok := p[len(p)-1].(cmp.MapIndex) - // Ignore these fields while comparing stats as these can differ per run - fields := []string{"bytes_returned", "total_execution_time"} - return ok && slices.Contains(fields, step.Key().String()) -} - func isEqualResultSetStats(got *ResultSetStats, want *ResultSetStats) error { if (got != nil && want == nil) || (got == nil && want != nil) { return fmt.Errorf("Stats: got: %+v, want: %+v", got, want) @@ -1543,8 +1530,12 @@ func isEqualResultSetStats(got *ResultSetStats, want *ResultSetStats) error { return fmt.Errorf("Stats.QueryPlan.PlanInfo: got: %+v, want: %+v", got.QueryPlan, want.QueryPlan) } - if !testutil.Equal(got.QueryStats, want.QueryStats, cmp.FilterPath(ignoreStatsFields, cmp.Ignore())) { - return fmt.Errorf("Stats.QueryPlan.QueryStats: got: %+v, want: %+v", got.QueryStats, want.QueryStats) + for wantK, wantV := range want.QueryStats { + gotV, ok := got.QueryStats[wantK] + if !ok || !testutil.Equal(gotV, wantV) { + fmt.Printf("Stats.QueryPlan.QueryStats: gotV: %+v, wantV: %+v", gotV, wantV) + return fmt.Errorf("Stats.QueryPlan.QueryStats: got: %+v, want: %+v", got.QueryStats, want.QueryStats) + } } } return nil diff --git a/datastore/query.go b/datastore/query.go index 55bf68eff456..39af49fdfbbf 100644 --- a/datastore/query.go +++ b/datastore/query.go @@ -833,7 +833,7 @@ func (c *Client) Run(ctx context.Context, q *Query, opts ...RunOption) *Iterator func (c *Client) RunAggregationQuery(ctx context.Context, aq *AggregationQuery) (ar AggregationResult, err error) { ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.Query.RunAggregationQuery") defer func() { trace.EndSpan(ctx, err) }() - aro, err := c.RunAggregationQueryWithOptions(ctx, aq, nil) + aro, err := c.RunAggregationQueryWithOptions(ctx, aq) return aro.Result, err } From 4276f61aa6cfde1bf065129115fe9c235df31a5f Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Tue, 16 Jan 2024 21:48:45 -0800 Subject: [PATCH 04/14] feat(datastore): Refactor integration test --- datastore/integration_test.go | 49 ++++++++++++++++++++++------------- datastore/query.go | 6 ++--- 2 files changed, 34 insertions(+), 21 deletions(-) diff --git a/datastore/integration_test.go b/datastore/integration_test.go index 18a4fc6e72f9..3ef8c73f8e70 100644 --- a/datastore/integration_test.go +++ b/datastore/integration_test.go @@ -1416,21 +1416,17 @@ func createTestEntities(ctx context.Context, t *testing.T, client *Client, parti } } -func TestIntegration_RunAndGetAllWithOptions(t *testing.T) { - ctx := context.Background() - client := newTestClient(ctx, t) - defer client.Close() +type runWithOptionsTestcase struct { + desc string + wantKeys []*Key + wantStats *ResultSetStats + wantEntities []SQChild + opts []RunOption +} - keys, entities, now, parent, cleanup := createTestEntities(ctx, t, client, "GetAllWithOptions", 3) - defer cleanup() - query := NewQuery("SQChild").Ancestor(parent).Filter("T=", now).Order("I") - for _, testcase := range []struct { - desc string - wantKeys []*Key - wantStats *ResultSetStats - wantEntities []SQChild - opts []RunOption - }{ +func getRunWithOptionsTestcases(ctx context.Context, t *testing.T, client *Client, partialNameKey string, count int) ([]runWithOptionsTestcase, int64, *Key, func()) { + keys, entities, now, parent, cleanup := createTestEntities(ctx, t, client, partialNameKey, count) + return []runWithOptionsTestcase{ { desc: "No mode", wantKeys: keys, @@ -1481,8 +1477,17 @@ func TestIntegration_RunAndGetAllWithOptions(t *testing.T) { }, wantEntities: entities, }, - } { - // Test GetAllWithOptions + }, now, parent, cleanup +} + +func TestIntegration_GetAllWithOptions(t *testing.T) { + ctx := context.Background() + client := newTestClient(ctx, t) + defer client.Close() + testcases, now, parent, cleanup := getRunWithOptionsTestcases(ctx, t, client, "GetAllWithOptions", 3) + defer cleanup() + query := NewQuery("SQChild").Ancestor(parent).Filter("T=", now).Order("I") + for _, testcase := range testcases { var gotSQChildsFromGetAll []SQChild gotRes, gotErr := client.GetAllWithOptions(ctx, query, &gotSQChildsFromGetAll, testcase.opts...) if gotErr != nil { @@ -1497,8 +1502,17 @@ func TestIntegration_RunAndGetAllWithOptions(t *testing.T) { if err := isEqualResultSetStats(gotRes.Stats, testcase.wantStats); err != nil { t.Errorf("%v %+v", testcase.desc, err) } + } +} - // Test Run() +func TestIntegration_RunWithOptions(t *testing.T) { + ctx := context.Background() + client := newTestClient(ctx, t) + defer client.Close() + testcases, now, parent, cleanup := getRunWithOptionsTestcases(ctx, t, client, "RunWithOptions", 3) + defer cleanup() + query := NewQuery("SQChild").Ancestor(parent).Filter("T=", now).Order("I") + for _, testcase := range testcases { var gotSQChildsFromRun []SQChild iter := client.Run(ctx, query, testcase.opts...) for { @@ -1533,7 +1547,6 @@ func isEqualResultSetStats(got *ResultSetStats, want *ResultSetStats) error { for wantK, wantV := range want.QueryStats { gotV, ok := got.QueryStats[wantK] if !ok || !testutil.Equal(gotV, wantV) { - fmt.Printf("Stats.QueryPlan.QueryStats: gotV: %+v, wantV: %+v", gotV, wantV) return fmt.Errorf("Stats.QueryPlan.QueryStats: got: %+v, want: %+v", got.QueryStats, want.QueryStats) } } diff --git a/datastore/query.go b/datastore/query.go index 39af49fdfbbf..7b1e13faa8a5 100644 --- a/datastore/query.go +++ b/datastore/query.go @@ -628,13 +628,13 @@ func (c *Client) Count(ctx context.Context, q *Query) (n int, err error) { const ( // QueryModeNormal is the default mode. Only the query results are returned. - QueryModeNormal QueryMode = iota + QueryModeNormal QueryMode = iota // = 0 // QueryModeExplain returns only the query plan, without any results or execution // statistics information. - QueryModeExplain + QueryModeExplain // = 1 // QueryModeExplainAnalyze returns both the query plan and the execution statistics along // with the results. - QueryModeExplainAnalyze + QueryModeExplainAnalyze // = 2 ) // RunOption lets the user provide options while running a query From 21923dc45389528d2e38425b3f50a21b037f5e1c Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Tue, 13 Feb 2024 21:13:04 -0800 Subject: [PATCH 05/14] feat(datastore): Use ExplainOptions instead of QueryMode --- datastore/integration_test.go | 164 ++++++++++++++++----------------- datastore/query.go | 169 ++++++++++++++++++++++++---------- datastore/query_test.go | 46 +++++---- 3 files changed, 224 insertions(+), 155 deletions(-) diff --git a/datastore/integration_test.go b/datastore/integration_test.go index 1627a07e0226..20aa16d9b71b 100644 --- a/datastore/integration_test.go +++ b/datastore/integration_test.go @@ -32,6 +32,7 @@ import ( "cloud.google.com/go/internal/testutil" "cloud.google.com/go/internal/uid" "cloud.google.com/go/rpcreplay" + "github.com/google/go-cmp/cmp/cmpopts" "google.golang.org/api/iterator" "google.golang.org/api/option" pb "google.golang.org/genproto/googleapis/datastore/v1" @@ -1005,6 +1006,7 @@ func TestIntegration_RunAggregationQueryWithOptions(t *testing.T) { "i_avg": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: 2}}, "count": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: 3}}, } + testCases := []struct { desc string wantFailure bool @@ -1013,59 +1015,46 @@ func TestIntegration_RunAggregationQueryWithOptions(t *testing.T) { opts []RunOption }{ { - desc: "No mode", + desc: "No options", wantRes: AggregationWithOptionsResult{ Result: wantAggResult, }, }, { - desc: "Normal mode", + desc: "ExplainOptions.Analyze is false", wantRes: AggregationWithOptionsResult{ - Result: wantAggResult, - }, - opts: []RunOption{QueryModeNormal}, - }, - { - desc: "Explain mode", - wantRes: AggregationWithOptionsResult{ - Stats: &ResultSetStats{ - QueryPlan: &QueryPlan{ - PlanInfo: map[string]interface{}{ - "indexes_used": []interface{}{ - map[string]interface{}{ - "properties": "(T ASC, I ASC, __name__ ASC)", - "query_scope": "Includes Ancestors", - }, - }, + Plan: &Plan{ + IndexesUsed: []*map[string]interface{}{ + { + "properties": "(T ASC, I ASC, __name__ ASC)", + "query_scope": "Includes Ancestors", }, }, }, }, - opts: []RunOption{QueryModeExplain}, + opts: []RunOption{ExplainOptions{}}, }, { - desc: "ExplainAnalyze mode", + desc: "ExplainAnalyze.Analyze is true", wantRes: AggregationWithOptionsResult{ Result: wantAggResult, - Stats: &ResultSetStats{ - QueryPlan: &QueryPlan{ - PlanInfo: map[string]interface{}{ - "indexes_used": []interface{}{ - map[string]interface{}{ - "properties": "(T ASC, I ASC, __name__ ASC)", - "query_scope": "Includes Ancestors", - }, - }, + Plan: &Plan{ + IndexesUsed: []*map[string]interface{}{ + { + "properties": "(T ASC, I ASC, __name__ ASC)", + "query_scope": "Includes Ancestors", }, }, - QueryStats: map[string]interface{}{ + }, + ExecutionStats: &ExecutionStats{ + ResultsReturned: 1, + DebugStats: &map[string]interface{}{ "documents_scanned": "0", "index_entries_scanned": "3", - "results_returned": "1", }, }, }, - opts: []RunOption{QueryModeExplainAnalyze}, + opts: []RunOption{ExplainOptions{Analyze: true}}, }, } @@ -1081,7 +1070,11 @@ func TestIntegration_RunAggregationQueryWithOptions(t *testing.T) { return } - if err := isEqualResultSetStats(gotRes.Stats, testcase.wantRes.Stats); err != nil { + if !testutil.Equal(gotRes.Plan, testcase.wantRes.Plan) { + t.Errorf("%v Plan: got: %+v, want: %+v", testcase.desc, gotRes.Plan, testcase.wantRes.Plan) + } + + if err := isEqualExecutionStats(gotRes.ExecutionStats, testcase.wantRes.ExecutionStats); err != nil { r.Errorf("%q: Mismatch in stats %+v", testcase.desc, err) } }) @@ -1432,62 +1425,51 @@ func createTestEntities(ctx context.Context, t *testing.T, client *Client, parti } type runWithOptionsTestcase struct { - desc string - wantKeys []*Key - wantStats *ResultSetStats - wantEntities []SQChild - opts []RunOption + desc string + wantKeys []*Key + wantExecutionStats *ExecutionStats + wantPlan *Plan + wantEntities []SQChild + opts []RunOption } func getRunWithOptionsTestcases(ctx context.Context, t *testing.T, client *Client, partialNameKey string, count int) ([]runWithOptionsTestcase, int64, *Key, func()) { keys, entities, now, parent, cleanup := createTestEntities(ctx, t, client, partialNameKey, count) return []runWithOptionsTestcase{ { - desc: "No mode", + desc: "No ExplainOptions", wantKeys: keys, wantEntities: entities, }, { - desc: "Normal query mode", - opts: []RunOption{QueryModeNormal}, - wantKeys: keys, - wantEntities: entities, - }, - { - desc: "Explain query mode", - opts: []RunOption{QueryModeExplain}, - wantStats: &ResultSetStats{ - QueryPlan: &QueryPlan{ - PlanInfo: map[string]interface{}{ - "indexes_used": []interface{}{ - map[string]interface{}{ - "properties": "(T ASC, I ASC, __name__ ASC)", - "query_scope": "Includes Ancestors", - }, - }, + desc: "ExplainOptions.Analyze is false", + opts: []RunOption{ExplainOptions{}}, + wantPlan: &Plan{ + IndexesUsed: []*map[string]interface{}{ + { + "properties": "(T ASC, I ASC, __name__ ASC)", + "query_scope": "Includes Ancestors", }, }, }, }, { - desc: "ExplainAnalyze query mode", - opts: []RunOption{QueryModeExplainAnalyze}, + desc: "ExplainAnalyze.Analyze is true", + opts: []RunOption{ExplainOptions{Analyze: true}}, wantKeys: keys, - wantStats: &ResultSetStats{ - QueryPlan: &QueryPlan{ - PlanInfo: map[string]interface{}{ - "indexes_used": []interface{}{ - map[string]interface{}{ - "properties": "(T ASC, I ASC, __name__ ASC)", - "query_scope": "Includes Ancestors", - }, - }, - }, - }, - QueryStats: map[string]interface{}{ + wantExecutionStats: &ExecutionStats{ + ResultsReturned: 1, + DebugStats: &map[string]interface{}{ "documents_scanned": "3", "index_entries_scanned": "3", - "results_returned": "3", + }, + }, + wantPlan: &Plan{ + IndexesUsed: []*map[string]interface{}{ + { + "properties": "(T ASC, I ASC, __name__ ASC)", + "query_scope": "Includes Ancestors", + }, }, }, wantEntities: entities, @@ -1514,7 +1496,10 @@ func TestIntegration_GetAllWithOptions(t *testing.T) { if !testutil.Equal(gotRes.Keys, testcase.wantKeys) { t.Errorf("%v keys: got: %+v, want: %+v", testcase.desc, gotRes.Keys, testcase.wantKeys) } - if err := isEqualResultSetStats(gotRes.Stats, testcase.wantStats); err != nil { + if !testutil.Equal(gotRes.Plan, testcase.wantPlan) { + t.Errorf("%v Plan: got: %+v, want: %+v", testcase.desc, gotRes.Plan, testcase.wantPlan) + } + if err := isEqualExecutionStats(gotRes.ExecutionStats, testcase.wantExecutionStats); err != nil { t.Errorf("%v %+v", testcase.desc, err) } } @@ -1544,28 +1529,39 @@ func TestIntegration_RunWithOptions(t *testing.T) { if !testutil.Equal(gotSQChildsFromRun, testcase.wantEntities) { t.Errorf("%v entities: got: %+v, want: %+v", testcase.desc, gotSQChildsFromRun, testcase.wantEntities) } - if err := isEqualResultSetStats(iter.Stats, testcase.wantStats); err != nil { + if !testutil.Equal(iter.Plan, testcase.wantPlan) { + t.Errorf("%v Plan: got: %+v, want: %+v", testcase.desc, iter.Plan, testcase.wantPlan) + } + if err := isEqualExecutionStats(iter.ExecutionStats, testcase.wantExecutionStats); err != nil { t.Errorf("%v %+v", testcase.desc, err) } } } -func isEqualResultSetStats(got *ResultSetStats, want *ResultSetStats) error { +func isEqualExecutionStats(got *ExecutionStats, want *ExecutionStats) error { if (got != nil && want == nil) || (got == nil && want != nil) { - return fmt.Errorf("Stats: got: %+v, want: %+v", got, want) + return fmt.Errorf("ExecutionStats: got: %+v, want: %+v", got, want) + } + if got == nil { + return nil } - if got != nil { - if !testutil.Equal(got.QueryPlan, want.QueryPlan) { - return fmt.Errorf("Stats.QueryPlan.PlanInfo: got: %+v, want: %+v", got.QueryPlan, want.QueryPlan) - } - for wantK, wantV := range want.QueryStats { - gotV, ok := got.QueryStats[wantK] - if !ok || !testutil.Equal(gotV, wantV) { - return fmt.Errorf("Stats.QueryPlan.QueryStats: got: %+v, want: %+v", got.QueryStats, want.QueryStats) - } + // Compare all fields except DebugStats + if !testutil.Equal(want, got, cmpopts.IgnoreFields(&ExecutionStats{}, "DebugStats")) { + return fmt.Errorf("ExecutionStats: mismatch (-want +got):\n%s:", testutil.Diff(want, got, cmpopts.IgnoreFields(&ExecutionStats{}, "DebugStats"))) + } + + // Compare DebugStats + gotDebugStats := *got.DebugStats + for wantK, wantV := range *want.DebugStats { + // ExecutionStats.Debugstats has some keys whose values cannot be predicted. So, those values have not been included in want + // Here, compare only those values included in want + gotV, ok := gotDebugStats[wantK] + if !ok || !testutil.Equal(gotV, wantV) { + return fmt.Errorf("ExecutionStats.DebugStats: wantKey: %v gotValue: %+v, wantValue: %+v", wantK, gotV, wantV) } } + return nil } diff --git a/datastore/query.go b/datastore/query.go index 664fdae7a0a6..eb0fbffdefed 100644 --- a/datastore/query.go +++ b/datastore/query.go @@ -23,6 +23,7 @@ import ( "reflect" "strconv" "strings" + "time" "cloud.google.com/go/internal/protostruct" "cloud.google.com/go/internal/trace" @@ -626,36 +627,31 @@ func (c *Client) Count(ctx context.Context, q *Query) (n int, err error) { } } -const ( - // QueryModeNormal is the default mode. Only the query results are returned. - QueryModeNormal QueryMode = iota // = 0 - // QueryModeExplain returns only the query plan, without any results or execution - // statistics information. - QueryModeExplain // = 1 - // QueryModeExplainAnalyze returns both the query plan and the execution statistics along - // with the results. - QueryModeExplainAnalyze // = 2 -) - // RunOption lets the user provide options while running a query type RunOption interface { apply(*runQuerySettings) error } -// QueryMode is the mode in which the query request must be processed. -type QueryMode int32 +// ExplainOptions is explain options for the query. +type ExplainOptions struct { + // When true, the query will be planned and executed, returning the full + // query results along with both planning and execution stage metrics. + Analyze bool +} -func (m QueryMode) apply(s *runQuerySettings) error { - if s.mode != nil { - return errors.New("datastore: only one mode can be specified") +func (e ExplainOptions) apply(s *runQuerySettings) error { + if s.explainOptions != nil { + return errors.New("datastore: ExplainOptions can be specified can be specified only once") + } + pbExplainOptions := pb.ExplainOptions{ + Analyze: e.Analyze, } - pbQueryMode := pb.QueryMode(m) - s.mode = &pbQueryMode + s.explainOptions = &pbExplainOptions return nil } type runQuerySettings struct { - mode *pb.QueryMode + explainOptions *pb.ExplainOptions } // newRunQuerySettings creates a runQuerySettings with a given RunOption slice. @@ -673,21 +669,55 @@ func newRunQuerySettings(opts []RunOption) (*runQuerySettings, error) { return s, nil } -// ResultSetStats is planning and execution statistics for the query. -type ResultSetStats struct { - QueryPlan *QueryPlan - QueryStats map[string]interface{} -} - -// QueryPlan is plan for the query. -type QueryPlan struct { - PlanInfo map[string]interface{} +// Planning phase information for the query. +type Plan struct { + // The indexes selected for the query. For example: + // + // [ + // {"query_scope": "Collection", "properties": "(foo ASC, __name__ ASC)"}, + // {"query_scope": "Collection", "properties": "(bar ASC, __name__ ASC)"} + // ] + IndexesUsed []*map[string]interface{} +} + +// Execution statistics for the query. +type ExecutionStats struct { + // Total number of results returned, including documents, projections, + // aggregation results, keys. + ResultsReturned int64 + // Total number of the bytes of the results returned. + BytesReturned int64 + // Total time to execute the query in the backend. + ExecutionDuration *time.Duration + // Total billable read operations. + ReadOperations int64 + // Debugging statistics from the execution of the query. Note that the + // debugging stats are subject to change as Firestore evolves. It could + // include: + // + // { + // "indexes_entries_scanned": "1000", + // "documents_scanned": "20", + // "billing_details" : { + // "documents_billable": "20", + // "index_entries_billable": "1000", + // "min_query_cost": "0" + // } + // } + DebugStats *map[string]interface{} } // GetAllWithOptionsResult is the result of call to GetAllWithOptions method type GetAllWithOptionsResult struct { - Keys []*Key - Stats *ResultSetStats + Keys []*Key + + // Planning phase information for the query. + // This is only present when ExplainOptions is used + Plan *Plan + + // Aggregated stats from the execution of the query. + // This is only present when ExplainOptions with Analyze as true is used + ExecutionStats *ExecutionStats } // GetAll runs the provided query in the given context and returns all keys @@ -743,7 +773,8 @@ func (c *Client) GetAllWithOptions(ctx context.Context, q *Query, dst interface{ for t := c.Run(ctx, q, opts...); ; { k, e, err := t.next() - res.Stats = t.Stats + res.ExecutionStats = t.ExecutionStats + res.Plan = t.Plan if err == iterator.Done { break } @@ -821,8 +852,8 @@ func (c *Client) Run(ctx context.Context, q *Query, opts ...RunOption) *Iterator return t } - if runSettings.mode != nil { - t.req.Mode = *runSettings.mode + if runSettings.explainOptions != nil { + t.req.ExplainOptions = runSettings.explainOptions } if err := q.toRunQueryRequest(t.req); err != nil { @@ -884,8 +915,8 @@ func (c *Client) RunAggregationQueryWithOptions(ctx context.Context, aq *Aggrega if err != nil { return ar, err } - if runSettings.mode != nil { - req.Mode = *runSettings.mode + if runSettings.explainOptions != nil { + req.ExplainOptions = runSettings.explainOptions } // Parse the read options. @@ -899,7 +930,7 @@ func (c *Client) RunAggregationQueryWithOptions(ctx context.Context, aq *Aggrega return ar, err } - if req.Mode != pb.QueryMode_PLAN { + if req.ExplainOptions != nil && req.ExplainOptions.Analyze { ar.Result = make(AggregationResult) // TODO(developer): change batch parsing logic if other aggregations are supported. for _, a := range resp.Batch.AggregationResults { @@ -908,7 +939,8 @@ func (c *Client) RunAggregationQueryWithOptions(ctx context.Context, aq *Aggrega } } } - ar.Stats = fromPbResultSetStats(resp.Stats) + ar.ExecutionStats = fromPbExecutionStats(resp.ExecutionStats) + ar.Plan = fromPbPlan(resp.Plan) return ar, nil } @@ -976,8 +1008,13 @@ type Iterator struct { // entityCursor is the compiled cursor of the next result. entityCursor []byte - // Stats contains query plan and execution statistics. - Stats *ResultSetStats + // Planning phase information for the query. + // This is only present when ExplainOptions is used + Plan *Plan + + // Aggregated stats from the execution of the query. + // This is only present when ExplainOptions with Analyze as true is used + ExecutionStats *ExecutionStats // trans records the transaction in which the query was run // Currently, this value is set but unused @@ -1058,10 +1095,10 @@ func (t *Iterator) nextBatch() error { return err } - if t.req.Mode == pb.QueryMode_PLAN { + if t.req.ExplainOptions != nil && !t.req.ExplainOptions.Analyze { // No results to process t.limit = 0 - t.Stats = fromPbResultSetStats(resp.Stats) + t.Plan = fromPbPlan(resp.Plan) return nil } @@ -1104,22 +1141,45 @@ func (t *Iterator) nextBatch() error { t.pageCursor = resp.Batch.EndCursor t.results = resp.Batch.EntityResults - t.Stats = fromPbResultSetStats(resp.Stats) + t.ExecutionStats = fromPbExecutionStats(resp.ExecutionStats) + t.Plan = fromPbPlan(resp.Plan) return nil } -func fromPbResultSetStats(pbstats *pb.ResultSetStats) *ResultSetStats { +func fromPbPlan(pbplan *pb.Plan) *Plan { + if pbplan == nil { + return nil + } + + plan := &Plan{} + indexesUsed := []*map[string]interface{}{} + for _, pbIndexUsed := range pbplan.GetIndexesUsed() { + indexUsed := protostruct.DecodeToMap(pbIndexUsed) + indexesUsed = append(indexesUsed, &indexUsed) + } + + plan.IndexesUsed = indexesUsed + return plan +} + +func fromPbExecutionStats(pbstats *pb.ExecutionStats) *ExecutionStats { if pbstats == nil { return nil } - planInfo := protostruct.DecodeToMap(pbstats.QueryPlan.PlanInfo) - queryStats := protostruct.DecodeToMap(pbstats.QueryStats) - return &ResultSetStats{ - QueryPlan: &QueryPlan{ - PlanInfo: planInfo, - }, - QueryStats: queryStats, + + executionStats := &ExecutionStats{ + ResultsReturned: pbstats.GetResultsReturned(), + BytesReturned: pbstats.GetBytesReturned(), + ReadOperations: pbstats.GetReadOperations(), } + + executionDuration := pbstats.GetExecutionDuration().AsDuration() + executionStats.ExecutionDuration = &executionDuration + + debugStats := protostruct.DecodeToMap(pbstats.GetDebugStats()) + executionStats.DebugStats = &debugStats + + return executionStats } // Cursor returns a cursor for the iterator's current location. @@ -1256,5 +1316,12 @@ type AggregationResult map[string]interface{} // AggregationWithOptionsResult contains the results of an aggregation query run with options. type AggregationWithOptionsResult struct { Result AggregationResult - Stats *ResultSetStats + + // Planning phase information for the query. + // This is only present when ExplainOptions is used + Plan *Plan + + // Aggregated stats from the execution of the query. + // This is only present when ExplainOptions with Analyze as true is used + ExecutionStats *ExecutionStats } diff --git a/datastore/query_test.go b/datastore/query_test.go index 9f4e2d932279..96bb49db078f 100644 --- a/datastore/query_test.go +++ b/datastore/query_test.go @@ -860,27 +860,33 @@ func TestAggregationQueryIsNil(t *testing.T) { } } -func TestQueryModeApply(t *testing.T) { - pbNormal := pb.QueryMode(QueryModeNormal) +func TestExplainOptionsApply(t *testing.T) { + pbExplainOptions := pb.ExplainOptions{ + Analyze: true, + } for _, testcase := range []struct { - desc string - existingMode *pb.QueryMode - newMode QueryMode - wantErrMsg string + desc string + existingOptions *pb.ExplainOptions + newOptions ExplainOptions + wantErrMsg string }{ { - desc: "Multiple modes", - existingMode: &pbNormal, - newMode: QueryModeExplain, - wantErrMsg: "only one mode can be specified", + desc: "ExplainOptions specified multiple times", + existingOptions: &pbExplainOptions, + newOptions: ExplainOptions{ + Analyze: true, + }, + wantErrMsg: "ExplainOptions can be specified can be specified only once", }, { - desc: "Single mode", - existingMode: nil, - newMode: QueryModeExplain, + desc: "ExplainOptions specified once", + existingOptions: nil, + newOptions: ExplainOptions{ + Analyze: true, + }, }, } { - gotErr := testcase.newMode.apply(&runQuerySettings{mode: testcase.existingMode}) + gotErr := testcase.newOptions.apply(&runQuerySettings{explainOptions: testcase.existingOptions}) if (gotErr == nil && testcase.wantErrMsg != "") || (gotErr != nil && !strings.Contains(gotErr.Error(), testcase.wantErrMsg)) { t.Errorf("%v: apply got: %v want: %v", testcase.desc, gotErr, testcase.wantErrMsg) @@ -896,24 +902,24 @@ func TestNewRunQuerySettings(t *testing.T) { }{ { desc: "nil RunOption", - opts: []RunOption{QueryModeNormal, nil}, + opts: []RunOption{ExplainOptions{Analyze: true}, nil}, wantErrMsg: "cannot be nil", }, { desc: "success RunOption", - opts: []RunOption{QueryModeNormal}, + opts: []RunOption{ExplainOptions{Analyze: true}}, }, { - desc: "multiple modes", - opts: []RunOption{QueryModeNormal, QueryModeExplainAnalyze}, - wantErrMsg: "only one mode can be specified", + desc: "ExplainOptions specified multiple times", + opts: []RunOption{ExplainOptions{Analyze: true}, ExplainOptions{Analyze: false}, ExplainOptions{Analyze: true}}, + wantErrMsg: "ExplainOptions can be specified can be specified only once", }, } { _, gotErr := newRunQuerySettings(testcase.opts) if (gotErr == nil && testcase.wantErrMsg != "") || (gotErr != nil && !strings.Contains(gotErr.Error(), testcase.wantErrMsg)) { t.Errorf("%v: newRunQuerySettings got: %v want: %v", testcase.desc, gotErr, testcase.wantErrMsg) - } + } } } From 365c625b6df0ab36bc035d7707b600bef35ee2fe Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Tue, 13 Feb 2024 21:14:53 -0800 Subject: [PATCH 06/14] feat(datastore): Refactoring code --- datastore/integration_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datastore/integration_test.go b/datastore/integration_test.go index 20aa16d9b71b..d53d4b0129ee 100644 --- a/datastore/integration_test.go +++ b/datastore/integration_test.go @@ -1035,7 +1035,7 @@ func TestIntegration_RunAggregationQueryWithOptions(t *testing.T) { opts: []RunOption{ExplainOptions{}}, }, { - desc: "ExplainAnalyze.Analyze is true", + desc: "ExplainOptions.Analyze is true", wantRes: AggregationWithOptionsResult{ Result: wantAggResult, Plan: &Plan{ @@ -1454,7 +1454,7 @@ func getRunWithOptionsTestcases(ctx context.Context, t *testing.T, client *Clien }, }, { - desc: "ExplainAnalyze.Analyze is true", + desc: "ExplainOptions.Analyze is true", opts: []RunOption{ExplainOptions{Analyze: true}}, wantKeys: keys, wantExecutionStats: &ExecutionStats{ From 5033f507a24fa11241128c980347aa5ea6884d36 Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Tue, 13 Feb 2024 21:18:52 -0800 Subject: [PATCH 07/14] feat(datastore): Additional comments --- datastore/query.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/datastore/query.go b/datastore/query.go index eb0fbffdefed..07e80aea4bd2 100644 --- a/datastore/query.go +++ b/datastore/query.go @@ -634,6 +634,8 @@ type RunOption interface { // ExplainOptions is explain options for the query. type ExplainOptions struct { + // When false (the default), the query will be planned, returning only + // metrics from the planning stages. // When true, the query will be planned and executed, returning the full // query results along with both planning and execution stage metrics. Analyze bool From ac6083002d9c2680e6dac30d0cfc64eb0630eb59 Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Tue, 13 Feb 2024 23:42:34 -0800 Subject: [PATCH 08/14] feat(datastore): Resolving vet failures --- datastore/integration_test.go | 2 +- datastore/query.go | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/datastore/integration_test.go b/datastore/integration_test.go index d53d4b0129ee..aafe89bfe8f0 100644 --- a/datastore/integration_test.go +++ b/datastore/integration_test.go @@ -1548,7 +1548,7 @@ func isEqualExecutionStats(got *ExecutionStats, want *ExecutionStats) error { // Compare all fields except DebugStats if !testutil.Equal(want, got, cmpopts.IgnoreFields(&ExecutionStats{}, "DebugStats")) { - return fmt.Errorf("ExecutionStats: mismatch (-want +got):\n%s:", testutil.Diff(want, got, cmpopts.IgnoreFields(&ExecutionStats{}, "DebugStats"))) + return fmt.Errorf("ExecutionStats: mismatch (-want +got):\n%s", testutil.Diff(want, got, cmpopts.IgnoreFields(&ExecutionStats{}, "DebugStats"))) } // Compare DebugStats diff --git a/datastore/query.go b/datastore/query.go index 07e80aea4bd2..2bcc1012bb56 100644 --- a/datastore/query.go +++ b/datastore/query.go @@ -671,7 +671,7 @@ func newRunQuerySettings(opts []RunOption) (*runQuerySettings, error) { return s, nil } -// Planning phase information for the query. +// Plan represents planning phase information for the query. type Plan struct { // The indexes selected for the query. For example: // @@ -682,7 +682,7 @@ type Plan struct { IndexesUsed []*map[string]interface{} } -// Execution statistics for the query. +// ExecutionStats represents execution statistics for the query. type ExecutionStats struct { // Total number of results returned, including documents, projections, // aggregation results, keys. From f0ba17f6f85766e5662ed1717414c0123b4acfdd Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Mon, 18 Mar 2024 09:49:02 +0530 Subject: [PATCH 09/14] feat(datastore): Updating to match latest protos --- datastore/integration_test.go | 121 +++++++++++++++++++--------------- datastore/query.go | 80 +++++++++++----------- 2 files changed, 109 insertions(+), 92 deletions(-) diff --git a/datastore/integration_test.go b/datastore/integration_test.go index aafe89bfe8f0..27ba2313c741 100644 --- a/datastore/integration_test.go +++ b/datastore/integration_test.go @@ -1023,11 +1023,13 @@ func TestIntegration_RunAggregationQueryWithOptions(t *testing.T) { { desc: "ExplainOptions.Analyze is false", wantRes: AggregationWithOptionsResult{ - Plan: &Plan{ - IndexesUsed: []*map[string]interface{}{ - { - "properties": "(T ASC, I ASC, __name__ ASC)", - "query_scope": "Includes Ancestors", + ExplainMetrics: &ExplainMetrics{ + PlanSummary: &PlanSummary{ + IndexesUsed: []*map[string]interface{}{ + { + "properties": "(T ASC, I ASC, __name__ ASC)", + "query_scope": "Includes Ancestors", + }, }, }, }, @@ -1038,19 +1040,22 @@ func TestIntegration_RunAggregationQueryWithOptions(t *testing.T) { desc: "ExplainOptions.Analyze is true", wantRes: AggregationWithOptionsResult{ Result: wantAggResult, - Plan: &Plan{ - IndexesUsed: []*map[string]interface{}{ - { - "properties": "(T ASC, I ASC, __name__ ASC)", - "query_scope": "Includes Ancestors", + ExplainMetrics: &ExplainMetrics{ + PlanSummary: &PlanSummary{ + IndexesUsed: []*map[string]interface{}{ + { + "properties": "(T ASC, I ASC, __name__ ASC)", + "query_scope": "Includes Ancestors", + }, }, }, - }, - ExecutionStats: &ExecutionStats{ - ResultsReturned: 1, - DebugStats: &map[string]interface{}{ - "documents_scanned": "0", - "index_entries_scanned": "3", + ExecutionStats: &ExecutionStats{ + ReadOperations: 1, + ResultsReturned: 1, + DebugStats: &map[string]interface{}{ + "documents_scanned": "0", + "index_entries_scanned": "3", + }, }, }, }, @@ -1065,17 +1070,14 @@ func TestIntegration_RunAggregationQueryWithOptions(t *testing.T) { r.Errorf("err: got %v, want: nil", gotErr) } - if gotErr == nil && !reflect.DeepEqual(gotRes.Result, testcase.wantRes.Result) { + if gotErr == nil && !testutil.Equal(gotRes.Result, testcase.wantRes.Result, + cmpopts.IgnoreFields(ExplainMetrics{})) { r.Errorf("%q: Mismatch in aggregation result got: %v, want: %v", testcase.desc, gotRes, testcase.wantRes) return } - if !testutil.Equal(gotRes.Plan, testcase.wantRes.Plan) { - t.Errorf("%v Plan: got: %+v, want: %+v", testcase.desc, gotRes.Plan, testcase.wantRes.Plan) - } - - if err := isEqualExecutionStats(gotRes.ExecutionStats, testcase.wantRes.ExecutionStats); err != nil { - r.Errorf("%q: Mismatch in stats %+v", testcase.desc, err) + if err := cmpExplainMetrics(gotRes.ExplainMetrics, testcase.wantRes.ExplainMetrics); err != nil { + r.Errorf("%q: Mismatch in ExplainMetrics %+v", testcase.desc, err) } }) } @@ -1427,8 +1429,7 @@ func createTestEntities(ctx context.Context, t *testing.T, client *Client, parti type runWithOptionsTestcase struct { desc string wantKeys []*Key - wantExecutionStats *ExecutionStats - wantPlan *Plan + wantExplainMetrics *ExplainMetrics wantEntities []SQChild opts []RunOption } @@ -1444,11 +1445,13 @@ func getRunWithOptionsTestcases(ctx context.Context, t *testing.T, client *Clien { desc: "ExplainOptions.Analyze is false", opts: []RunOption{ExplainOptions{}}, - wantPlan: &Plan{ - IndexesUsed: []*map[string]interface{}{ - { - "properties": "(T ASC, I ASC, __name__ ASC)", - "query_scope": "Includes Ancestors", + wantExplainMetrics: &ExplainMetrics{ + PlanSummary: &PlanSummary{ + IndexesUsed: []*map[string]interface{}{ + { + "properties": "(T ASC, I ASC, __name__ ASC)", + "query_scope": "Includes Ancestors", + }, }, }, }, @@ -1457,18 +1460,21 @@ func getRunWithOptionsTestcases(ctx context.Context, t *testing.T, client *Clien desc: "ExplainOptions.Analyze is true", opts: []RunOption{ExplainOptions{Analyze: true}}, wantKeys: keys, - wantExecutionStats: &ExecutionStats{ - ResultsReturned: 1, - DebugStats: &map[string]interface{}{ - "documents_scanned": "3", - "index_entries_scanned": "3", + wantExplainMetrics: &ExplainMetrics{ + ExecutionStats: &ExecutionStats{ + ReadOperations: int64(count), + ResultsReturned: int64(count), + DebugStats: &map[string]interface{}{ + "documents_scanned": fmt.Sprint(count), + "index_entries_scanned": fmt.Sprint(count), + }, }, - }, - wantPlan: &Plan{ - IndexesUsed: []*map[string]interface{}{ - { - "properties": "(T ASC, I ASC, __name__ ASC)", - "query_scope": "Includes Ancestors", + PlanSummary: &PlanSummary{ + IndexesUsed: []*map[string]interface{}{ + { + "properties": "(T ASC, I ASC, __name__ ASC)", + "query_scope": "Includes Ancestors", + }, }, }, }, @@ -1496,10 +1502,7 @@ func TestIntegration_GetAllWithOptions(t *testing.T) { if !testutil.Equal(gotRes.Keys, testcase.wantKeys) { t.Errorf("%v keys: got: %+v, want: %+v", testcase.desc, gotRes.Keys, testcase.wantKeys) } - if !testutil.Equal(gotRes.Plan, testcase.wantPlan) { - t.Errorf("%v Plan: got: %+v, want: %+v", testcase.desc, gotRes.Plan, testcase.wantPlan) - } - if err := isEqualExecutionStats(gotRes.ExecutionStats, testcase.wantExecutionStats); err != nil { + if err := cmpExplainMetrics(gotRes.ExplainMetrics, testcase.wantExplainMetrics); err != nil { t.Errorf("%v %+v", testcase.desc, err) } } @@ -1529,16 +1532,30 @@ func TestIntegration_RunWithOptions(t *testing.T) { if !testutil.Equal(gotSQChildsFromRun, testcase.wantEntities) { t.Errorf("%v entities: got: %+v, want: %+v", testcase.desc, gotSQChildsFromRun, testcase.wantEntities) } - if !testutil.Equal(iter.Plan, testcase.wantPlan) { - t.Errorf("%v Plan: got: %+v, want: %+v", testcase.desc, iter.Plan, testcase.wantPlan) - } - if err := isEqualExecutionStats(iter.ExecutionStats, testcase.wantExecutionStats); err != nil { + + if err := cmpExplainMetrics(iter.ExplainMetrics, testcase.wantExplainMetrics); err != nil { t.Errorf("%v %+v", testcase.desc, err) } } } -func isEqualExecutionStats(got *ExecutionStats, want *ExecutionStats) error { +func cmpExplainMetrics(got *ExplainMetrics, want *ExplainMetrics) error { + if (got != nil && want == nil) || (got == nil && want != nil) { + return fmt.Errorf("ExplainMetrics: got: %+v, want: %+v", got, want) + } + if got == nil { + return nil + } + if !testutil.Equal(got.PlanSummary, want.PlanSummary) { + return fmt.Errorf("Plan: got: %+v, want: %+v", got.PlanSummary, want.PlanSummary) + } + if err := cmpExecutionStats(got.ExecutionStats, want.ExecutionStats); err != nil { + return err + } + return nil +} + +func cmpExecutionStats(got *ExecutionStats, want *ExecutionStats) error { if (got != nil && want == nil) || (got == nil && want != nil) { return fmt.Errorf("ExecutionStats: got: %+v, want: %+v", got, want) } @@ -1547,8 +1564,8 @@ func isEqualExecutionStats(got *ExecutionStats, want *ExecutionStats) error { } // Compare all fields except DebugStats - if !testutil.Equal(want, got, cmpopts.IgnoreFields(&ExecutionStats{}, "DebugStats")) { - return fmt.Errorf("ExecutionStats: mismatch (-want +got):\n%s", testutil.Diff(want, got, cmpopts.IgnoreFields(&ExecutionStats{}, "DebugStats"))) + if !testutil.Equal(want, got, cmpopts.IgnoreFields(ExecutionStats{}, "DebugStats", "ExecutionDuration")) { + return fmt.Errorf("ExecutionStats: mismatch (-want +got):\n%s", testutil.Diff(want, got, cmpopts.IgnoreFields(ExecutionStats{}, "DebugStats"))) } // Compare DebugStats diff --git a/datastore/query.go b/datastore/query.go index 105b57e58c20..f91fee164063 100644 --- a/datastore/query.go +++ b/datastore/query.go @@ -673,8 +673,18 @@ func newRunQuerySettings(opts []RunOption) (*runQuerySettings, error) { return s, nil } -// Plan represents planning phase information for the query. -type Plan struct { +// ExplainMetrics for the query. +type ExplainMetrics struct { + // Planning phase information for the query. + PlanSummary *PlanSummary + + // Aggregated stats from the execution of the query. Only present when + // ExplainOptions.Analyze is set to true. + ExecutionStats *ExecutionStats +} + +// PlanSummary represents planning phase information for the query. +type PlanSummary struct { // The indexes selected for the query. For example: // // [ @@ -689,8 +699,6 @@ type ExecutionStats struct { // Total number of results returned, including documents, projections, // aggregation results, keys. ResultsReturned int64 - // Total number of the bytes of the results returned. - BytesReturned int64 // Total time to execute the query in the backend. ExecutionDuration *time.Duration // Total billable read operations. @@ -715,13 +723,8 @@ type ExecutionStats struct { type GetAllWithOptionsResult struct { Keys []*Key - // Planning phase information for the query. - // This is only present when ExplainOptions is used - Plan *Plan - - // Aggregated stats from the execution of the query. - // This is only present when ExplainOptions with Analyze as true is used - ExecutionStats *ExecutionStats + // Query explain metrics. This is only present when ExplainOptions is provided. + ExplainMetrics *ExplainMetrics } // GetAll runs the provided query in the given context and returns all keys @@ -777,8 +780,7 @@ func (c *Client) GetAllWithOptions(ctx context.Context, q *Query, dst interface{ for t := c.Run(ctx, q, opts...); ; { k, e, err := t.next() - res.ExecutionStats = t.ExecutionStats - res.Plan = t.Plan + res.ExplainMetrics = t.ExplainMetrics if err == iterator.Done { break } @@ -934,7 +936,7 @@ func (c *Client) RunAggregationQueryWithOptions(ctx context.Context, aq *Aggrega return ar, err } - if req.ExplainOptions != nil && req.ExplainOptions.Analyze { + if req.ExplainOptions == nil || req.ExplainOptions.Analyze { ar.Result = make(AggregationResult) // TODO(developer): change batch parsing logic if other aggregations are supported. for _, a := range resp.Batch.AggregationResults { @@ -943,9 +945,8 @@ func (c *Client) RunAggregationQueryWithOptions(ctx context.Context, aq *Aggrega } } } - ar.ExecutionStats = fromPbExecutionStats(resp.ExecutionStats) - ar.Plan = fromPbPlan(resp.Plan) + ar.ExplainMetrics = fromPbExplainMetrics(resp.GetExplainMetrics()) return ar, nil } @@ -1012,13 +1013,8 @@ type Iterator struct { // entityCursor is the compiled cursor of the next result. entityCursor []byte - // Planning phase information for the query. - // This is only present when ExplainOptions is used - Plan *Plan - - // Aggregated stats from the execution of the query. - // This is only present when ExplainOptions with Analyze as true is used - ExecutionStats *ExecutionStats + // Query explain metrics. This is only present when ExplainOptions is used. + ExplainMetrics *ExplainMetrics // trans records the transaction in which the query was run // Currently, this value is set but unused @@ -1102,7 +1098,7 @@ func (t *Iterator) nextBatch() error { if t.req.ExplainOptions != nil && !t.req.ExplainOptions.Analyze { // No results to process t.limit = 0 - t.Plan = fromPbPlan(resp.Plan) + t.ExplainMetrics = fromPbExplainMetrics(resp.GetExplainMetrics()) return nil } @@ -1145,25 +1141,35 @@ func (t *Iterator) nextBatch() error { t.pageCursor = resp.Batch.EndCursor t.results = resp.Batch.EntityResults - t.ExecutionStats = fromPbExecutionStats(resp.ExecutionStats) - t.Plan = fromPbPlan(resp.Plan) + t.ExplainMetrics = fromPbExplainMetrics(resp.GetExplainMetrics()) return nil } -func fromPbPlan(pbplan *pb.Plan) *Plan { - if pbplan == nil { +func fromPbExplainMetrics(pbExplainMetrics *pb.ExplainMetrics) *ExplainMetrics { + if pbExplainMetrics == nil { return nil } + explainMetrics := &ExplainMetrics{ + PlanSummary: fromPbPlanSummary(pbExplainMetrics.PlanSummary), + ExecutionStats: fromPbExecutionStats(pbExplainMetrics.ExecutionStats), + } + return explainMetrics +} - plan := &Plan{} +func fromPbPlanSummary(pbPlanSummary *pb.PlanSummary) *PlanSummary { + if pbPlanSummary == nil { + return nil + } + + planSummary := &PlanSummary{} indexesUsed := []*map[string]interface{}{} - for _, pbIndexUsed := range pbplan.GetIndexesUsed() { + for _, pbIndexUsed := range pbPlanSummary.GetIndexesUsed() { indexUsed := protostruct.DecodeToMap(pbIndexUsed) indexesUsed = append(indexesUsed, &indexUsed) } - plan.IndexesUsed = indexesUsed - return plan + planSummary.IndexesUsed = indexesUsed + return planSummary } func fromPbExecutionStats(pbstats *pb.ExecutionStats) *ExecutionStats { @@ -1173,7 +1179,6 @@ func fromPbExecutionStats(pbstats *pb.ExecutionStats) *ExecutionStats { executionStats := &ExecutionStats{ ResultsReturned: pbstats.GetResultsReturned(), - BytesReturned: pbstats.GetBytesReturned(), ReadOperations: pbstats.GetReadOperations(), } @@ -1321,11 +1326,6 @@ type AggregationResult map[string]interface{} type AggregationWithOptionsResult struct { Result AggregationResult - // Planning phase information for the query. - // This is only present when ExplainOptions is used - Plan *Plan - - // Aggregated stats from the execution of the query. - // This is only present when ExplainOptions with Analyze as true is used - ExecutionStats *ExecutionStats + // Query explain metrics. This is only present when ExplainOptions is provided. + ExplainMetrics *ExplainMetrics } From 54e092d11091f0f77ff17eb55d4da818393163e5 Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Tue, 9 Apr 2024 12:15:48 -0700 Subject: [PATCH 10/14] feat(datastore): Add RunWithOptions --- datastore/integration_test.go | 2 +- datastore/query.go | 9 +++++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/datastore/integration_test.go b/datastore/integration_test.go index 27ba2313c741..da5eda2cf940 100644 --- a/datastore/integration_test.go +++ b/datastore/integration_test.go @@ -1517,7 +1517,7 @@ func TestIntegration_RunWithOptions(t *testing.T) { query := NewQuery("SQChild").Ancestor(parent).Filter("T=", now).Order("I") for _, testcase := range testcases { var gotSQChildsFromRun []SQChild - iter := client.Run(ctx, query, testcase.opts...) + iter := client.RunWithOptions(ctx, query, testcase.opts...) for { var gotSQChild SQChild _, err := iter.Next(&gotSQChild) diff --git a/datastore/query.go b/datastore/query.go index f91fee164063..0705ca3f466f 100644 --- a/datastore/query.go +++ b/datastore/query.go @@ -778,7 +778,7 @@ func (c *Client) GetAllWithOptions(ctx context.Context, q *Query, dst interface{ } } - for t := c.Run(ctx, q, opts...); ; { + for t := c.RunWithOptions(ctx, q, opts...); ; { k, e, err := t.next() res.ExplainMetrics = t.ExplainMetrics if err == iterator.Done { @@ -825,8 +825,13 @@ func (c *Client) GetAllWithOptions(ctx context.Context, q *Query, dst interface{ return res, errFieldMismatch } +// Run runs the given query in the given context +func (c *Client) Run(ctx context.Context, q *Query) *Iterator { + return c.RunWithOptions(ctx, q) +} + // Run runs the given query in the given context with the provided options -func (c *Client) Run(ctx context.Context, q *Query, opts ...RunOption) *Iterator { +func (c *Client) RunWithOptions(ctx context.Context, q *Query, opts ...RunOption) *Iterator { if q.err != nil { return &Iterator{err: q.err} } From 6f9f76a7f305cd78a8a5cd6dc4914f721b0326b9 Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Thu, 18 Apr 2024 00:36:11 +0000 Subject: [PATCH 11/14] fix(datastore): Resolving vet failures --- datastore/query.go | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/datastore/query.go b/datastore/query.go index 22f0c0c8ef3a..10ace79f04a2 100644 --- a/datastore/query.go +++ b/datastore/query.go @@ -825,19 +825,18 @@ func (c *Client) GetAllWithOptions(ctx context.Context, q *Query, dst interface{ return res, errFieldMismatch } - // Run runs the given query in the given context -func (c *Client) Run(ctx context.Context, q *Query) *Iterator { - ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.Query.Run") +func (c *Client) Run(ctx context.Context, q *Query) (it *Iterator) { + ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.Query.Run") defer func() { trace.EndSpan(ctx, it.err) }() return c.run(ctx, q) } // RunWithOptions runs the given query in the given context with the provided options -func (c *Client) RunWithOptions(ctx context.Context, q *Query, opts ...RunOption) *Iterator { +func (c *Client) RunWithOptions(ctx context.Context, q *Query, opts ...RunOption) (it *Iterator) { ctx = trace.StartSpan(ctx, "cloud.google.com/go/datastore.Query.RunWithOptions") defer func() { trace.EndSpan(ctx, it.err) }() - return c.run(ctx, q, opts) + return c.run(ctx, q, opts...) } // run runs the given query in the given context with the provided options From 6d23653ec93a57f45a30ff1921664cbb10c09c06 Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Thu, 18 Apr 2024 17:59:11 +0000 Subject: [PATCH 12/14] doc(datastore): Adde comment for preview feature --- datastore/query.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/datastore/query.go b/datastore/query.go index 255a583f743f..c4a81bcc8e0b 100644 --- a/datastore/query.go +++ b/datastore/query.go @@ -635,6 +635,9 @@ type RunOption interface { } // ExplainOptions is explain options for the query. +// +// Query Explain feature is still in preview and not yet publicly available. +// Pre-GA features might have limited support and can change at any time. type ExplainOptions struct { // When false (the default), the query will be planned, returning only // metrics from the planning stages. From 700e3a1e2e12391bdcb0659c5a08231472b44233 Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Mon, 22 Apr 2024 20:07:15 +0000 Subject: [PATCH 13/14] feat(datastore): Correcting the error message --- datastore/query.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datastore/query.go b/datastore/query.go index c4a81bcc8e0b..e701d9d42bb3 100644 --- a/datastore/query.go +++ b/datastore/query.go @@ -648,7 +648,7 @@ type ExplainOptions struct { func (e ExplainOptions) apply(s *runQuerySettings) error { if s.explainOptions != nil { - return errors.New("datastore: ExplainOptions can be specified can be specified only once") + return errors.New("datastore: ExplainOptions can be specified only once") } pbExplainOptions := pb.ExplainOptions{ Analyze: e.Analyze, From 4730d3eefb471f2459a31cc534f52d047273f898 Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Tue, 7 May 2024 12:21:14 +0000 Subject: [PATCH 14/14] feat(datastore): Resolving test failures --- datastore/integration_test.go | 8 ++++---- datastore/query.go | 10 +++++----- datastore/query_test.go | 6 +++--- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/datastore/integration_test.go b/datastore/integration_test.go index 839bc2fbf2a2..318f2a4d0cdc 100644 --- a/datastore/integration_test.go +++ b/datastore/integration_test.go @@ -1278,7 +1278,7 @@ func TestIntegration_RunAggregationQueryWithOptions(t *testing.T) { IndexesUsed: []*map[string]interface{}{ { "properties": "(T ASC, I ASC, __name__ ASC)", - "query_scope": "Includes Ancestors", + "query_scope": "Includes ancestors", }, }, }, @@ -1295,7 +1295,7 @@ func TestIntegration_RunAggregationQueryWithOptions(t *testing.T) { IndexesUsed: []*map[string]interface{}{ { "properties": "(T ASC, I ASC, __name__ ASC)", - "query_scope": "Includes Ancestors", + "query_scope": "Includes ancestors", }, }, }, @@ -1700,7 +1700,7 @@ func getRunWithOptionsTestcases(ctx context.Context, t *testing.T, client *Clien IndexesUsed: []*map[string]interface{}{ { "properties": "(T ASC, I ASC, __name__ ASC)", - "query_scope": "Includes Ancestors", + "query_scope": "Includes ancestors", }, }, }, @@ -1723,7 +1723,7 @@ func getRunWithOptionsTestcases(ctx context.Context, t *testing.T, client *Clien IndexesUsed: []*map[string]interface{}{ { "properties": "(T ASC, I ASC, __name__ ASC)", - "query_scope": "Includes Ancestors", + "query_scope": "Includes ancestors", }, }, }, diff --git a/datastore/query.go b/datastore/query.go index ea17a5979c50..d4dec92a8159 100644 --- a/datastore/query.go +++ b/datastore/query.go @@ -958,8 +958,8 @@ func (c *Client) RunAggregationQueryWithOptions(ctx context.Context, aq *Aggrega return ar, err } - if txn != nil && txn.state == transactionStateNotStarted { - txn.setToInProgress(res.Transaction) + if txn != nil && txn.state == transactionStateNotStarted { + txn.setToInProgress(resp.Transaction) } if req.ExplainOptions == nil || req.ExplainOptions.Analyze { @@ -1131,13 +1131,13 @@ func (t *Iterator) nextBatch() error { if txn != nil && txn.state == transactionStateNotStarted { txn.setToInProgress(resp.Transaction) } - - if t.req.ExplainOptions != nil && !t.req.ExplainOptions.Analyze { + + if t.req.ExplainOptions != nil && !t.req.ExplainOptions.Analyze { // No results to process t.limit = 0 t.ExplainMetrics = fromPbExplainMetrics(resp.GetExplainMetrics()) return nil - } + } // Adjust any offset from skipped results. skip := resp.Batch.SkippedResults diff --git a/datastore/query_test.go b/datastore/query_test.go index f5531fd1b796..9c6d60ab1162 100644 --- a/datastore/query_test.go +++ b/datastore/query_test.go @@ -902,7 +902,7 @@ func TestExplainOptionsApply(t *testing.T) { newOptions: ExplainOptions{ Analyze: true, }, - wantErrMsg: "ExplainOptions can be specified can be specified only once", + wantErrMsg: "ExplainOptions can be specified only once", }, { desc: "ExplainOptions specified once", @@ -915,7 +915,7 @@ func TestExplainOptionsApply(t *testing.T) { gotErr := testcase.newOptions.apply(&runQuerySettings{explainOptions: testcase.existingOptions}) if (gotErr == nil && testcase.wantErrMsg != "") || (gotErr != nil && !strings.Contains(gotErr.Error(), testcase.wantErrMsg)) { - t.Errorf("%v: apply got: %v want: %v", testcase.desc, gotErr, testcase.wantErrMsg) + t.Errorf("%v: apply got: %v want: %v", testcase.desc, gotErr.Error(), testcase.wantErrMsg) } } } @@ -938,7 +938,7 @@ func TestNewRunQuerySettings(t *testing.T) { { desc: "ExplainOptions specified multiple times", opts: []RunOption{ExplainOptions{Analyze: true}, ExplainOptions{Analyze: false}, ExplainOptions{Analyze: true}}, - wantErrMsg: "ExplainOptions can be specified can be specified only once", + wantErrMsg: "ExplainOptions can be specified only once", }, } { _, gotErr := newRunQuerySettings(testcase.opts)