Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(datastore): Add query profiling #9200

Merged
merged 23 commits into from
May 7, 2024
Merged
Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
730d947
feat(datastore): Query profiling
bhshkh Dec 18, 2023
5825267
feat(datastore): remove comment
bhshkh Dec 18, 2023
bf802dc
Merge branch 'main' into feature/ds-query-profiling
bhshkh Dec 27, 2023
4d8a1f1
feat(datastore): Refactor test
bhshkh Dec 27, 2023
1f507c6
Merge branch 'main' into feature/ds-query-profiling
bhshkh Jan 10, 2024
4276f61
feat(datastore): Refactor integration test
bhshkh Jan 17, 2024
8420179
Merge branch 'main' into feature/ds-query-profiling
bhshkh Feb 13, 2024
21923dc
feat(datastore): Use ExplainOptions instead of QueryMode
bhshkh Feb 14, 2024
365c625
feat(datastore): Refactoring code
bhshkh Feb 14, 2024
5033f50
feat(datastore): Additional comments
bhshkh Feb 14, 2024
ac60830
feat(datastore): Resolving vet failures
bhshkh Feb 14, 2024
252f0e9
Merge branch 'googleapis:main' into feature/ds-query-profiling
bhshkh Mar 18, 2024
f0ba17f
feat(datastore): Updating to match latest protos
bhshkh Mar 18, 2024
edb28f2
Merge branch 'main' into feature/ds-query-profiling
bhshkh Apr 8, 2024
54e092d
feat(datastore): Add RunWithOptions
bhshkh Apr 9, 2024
1cfdc14
Merge branch 'main' into feature/ds-query-profiling
bhshkh Apr 17, 2024
6f9f76a
fix(datastore): Resolving vet failures
bhshkh Apr 18, 2024
cf4823b
Merge branch 'main' into feature/ds-query-profiling
bhshkh Apr 18, 2024
0ceff9f
Merge branch 'main' into feature/ds-query-profiling
bhshkh Apr 18, 2024
6d23653
doc(datastore): Adde comment for preview feature
bhshkh Apr 18, 2024
700e3a1
feat(datastore): Correcting the error message
bhshkh Apr 22, 2024
144e73e
Merge branch 'main' into feature/ds-query-profiling
bhshkh May 7, 2024
4730d3e
feat(datastore): Resolving test failures
bhshkh May 7, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
263 changes: 263 additions & 0 deletions datastore/integration_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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)
Expand Down
Loading
Loading