feat(datastore): Add query profiling

datastore/integration_test.go

@@ -990,6 +990,104 @@ 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(ctx, t, 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{}{
+						"documents_scanned":     "0",
+						"index_entries_scanned": "3",
+						"results_returned":      "1",
+					},
+				},
+			},
+			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 +1389,171 @@ func TestIntegration_GetAllWithFieldMismatch(t *testing.T) {
 	}
 }
 
+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()
+
+	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)
+		}
+	}
+}
+
+type runWithOptionsTestcase 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,
+			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{}{
+					"documents_scanned":     "3",
+					"index_entries_scanned": "3",
+					"results_returned":      "3",
+				},
+			},
+			wantEntities: entities,
+		},
+	}, 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 {
+			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)
+		}
+	}
+}
+
+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 {
+			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 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)
+		}
+
+		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)
+			}
+		}
+	}
+	return nil
+}
+
 func TestIntegration_KindlessQueries(t *testing.T) {
 	ctx := context.Background()
 	client := newTestClient(ctx, t)