From 58052a2eefd56b3129e04f177398b3ffb688d4d7 Mon Sep 17 00:00:00 2001 From: Baha Aiman Date: Thu, 29 Aug 2024 13:36:30 -0700 Subject: [PATCH] feat(firestore): Query profiling (#10164) --- firestore/integration_test.go | 365 +++++++++++++++- firestore/options.go | 49 +++ firestore/query.go | 231 +++++++++- firestore/query_test.go | 793 ++++++++++++++++++++++++---------- firestore/transaction.go | 1 + 5 files changed, 1185 insertions(+), 254 deletions(-) diff --git a/firestore/integration_test.go b/firestore/integration_test.go index 36ec9d2aaadc..b7fe0ab6e4c4 100644 --- a/firestore/integration_test.go +++ b/firestore/integration_test.go @@ -72,6 +72,7 @@ const ( envProjID = "GCLOUD_TESTS_GOLANG_FIRESTORE_PROJECT_ID" envPrivateKey = "GCLOUD_TESTS_GOLANG_FIRESTORE_KEY" envDatabases = "GCLOUD_TESTS_GOLANG_FIRESTORE_DATABASES" + envEmulator = "FIRESTORE_EMULATOR_HOST" ) var ( @@ -82,6 +83,7 @@ var ( wantDBPath string testParams map[string]interface{} seededFirstIndex bool + useEmulator bool ) func initIntegrationTest() { @@ -91,6 +93,9 @@ func initIntegrationTest() { if testing.Short() { return } + if addr := os.Getenv(envEmulator); addr != "" { + useEmulator = true + } ctx := context.Background() testProjectID := os.Getenv(envProjID) if testProjectID == "" { @@ -559,6 +564,163 @@ func TestIntegration_GetAll(t *testing.T) { }) } +type runWithOptionsTestcase struct { + desc string + wantExplainMetrics *ExplainMetrics + wantSnapshots bool + opts []RunOption +} + +func getRunWithOptionsTestcases(t *testing.T) ([]runWithOptionsTestcase, []*DocumentRef) { + type getAll struct{ N int } + + count := 5 + h := testHelper{t} + coll := integrationColl(t) + var wantDocRefs []*DocumentRef + for i := 0; i < count; i++ { + doc := coll.Doc("getRunWithOptionsTestcases" + fmt.Sprint(i)) + wantDocRefs = append(wantDocRefs, doc) + h.mustCreate(doc, getAll{N: i}) + } + + wantPlanSummary := &PlanSummary{ + IndexesUsed: []*map[string]interface{}{ + { + "properties": "(__name__ ASC)", + "query_scope": "Collection", + }, + }, + } + return []runWithOptionsTestcase{ + { + desc: "No ExplainOptions", + wantSnapshots: true, + }, + + { + desc: "ExplainOptions.Analyze is false", + opts: []RunOption{ExplainOptions{}}, + wantExplainMetrics: &ExplainMetrics{ + PlanSummary: wantPlanSummary, + }, + }, + { + desc: "ExplainOptions.Analyze is true", + opts: []RunOption{ExplainOptions{Analyze: true}}, + 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), + }, + }, + PlanSummary: wantPlanSummary, + }, + wantSnapshots: true, + }, + }, wantDocRefs +} + +func TestIntegration_GetAll_WithRunOptions(t *testing.T) { + if useEmulator { + t.Skip("Skipping. Query profiling not supported in emulator.") + } + coll := integrationColl(t) + ctx := context.Background() + testcases, wantDocRefs := getRunWithOptionsTestcases(t) + snapshotRefIDs := []string{} + for _, wantRef := range wantDocRefs { + snapshotRefIDs = append(snapshotRefIDs, wantRef.ID) + } + + defer func() { + deleteDocuments(wantDocRefs) + }() + + for _, testcase := range testcases { + t.Run(testcase.desc, func(t *testing.T) { + docIter := coll.WithRunOptions(testcase.opts...).Documents(ctx) + gotDocSnaps, gotErr := docIter.GetAll() + if gotErr != nil { + t.Fatalf("err: got: %+v, want: nil", gotErr) + } + + gotExpM, gotExpMErr := docIter.ExplainMetrics() + if gotExpMErr != nil { + t.Fatalf("ExplainMetrics() err: got: %+v, want: nil", gotExpMErr) + } + + gotIDs := []string{} + for _, gotSnapshot := range gotDocSnaps { + gotIDs = append(gotIDs, gotSnapshot.Ref.ID) + } + + if testcase.wantSnapshots && !testutil.Equal(gotIDs, snapshotRefIDs) { + t.Errorf("snapshots ID: got: %+v, want: %+v", gotIDs, snapshotRefIDs) + } + if !testcase.wantSnapshots && len(gotIDs) != 0 { + t.Errorf("snapshots ID: got: %+v, want: %+v", gotIDs, nil) + } + + if err := cmpExplainMetrics(gotExpM, testcase.wantExplainMetrics); err != nil { + t.Error(err) + } + + if gotExpM != nil && gotExpM.PlanSummary != nil && len(gotExpM.PlanSummary.IndexesUsed) != 0 { + indexesUsed := *gotExpM.PlanSummary.IndexesUsed[0] + fmt.Printf("type=%T\n", indexesUsed["properties"]) + } + }) + } +} + +func TestIntegration_Query_WithRunOptions(t *testing.T) { + if useEmulator { + t.Skip("Skipping. Query profiling not supported in emulator.") + } + coll := integrationColl(t) + ctx := context.Background() + testcases, wantDocRefs := getRunWithOptionsTestcases(t) + snapshotRefIDs := []string{} + for _, wantRef := range wantDocRefs { + snapshotRefIDs = append(snapshotRefIDs, wantRef.ID) + } + + defer func() { + deleteDocuments(wantDocRefs) + }() + + for _, testcase := range testcases { + gotIDs := []string{} + gotDocIter := coll.WithRunOptions(testcase.opts...).Documents(ctx) + for { + gotDocSnap, err := gotDocIter.Next() + if err == iterator.Done { + break + } + if err != nil { + t.Fatalf("%v: Failed to get next document: %+v\n", testcase.desc, err) + } + gotIDs = append(gotIDs, gotDocSnap.Ref.ID) + } + + if (testcase.wantSnapshots && !testutil.Equal(gotIDs, snapshotRefIDs)) || (!testcase.wantSnapshots && len(gotIDs) != 0) { + t.Errorf("%v: snapshots ID: got: %+v, want: %+v", testcase.desc, gotIDs, snapshotRefIDs) + } + + gotExp, gotExpErr := gotDocIter.ExplainMetrics() + if gotExpErr != nil { + t.Fatalf("%v: Failed to get explain metrics: %+v\n", testcase.desc, gotExpErr) + } + if err := cmpExplainMetrics(gotExp, testcase.wantExplainMetrics); err != nil { + t.Errorf("%v: %+v", testcase.desc, err) + } + + } +} func TestIntegration_Add(t *testing.T) { start := time.Now() docRef, wr, err := integrationColl(t).Add(context.Background(), integrationTestMap) @@ -924,7 +1086,8 @@ func TestIntegration_QueryDocuments_WhereEntity(t *testing.T) { indexFields := [][]string{ {"updatedAt", "weight", "height"}, - {"weight", "height"}} + {"weight", "height"}, + } adminCtx, cancel := context.WithTimeout(context.Background(), 15*time.Minute) defer cancel() indexNames := createIndexes(adminCtx, wantDBPath, indexFields) @@ -1396,6 +1559,59 @@ func TestIntegration_RunTransaction(t *testing.T) { }) } +func TestIntegration_RunTransaction_WithRunOptions(t *testing.T) { + if useEmulator { + t.Skip("Skipping. Query profiling not supported in emulator.") + } + ctx := context.Background() + client := integrationClient(t) + testcases, wantDocRefs := getRunWithOptionsTestcases(t) + numDocs := len(wantDocRefs) + for _, testcase := range testcases { + t.Run(testcase.desc, func(t *testing.T) { + err := client.RunTransaction(ctx, func(_ context.Context, tx *Transaction) error { + docIter := tx.Documents(iColl.WithRunOptions(testcase.opts...)) + docsRead := 0 + for { + _, err := docIter.Next() + if err == iterator.Done { + break + } + if err != nil { + return fmt.Errorf("Next got %+v, want %+v", err, nil) + } + + docsRead++ + + // There are documents available in the iterator, + // error should be received + _, gotExpErr := docIter.ExplainMetrics() + if docsRead < numDocs && (gotExpErr == nil || !strings.Contains(errMetricsBeforeEnd.Error(), gotExpErr.Error())) { + fmt.Printf("Error thrown from here %v %v\n", gotExpErr == nil, strings.Contains(errMetricsBeforeEnd.Error(), gotExpErr.Error())) + return fmt.Errorf("ExplainMetrics got %+v, want %+v", gotExpErr, errMetricsBeforeEnd) + } + } + + gotExp, gotExpErr := docIter.ExplainMetrics() + if gotExpErr != nil { + return fmt.Errorf("ExplainMetrics got %+v, want %+v", gotExpErr, nil) + } + if err := cmpExplainMetrics(gotExp, testcase.wantExplainMetrics); err != nil { + return fmt.Errorf("ExplainMetrics %+v", err) + } + return nil + }) + if err != nil { + t.Fatal(err) + } + }) + } + + t.Cleanup(func() { + deleteDocuments(wantDocRefs) + }) +} + func TestIntegration_TransactionGetAll(t *testing.T) { ctx := context.Background() h := testHelper{t} @@ -2739,6 +2955,153 @@ func TestIntegration_AggregationQueries(t *testing.T) { } } +func TestIntegration_AggregationQueries_WithRunOptions(t *testing.T) { + if useEmulator { + t.Skip("Skipping. Query profiling not supported in emulator.") + } + ctx := context.Background() + coll := integrationColl(t) + + h := testHelper{t} + docs := []map[string]interface{}{ + {"weight": 0.5, "height": 99, "model": "A"}, + {"weight": 0.5, "height": 98, "model": "A"}, + {"weight": 0.5, "height": 97, "model": "B"}, + } + for _, doc := range docs { + newDoc := coll.NewDoc() + h.mustCreate(newDoc, doc) + } + + aggResult := map[string]interface{}{ + "count1": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: int64(3)}}, + "weight_sum1": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: float64(1.5)}}, + "weight_avg1": &pb.Value{ValueType: &pb.Value_DoubleValue{DoubleValue: float64(0.5)}}, + } + wantPlanSummary := &PlanSummary{ + IndexesUsed: []*map[string]interface{}{ + { + "properties": "(weight ASC, __name__ ASC)", + "query_scope": "Collection", + }, + }, + } + + testcases := []struct { + desc string + wantRes *AggregationResponse + wantErrMsg string + query Query + }{ + { + desc: "no options", + query: coll.Where("weight", "<=", 1), + wantRes: &AggregationResponse{ + Result: aggResult, + }, + }, + { + desc: "ExplainOptions.Analyze is false", + query: coll.Where("weight", "<=", 1).WithRunOptions(ExplainOptions{Analyze: false}), + wantRes: &AggregationResponse{ + ExplainMetrics: &ExplainMetrics{ + PlanSummary: wantPlanSummary, + }, + }, + }, + { + desc: "ExplainOptions.Analyze is true", + query: coll.Where("weight", "<=", 1).WithRunOptions(ExplainOptions{Analyze: true}), + wantRes: &AggregationResponse{ + Result: aggResult, + ExplainMetrics: &ExplainMetrics{ + ExecutionStats: &ExecutionStats{ + ReadOperations: int64(1), + ResultsReturned: int64(1), + DebugStats: &map[string]interface{}{ + "documents_scanned": fmt.Sprint(0), + "index_entries_scanned": fmt.Sprint(3), + }, + }, + PlanSummary: wantPlanSummary, + }, + }, + }, + } + + for _, testcase := range testcases { + testutil.Retry(t, 10, time.Second, func(r *testutil.R) { + aq := testcase.query.NewAggregationQuery().WithCount("count1"). + WithAvg("weight", "weight_avg1"). + WithSum("weight", "weight_sum1") + gotRes, gotErr := aq.GetResponse(ctx) + + gotErrMsg := "" + if gotErr != nil { + gotErrMsg = gotErr.Error() + } + + gotFailed := gotErr != nil + wantFailed := len(testcase.wantErrMsg) != 0 + if gotFailed != wantFailed || !strings.Contains(gotErrMsg, testcase.wantErrMsg) { + r.Errorf("%s: Mismatch in error got: %v, want: %v", testcase.desc, gotErr, testcase.wantErrMsg) + return + } + if !gotFailed && !testutil.Equal(gotRes.Result, testcase.wantRes.Result) { + r.Errorf("%q: Mismatch in aggregation result got: %v, want: %v", testcase.desc, gotRes.Result, testcase.wantRes.Result) + return + } + + if err := cmpExplainMetrics(gotRes.ExplainMetrics, testcase.wantRes.ExplainMetrics); err != nil { + r.Errorf("%q: Mismatch in ExplainMetrics %+v", testcase.desc, err) + } + }) + } +} + +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("PlanSummary diff (-want +got): %+v", testutil.Diff(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) + } + if got == nil { + return nil + } + + // Compare all fields except 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", "ExecutionDuration"))) + } + + // 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 +} + func TestIntegration_CountAggregationQuery(t *testing.T) { str := uid.NewSpace("firestore-count", &uid.Options{}) datum := str.New() diff --git a/firestore/options.go b/firestore/options.go index a1b642ab876e..1c78b5c0220b 100644 --- a/firestore/options.go +++ b/firestore/options.go @@ -170,3 +170,52 @@ func processSetOptions(opts []SetOption) (fps []FieldPath, all bool, err error) return nil, false, fmt.Errorf("conflicting options: %+v", opts) } } + +type runQuerySettings struct { + // Explain options for the query. If set, additional query + // statistics will be returned. If not, only query results will be returned. + explainOptions *pb.ExplainOptions +} + +// 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("firestore: RunOption cannot be nil") + } + err := o.apply(s) + if err != nil { + return nil, err + } + } + return s, nil +} + +// RunOption are options used while running a query +type RunOption interface { + apply(*runQuerySettings) error +} + +// ExplainOptions are Query Explain options. +// Query Explain allows you to submit Cloud Firestore queries to the backend and +// receive detailed performance statistics on backend query execution in return. +type ExplainOptions struct { + // When false (the default), Query Explain plans the query, but skips over the + // execution stage. This will return planner stage information. + // + // When true, Query Explain both plans and executes the query. This returns all + // the planner information along with statistics from the query execution runtime. + // This will include the billing information of the query along with system-level + // insights into the query execution. + Analyze bool +} + +func (e ExplainOptions) apply(s *runQuerySettings) error { + if s.explainOptions != nil { + return errors.New("firestore: ExplainOptions can be specified only once") + } + pbExplainOptions := pb.ExplainOptions{Analyze: e.Analyze} + s.explainOptions = &pbExplainOptions + return nil +} diff --git a/firestore/query.go b/firestore/query.go index dbcd4fbab856..9e77e2a4f6ae 100644 --- a/firestore/query.go +++ b/firestore/query.go @@ -25,12 +25,17 @@ import ( pb "cloud.google.com/go/firestore/apiv1/firestorepb" "cloud.google.com/go/internal/btree" + "cloud.google.com/go/internal/protostruct" "cloud.google.com/go/internal/trace" "google.golang.org/api/iterator" "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/wrapperspb" ) +var ( + errMetricsBeforeEnd = errors.New("firestore: ExplainMetrics are available only after the iterator reaches the end") +) + // Query represents a Firestore query. // // Query values are immutable. Each Query method creates @@ -62,9 +67,104 @@ type Query struct { // e.g. read time readSettings *readSettings + // readOptions specifies constraints for running the query + // e.g. explainOptions + runQuerySettings *runQuerySettings + findNearest *pb.StructuredQuery_FindNearest } +// ExplainMetrics represents explain metrics 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: + // + // [ + // {"query_scope": "Collection", "properties": "(foo ASC, __name__ ASC)"}, + // {"query_scope": "Collection", "properties": "(bar ASC, __name__ ASC)"} + // ] + IndexesUsed []*map[string]any +} + +// ExecutionStats represents execution statistics for the query. +type ExecutionStats struct { + // Total number of results returned, including documents, projections, + // aggregation results, keys. + ResultsReturned 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]any +} + +func fromExplainMetricsProto(pbExplainMetrics *pb.ExplainMetrics) *ExplainMetrics { + if pbExplainMetrics == nil { + return nil + } + return &ExplainMetrics{ + PlanSummary: fromPlanSummaryProto(pbExplainMetrics.PlanSummary), + ExecutionStats: fromExecutionStatsProto(pbExplainMetrics.ExecutionStats), + } +} + +func fromPlanSummaryProto(pbPlanSummary *pb.PlanSummary) *PlanSummary { + if pbPlanSummary == nil { + return nil + } + + planSummary := &PlanSummary{} + indexesUsed := []*map[string]any{} + for _, pbIndexUsed := range pbPlanSummary.GetIndexesUsed() { + indexUsed := protostruct.DecodeToMap(pbIndexUsed) + indexesUsed = append(indexesUsed, &indexUsed) + } + + planSummary.IndexesUsed = indexesUsed + return planSummary +} + +func fromExecutionStatsProto(pbstats *pb.ExecutionStats) *ExecutionStats { + if pbstats == nil { + return nil + } + + executionStats := &ExecutionStats{ + ResultsReturned: pbstats.GetResultsReturned(), + ReadOperations: pbstats.GetReadOperations(), + } + + executionDuration := pbstats.GetExecutionDuration().AsDuration() + executionStats.ExecutionDuration = &executionDuration + + debugStats := protostruct.DecodeToMap(pbstats.GetDebugStats()) + executionStats.DebugStats = &debugStats + + return executionStats +} + // DocumentID is the special field name representing the ID of a document // in queries. const DocumentID = "__name__" @@ -284,6 +384,19 @@ func (q Query) EndBefore(docSnapshotOrFieldValues ...interface{}) Query { return q } +// WithRunOptions allows passing options to the query +// Calling WithRunOptions overrides a previous call to WithRunOptions. +func (q Query) WithRunOptions(opts ...RunOption) Query { + settings, err := newRunQuerySettings(opts) + if err != nil { + q.err = err + return q + } + + q.runQuerySettings = settings + return q +} + func (q *Query) processCursorArg(name string, docSnapshotOrFieldValues []interface{}) ([]interface{}, *DocumentSnapshot, error) { for _, e := range docSnapshotOrFieldValues { if ds, ok := e.(*DocumentSnapshot); ok { @@ -340,17 +453,11 @@ func (q Query) query() *Query { return &q } // This could be useful, for instance, if executing a query formed in one // process in another. func (q Query) Serialize() ([]byte, error) { - structuredQuery, err := q.toProto() + req, err := q.toRunQueryRequestProto() if err != nil { return nil, err } - - p := &pb.RunQueryRequest{ - Parent: q.parentPath, - QueryType: &pb.RunQueryRequest_StructuredQuery{StructuredQuery: structuredQuery}, - } - - return proto.Marshal(p) + return proto.Marshal(req) } // Deserialize takes a slice of bytes holding the wire-format message of RunQueryRequest, @@ -366,6 +473,24 @@ func (q Query) Deserialize(bytes []byte) (Query, error) { return q.fromProto(&runQueryRequest) } +func (q Query) toRunQueryRequestProto() (*pb.RunQueryRequest, error) { + structuredQuery, err := q.toProto() + if err != nil { + return nil, err + } + + var explainOptions *pb.ExplainOptions + if q.runQuerySettings != nil && q.runQuerySettings.explainOptions != nil { + explainOptions = q.runQuerySettings.explainOptions + } + p := &pb.RunQueryRequest{ + Parent: q.parentPath, + ExplainOptions: explainOptions, + QueryType: &pb.RunQueryRequest_StructuredQuery{StructuredQuery: structuredQuery}, + } + return p, nil +} + // DistanceMeasure is the distance measure to use when comparing vectors with [Query.FindNearest] or [Query.FindNearestPath]. type DistanceMeasure int32 @@ -581,6 +706,13 @@ func (q Query) fromProto(pbQuery *pb.RunQueryRequest) (Query, error) { q.limit = limit } + var err error + q.runQuerySettings, err = newRunQuerySettings(nil) + if err != nil { + q.err = err + return q, q.err + } + q.runQuerySettings.explainOptions = pbQuery.GetExplainOptions() q.findNearest = pbq.GetFindNearest() // NOTE: limit to last isn't part of the proto, this is a client-side concept @@ -1106,9 +1238,24 @@ type DocumentIterator struct { // is an internal detail. type docIterator interface { next() (*DocumentSnapshot, error) + getExplainMetrics() (*ExplainMetrics, error) stop() } +// ExplainMetrics returns query explain metrics. +// This is only present when [ExplainOptions] is added to the query +// (see [Query.WithRunOptions]), and after the iterator reaches the end. +// An error is returned if either of those conditions does not hold. +func (it *DocumentIterator) ExplainMetrics() (*ExplainMetrics, error) { + if it == nil { + return nil, errors.New("firestore: iterator is nil") + } + if it.err == nil || it.err != iterator.Done { + return nil, errMetricsBeforeEnd + } + return it.iter.getExplainMetrics() +} + // Next returns the next result. Its second return value is iterator.Done if there // are no more results. Once Next returns Done, all subsequent calls will return // Done. @@ -1119,10 +1266,12 @@ func (it *DocumentIterator) Next() (*DocumentSnapshot, error) { if it.q.limitToLast { return nil, errors.New("firestore: queries that include limitToLast constraints cannot be streamed. Use DocumentIterator.GetAll() instead") } + ds, err := it.iter.next() if err != nil { it.err = err } + return ds, err } @@ -1179,6 +1328,9 @@ type queryDocumentIterator struct { tid []byte // transaction ID, if any streamClient pb.Firestore_RunQueryClient readSettings *readSettings // readOptions, if any + + // Query explain metrics. This is only present when ExplainOptions is used. + explainMetrics *ExplainMetrics } func newQueryDocumentIterator(ctx context.Context, q *Query, tid []byte, rs *readSettings) *queryDocumentIterator { @@ -1192,6 +1344,7 @@ func newQueryDocumentIterator(ctx context.Context, q *Query, tid []byte, rs *rea } } +// opts override the options stored in it.q.runQuerySettings func (it *queryDocumentIterator) next() (_ *DocumentSnapshot, err error) { client := it.q.c if it.streamClient == nil { @@ -1204,14 +1357,10 @@ func (it *queryDocumentIterator) next() (_ *DocumentSnapshot, err error) { } }() - sq, err := it.q.toProto() + req, err := it.q.toRunQueryRequestProto() if err != nil { return nil, err } - req := &pb.RunQueryRequest{ - Parent: it.q.parentPath, - QueryType: &pb.RunQueryRequest_StructuredQuery{StructuredQuery: sq}, - } // Respect transactions first and read options (read time) second if rt, hasOpts := parseReadTime(client, it.readSettings); hasOpts { @@ -1238,7 +1387,11 @@ func (it *queryDocumentIterator) next() (_ *DocumentSnapshot, err error) { break } // No document => partial progress; keep receiving. + it.explainMetrics = fromExplainMetricsProto(res.GetExplainMetrics()) } + + it.explainMetrics = fromExplainMetricsProto(res.GetExplainMetrics()) + docRef, err := pathToDoc(res.Document.Name, client) if err != nil { return nil, err @@ -1250,6 +1403,13 @@ func (it *queryDocumentIterator) next() (_ *DocumentSnapshot, err error) { return doc, nil } +func (it *queryDocumentIterator) getExplainMetrics() (*ExplainMetrics, error) { + if it == nil { + return nil, fmt.Errorf("firestore: iterator is nil") + } + return it.explainMetrics, nil +} + func (it *queryDocumentIterator) stop() { it.cancel() } @@ -1343,6 +1503,9 @@ func (it *btreeDocumentIterator) next() (*DocumentSnapshot, error) { } func (*btreeDocumentIterator) stop() {} +func (*btreeDocumentIterator) getExplainMetrics() (*ExplainMetrics, error) { + return nil, nil +} // WithReadOptions specifies constraints for accessing documents from the database, // e.g. at what time snapshot to read the documents. @@ -1478,12 +1641,21 @@ func (a *AggregationQuery) WithAvg(path string, alias string) *AggregationQuery // Get retrieves the aggregation query results from the service. func (a *AggregationQuery) Get(ctx context.Context) (AggregationResult, error) { + aro, err := a.GetResponse(ctx) + if aro != nil { + return aro.Result, err + } + return nil, err +} + +// GetResponse runs the aggregation with the options provided in the query +func (a *AggregationQuery) GetResponse(ctx context.Context) (aro *AggregationResponse, err error) { a.query.processLimitToLast() client := a.query.c.c q, err := a.query.toProto() if err != nil { - return nil, err + return aro, err } req := &pb.RunAggregationQueryRequest{ @@ -1498,6 +1670,10 @@ func (a *AggregationQuery) Get(ctx context.Context) (AggregationResult, error) { }, } + if a.query.runQuerySettings != nil { + req.ExplainOptions = a.query.runQuerySettings.explainOptions + } + if a.tx != nil { req.ConsistencySelector = &pb.RunAggregationQueryRequest_Transaction{ Transaction: a.tx.id, @@ -1510,7 +1686,8 @@ func (a *AggregationQuery) Get(ctx context.Context) (AggregationResult, error) { return nil, err } - resp := make(AggregationResult) + aro = &AggregationResponse{} + var resp AggregationResult for { res, err := stream.Recv() @@ -1520,15 +1697,29 @@ func (a *AggregationQuery) Get(ctx context.Context) (AggregationResult, error) { if err != nil { return nil, err } + if res.Result != nil { + if resp == nil { + resp = make(AggregationResult) + } + f := res.Result.AggregateFields - f := res.Result.AggregateFields - - for k, v := range f { - resp[k] = v + for k, v := range f { + resp[k] = v + } } + aro.ExplainMetrics = fromExplainMetricsProto(res.GetExplainMetrics()) } - return resp, nil + aro.Result = resp + return aro, nil } // AggregationResult contains the results of an aggregation query. type AggregationResult map[string]interface{} + +// AggregationResponse contains AggregationResult and response from the run options in the query +type AggregationResponse struct { + Result AggregationResult + + // Query explain metrics. This is only present when ExplainOptions is provided. + ExplainMetrics *ExplainMetrics +} diff --git a/firestore/query_test.go b/firestore/query_test.go index e04bb9278d9f..5d369198546c 100644 --- a/firestore/query_test.go +++ b/firestore/query_test.go @@ -16,8 +16,10 @@ package firestore import ( "context" + "fmt" "math" "sort" + "strings" "testing" pb "cloud.google.com/go/firestore/apiv1/firestorepb" @@ -239,7 +241,7 @@ func TestFilterToProto(t *testing.T) { type toProtoScenario struct { desc string in Query - want *pb.StructuredQuery + want *pb.RunQueryRequest } // Creates protos used to test toProto, FromProto, ToProto funcs. @@ -266,84 +268,185 @@ func createTestScenarios(t *testing.T) []toProtoScenario { { desc: "q.Select()", in: q.Select(), - want: &pb.StructuredQuery{ - Select: &pb.StructuredQuery_Projection{ - Fields: []*pb.StructuredQuery_FieldReference{fref1("__name__")}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Select: &pb.StructuredQuery_Projection{ + Fields: []*pb.StructuredQuery_FieldReference{fref1("__name__")}, + }, + }, }, }, }, { desc: `q.Select("a", "b")`, in: q.Select("a", "b"), - want: &pb.StructuredQuery{ - Select: &pb.StructuredQuery_Projection{ - Fields: []*pb.StructuredQuery_FieldReference{fref1("a"), fref1("b")}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Select: &pb.StructuredQuery_Projection{ + Fields: []*pb.StructuredQuery_FieldReference{fref1("a"), fref1("b")}, + }, + }, }, }, }, { desc: `q.Select("a", "b").Select("c")`, in: q.Select("a", "b").Select("c"), // last wins - want: &pb.StructuredQuery{ - Select: &pb.StructuredQuery_Projection{ - Fields: []*pb.StructuredQuery_FieldReference{fref1("c")}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Select: &pb.StructuredQuery_Projection{ + Fields: []*pb.StructuredQuery_FieldReference{fref1("c")}, + }, + }, }, }, }, { desc: `q.SelectPaths([]string{"*"}, []string{"/"})`, in: q.SelectPaths([]string{"*"}, []string{"/"}), - want: &pb.StructuredQuery{ - Select: &pb.StructuredQuery_Projection{ - Fields: []*pb.StructuredQuery_FieldReference{fref1("*"), fref1("/")}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Select: &pb.StructuredQuery_Projection{ + Fields: []*pb.StructuredQuery_FieldReference{fref1("*"), fref1("/")}, + }, + }, }, }, }, { desc: `q.Where("a", ">", 5)`, in: q.Where("a", ">", 5), - want: &pb.StructuredQuery{Where: filtr([]string{"a"}, ">", 5)}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{Where: filtr([]string{"a"}, ">", 5)}, + }, + }, }, { desc: `q.Where("a", "==", NaN)`, in: q.Where("a", "==", float32(math.NaN())), - want: &pb.StructuredQuery{Where: filtr([]string{"a"}, "==", math.NaN())}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{Where: filtr([]string{"a"}, "==", math.NaN())}, + }, + }, }, { desc: `q.Where("a", "!=", 3)`, in: q.Where("a", "!=", 3), - want: &pb.StructuredQuery{Where: filtr([]string{"a"}, "!=", 3)}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{Where: filtr([]string{"a"}, "!=", 3)}, + }, + }, }, { desc: `q.Where("a", "in", []int{7, 8})`, in: q.Where("a", "in", []int{7, 8}), - want: &pb.StructuredQuery{Where: filtr([]string{"a"}, "in", []int{7, 8})}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{Where: filtr([]string{"a"}, "in", []int{7, 8})}, + }, + }, }, { desc: `q.Where("a", "not-in", []int{9})`, in: q.Where("a", "not-in", []int{9}), - want: &pb.StructuredQuery{Where: filtr([]string{"a"}, "not-in", []int{9})}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{Where: filtr([]string{"a"}, "not-in", []int{9})}, + }, + }, }, { desc: `q.Where("c", "array-contains", 1)`, in: q.Where("c", "array-contains", 1), - want: &pb.StructuredQuery{Where: filtr([]string{"c"}, "array-contains", 1)}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{Where: filtr([]string{"c"}, "array-contains", 1)}, + }, + }, }, { desc: `q.Where("c", "array-contains-any", []int{1, 2})`, in: q.Where("c", "array-contains-any", []int{1, 2}), - want: &pb.StructuredQuery{Where: filtr([]string{"c"}, "array-contains-any", []int{1, 2})}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{Where: filtr([]string{"c"}, "array-contains-any", []int{1, 2})}, + }, + }, + }, + { + desc: `q.Where("c", "array-contains-any", []int{1, 2}).RunOptions(ExplainOptions{Analyze: true})`, + in: q.Where("c", "array-contains-any", []int{1, 2}).WithRunOptions(ExplainOptions{Analyze: true}), + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{Where: filtr([]string{"c"}, "array-contains-any", []int{1, 2})}, + }, + ExplainOptions: &pb.ExplainOptions{ + Analyze: true, + }, + }, + }, + { + desc: `q.Where("c", "array-contains-any", []int{1, 2}).RunOptions(ExplainOptions{Analyze: false})`, + in: q.Where("c", "array-contains-any", []int{1, 2}).WithRunOptions(ExplainOptions{Analyze: false}), + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{Where: filtr([]string{"c"}, "array-contains-any", []int{1, 2})}, + }, + ExplainOptions: &pb.ExplainOptions{ + Analyze: false, + }, + }, + }, + { + desc: `q.Where("c", "array-contains-any", []int{1, 2}) RunOptions invoked multiple times`, + in: q.Where("c", "array-contains-any", []int{1, 2}). + WithRunOptions(ExplainOptions{Analyze: false}). + WithRunOptions(ExplainOptions{Analyze: true}), + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{Where: filtr([]string{"c"}, "array-contains-any", []int{1, 2})}, + }, + ExplainOptions: &pb.ExplainOptions{ + Analyze: true, + }, + }, }, { desc: `q.Where("a", ">", 5).Where("b", "<", "foo")`, in: q.Where("a", ">", 5).Where("b", "<", "foo"), - want: &pb.StructuredQuery{ - Where: &pb.StructuredQuery_Filter{ - FilterType: &pb.StructuredQuery_Filter_CompositeFilter{ - CompositeFilter: &pb.StructuredQuery_CompositeFilter{ - Op: pb.StructuredQuery_CompositeFilter_AND, - Filters: []*pb.StructuredQuery_Filter{ - filtr([]string{"a"}, ">", 5), filtr([]string{"b"}, "<", "foo"), + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Where: &pb.StructuredQuery_Filter{ + FilterType: &pb.StructuredQuery_Filter_CompositeFilter{ + CompositeFilter: &pb.StructuredQuery_CompositeFilter{ + Op: pb.StructuredQuery_CompositeFilter_AND, + Filters: []*pb.StructuredQuery_Filter{ + filtr([]string{"a"}, ">", 5), filtr([]string{"b"}, "<", "foo"), + }, + }, }, }, }, @@ -368,27 +471,32 @@ func createTestScenarios(t *testing.T) []toProtoScenario { }, }, ), - want: &pb.StructuredQuery{ - Where: &pb.StructuredQuery_Filter{ - FilterType: &pb.StructuredQuery_Filter_CompositeFilter{ - CompositeFilter: &pb.StructuredQuery_CompositeFilter{ - Op: pb.StructuredQuery_CompositeFilter_AND, - Filters: []*pb.StructuredQuery_Filter{ - { - FilterType: &pb.StructuredQuery_Filter_FieldFilter{ - FieldFilter: &pb.StructuredQuery_FieldFilter{ - Field: &pb.StructuredQuery_FieldReference{FieldPath: "a"}, - Op: pb.StructuredQuery_FieldFilter_GREATER_THAN, - Value: intval(5), + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Where: &pb.StructuredQuery_Filter{ + FilterType: &pb.StructuredQuery_Filter_CompositeFilter{ + CompositeFilter: &pb.StructuredQuery_CompositeFilter{ + Op: pb.StructuredQuery_CompositeFilter_AND, + Filters: []*pb.StructuredQuery_Filter{ + { + FilterType: &pb.StructuredQuery_Filter_FieldFilter{ + FieldFilter: &pb.StructuredQuery_FieldFilter{ + Field: &pb.StructuredQuery_FieldReference{FieldPath: "a"}, + Op: pb.StructuredQuery_FieldFilter_GREATER_THAN, + Value: intval(5), + }, + }, }, - }, - }, - { - FilterType: &pb.StructuredQuery_Filter_FieldFilter{ - FieldFilter: &pb.StructuredQuery_FieldFilter{ - Field: &pb.StructuredQuery_FieldReference{FieldPath: "b"}, - Op: pb.StructuredQuery_FieldFilter_LESS_THAN, - Value: strval("foo"), + { + FilterType: &pb.StructuredQuery_Filter_FieldFilter{ + FieldFilter: &pb.StructuredQuery_FieldFilter{ + Field: &pb.StructuredQuery_FieldReference{FieldPath: "b"}, + Op: pb.StructuredQuery_FieldFilter_LESS_THAN, + Value: strval("foo"), + }, + }, }, }, }, @@ -401,118 +509,172 @@ func createTestScenarios(t *testing.T) []toProtoScenario { { desc: ` q.WherePath([]string{"/", "*"}, ">", 5)`, in: q.WherePath([]string{"/", "*"}, ">", 5), - want: &pb.StructuredQuery{Where: filtr([]string{"/", "*"}, ">", 5)}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{Where: filtr([]string{"/", "*"}, ">", 5)}, + }, + }, }, { desc: `q.OrderBy("b", Asc).OrderBy("a", Desc).OrderByPath([]string{"~"}, Asc)`, in: q.OrderBy("b", Asc).OrderBy("a", Desc).OrderByPath([]string{"~"}, Asc), - want: &pb.StructuredQuery{ - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("b"), Direction: pb.StructuredQuery_ASCENDING}, - {Field: fref1("a"), Direction: pb.StructuredQuery_DESCENDING}, - {Field: fref1("~"), Direction: pb.StructuredQuery_ASCENDING}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("b"), Direction: pb.StructuredQuery_ASCENDING}, + {Field: fref1("a"), Direction: pb.StructuredQuery_DESCENDING}, + {Field: fref1("~"), Direction: pb.StructuredQuery_ASCENDING}, + }, + }, }, }, }, { desc: `q.Offset(2).Limit(3)`, in: q.Offset(2).Limit(3), - want: &pb.StructuredQuery{ - Offset: 2, - Limit: &wrapperspb.Int32Value{Value: 3}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Offset: 2, + Limit: &wrapperspb.Int32Value{Value: 3}, + }, + }, }, }, { desc: `q.Offset(2).Limit(3).Limit(4).Offset(5)`, in: q.Offset(2).Limit(3).Limit(4).Offset(5), // last wins - want: &pb.StructuredQuery{ - Offset: 5, - Limit: &wrapperspb.Int32Value{Value: 4}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Offset: 5, + Limit: &wrapperspb.Int32Value{Value: 4}, + }, + }, }, }, { desc: `q.OrderBy("a", Asc).StartAt(7).EndBefore(9)`, in: q.OrderBy("a", Asc).StartAt(7).EndBefore(9), - want: &pb.StructuredQuery{ - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{intval(7)}, - Before: true, - }, - EndAt: &pb.Cursor{ - Values: []*pb.Value{intval(9)}, - Before: true, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{intval(7)}, + Before: true, + }, + EndAt: &pb.Cursor{ + Values: []*pb.Value{intval(9)}, + Before: true, + }, + }, }, }, }, { desc: `q.OrderBy("a", Asc).StartAt(7).EndAt(9)`, in: q.OrderBy("a", Asc).StartAt(7).EndAt(9), - want: &pb.StructuredQuery{ - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{intval(7)}, - Before: true, - }, - EndAt: &pb.Cursor{ - Values: []*pb.Value{intval(9)}, - Before: false, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{intval(7)}, + Before: true, + }, + EndAt: &pb.Cursor{ + Values: []*pb.Value{intval(9)}, + Before: false, + }, + }, }, }, }, { desc: `q.OrderBy("a", Asc).StartAfter(7).EndAt(9)`, in: q.OrderBy("a", Asc).StartAfter(7).EndAt(9), - want: &pb.StructuredQuery{ - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{intval(7)}, - Before: false, - }, - EndAt: &pb.Cursor{ - Values: []*pb.Value{intval(9)}, - Before: false, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{intval(7)}, + Before: false, + }, + EndAt: &pb.Cursor{ + Values: []*pb.Value{intval(9)}, + Before: false, + }, + }, }, }, }, { desc: `q.OrderBy(DocumentID, Asc).StartAfter("foo").EndBefore("bar")`, in: q.OrderBy(DocumentID, Asc).StartAfter("foo").EndBefore("bar"), - want: &pb.StructuredQuery{ - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{refval(coll.parentPath + "/C/foo")}, - Before: false, - }, - EndAt: &pb.Cursor{ - Values: []*pb.Value{refval(coll.parentPath + "/C/bar")}, - Before: true, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{refval(coll.parentPath + "/C/foo")}, + Before: false, + }, + EndAt: &pb.Cursor{ + Values: []*pb.Value{refval(coll.parentPath + "/C/bar")}, + Before: true, + }, + }, }, }, }, { desc: `q.OrderBy("a", Asc).OrderBy("b", Desc).StartAfter(7, 8).EndAt(9, 10)`, in: q.OrderBy("a", Asc).OrderBy("b", Desc).StartAfter(7, 8).EndAt(9, 10), - want: &pb.StructuredQuery{ - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, - {Field: fref1("b"), Direction: pb.StructuredQuery_DESCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{intval(7), intval(8)}, - Before: false, - }, - EndAt: &pb.Cursor{ - Values: []*pb.Value{intval(9), intval(10)}, - Before: false, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, + {Field: fref1("b"), Direction: pb.StructuredQuery_DESCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{intval(7), intval(8)}, + Before: false, + }, + EndAt: &pb.Cursor{ + Values: []*pb.Value{intval(9), intval(10)}, + Before: false, + }, + }, }, }, }, @@ -522,17 +684,23 @@ func createTestScenarios(t *testing.T) []toProtoScenario { in: q.OrderBy("a", Asc). StartAfter(1).StartAt(2). EndAt(3).EndBefore(4), - want: &pb.StructuredQuery{ - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{intval(2)}, - Before: true, - }, - EndAt: &pb.Cursor{ - Values: []*pb.Value{intval(4)}, - Before: true, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{intval(2)}, + Before: true, + }, + EndAt: &pb.Cursor{ + Values: []*pb.Value{intval(4)}, + Before: true, + }, + }, }, }, }, @@ -541,144 +709,191 @@ func createTestScenarios(t *testing.T) []toProtoScenario { { desc: `q.StartAt(docsnap)`, in: q.StartAt(docsnap), - want: &pb.StructuredQuery{ - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{refval(coll.parentPath + "/C/D")}, - Before: true, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{refval(coll.parentPath + "/C/D")}, + Before: true, + }, + }, }, }, }, { desc: `q.OrderBy("a", Asc).StartAt(docsnap)`, in: q.OrderBy("a", Asc).StartAt(docsnap), - want: &pb.StructuredQuery{ - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, - {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{intval(7), refval(coll.parentPath + "/C/D")}, - Before: true, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, + {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{intval(7), refval(coll.parentPath + "/C/D")}, + Before: true, + }, + }, }, }, }, - { desc: `q.OrderBy("a", Desc).StartAt(docsnap)`, in: q.OrderBy("a", Desc).StartAt(docsnap), - want: &pb.StructuredQuery{ - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("a"), Direction: pb.StructuredQuery_DESCENDING}, - {Field: fref1("__name__"), Direction: pb.StructuredQuery_DESCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{intval(7), refval(coll.parentPath + "/C/D")}, - Before: true, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("a"), Direction: pb.StructuredQuery_DESCENDING}, + {Field: fref1("__name__"), Direction: pb.StructuredQuery_DESCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{intval(7), refval(coll.parentPath + "/C/D")}, + Before: true, + }, + }, }, }, }, { desc: `q.OrderBy("a", Desc).OrderBy("b", Asc).StartAt(docsnap)`, in: q.OrderBy("a", Desc).OrderBy("b", Asc).StartAt(docsnap), - want: &pb.StructuredQuery{ - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("a"), Direction: pb.StructuredQuery_DESCENDING}, - {Field: fref1("b"), Direction: pb.StructuredQuery_ASCENDING}, - {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{intval(7), intval(8), refval(coll.parentPath + "/C/D")}, - Before: true, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("a"), Direction: pb.StructuredQuery_DESCENDING}, + {Field: fref1("b"), Direction: pb.StructuredQuery_ASCENDING}, + {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{intval(7), intval(8), refval(coll.parentPath + "/C/D")}, + Before: true, + }, + }, }, }, }, { desc: `q.Where("a", "==", 3).StartAt(docsnap)`, in: q.Where("a", "==", 3).StartAt(docsnap), - want: &pb.StructuredQuery{ - Where: filtr([]string{"a"}, "==", 3), - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{refval(coll.parentPath + "/C/D")}, - Before: true, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Where: filtr([]string{"a"}, "==", 3), + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{refval(coll.parentPath + "/C/D")}, + Before: true, + }, + }, }, }, }, { desc: `q.Where("a", "<", 3).StartAt(docsnap)`, in: q.Where("a", "<", 3).StartAt(docsnap), - want: &pb.StructuredQuery{ - Where: filtr([]string{"a"}, "<", 3), - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, - {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{intval(7), refval(coll.parentPath + "/C/D")}, - Before: true, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Where: filtr([]string{"a"}, "<", 3), + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, + {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{intval(7), refval(coll.parentPath + "/C/D")}, + Before: true, + }, + }, }, }, }, { desc: `q.Where("b", "==", 1).Where("a", "<", 3).StartAt(docsnap)`, in: q.Where("b", "==", 1).Where("a", "<", 3).StartAt(docsnap), - want: &pb.StructuredQuery{ - Where: &pb.StructuredQuery_Filter{ - FilterType: &pb.StructuredQuery_Filter_CompositeFilter{ - CompositeFilter: &pb.StructuredQuery_CompositeFilter{ - Op: pb.StructuredQuery_CompositeFilter_AND, - Filters: []*pb.StructuredQuery_Filter{ - filtr([]string{"b"}, "==", 1), - filtr([]string{"a"}, "<", 3), + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Where: &pb.StructuredQuery_Filter{ + FilterType: &pb.StructuredQuery_Filter_CompositeFilter{ + CompositeFilter: &pb.StructuredQuery_CompositeFilter{ + Op: pb.StructuredQuery_CompositeFilter_AND, + Filters: []*pb.StructuredQuery_Filter{ + filtr([]string{"b"}, "==", 1), + filtr([]string{"a"}, "<", 3), + }, + }, }, }, + OrderBy: []*pb.StructuredQuery_Order{ + {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, + {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, + }, + StartAt: &pb.Cursor{ + Values: []*pb.Value{intval(7), refval(coll.parentPath + "/C/D")}, + Before: true, + }, }, }, - OrderBy: []*pb.StructuredQuery_Order{ - {Field: fref1("a"), Direction: pb.StructuredQuery_ASCENDING}, - {Field: fref1("__name__"), Direction: pb.StructuredQuery_ASCENDING}, - }, - StartAt: &pb.Cursor{ - Values: []*pb.Value{intval(7), refval(coll.parentPath + "/C/D")}, - Before: true, - }, }, }, { desc: `q.Where("a", ">", 5).FindNearest float64 vector`, in: q.Where("a", ">", 5). FindNearest("embeddedField", []float64{100, 200, 300}, 2, DistanceMeasureEuclidean, nil).q, - want: &pb.StructuredQuery{ - Where: filtr([]string{"a"}, ">", 5), - FindNearest: &pb.StructuredQuery_FindNearest{ - VectorField: fref1("embeddedField"), - QueryVector: &pb.Value{ - ValueType: &pb.Value_MapValue{ - MapValue: &pb.MapValue{ - Fields: map[string]*pb.Value{ - typeKey: stringToProtoValue(typeValVector), - valueKey: { - ValueType: &pb.Value_ArrayValue{ - ArrayValue: &pb.ArrayValue{ - Values: []*pb.Value{ - {ValueType: &pb.Value_DoubleValue{DoubleValue: 100}}, - {ValueType: &pb.Value_DoubleValue{DoubleValue: 200}}, - {ValueType: &pb.Value_DoubleValue{DoubleValue: 300}}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Where: filtr([]string{"a"}, ">", 5), + FindNearest: &pb.StructuredQuery_FindNearest{ + VectorField: fref1("embeddedField"), + QueryVector: &pb.Value{ + ValueType: &pb.Value_MapValue{ + MapValue: &pb.MapValue{ + Fields: map[string]*pb.Value{ + typeKey: stringToProtoValue(typeValVector), + valueKey: { + ValueType: &pb.Value_ArrayValue{ + ArrayValue: &pb.ArrayValue{ + Values: []*pb.Value{ + {ValueType: &pb.Value_DoubleValue{DoubleValue: 100}}, + {ValueType: &pb.Value_DoubleValue{DoubleValue: 200}}, + {ValueType: &pb.Value_DoubleValue{DoubleValue: 300}}, + }, + }, }, }, }, }, }, }, + Limit: &wrapperspb.Int32Value{Value: trunc32(2)}, + DistanceMeasure: pb.StructuredQuery_FindNearest_EUCLIDEAN, }, }, - Limit: &wrapperspb.Int32Value{Value: trunc32(2)}, - DistanceMeasure: pb.StructuredQuery_FindNearest_EUCLIDEAN, }, }, }, @@ -686,32 +901,38 @@ func createTestScenarios(t *testing.T) []toProtoScenario { desc: `q.Where("a", ">", 5).FindNearest float32 vector`, in: q.Where("a", ">", 5). FindNearest("embeddedField", []float32{100, 200, 300}, 2, DistanceMeasureEuclidean, nil).q, - want: &pb.StructuredQuery{ - Where: filtr([]string{"a"}, ">", 5), - FindNearest: &pb.StructuredQuery_FindNearest{ - VectorField: fref1("embeddedField"), - QueryVector: &pb.Value{ - ValueType: &pb.Value_MapValue{ - MapValue: &pb.MapValue{ - Fields: map[string]*pb.Value{ - typeKey: stringToProtoValue(typeValVector), - valueKey: { - ValueType: &pb.Value_ArrayValue{ - ArrayValue: &pb.ArrayValue{ - Values: []*pb.Value{ - {ValueType: &pb.Value_DoubleValue{DoubleValue: 100}}, - {ValueType: &pb.Value_DoubleValue{DoubleValue: 200}}, - {ValueType: &pb.Value_DoubleValue{DoubleValue: 300}}, + want: &pb.RunQueryRequest{ + Parent: q.parentPath, + + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + Where: filtr([]string{"a"}, ">", 5), + FindNearest: &pb.StructuredQuery_FindNearest{ + VectorField: fref1("embeddedField"), + QueryVector: &pb.Value{ + ValueType: &pb.Value_MapValue{ + MapValue: &pb.MapValue{ + Fields: map[string]*pb.Value{ + typeKey: stringToProtoValue(typeValVector), + valueKey: { + ValueType: &pb.Value_ArrayValue{ + ArrayValue: &pb.ArrayValue{ + Values: []*pb.Value{ + {ValueType: &pb.Value_DoubleValue{DoubleValue: 100}}, + {ValueType: &pb.Value_DoubleValue{DoubleValue: 200}}, + {ValueType: &pb.Value_DoubleValue{DoubleValue: 300}}, + }, + }, }, }, }, }, }, }, + Limit: &wrapperspb.Int32Value{Value: trunc32(2)}, + DistanceMeasure: pb.StructuredQuery_FindNearest_EUCLIDEAN, }, }, - Limit: &wrapperspb.Int32Value{Value: trunc32(2)}, - DistanceMeasure: pb.StructuredQuery_FindNearest_EUCLIDEAN, }, }, }, @@ -720,13 +941,15 @@ func createTestScenarios(t *testing.T) []toProtoScenario { func TestQueryToProto(t *testing.T) { for _, test := range createTestScenarios(t) { - got, err := test.in.toProto() + got, err := test.in.toRunQueryRequestProto() if err != nil { t.Fatalf("%s: %v", test.desc, err) } - test.want.From = []*pb.StructuredQuery_CollectionSelector{{CollectionId: "C"}} + pbStructuredQuery := test.want.QueryType.(*pb.RunQueryRequest_StructuredQuery) + pbStructuredQuery.StructuredQuery.From = []*pb.StructuredQuery_CollectionSelector{{CollectionId: "C"}} if !testEqual(got, test.want) { - t.Fatalf("%s:\ngot\n%v\nwant\n%v", test.desc, pretty.Value(got), pretty.Value(test.want)) + + t.Fatalf("%s:\ngot\n%v\nwant\n%v\ndiff\n%v", test.desc, pretty.Value(got), pretty.Value(test.want), testDiff(got, test.want)) } } } @@ -747,9 +970,9 @@ func TestQueryFromProtoRoundTrip(t *testing.T) { if err != nil { t.Fatal(err) } - got, err := gotq.toProto() + got, err := gotq.toRunQueryRequestProto() want := test.want - want.From = []*pb.StructuredQuery_CollectionSelector{{CollectionId: "C"}} + want.QueryType.(*pb.RunQueryRequest_StructuredQuery).StructuredQuery.From = []*pb.StructuredQuery_CollectionSelector{{CollectionId: "C"}} if diff := cmp.Diff(want, got, protocmp.Transform()); diff != "" { t.Errorf("mismatch (-want, +got)\n: %s", diff) } @@ -915,6 +1138,7 @@ func TestQueryGetAll(t *testing.T) { const dbPath = "projects/projectID/databases/(default)" ctx := context.Background() c, srv, cleanup := newMock(t) + srv.reset() defer cleanup() docNames := []string{"C/a", "C/b"} @@ -941,6 +1165,7 @@ func TestQueryGetAll(t *testing.T) { if err != nil { t.Fatal(err) } + fmt.Printf("gotDocs: %+v\n", gotDocs) if got, want := len(gotDocs), len(wantPBDocs); got != want { t.Errorf("got %d docs, wanted %d", got, want) } @@ -1421,6 +1646,108 @@ func TestWithAvgPath(t *testing.T) { } } +func TestExplainOptionsApply(t *testing.T) { + pbExplainOptions := pb.ExplainOptions{Analyze: true} + for _, testcase := range []struct { + desc string + existingOptions *pb.ExplainOptions + newOptions ExplainOptions + wantErrMsg string + }{ + { + desc: "ExplainOptions specified multiple times", + existingOptions: &pbExplainOptions, + newOptions: ExplainOptions{Analyze: true}, + wantErrMsg: "ExplainOptions can be specified only once", + }, + { + desc: "ExplainOptions specified once", + existingOptions: nil, + newOptions: ExplainOptions{Analyze: true}, + }, + } { + 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.Error(), testcase.wantErrMsg) + } + } +} + +func TestNewRunQuerySettings(t *testing.T) { + for _, testcase := range []struct { + desc string + opts []RunOption + wantErrMsg string + }{ + { + desc: "nil RunOption", + opts: []RunOption{ExplainOptions{Analyze: true}, nil}, + wantErrMsg: "cannot be nil", + }, + { + desc: "success RunOption", + opts: []RunOption{ExplainOptions{Analyze: true}}, + }, + { + desc: "ExplainOptions specified multiple times", + opts: []RunOption{ExplainOptions{Analyze: true}, ExplainOptions{Analyze: false}, ExplainOptions{Analyze: true}}, + wantErrMsg: "ExplainOptions 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) + } + } +} + +func TestQueryRunOptionsAndGetAllWithOptions(t *testing.T) { + ctx := context.Background() + client, srv, cleanup := newMock(t) + defer cleanup() + + dbPath := "projects/projectID/databases/(default)" + collectionName := "collection01" + wantReq := &pb.RunQueryRequest{ + Parent: fmt.Sprintf("%v/documents", dbPath), + QueryType: &pb.RunQueryRequest_StructuredQuery{ + StructuredQuery: &pb.StructuredQuery{ + From: []*pb.StructuredQuery_CollectionSelector{ + {CollectionId: collectionName}, + }, + }, + }, + ExplainOptions: &pb.ExplainOptions{ + Analyze: false, + }, + } + + srv.addRPC(wantReq, []interface{}{ + &pb.RunQueryResponse{ + Document: &pb.Document{ + Name: fmt.Sprintf("%v/documents/%v/doc1", dbPath, collectionName), + CreateTime: aTimestamp, + UpdateTime: aTimestamp, + Fields: map[string]*pb.Value{"f": intval(2)}, + }, + ReadTime: aTimestamp, + }, + }) + + _, err := client.Collection(collectionName). + WithRunOptions(ExplainOptions{Analyze: false}). + WithRunOptions(ExplainOptions{Analyze: true}). + WithRunOptions(ExplainOptions{Analyze: false}). + Documents(ctx). + GetAll() + + if err != nil { + t.Fatal(err) + } +} + func TestFindNearest(t *testing.T) { ctx := context.Background() c, srv, cleanup := newMock(t) diff --git a/firestore/transaction.go b/firestore/transaction.go index 3ddc8beca84c..5cf07eab5564 100644 --- a/firestore/transaction.go +++ b/firestore/transaction.go @@ -35,6 +35,7 @@ type Transaction struct { readOnly bool readAfterWrite bool readSettings *readSettings + explainOptions *ExplainOptions } // A TransactionOption is an option passed to Client.Transaction.