From a9b67cfc95b567d29358501ec7c5883b1f90bd3e Mon Sep 17 00:00:00 2001
From: Baha Aiman <bahaaiman@google.com>
Date: Thu, 6 Jul 2023 13:59:15 -0700
Subject: [PATCH] fix(datastore): Return error from RunAggregationQuery (#8222)

---
 datastore/client.go           |  2 +-
 datastore/integration_test.go | 77 +++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/datastore/client.go b/datastore/client.go
index 2d2bdc92802..c10791294f7 100644
--- a/datastore/client.go
+++ b/datastore/client.go
@@ -81,7 +81,7 @@ func (dc *datastoreClient) RunAggregationQuery(ctx context.Context, in *pb.RunAg
 		res, err = dc.c.RunAggregationQuery(ctx, in, opts...)
 		return err
 	})
-	return res, nil
+	return res, err
 }
 
 func (dc *datastoreClient) BeginTransaction(ctx context.Context, in *pb.BeginTransactionRequest, opts ...grpc.CallOption) (res *pb.BeginTransactionResponse, err error) {
diff --git a/datastore/integration_test.go b/datastore/integration_test.go
index 5af0a780cf8..39783578e1a 100644
--- a/datastore/integration_test.go
+++ b/datastore/integration_test.go
@@ -33,6 +33,7 @@ import (
 	"cloud.google.com/go/rpcreplay"
 	"google.golang.org/api/iterator"
 	"google.golang.org/api/option"
+	pb "google.golang.org/genproto/googleapis/datastore/v1"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -651,6 +652,82 @@ func TestIntegration_Filters(t *testing.T) {
 	})
 }
 
+func TestIntegration_AggregationQueries(t *testing.T) {
+	ctx := context.Background()
+	client := newTestClient(ctx, t)
+	defer client.Close()
+
+	parent := NameKey("SQParent", "TestIntegration_Filters"+suffix, nil)
+	now := timeNow.Truncate(time.Millisecond).Unix()
+	children := []*SQChild{
+		{I: 0, T: now, U: now},
+		{I: 1, T: now, U: now},
+		{I: 2, T: now, U: now},
+		{I: 3, T: now, U: now},
+		{I: 4, T: now, U: now},
+		{I: 5, T: now, U: now},
+		{I: 6, T: now, U: now},
+		{I: 7, T: now, U: now},
+	}
+
+	keys := make([]*Key, len(children))
+	for i := range keys {
+		keys[i] = IncompleteKey("SQChild", parent)
+	}
+	keys, err := client.PutMulti(ctx, keys, children)
+	if err != nil {
+		t.Fatalf("client.PutMulti: %v", err)
+	}
+	defer func() {
+		err := client.DeleteMulti(ctx, keys)
+		if err != nil {
+			t.Errorf("client.DeleteMulti: %v", err)
+		}
+	}()
+
+	baseQuery := NewQuery("SQChild").Ancestor(parent)
+	testCases := []struct {
+		desc          string
+		aggQuery      *AggregationQuery
+		wantFailure   bool
+		wantErrMsg    string
+		wantAggResult AggregationResult
+	}{
+		{
+			desc:          "Count Failure - Missing index",
+			aggQuery:      baseQuery.Filter("T>=", now).NewAggregationQuery().WithCount("count"),
+			wantFailure:   true,
+			wantErrMsg:    "no matching index found",
+			wantAggResult: nil,
+		},
+		{
+			desc:        "Count Success",
+			aggQuery:    baseQuery.Filter("T=", now).Filter("I>=", 3).NewAggregationQuery().WithCount("count"),
+			wantFailure: false,
+			wantErrMsg:  "",
+			wantAggResult: map[string]interface{}{
+				"count": &pb.Value{ValueType: &pb.Value_IntegerValue{IntegerValue: 5}},
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		gotAggResult, gotErr := client.RunAggregationQuery(ctx, testCase.aggQuery)
+		gotFailure := gotErr != nil
+
+		if gotFailure != testCase.wantFailure ||
+			(gotErr != nil && !strings.Contains(gotErr.Error(), testCase.wantErrMsg)) {
+			t.Errorf("%q: Mismatch in error got: %v, want: %q", testCase.desc, gotErr, testCase.wantErrMsg)
+			continue
+		}
+		if !reflect.DeepEqual(gotAggResult, testCase.wantAggResult) {
+			t.Errorf("%q: Mismatch in aggregation result got: %v, want: %v", testCase.desc, gotAggResult, testCase.wantAggResult)
+			continue
+		}
+	}
+
+}
+
 type ckey struct{}
 
 func TestIntegration_LargeQuery(t *testing.T) {