From ce4c149c88e06c215cae7959a463fb958d27244c Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Mon, 29 Aug 2022 14:14:17 +0530 Subject: [PATCH 01/82] Add method in Datastore client to invoke rpc for aggregation query --- .../java/com/google/datastore/v1/client/Datastore.java | 10 ++++++++++ .../datastore/v1/client/DatastoreClientTest.java | 9 +++++++++ 2 files changed, 19 insertions(+) diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java index db117142f..4de2cc96e 100644 --- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java +++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java @@ -27,6 +27,8 @@ import com.google.datastore.v1.ReserveIdsResponse; import com.google.datastore.v1.RollbackRequest; import com.google.datastore.v1.RollbackResponse; +import com.google.datastore.v1.RunAggregationQueryRequest; +import com.google.datastore.v1.RunAggregationQueryResponse; import com.google.datastore.v1.RunQueryRequest; import com.google.datastore.v1.RunQueryResponse; import com.google.rpc.Code; @@ -120,4 +122,12 @@ public RunQueryResponse runQuery(RunQueryRequest request) throws DatastoreExcept throw invalidResponseException("runQuery", exception); } } + + public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) throws DatastoreException { + try (InputStream is = remoteRpc.call("runAggregationQuery", request)) { + return RunAggregationQueryResponse.parseFrom(is); + } catch (IOException exception) { + throw invalidResponseException("runAggregationQuery", exception); + } + } } diff --git a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/DatastoreClientTest.java b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/DatastoreClientTest.java index 2ab2c89f8..16a6303bb 100644 --- a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/DatastoreClientTest.java +++ b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/DatastoreClientTest.java @@ -38,6 +38,8 @@ import com.google.datastore.v1.ReserveIdsResponse; import com.google.datastore.v1.RollbackRequest; import com.google.datastore.v1.RollbackResponse; +import com.google.datastore.v1.RunAggregationQueryRequest; +import com.google.datastore.v1.RunAggregationQueryResponse; import com.google.datastore.v1.RunQueryRequest; import com.google.datastore.v1.RunQueryResponse; import com.google.datastore.v1.client.testing.MockCredential; @@ -336,6 +338,13 @@ public void runQuery() throws Exception { expectRpc("runQuery", request.build(), response.build()); } + @Test + public void runAggregationQuery() throws Exception { + RunAggregationQueryRequest.Builder request = RunAggregationQueryRequest.newBuilder(); + RunAggregationQueryResponse.Builder response = RunAggregationQueryResponse.newBuilder(); + expectRpc("runAggregationQuery", request.build(), response.build()); + } + private void expectRpc(String methodName, Message request, Message response) throws Exception { Datastore datastore = factory.create(options.build()); MockDatastoreFactory mockClient = (MockDatastoreFactory) factory; From 100115f90173fcbf56c77bf75a5fa1759f5478c3 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Mon, 29 Aug 2022 18:08:27 +0530 Subject: [PATCH 02/82] Creating count aggregation and using it to populate Aggregation proto --- .../google/cloud/datastore/Aggregation.java | 26 ++++++++ .../aggregation/CountAggregation.java | 66 +++++++++++++++++++ .../aggregation/DatastoreAggregation.java | 34 ++++++++++ .../DatastoreAggregationBuilder.java | 23 +++++++ .../aggregation/CountAggregationTest.java | 66 +++++++++++++++++++ 5 files changed, 215 insertions(+) create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/Aggregation.java create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregation.java create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregationBuilder.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Aggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Aggregation.java new file mode 100644 index 000000000..5c7b4e428 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Aggregation.java @@ -0,0 +1,26 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.datastore; + +import com.google.cloud.datastore.aggregation.CountAggregation; + +public class Aggregation { + + public static CountAggregation.Builder count() { + return new CountAggregation.Builder(); + } +} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java new file mode 100644 index 000000000..35e28bedf --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java @@ -0,0 +1,66 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.datastore.aggregation; + +import com.google.datastore.v1.AggregationQuery.Aggregation; +import com.google.datastore.v1.AggregationQuery.Aggregation.Count; +import com.google.protobuf.Int64Value; + +public class CountAggregation extends DatastoreAggregation { + + private final long limit; + + private CountAggregation(String alias, long limit) { + super(alias); + this.limit = limit; + } + + @Override + Aggregation toPb() { + Count count = Count.newBuilder() + .setUpTo(Int64Value.of(limit)) + .build(); + Aggregation.Builder aggregationBuilder = Aggregation.newBuilder() + .setCount(count); + if (this.getAlias() != null) { + aggregationBuilder.setAlias(this.getAlias()); + } + return aggregationBuilder.build(); + } + + public static class Builder implements DatastoreAggregationBuilder { + + private String alias; + private long limit; + + public Builder as(String alias) { + this.alias = alias; + return this; + } + + public Builder limit(long limit) { + this.limit = limit; + return this; + } + + @Override + public CountAggregation build() { + return new CountAggregation(alias, limit); + } + } + +} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregation.java new file mode 100644 index 000000000..07d1554da --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregation.java @@ -0,0 +1,34 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.datastore.aggregation; + +import com.google.datastore.v1.AggregationQuery.Aggregation; + +abstract class DatastoreAggregation { + + private final String alias; + + public DatastoreAggregation(String alias) { + this.alias = alias; + } + + abstract Aggregation toPb(); + + public String getAlias() { + return alias; + } +} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregationBuilder.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregationBuilder.java new file mode 100644 index 000000000..cf5248045 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregationBuilder.java @@ -0,0 +1,23 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.datastore.aggregation; + +public interface DatastoreAggregationBuilder { + + A build(); + +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java new file mode 100644 index 000000000..866b873f7 --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java @@ -0,0 +1,66 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.datastore.aggregation; + +import static com.google.cloud.datastore.Aggregation.count; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +import com.google.datastore.v1.AggregationQuery; +import org.junit.Test; + +public class CountAggregationTest { + + @Test + public void testCountAggregationWithDefaultValues() { + AggregationQuery.Aggregation countAggregationPb = count().build().toPb(); + + assertThat(countAggregationPb.getCount().getUpTo().getValue(), equalTo(0L)); + assertThat(countAggregationPb.getAlias(), equalTo("")); + } + + @Test + public void testCountAggregationWithLimit() { + AggregationQuery.Aggregation countAggregationPb = count() + .limit(100) + .build().toPb(); + + assertThat(countAggregationPb.getCount().getUpTo().getValue(), equalTo(100L)); + assertThat(countAggregationPb.getAlias(), equalTo("")); + } + + @Test + public void testCountAggregationWithAlias() { + AggregationQuery.Aggregation countAggregationPb = count() + .as("column_1") + .build().toPb(); + + assertThat(countAggregationPb.getCount().getUpTo().getValue(), equalTo(0L)); + assertThat(countAggregationPb.getAlias(), equalTo("column_1")); + } + + @Test + public void testCountAggregationWithAliasAndLimit() { + AggregationQuery.Aggregation countAggregationPb = count() + .as("column_1") + .limit(100) + .build().toPb(); + + assertThat(countAggregationPb.getCount().getUpTo().getValue(), equalTo(100L)); + assertThat(countAggregationPb.getAlias(), equalTo("column_1")); + } +} \ No newline at end of file From c1956d8107a7dfe2b79ee66dd74bd9394e434475 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Tue, 30 Aug 2022 11:21:26 +0530 Subject: [PATCH 03/82] Moving aggregation builder method to root level aggregation class --- .../google/cloud/datastore/Aggregation.java | 26 ------------------- ...storeAggregation.java => Aggregation.java} | 14 ++++++---- ...onBuilder.java => AggregationBuilder.java} | 2 +- .../aggregation/CountAggregation.java | 10 +++---- .../aggregation/CountAggregationTest.java | 2 +- 5 files changed, 16 insertions(+), 38 deletions(-) delete mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/Aggregation.java rename google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/{DatastoreAggregation.java => Aggregation.java} (74%) rename google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/{DatastoreAggregationBuilder.java => AggregationBuilder.java} (89%) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Aggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Aggregation.java deleted file mode 100644 index 5c7b4e428..000000000 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Aggregation.java +++ /dev/null @@ -1,26 +0,0 @@ -/* - * Copyright 2022 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.cloud.datastore; - -import com.google.cloud.datastore.aggregation.CountAggregation; - -public class Aggregation { - - public static CountAggregation.Builder count() { - return new CountAggregation.Builder(); - } -} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java similarity index 74% rename from google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregation.java rename to google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java index 07d1554da..64554051b 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregation.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java @@ -16,19 +16,23 @@ package com.google.cloud.datastore.aggregation; -import com.google.datastore.v1.AggregationQuery.Aggregation; +import com.google.datastore.v1.AggregationQuery; -abstract class DatastoreAggregation { +abstract class Aggregation { private final String alias; - public DatastoreAggregation(String alias) { + public Aggregation(String alias) { this.alias = alias; } - abstract Aggregation toPb(); - public String getAlias() { return alias; } + + abstract AggregationQuery.Aggregation toPb(); + + public static CountAggregation.Builder count() { + return new CountAggregation.Builder(); + } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregationBuilder.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java similarity index 89% rename from google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregationBuilder.java rename to google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java index cf5248045..99179497a 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/DatastoreAggregationBuilder.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java @@ -16,7 +16,7 @@ package com.google.cloud.datastore.aggregation; -public interface DatastoreAggregationBuilder { +public interface AggregationBuilder { A build(); diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java index 35e28bedf..dc2ce18e3 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java @@ -16,11 +16,11 @@ package com.google.cloud.datastore.aggregation; -import com.google.datastore.v1.AggregationQuery.Aggregation; +import com.google.datastore.v1.AggregationQuery; import com.google.datastore.v1.AggregationQuery.Aggregation.Count; import com.google.protobuf.Int64Value; -public class CountAggregation extends DatastoreAggregation { +public class CountAggregation extends Aggregation { private final long limit; @@ -30,11 +30,11 @@ private CountAggregation(String alias, long limit) { } @Override - Aggregation toPb() { + AggregationQuery.Aggregation toPb() { Count count = Count.newBuilder() .setUpTo(Int64Value.of(limit)) .build(); - Aggregation.Builder aggregationBuilder = Aggregation.newBuilder() + AggregationQuery.Aggregation.Builder aggregationBuilder = AggregationQuery.Aggregation.newBuilder() .setCount(count); if (this.getAlias() != null) { aggregationBuilder.setAlias(this.getAlias()); @@ -42,7 +42,7 @@ Aggregation toPb() { return aggregationBuilder.build(); } - public static class Builder implements DatastoreAggregationBuilder { + public static class Builder implements AggregationBuilder { private String alias; private long limit; diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java index 866b873f7..93a5d6372 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java @@ -16,7 +16,7 @@ package com.google.cloud.datastore.aggregation; -import static com.google.cloud.datastore.Aggregation.count; +import static com.google.cloud.datastore.aggregation.Aggregation.count; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; From 493edfa057e4cfaac35dd18cd6a55444919f7e64 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 7 Sep 2022 17:53:42 +0530 Subject: [PATCH 04/82] Introducing RecordQuery to represent queries which returns entity records when executed --- .../google/cloud/datastore/DatastoreImpl.java | 6 +-- .../com/google/cloud/datastore/GqlQuery.java | 26 ++++++++++--- .../com/google/cloud/datastore/Query.java | 28 +------------- .../cloud/datastore/QueryResultsImpl.java | 4 +- .../google/cloud/datastore/RecordQuery.java | 37 +++++++++++++++++++ .../cloud/datastore/StructuredQuery.java | 24 ++++++++++-- .../cloud/datastore/TransactionImpl.java | 2 +- .../google/cloud/datastore/DatastoreTest.java | 4 +- 8 files changed, 88 insertions(+), 43 deletions(-) create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index 9892e1517..06151a664 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -172,15 +172,15 @@ public T runInTransaction( @Override public QueryResults run(Query query) { - return run(null, query); + return run(null, (RecordQuery) query); } @Override public QueryResults run(Query query, ReadOption... options) { - return run(toReadOptionsPb(options), query); + return run(toReadOptionsPb(options), (RecordQuery) query); } - QueryResults run(com.google.datastore.v1.ReadOptions readOptionsPb, Query query) { + QueryResults run(com.google.datastore.v1.ReadOptions readOptionsPb, RecordQuery query) { return new QueryResultsImpl<>(this, readOptionsPb, query); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java index 2b99fd0a9..ba5ff178a 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java @@ -20,6 +20,7 @@ import static com.google.common.base.Preconditions.checkNotNull; import com.google.cloud.Timestamp; +import com.google.cloud.datastore.Query.ResultType; import com.google.common.base.MoreObjects; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -71,7 +72,7 @@ * @param the type of the result values this query will produce * @see GQL Reference */ -public final class GqlQuery extends Query { +public final class GqlQuery implements RecordQuery { private static final long serialVersionUID = -5514894742849230793L; @@ -80,6 +81,10 @@ public final class GqlQuery extends Query { private final ImmutableMap namedBindings; private final ImmutableList positionalBindings; + private final ResultType resultType; + private final String namespace; + + static final class Binding implements Serializable { private static final long serialVersionUID = 2344746877591371548L; @@ -423,7 +428,8 @@ private static Binding toBinding( } private GqlQuery(Builder builder) { - super(builder.resultType, builder.namespace); + resultType = checkNotNull(builder.resultType); + namespace = builder.namespace; queryString = builder.queryString; allowLiteral = builder.allowLiteral; namedBindings = ImmutableMap.copyOf(builder.namedBindings); @@ -461,9 +467,19 @@ public List getNumberArgs() { return builder.build(); } + @Override + public String getNamespace() { + return namespace; + } + + @Override + public ResultType getType() { + return resultType; + } + @Override public String toString() { - return super.toStringHelper() + return toStringHelper() .add("queryString", queryString) .add("allowLiteral", allowLiteral) .add("namedBindings", namedBindings) @@ -508,12 +524,12 @@ com.google.datastore.v1.GqlQuery toPb() { } @Override - void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb) { + public void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb) { requestPb.setGqlQuery(toPb()); } @Override - Query nextQuery(com.google.datastore.v1.RunQueryResponse responsePb) { + public RecordQuery nextQuery(com.google.datastore.v1.RunQueryResponse responsePb) { return StructuredQuery.fromPb(getType(), getNamespace(), responsePb.getQuery()) .nextQuery(responsePb); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java index 00aa6f17c..f56d7fad8 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java @@ -35,12 +35,9 @@ * @param the type of the values returned by this query. * @see Datastore Queries */ -public abstract class Query implements Serializable { +public interface Query extends Serializable { - private static final long serialVersionUID = 7967659059395653941L; - - private final ResultType resultType; - private final String namespace; + String getNamespace(); /** * This class represents the expected type of the result. ENTITY: A full entity represented by @@ -156,27 +153,6 @@ static ResultType fromPb(com.google.datastore.v1.EntityResult.ResultType type } } - Query(ResultType resultType, String namespace) { - this.resultType = checkNotNull(resultType); - this.namespace = namespace; - } - - ResultType getType() { - return resultType; - } - - public String getNamespace() { - return namespace; - } - - ToStringHelper toStringHelper() { - return MoreObjects.toStringHelper(this).add("type", resultType).add("namespace", namespace); - } - - abstract void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb); - - abstract Query nextQuery(com.google.datastore.v1.RunQueryResponse responsePb); - /** * Returns a new {@link GqlQuery} builder. * diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java index 9ed822985..0c58e79c7 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java @@ -30,7 +30,7 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults private final com.google.datastore.v1.ReadOptions readOptionsPb; private final com.google.datastore.v1.PartitionId partitionIdPb; private final ResultType queryResultType; - private Query query; + private RecordQuery query; private ResultType actualResultType; private com.google.datastore.v1.RunQueryResponse runQueryResponsePb; private com.google.datastore.v1.Query mostRecentQueryPb; @@ -40,7 +40,7 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults private MoreResultsType moreResults; QueryResultsImpl( - DatastoreImpl datastore, com.google.datastore.v1.ReadOptions readOptionsPb, Query query) { + DatastoreImpl datastore, com.google.datastore.v1.ReadOptions readOptionsPb, RecordQuery query) { this.datastore = datastore; this.readOptionsPb = readOptionsPb; this.query = query; diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java new file mode 100644 index 000000000..bed1a0361 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java @@ -0,0 +1,37 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import com.google.api.core.InternalApi; +import com.google.common.base.MoreObjects; +import com.google.common.base.MoreObjects.ToStringHelper; + +@InternalApi +public interface RecordQuery extends Query{ + + ResultType getType(); + + @InternalApi + void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb); + + @InternalApi + RecordQuery nextQuery(com.google.datastore.v1.RunQueryResponse responsePb); + + default ToStringHelper toStringHelper() { + return MoreObjects.toStringHelper(this).add("type", getType()).add("namespace", getNamespace()); + } + +} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java index 8e50d0867..c0e4daa21 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java @@ -29,6 +29,7 @@ import com.google.cloud.StringEnumType; import com.google.cloud.StringEnumValue; import com.google.cloud.Timestamp; +import com.google.cloud.datastore.Query.ResultType; import com.google.common.base.MoreObjects; import com.google.common.base.MoreObjects.ToStringHelper; import com.google.common.base.Preconditions; @@ -85,7 +86,7 @@ * @see Datastore * queries */ -public abstract class StructuredQuery extends Query { +public abstract class StructuredQuery implements RecordQuery { private static final long serialVersionUID = 546838955624019594L; static final String KEY_PROPERTY_NAME = "__key__"; @@ -100,6 +101,10 @@ public abstract class StructuredQuery extends Query { private final int offset; private final Integer limit; + private final ResultType resultType; + private final String namespace; + + public abstract static class Filter implements Serializable { private static final long serialVersionUID = -6443285436239990860L; @@ -899,7 +904,8 @@ B mergeFrom(com.google.datastore.v1.Query queryPb) { } StructuredQuery(BuilderImpl builder) { - super(builder.resultType, builder.namespace); + resultType = checkNotNull(builder.resultType); + namespace = builder.namespace; kind = builder.kind; projection = ImmutableList.copyOf(builder.projection); filter = builder.filter; @@ -1014,12 +1020,22 @@ public Integer getLimit() { public abstract Builder toBuilder(); @Override - void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb) { + public String getNamespace() { + return namespace; + } + + @Override + public ResultType getType() { + return resultType; + } + + @Override + public void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb) { requestPb.setQuery(toPb()); } @Override - StructuredQuery nextQuery(com.google.datastore.v1.RunQueryResponse responsePb) { + public StructuredQuery nextQuery(com.google.datastore.v1.RunQueryResponse responsePb) { Builder builder = toBuilder(); builder.setStartCursor(new Cursor(responsePb.getBatch().getEndCursor())); if (offset > 0 && responsePb.getBatch().getSkippedResults() < offset) { diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java index 3318ec866..a9a338f76 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java @@ -93,7 +93,7 @@ public QueryResults run(Query query) { com.google.datastore.v1.ReadOptions.Builder readOptionsPb = com.google.datastore.v1.ReadOptions.newBuilder(); readOptionsPb.setTransaction(transactionId); - return datastore.run(readOptionsPb.build(), query); + return datastore.run(readOptionsPb.build(), (RecordQuery) query); } @Override diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java index fa077bc61..72067fd20 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java @@ -613,7 +613,7 @@ private List buildResponsesForQueryPagination() { Entity entity5 = Entity.newBuilder(KEY5).set("value", "value").build(); datastore.add(ENTITY3, entity4, entity5); List responses = new ArrayList<>(); - Query query = Query.newKeyQueryBuilder().build(); + RecordQuery query = Query.newKeyQueryBuilder().build(); RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder(); query.populatePb(requestPb); QueryResultBatch queryResultBatchPb = @@ -722,7 +722,7 @@ private List buildResponsesForQueryPaginationWithLimit() { datastore.add(ENTITY3, entity4, entity5); DatastoreRpc datastoreRpc = datastore.getOptions().getDatastoreRpcV1(); List responses = new ArrayList<>(); - Query query = Query.newEntityQueryBuilder().build(); + RecordQuery query = Query.newEntityQueryBuilder().build(); RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder(); query.populatePb(requestPb); QueryResultBatch queryResultBatchPb = From 9c7c4ba9648bc44cf981075da81ac6c262636445 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 7 Sep 2022 17:54:42 +0530 Subject: [PATCH 05/82] Updating gitignore with patch extension --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 6f5e4424d..d1e66c922 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,5 @@ target/ *.iml __pycache__/ -.flattened-pom.xml \ No newline at end of file +.flattened-pom.xml +*.patch \ No newline at end of file From cf49a374e7f9261aa9f2e92fe96a6fec660d9e7e Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Thu, 8 Sep 2022 19:16:45 +0530 Subject: [PATCH 06/82] Setting up structure of Aggregation query and its builder --- .../cloud/datastore/AggregationQuery.java | 84 +++++++++++++++++++ .../cloud/datastore/AggregationResult.java | 20 +++++ .../com/google/cloud/datastore/Query.java | 4 + .../datastore/aggregation/Aggregation.java | 2 +- .../aggregation/CountAggregation.java | 18 ++++ .../cloud/datastore/AggregationQueryTest.java | 84 +++++++++++++++++++ .../aggregation/CountAggregationTest.java | 20 +++++ 7 files changed, 231 insertions(+), 1 deletion(-) create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java new file mode 100644 index 000000000..09c6aebdd --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java @@ -0,0 +1,84 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.cloud.datastore.aggregation.Aggregation; +import com.google.cloud.datastore.aggregation.AggregationBuilder; +import java.util.ArrayList; +import java.util.List; + +public class AggregationQuery implements Query { + + private final String namespace; + private final List aggregations; + private final RecordQuery nestedQuery; + + AggregationQuery(String namespace, List aggregations, RecordQuery nestedQuery) { + checkArgument(nestedQuery != null, + "Nested query is required for an aggregation query to run"); + checkArgument(!aggregations.isEmpty(), + "At least one aggregation is required for an aggregation query to run"); + this.namespace = checkNotNull(namespace); + this.aggregations = aggregations; + this.nestedQuery = nestedQuery; + } + + @Override + public String getNamespace() { + return namespace; + } + + public List getAggregations() { + return aggregations; + } + + public RecordQuery getNestedQuery() { + return nestedQuery; + } + + public static class Builder { + + private String namespace; + private List aggregations; + private RecordQuery nestedQuery; + + public Builder() { + this.aggregations = new ArrayList<>(); + } + + public Builder setNamespace(String namespace) { + this.namespace = namespace; + return this; + } + + public Builder addAggregation(AggregationBuilder aggregationBuilder) { + this.aggregations.add(aggregationBuilder.build()); + return this; + } + + public Builder over(RecordQuery nestedQuery) { + this.nestedQuery = nestedQuery; + return this; + } + + public AggregationQuery build() { + return new AggregationQuery(namespace, aggregations, nestedQuery); + } + } +} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java new file mode 100644 index 000000000..bdf2b12f4 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java @@ -0,0 +1,20 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +public class AggregationResult { + +} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java index f56d7fad8..7893a9d1a 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java @@ -242,4 +242,8 @@ public static KeyQuery.Builder newKeyQueryBuilder() { public static ProjectionEntityQuery.Builder newProjectionEntityQueryBuilder() { return new ProjectionEntityQuery.Builder(); } + + public static AggregationQuery.Builder newAggregationQueryBuilder() { + return new AggregationQuery.Builder(); + } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java index 64554051b..7d10c90f9 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java @@ -18,7 +18,7 @@ import com.google.datastore.v1.AggregationQuery; -abstract class Aggregation { +public abstract class Aggregation { private final String alias; diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java index dc2ce18e3..413cf27b6 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java @@ -19,6 +19,7 @@ import com.google.datastore.v1.AggregationQuery; import com.google.datastore.v1.AggregationQuery.Aggregation.Count; import com.google.protobuf.Int64Value; +import java.util.Objects; public class CountAggregation extends Aggregation { @@ -42,6 +43,23 @@ AggregationQuery.Aggregation toPb() { return aggregationBuilder.build(); } + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + CountAggregation that = (CountAggregation) o; + return limit == that.limit && getAlias().equals(that.getAlias()); + } + + @Override + public int hashCode() { + return Objects.hash(limit, getAlias()); + } + public static class Builder implements AggregationBuilder { private String alias; diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java new file mode 100644 index 000000000..a7185fa5a --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java @@ -0,0 +1,84 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import static com.google.cloud.datastore.StructuredQuery.PropertyFilter.eq; +import static com.google.cloud.datastore.aggregation.Aggregation.count; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; + +public class AggregationQueryTest { + + private static final String KIND = "Task"; + private static final String NAMESPACE = "ns"; + private static final EntityQuery COMPLETED_TASK_QUERY = Query.newEntityQueryBuilder() + .setNamespace(NAMESPACE) + .setKind(KIND) + .setFilter(eq("done", true)) + .build(); + + @Rule + public ExpectedException exceptionRule = ExpectedException.none(); + + @Test + public void testAggregationBuilder() { + AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .addAggregation(count().as("total")) + .addAggregation(count().limit(100).as("total_upto_100")) + .over(COMPLETED_TASK_QUERY) + .build(); + + assertThat(aggregationQuery.getNamespace(), equalTo(NAMESPACE)); + assertThat(aggregationQuery.getAggregations().get(0), equalTo(count().as("total").build())); + assertThat(aggregationQuery.getAggregations().get(1), + equalTo(count().limit(100).as("total_upto_100").build())); + assertThat(aggregationQuery.getNestedQuery(), equalTo(COMPLETED_TASK_QUERY)); + } + + @Test + public void testAggregationBuilderWithoutNamespace() { + exceptionRule.expect(NullPointerException.class); + Query.newAggregationQueryBuilder() + .addAggregation(count().as("total")) + .over(COMPLETED_TASK_QUERY) + .build(); + } + + @Test + public void testAggregationBuilderWithoutNestedQuery() { + exceptionRule.expect(IllegalArgumentException.class); + exceptionRule.expectMessage("Nested query is required for an aggregation query to run"); + Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .addAggregation(count().as("total")) + .build(); + } + + @Test + public void testAggregationBuilderWithoutAggregation() { + exceptionRule.expect(IllegalArgumentException.class); + exceptionRule.expectMessage("At least one aggregation is required for an aggregation query to run"); + Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .over(COMPLETED_TASK_QUERY) + .build(); + } +} \ No newline at end of file diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java index 93a5d6372..e3275987f 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java @@ -19,6 +19,9 @@ import static com.google.cloud.datastore.aggregation.Aggregation.count; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertTrue; import com.google.datastore.v1.AggregationQuery; import org.junit.Test; @@ -63,4 +66,21 @@ public void testCountAggregationWithAliasAndLimit() { assertThat(countAggregationPb.getCount().getUpTo().getValue(), equalTo(100L)); assertThat(countAggregationPb.getAlias(), equalTo("column_1")); } + + @Test + public void testEquals() { + CountAggregation.Builder aggregation1 = count() + .as("total") + .limit(100); + + CountAggregation.Builder aggregation2 = count() + .as("total") + .limit(100); + + assertEquals(aggregation1.build(), aggregation2.build()); + assertEquals(aggregation2.build(), aggregation1.build()); + + assertNotEquals(aggregation1.as("new-alias").build(), aggregation2.build()); + assertNotEquals(aggregation1.limit(399).build(), aggregation2.build()); + } } \ No newline at end of file From 78fd6c96681cdc1804abb1f6371ff40b3041c2b5 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Fri, 9 Sep 2022 12:01:54 +0530 Subject: [PATCH 07/82] Introducing ProtoPreparer to populate the request protos --- .../AggregationQueryRequestProtoPreparer.java | 43 +++++++++++++ .../execution/request/ProtoPreparer.java | 20 ++++++ ...regationQueryRequestProtoPreparerTest.java | 64 +++++++++++++++++++ 3 files changed, 127 insertions(+) create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/ProtoPreparer.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java new file mode 100644 index 000000000..691b6ec5d --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java @@ -0,0 +1,43 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore.execution.request; + +import com.google.cloud.datastore.AggregationQuery; +import com.google.cloud.datastore.DatastoreOptions; +import com.google.datastore.v1.PartitionId; +import com.google.datastore.v1.RunAggregationQueryRequest; + +public class AggregationQueryRequestProtoPreparer implements + ProtoPreparer { + + private DatastoreOptions datastoreOptions; + + public AggregationQueryRequestProtoPreparer(DatastoreOptions datastoreOptions) { + this.datastoreOptions = datastoreOptions; + } + + @Override + public RunAggregationQueryRequest prepare(AggregationQuery aggregationQuery) { + PartitionId partitionId = PartitionId.newBuilder() + .setProjectId(datastoreOptions.getProjectId()) + .setNamespaceId(aggregationQuery.getNamespace()) + .build(); + return RunAggregationQueryRequest.newBuilder() + .setPartitionId(partitionId) + .setProjectId(datastoreOptions.getProjectId()) + .build(); + } +} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/ProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/ProtoPreparer.java new file mode 100644 index 000000000..2f3302d39 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/ProtoPreparer.java @@ -0,0 +1,20 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore.execution.request; + +public interface ProtoPreparer { + OUTPUT prepare(INPUT input); +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java new file mode 100644 index 000000000..e65e7e3fc --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java @@ -0,0 +1,64 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore.execution.request; + +import static com.google.cloud.datastore.StructuredQuery.PropertyFilter.eq; +import static com.google.cloud.datastore.aggregation.Aggregation.count; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +import com.google.cloud.datastore.AggregationQuery; +import com.google.cloud.datastore.DatastoreOptions; +import com.google.cloud.datastore.EntityQuery; +import com.google.cloud.datastore.Query; +import com.google.datastore.v1.RunAggregationQueryRequest; +import org.junit.Test; + +public class AggregationQueryRequestProtoPreparerTest { + + private static final String KIND = "Task"; + private static final String NAMESPACE = "ns"; + private static final String PROJECT_ID = "project-id"; + private static final DatastoreOptions DATASTORE_OPTIONS = DatastoreOptions.newBuilder() + .setProjectId(PROJECT_ID) + .setNamespace(NAMESPACE) + .build(); + private static final EntityQuery COMPLETED_TASK_QUERY = Query.newEntityQueryBuilder() + .setNamespace(NAMESPACE).setKind(KIND).setFilter(eq("done", true)).build(); + + private ProtoPreparer protoPreparer = new AggregationQueryRequestProtoPreparer(DATASTORE_OPTIONS); + + @Test + public void shouldPrepareAggregationQueryRequestWithGivenStructuredQuery() { + AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .addAggregation(count().as("total")) + .over(COMPLETED_TASK_QUERY) + .build(); + + RunAggregationQueryRequest runAggregationQueryRequest = protoPreparer.prepare(aggregationQuery); + + assertThat(runAggregationQueryRequest.getProjectId(), equalTo(PROJECT_ID)); + + assertThat(runAggregationQueryRequest.getPartitionId().getProjectId(), equalTo(PROJECT_ID)); + assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId(), equalTo(NAMESPACE)); + + + //TODO verify the complete aggregation query + assertThat(runAggregationQueryRequest.getAggregationQuery(), equalTo(com.google.datastore.v1.AggregationQuery.newBuilder().build())); + + } +} \ No newline at end of file From 2ae957a95b4f026b1f99fb74e0e6c7409636f44c Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Fri, 9 Sep 2022 15:01:40 +0530 Subject: [PATCH 08/82] Delegating responsibility of preparing query proto to QueryPreparer --- .../cloud/datastore/StructuredQuery.java | 36 +---- .../StructuredQueryProtoPreparer.java | 62 ++++++++ .../StructuredQueryProtoPreparerTest.java | 148 ++++++++++++++++++ 3 files changed, 212 insertions(+), 34 deletions(-) create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQueryProtoPreparer.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java index c0e4daa21..3a167a1d8 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java @@ -1050,40 +1050,8 @@ public StructuredQuery nextQuery(com.google.datastore.v1.RunQueryResponse res } com.google.datastore.v1.Query toPb() { - com.google.datastore.v1.Query.Builder queryPb = com.google.datastore.v1.Query.newBuilder(); - if (kind != null) { - queryPb.addKindBuilder().setName(kind); - } - if (startCursor != null) { - queryPb.setStartCursor(startCursor.getByteString()); - } - if (endCursor != null) { - queryPb.setEndCursor(endCursor.getByteString()); - } - if (offset > 0) { - queryPb.setOffset(offset); - } - if (limit != null) { - queryPb.setLimit(com.google.protobuf.Int32Value.newBuilder().setValue(limit)); - } - if (filter != null) { - queryPb.setFilter(filter.toPb()); - } - for (OrderBy value : orderBy) { - queryPb.addOrder(value.toPb()); - } - for (String value : distinctOn) { - queryPb.addDistinctOn( - com.google.datastore.v1.PropertyReference.newBuilder().setName(value).build()); - } - for (String value : projection) { - com.google.datastore.v1.Projection.Builder expressionPb = - com.google.datastore.v1.Projection.newBuilder(); - expressionPb.setProperty( - com.google.datastore.v1.PropertyReference.newBuilder().setName(value).build()); - queryPb.addProjection(expressionPb.build()); - } - return queryPb.build(); + StructuredQueryProtoPreparer protoPreparer = new StructuredQueryProtoPreparer(); + return protoPreparer.prepare(this); } @SuppressWarnings("unchecked") diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQueryProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQueryProtoPreparer.java new file mode 100644 index 000000000..030d8e114 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQueryProtoPreparer.java @@ -0,0 +1,62 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import com.google.cloud.datastore.StructuredQuery.OrderBy; +import com.google.cloud.datastore.execution.request.ProtoPreparer; +import com.google.datastore.v1.Query; +import com.google.protobuf.Int32Value; + +public class StructuredQueryProtoPreparer implements ProtoPreparer, Query> { + + @Override + public Query prepare(StructuredQuery query) { + com.google.datastore.v1.Query.Builder queryPb = com.google.datastore.v1.Query.newBuilder(); + if (query.getKind() != null) { + queryPb.addKindBuilder().setName(query.getKind()); + } + if (query.getStartCursor() != null) { + queryPb.setStartCursor(query.getStartCursor().getByteString()); + } + if (query.getEndCursor() != null) { + queryPb.setEndCursor(query.getEndCursor().getByteString()); + } + if (query.getOffset() > 0) { + queryPb.setOffset(query.getOffset()); + } + if (query.getLimit() != null) { + queryPb.setLimit(Int32Value.of(query.getLimit())); + } + if (query.getFilter() != null) { + queryPb.setFilter(query.getFilter().toPb()); + } + for (OrderBy value : query.getOrderBy()) { + queryPb.addOrder(value.toPb()); + } + for (String value : query.getDistinctOn()) { + queryPb.addDistinctOn( + com.google.datastore.v1.PropertyReference.newBuilder().setName(value).build()); + } + for (String value : query.getProjection()) { + com.google.datastore.v1.Projection expressionPb = com.google.datastore.v1.Projection.newBuilder() + .setProperty(com.google.datastore.v1.PropertyReference.newBuilder().setName(value).build()) + .build(); + queryPb.addProjection(expressionPb); + } + + return queryPb.build(); + } +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java new file mode 100644 index 000000000..6b3adf751 --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java @@ -0,0 +1,148 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import static com.google.cloud.datastore.Query.newEntityQueryBuilder; +import static com.google.datastore.v1.PropertyFilter.Operator.EQUAL; +import static com.google.datastore.v1.PropertyOrder.Direction.ASCENDING; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +import com.google.cloud.datastore.StructuredQuery.OrderBy; +import com.google.cloud.datastore.StructuredQuery.PropertyFilter; +import com.google.datastore.v1.Filter; +import com.google.datastore.v1.KindExpression; +import com.google.datastore.v1.Projection; +import com.google.datastore.v1.PropertyFilter.Operator; +import com.google.datastore.v1.PropertyOrder; +import com.google.datastore.v1.PropertyReference; +import com.google.datastore.v1.Query; +import com.google.datastore.v1.Value; +import com.google.protobuf.ByteString; +import com.google.protobuf.Int32Value; +import org.junit.Test; + +public class StructuredQueryProtoPreparerTest { + + private final StructuredQueryProtoPreparer protoPreparer = new StructuredQueryProtoPreparer(); + + @Test + public void testKind() { + Query queryProto = protoPreparer.prepare(newEntityQueryBuilder().setKind("kind").build()); + + assertThat(queryProto.getKind(0), equalTo(KindExpression.newBuilder().setName("kind").build())); + } + + @Test + public void testStartCursor() { + byte[] bytes = {1, 2}; + Query queryProto = protoPreparer.prepare(newEntityQueryBuilder().setStartCursor( + Cursor.copyFrom(bytes)).build()); + + assertThat(queryProto.getStartCursor(), equalTo(ByteString.copyFrom(bytes))); + } + + @Test + public void testEndCursor() { + byte[] bytes = {1, 2}; + Query queryProto = protoPreparer.prepare(newEntityQueryBuilder().setEndCursor( + Cursor.copyFrom(bytes)).build()); + + assertThat(queryProto.getEndCursor(), equalTo(ByteString.copyFrom(bytes))); + } + + @Test + public void testOffset() { + Query queryProto = protoPreparer.prepare(newEntityQueryBuilder().setOffset(5).build()); + + assertThat(queryProto.getOffset(), equalTo(5)); + } + + @Test + public void testLimit() { + Query queryProto = protoPreparer.prepare(newEntityQueryBuilder().setLimit(5).build()); + + assertThat(queryProto.getLimit(), equalTo(Int32Value.of(5))); + } + + @Test + public void testFilter() { + Query queryProto = protoPreparer.prepare(newEntityQueryBuilder() + .setFilter(PropertyFilter.eq("done", true)) + .build()); + + assertThat(queryProto.getFilter(), equalTo( + propertyFilter("done", EQUAL, Value.newBuilder().setBooleanValue(true).build()) + )); + } + + @Test + public void testOrderBy() { + Query queryProto = protoPreparer.prepare(newEntityQueryBuilder() + .setOrderBy(OrderBy.asc("dept-id"), OrderBy.asc("rank")) + .build()); + + assertThat(queryProto.getOrder(0), equalTo(propertyOrder("dept-id"))); + assertThat(queryProto.getOrder(1), equalTo(propertyOrder("rank"))); + } + + @Test + public void testDistinctOn() { + Query queryProto = protoPreparer.prepare(newEntityQueryBuilder() + .setDistinctOn("dept-id", "rank") + .build()); + + assertThat(queryProto.getDistinctOn(0), equalTo(propertyReference("dept-id"))); + assertThat(queryProto.getDistinctOn(1), equalTo(propertyReference("rank"))); + } + + @Test + public void testProjections() { + Query queryProto = protoPreparer.prepare(newEntityQueryBuilder() + .setProjection("dept-id", "rank") + .build()); + + assertThat(queryProto.getProjection(0), equalTo(projection("dept-id"))); + assertThat(queryProto.getProjection(1), equalTo(projection("rank"))); + } + + private Filter propertyFilter(String propertyName, Operator operator, Value value) { + return Filter.newBuilder().setPropertyFilter( + com.google.datastore.v1.PropertyFilter.newBuilder() + .setProperty(propertyReference(propertyName)) + .setOp(operator) + .setValue(value) + .build() + ).build(); + } + + private PropertyOrder propertyOrder(String value) { + return PropertyOrder.newBuilder() + .setProperty(propertyReference(value)) + .setDirection(ASCENDING) + .build(); + } + + private Projection projection(String value) { + return Projection.newBuilder() + .setProperty(propertyReference(value)) + .build(); + } + + private PropertyReference propertyReference(String value) { + return PropertyReference.newBuilder().setName(value).build(); + } +} \ No newline at end of file From 14ac55f3084717e2fa6132e44091a36bfa75a709 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Fri, 9 Sep 2022 16:14:21 +0530 Subject: [PATCH 09/82] Populating aggregation query with nested structured query --- .../datastore/aggregation/Aggregation.java | 4 +- .../aggregation/CountAggregation.java | 14 ++-- .../AggregationQueryRequestProtoPreparer.java | 35 ++++++-- .../execution/request/ProtoPreparer.java | 3 + .../google/cloud/datastore/ProtoTestData.java | 82 +++++++++++++++++++ .../StructuredQueryProtoPreparerTest.java | 41 ++-------- ...regationQueryRequestProtoPreparerTest.java | 27 ++++-- 7 files changed, 155 insertions(+), 51 deletions(-) create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/ProtoTestData.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java index 7d10c90f9..f12cab60b 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java @@ -16,6 +16,7 @@ package com.google.cloud.datastore.aggregation; +import com.google.api.core.InternalApi; import com.google.datastore.v1.AggregationQuery; public abstract class Aggregation { @@ -30,7 +31,8 @@ public String getAlias() { return alias; } - abstract AggregationQuery.Aggregation toPb(); + @InternalApi + public abstract AggregationQuery.Aggregation toPb(); public static CountAggregation.Builder count() { return new CountAggregation.Builder(); diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java index 413cf27b6..7cb7d57c4 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java @@ -20,6 +20,8 @@ import com.google.datastore.v1.AggregationQuery.Aggregation.Count; import com.google.protobuf.Int64Value; import java.util.Objects; +import java.util.Optional; +import java.util.function.Consumer; public class CountAggregation extends Aggregation { @@ -31,12 +33,14 @@ private CountAggregation(String alias, long limit) { } @Override - AggregationQuery.Aggregation toPb() { - Count count = Count.newBuilder() - .setUpTo(Int64Value.of(limit)) - .build(); + public AggregationQuery.Aggregation toPb() { + Count.Builder countBuilder = Count.newBuilder(); + if(limit > 0) { + countBuilder.setUpTo(Int64Value.of(limit)); + } + AggregationQuery.Aggregation.Builder aggregationBuilder = AggregationQuery.Aggregation.newBuilder() - .setCount(count); + .setCount(countBuilder); if (this.getAlias() != null) { aggregationBuilder.setAlias(this.getAlias()); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java index 691b6ec5d..5f46faadd 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java @@ -17,27 +17,52 @@ import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.DatastoreOptions; +import com.google.cloud.datastore.StructuredQuery; +import com.google.cloud.datastore.StructuredQueryProtoPreparer; +import com.google.cloud.datastore.aggregation.Aggregation; import com.google.datastore.v1.PartitionId; +import com.google.datastore.v1.Query; import com.google.datastore.v1.RunAggregationQueryRequest; public class AggregationQueryRequestProtoPreparer implements ProtoPreparer { - private DatastoreOptions datastoreOptions; + private final DatastoreOptions datastoreOptions; + private final StructuredQueryProtoPreparer structuredQueryProtoPreparer; public AggregationQueryRequestProtoPreparer(DatastoreOptions datastoreOptions) { this.datastoreOptions = datastoreOptions; + this.structuredQueryProtoPreparer = new StructuredQueryProtoPreparer(); } @Override public RunAggregationQueryRequest prepare(AggregationQuery aggregationQuery) { - PartitionId partitionId = PartitionId.newBuilder() - .setProjectId(datastoreOptions.getProjectId()) - .setNamespaceId(aggregationQuery.getNamespace()) - .build(); + PartitionId partitionId = getPartitionId(aggregationQuery); + com.google.datastore.v1.AggregationQuery aggregationQueryProtoBuilder = getAggregationQuery(aggregationQuery); + return RunAggregationQueryRequest.newBuilder() .setPartitionId(partitionId) .setProjectId(datastoreOptions.getProjectId()) + .setAggregationQuery(aggregationQueryProtoBuilder) + .build(); + } + + private com.google.datastore.v1.AggregationQuery getAggregationQuery(AggregationQuery aggregationQuery) { + Query nestedQueryProto = structuredQueryProtoPreparer.prepare( + (StructuredQuery) aggregationQuery.getNestedQuery()); + + com.google.datastore.v1.AggregationQuery.Builder aggregationQueryProtoBuilder = com.google.datastore.v1.AggregationQuery.newBuilder() + .setNestedQuery(nestedQueryProto); + for (Aggregation aggregation : aggregationQuery.getAggregations()) { + aggregationQueryProtoBuilder.addAggregations(aggregation.toPb()); + } + return aggregationQueryProtoBuilder.build(); + } + + private PartitionId getPartitionId(AggregationQuery aggregationQuery) { + return PartitionId.newBuilder() + .setProjectId(datastoreOptions.getProjectId()) + .setNamespaceId(aggregationQuery.getNamespace()) .build(); } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/ProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/ProtoPreparer.java index 2f3302d39..fa56c9e99 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/ProtoPreparer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/ProtoPreparer.java @@ -15,6 +15,9 @@ */ package com.google.cloud.datastore.execution.request; +import com.google.api.core.InternalApi; + +@InternalApi public interface ProtoPreparer { OUTPUT prepare(INPUT input); } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ProtoTestData.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ProtoTestData.java new file mode 100644 index 000000000..d2b4b9a5f --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ProtoTestData.java @@ -0,0 +1,82 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import static com.google.datastore.v1.PropertyOrder.Direction.ASCENDING; + +import com.google.datastore.v1.AggregationQuery.Aggregation; +import com.google.datastore.v1.AggregationQuery.Aggregation.Count; +import com.google.datastore.v1.Filter; +import com.google.datastore.v1.KindExpression; +import com.google.datastore.v1.Projection; +import com.google.datastore.v1.PropertyFilter.Operator; +import com.google.datastore.v1.PropertyOrder; +import com.google.datastore.v1.PropertyReference; +import com.google.datastore.v1.Value; +import com.google.protobuf.Int64Value; + +public class ProtoTestData { + + public static Value booleanValue(boolean value) { + return Value.newBuilder().setBooleanValue(value).build(); + } + + public static KindExpression kind(String kind) { + return KindExpression.newBuilder().setName(kind).build(); + } + + public static Filter propertyFilter(String propertyName, Operator operator, Value value) { + return Filter.newBuilder().setPropertyFilter( + com.google.datastore.v1.PropertyFilter.newBuilder() + .setProperty(propertyReference(propertyName)) + .setOp(operator) + .setValue(value) + .build() + ).build(); + } + + public static PropertyReference propertyReference(String value) { + return PropertyReference.newBuilder().setName(value).build(); + } + + public static Aggregation countAggregation(String alias) { + return Aggregation.newBuilder() + .setAlias(alias) + .setCount(Count.newBuilder().build()) + .build(); + } + + public static Aggregation countAggregation(String alias, long limit) { + return Aggregation.newBuilder() + .setAlias(alias) + .setCount(Count.newBuilder().setUpTo(Int64Value.of(limit))) + .build(); + } + + public static PropertyOrder propertyOrder(String value) { + return PropertyOrder.newBuilder() + .setProperty(propertyReference(value)) + .setDirection(ASCENDING) + .build(); + } + + public static Projection projection(String value) { + return Projection.newBuilder() + .setProperty(propertyReference(value)) + .build(); + } + +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java index 6b3adf751..38deee0c5 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java @@ -15,22 +15,20 @@ */ package com.google.cloud.datastore; +import static com.google.cloud.datastore.ProtoTestData.booleanValue; +import static com.google.cloud.datastore.ProtoTestData.projection; +import static com.google.cloud.datastore.ProtoTestData.propertyFilter; +import static com.google.cloud.datastore.ProtoTestData.propertyOrder; +import static com.google.cloud.datastore.ProtoTestData.propertyReference; import static com.google.cloud.datastore.Query.newEntityQueryBuilder; import static com.google.datastore.v1.PropertyFilter.Operator.EQUAL; -import static com.google.datastore.v1.PropertyOrder.Direction.ASCENDING; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; import com.google.cloud.datastore.StructuredQuery.OrderBy; import com.google.cloud.datastore.StructuredQuery.PropertyFilter; -import com.google.datastore.v1.Filter; import com.google.datastore.v1.KindExpression; -import com.google.datastore.v1.Projection; -import com.google.datastore.v1.PropertyFilter.Operator; -import com.google.datastore.v1.PropertyOrder; -import com.google.datastore.v1.PropertyReference; import com.google.datastore.v1.Query; -import com.google.datastore.v1.Value; import com.google.protobuf.ByteString; import com.google.protobuf.Int32Value; import org.junit.Test; @@ -85,7 +83,7 @@ public void testFilter() { .build()); assertThat(queryProto.getFilter(), equalTo( - propertyFilter("done", EQUAL, Value.newBuilder().setBooleanValue(true).build()) + propertyFilter("done", EQUAL, booleanValue(true)) )); } @@ -118,31 +116,4 @@ public void testProjections() { assertThat(queryProto.getProjection(0), equalTo(projection("dept-id"))); assertThat(queryProto.getProjection(1), equalTo(projection("rank"))); } - - private Filter propertyFilter(String propertyName, Operator operator, Value value) { - return Filter.newBuilder().setPropertyFilter( - com.google.datastore.v1.PropertyFilter.newBuilder() - .setProperty(propertyReference(propertyName)) - .setOp(operator) - .setValue(value) - .build() - ).build(); - } - - private PropertyOrder propertyOrder(String value) { - return PropertyOrder.newBuilder() - .setProperty(propertyReference(value)) - .setDirection(ASCENDING) - .build(); - } - - private Projection projection(String value) { - return Projection.newBuilder() - .setProperty(propertyReference(value)) - .build(); - } - - private PropertyReference propertyReference(String value) { - return PropertyReference.newBuilder().setName(value).build(); - } } \ No newline at end of file diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java index e65e7e3fc..3d897fb47 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java @@ -15,16 +15,25 @@ */ package com.google.cloud.datastore.execution.request; +import static com.google.cloud.datastore.ProtoTestData.booleanValue; +import static com.google.cloud.datastore.ProtoTestData.countAggregation; +import static com.google.cloud.datastore.ProtoTestData.kind; +import static com.google.cloud.datastore.ProtoTestData.propertyFilter; import static com.google.cloud.datastore.StructuredQuery.PropertyFilter.eq; import static com.google.cloud.datastore.aggregation.Aggregation.count; +import static com.google.datastore.v1.PropertyFilter.Operator.EQUAL; +import static java.util.Arrays.asList; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.EntityQuery; +import com.google.cloud.datastore.ProtoTestData; import com.google.cloud.datastore.Query; import com.google.datastore.v1.RunAggregationQueryRequest; +import com.google.datastore.v1.Value; +import java.util.Arrays; import org.junit.Test; public class AggregationQueryRequestProtoPreparerTest { @@ -39,13 +48,15 @@ public class AggregationQueryRequestProtoPreparerTest { private static final EntityQuery COMPLETED_TASK_QUERY = Query.newEntityQueryBuilder() .setNamespace(NAMESPACE).setKind(KIND).setFilter(eq("done", true)).build(); - private ProtoPreparer protoPreparer = new AggregationQueryRequestProtoPreparer(DATASTORE_OPTIONS); + private final AggregationQueryRequestProtoPreparer protoPreparer = new AggregationQueryRequestProtoPreparer( + DATASTORE_OPTIONS); @Test public void shouldPrepareAggregationQueryRequestWithGivenStructuredQuery() { AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() .setNamespace(NAMESPACE) .addAggregation(count().as("total")) + .addAggregation(count().limit(100).as("total_upto_100")) .over(COMPLETED_TASK_QUERY) .build(); @@ -56,9 +67,15 @@ public void shouldPrepareAggregationQueryRequestWithGivenStructuredQuery() { assertThat(runAggregationQueryRequest.getPartitionId().getProjectId(), equalTo(PROJECT_ID)); assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId(), equalTo(NAMESPACE)); - - //TODO verify the complete aggregation query - assertThat(runAggregationQueryRequest.getAggregationQuery(), equalTo(com.google.datastore.v1.AggregationQuery.newBuilder().build())); - + com.google.datastore.v1.AggregationQuery aggregationQueryProto = runAggregationQueryRequest.getAggregationQuery(); + assertThat(aggregationQueryProto.getNestedQuery(), + equalTo(com.google.datastore.v1.Query.newBuilder() + .addKind(kind(KIND)) + .setFilter(propertyFilter("done", EQUAL, booleanValue(true))) + .build())); + assertThat(aggregationQueryProto.getAggregationsList(), equalTo(asList( + countAggregation("total"), + countAggregation("total_upto_100", 100) + ))); } } \ No newline at end of file From 3ff182ffd49913b96f6398530faea46db87d97d6 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Fri, 9 Sep 2022 18:01:57 +0530 Subject: [PATCH 10/82] Delegating responsibility of preparing query proto in GqlQuery to QueryPreparer --- .../com/google/cloud/datastore/GqlQuery.java | 22 +++-- .../datastore/GqlQueryProtoPreparer.java | 41 +++++++++ .../datastore/GqlQueryProtoPreparerTest.java | 85 +++++++++++++++++++ .../google/cloud/datastore/ProtoTestData.java | 13 +++ 4 files changed, 149 insertions(+), 12 deletions(-) create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQueryProtoPreparer.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java index ba5ff178a..3befcb638 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java @@ -20,7 +20,6 @@ import static com.google.common.base.Preconditions.checkNotNull; import com.google.cloud.Timestamp; -import com.google.cloud.datastore.Query.ResultType; import com.google.common.base.MoreObjects; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -458,6 +457,14 @@ public Map getNamedBindings() { return builder.buildOrThrow(); } + public Map getNamedBindingsMap() { + return namedBindings; + } + + public List getPositionalBindingsMap() { + return positionalBindings; + } + /** Returns an immutable list of positional bindings (using original order). */ public List getNumberArgs() { ImmutableList.Builder builder = ImmutableList.builder(); @@ -510,17 +517,8 @@ public boolean equals(Object obj) { } com.google.datastore.v1.GqlQuery toPb() { - com.google.datastore.v1.GqlQuery.Builder queryPb = - com.google.datastore.v1.GqlQuery.newBuilder(); - queryPb.setQueryString(queryString); - queryPb.setAllowLiterals(allowLiteral); - for (Map.Entry entry : namedBindings.entrySet()) { - queryPb.putNamedBindings(entry.getKey(), entry.getValue().toPb()); - } - for (Binding argument : positionalBindings) { - queryPb.addPositionalBindings(argument.toPb()); - } - return queryPb.build(); + GqlQueryProtoPreparer protoPreparer = new GqlQueryProtoPreparer(); + return protoPreparer.prepare(this); } @Override diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQueryProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQueryProtoPreparer.java new file mode 100644 index 000000000..a81cb3640 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQueryProtoPreparer.java @@ -0,0 +1,41 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import com.google.cloud.datastore.GqlQuery.Binding; +import com.google.cloud.datastore.execution.request.ProtoPreparer; +import java.util.Map; + +public class GqlQueryProtoPreparer implements + ProtoPreparer, com.google.datastore.v1.GqlQuery> { + + @Override + public com.google.datastore.v1.GqlQuery prepare(GqlQuery gqlQuery) { + com.google.datastore.v1.GqlQuery.Builder queryPb = + com.google.datastore.v1.GqlQuery.newBuilder(); + + queryPb.setQueryString(gqlQuery.getQueryString()); + queryPb.setAllowLiterals(gqlQuery.allowLiteral()); + for (Map.Entry entry : gqlQuery.getNamedBindingsMap().entrySet()) { + queryPb.putNamedBindings(entry.getKey(), entry.getValue().toPb()); + } + for (Binding argument : gqlQuery.getPositionalBindingsMap()) { + queryPb.addPositionalBindings(argument.toPb()); + } + + return queryPb.build(); + } +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java new file mode 100644 index 000000000..1ed707228 --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java @@ -0,0 +1,85 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import static com.google.cloud.datastore.ProtoTestData.gqlQueryParameter; +import static com.google.cloud.datastore.ProtoTestData.intValue; +import static com.google.cloud.datastore.ProtoTestData.stringValue; +import static com.google.cloud.datastore.Query.newGqlQueryBuilder; +import static java.util.Arrays.asList; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import com.google.datastore.v1.GqlQueryParameter; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import org.junit.Test; + +public class GqlQueryProtoPreparerTest { + + private final GqlQueryProtoPreparer protoPreparer = new GqlQueryProtoPreparer(); + private final GqlQuery.Builder gqlQueryBuilder = newGqlQueryBuilder("SELECT * from Character"); + + @Test + public void testQueryString() { + com.google.datastore.v1.GqlQuery gqlQuery = protoPreparer.prepare(gqlQueryBuilder.build()); + + assertThat(gqlQuery.getQueryString(), equalTo("SELECT * from Character")); + } + + @Test + public void testAllowLiteral() { + assertTrue(protoPreparer.prepare( + gqlQueryBuilder.setAllowLiteral(true).build()).getAllowLiterals()); + assertFalse(protoPreparer.prepare( + gqlQueryBuilder.setAllowLiteral(false).build()).getAllowLiterals()); + } + + @Test + public void testNamedBinding() { + com.google.datastore.v1.GqlQuery gqlQuery = protoPreparer.prepare( + gqlQueryBuilder + .setBinding("name", "John Doe") + .setBinding("age", 27) + .build() + ); + + assertThat(gqlQuery.getNamedBindingsMap(), equalTo(new HashMap() {{ + put("name", gqlQueryParameter(stringValue("John Doe"))); + put("age", gqlQueryParameter(intValue(27))); + }})); + } + + @Test + public void testPositionalBinding() { + com.google.datastore.v1.GqlQuery gqlQuery = protoPreparer.prepare( + gqlQueryBuilder + .addBinding("John Doe") + .addBinding(27) + .build() + ); + + assertThat(gqlQuery.getPositionalBindingsList(), equalTo(asList( + gqlQueryParameter(stringValue("John Doe")), + gqlQueryParameter(intValue(27)) + ))); + } + +} \ No newline at end of file diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ProtoTestData.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ProtoTestData.java index d2b4b9a5f..5cff2cf11 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ProtoTestData.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ProtoTestData.java @@ -20,6 +20,7 @@ import com.google.datastore.v1.AggregationQuery.Aggregation; import com.google.datastore.v1.AggregationQuery.Aggregation.Count; import com.google.datastore.v1.Filter; +import com.google.datastore.v1.GqlQueryParameter; import com.google.datastore.v1.KindExpression; import com.google.datastore.v1.Projection; import com.google.datastore.v1.PropertyFilter.Operator; @@ -34,6 +35,18 @@ public static Value booleanValue(boolean value) { return Value.newBuilder().setBooleanValue(value).build(); } + public static Value stringValue(String value) { + return Value.newBuilder().setStringValue(value).build(); + } + + public static Value intValue(long value) { + return Value.newBuilder().setIntegerValue(value).build(); + } + + public static GqlQueryParameter gqlQueryParameter(Value value) { + return GqlQueryParameter.newBuilder().setValue(value).build(); + } + public static KindExpression kind(String kind) { return KindExpression.newBuilder().setName(kind).build(); } From 7de95e67f5cd056baf41e99d9f9f9547d2b52ba4 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Mon, 12 Sep 2022 13:16:50 +0530 Subject: [PATCH 11/82] Removing RecordQuery from the query hierarchy and making it a standalone interface for now --- .../cloud/datastore/AggregationQuery.java | 9 ++------- .../google/cloud/datastore/DatastoreImpl.java | 10 ++++++---- .../com/google/cloud/datastore/GqlQuery.java | 10 ++-------- .../com/google/cloud/datastore/Query.java | 20 +++++++++++++++++-- .../cloud/datastore/QueryResultsImpl.java | 6 +++--- .../google/cloud/datastore/RecordQuery.java | 9 ++------- .../cloud/datastore/StructuredQuery.java | 11 +++------- .../cloud/datastore/TransactionImpl.java | 2 +- 8 files changed, 37 insertions(+), 40 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java index 09c6aebdd..4e3bc8024 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java @@ -23,26 +23,21 @@ import java.util.ArrayList; import java.util.List; -public class AggregationQuery implements Query { +public class AggregationQuery extends Query { - private final String namespace; private final List aggregations; private final RecordQuery nestedQuery; AggregationQuery(String namespace, List aggregations, RecordQuery nestedQuery) { + super(checkNotNull(namespace)); checkArgument(nestedQuery != null, "Nested query is required for an aggregation query to run"); checkArgument(!aggregations.isEmpty(), "At least one aggregation is required for an aggregation query to run"); - this.namespace = checkNotNull(namespace); this.aggregations = aggregations; this.nestedQuery = nestedQuery; } - @Override - public String getNamespace() { - return namespace; - } public List getAggregations() { return aggregations; diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index 06151a664..283788728 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -172,16 +172,18 @@ public T runInTransaction( @Override public QueryResults run(Query query) { - return run(null, (RecordQuery) query); + return run(null, query); } @Override public QueryResults run(Query query, ReadOption... options) { - return run(toReadOptionsPb(options), (RecordQuery) query); + return run(toReadOptionsPb(options), query); } - QueryResults run(com.google.datastore.v1.ReadOptions readOptionsPb, RecordQuery query) { - return new QueryResultsImpl<>(this, readOptionsPb, query); + //TODO- Need to check whether we can handle this typecasting in a better way + @SuppressWarnings("unchecked") + QueryResults run(com.google.datastore.v1.ReadOptions readOptionsPb, Query query) { + return new QueryResultsImpl(this, readOptionsPb, (RecordQuery) query, query.getNamespace()); } com.google.datastore.v1.RunQueryResponse runQuery( diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java index 3befcb638..5e0c0def5 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java @@ -71,7 +71,7 @@ * @param the type of the result values this query will produce * @see GQL Reference */ -public final class GqlQuery implements RecordQuery { +public final class GqlQuery extends Query implements RecordQuery { private static final long serialVersionUID = -5514894742849230793L; @@ -81,7 +81,6 @@ public final class GqlQuery implements RecordQuery { private final ImmutableList positionalBindings; private final ResultType resultType; - private final String namespace; static final class Binding implements Serializable { @@ -427,8 +426,8 @@ private static Binding toBinding( } private GqlQuery(Builder builder) { + super(builder.namespace); resultType = checkNotNull(builder.resultType); - namespace = builder.namespace; queryString = builder.queryString; allowLiteral = builder.allowLiteral; namedBindings = ImmutableMap.copyOf(builder.namedBindings); @@ -474,11 +473,6 @@ public List getNumberArgs() { return builder.build(); } - @Override - public String getNamespace() { - return namespace; - } - @Override public ResultType getType() { return resultType; diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java index 7893a9d1a..a34ae46f4 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java @@ -20,6 +20,7 @@ import com.google.common.base.MoreObjects; import com.google.common.base.MoreObjects.ToStringHelper; +import com.google.common.base.Preconditions; import com.google.common.collect.Maps; import java.io.Serializable; import java.util.Map; @@ -35,9 +36,11 @@ * @param the type of the values returned by this query. * @see Datastore Queries */ -public interface Query extends Serializable { +public abstract class Query implements Serializable { - String getNamespace(); + private static final long serialVersionUID = 7967659059395653941L; + + private final String namespace; /** * This class represents the expected type of the result. ENTITY: A full entity represented by @@ -153,6 +156,19 @@ static ResultType fromPb(com.google.datastore.v1.EntityResult.ResultType type } } + Query(String namespace) { + this.namespace = namespace; + } + + public String getNamespace() { + return namespace; + } + + ToStringHelper toStringHelper() { + return MoreObjects.toStringHelper(this).add("namespace", namespace); + } + + /** * Returns a new {@link GqlQuery} builder. * diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java index 0c58e79c7..09c319aff 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java @@ -40,7 +40,7 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults private MoreResultsType moreResults; QueryResultsImpl( - DatastoreImpl datastore, com.google.datastore.v1.ReadOptions readOptionsPb, RecordQuery query) { + DatastoreImpl datastore, com.google.datastore.v1.ReadOptions readOptionsPb, RecordQuery query, String namespace) { this.datastore = datastore; this.readOptionsPb = readOptionsPb; this.query = query; @@ -48,8 +48,8 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults com.google.datastore.v1.PartitionId.Builder pbBuilder = com.google.datastore.v1.PartitionId.newBuilder(); pbBuilder.setProjectId(datastore.getOptions().getProjectId()); - if (query.getNamespace() != null) { - pbBuilder.setNamespaceId(query.getNamespace()); + if (namespace != null) { + pbBuilder.setNamespaceId(namespace); } else if (datastore.getOptions().getNamespace() != null) { pbBuilder.setNamespaceId(datastore.getOptions().getNamespace()); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java index bed1a0361..8fb1c29eb 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java @@ -16,11 +16,10 @@ package com.google.cloud.datastore; import com.google.api.core.InternalApi; -import com.google.common.base.MoreObjects; -import com.google.common.base.MoreObjects.ToStringHelper; +import com.google.cloud.datastore.Query.ResultType; @InternalApi -public interface RecordQuery extends Query{ +public interface RecordQuery{ ResultType getType(); @@ -30,8 +29,4 @@ public interface RecordQuery extends Query{ @InternalApi RecordQuery nextQuery(com.google.datastore.v1.RunQueryResponse responsePb); - default ToStringHelper toStringHelper() { - return MoreObjects.toStringHelper(this).add("type", getType()).add("namespace", getNamespace()); - } - } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java index 3a167a1d8..5ccf3e4fb 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java @@ -86,7 +86,7 @@ * @see Datastore * queries */ -public abstract class StructuredQuery implements RecordQuery { +public abstract class StructuredQuery extends Query implements RecordQuery { private static final long serialVersionUID = 546838955624019594L; static final String KEY_PROPERTY_NAME = "__key__"; @@ -102,7 +102,6 @@ public abstract class StructuredQuery implements RecordQuery { private final Integer limit; private final ResultType resultType; - private final String namespace; public abstract static class Filter implements Serializable { @@ -904,8 +903,8 @@ B mergeFrom(com.google.datastore.v1.Query queryPb) { } StructuredQuery(BuilderImpl builder) { + super(builder.namespace); resultType = checkNotNull(builder.resultType); - namespace = builder.namespace; kind = builder.kind; projection = ImmutableList.copyOf(builder.projection); filter = builder.filter; @@ -920,6 +919,7 @@ B mergeFrom(com.google.datastore.v1.Query queryPb) { @Override public String toString() { return toStringHelper() + .add("type", getType()) .add("kind", kind) .add("startCursor", startCursor) .add("endCursor", endCursor) @@ -1019,11 +1019,6 @@ public Integer getLimit() { public abstract Builder toBuilder(); - @Override - public String getNamespace() { - return namespace; - } - @Override public ResultType getType() { return resultType; diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java index a9a338f76..3318ec866 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java @@ -93,7 +93,7 @@ public QueryResults run(Query query) { com.google.datastore.v1.ReadOptions.Builder readOptionsPb = com.google.datastore.v1.ReadOptions.newBuilder(); readOptionsPb.setTransaction(transactionId); - return datastore.run(readOptionsPb.build(), (RecordQuery) query); + return datastore.run(readOptionsPb.build(), query); } @Override From d62964f2727494dd47cd2c47c0fd7b3a154b4a9d Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Mon, 12 Sep 2022 15:44:12 +0530 Subject: [PATCH 12/82] Populating aggregation query with nested gql query --- .../cloud/datastore/AggregationQuery.java | 64 +++++++++++++++---- .../AggregationQueryRequestProtoPreparer.java | 31 ++++++--- .../cloud/datastore/AggregationQueryTest.java | 32 +++++++++- ...regationQueryRequestProtoPreparerTest.java | 47 ++++++++++++-- 4 files changed, 146 insertions(+), 28 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java index 4e3bc8024..db8825a65 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java @@ -15,6 +15,8 @@ */ package com.google.cloud.datastore; +import static com.google.cloud.datastore.AggregationQuery.Mode.GQL; +import static com.google.cloud.datastore.AggregationQuery.Mode.STRUCTURED; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkNotNull; @@ -25,33 +27,50 @@ public class AggregationQuery extends Query { - private final List aggregations; - private final RecordQuery nestedQuery; + private List aggregations; + private StructuredQuery nestedStructuredQuery; + private final Mode mode; + private GqlQuery nestedGqlQuery; - AggregationQuery(String namespace, List aggregations, RecordQuery nestedQuery) { + AggregationQuery(String namespace, List aggregations, + StructuredQuery nestedQuery) { super(checkNotNull(namespace)); - checkArgument(nestedQuery != null, - "Nested query is required for an aggregation query to run"); checkArgument(!aggregations.isEmpty(), "At least one aggregation is required for an aggregation query to run"); this.aggregations = aggregations; - this.nestedQuery = nestedQuery; + this.nestedStructuredQuery = nestedQuery; + this.mode = STRUCTURED; } + AggregationQuery(String namespace, GqlQuery gqlQuery) { + super(checkNotNull(namespace)); + this.nestedGqlQuery = gqlQuery; + this.mode = GQL; + } public List getAggregations() { return aggregations; } - public RecordQuery getNestedQuery() { - return nestedQuery; + public StructuredQuery getNestedStructuredQuery() { + return nestedStructuredQuery; + } + + public GqlQuery getNestedGqlQuery() { + return nestedGqlQuery; + } + + public Mode getMode() { + return mode; } public static class Builder { private String namespace; - private List aggregations; - private RecordQuery nestedQuery; + private Mode mode; + private final List aggregations; + private StructuredQuery nestedStructuredQuery; + private GqlQuery nestedGqlQuery; public Builder() { this.aggregations = new ArrayList<>(); @@ -67,13 +86,32 @@ public Builder addAggregation(AggregationBuilder aggregationBuilder) { return this; } - public Builder over(RecordQuery nestedQuery) { - this.nestedQuery = nestedQuery; + public Builder over(StructuredQuery nestedQuery) { + this.nestedStructuredQuery = nestedQuery; + this.mode = STRUCTURED; + return this; + } + + public Builder over(GqlQuery nestedQuery) { + this.nestedGqlQuery = nestedQuery; + this.mode = GQL; return this; } public AggregationQuery build() { - return new AggregationQuery(namespace, aggregations, nestedQuery); + boolean nestedQueryProvided = nestedGqlQuery != null || nestedStructuredQuery != null; + checkArgument(nestedQueryProvided, + "Nested query is required for an aggregation query to run"); + + if (mode == GQL) { + return new AggregationQuery(namespace, nestedGqlQuery); + } + return new AggregationQuery(namespace, aggregations, nestedStructuredQuery); } } + + public enum Mode { + STRUCTURED, + GQL, + } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java index 5f46faadd..c80ccd2dd 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java @@ -15,11 +15,14 @@ */ package com.google.cloud.datastore.execution.request; +import static com.google.cloud.datastore.AggregationQuery.Mode.GQL; + import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.DatastoreOptions; -import com.google.cloud.datastore.StructuredQuery; +import com.google.cloud.datastore.GqlQueryProtoPreparer; import com.google.cloud.datastore.StructuredQueryProtoPreparer; import com.google.cloud.datastore.aggregation.Aggregation; +import com.google.datastore.v1.GqlQuery; import com.google.datastore.v1.PartitionId; import com.google.datastore.v1.Query; import com.google.datastore.v1.RunAggregationQueryRequest; @@ -29,27 +32,39 @@ public class AggregationQueryRequestProtoPreparer implements private final DatastoreOptions datastoreOptions; private final StructuredQueryProtoPreparer structuredQueryProtoPreparer; + private final GqlQueryProtoPreparer gqlQueryProtoPreparer; public AggregationQueryRequestProtoPreparer(DatastoreOptions datastoreOptions) { this.datastoreOptions = datastoreOptions; this.structuredQueryProtoPreparer = new StructuredQueryProtoPreparer(); + this.gqlQueryProtoPreparer = new GqlQueryProtoPreparer(); } @Override public RunAggregationQueryRequest prepare(AggregationQuery aggregationQuery) { PartitionId partitionId = getPartitionId(aggregationQuery); - com.google.datastore.v1.AggregationQuery aggregationQueryProtoBuilder = getAggregationQuery(aggregationQuery); - - return RunAggregationQueryRequest.newBuilder() + RunAggregationQueryRequest.Builder aggregationQueryRequestBuilder = RunAggregationQueryRequest.newBuilder() .setPartitionId(partitionId) - .setProjectId(datastoreOptions.getProjectId()) - .setAggregationQuery(aggregationQueryProtoBuilder) + .setProjectId(datastoreOptions.getProjectId()); + + if (aggregationQuery.getMode() == GQL) { + return aggregationQueryRequestBuilder + .setGqlQuery(buildGqlQuery(aggregationQuery)) + .build(); + } + return aggregationQueryRequestBuilder + .setAggregationQuery(getAggregationQuery(aggregationQuery)) .build(); } - private com.google.datastore.v1.AggregationQuery getAggregationQuery(AggregationQuery aggregationQuery) { + private GqlQuery buildGqlQuery(AggregationQuery aggregationQuery) { + return gqlQueryProtoPreparer.prepare(aggregationQuery.getNestedGqlQuery()); + } + + private com.google.datastore.v1.AggregationQuery getAggregationQuery( + AggregationQuery aggregationQuery) { Query nestedQueryProto = structuredQueryProtoPreparer.prepare( - (StructuredQuery) aggregationQuery.getNestedQuery()); + aggregationQuery.getNestedStructuredQuery()); com.google.datastore.v1.AggregationQuery.Builder aggregationQueryProtoBuilder = com.google.datastore.v1.AggregationQuery.newBuilder() .setNestedQuery(nestedQueryProto); diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java index a7185fa5a..b3920525a 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java @@ -15,6 +15,8 @@ */ package com.google.cloud.datastore; +import static com.google.cloud.datastore.AggregationQuery.Mode.GQL; +import static com.google.cloud.datastore.AggregationQuery.Mode.STRUCTURED; import static com.google.cloud.datastore.StructuredQuery.PropertyFilter.eq; import static com.google.cloud.datastore.aggregation.Aggregation.count; import static org.hamcrest.CoreMatchers.equalTo; @@ -50,7 +52,8 @@ public void testAggregationBuilder() { assertThat(aggregationQuery.getAggregations().get(0), equalTo(count().as("total").build())); assertThat(aggregationQuery.getAggregations().get(1), equalTo(count().limit(100).as("total_upto_100").build())); - assertThat(aggregationQuery.getNestedQuery(), equalTo(COMPLETED_TASK_QUERY)); + assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY)); + assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED)); } @Test @@ -75,10 +78,35 @@ public void testAggregationBuilderWithoutNestedQuery() { @Test public void testAggregationBuilderWithoutAggregation() { exceptionRule.expect(IllegalArgumentException.class); - exceptionRule.expectMessage("At least one aggregation is required for an aggregation query to run"); + exceptionRule.expectMessage( + "At least one aggregation is required for an aggregation query to run"); Query.newAggregationQueryBuilder() .setNamespace(NAMESPACE) .over(COMPLETED_TASK_QUERY) .build(); } + + @Test + public void testAggregationBuilderWithGqlQuery() { + GqlQuery gqlQuery = Query.newGqlQueryBuilder("SELECT * FROM Task WHERE done = true").build(); + + AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .over(gqlQuery) + .build(); + + assertThat(aggregationQuery.getNestedGqlQuery(), equalTo(gqlQuery)); + assertThat(aggregationQuery.getMode(), equalTo(GQL)); + } + + @Test + public void testAggregationBuilderWithoutProvidingAnyNestedQuery() { + exceptionRule.expect(IllegalArgumentException.class); + exceptionRule.expectMessage( + "Nested query is required for an aggregation query to run"); + Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .build(); + } + } \ No newline at end of file diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java index 3d897fb47..c2cd0635a 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java @@ -17,8 +17,11 @@ import static com.google.cloud.datastore.ProtoTestData.booleanValue; import static com.google.cloud.datastore.ProtoTestData.countAggregation; +import static com.google.cloud.datastore.ProtoTestData.gqlQueryParameter; +import static com.google.cloud.datastore.ProtoTestData.intValue; import static com.google.cloud.datastore.ProtoTestData.kind; import static com.google.cloud.datastore.ProtoTestData.propertyFilter; +import static com.google.cloud.datastore.ProtoTestData.stringValue; import static com.google.cloud.datastore.StructuredQuery.PropertyFilter.eq; import static com.google.cloud.datastore.aggregation.Aggregation.count; import static com.google.datastore.v1.PropertyFilter.Operator.EQUAL; @@ -29,11 +32,11 @@ import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.EntityQuery; -import com.google.cloud.datastore.ProtoTestData; +import com.google.cloud.datastore.GqlQuery; import com.google.cloud.datastore.Query; +import com.google.datastore.v1.GqlQueryParameter; import com.google.datastore.v1.RunAggregationQueryRequest; -import com.google.datastore.v1.Value; -import java.util.Arrays; +import java.util.HashMap; import org.junit.Test; public class AggregationQueryRequestProtoPreparerTest { @@ -45,9 +48,17 @@ public class AggregationQueryRequestProtoPreparerTest { .setProjectId(PROJECT_ID) .setNamespace(NAMESPACE) .build(); - private static final EntityQuery COMPLETED_TASK_QUERY = Query.newEntityQueryBuilder() + private static final EntityQuery COMPLETED_TASK_STRUCTURED_QUERY = Query.newEntityQueryBuilder() .setNamespace(NAMESPACE).setKind(KIND).setFilter(eq("done", true)).build(); + private static final GqlQuery COMPLETED_TASK_GQL_QUERY = Query.newGqlQueryBuilder( + "AGGREGATE COUNT AS total_characters OVER (" + + "SELECT * FROM Character WHERE name = @name and age > @1" + + ")") + .setBinding("name", "John Doe") + .addBinding(27) + .build(); + private final AggregationQueryRequestProtoPreparer protoPreparer = new AggregationQueryRequestProtoPreparer( DATASTORE_OPTIONS); @@ -57,7 +68,7 @@ public void shouldPrepareAggregationQueryRequestWithGivenStructuredQuery() { .setNamespace(NAMESPACE) .addAggregation(count().as("total")) .addAggregation(count().limit(100).as("total_upto_100")) - .over(COMPLETED_TASK_QUERY) + .over(COMPLETED_TASK_STRUCTURED_QUERY) .build(); RunAggregationQueryRequest runAggregationQueryRequest = protoPreparer.prepare(aggregationQuery); @@ -78,4 +89,30 @@ public void shouldPrepareAggregationQueryRequestWithGivenStructuredQuery() { countAggregation("total_upto_100", 100) ))); } + + @Test + public void shouldPrepareAggregationQueryRequestWithGivenGqlQuery() { + AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .over(COMPLETED_TASK_GQL_QUERY) + .build(); + + RunAggregationQueryRequest runAggregationQueryRequest = protoPreparer.prepare(aggregationQuery); + + assertThat(runAggregationQueryRequest.getProjectId(), equalTo(PROJECT_ID)); + + assertThat(runAggregationQueryRequest.getPartitionId().getProjectId(), equalTo(PROJECT_ID)); + assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId(), equalTo(NAMESPACE)); + + com.google.datastore.v1.GqlQuery gqlQueryProto = runAggregationQueryRequest.getGqlQuery(); + + assertThat(gqlQueryProto.getQueryString(), equalTo(COMPLETED_TASK_GQL_QUERY.getQueryString())); + assertThat(gqlQueryProto.getNamedBindingsMap(), + equalTo(new HashMap() {{ + put("name", gqlQueryParameter(stringValue("John Doe"))); + }})); + assertThat(gqlQueryProto.getPositionalBindingsList(), equalTo(asList( + gqlQueryParameter(intValue(27)) + ))); + } } \ No newline at end of file From 1f5b1a4c78b63cda73d33e4818a43a55d5a9150d Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Mon, 12 Sep 2022 15:50:45 +0530 Subject: [PATCH 13/82] Removing deprecation warning by using assertThrows instead of ExpectedException rule --- .../cloud/datastore/AggregationQueryTest.java | 49 ++++++++++--------- 1 file changed, 25 insertions(+), 24 deletions(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java index b3920525a..f6fa2ee37 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java @@ -21,9 +21,12 @@ import static com.google.cloud.datastore.aggregation.Aggregation.count; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertThrows; +import org.junit.Assert; import org.junit.Rule; import org.junit.Test; +import org.junit.function.ThrowingRunnable; import org.junit.rules.ExpectedException; public class AggregationQueryTest { @@ -58,32 +61,31 @@ public void testAggregationBuilder() { @Test public void testAggregationBuilderWithoutNamespace() { - exceptionRule.expect(NullPointerException.class); - Query.newAggregationQueryBuilder() - .addAggregation(count().as("total")) - .over(COMPLETED_TASK_QUERY) - .build(); + assertThrows(NullPointerException.class, () -> + Query.newAggregationQueryBuilder() + .addAggregation(count().as("total")) + .over(COMPLETED_TASK_QUERY) + .build()); } @Test public void testAggregationBuilderWithoutNestedQuery() { - exceptionRule.expect(IllegalArgumentException.class); - exceptionRule.expectMessage("Nested query is required for an aggregation query to run"); - Query.newAggregationQueryBuilder() - .setNamespace(NAMESPACE) - .addAggregation(count().as("total")) - .build(); + assertThrows("Nested query is required for an aggregation query to run", + IllegalArgumentException.class, + () -> Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .addAggregation(count().as("total")) + .build()); } @Test public void testAggregationBuilderWithoutAggregation() { - exceptionRule.expect(IllegalArgumentException.class); - exceptionRule.expectMessage( - "At least one aggregation is required for an aggregation query to run"); - Query.newAggregationQueryBuilder() - .setNamespace(NAMESPACE) - .over(COMPLETED_TASK_QUERY) - .build(); + assertThrows("At least one aggregation is required for an aggregation query to run", + IllegalArgumentException.class, + () -> Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .over(COMPLETED_TASK_QUERY) + .build()); } @Test @@ -101,12 +103,11 @@ public void testAggregationBuilderWithGqlQuery() { @Test public void testAggregationBuilderWithoutProvidingAnyNestedQuery() { - exceptionRule.expect(IllegalArgumentException.class); - exceptionRule.expectMessage( - "Nested query is required for an aggregation query to run"); - Query.newAggregationQueryBuilder() - .setNamespace(NAMESPACE) - .build(); + assertThrows("Nested query is required for an aggregation query to run", + IllegalArgumentException.class, + () -> Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .build()); } } \ No newline at end of file From b69b1664e8b0d201d97002986176e639fda67d93 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Mon, 12 Sep 2022 18:01:36 +0530 Subject: [PATCH 14/82] Making DatastoreRpc call aggregation query method on client --- .../google/cloud/datastore/spi/v1/DatastoreRpc.java | 9 +++++++++ .../cloud/datastore/spi/v1/HttpDatastoreRpc.java | 11 +++++++++++ 2 files changed, 20 insertions(+) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java index 5e64c9255..6b542da1a 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java @@ -30,6 +30,8 @@ import com.google.datastore.v1.ReserveIdsResponse; import com.google.datastore.v1.RollbackRequest; import com.google.datastore.v1.RollbackResponse; +import com.google.datastore.v1.RunAggregationQueryRequest; +import com.google.datastore.v1.RunAggregationQueryResponse; import com.google.datastore.v1.RunQueryRequest; import com.google.datastore.v1.RunQueryResponse; @@ -85,4 +87,11 @@ BeginTransactionResponse beginTransaction(BeginTransactionRequest request) * @throws DatastoreException upon failure */ RunQueryResponse runQuery(RunQueryRequest request); + + /** + * Sends a request to run an aggregation query. + * + * @throws DatastoreException upon failure + */ + RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java index 4f13b4600..fd3cdc658 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java @@ -36,6 +36,8 @@ import com.google.datastore.v1.ReserveIdsResponse; import com.google.datastore.v1.RollbackRequest; import com.google.datastore.v1.RollbackResponse; +import com.google.datastore.v1.RunAggregationQueryRequest; +import com.google.datastore.v1.RunAggregationQueryResponse; import com.google.datastore.v1.RunQueryRequest; import com.google.datastore.v1.RunQueryResponse; import java.io.IOException; @@ -200,4 +202,13 @@ public RunQueryResponse runQuery(RunQueryRequest request) { throw translate(ex); } } + + @Override + public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) { + try { + return client.runAggregationQuery(request); + } catch (com.google.datastore.v1.client.DatastoreException ex) { + throw translate(ex); + } + } } From 9589f5c8ceb467953fdf2f92024e82b0a1401a24 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Tue, 13 Sep 2022 10:58:41 +0530 Subject: [PATCH 15/82] Creating response transformer to transform aggregation query response into domain objects --- .../cloud/datastore/AggregationQuery.java | 2 +- .../cloud/datastore/AggregationResult.java | 25 ++++++ .../cloud/datastore/AggregationResults.java | 43 ++++++++++ .../AggregationQueryResponseTransformer.java | 54 +++++++++++++ .../response/ResponseTransformer.java | 20 +++++ ...gregationQueryResponseTransformerTest.java | 81 +++++++++++++++++++ 6 files changed, 224 insertions(+), 1 deletion(-) create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/ResponseTransformer.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java index db8825a65..1dfb11f42 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java @@ -25,7 +25,7 @@ import java.util.ArrayList; import java.util.List; -public class AggregationQuery extends Query { +public class AggregationQuery extends Query { private List aggregations; private StructuredQuery nestedStructuredQuery; diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java index bdf2b12f4..64a28cb04 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java @@ -15,6 +15,31 @@ */ package com.google.cloud.datastore; +import java.util.Map; +import java.util.Objects; + public class AggregationResult { + private final Map properties; + + public AggregationResult(Map properties) { + this.properties = properties; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + AggregationResult that = (AggregationResult) o; + return properties.equals(that.properties); + } + + @Override + public int hashCode() { + return Objects.hash(properties); + } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java new file mode 100644 index 000000000..97e783aaf --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java @@ -0,0 +1,43 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import com.google.api.core.InternalApi; +import java.util.Iterator; +import java.util.List; + +public class AggregationResults implements Iterable { + + private final List aggregationResults; + + public AggregationResults(List aggregationResults) { + this.aggregationResults = aggregationResults; + } + + @Override + public Iterator iterator() { + return this.aggregationResults.iterator(); + } + + public int size() { + return this.aggregationResults.size(); + } + + @InternalApi + public AggregationResult get(int index) { + return this.aggregationResults.get(index); + } +} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java new file mode 100644 index 000000000..bab27ab80 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java @@ -0,0 +1,54 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore.execution.response; + +import com.google.cloud.datastore.AggregationResult; +import com.google.cloud.datastore.AggregationResults; +import com.google.cloud.datastore.LongValue; +import com.google.datastore.v1.RunAggregationQueryResponse; +import com.google.datastore.v1.Value; +import java.util.AbstractMap.SimpleEntry; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.function.Function; +import java.util.stream.Collectors; + +public class AggregationQueryResponseTransformer implements + ResponseTransformer { + + @Override + public AggregationResults transform(RunAggregationQueryResponse response) { + List aggregationResults = response + .getBatch() + .getAggregationResultsList() + .stream() + .map(aggregationResult -> new AggregationResult(resultWithLongValues(aggregationResult))) + .collect(Collectors.toCollection(LinkedList::new)); + return new AggregationResults(aggregationResults); + } + + private Map resultWithLongValues(com.google.datastore.v1.AggregationResult aggregationResult) { + return aggregationResult.getAggregatePropertiesMap() + .entrySet() + .stream() + .map((Function, Entry>) entry -> + new SimpleEntry<>(entry.getKey(), (LongValue) LongValue.fromPb(entry.getValue()))) + .collect(Collectors.toMap(Entry::getKey, Entry::getValue)); + } +} + diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/ResponseTransformer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/ResponseTransformer.java new file mode 100644 index 000000000..3e5c10b66 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/ResponseTransformer.java @@ -0,0 +1,20 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore.execution.response; + +public interface ResponseTransformer { + OUTPUT transform(INPUT response); +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java new file mode 100644 index 000000000..5198c658f --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java @@ -0,0 +1,81 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore.execution.response; + +import static com.google.cloud.datastore.ProtoTestData.intValue; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +import com.google.cloud.datastore.AggregationResult; +import com.google.cloud.datastore.AggregationResults; +import com.google.cloud.datastore.LongValue; +import com.google.datastore.v1.AggregationResultBatch; +import com.google.datastore.v1.RunAggregationQueryResponse; +import com.google.datastore.v1.Value; +import java.util.AbstractMap.SimpleEntry; +import java.util.HashMap; +import java.util.Map; +import java.util.Map.Entry; +import java.util.function.Function; +import java.util.stream.Collectors; +import org.junit.Test; + +public class AggregationQueryResponseTransformerTest { + + + private AggregationQueryResponseTransformer responseTransformer = new AggregationQueryResponseTransformer(); + + @Test + public void shouldTransformAggregationQueryResponse() { + Map result1 = new HashMap() {{ + put("count", intValue(209)); + put("count_upto_100", intValue(100)); + }}; + + Map result2 = new HashMap() {{ + put("count", intValue(509)); + put("count_upto_100", intValue(100)); + }}; + + AggregationResultBatch resultBatch = AggregationResultBatch.newBuilder() + .addAggregationResults(com.google.datastore.v1.AggregationResult.newBuilder() + .putAllAggregateProperties(result1).build()) + .addAggregationResults(com.google.datastore.v1.AggregationResult.newBuilder() + .putAllAggregateProperties(result2).build()) + .build(); + RunAggregationQueryResponse runAggregationQueryResponse = RunAggregationQueryResponse.newBuilder() + .setBatch(resultBatch) + .build(); + + AggregationResults aggregationResults = responseTransformer.transform( + runAggregationQueryResponse); + + assertThat(aggregationResults.size(), equalTo(2)); + assertThat(aggregationResults.get(0), equalTo(new AggregationResult(toDomainValues(result1)))); + assertThat(aggregationResults.get(1), equalTo(new AggregationResult(toDomainValues(result2)))); + + } + + private Map toDomainValues(Map map) { + + return map + .entrySet() + .stream() + .map((Function, Entry>) entry -> + new SimpleEntry<>(entry.getKey(), (LongValue) LongValue.fromPb(entry.getValue()))) + .collect(Collectors.toMap(Entry::getKey, Entry::getValue)); + } +} \ No newline at end of file From 2103b800caed606d5a9c79d4c1db3b4f78b6991e Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Tue, 13 Sep 2022 13:47:10 +0530 Subject: [PATCH 16/82] Implementing aggregation query executor to execute AggergationQuery --- .../execution/AggregationQueryExecutor.java | 47 ++++++++ .../datastore/execution/QueryExecutor.java | 22 ++++ .../AggregationQueryExecutorTest.java | 103 ++++++++++++++++++ 3 files changed, 172 insertions(+) create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java new file mode 100644 index 000000000..0122b0f96 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java @@ -0,0 +1,47 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore.execution; + +import com.google.cloud.datastore.AggregationQuery; +import com.google.cloud.datastore.AggregationResults; +import com.google.cloud.datastore.Datastore; +import com.google.cloud.datastore.DatastoreOptions; +import com.google.cloud.datastore.execution.request.AggregationQueryRequestProtoPreparer; +import com.google.cloud.datastore.execution.response.AggregationQueryResponseTransformer; +import com.google.cloud.datastore.spi.v1.DatastoreRpc; +import com.google.datastore.v1.RunAggregationQueryRequest; +import com.google.datastore.v1.RunAggregationQueryResponse; + +public class AggregationQueryExecutor implements QueryExecutor{ + + private final DatastoreRpc datastoreRpc; + private final AggregationQueryRequestProtoPreparer protoPreparer; + private final AggregationQueryResponseTransformer responseTransformer; + + public AggregationQueryExecutor(DatastoreRpc datastoreRpc, DatastoreOptions datastoreOptions) { + this.datastoreRpc = datastoreRpc; + this.protoPreparer = new AggregationQueryRequestProtoPreparer(datastoreOptions); + this.responseTransformer = new AggregationQueryResponseTransformer(); + } + + @Override + public AggregationResults execute(AggregationQuery query) { + RunAggregationQueryRequest runAggregationQueryRequest = this.protoPreparer.prepare(query); + RunAggregationQueryResponse runAggregationQueryResponse = this.datastoreRpc.runAggregationQuery( + runAggregationQueryRequest); + return this.responseTransformer.transform(runAggregationQueryResponse); + } +} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java new file mode 100644 index 000000000..518e286f8 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java @@ -0,0 +1,22 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore.execution; + +import com.google.cloud.datastore.Query; + +public interface QueryExecutor, OUTPUT> { + OUTPUT execute(INPUT query); +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java new file mode 100644 index 000000000..cb30dd16a --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java @@ -0,0 +1,103 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore.execution; + +import static com.google.cloud.datastore.ProtoTestData.intValue; +import static com.google.cloud.datastore.StructuredQuery.PropertyFilter.eq; +import static com.google.cloud.datastore.aggregation.Aggregation.count; +import static org.easymock.EasyMock.anyObject; +import static org.easymock.EasyMock.expect; +import static org.easymock.EasyMock.replay; +import static org.easymock.EasyMock.verify; + +import com.google.cloud.datastore.AggregationQuery; +import com.google.cloud.datastore.AggregationResults; +import com.google.cloud.datastore.DatastoreOptions; +import com.google.cloud.datastore.EntityQuery; +import com.google.cloud.datastore.Query; +import com.google.cloud.datastore.spi.v1.DatastoreRpc; +import com.google.datastore.v1.AggregationResultBatch; +import com.google.datastore.v1.RunAggregationQueryRequest; +import com.google.datastore.v1.RunAggregationQueryResponse; +import com.google.datastore.v1.Value; +import java.util.HashMap; +import java.util.Map; +import org.easymock.EasyMock; +import org.junit.Before; +import org.junit.Test; + +public class AggregationQueryExecutorTest { + + private static final String KIND = "Task"; + private static final String NAMESPACE = "ns"; + + private DatastoreRpc mockRpc; + private AggregationQueryExecutor queryExecutor; + private DatastoreOptions datastoreOptions; + + @Before + public void setUp() throws Exception { + mockRpc = EasyMock.createStrictMock(DatastoreRpc.class); + datastoreOptions = DatastoreOptions.newBuilder().setNamespace(NAMESPACE).build(); + queryExecutor = new AggregationQueryExecutor(mockRpc, datastoreOptions); + } + + @Test + public void shouldExecuteAggregationQuery() { + EntityQuery nestedQuery = Query.newEntityQueryBuilder() + .setNamespace(NAMESPACE) + .setKind(KIND) + .setFilter(eq("done", true)) + .build(); + + AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .addAggregation(count().as("total")) + .over(nestedQuery) + .build(); + + expect(mockRpc.runAggregationQuery(anyObject(RunAggregationQueryRequest.class))).andReturn(dummyAggregationQueryResponse()); + + replay(mockRpc); + + AggregationResults aggregationResults = queryExecutor.execute(aggregationQuery); + + verify(mockRpc); + System.out.println(aggregationResults); + } + + private RunAggregationQueryResponse dummyAggregationQueryResponse() { + Map result1 = new HashMap() {{ + put("count", intValue(209)); + put("count_upto_100", intValue(100)); + }}; + + Map result2 = new HashMap() {{ + put("count", intValue(509)); + put("count_upto_100", intValue(100)); + }}; + + AggregationResultBatch resultBatch = AggregationResultBatch.newBuilder() + .addAggregationResults(com.google.datastore.v1.AggregationResult.newBuilder() + .putAllAggregateProperties(result1).build()) + .addAggregationResults(com.google.datastore.v1.AggregationResult.newBuilder() + .putAllAggregateProperties(result2).build()) + .build(); + return RunAggregationQueryResponse.newBuilder() + .setBatch(resultBatch) + .build(); + } +} \ No newline at end of file From 775c341ae5edfa5817f04a53f47770fb545c457a Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Tue, 13 Sep 2022 22:15:09 +0530 Subject: [PATCH 17/82] Adding missing assertion statements --- .../cloud/datastore/AggregationResults.java | 18 ++++++++++++++++++ .../AggregationQueryExecutorTest.java | 12 +++++++++++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java index 97e783aaf..bd6d905cb 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java @@ -18,6 +18,7 @@ import com.google.api.core.InternalApi; import java.util.Iterator; import java.util.List; +import java.util.Objects; public class AggregationResults implements Iterable { @@ -40,4 +41,21 @@ public int size() { public AggregationResult get(int index) { return this.aggregationResults.get(index); } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + AggregationResults that = (AggregationResults) o; + return Objects.equals(aggregationResults, that.aggregationResults); + } + + @Override + public int hashCode() { + return Objects.hash(aggregationResults); + } } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java index cb30dd16a..592b3415a 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java @@ -18,21 +18,28 @@ import static com.google.cloud.datastore.ProtoTestData.intValue; import static com.google.cloud.datastore.StructuredQuery.PropertyFilter.eq; import static com.google.cloud.datastore.aggregation.Aggregation.count; +import static java.util.Arrays.asList; import static org.easymock.EasyMock.anyObject; import static org.easymock.EasyMock.expect; import static org.easymock.EasyMock.replay; import static org.easymock.EasyMock.verify; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; import com.google.cloud.datastore.AggregationQuery; +import com.google.cloud.datastore.AggregationResult; import com.google.cloud.datastore.AggregationResults; import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.EntityQuery; +import com.google.cloud.datastore.LongValue; import com.google.cloud.datastore.Query; import com.google.cloud.datastore.spi.v1.DatastoreRpc; +import com.google.common.collect.ImmutableMap; import com.google.datastore.v1.AggregationResultBatch; import com.google.datastore.v1.RunAggregationQueryRequest; import com.google.datastore.v1.RunAggregationQueryResponse; import com.google.datastore.v1.Value; +import java.util.Arrays; import java.util.HashMap; import java.util.Map; import org.easymock.EasyMock; @@ -76,7 +83,10 @@ public void shouldExecuteAggregationQuery() { AggregationResults aggregationResults = queryExecutor.execute(aggregationQuery); verify(mockRpc); - System.out.println(aggregationResults); + assertThat(aggregationResults, equalTo(new AggregationResults(asList( + new AggregationResult(ImmutableMap.of("count", LongValue.of(209), "count_upto_100", LongValue.of(100))), + new AggregationResult(ImmutableMap.of("count", LongValue.of(509), "count_upto_100", LongValue.of(100))) + )))); } private RunAggregationQueryResponse dummyAggregationQueryResponse() { From 8872a55cd1b1e22643222aa62a17049169dee71f Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 14 Sep 2022 10:02:22 +0530 Subject: [PATCH 18/82] Creating RetryExecutor to inject it as a dependency in other components --- .../google/cloud/datastore/RetryExecutor.java | 31 +++++++++ .../cloud/datastore/RetryExecutorTest.java | 68 +++++++++++++++++++ 2 files changed, 99 insertions(+) create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java new file mode 100644 index 000000000..0d96c823f --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java @@ -0,0 +1,31 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import com.google.api.core.ApiClock; +import com.google.api.gax.retrying.ResultRetryAlgorithm; +import com.google.api.gax.retrying.RetrySettings; +import com.google.cloud.RetryHelper; +import java.util.concurrent.Callable; + +public class RetryExecutor { + + public T execute(Callable block, RetrySettings retrySettings, + ResultRetryAlgorithm resultRetryAlgorithm, ApiClock clock) { + return RetryHelper.runWithRetries(block, retrySettings, resultRetryAlgorithm, clock); + } + +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java new file mode 100644 index 000000000..79d3f3f4a --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java @@ -0,0 +1,68 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +import com.google.api.core.NanoClock; +import com.google.api.gax.retrying.BasicResultRetryAlgorithm; +import com.google.api.gax.retrying.RetrySettings; +import java.util.concurrent.Callable; +import org.junit.Test; + +public class RetryExecutorTest { + + private RetryExecutor retryExecutor = new RetryExecutor(); + + @Test + public void shouldRetryWhenErrorOccurred() { + RetrySettings retrySettings = RetrySettings.newBuilder() + .setMaxAttempts(4) + .build(); + CustomErrorCallable callable = new CustomErrorCallable(3); + + retryExecutor.execute(callable, retrySettings, new BasicResultRetryAlgorithm<>(), + NanoClock.getDefaultClock()); + + assertThat(callable.invokedTimes(), equalTo(4)); + } + + private static class CustomErrorCallable implements Callable { + + private final int errorThreshold; + private int numberOfInvocations; + + public CustomErrorCallable(int errorThreshold) { + this.errorThreshold = errorThreshold; + this.numberOfInvocations = 0; + } + + @Override + public String call() throws Exception { + this.numberOfInvocations++; + if (this.numberOfInvocations <= errorThreshold) { + throw new RuntimeException("error"); + } + return "success"; + } + + public int invokedTimes() { + return this.numberOfInvocations; + } + } +} + From 1dfafb7f8adcdd4ec566ac6fc712d2233623d33b Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 14 Sep 2022 10:18:32 +0530 Subject: [PATCH 19/82] Making RetryExecutor accept RetrySettings when creating it --- .../google/cloud/datastore/RetryExecutor.java | 18 ++++++++++++++---- .../cloud/datastore/RetryExecutorTest.java | 16 +++++++++------- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java index 0d96c823f..73a7b1364 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java @@ -15,7 +15,8 @@ */ package com.google.cloud.datastore; -import com.google.api.core.ApiClock; +import static com.google.cloud.RetryHelper.runWithRetries; + import com.google.api.gax.retrying.ResultRetryAlgorithm; import com.google.api.gax.retrying.RetrySettings; import com.google.cloud.RetryHelper; @@ -23,9 +24,18 @@ public class RetryExecutor { - public T execute(Callable block, RetrySettings retrySettings, - ResultRetryAlgorithm resultRetryAlgorithm, ApiClock clock) { - return RetryHelper.runWithRetries(block, retrySettings, resultRetryAlgorithm, clock); + private final RetrySettings retrySettings; + private final DatastoreOptions datastoreOptions; + + public RetryExecutor(RetrySettings retrySettings, + DatastoreOptions datastoreOptions) { + this.retrySettings = retrySettings; + this.datastoreOptions = datastoreOptions; + } + + public T execute(Callable block, ResultRetryAlgorithm resultRetryAlgorithm) { + return runWithRetries(block, retrySettings, resultRetryAlgorithm, + this.datastoreOptions.getClock()); } } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java index 79d3f3f4a..d5bcf656f 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java @@ -18,7 +18,6 @@ import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; -import com.google.api.core.NanoClock; import com.google.api.gax.retrying.BasicResultRetryAlgorithm; import com.google.api.gax.retrying.RetrySettings; import java.util.concurrent.Callable; @@ -26,17 +25,20 @@ public class RetryExecutorTest { - private RetryExecutor retryExecutor = new RetryExecutor(); + private static final int MAX_ATTEMPTS = 4; + + private final RetrySettings retrySettings = RetrySettings.newBuilder() + .setMaxAttempts(MAX_ATTEMPTS) + .build(); + private final DatastoreOptions datastoreOptions = DatastoreOptions.getDefaultInstance(); + private final RetryExecutor retryExecutor = new RetryExecutor(retrySettings, datastoreOptions); + @Test public void shouldRetryWhenErrorOccurred() { - RetrySettings retrySettings = RetrySettings.newBuilder() - .setMaxAttempts(4) - .build(); CustomErrorCallable callable = new CustomErrorCallable(3); - retryExecutor.execute(callable, retrySettings, new BasicResultRetryAlgorithm<>(), - NanoClock.getDefaultClock()); + retryExecutor.execute(callable, new BasicResultRetryAlgorithm<>()); assertThat(callable.invokedTimes(), equalTo(4)); } From f6bf0cef632fffee4e85d00bb616c3a0f1ce025a Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 14 Sep 2022 10:41:23 +0530 Subject: [PATCH 20/82] Revert "Making RetryExecutor accept RetrySettings when creating it" This reverts commit 1dfafb7f8adcdd4ec566ac6fc712d2233623d33b. --- .../google/cloud/datastore/RetryExecutor.java | 18 ++++-------------- .../cloud/datastore/RetryExecutorTest.java | 16 +++++++--------- 2 files changed, 11 insertions(+), 23 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java index 73a7b1364..0d96c823f 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java @@ -15,8 +15,7 @@ */ package com.google.cloud.datastore; -import static com.google.cloud.RetryHelper.runWithRetries; - +import com.google.api.core.ApiClock; import com.google.api.gax.retrying.ResultRetryAlgorithm; import com.google.api.gax.retrying.RetrySettings; import com.google.cloud.RetryHelper; @@ -24,18 +23,9 @@ public class RetryExecutor { - private final RetrySettings retrySettings; - private final DatastoreOptions datastoreOptions; - - public RetryExecutor(RetrySettings retrySettings, - DatastoreOptions datastoreOptions) { - this.retrySettings = retrySettings; - this.datastoreOptions = datastoreOptions; - } - - public T execute(Callable block, ResultRetryAlgorithm resultRetryAlgorithm) { - return runWithRetries(block, retrySettings, resultRetryAlgorithm, - this.datastoreOptions.getClock()); + public T execute(Callable block, RetrySettings retrySettings, + ResultRetryAlgorithm resultRetryAlgorithm, ApiClock clock) { + return RetryHelper.runWithRetries(block, retrySettings, resultRetryAlgorithm, clock); } } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java index d5bcf656f..79d3f3f4a 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java @@ -18,6 +18,7 @@ import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import com.google.api.core.NanoClock; import com.google.api.gax.retrying.BasicResultRetryAlgorithm; import com.google.api.gax.retrying.RetrySettings; import java.util.concurrent.Callable; @@ -25,20 +26,17 @@ public class RetryExecutorTest { - private static final int MAX_ATTEMPTS = 4; - - private final RetrySettings retrySettings = RetrySettings.newBuilder() - .setMaxAttempts(MAX_ATTEMPTS) - .build(); - private final DatastoreOptions datastoreOptions = DatastoreOptions.getDefaultInstance(); - private final RetryExecutor retryExecutor = new RetryExecutor(retrySettings, datastoreOptions); - + private RetryExecutor retryExecutor = new RetryExecutor(); @Test public void shouldRetryWhenErrorOccurred() { + RetrySettings retrySettings = RetrySettings.newBuilder() + .setMaxAttempts(4) + .build(); CustomErrorCallable callable = new CustomErrorCallable(3); - retryExecutor.execute(callable, new BasicResultRetryAlgorithm<>()); + retryExecutor.execute(callable, retrySettings, new BasicResultRetryAlgorithm<>(), + NanoClock.getDefaultClock()); assertThat(callable.invokedTimes(), equalTo(4)); } From deb63bf097ebe7b682f2122716ef5e348cf18509 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 14 Sep 2022 10:41:40 +0530 Subject: [PATCH 21/82] Revert "Creating RetryExecutor to inject it as a dependency in other components" This reverts commit 8872a55cd1b1e22643222aa62a17049169dee71f. --- .../google/cloud/datastore/RetryExecutor.java | 31 --------- .../cloud/datastore/RetryExecutorTest.java | 68 ------------------- 2 files changed, 99 deletions(-) delete mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java delete mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java deleted file mode 100644 index 0d96c823f..000000000 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java +++ /dev/null @@ -1,31 +0,0 @@ -/* - * Copyright 2022 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.google.cloud.datastore; - -import com.google.api.core.ApiClock; -import com.google.api.gax.retrying.ResultRetryAlgorithm; -import com.google.api.gax.retrying.RetrySettings; -import com.google.cloud.RetryHelper; -import java.util.concurrent.Callable; - -public class RetryExecutor { - - public T execute(Callable block, RetrySettings retrySettings, - ResultRetryAlgorithm resultRetryAlgorithm, ApiClock clock) { - return RetryHelper.runWithRetries(block, retrySettings, resultRetryAlgorithm, clock); - } - -} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java deleted file mode 100644 index 79d3f3f4a..000000000 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java +++ /dev/null @@ -1,68 +0,0 @@ -/* - * Copyright 2022 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.google.cloud.datastore; - -import static org.hamcrest.CoreMatchers.equalTo; -import static org.hamcrest.MatcherAssert.assertThat; - -import com.google.api.core.NanoClock; -import com.google.api.gax.retrying.BasicResultRetryAlgorithm; -import com.google.api.gax.retrying.RetrySettings; -import java.util.concurrent.Callable; -import org.junit.Test; - -public class RetryExecutorTest { - - private RetryExecutor retryExecutor = new RetryExecutor(); - - @Test - public void shouldRetryWhenErrorOccurred() { - RetrySettings retrySettings = RetrySettings.newBuilder() - .setMaxAttempts(4) - .build(); - CustomErrorCallable callable = new CustomErrorCallable(3); - - retryExecutor.execute(callable, retrySettings, new BasicResultRetryAlgorithm<>(), - NanoClock.getDefaultClock()); - - assertThat(callable.invokedTimes(), equalTo(4)); - } - - private static class CustomErrorCallable implements Callable { - - private final int errorThreshold; - private int numberOfInvocations; - - public CustomErrorCallable(int errorThreshold) { - this.errorThreshold = errorThreshold; - this.numberOfInvocations = 0; - } - - @Override - public String call() throws Exception { - this.numberOfInvocations++; - if (this.numberOfInvocations <= errorThreshold) { - throw new RuntimeException("error"); - } - return "success"; - } - - public int invokedTimes() { - return this.numberOfInvocations; - } - } -} - From 428e9e356365b5251d052f0518eb0530b8e7d391 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 14 Sep 2022 12:22:13 +0530 Subject: [PATCH 22/82] Introducing RetryAndTraceDatastoreRpcDecorator to have retry and traceability logic on top of another DatastoreRpc --- .../RetryAndTraceDatastoreRpcDecorator.java | 113 ++++++++++++++++++ .../com/google/cloud/datastore/TraceUtil.java | 1 + ...etryAndTraceDatastoreRpcDecoratorTest.java | 78 ++++++++++++ 3 files changed, 192 insertions(+) create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java new file mode 100644 index 000000000..ad3198fcd --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java @@ -0,0 +1,113 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import static com.google.cloud.BaseService.EXCEPTION_HANDLER; +import static com.google.cloud.datastore.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY; + +import com.google.api.gax.retrying.RetrySettings; +import com.google.cloud.RetryHelper; +import com.google.cloud.RetryHelper.RetryHelperException; +import com.google.cloud.datastore.spi.v1.DatastoreRpc; +import com.google.datastore.v1.AllocateIdsRequest; +import com.google.datastore.v1.AllocateIdsResponse; +import com.google.datastore.v1.BeginTransactionRequest; +import com.google.datastore.v1.BeginTransactionResponse; +import com.google.datastore.v1.CommitRequest; +import com.google.datastore.v1.CommitResponse; +import com.google.datastore.v1.LookupRequest; +import com.google.datastore.v1.LookupResponse; +import com.google.datastore.v1.ReserveIdsRequest; +import com.google.datastore.v1.ReserveIdsResponse; +import com.google.datastore.v1.RollbackRequest; +import com.google.datastore.v1.RollbackResponse; +import com.google.datastore.v1.RunAggregationQueryRequest; +import com.google.datastore.v1.RunAggregationQueryResponse; +import com.google.datastore.v1.RunQueryRequest; +import com.google.datastore.v1.RunQueryResponse; +import io.opencensus.common.Scope; +import io.opencensus.trace.Span; +import io.opencensus.trace.Status; +import java.util.concurrent.Callable; + +public class RetryAndTraceDatastoreRpcDecorator implements DatastoreRpc { + + private final DatastoreRpc datastoreRpc; + private final TraceUtil traceUtil; + private final RetrySettings retrySettings; + private final DatastoreOptions datastoreOptions; + + public RetryAndTraceDatastoreRpcDecorator(DatastoreRpc datastoreRpc, TraceUtil traceUtil, + RetrySettings retrySettings, DatastoreOptions datastoreOptions) { + this.datastoreRpc = datastoreRpc; + this.traceUtil = traceUtil; + this.retrySettings = retrySettings; + this.datastoreOptions = datastoreOptions; + } + + @Override + public AllocateIdsResponse allocateIds(AllocateIdsRequest request) { + throw new UnsupportedOperationException("Not implemented."); + } + + @Override + public BeginTransactionResponse beginTransaction(BeginTransactionRequest request) + throws DatastoreException { + throw new UnsupportedOperationException("Not implemented."); + } + + @Override + public CommitResponse commit(CommitRequest request) { + throw new UnsupportedOperationException("Not implemented."); + } + + @Override + public LookupResponse lookup(LookupRequest request) { + throw new UnsupportedOperationException("Not implemented."); + } + + @Override + public ReserveIdsResponse reserveIds(ReserveIdsRequest request) { + throw new UnsupportedOperationException("Not implemented."); + } + + @Override + public RollbackResponse rollback(RollbackRequest request) { + throw new UnsupportedOperationException("Not implemented."); + } + + @Override + public RunQueryResponse runQuery(RunQueryRequest request) { + throw new UnsupportedOperationException("Not implemented."); + } + + @Override + public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) { + return invokeRpc(() -> datastoreRpc.runAggregationQuery(request), SPAN_NAME_RUN_AGGREGATION_QUERY); + } + + public O invokeRpc(Callable block, String startSpan) { + Span span = traceUtil.startSpan(startSpan); + try (Scope scope = traceUtil.getTracer().withSpan(span)) { + return RetryHelper.runWithRetries(block, this.retrySettings, EXCEPTION_HANDLER, this.datastoreOptions.getClock()); + } catch (RetryHelperException e) { + span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); + throw DatastoreException.translateAndThrow(e); + } finally { + span.end(TraceUtil.END_SPAN_OPTIONS); + } + } +} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java index 1f28b2e80..e6909b4df 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java @@ -39,6 +39,7 @@ public class TraceUtil { static final String SPAN_NAME_RESERVEIDS = "CloudDatastoreOperation.reserveIds"; static final String SPAN_NAME_ROLLBACK = "CloudDatastoreOperation.rollback"; static final String SPAN_NAME_RUNQUERY = "CloudDatastoreOperation.runQuery"; + static final String SPAN_NAME_RUN_AGGREGATION_QUERY = "CloudDatastoreOperation.runAggregationQuery"; static final EndSpanOptions END_SPAN_OPTIONS = EndSpanOptions.builder().setSampleToLocalSpanStore(true).build(); diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java new file mode 100644 index 000000000..b92233187 --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java @@ -0,0 +1,78 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import static com.google.cloud.datastore.TraceUtil.END_SPAN_OPTIONS; +import static com.google.cloud.datastore.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY; +import static com.google.rpc.Code.UNAVAILABLE; +import static org.easymock.EasyMock.createNiceMock; +import static org.easymock.EasyMock.createStrictMock; +import static org.easymock.EasyMock.expect; +import static org.easymock.EasyMock.replay; +import static org.easymock.EasyMock.verify; +import static org.hamcrest.CoreMatchers.sameInstance; +import static org.hamcrest.MatcherAssert.assertThat; + +import com.google.api.gax.retrying.RetrySettings; +import com.google.cloud.datastore.spi.v1.DatastoreRpc; +import com.google.datastore.v1.RunAggregationQueryRequest; +import com.google.datastore.v1.RunAggregationQueryResponse; +import io.opencensus.trace.Span; +import io.opencensus.trace.Tracer; +import org.junit.Before; +import org.junit.Test; + +public class RetryAndTraceDatastoreRpcDecoratorTest { + + public static final int MAX_ATTEMPTS = 3; + private DatastoreRpc mockDatastoreRpc; + private TraceUtil mockTraceUtil; + private DatastoreOptions datastoreOptions = DatastoreOptions.getDefaultInstance(); + private RetrySettings retrySettings = RetrySettings.newBuilder() + .setMaxAttempts(MAX_ATTEMPTS) + .build(); + + private RetryAndTraceDatastoreRpcDecorator datastoreRpcDecorator; + + @Before + public void setUp() throws Exception { + mockDatastoreRpc = createStrictMock(DatastoreRpc.class); + mockTraceUtil = createStrictMock(TraceUtil.class); + datastoreRpcDecorator = new RetryAndTraceDatastoreRpcDecorator(mockDatastoreRpc, mockTraceUtil, retrySettings, datastoreOptions); + } + + @Test + public void testRunAggregationQuery() { + Span mockSpan = createStrictMock(Span.class); + RunAggregationQueryRequest aggregationQueryRequest = RunAggregationQueryRequest.getDefaultInstance(); + RunAggregationQueryResponse aggregationQueryResponse = RunAggregationQueryResponse.getDefaultInstance(); + + expect(mockDatastoreRpc.runAggregationQuery(aggregationQueryRequest)) + .andThrow(new DatastoreException(UNAVAILABLE.getNumber(), "API not accessible currently", UNAVAILABLE.name())).times(2) + .andReturn(aggregationQueryResponse); + expect(mockTraceUtil.startSpan(SPAN_NAME_RUN_AGGREGATION_QUERY)).andReturn(mockSpan); + expect(mockTraceUtil.getTracer()).andReturn(createNiceMock(Tracer.class)); + mockSpan.end(END_SPAN_OPTIONS); + + replay(mockDatastoreRpc, mockTraceUtil, mockSpan); + + RunAggregationQueryResponse actualAggregationQueryResponse = datastoreRpcDecorator.runAggregationQuery( + aggregationQueryRequest); + + assertThat(actualAggregationQueryResponse, sameInstance(aggregationQueryResponse)); + verify(mockDatastoreRpc, mockTraceUtil, mockSpan); + } +} \ No newline at end of file From c86913efe0f56b036abe498b7326e0ad9c143d20 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 14 Sep 2022 14:40:18 +0530 Subject: [PATCH 23/82] Extracting out the responsibility of preparing ReadOption in it's own ProtoPreparer --- .../google/cloud/datastore/DatastoreImpl.java | 34 +++-------- .../google/cloud/datastore/ReadOption.java | 9 +++ .../datastore/ReadOptionProtoPreparer.java | 55 +++++++++++++++++ .../ReadOptionProtoPreparerTest.java | 59 +++++++++++++++++++ 4 files changed, 131 insertions(+), 26 deletions(-) create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index 283788728..9075563fc 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -57,13 +57,17 @@ final class DatastoreImpl extends BaseService implements Datas TransactionExceptionHandler.build(); private static final ExceptionHandler TRANSACTION_OPERATION_EXCEPTION_HANDLER = TransactionOperationExceptionHandler.build(); - private final TraceUtil traceUtil = TraceUtil.getInstance();; + private final TraceUtil traceUtil = TraceUtil.getInstance(); + + private final ReadOptionProtoPreparer readOptionProtoPreparer; DatastoreImpl(DatastoreOptions options) { super(options); this.datastoreRpc = options.getDatastoreRpcV1(); retrySettings = MoreObjects.firstNonNull(options.getRetrySettings(), ServiceOptions.getNoRetrySettings()); + + readOptionProtoPreparer = new ReadOptionProtoPreparer(); } @Override @@ -339,33 +343,11 @@ public Iterator get(Iterable keys, ReadOption... options) { return get(toReadOptionsPb(options), Iterables.toArray(keys, Key.class)); } - private static com.google.datastore.v1.ReadOptions toReadOptionsPb(ReadOption... options) { - com.google.datastore.v1.ReadOptions readOptionsPb = null; + private com.google.datastore.v1.ReadOptions toReadOptionsPb(ReadOption... options) { if (options != null) { - Map, ReadOption> optionsByType = - ReadOption.asImmutableMap(options); - - if (optionsByType.containsKey(EventualConsistency.class) - && optionsByType.containsKey(ReadTime.class)) { - throw DatastoreException.throwInvalidRequest( - "Can not use eventual consistency read with read time."); - } - - if (optionsByType.containsKey(EventualConsistency.class)) { - readOptionsPb = - com.google.datastore.v1.ReadOptions.newBuilder() - .setReadConsistency(ReadConsistency.EVENTUAL) - .build(); - } - - if (optionsByType.containsKey(ReadTime.class)) { - readOptionsPb = - com.google.datastore.v1.ReadOptions.newBuilder() - .setReadTime(((ReadTime) optionsByType.get(ReadTime.class)).time().toProto()) - .build(); - } + return this.readOptionProtoPreparer.prepare(Arrays.asList(options)); } - return readOptionsPb; + return null; } @Override diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java index a30533e2d..75cc77c0d 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java @@ -20,6 +20,7 @@ import com.google.cloud.Timestamp; import com.google.common.collect.ImmutableMap; import java.io.Serializable; +import java.util.List; import java.util.Map; /** @@ -95,4 +96,12 @@ static Map, ReadOption> asImmutableMap(ReadOption... } return builder.buildOrThrow(); } + + static Map, ReadOption> asImmutableMap(List options) { + ImmutableMap.Builder, ReadOption> builder = ImmutableMap.builder(); + for (ReadOption option : options) { + builder.put(option.getClass(), option); + } + return builder.buildOrThrow(); + } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java new file mode 100644 index 000000000..27f1c31f6 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java @@ -0,0 +1,55 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import com.google.cloud.datastore.ReadOption.EventualConsistency; +import com.google.cloud.datastore.ReadOption.ReadTime; +import com.google.cloud.datastore.execution.request.ProtoPreparer; +import com.google.datastore.v1.ReadOptions; +import com.google.datastore.v1.ReadOptions.ReadConsistency; +import java.util.List; +import java.util.Map; + +public class ReadOptionProtoPreparer implements ProtoPreparer, ReadOptions> { + + @Override + public ReadOptions prepare(List options) { + com.google.datastore.v1.ReadOptions readOptionsPb = null; + if (options != null && !options.isEmpty()) { + Map, ReadOption> optionsByType = + ReadOption.asImmutableMap(options); + + if (optionsByType.containsKey(EventualConsistency.class) + && optionsByType.containsKey(ReadTime.class)) { + throw DatastoreException.throwInvalidRequest( + "Can not use eventual consistency read with read time."); + } + + if (optionsByType.containsKey(EventualConsistency.class)) { + readOptionsPb = ReadOptions.newBuilder() + .setReadConsistency(ReadConsistency.EVENTUAL) + .build(); + } + + if (optionsByType.containsKey(ReadTime.class)) { + readOptionsPb = ReadOptions.newBuilder() + .setReadTime(((ReadTime) optionsByType.get(ReadTime.class)).time().toProto()) + .build(); + } + } + return readOptionsPb; + } +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java new file mode 100644 index 000000000..0f2cac58d --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java @@ -0,0 +1,59 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import static com.google.cloud.datastore.ReadOption.eventualConsistency; +import static com.google.cloud.datastore.ReadOption.readTime; +import static com.google.datastore.v1.ReadOptions.ReadConsistency.EVENTUAL; +import static java.util.Collections.singletonList; +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.cloud.Timestamp; +import com.google.common.collect.ImmutableList; +import com.google.datastore.v1.ReadOptions; +import java.util.Arrays; +import java.util.Collections; +import org.junit.Test; + +public class ReadOptionProtoPreparerTest { + + private ReadOptionProtoPreparer protoPreparer = new ReadOptionProtoPreparer(); + + @Test + public void shouldThrowErrorWhenUsingEventualConsistencyWithReadTime() { + assertThrows("Can not use eventual consistency read with read time.", + DatastoreException.class, + () -> protoPreparer.prepare( + Arrays.asList(eventualConsistency(), readTime(Timestamp.now())))); + } + + @Test + public void shouldPrepareReadOptionsWithEventualConsistency() { + ReadOptions readOptions = protoPreparer.prepare(singletonList(eventualConsistency())); + + assertThat(readOptions.getReadConsistency(), is(EVENTUAL)); + } + + @Test + public void shouldPrepareReadOptionsWithReadTime() { + Timestamp timestamp = Timestamp.now(); + ReadOptions readOptions = protoPreparer.prepare(singletonList(readTime(timestamp))); + + assertThat(Timestamp.fromProto(readOptions.getReadTime()), is(timestamp)); + } +} \ No newline at end of file From b235935087431a14ea99be6b1af5450fd7dede48 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 14 Sep 2022 17:27:22 +0530 Subject: [PATCH 24/82] Making QueryExecutor to execute query with provided ReadOptions --- .../google/cloud/datastore/ReadOption.java | 37 +++++++++- .../execution/AggregationQueryExecutor.java | 20 ++++- .../datastore/execution/QueryExecutor.java | 3 +- .../AggregationQueryRequestProtoPreparer.java | 28 +++++-- .../ReadOptionProtoPreparerTest.java | 11 +++ .../com/google/cloud/datastore/TestUtils.java | 41 ++++++++++ .../AggregationQueryExecutorTest.java | 55 +++++++++++++- ...regationQueryRequestProtoPreparerTest.java | 74 +++++++++++++++---- 8 files changed, 238 insertions(+), 31 deletions(-) create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/TestUtils.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java index 75cc77c0d..2cf91d195 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java @@ -20,6 +20,7 @@ import com.google.cloud.Timestamp; import com.google.common.collect.ImmutableMap; import java.io.Serializable; +import java.util.Collections; import java.util.List; import java.util.Map; @@ -69,7 +70,8 @@ public Timestamp time() { } } - private ReadOption() {} + private ReadOption() { + } /** * Returns a {@code ReadOption} that specifies eventual consistency, allowing Datastore to return @@ -104,4 +106,37 @@ static Map, ReadOption> asImmutableMap(List> { + + Q query; + List readOptions; + + private QueryAndReadOptions(Q query, List readOptions) { + this.query = query; + this.readOptions = readOptions; + } + + private QueryAndReadOptions(Q query) { + this.query = query; + this.readOptions = Collections.emptyList(); + } + + public Q getQuery() { + return query; + } + + public List getReadOptions() { + return readOptions; + } + + public static > QueryAndReadOptions create(Q query) { + return new QueryAndReadOptions<>(query); + } + + public static > QueryAndReadOptions create(Q query, List readOptions) { + return new QueryAndReadOptions<>(query, readOptions); + } + } + } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java index 0122b0f96..446bfc887 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java @@ -17,15 +17,18 @@ import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.AggregationResults; -import com.google.cloud.datastore.Datastore; import com.google.cloud.datastore.DatastoreOptions; +import com.google.cloud.datastore.ReadOption; +import com.google.cloud.datastore.ReadOption.QueryAndReadOptions; import com.google.cloud.datastore.execution.request.AggregationQueryRequestProtoPreparer; import com.google.cloud.datastore.execution.response.AggregationQueryResponseTransformer; import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.datastore.v1.RunAggregationQueryRequest; import com.google.datastore.v1.RunAggregationQueryResponse; +import java.util.Arrays; -public class AggregationQueryExecutor implements QueryExecutor{ +public class AggregationQueryExecutor implements + QueryExecutor { private final DatastoreRpc datastoreRpc; private final AggregationQueryRequestProtoPreparer protoPreparer; @@ -38,10 +41,19 @@ public AggregationQueryExecutor(DatastoreRpc datastoreRpc, DatastoreOptions data } @Override - public AggregationResults execute(AggregationQuery query) { - RunAggregationQueryRequest runAggregationQueryRequest = this.protoPreparer.prepare(query); + public AggregationResults execute(AggregationQuery query, ReadOption... readOptions) { + RunAggregationQueryRequest runAggregationQueryRequest = getRunAggregationQueryRequest( + query, readOptions); RunAggregationQueryResponse runAggregationQueryResponse = this.datastoreRpc.runAggregationQuery( runAggregationQueryRequest); return this.responseTransformer.transform(runAggregationQueryResponse); } + + private RunAggregationQueryRequest getRunAggregationQueryRequest(AggregationQuery query, + ReadOption... readOptions) { + QueryAndReadOptions queryAndReadOptions = readOptions == null ? + QueryAndReadOptions.create(query) : + QueryAndReadOptions.create(query, Arrays.asList(readOptions)); + return this.protoPreparer.prepare(queryAndReadOptions); + } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java index 518e286f8..9bb2c3499 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java @@ -16,7 +16,8 @@ package com.google.cloud.datastore.execution; import com.google.cloud.datastore.Query; +import com.google.cloud.datastore.ReadOption; public interface QueryExecutor, OUTPUT> { - OUTPUT execute(INPUT query); + OUTPUT execute(INPUT query, ReadOption... readOptions); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java index c80ccd2dd..e4bdfcf74 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java @@ -20,41 +20,53 @@ import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.GqlQueryProtoPreparer; +import com.google.cloud.datastore.ReadOption; +import com.google.cloud.datastore.ReadOption.QueryAndReadOptions; +import com.google.cloud.datastore.ReadOptionProtoPreparer; import com.google.cloud.datastore.StructuredQueryProtoPreparer; import com.google.cloud.datastore.aggregation.Aggregation; import com.google.datastore.v1.GqlQuery; import com.google.datastore.v1.PartitionId; import com.google.datastore.v1.Query; +import com.google.datastore.v1.ReadOptions; import com.google.datastore.v1.RunAggregationQueryRequest; +import java.util.List; public class AggregationQueryRequestProtoPreparer implements - ProtoPreparer { + ProtoPreparer, RunAggregationQueryRequest> { private final DatastoreOptions datastoreOptions; private final StructuredQueryProtoPreparer structuredQueryProtoPreparer; private final GqlQueryProtoPreparer gqlQueryProtoPreparer; + private final ReadOptionProtoPreparer readOptionProtoPreparer; public AggregationQueryRequestProtoPreparer(DatastoreOptions datastoreOptions) { this.datastoreOptions = datastoreOptions; this.structuredQueryProtoPreparer = new StructuredQueryProtoPreparer(); this.gqlQueryProtoPreparer = new GqlQueryProtoPreparer(); + this.readOptionProtoPreparer = new ReadOptionProtoPreparer(); } @Override - public RunAggregationQueryRequest prepare(AggregationQuery aggregationQuery) { + public RunAggregationQueryRequest prepare(QueryAndReadOptions aggregationQueryAndReadOptions) { + AggregationQuery aggregationQuery = aggregationQueryAndReadOptions.getQuery(); + List readOptions = aggregationQueryAndReadOptions.getReadOptions(); PartitionId partitionId = getPartitionId(aggregationQuery); RunAggregationQueryRequest.Builder aggregationQueryRequestBuilder = RunAggregationQueryRequest.newBuilder() .setPartitionId(partitionId) .setProjectId(datastoreOptions.getProjectId()); if (aggregationQuery.getMode() == GQL) { - return aggregationQueryRequestBuilder - .setGqlQuery(buildGqlQuery(aggregationQuery)) - .build(); + aggregationQueryRequestBuilder.setGqlQuery(buildGqlQuery(aggregationQuery)); + } else { + aggregationQueryRequestBuilder.setAggregationQuery(getAggregationQuery(aggregationQuery)); } - return aggregationQueryRequestBuilder - .setAggregationQuery(getAggregationQuery(aggregationQuery)) - .build(); + + ReadOptions readOptionsPb = readOptionProtoPreparer.prepare(readOptions); + if (readOptionsPb != null) { + aggregationQueryRequestBuilder.setReadOptions(readOptionsPb); + } + return aggregationQueryRequestBuilder.build(); } private GqlQuery buildGqlQuery(AggregationQuery aggregationQuery) { diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java index 0f2cac58d..847d60fb6 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java @@ -21,6 +21,7 @@ import static java.util.Collections.singletonList; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; import com.google.cloud.Timestamp; @@ -56,4 +57,14 @@ public void shouldPrepareReadOptionsWithReadTime() { assertThat(Timestamp.fromProto(readOptions.getReadTime()), is(timestamp)); } + + @Test + public void shouldReturnNullWhenReadOptionsIsNull() { + assertNull(protoPreparer.prepare(null)); + } + + @Test + public void shouldReturnNullWhenReadOptionsIsAnEmptyList() { + assertNull(protoPreparer.prepare(ImmutableList.of())); + } } \ No newline at end of file diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/TestUtils.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/TestUtils.java new file mode 100644 index 000000000..9c1ea8ffd --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/TestUtils.java @@ -0,0 +1,41 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import com.google.datastore.v1.RunAggregationQueryRequest; +import java.util.function.Predicate; +import org.easymock.EasyMock; +import org.easymock.IArgumentMatcher; + +public class TestUtils { + + public static T matches(Predicate predicate) { + EasyMock.reportMatcher(new IArgumentMatcher() { + @Override + public boolean matches(Object argument) { + return predicate.test(((T) argument)); + } + + @Override + public void appendTo(StringBuffer buffer) { + buffer.append("matches(\"").append(predicate).append("\")"); + } + }); + return null; + } + + +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java index 592b3415a..34091e719 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java @@ -16,8 +16,11 @@ package com.google.cloud.datastore.execution; import static com.google.cloud.datastore.ProtoTestData.intValue; +import static com.google.cloud.datastore.ReadOption.eventualConsistency; import static com.google.cloud.datastore.StructuredQuery.PropertyFilter.eq; +import static com.google.cloud.datastore.TestUtils.matches; import static com.google.cloud.datastore.aggregation.Aggregation.count; +import static com.google.datastore.v1.ReadOptions.ReadConsistency.EVENTUAL; import static java.util.Arrays.asList; import static org.easymock.EasyMock.anyObject; import static org.easymock.EasyMock.expect; @@ -33,16 +36,22 @@ import com.google.cloud.datastore.EntityQuery; import com.google.cloud.datastore.LongValue; import com.google.cloud.datastore.Query; +import com.google.cloud.datastore.ReadOption; +import com.google.cloud.datastore.TestUtils; import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.common.collect.ImmutableMap; import com.google.datastore.v1.AggregationResultBatch; +import com.google.datastore.v1.ReadOptions; +import com.google.datastore.v1.ReadOptions.ReadConsistency; import com.google.datastore.v1.RunAggregationQueryRequest; import com.google.datastore.v1.RunAggregationQueryResponse; import com.google.datastore.v1.Value; import java.util.Arrays; import java.util.HashMap; import java.util.Map; +import java.util.function.Predicate; import org.easymock.EasyMock; +import org.easymock.IArgumentMatcher; import org.junit.Before; import org.junit.Test; @@ -76,7 +85,8 @@ public void shouldExecuteAggregationQuery() { .over(nestedQuery) .build(); - expect(mockRpc.runAggregationQuery(anyObject(RunAggregationQueryRequest.class))).andReturn(dummyAggregationQueryResponse()); + expect(mockRpc.runAggregationQuery(anyObject(RunAggregationQueryRequest.class))).andReturn( + dummyAggregationQueryResponse()); replay(mockRpc); @@ -84,8 +94,41 @@ public void shouldExecuteAggregationQuery() { verify(mockRpc); assertThat(aggregationResults, equalTo(new AggregationResults(asList( - new AggregationResult(ImmutableMap.of("count", LongValue.of(209), "count_upto_100", LongValue.of(100))), - new AggregationResult(ImmutableMap.of("count", LongValue.of(509), "count_upto_100", LongValue.of(100))) + new AggregationResult( + ImmutableMap.of("count", LongValue.of(209), "count_upto_100", LongValue.of(100))), + new AggregationResult( + ImmutableMap.of("count", LongValue.of(509), "count_upto_100", LongValue.of(100))) + )))); + } + + @Test + public void shouldExecuteAggregationQueryWithReadOptions() { + EntityQuery nestedQuery = Query.newEntityQueryBuilder() + .setNamespace(NAMESPACE) + .setKind(KIND) + .setFilter(eq("done", true)) + .build(); + + AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .addAggregation(count().as("total")) + .over(nestedQuery) + .build(); + + expect(mockRpc.runAggregationQuery(matches(runAggregationRequestWithEventualConsistency()))) + .andReturn(dummyAggregationQueryResponse()); + + replay(mockRpc); + + AggregationResults aggregationResults = queryExecutor.execute(aggregationQuery, + eventualConsistency()); + + verify(mockRpc); + assertThat(aggregationResults, equalTo(new AggregationResults(asList( + new AggregationResult( + ImmutableMap.of("count", LongValue.of(209), "count_upto_100", LongValue.of(100))), + new AggregationResult( + ImmutableMap.of("count", LongValue.of(509), "count_upto_100", LongValue.of(100))) )))); } @@ -110,4 +153,10 @@ private RunAggregationQueryResponse dummyAggregationQueryResponse() { .setBatch(resultBatch) .build(); } + + private Predicate runAggregationRequestWithEventualConsistency() { + return runAggregationQueryRequest -> + runAggregationQueryRequest.getReadOptions().getReadConsistency() == EVENTUAL; + } + } \ No newline at end of file diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java index c2cd0635a..98433e984 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java @@ -22,18 +22,25 @@ import static com.google.cloud.datastore.ProtoTestData.kind; import static com.google.cloud.datastore.ProtoTestData.propertyFilter; import static com.google.cloud.datastore.ProtoTestData.stringValue; +import static com.google.cloud.datastore.ReadOption.eventualConsistency; import static com.google.cloud.datastore.StructuredQuery.PropertyFilter.eq; import static com.google.cloud.datastore.aggregation.Aggregation.count; import static com.google.datastore.v1.PropertyFilter.Operator.EQUAL; +import static com.google.datastore.v1.ReadOptions.ReadConsistency.EVENTUAL; import static java.util.Arrays.asList; +import static java.util.Collections.singletonList; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import com.google.cloud.Timestamp; import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.EntityQuery; import com.google.cloud.datastore.GqlQuery; import com.google.cloud.datastore.Query; +import com.google.cloud.datastore.ReadOption; +import com.google.cloud.datastore.ReadOption.EventualConsistency; +import com.google.cloud.datastore.ReadOption.QueryAndReadOptions; import com.google.datastore.v1.GqlQueryParameter; import com.google.datastore.v1.RunAggregationQueryRequest; import java.util.HashMap; @@ -59,19 +66,26 @@ public class AggregationQueryRequestProtoPreparerTest { .addBinding(27) .build(); + private final AggregationQuery AGGREGATION_OVER_STRUCTURED_QUERY = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .addAggregation(count().as("total")) + .addAggregation(count().limit(100).as("total_upto_100")) + .over(COMPLETED_TASK_STRUCTURED_QUERY) + .build(); + + private final AggregationQuery AGGREGATION_OVER_GQL_QUERY = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .over(COMPLETED_TASK_GQL_QUERY) + .build(); + + private final AggregationQueryRequestProtoPreparer protoPreparer = new AggregationQueryRequestProtoPreparer( DATASTORE_OPTIONS); @Test public void shouldPrepareAggregationQueryRequestWithGivenStructuredQuery() { - AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() - .setNamespace(NAMESPACE) - .addAggregation(count().as("total")) - .addAggregation(count().limit(100).as("total_upto_100")) - .over(COMPLETED_TASK_STRUCTURED_QUERY) - .build(); - - RunAggregationQueryRequest runAggregationQueryRequest = protoPreparer.prepare(aggregationQuery); + RunAggregationQueryRequest runAggregationQueryRequest = protoPreparer.prepare( + QueryAndReadOptions.create(AGGREGATION_OVER_STRUCTURED_QUERY)); assertThat(runAggregationQueryRequest.getProjectId(), equalTo(PROJECT_ID)); @@ -92,12 +106,9 @@ public void shouldPrepareAggregationQueryRequestWithGivenStructuredQuery() { @Test public void shouldPrepareAggregationQueryRequestWithGivenGqlQuery() { - AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() - .setNamespace(NAMESPACE) - .over(COMPLETED_TASK_GQL_QUERY) - .build(); - - RunAggregationQueryRequest runAggregationQueryRequest = protoPreparer.prepare(aggregationQuery); + RunAggregationQueryRequest runAggregationQueryRequest = protoPreparer.prepare( + QueryAndReadOptions.create( + AGGREGATION_OVER_GQL_QUERY)); assertThat(runAggregationQueryRequest.getProjectId(), equalTo(PROJECT_ID)); @@ -115,4 +126,39 @@ public void shouldPrepareAggregationQueryRequestWithGivenGqlQuery() { gqlQueryParameter(intValue(27)) ))); } + + @Test + public void shouldPrepareReadOptionsWithGivenStructuredQuery() { + RunAggregationQueryRequest eventualConsistencyAggregationRequest = prepareQuery( + AGGREGATION_OVER_STRUCTURED_QUERY, eventualConsistency()); + assertThat(eventualConsistencyAggregationRequest.getReadOptions().getReadConsistency(), + equalTo(EVENTUAL)); + + Timestamp now = Timestamp.now(); + RunAggregationQueryRequest readTimeAggregationRequest = prepareQuery( + AGGREGATION_OVER_STRUCTURED_QUERY, ReadOption.readTime(now)); + assertThat(Timestamp.fromProto(readTimeAggregationRequest.getReadOptions().getReadTime()), + equalTo(now)); + } + + @Test + public void shouldPrepareReadOptionsWithGivenGqlQuery() { + RunAggregationQueryRequest eventualConsistencyAggregationRequest = prepareQuery( + AGGREGATION_OVER_GQL_QUERY, eventualConsistency()); + assertThat(eventualConsistencyAggregationRequest.getReadOptions().getReadConsistency(), + equalTo(EVENTUAL)); + + Timestamp now = Timestamp.now(); + RunAggregationQueryRequest readTimeAggregationRequest = prepareQuery( + AGGREGATION_OVER_GQL_QUERY, ReadOption.readTime(now)); + assertThat(Timestamp.fromProto(readTimeAggregationRequest.getReadOptions().getReadTime()), + equalTo(now)); + } + + private RunAggregationQueryRequest prepareQuery(AggregationQuery query, ReadOption readOption) { + return protoPreparer.prepare( + QueryAndReadOptions.create(query, + singletonList(readOption))); + } + } \ No newline at end of file From f1c3379fa4cf64256390fcee7891d69d5f9b515b Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 14 Sep 2022 17:49:33 +0530 Subject: [PATCH 25/82] Exposing readTime to the user --- .../cloud/datastore/AggregationResults.java | 10 +++++++++- .../AggregationQueryResponseTransformer.java | 4 +++- .../execution/AggregationQueryExecutorTest.java | 17 +++++++---------- ...AggregationQueryResponseTransformerTest.java | 7 +++++-- 4 files changed, 24 insertions(+), 14 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java index bd6d905cb..ebb8ad24b 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java @@ -16,6 +16,7 @@ package com.google.cloud.datastore; import com.google.api.core.InternalApi; +import com.google.cloud.Timestamp; import java.util.Iterator; import java.util.List; import java.util.Objects; @@ -23,9 +24,12 @@ public class AggregationResults implements Iterable { private final List aggregationResults; + private Timestamp readTime; - public AggregationResults(List aggregationResults) { + public AggregationResults(List aggregationResults, + Timestamp readTime) { this.aggregationResults = aggregationResults; + this.readTime = readTime; } @Override @@ -42,6 +46,10 @@ public AggregationResult get(int index) { return this.aggregationResults.get(index); } + public Timestamp getReadTime() { + return this.readTime; + } + @Override public boolean equals(Object o) { if (this == o) { diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java index bab27ab80..5b01b4185 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java @@ -15,6 +15,7 @@ */ package com.google.cloud.datastore.execution.response; +import com.google.cloud.Timestamp; import com.google.cloud.datastore.AggregationResult; import com.google.cloud.datastore.AggregationResults; import com.google.cloud.datastore.LongValue; @@ -33,13 +34,14 @@ public class AggregationQueryResponseTransformer implements @Override public AggregationResults transform(RunAggregationQueryResponse response) { + Timestamp readTime = Timestamp.fromProto(response.getBatch().getReadTime()); List aggregationResults = response .getBatch() .getAggregationResultsList() .stream() .map(aggregationResult -> new AggregationResult(resultWithLongValues(aggregationResult))) .collect(Collectors.toCollection(LinkedList::new)); - return new AggregationResults(aggregationResults); + return new AggregationResults(aggregationResults, readTime); } private Map resultWithLongValues(com.google.datastore.v1.AggregationResult aggregationResult) { diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java index 34091e719..f34257250 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java @@ -29,6 +29,7 @@ import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import com.google.cloud.Timestamp; import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.AggregationResult; import com.google.cloud.datastore.AggregationResults; @@ -36,22 +37,16 @@ import com.google.cloud.datastore.EntityQuery; import com.google.cloud.datastore.LongValue; import com.google.cloud.datastore.Query; -import com.google.cloud.datastore.ReadOption; -import com.google.cloud.datastore.TestUtils; import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.common.collect.ImmutableMap; import com.google.datastore.v1.AggregationResultBatch; -import com.google.datastore.v1.ReadOptions; -import com.google.datastore.v1.ReadOptions.ReadConsistency; import com.google.datastore.v1.RunAggregationQueryRequest; import com.google.datastore.v1.RunAggregationQueryResponse; import com.google.datastore.v1.Value; -import java.util.Arrays; import java.util.HashMap; import java.util.Map; import java.util.function.Predicate; import org.easymock.EasyMock; -import org.easymock.IArgumentMatcher; import org.junit.Before; import org.junit.Test; @@ -85,8 +80,9 @@ public void shouldExecuteAggregationQuery() { .over(nestedQuery) .build(); + RunAggregationQueryResponse runAggregationQueryResponse = dummyAggregationQueryResponse(); expect(mockRpc.runAggregationQuery(anyObject(RunAggregationQueryRequest.class))).andReturn( - dummyAggregationQueryResponse()); + runAggregationQueryResponse); replay(mockRpc); @@ -98,7 +94,7 @@ public void shouldExecuteAggregationQuery() { ImmutableMap.of("count", LongValue.of(209), "count_upto_100", LongValue.of(100))), new AggregationResult( ImmutableMap.of("count", LongValue.of(509), "count_upto_100", LongValue.of(100))) - )))); + ), Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime())))); } @Test @@ -115,8 +111,9 @@ public void shouldExecuteAggregationQueryWithReadOptions() { .over(nestedQuery) .build(); + RunAggregationQueryResponse runAggregationQueryResponse = dummyAggregationQueryResponse(); expect(mockRpc.runAggregationQuery(matches(runAggregationRequestWithEventualConsistency()))) - .andReturn(dummyAggregationQueryResponse()); + .andReturn(runAggregationQueryResponse); replay(mockRpc); @@ -129,7 +126,7 @@ public void shouldExecuteAggregationQueryWithReadOptions() { ImmutableMap.of("count", LongValue.of(209), "count_upto_100", LongValue.of(100))), new AggregationResult( ImmutableMap.of("count", LongValue.of(509), "count_upto_100", LongValue.of(100))) - )))); + ), Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime())))); } private RunAggregationQueryResponse dummyAggregationQueryResponse() { diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java index 5198c658f..a9570c1c6 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java @@ -19,6 +19,7 @@ import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import com.google.cloud.Timestamp; import com.google.cloud.datastore.AggregationResult; import com.google.cloud.datastore.AggregationResults; import com.google.cloud.datastore.LongValue; @@ -35,8 +36,7 @@ public class AggregationQueryResponseTransformerTest { - - private AggregationQueryResponseTransformer responseTransformer = new AggregationQueryResponseTransformer(); + private final AggregationQueryResponseTransformer responseTransformer = new AggregationQueryResponseTransformer(); @Test public void shouldTransformAggregationQueryResponse() { @@ -49,12 +49,14 @@ public void shouldTransformAggregationQueryResponse() { put("count", intValue(509)); put("count_upto_100", intValue(100)); }}; + Timestamp readTime = Timestamp.now(); AggregationResultBatch resultBatch = AggregationResultBatch.newBuilder() .addAggregationResults(com.google.datastore.v1.AggregationResult.newBuilder() .putAllAggregateProperties(result1).build()) .addAggregationResults(com.google.datastore.v1.AggregationResult.newBuilder() .putAllAggregateProperties(result2).build()) + .setReadTime(readTime.toProto()) .build(); RunAggregationQueryResponse runAggregationQueryResponse = RunAggregationQueryResponse.newBuilder() .setBatch(resultBatch) @@ -66,6 +68,7 @@ public void shouldTransformAggregationQueryResponse() { assertThat(aggregationResults.size(), equalTo(2)); assertThat(aggregationResults.get(0), equalTo(new AggregationResult(toDomainValues(result1)))); assertThat(aggregationResults.get(1), equalTo(new AggregationResult(toDomainValues(result2)))); + assertThat(aggregationResults.getReadTime(), equalTo(readTime)); } From 168b864f074405b0a1a75586c2f667a1acbffe11 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 14 Sep 2022 19:52:43 +0530 Subject: [PATCH 26/82] Ignoring runAggregationQuery method from clirr check --- google-cloud-datastore/clirr-ignored-differences.xml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/google-cloud-datastore/clirr-ignored-differences.xml b/google-cloud-datastore/clirr-ignored-differences.xml index 110f22f73..1cc969163 100644 --- a/google-cloud-datastore/clirr-ignored-differences.xml +++ b/google-cloud-datastore/clirr-ignored-differences.xml @@ -11,4 +11,9 @@ com.google.datastore.v1.ReserveIdsResponse reserveIds(com.google.datastore.v1.ReserveIdsRequest) 7012 + + com/google/cloud/datastore/spi/v1/DatastoreRpc + com.google.datastore.v1.RunAggregationQueryResponse runAggregationQuery(com.google.datastore.v1.RunAggregationQueryRequest) + 7012 + From 005429bd1a2dbc1dd1afe92594b8b9273e53b9ed Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Thu, 15 Sep 2022 11:15:38 +0530 Subject: [PATCH 27/82] Making readTime final --- .../java/com/google/cloud/datastore/AggregationResults.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java index ebb8ad24b..bfb8f5fb6 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java @@ -24,7 +24,7 @@ public class AggregationResults implements Iterable { private final List aggregationResults; - private Timestamp readTime; + private final Timestamp readTime; public AggregationResults(List aggregationResults, Timestamp readTime) { From dfbdf6009d40e13966ce9a2ff8d21446946dc484 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Thu, 15 Sep 2022 11:51:58 +0530 Subject: [PATCH 28/82] Allowing namespace to be optional in AggregationQuery --- .../cloud/datastore/AggregationQuery.java | 4 ++-- .../AggregationQueryRequestProtoPreparer.java | 13 ++++++----- .../cloud/datastore/AggregationQueryTest.java | 15 ++++++++----- ...regationQueryRequestProtoPreparerTest.java | 22 +++++++++++++++++++ 4 files changed, 42 insertions(+), 12 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java index 1dfb11f42..a4210babc 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java @@ -34,7 +34,7 @@ public class AggregationQuery extends Query { AggregationQuery(String namespace, List aggregations, StructuredQuery nestedQuery) { - super(checkNotNull(namespace)); + super(namespace); checkArgument(!aggregations.isEmpty(), "At least one aggregation is required for an aggregation query to run"); this.aggregations = aggregations; @@ -43,7 +43,7 @@ public class AggregationQuery extends Query { } AggregationQuery(String namespace, GqlQuery gqlQuery) { - super(checkNotNull(namespace)); + super(namespace); this.nestedGqlQuery = gqlQuery; this.mode = GQL; } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java index e4bdfcf74..3d1b468fa 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java @@ -48,7 +48,8 @@ public AggregationQueryRequestProtoPreparer(DatastoreOptions datastoreOptions) { } @Override - public RunAggregationQueryRequest prepare(QueryAndReadOptions aggregationQueryAndReadOptions) { + public RunAggregationQueryRequest prepare( + QueryAndReadOptions aggregationQueryAndReadOptions) { AggregationQuery aggregationQuery = aggregationQueryAndReadOptions.getQuery(); List readOptions = aggregationQueryAndReadOptions.getReadOptions(); PartitionId partitionId = getPartitionId(aggregationQuery); @@ -87,9 +88,11 @@ private com.google.datastore.v1.AggregationQuery getAggregationQuery( } private PartitionId getPartitionId(AggregationQuery aggregationQuery) { - return PartitionId.newBuilder() - .setProjectId(datastoreOptions.getProjectId()) - .setNamespaceId(aggregationQuery.getNamespace()) - .build(); + PartitionId.Builder builder = PartitionId.newBuilder() + .setProjectId(datastoreOptions.getProjectId()); + if (aggregationQuery.getNamespace() != null) { + builder.setNamespaceId(aggregationQuery.getNamespace()); + } + return builder.build(); } } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java index f6fa2ee37..e5e13a614 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java @@ -21,6 +21,7 @@ import static com.google.cloud.datastore.aggregation.Aggregation.count; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; import org.junit.Assert; @@ -61,11 +62,15 @@ public void testAggregationBuilder() { @Test public void testAggregationBuilderWithoutNamespace() { - assertThrows(NullPointerException.class, () -> - Query.newAggregationQueryBuilder() - .addAggregation(count().as("total")) - .over(COMPLETED_TASK_QUERY) - .build()); + AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() + .addAggregation(count().as("total")) + .over(COMPLETED_TASK_QUERY) + .build(); + + assertNull(aggregationQuery.getNamespace()); + assertThat(aggregationQuery.getAggregations().get(0), equalTo(count().as("total").build())); + assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY)); + assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED)); } @Test diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java index 98433e984..0d43f8099 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java @@ -30,7 +30,9 @@ import static java.util.Arrays.asList; import static java.util.Collections.singletonList; import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertNull; import com.google.cloud.Timestamp; import com.google.cloud.datastore.AggregationQuery; @@ -155,6 +157,26 @@ public void shouldPrepareReadOptionsWithGivenGqlQuery() { equalTo(now)); } + @Test + public void shouldPrepareAggregationQueryWithoutNamespace() { + AggregationQuery structuredQueryWithoutNamespace = Query.newAggregationQueryBuilder() + .addAggregation(count().as("total")) + .over(COMPLETED_TASK_STRUCTURED_QUERY) + .build(); + AggregationQuery gqlQueryWithoutNamespace = Query.newAggregationQueryBuilder() + .over(COMPLETED_TASK_GQL_QUERY) + .build(); + + RunAggregationQueryRequest runAggregationQueryFromStructuredQuery = protoPreparer.prepare( + QueryAndReadOptions.create(structuredQueryWithoutNamespace)); + RunAggregationQueryRequest runAggregationQueryFromGqlQuery = protoPreparer.prepare( + QueryAndReadOptions.create(gqlQueryWithoutNamespace)); + + assertThat(runAggregationQueryFromStructuredQuery.getPartitionId().getNamespaceId(), is("")); + assertThat(runAggregationQueryFromGqlQuery.getPartitionId().getNamespaceId(), is("")); + + } + private RunAggregationQueryRequest prepareQuery(AggregationQuery query, ReadOption readOption) { return protoPreparer.prepare( QueryAndReadOptions.create(query, From 0565ba66a36a30f0ab219c72da485c3ebaf00e18 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Thu, 15 Sep 2022 12:15:46 +0530 Subject: [PATCH 29/82] Add capability to fetch aggregation result by passing alias --- .../cloud/datastore/AggregationResult.java | 16 +++++++++ .../datastore/AggregationResultTest.java | 36 +++++++++++++++++++ 2 files changed, 52 insertions(+) create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java index 64a28cb04..0ccea0d9d 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java @@ -15,7 +15,10 @@ */ package com.google.cloud.datastore; +import com.google.common.base.MoreObjects; +import com.google.common.base.MoreObjects.ToStringHelper; import java.util.Map; +import java.util.Map.Entry; import java.util.Objects; public class AggregationResult { @@ -26,6 +29,10 @@ public AggregationResult(Map properties) { this.properties = properties; } + public Long get(String alias) { + return properties.get(alias).get(); + } + @Override public boolean equals(Object o) { if (this == o) { @@ -42,4 +49,13 @@ public boolean equals(Object o) { public int hashCode() { return Objects.hash(properties); } + + @Override + public String toString() { + ToStringHelper toStringHelper = MoreObjects.toStringHelper(this); + for (Entry entry : properties.entrySet()) { + toStringHelper.add(entry.getKey(), entry.getValue().get()); + } + return toStringHelper.toString(); + } } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java new file mode 100644 index 000000000..776fae8ab --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java @@ -0,0 +1,36 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; + +import com.google.common.collect.ImmutableMap; +import org.junit.Test; + +public class AggregationResultTest { + + @Test + public void shouldGetAggregationResultValueByAlias() { + AggregationResult aggregationResult = new AggregationResult(ImmutableMap.of( + "count", LongValue.of(45), + "count_upto_30", LongValue.of(30) + )); + + assertThat(aggregationResult.get("count"), equalTo(45L)); + assertThat(aggregationResult.get("count_upto_30"), equalTo(30L)); + } +} \ No newline at end of file From 0e8e7b97cfd786675867117afc71f9b9f52da012 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Thu, 15 Sep 2022 16:17:21 +0530 Subject: [PATCH 30/82] Implementing User facing datastore.runAggrgation method to run aggregation query --- .../com/google/cloud/datastore/Datastore.java | 2 ++ .../google/cloud/datastore/DatastoreImpl.java | 14 +++++++++++- .../google/cloud/datastore/DatastoreTest.java | 22 +++++++++++++++++++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java index bb115995e..06c5e3a91 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java @@ -461,4 +461,6 @@ interface TransactionCallable { * @throws DatastoreException upon failure */ QueryResults run(Query query, ReadOption... options); + + AggregationResults runAggregation(AggregationQuery query, ReadOption... options); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index 9075563fc..22415909a 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -24,6 +24,7 @@ import com.google.cloud.ServiceOptions; import com.google.cloud.datastore.ReadOption.EventualConsistency; import com.google.cloud.datastore.ReadOption.ReadTime; +import com.google.cloud.datastore.execution.AggregationQueryExecutor; import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; @@ -60,6 +61,7 @@ final class DatastoreImpl extends BaseService implements Datas private final TraceUtil traceUtil = TraceUtil.getInstance(); private final ReadOptionProtoPreparer readOptionProtoPreparer; + private final AggregationQueryExecutor aggregationQueryExecutor; DatastoreImpl(DatastoreOptions options) { super(options); @@ -68,6 +70,10 @@ final class DatastoreImpl extends BaseService implements Datas MoreObjects.firstNonNull(options.getRetrySettings(), ServiceOptions.getNoRetrySettings()); readOptionProtoPreparer = new ReadOptionProtoPreparer(); + aggregationQueryExecutor = new AggregationQueryExecutor( + new RetryAndTraceDatastoreRpcDecorator(datastoreRpc, traceUtil, retrySettings, options), + options + ); } @Override @@ -86,6 +92,7 @@ public Transaction newTransaction() { } static class ReadWriteTransactionCallable implements Callable { + private final Datastore datastore; private final TransactionCallable callable; private volatile TransactionOptions options; @@ -176,7 +183,7 @@ public T runInTransaction( @Override public QueryResults run(Query query) { - return run(null, query); + return run(null, query); } @Override @@ -190,6 +197,11 @@ QueryResults run(com.google.datastore.v1.ReadOptions readOptionsPb, Query return new QueryResultsImpl(this, readOptionsPb, (RecordQuery) query, query.getNamespace()); } + @Override + public AggregationResults runAggregation(AggregationQuery query, ReadOption... options){ + return aggregationQueryExecutor.execute(query, options); + } + com.google.datastore.v1.RunQueryResponse runQuery( final com.google.datastore.v1.RunQueryRequest requestPb) { Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RUNQUERY); diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java index 72067fd20..fb8d11dc7 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java @@ -16,9 +16,12 @@ package com.google.cloud.datastore; +import static com.google.cloud.datastore.aggregation.Aggregation.count; import static org.easymock.EasyMock.createStrictMock; import static org.easymock.EasyMock.replay; import static org.easymock.EasyMock.verify; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotEquals; @@ -37,6 +40,7 @@ import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.cloud.datastore.testing.LocalDatastoreHelper; import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; import com.google.common.collect.Iterators; import com.google.common.collect.Lists; import com.google.datastore.v1.BeginTransactionRequest; @@ -72,6 +76,7 @@ import org.junit.Assert; import org.junit.Before; import org.junit.BeforeClass; +import org.junit.Ignore; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -526,6 +531,23 @@ public void testGqlQueryPagination() throws DatastoreException { EasyMock.verify(rpcFactoryMock, rpcMock); } + @Test + @Ignore + public void testRunAggregationQuery() { + EntityQuery selectAllQuery = Query.newEntityQueryBuilder().build(); + AggregationQuery getCountQuery = Query.newAggregationQueryBuilder() + .addAggregation(count().as("total_count").limit(100)) + .over(selectAllQuery) + .build(); + AggregationResult resultBeforeInsert = Iterables.getOnlyElement(datastore.runAggregation(getCountQuery)); + assertThat(resultBeforeInsert.get("total_count"), equalTo(2L)); + + datastore.put(ENTITY3); + + AggregationResult resultAfterInsert = Iterables.getOnlyElement(datastore.runAggregation(getCountQuery)); + assertThat(resultAfterInsert.get("total_count"), equalTo(3L)); + } + @Test public void testRunStructuredQuery() { Query query = From fc0f0ddd9e05a8963fb910c76e1cd836e725f54b Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Thu, 15 Sep 2022 19:08:16 +0530 Subject: [PATCH 31/82] Add integration test for count aggregation --- .../cloud/datastore/it/ITDatastoreTest.java | 167 +++++++++++++++++- 1 file changed, 165 insertions(+), 2 deletions(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java index 869984932..689ee54d2 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java @@ -16,6 +16,10 @@ package com.google.cloud.datastore.it; +import static com.google.cloud.datastore.aggregation.Aggregation.count; +import static com.google.common.collect.Iterables.getOnlyElement; +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotEquals; @@ -26,6 +30,7 @@ import static org.junit.Assert.fail; import com.google.cloud.Timestamp; +import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.Batch; import com.google.cloud.datastore.BooleanValue; import com.google.cloud.datastore.Cursor; @@ -34,6 +39,7 @@ import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.DatastoreReaderWriter; import com.google.cloud.datastore.Entity; +import com.google.cloud.datastore.EntityQuery; import com.google.cloud.datastore.EntityValue; import com.google.cloud.datastore.FullEntity; import com.google.cloud.datastore.GqlQuery; @@ -59,8 +65,11 @@ import com.google.cloud.datastore.Transaction; import com.google.cloud.datastore.Value; import com.google.cloud.datastore.ValueType; +import com.google.cloud.datastore.aggregation.Aggregation; import com.google.cloud.datastore.testing.RemoteDatastoreHelper; +import com.google.common.base.Function; import com.google.common.base.Preconditions; +import com.google.common.collect.Iterables; import com.google.datastore.v1.TransactionOptions; import java.util.ArrayList; import java.util.Collections; @@ -68,6 +77,8 @@ import java.util.Iterator; import java.util.List; import java.util.Set; +import java.util.concurrent.Callable; +import java.util.function.Consumer; import org.junit.After; import org.junit.AfterClass; import org.junit.Before; @@ -154,9 +165,11 @@ public class ITDatastoreTest { .set("partial2", ENTITY2) .build(); - @Rule public Timeout globalTimeout = Timeout.seconds(100); + @Rule + public Timeout globalTimeout = Timeout.seconds(100); - @Rule public MultipleAttemptsRule multipleAttemptsRule = new MultipleAttemptsRule(3); + @Rule + public MultipleAttemptsRule multipleAttemptsRule = new MultipleAttemptsRule(3); @AfterClass public static void afterClass() { @@ -506,6 +519,91 @@ public void testRunGqlQueryWithCasting() throws InterruptedException { assertFalse(results3.hasNext()); } + @Test + public void testRunAggregationQuery() { + // verifying aggregation with an entity query + testCountAggregationWith(builder -> + builder + .addAggregation(count().as("total_count")) + .over(Query.newEntityQueryBuilder() + .setNamespace(NAMESPACE) + .setKind(KIND1) + .build() + )); + + // verifying aggregation with a projection query + testCountAggregationWith(builder -> + builder + .addAggregation(count().as("total_count")) + .over(Query.newProjectionEntityQueryBuilder() + .setProjection("str") + .setNamespace(NAMESPACE) + .setKind(KIND1) + .build() + )); + + // verifying aggregation with a key query + testCountAggregationWith(builder -> + builder + .addAggregation(count().as("total_count")) + .over(Query.newKeyQueryBuilder() + .setNamespace(NAMESPACE) + .setKind(KIND1) + .build() + )); + + // verifying aggregation with a GQL query + testCountAggregationWith(builder -> + builder + .over(Query.newGqlQueryBuilder( + "AGGREGATE COUNT(*) AS total_count OVER (SELECT * FROM kind1)") + .setNamespace(NAMESPACE) + .build() + )); + } + + @Test + public void testRunAggregationQueryWithReadTime() throws InterruptedException { + String alias = "total_count"; + + // verifying aggregation readTime with an entity query + testCountAggregationReadTimeWith(builder -> + builder + .over( + Query.newEntityQueryBuilder().setKind("new_kind").build()) + .addAggregation(count().as(alias)) + ); + + // verifying aggregation readTime with a projection query + testCountAggregationReadTimeWith(builder -> + builder + .over( + Query.newProjectionEntityQueryBuilder() + .setProjection("name") + .setKind("new_kind").build()) + .addAggregation(count().as(alias)) + ); + + // verifying aggregation readTime with a key query + testCountAggregationReadTimeWith(builder -> + builder + .over( + Query.newKeyQueryBuilder() + .setKind("new_kind").build()) + .addAggregation(count().as(alias)) + ); + + // verifying aggregation readTime with a GQL query + testCountAggregationReadTimeWith(builder -> + builder + .over( + Query.newGqlQueryBuilder( + "AGGREGATE COUNT(*) AS total_count OVER (SELECT * FROM new_kind)") + .build()) + .addAggregation(count().as(alias)) + ); + } + @Test public void testRunStructuredQuery() throws InterruptedException { Query query = @@ -1067,4 +1165,69 @@ public void testQueryWithReadTime() throws InterruptedException { DATASTORE.delete(entity1.getKey(), entity2.getKey(), entity3.getKey()); } } + + private void testCountAggregationWith(Consumer configurer) { + AggregationQuery.Builder builder = Query.newAggregationQueryBuilder().setNamespace(NAMESPACE); + configurer.accept(builder); + AggregationQuery aggregationQuery = builder.build(); + String alias = "total_count"; + + Long countBeforeAdd = getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get(alias); + long expectedCount = countBeforeAdd + 1; + + Entity newEntity = Entity.newBuilder(ENTITY1) + .setKey(Key.newBuilder(KEY3, KIND1, 1).build()) + .set("null", NULL_VALUE) + .set("partial1", PARTIAL_ENTITY2) + .set("partial2", ENTITY2) + .build(); + DATASTORE.put(newEntity); + + Long countAfterAdd = getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get(alias); + assertThat(countAfterAdd, equalTo(expectedCount)); + + DATASTORE.delete(newEntity.getKey()); + } + + private void testCountAggregationReadTimeWith(Consumer configurer) + throws InterruptedException { + Entity entity1 = + Entity.newBuilder( + Key.newBuilder(PROJECT_ID, "new_kind", "name-01").setNamespace(NAMESPACE).build()) + .set("name", "Tyrion Lannister") + .build(); + Entity entity2 = + Entity.newBuilder( + Key.newBuilder(PROJECT_ID, "new_kind", "name-02").setNamespace(NAMESPACE).build()) + .set("name", "Jaime Lannister") + .build(); + Entity entity3 = + Entity.newBuilder( + Key.newBuilder(PROJECT_ID, "new_kind", "name-03").setNamespace(NAMESPACE).build()) + .set("name", "Cersei Lannister") + .build(); + + DATASTORE.put(entity1, entity2); + Thread.sleep(1000); + Timestamp now = Timestamp.now(); + Thread.sleep(1000); + DATASTORE.put(entity3); + + try { + AggregationQuery.Builder builder = Query.newAggregationQueryBuilder().setNamespace(NAMESPACE); + configurer.accept(builder); + AggregationQuery countAggregationQuery = builder.build(); + + Long latestCount = getOnlyElement(DATASTORE.runAggregation(countAggregationQuery)) + .get("total_count"); + assertThat(latestCount, equalTo(3L)); + + Long oldCount = getOnlyElement( + DATASTORE.runAggregation(countAggregationQuery, ReadOption.readTime(now)) + ).get("total_count"); + assertThat(oldCount, equalTo(2L)); + } finally { + DATASTORE.delete(entity1.getKey(), entity2.getKey(), entity3.getKey()); + } + } } From e6cbca6a96a824ea71bde91ca4d286b115614206 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Mon, 19 Sep 2022 11:52:05 +0530 Subject: [PATCH 32/82] Add transaction Id support in ReadOptionsProtoPreparer --- .../google/cloud/datastore/ReadOption.java | 20 +++++++++++++++ .../datastore/ReadOptionProtoPreparer.java | 25 ++++++++++++++++--- .../ReadOptionProtoPreparerTest.java | 25 ++++++++++++++++--- 3 files changed, 64 insertions(+), 6 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java index 2cf91d195..5934c098a 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java @@ -17,8 +17,10 @@ package com.google.cloud.datastore; import com.google.api.core.BetaApi; +import com.google.api.core.InternalApi; import com.google.cloud.Timestamp; import com.google.common.collect.ImmutableMap; +import com.google.protobuf.ByteString; import java.io.Serializable; import java.util.Collections; import java.util.List; @@ -70,6 +72,19 @@ public Timestamp time() { } } + @InternalApi + static class TransactionId extends ReadOption { + private final ByteString transactionId; + + TransactionId(ByteString transactionId) { + this.transactionId = transactionId; + } + + public ByteString getTransactionId() { + return transactionId; + } + } + private ReadOption() { } @@ -91,6 +106,11 @@ public static ReadTime readTime(Timestamp time) { return new ReadTime(time); } + @InternalApi + public static ReadOption transactionId(String transactionId) { + return new TransactionId(ByteString.copyFrom(transactionId.getBytes())); + } + static Map, ReadOption> asImmutableMap(ReadOption... options) { ImmutableMap.Builder, ReadOption> builder = ImmutableMap.builder(); for (ReadOption option : options) { diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java index 27f1c31f6..b00ad2bad 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java @@ -17,11 +17,14 @@ import com.google.cloud.datastore.ReadOption.EventualConsistency; import com.google.cloud.datastore.ReadOption.ReadTime; +import com.google.cloud.datastore.ReadOption.TransactionId; import com.google.cloud.datastore.execution.request.ProtoPreparer; import com.google.datastore.v1.ReadOptions; import com.google.datastore.v1.ReadOptions.ReadConsistency; import java.util.List; import java.util.Map; +import java.util.function.Function; +import java.util.stream.Collectors; public class ReadOptionProtoPreparer implements ProtoPreparer, ReadOptions> { @@ -32,10 +35,10 @@ public ReadOptions prepare(List options) { Map, ReadOption> optionsByType = ReadOption.asImmutableMap(options); - if (optionsByType.containsKey(EventualConsistency.class) - && optionsByType.containsKey(ReadTime.class)) { + boolean moreThanOneReadOption = optionsByType.keySet().size() > 1; + if (moreThanOneReadOption) { throw DatastoreException.throwInvalidRequest( - "Can not use eventual consistency read with read time."); + String.format("Can not use %s together.", getInvalidOptions(optionsByType))); } if (optionsByType.containsKey(EventualConsistency.class)) { @@ -49,7 +52,23 @@ public ReadOptions prepare(List options) { .setReadTime(((ReadTime) optionsByType.get(ReadTime.class)).time().toProto()) .build(); } + + if (optionsByType.containsKey(TransactionId.class)) { + readOptionsPb = ReadOptions.newBuilder() + .setTransaction(((TransactionId) optionsByType.get(TransactionId.class)).getTransactionId()) + .build(); + } } return readOptionsPb; } + + private String getInvalidOptions(Map, ReadOption> optionsByType) { + String regex = "([a-z])([A-Z]+)"; + String replacement = "$1 $2"; + return optionsByType + .keySet() + .stream().map(Class::getSimpleName) + .map(s -> s.replaceAll(regex, replacement).toLowerCase()) + .collect(Collectors.joining(", ")); + } } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java index 847d60fb6..90df8885e 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java @@ -17,6 +17,7 @@ import static com.google.cloud.datastore.ReadOption.eventualConsistency; import static com.google.cloud.datastore.ReadOption.readTime; +import static com.google.cloud.datastore.ReadOption.transactionId; import static com.google.datastore.v1.ReadOptions.ReadConsistency.EVENTUAL; import static java.util.Collections.singletonList; import static org.hamcrest.CoreMatchers.is; @@ -27,6 +28,7 @@ import com.google.cloud.Timestamp; import com.google.common.collect.ImmutableList; import com.google.datastore.v1.ReadOptions; +import com.google.protobuf.ByteString; import java.util.Arrays; import java.util.Collections; import org.junit.Test; @@ -36,11 +38,19 @@ public class ReadOptionProtoPreparerTest { private ReadOptionProtoPreparer protoPreparer = new ReadOptionProtoPreparer(); @Test - public void shouldThrowErrorWhenUsingEventualConsistencyWithReadTime() { - assertThrows("Can not use eventual consistency read with read time.", - DatastoreException.class, + public void shouldThrowErrorWhenUsingMultipleReadOptions() { + assertThrows(DatastoreException.class, () -> protoPreparer.prepare( Arrays.asList(eventualConsistency(), readTime(Timestamp.now())))); + assertThrows(DatastoreException.class, + () -> protoPreparer.prepare( + Arrays.asList(eventualConsistency(), transactionId("transaction-id")))); + assertThrows(DatastoreException.class, + () -> protoPreparer.prepare( + Arrays.asList(transactionId("transaction-id"), readTime(Timestamp.now())))); + assertThrows(DatastoreException.class, + () -> protoPreparer.prepare( + Arrays.asList(eventualConsistency(), readTime(Timestamp.now()), transactionId("transaction-id")))); } @Test @@ -58,6 +68,15 @@ public void shouldPrepareReadOptionsWithReadTime() { assertThat(Timestamp.fromProto(readOptions.getReadTime()), is(timestamp)); } + @Test + public void shouldPrepareReadOptionsWithTransactionId() { + String dummyTransactionId = "transaction-id"; + ReadOptions readOptions = protoPreparer.prepare(singletonList(transactionId( + dummyTransactionId))); + + assertThat(readOptions.getTransaction().toStringUtf8(), is(dummyTransactionId)); + } + @Test public void shouldReturnNullWhenReadOptionsIsNull() { assertNull(protoPreparer.prepare(null)); From d205997c44ae4372830dc92cb80f44187a02a86a Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Mon, 19 Sep 2022 15:37:23 +0530 Subject: [PATCH 33/82] Supporting aggregation query with transactions --- .../google/cloud/datastore/DatastoreImpl.java | 5 +++ .../cloud/datastore/DatastoreReader.java | 3 ++ .../google/cloud/datastore/ReadOption.java | 5 +++ .../cloud/datastore/TransactionImpl.java | 7 ++++ .../cloud/datastore/it/ITDatastoreTest.java | 32 ++++++++++++++++--- 5 files changed, 48 insertions(+), 4 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index 22415909a..9c2af79b6 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -197,6 +197,11 @@ QueryResults run(com.google.datastore.v1.ReadOptions readOptionsPb, Query return new QueryResultsImpl(this, readOptionsPb, (RecordQuery) query, query.getNamespace()); } + @Override + public AggregationResults runAggregation(AggregationQuery query){ + return aggregationQueryExecutor.execute(query); + } + @Override public AggregationResults runAggregation(AggregationQuery query, ReadOption... options){ return aggregationQueryExecutor.execute(query, options); diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java index 3d5b7cd3e..34ace2182 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java @@ -53,4 +53,7 @@ public interface DatastoreReader { * @throws DatastoreException upon failure */ QueryResults run(Query query); + + + AggregationResults runAggregation(AggregationQuery query); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java index 5934c098a..200519bac 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java @@ -111,6 +111,11 @@ public static ReadOption transactionId(String transactionId) { return new TransactionId(ByteString.copyFrom(transactionId.getBytes())); } + @InternalApi + public static ReadOption transactionId(ByteString transactionId) { + return new TransactionId(transactionId); + } + static Map, ReadOption> asImmutableMap(ReadOption... options) { ImmutableMap.Builder, ReadOption> builder = ImmutableMap.builder(); for (ReadOption option : options) { diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java index 3318ec866..c13439460 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java @@ -16,6 +16,8 @@ package com.google.cloud.datastore; +import static com.google.cloud.datastore.ReadOption.transactionId; + import com.google.datastore.v1.TransactionOptions; import com.google.protobuf.ByteString; import java.util.ArrayList; @@ -96,6 +98,11 @@ public QueryResults run(Query query) { return datastore.run(readOptionsPb.build(), query); } + @Override + public AggregationResults runAggregation(AggregationQuery query) { + return datastore.runAggregation(query, transactionId(transactionId)); + } + @Override public Transaction.Response commit() { validateActive(); diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java index 689ee54d2..708d44dd4 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java @@ -65,11 +65,8 @@ import com.google.cloud.datastore.Transaction; import com.google.cloud.datastore.Value; import com.google.cloud.datastore.ValueType; -import com.google.cloud.datastore.aggregation.Aggregation; import com.google.cloud.datastore.testing.RemoteDatastoreHelper; -import com.google.common.base.Function; import com.google.common.base.Preconditions; -import com.google.common.collect.Iterables; import com.google.datastore.v1.TransactionOptions; import java.util.ArrayList; import java.util.Collections; @@ -77,7 +74,6 @@ import java.util.Iterator; import java.util.List; import java.util.Set; -import java.util.concurrent.Callable; import java.util.function.Consumer; import org.junit.After; import org.junit.AfterClass; @@ -562,6 +558,34 @@ public void testRunAggregationQuery() { )); } + @Test + public void testRunAggregationQueryInTransaction() { + EntityQuery entityQuery = Query.newEntityQueryBuilder() + .setNamespace(NAMESPACE) + .setFilter(PropertyFilter.hasAncestor(KEY1)) + .build(); + + AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .over(entityQuery) + .addAggregation(count().as("count")) + .build(); + + Transaction transaction = DATASTORE.newTransaction(); + assertThat(getOnlyElement(transaction.runAggregation(aggregationQuery)).get("count"), equalTo(2L)); + + Entity aNewEntity = Entity.newBuilder(ENTITY2) + .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build()) + .set("v_int", 10) + .build(); + transaction.put(aNewEntity); + + assertThat(getOnlyElement(transaction.runAggregation(aggregationQuery)).get("count"), equalTo(2L)); + assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(2L)); + transaction.commit(); + assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(3L)); + } + @Test public void testRunAggregationQueryWithReadTime() throws InterruptedException { String alias = "total_count"; From 5097cd69fc6c1684103fa4500c289b86fe81d1ce Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Mon, 19 Sep 2022 17:11:25 +0530 Subject: [PATCH 34/82] Allowing user to create Aggregation directly without involving its builder --- .../cloud/datastore/AggregationQuery.java | 5 ++++ .../aggregation/CountAggregation.java | 2 +- .../cloud/datastore/AggregationQueryTest.java | 28 +++++++++++++++---- 3 files changed, 28 insertions(+), 7 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java index a4210babc..629fe04ea 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java @@ -86,6 +86,11 @@ public Builder addAggregation(AggregationBuilder aggregationBuilder) { return this; } + public Builder addAggregation(Aggregation aggregation) { + this.aggregations.add(aggregation); + return this; + } + public Builder over(StructuredQuery nestedQuery) { this.nestedStructuredQuery = nestedQuery; this.mode = STRUCTURED; diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java index 7cb7d57c4..16312cab9 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java @@ -27,7 +27,7 @@ public class CountAggregation extends Aggregation { private final long limit; - private CountAggregation(String alias, long limit) { + public CountAggregation(String alias, long limit) { super(alias); this.limit = limit; } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java index e5e13a614..ba3a82a52 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java @@ -24,6 +24,7 @@ import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; +import com.google.cloud.datastore.aggregation.CountAggregation; import org.junit.Assert; import org.junit.Rule; import org.junit.Test; @@ -43,6 +44,22 @@ public class AggregationQueryTest { @Rule public ExpectedException exceptionRule = ExpectedException.none(); + @Test + public void testAggregations() { + AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .addAggregation(new CountAggregation("total_upto_100", 100)) + .over(COMPLETED_TASK_QUERY) + .build(); + + assertThat(aggregationQuery.getNamespace(), equalTo(NAMESPACE)); + assertThat(aggregationQuery.getAggregations().get(0), + equalTo(count().limit(100).as("total_upto_100").build())); + assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY)); + assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED)); + } + + @Test public void testAggregationBuilder() { AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() @@ -61,7 +78,7 @@ public void testAggregationBuilder() { } @Test - public void testAggregationBuilderWithoutNamespace() { + public void testAggregationQueryBuilderWithoutNamespace() { AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() .addAggregation(count().as("total")) .over(COMPLETED_TASK_QUERY) @@ -74,7 +91,7 @@ public void testAggregationBuilderWithoutNamespace() { } @Test - public void testAggregationBuilderWithoutNestedQuery() { + public void testAggregationQueryBuilderWithoutNestedQuery() { assertThrows("Nested query is required for an aggregation query to run", IllegalArgumentException.class, () -> Query.newAggregationQueryBuilder() @@ -84,7 +101,7 @@ public void testAggregationBuilderWithoutNestedQuery() { } @Test - public void testAggregationBuilderWithoutAggregation() { + public void testAggregationQueryBuilderWithoutAggregation() { assertThrows("At least one aggregation is required for an aggregation query to run", IllegalArgumentException.class, () -> Query.newAggregationQueryBuilder() @@ -94,7 +111,7 @@ public void testAggregationBuilderWithoutAggregation() { } @Test - public void testAggregationBuilderWithGqlQuery() { + public void testAggregationQueryBuilderWithGqlQuery() { GqlQuery gqlQuery = Query.newGqlQueryBuilder("SELECT * FROM Task WHERE done = true").build(); AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() @@ -107,12 +124,11 @@ public void testAggregationBuilderWithGqlQuery() { } @Test - public void testAggregationBuilderWithoutProvidingAnyNestedQuery() { + public void testAggregationQueryBuilderWithoutProvidingAnyNestedQuery() { assertThrows("Nested query is required for an aggregation query to run", IllegalArgumentException.class, () -> Query.newAggregationQueryBuilder() .setNamespace(NAMESPACE) .build()); } - } \ No newline at end of file From 6508ff6dc97171ee652a8c9b189f7425e9d4bd45 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Mon, 19 Sep 2022 17:30:52 +0530 Subject: [PATCH 35/82] Preventing creating duplicated aggregation when creating an aggregation query --- .../cloud/datastore/AggregationQuery.java | 12 ++++--- .../cloud/datastore/AggregationResults.java | 4 +++ .../cloud/datastore/AggregationQueryTest.java | 33 +++++++++++++++---- 3 files changed, 38 insertions(+), 11 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java index 629fe04ea..84259e1b2 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java @@ -23,16 +23,18 @@ import com.google.cloud.datastore.aggregation.Aggregation; import com.google.cloud.datastore.aggregation.AggregationBuilder; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; +import java.util.Set; public class AggregationQuery extends Query { - private List aggregations; + private Set aggregations; private StructuredQuery nestedStructuredQuery; private final Mode mode; private GqlQuery nestedGqlQuery; - AggregationQuery(String namespace, List aggregations, + AggregationQuery(String namespace, Set aggregations, StructuredQuery nestedQuery) { super(namespace); checkArgument(!aggregations.isEmpty(), @@ -48,7 +50,7 @@ public class AggregationQuery extends Query { this.mode = GQL; } - public List getAggregations() { + public Set getAggregations() { return aggregations; } @@ -68,12 +70,12 @@ public static class Builder { private String namespace; private Mode mode; - private final List aggregations; + private final Set aggregations; private StructuredQuery nestedStructuredQuery; private GqlQuery nestedGqlQuery; public Builder() { - this.aggregations = new ArrayList<>(); + this.aggregations = new HashSet<>(); } public Builder setNamespace(String namespace) { diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java index bfb8f5fb6..039c31c08 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java @@ -15,6 +15,8 @@ */ package com.google.cloud.datastore; +import static com.google.api.client.util.Preconditions.checkNotNull; + import com.google.api.core.InternalApi; import com.google.cloud.Timestamp; import java.util.Iterator; @@ -28,6 +30,8 @@ public class AggregationResults implements Iterable { public AggregationResults(List aggregationResults, Timestamp readTime) { + checkNotNull(aggregationResults, "Aggregation results cannot be null"); + checkNotNull(readTime, "readTime cannot be null"); this.aggregationResults = aggregationResults; this.readTime = readTime; } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java index ba3a82a52..53cab0369 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java @@ -25,6 +25,7 @@ import static org.junit.Assert.assertThrows; import com.google.cloud.datastore.aggregation.CountAggregation; +import com.google.common.collect.ImmutableSet; import org.junit.Assert; import org.junit.Rule; import org.junit.Test; @@ -53,8 +54,8 @@ public void testAggregations() { .build(); assertThat(aggregationQuery.getNamespace(), equalTo(NAMESPACE)); - assertThat(aggregationQuery.getAggregations().get(0), - equalTo(count().limit(100).as("total_upto_100").build())); + assertThat(aggregationQuery.getAggregations(), equalTo( + ImmutableSet.of(count().limit(100).as("total_upto_100").build()))); assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY)); assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED)); } @@ -70,9 +71,27 @@ public void testAggregationBuilder() { .build(); assertThat(aggregationQuery.getNamespace(), equalTo(NAMESPACE)); - assertThat(aggregationQuery.getAggregations().get(0), equalTo(count().as("total").build())); - assertThat(aggregationQuery.getAggregations().get(1), - equalTo(count().limit(100).as("total_upto_100").build())); + assertThat(aggregationQuery.getAggregations(), equalTo(ImmutableSet.of( + count().as("total").build(), + count().limit(100).as("total_upto_100").build() + ))); + assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY)); + assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED)); + } + + @Test + public void testAggregationBuilderWithDuplicateAggregations() { + AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .addAggregation(count().limit(100).as("total_upto_100")) + .addAggregation(count().limit(100).as("total_upto_100")) + .over(COMPLETED_TASK_QUERY) + .build(); + + assertThat(aggregationQuery.getNamespace(), equalTo(NAMESPACE)); + assertThat(aggregationQuery.getAggregations(), equalTo(ImmutableSet.of( + count().limit(100).as("total_upto_100").build() + ))); assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY)); assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED)); } @@ -85,7 +104,9 @@ public void testAggregationQueryBuilderWithoutNamespace() { .build(); assertNull(aggregationQuery.getNamespace()); - assertThat(aggregationQuery.getAggregations().get(0), equalTo(count().as("total").build())); + assertThat(aggregationQuery.getAggregations(), equalTo(ImmutableSet.of( + count().as("total").build() + ))); assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY)); assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED)); } From 5f1e419a3e30aca687e49eadccd0a135a8f1b2e0 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Mon, 19 Sep 2022 17:50:48 +0530 Subject: [PATCH 36/82] Marking RecordQuery implemented method as InternalApi --- .../src/main/java/com/google/cloud/datastore/GqlQuery.java | 3 +++ .../src/main/java/com/google/cloud/datastore/RecordQuery.java | 1 - .../main/java/com/google/cloud/datastore/StructuredQuery.java | 3 +++ 3 files changed, 6 insertions(+), 1 deletion(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java index 5e0c0def5..814e68ae2 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java @@ -19,6 +19,7 @@ import static com.google.cloud.datastore.Validator.validateNamespace; import static com.google.common.base.Preconditions.checkNotNull; +import com.google.api.core.InternalApi; import com.google.cloud.Timestamp; import com.google.common.base.MoreObjects; import com.google.common.collect.ImmutableList; @@ -515,11 +516,13 @@ com.google.datastore.v1.GqlQuery toPb() { return protoPreparer.prepare(this); } + @InternalApi @Override public void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb) { requestPb.setGqlQuery(toPb()); } + @InternalApi @Override public RecordQuery nextQuery(com.google.datastore.v1.RunQueryResponse responsePb) { return StructuredQuery.fromPb(getType(), getNamespace(), responsePb.getQuery()) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java index 8fb1c29eb..6b158bc15 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java @@ -28,5 +28,4 @@ public interface RecordQuery{ @InternalApi RecordQuery nextQuery(com.google.datastore.v1.RunQueryResponse responsePb); - } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java index 5ccf3e4fb..8db6a506d 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java @@ -26,6 +26,7 @@ import static com.google.common.base.Preconditions.checkNotNull; import com.google.api.core.ApiFunction; +import com.google.api.core.InternalApi; import com.google.cloud.StringEnumType; import com.google.cloud.StringEnumValue; import com.google.cloud.Timestamp; @@ -1024,11 +1025,13 @@ public ResultType getType() { return resultType; } + @InternalApi @Override public void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb) { requestPb.setQuery(toPb()); } + @InternalApi @Override public StructuredQuery nextQuery(com.google.datastore.v1.RunQueryResponse responsePb) { Builder builder = toBuilder(); From 5e2fcb3d4aeee19181d50d6b64b4d54d7363a95b Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Tue, 20 Sep 2022 19:05:56 +0530 Subject: [PATCH 37/82] Writing comments and JavaDoc for aggregation query related class --- .../cloud/datastore/AggregationQuery.java | 54 +++++++++++++++++++ .../cloud/datastore/AggregationResult.java | 9 ++++ .../cloud/datastore/AggregationResults.java | 18 +++++++ .../com/google/cloud/datastore/Datastore.java | 46 ++++++++++++++++ .../cloud/datastore/DatastoreReader.java | 5 ++ .../datastore/GqlQueryProtoPreparer.java | 2 + .../com/google/cloud/datastore/Query.java | 40 +++++++++++++- .../google/cloud/datastore/ReadOption.java | 16 +++++- .../datastore/ReadOptionProtoPreparer.java | 2 + .../google/cloud/datastore/RecordQuery.java | 3 ++ .../RetryAndTraceDatastoreRpcDecorator.java | 4 ++ .../StructuredQueryProtoPreparer.java | 2 + .../datastore/aggregation/Aggregation.java | 9 ++++ .../aggregation/AggregationBuilder.java | 8 ++- .../aggregation/CountAggregation.java | 10 ++++ .../execution/AggregationQueryExecutor.java | 6 +++ .../datastore/execution/QueryExecutor.java | 17 ++++++ .../AggregationQueryRequestProtoPreparer.java | 2 + .../execution/request/ProtoPreparer.java | 7 +++ .../AggregationQueryResponseTransformer.java | 2 + .../response/ResponseTransformer.java | 10 ++++ 21 files changed, 267 insertions(+), 5 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java index 84259e1b2..1d7e82428 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java @@ -27,6 +27,52 @@ import java.util.List; import java.util.Set; +/** + * An implementation of a Google Cloud Datastore Query that returns {@link AggregationResults}, It + * can be constructed by providing a nested query ({@link StructuredQuery} or {@link GqlQuery}) + * to run the aggregations on and a set of {@link Aggregation}. + + *

{@link StructuredQuery} example:

+ * 
{@code
+ * import static com.google.cloud.datastore.aggregation.Aggregation.count;
+ *
+ * EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
+ *    .setKind("Task")
+ *    .build();
+ * AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
+ *    .addAggregation(count().as("total_count"))
+ *    .addAggregation(count().limit(100).as("count_upto_100"))
+ *    .over(selectAllQuery)
+ *    .build();
+ * AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery);
+ * for (AggregationResult aggregationResult : aggregationResults) {
+ *     System.out.println(aggregationResult.get("total_count"));
+ *     System.out.println(aggregationResult.get("count_upto_100"));
+ * }
+ * }
+ * + *

{@link GqlQuery} example:

+ * 
{@code
+ * import static com.google.cloud.datastore.aggregation.Aggregation.count;
+ *
+ * GqlQuery selectAllGqlQuery = Query.newGqlQueryBuilder(
+*         "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(100) AS count_upto_100 OVER(SELECT * FROM Task)"
+ *     )
+ *     .setAllowLiteral(true)
+ *     .build();
+ * AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
+ *     .over(selectAllGqlQuery)
+ *     .build();
+ * AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery);
+ * for (AggregationResult aggregationResult : aggregationResults) {
+ *   System.out.println(aggregationResult.get("total_count"));
+ *   System.out.println(aggregationResult.get("count_upto_100"));
+ * }
+ * }
+ * + * @see Datastore + * queries + */ public class AggregationQuery extends Query { private Set aggregations; @@ -50,18 +96,26 @@ public class AggregationQuery extends Query { this.mode = GQL; } + /** Returns the {@link Aggregation}(s) for this Query. */ public Set getAggregations() { return aggregations; } + /** Returns the underlying {@link StructuredQuery for this Query}. + * Returns null if created with {@link GqlQuery} + * */ public StructuredQuery getNestedStructuredQuery() { return nestedStructuredQuery; } + /** Returns the underlying {@link GqlQuery for this Query}. + * Returns null if created with {@link StructuredQuery} + * */ public GqlQuery getNestedGqlQuery() { return nestedGqlQuery; } + /** Returns the {@link Mode} for this query. */ public Mode getMode() { return mode; } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java index 0ccea0d9d..72421a5de 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java @@ -17,10 +17,14 @@ import com.google.common.base.MoreObjects; import com.google.common.base.MoreObjects.ToStringHelper; +import java.util.Iterator; import java.util.Map; import java.util.Map.Entry; import java.util.Objects; +/** + * Represents a result of an {@link AggregationQuery} query submission. + */ public class AggregationResult { private final Map properties; @@ -29,6 +33,11 @@ public AggregationResult(Map properties) { this.properties = properties; } + /** + * Returns a result value for the given alias. + * @param alias A custom alias provided in the query or an autogenerated alias in the form of 'property_\d' + * @return An aggregation result value for the given alias. + */ public Long get(String alias) { return properties.get(alias).get(); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java index 039c31c08..9b6148175 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java @@ -23,6 +23,18 @@ import java.util.List; import java.util.Objects; +/** + * The result of an {@link AggregationQuery} query submission. Contains a + * {@link List} and readTime {@link Timestamp} in it. + * + * This can be used to iterate over underlying {@link List} directly. + * + * Though {@link com.google.cloud.datastore.aggregation.CountAggregation} is guaranteed to return + * only one {@link AggregationResult} as part of its execution. + * + * In the future, we might support more complex {@link AggregationQuery} that might result in + * multiple {@link AggregationResult} + */ public class AggregationResults implements Iterable { private final List aggregationResults; @@ -36,6 +48,9 @@ public AggregationResults(List aggregationResults, this.readTime = readTime; } + /** + * Returns {@link Iterator} for underlying {@link List}. + */ @Override public Iterator iterator() { return this.aggregationResults.iterator(); @@ -50,6 +65,9 @@ public AggregationResult get(int index) { return this.aggregationResults.get(index); } + /** + * Returns read timestamp this result batch was returned from. + */ public Timestamp getReadTime() { return this.readTime; } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java index 06c5e3a91..d59cf59c9 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java @@ -462,5 +462,51 @@ interface TransactionCallable { */ QueryResults run(Query query, ReadOption... options); + /** + * Submits a {@link AggregationQuery} and returns {@link AggregationResults}. + * {@link ReadOption}s can be specified if desired. + * + *

Example of running an {@link AggregationQuery} to find the count of entities of one kind. + * + *

{@link StructuredQuery} example:

+ * 
{@code
+   * import static com.google.cloud.datastore.aggregation.Aggregation.count;
+   *
+   * EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
+   *    .setKind("Task")
+   *    .build();
+   * AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
+   *    .addAggregation(count().as("total_count"))
+   *    .addAggregation(count().limit(100).as("count_upto_100"))
+   *    .over(selectAllQuery)
+   *    .build();
+   * AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery);
+   * for (AggregationResult aggregationResult : aggregationResults) {
+   *     System.out.println(aggregationResult.get("total_count"));
+   *     System.out.println(aggregationResult.get("count_upto_100"));
+   * }
+   * }
+ * + *

{@link GqlQuery} example:

+ * 
{@code
+   * import static com.google.cloud.datastore.aggregation.Aggregation.count;
+   *
+   * GqlQuery selectAllGqlQuery = Query.newGqlQueryBuilder(
+   *         "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(100) AS count_upto_100 OVER(SELECT * FROM Task)"
+   *     )
+   *     .setAllowLiteral(true)
+   *     .build();
+   * AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
+   *     .over(selectAllGqlQuery)
+   *     .build();
+   * AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery);
+   * for (AggregationResult aggregationResult : aggregationResults) {
+   *   System.out.println(aggregationResult.get("total_count"));
+   *   System.out.println(aggregationResult.get("count_upto_100"));
+   * }
+   * }
+ * @throws DatastoreException upon failure + * @return {@link AggregationResults} + */ AggregationResults runAggregation(AggregationQuery query, ReadOption... options); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java index 34ace2182..51e64f28f 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java @@ -55,5 +55,10 @@ public interface DatastoreReader { QueryResults run(Query query); + /** + * Submits a {@link AggregationQuery} and returns {@link AggregationResults}. + * + * @throws DatastoreException upon failure + */ AggregationResults runAggregation(AggregationQuery query); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQueryProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQueryProtoPreparer.java index a81cb3640..90837d75e 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQueryProtoPreparer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQueryProtoPreparer.java @@ -15,10 +15,12 @@ */ package com.google.cloud.datastore; +import com.google.api.core.InternalApi; import com.google.cloud.datastore.GqlQuery.Binding; import com.google.cloud.datastore.execution.request.ProtoPreparer; import java.util.Map; +@InternalApi public class GqlQueryProtoPreparer implements ProtoPreparer, com.google.datastore.v1.GqlQuery> { diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java index a34ae46f4..f27fca572 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java @@ -26,8 +26,8 @@ import java.util.Map; /** - * A Google Cloud Datastore query. For usage examples see {@link GqlQuery} and {@link - * StructuredQuery}. + * A Google Cloud Datastore query. For usage examples see {@link GqlQuery}, {@link + * StructuredQuery} and {@link AggregationQuery}. * *

Note that queries require proper indexing. See Cloud Datastore Index @@ -259,6 +259,42 @@ public static ProjectionEntityQuery.Builder newProjectionEntityQueryBuilder() { return new ProjectionEntityQuery.Builder(); } + /** + * Returns a new {@link AggregationQuery} builder. + * + *

Example of creating and running an {@link AggregationQuery}. + * + *

{@link StructuredQuery} example:

+ * 
{@code
+   * import static com.google.cloud.datastore.aggregation.Aggregation.count;
+   *
+   * EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
+   *    .setKind("Task")
+   *    .build();
+   * AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
+   *    .addAggregation(count().as("total_count"))
+   *    .over(selectAllQuery)
+   *    .build();
+   * AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery);
+   * // Use aggregationResults
+   * }
+ * + *

{@link GqlQuery} example:

+ * 
{@code
+   * import static com.google.cloud.datastore.aggregation.Aggregation.count;
+   *
+   * GqlQuery selectAllGqlQuery = Query.newGqlQueryBuilder(
+   *         "AGGREGATE COUNT(*) AS total_count OVER(SELECT * FROM Task)"
+   *     )
+   *     .setAllowLiteral(true)
+   *     .build();
+   * AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
+   *     .over(selectAllGqlQuery)
+   *     .build();
+   * AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery);
+   * // Use aggregationResults
+   * }
+ */ public static AggregationQuery.Builder newAggregationQueryBuilder() { return new AggregationQuery.Builder(); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java index 200519bac..1fad1842b 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java @@ -72,8 +72,12 @@ public Timestamp time() { } } + /** + * Specifies transaction to be used when running a {@link Query}. + */ @InternalApi static class TransactionId extends ReadOption { + private final ByteString transactionId; TransactionId(ByteString transactionId) { @@ -106,11 +110,19 @@ public static ReadTime readTime(Timestamp time) { return new ReadTime(time); } + /** + * Returns a {@code ReadOption} that specifies transaction id, allowing Datastore to execute a + * {@link Query} in this transaction. + */ @InternalApi public static ReadOption transactionId(String transactionId) { return new TransactionId(ByteString.copyFrom(transactionId.getBytes())); } + /** + * Returns a {@code ReadOption} that specifies transaction id, allowing Datastore to execute a + * {@link Query} in this transaction. + */ @InternalApi public static ReadOption transactionId(ByteString transactionId) { return new TransactionId(transactionId); @@ -132,6 +144,7 @@ static Map, ReadOption> asImmutableMap(List> { Q query; @@ -159,7 +172,8 @@ public static > QueryAndReadOptions create(Q query) { return new QueryAndReadOptions<>(query); } - public static > QueryAndReadOptions create(Q query, List readOptions) { + public static > QueryAndReadOptions create(Q query, + List readOptions) { return new QueryAndReadOptions<>(query, readOptions); } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java index b00ad2bad..d332d3763 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java @@ -15,6 +15,7 @@ */ package com.google.cloud.datastore; +import com.google.api.core.InternalApi; import com.google.cloud.datastore.ReadOption.EventualConsistency; import com.google.cloud.datastore.ReadOption.ReadTime; import com.google.cloud.datastore.ReadOption.TransactionId; @@ -26,6 +27,7 @@ import java.util.function.Function; import java.util.stream.Collectors; +@InternalApi public class ReadOptionProtoPreparer implements ProtoPreparer, ReadOptions> { @Override diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java index 6b158bc15..6b9bc9125 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java @@ -18,6 +18,9 @@ import com.google.api.core.InternalApi; import com.google.cloud.datastore.Query.ResultType; +/** + * An internal marker interface to represent {@link Query} that returns the entity records. + */ @InternalApi public interface RecordQuery{ diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java index ad3198fcd..7d15df188 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java @@ -43,6 +43,10 @@ import io.opencensus.trace.Status; import java.util.concurrent.Callable; +/** + * An implementation of {@link DatastoreRpc} which acts as a Decorator and decorates the underlying + * {@link DatastoreRpc} with the logic of retry and Traceability. + */ public class RetryAndTraceDatastoreRpcDecorator implements DatastoreRpc { private final DatastoreRpc datastoreRpc; diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQueryProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQueryProtoPreparer.java index 030d8e114..ad07dd838 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQueryProtoPreparer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQueryProtoPreparer.java @@ -15,11 +15,13 @@ */ package com.google.cloud.datastore; +import com.google.api.core.InternalApi; import com.google.cloud.datastore.StructuredQuery.OrderBy; import com.google.cloud.datastore.execution.request.ProtoPreparer; import com.google.datastore.v1.Query; import com.google.protobuf.Int32Value; +@InternalApi public class StructuredQueryProtoPreparer implements ProtoPreparer, Query> { @Override diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java index f12cab60b..b2b48726d 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java @@ -19,6 +19,9 @@ import com.google.api.core.InternalApi; import com.google.datastore.v1.AggregationQuery; +/** + * Represents a Google Cloud Datastore Aggregation which is used with an {@link AggregationQuery}. + */ public abstract class Aggregation { private final String alias; @@ -27,6 +30,9 @@ public Aggregation(String alias) { this.alias = alias; } + /** + * Returns the alias for this aggregation. + */ public String getAlias() { return alias; } @@ -34,6 +40,9 @@ public String getAlias() { @InternalApi public abstract AggregationQuery.Aggregation toPb(); + /** + * Returns a {@link CountAggregation} builder. + */ public static CountAggregation.Builder count() { return new CountAggregation.Builder(); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java index 99179497a..3a6c826f9 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java @@ -16,8 +16,12 @@ package com.google.cloud.datastore.aggregation; +/** + * An interface to represent the builders which build and customize {@link Aggregation} for + * {@link com.google.cloud.datastore.AggregationQuery}. + * + * Used by {@link com.google.cloud.datastore.AggregationQuery.Builder#addAggregation(AggregationBuilder)}. + */ public interface AggregationBuilder { - A build(); - } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java index 16312cab9..eef7c16af 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java @@ -23,10 +23,17 @@ import java.util.Optional; import java.util.function.Consumer; +/** + * Represents an {@link Aggregation} which returns count. + */ public class CountAggregation extends Aggregation { private final long limit; + /** + * @param alias Alias to used when running this aggregation. + * @param limit Specify the number of item to scan to reduce latency and cost. + */ public CountAggregation(String alias, long limit) { super(alias); this.limit = limit; @@ -64,6 +71,9 @@ public int hashCode() { return Objects.hash(limit, getAlias()); } + /** + * A builder class to create and customize a {@link CountAggregation}. + */ public static class Builder implements AggregationBuilder { private String alias; diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java index 446bfc887..945683dc6 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java @@ -15,6 +15,7 @@ */ package com.google.cloud.datastore.execution; +import com.google.api.core.InternalApi; import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.AggregationResults; import com.google.cloud.datastore.DatastoreOptions; @@ -27,6 +28,11 @@ import com.google.datastore.v1.RunAggregationQueryResponse; import java.util.Arrays; +/** + * An implementation of {@link QueryExecutor} which executes {@link AggregationQuery} and returns + * {@link AggregationResults}. + */ +@InternalApi public class AggregationQueryExecutor implements QueryExecutor { diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java index 9bb2c3499..b737da522 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java @@ -15,9 +15,26 @@ */ package com.google.cloud.datastore.execution; +import com.google.api.core.InternalApi; import com.google.cloud.datastore.Query; import com.google.cloud.datastore.ReadOption; +/** + * An internal functional interface whose implementation has the responsibility to execute a + * {@link Query} and returns the result. This class will have the responsibility to orchestrate + * between {@link com.google.cloud.datastore.execution.request.ProtoPreparer}, + * {@link com.google.cloud.datastore.spi.v1.DatastoreRpc} and + * {@link com.google.cloud.datastore.execution.response.ResponseTransformer} layers. + * + * @param A {@link Query} to execute. + * @param the type of result produced by Query. + */ +@InternalApi public interface QueryExecutor, OUTPUT> { + + /** + * @param query A {@link Query} to execute. + * @param readOptions Optional {@link ReadOption}s to be used when executing {@link Query}. + */ OUTPUT execute(INPUT query, ReadOption... readOptions); } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java index 3d1b468fa..2169cd7bb 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java @@ -17,6 +17,7 @@ import static com.google.cloud.datastore.AggregationQuery.Mode.GQL; +import com.google.api.core.InternalApi; import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.GqlQueryProtoPreparer; @@ -32,6 +33,7 @@ import com.google.datastore.v1.RunAggregationQueryRequest; import java.util.List; +@InternalApi public class AggregationQueryRequestProtoPreparer implements ProtoPreparer, RunAggregationQueryRequest> { diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/ProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/ProtoPreparer.java index fa56c9e99..270169965 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/ProtoPreparer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/ProtoPreparer.java @@ -17,6 +17,13 @@ import com.google.api.core.InternalApi; +/** + * An internal functional interface whose implementation has the responsibility to populate a Proto + * object from a domain object. + * + * @param the type of domain object. + * @param the type of proto object + */ @InternalApi public interface ProtoPreparer { OUTPUT prepare(INPUT input); diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java index 5b01b4185..dba180d22 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java @@ -15,6 +15,7 @@ */ package com.google.cloud.datastore.execution.response; +import com.google.api.core.InternalApi; import com.google.cloud.Timestamp; import com.google.cloud.datastore.AggregationResult; import com.google.cloud.datastore.AggregationResults; @@ -29,6 +30,7 @@ import java.util.function.Function; import java.util.stream.Collectors; +@InternalApi public class AggregationQueryResponseTransformer implements ResponseTransformer { diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/ResponseTransformer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/ResponseTransformer.java index 3e5c10b66..b17da3f79 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/ResponseTransformer.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/ResponseTransformer.java @@ -15,6 +15,16 @@ */ package com.google.cloud.datastore.execution.response; +import com.google.api.core.InternalApi; + +/** + * An internal functional interface whose implementation has the responsibility to populate a Domain + * object from a proto response. + * + * @param the type of proto response object. + * @param the type of domain object. + */ +@InternalApi public interface ResponseTransformer { OUTPUT transform(INPUT response); } From 372a2d7216ea897893e5f67b73a4fb315126f414 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Thu, 22 Sep 2022 10:48:03 +0530 Subject: [PATCH 38/82] Adding a default implementation in the public interfaces to avoid compile time failures --- google-cloud-datastore/clirr-ignored-differences.xml | 10 ++++++++++ .../java/com/google/cloud/datastore/Datastore.java | 4 +++- .../com/google/cloud/datastore/DatastoreReader.java | 4 +++- .../google/cloud/datastore/spi/v1/DatastoreRpc.java | 4 +++- 4 files changed, 19 insertions(+), 3 deletions(-) diff --git a/google-cloud-datastore/clirr-ignored-differences.xml b/google-cloud-datastore/clirr-ignored-differences.xml index 1cc969163..018afb17e 100644 --- a/google-cloud-datastore/clirr-ignored-differences.xml +++ b/google-cloud-datastore/clirr-ignored-differences.xml @@ -16,4 +16,14 @@ com.google.datastore.v1.RunAggregationQueryResponse runAggregationQuery(com.google.datastore.v1.RunAggregationQueryRequest) 7012 + + com/google/cloud/datastore/Datastore + com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.ReadOption[]) + 7012 + + + com/google/cloud/datastore/DatastoreReader + com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery) + 7012 + diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java index d59cf59c9..17726da2e 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java @@ -508,5 +508,7 @@ interface TransactionCallable { * @throws DatastoreException upon failure * @return {@link AggregationResults} */ - AggregationResults runAggregation(AggregationQuery query, ReadOption... options); + default AggregationResults runAggregation(AggregationQuery query, ReadOption... options){ + throw new UnsupportedOperationException("Not implemented."); + } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java index 51e64f28f..86a18ca3d 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java @@ -60,5 +60,7 @@ public interface DatastoreReader { * * @throws DatastoreException upon failure */ - AggregationResults runAggregation(AggregationQuery query); + default AggregationResults runAggregation(AggregationQuery query){ + throw new UnsupportedOperationException("Not implemented."); + } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java index 6b542da1a..cdaab95a3 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java @@ -93,5 +93,7 @@ BeginTransactionResponse beginTransaction(BeginTransactionRequest request) * * @throws DatastoreException upon failure */ - RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request); + default RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request){ + throw new UnsupportedOperationException("Not implemented."); + } } From 2c2cf63dd44bb8f2faea44b4fcd324b25b02fdf8 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Thu, 22 Sep 2022 15:59:43 +0530 Subject: [PATCH 39/82] covering a scenario to maintain consistent snapshot when executing aggregation query in a transaction --- .../cloud/datastore/it/ITDatastoreTest.java | 98 ++++++++++++++++--- 1 file changed, 87 insertions(+), 11 deletions(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java index 708d44dd4..ca35057f4 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java @@ -18,6 +18,7 @@ import static com.google.cloud.datastore.aggregation.Aggregation.count; import static com.google.common.collect.Iterables.getOnlyElement; +import static java.util.concurrent.TimeUnit.SECONDS; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.Assert.assertEquals; @@ -26,6 +27,7 @@ import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; @@ -35,6 +37,7 @@ import com.google.cloud.datastore.BooleanValue; import com.google.cloud.datastore.Cursor; import com.google.cloud.datastore.Datastore; +import com.google.cloud.datastore.Datastore.TransactionCallable; import com.google.cloud.datastore.DatastoreException; import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.DatastoreReaderWriter; @@ -74,6 +77,10 @@ import java.util.Iterator; import java.util.List; import java.util.Set; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeoutException; import java.util.function.Consumer; import org.junit.After; import org.junit.AfterClass; @@ -558,8 +565,15 @@ public void testRunAggregationQuery() { )); } + /** + * if an entity is modified or deleted within a transaction, a query or lookup returns the + * original version of the entity as of the beginning of the transaction, + * or nothing if the entity did not exist then. + * @see + * Source + */ @Test - public void testRunAggregationQueryInTransaction() { + public void testRunAggregationQueryInTransactionShouldReturnAConsistentSnapshot() { EntityQuery entityQuery = Query.newEntityQueryBuilder() .setNamespace(NAMESPACE) .setFilter(PropertyFilter.hasAncestor(KEY1)) @@ -571,19 +585,81 @@ public void testRunAggregationQueryInTransaction() { .addAggregation(count().as("count")) .build(); - Transaction transaction = DATASTORE.newTransaction(); - assertThat(getOnlyElement(transaction.runAggregation(aggregationQuery)).get("count"), equalTo(2L)); + // original entity count is 2 + assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(2L)); - Entity aNewEntity = Entity.newBuilder(ENTITY2) - .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build()) - .set("v_int", 10) - .build(); - transaction.put(aNewEntity); + // FIRST TRANSACTION + DATASTORE.runInTransaction((TransactionCallable) inFirstTransaction -> { + // creating a new entity + Entity aNewEntity = Entity.newBuilder(ENTITY2) + .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build()) + .set("v_int", 10) + .build(); + inFirstTransaction.put(aNewEntity); + + // count remains 2 + assertThat(getOnlyElement(inFirstTransaction.runAggregation(aggregationQuery)).get("count"), equalTo(2L)); + assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(2L)); + return null; + }); + // after first transaction is committed, count is updated to 3 now. + assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(3L)); - assertThat(getOnlyElement(transaction.runAggregation(aggregationQuery)).get("count"), equalTo(2L)); + // SECOND TRANSACTION + DATASTORE.runInTransaction((TransactionCallable) inSecondTransaction -> { + // deleting ENTITY2 + inSecondTransaction.delete(ENTITY2.getKey()); + + // count remains 3 + assertThat(getOnlyElement(inSecondTransaction.runAggregation(aggregationQuery)).get("count"), equalTo(3L)); + assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(3L)); + return null; + }); + // after second transaction is committed, count is updated to 2 now. assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(2L)); - transaction.commit(); - assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(3L)); + } + + /** + * Data read or modified by a transaction cannot be concurrently modified. + * @see + * Source + */ + @Test + public void testRunAggregationQueryInATransactionShouldLockTheCountedDocuments() throws Exception { + ExecutorService executor = Executors.newSingleThreadExecutor(); + EntityQuery entityQuery = Query.newEntityQueryBuilder() + .setNamespace(NAMESPACE) + .setFilter(PropertyFilter.hasAncestor(KEY1)) + .build(); + AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .over(entityQuery) + .addAggregation(count().as("count")) + .build(); + + Transaction insideTransaction = DATASTORE.newTransaction(); + + // acquiring lock by executing query in transaction + assertThat(getOnlyElement(insideTransaction.runAggregation(aggregationQuery)).get("count"), + equalTo(2L)); + + // Waiting task will be blocked by ongoing transaction. + Future addNewEntityTaskOutsideTransaction = executor.submit(() -> { + Entity aNewEntity = Entity.newBuilder(ENTITY2) + .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build()) + .set("v_int", 10) + .build(); + DATASTORE.put(aNewEntity); + return null; + }); + + // should throw Timeout exception as we haven't yet committed the transaction + assertThrows(TimeoutException.class, () -> addNewEntityTaskOutsideTransaction.get(3, SECONDS)); + + //cleanup + insideTransaction.commit(); + addNewEntityTaskOutsideTransaction.cancel(true); + executor.shutdownNow(); } @Test From 4719ae8fb623a39eb4a149fb50459beece656dbb Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Fri, 23 Sep 2022 17:57:45 +0530 Subject: [PATCH 40/82] Creating emulator proxy to simulate AggregationQuery response from emulator --- google-cloud-datastore/pom.xml | 6 + .../google/cloud/datastore/DatastoreTest.java | 27 +++- .../datastore/emulator/EmulatorProxy.java | 51 ++++++++ .../datastore/emulator/ProxyDispatcher.java | 118 ++++++++++++++++++ 4 files changed, 198 insertions(+), 4 deletions(-) create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/EmulatorProxy.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/ProxyDispatcher.java diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml index 42eaa95e4..1590c69c8 100644 --- a/google-cloud-datastore/pom.xml +++ b/google-cloud-datastore/pom.xml @@ -143,6 +143,12 @@ easymock test + + com.squareup.okhttp3 + mockwebserver + 4.10.0 + test + diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java index fb8d11dc7..5256ffd3a 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java @@ -17,6 +17,7 @@ package com.google.cloud.datastore; import static com.google.cloud.datastore.aggregation.Aggregation.count; +import static com.google.common.collect.Iterables.getOnlyElement; import static org.easymock.EasyMock.createStrictMock; import static org.easymock.EasyMock.replay; import static org.easymock.EasyMock.verify; @@ -36,6 +37,7 @@ import com.google.cloud.datastore.Query.ResultType; import com.google.cloud.datastore.StructuredQuery.OrderBy; import com.google.cloud.datastore.StructuredQuery.PropertyFilter; +import com.google.cloud.datastore.emulator.EmulatorProxy; import com.google.cloud.datastore.spi.DatastoreRpcFactory; import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.cloud.datastore.testing.LocalDatastoreHelper; @@ -76,7 +78,6 @@ import org.junit.Assert; import org.junit.Before; import org.junit.BeforeClass; -import org.junit.Ignore; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -88,6 +89,8 @@ public class DatastoreTest { private static LocalDatastoreHelper helper = LocalDatastoreHelper.create(1.0); private static final DatastoreOptions options = helper.getOptions(); private static final Datastore datastore = options.getService(); + private static final EmulatorProxy emulatorProxy = new EmulatorProxy(options); + private static final Datastore datastoreEmulatorProxy = emulatorProxy.getOptions().getService(); private static final String PROJECT_ID = options.getProjectId(); private static final String KIND1 = "kind1"; private static final String KIND2 = "kind2"; @@ -185,6 +188,7 @@ public void setUp() { @AfterClass public static void afterClass() throws IOException, InterruptedException, TimeoutException { helper.stop(Duration.ofMinutes(1)); + emulatorProxy.stop(); } @Test @@ -532,22 +536,37 @@ public void testGqlQueryPagination() throws DatastoreException { } @Test - @Ignore public void testRunAggregationQuery() { EntityQuery selectAllQuery = Query.newEntityQueryBuilder().build(); AggregationQuery getCountQuery = Query.newAggregationQueryBuilder() .addAggregation(count().as("total_count").limit(100)) .over(selectAllQuery) .build(); - AggregationResult resultBeforeInsert = Iterables.getOnlyElement(datastore.runAggregation(getCountQuery)); + AggregationResult resultBeforeInsert = getOnlyElement( + datastoreEmulatorProxy.runAggregation(getCountQuery)); assertThat(resultBeforeInsert.get("total_count"), equalTo(2L)); datastore.put(ENTITY3); - AggregationResult resultAfterInsert = Iterables.getOnlyElement(datastore.runAggregation(getCountQuery)); + AggregationResult resultAfterInsert = getOnlyElement( + datastoreEmulatorProxy.runAggregation(getCountQuery)); assertThat(resultAfterInsert.get("total_count"), equalTo(3L)); } + @Test + public void testRunAggregationQueryWithUptoOption() { + datastore.put(ENTITY3); + + EntityQuery selectAllQuery = Query.newEntityQueryBuilder().build(); + AggregationQuery getCountQuery = Query.newAggregationQueryBuilder() + .addAggregation(count().as("count_upto_2").limit(2)) + .over(selectAllQuery) + .build(); + AggregationResult resultBeforeInsert = getOnlyElement( + datastoreEmulatorProxy.runAggregation(getCountQuery)); + assertThat(resultBeforeInsert.get("count_upto_2"), equalTo(2L)); + } + @Test public void testRunStructuredQuery() { Query query = diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/EmulatorProxy.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/EmulatorProxy.java new file mode 100644 index 000000000..726f3fe11 --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/EmulatorProxy.java @@ -0,0 +1,51 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore.emulator; + +import com.google.cloud.datastore.DatastoreOptions; +import com.google.cloud.datastore.spi.v1.HttpDatastoreRpc; +import java.io.IOException; +import okhttp3.mockwebserver.MockWebServer; + +public class EmulatorProxy { + + private final MockWebServer mockWebServer; + private final DatastoreOptions emulatorDataStoreOptions; + + public EmulatorProxy(DatastoreOptions emulatorDataStoreOptions) { + this.emulatorDataStoreOptions = emulatorDataStoreOptions; + this.mockWebServer = new MockWebServer(); + init(); + } + + private void init() { + this.mockWebServer.setDispatcher(new ProxyDispatcher(new HttpDatastoreRpc(emulatorDataStoreOptions))); + } + + public void start() throws IOException { + this.mockWebServer.start(); + } + + public void stop() throws IOException { + this.mockWebServer.shutdown(); + } + + public DatastoreOptions getOptions() { + DatastoreOptions.Builder builder = this.emulatorDataStoreOptions.toBuilder(); + builder.setHost(this.mockWebServer.getHostName() + ":" + this.mockWebServer.getPort()); + return builder.build(); + } +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/ProxyDispatcher.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/ProxyDispatcher.java new file mode 100644 index 000000000..2a8792e5d --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/ProxyDispatcher.java @@ -0,0 +1,118 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.cloud.datastore.emulator; + +import static com.google.cloud.datastore.ProtoTestData.intValue; + +import com.google.cloud.datastore.spi.v1.DatastoreRpc; +import com.google.datastore.v1.AggregationQuery.Aggregation; +import com.google.datastore.v1.AggregationQuery.Aggregation.Count; +import com.google.datastore.v1.AggregationResult; +import com.google.datastore.v1.AggregationResultBatch; +import com.google.datastore.v1.Query; +import com.google.datastore.v1.RunAggregationQueryRequest; +import com.google.datastore.v1.RunAggregationQueryResponse; +import com.google.datastore.v1.RunQueryRequest; +import com.google.datastore.v1.RunQueryResponse; +import com.google.datastore.v1.Value; +import java.io.IOException; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import okhttp3.mockwebserver.Dispatcher; +import okhttp3.mockwebserver.MockResponse; +import okhttp3.mockwebserver.RecordedRequest; +import okio.Buffer; +import org.jetbrains.annotations.NotNull; + +public class ProxyDispatcher extends Dispatcher { + + private final DatastoreRpc datastoreRpc; + + public ProxyDispatcher(DatastoreRpc datastoreRpc) { + this.datastoreRpc = datastoreRpc; + } + + @NotNull + @Override + public MockResponse dispatch(@NotNull RecordedRequest recordedRequest) { + String methodName = recordedRequest.getPath().split(":")[1]; + if (!"runAggregationQuery".equals(methodName)) { + throw new IllegalStateException( + String.format("Proxy only supports RunAggregationQuery method, Found %s method.", + methodName)); + } + + try { + RunAggregationQueryRequest runAggregationQueryRequest = RunAggregationQueryRequest.parseFrom( + recordedRequest.getBody().inputStream()); + + RunQueryRequest runQueryRequest = getRunQueryRequest(runAggregationQueryRequest); + RunQueryResponse runQueryResponse = this.datastoreRpc.runQuery(runQueryRequest); + + RunAggregationQueryResponse runAggregationQueryResponse = getRunAggregationQueryResponse( + runAggregationQueryRequest, runQueryResponse); + Buffer buffer = new Buffer(); + runAggregationQueryResponse.writeTo(buffer.outputStream()); + return new MockResponse().setBody(buffer); + } catch (IOException e) { + e.printStackTrace(); + } + return new MockResponse(); + } + + @NotNull + private RunAggregationQueryResponse getRunAggregationQueryResponse( + RunAggregationQueryRequest runAggregationQueryRequest, RunQueryResponse runQueryResponse) { + AggregationResult aggregationResult = AggregationResult.newBuilder() + .putAllAggregateProperties(getProperties( + runAggregationQueryRequest.getAggregationQuery().getAggregationsList(), + runQueryResponse)) + .build(); + return RunAggregationQueryResponse.newBuilder() + .setBatch( + AggregationResultBatch.newBuilder().addAggregationResults(aggregationResult).build()) + .build(); + } + + @NotNull + private RunQueryRequest getRunQueryRequest(RunAggregationQueryRequest runAggregationQueryRequest) { + Query nestedQuery = runAggregationQueryRequest.getAggregationQuery().getNestedQuery(); + RunQueryRequest.Builder builder = RunQueryRequest.newBuilder().setQuery(nestedQuery); + if (runAggregationQueryRequest.hasReadOptions()) { + builder.setReadOptions(runAggregationQueryRequest.getReadOptions()); + } + return builder.build(); + } + + private Map getProperties(List aggregationsList, + RunQueryResponse runQueryResponse) { + HashMap map = new HashMap<>(); + for (Aggregation aggregation : aggregationsList) { + map.put(aggregation.getAlias(), getValue(runQueryResponse, aggregation.getCount())); + } + return map; + } + + private Value getValue(RunQueryResponse runQueryResponse, Count count) { + int totalEntityCount = runQueryResponse.getBatch().getEntityResultsCount(); + if (count.hasUpTo()) { + return intValue(Math.min(count.getUpTo().getValue(), totalEntityCount)); + } else { + return intValue(totalEntityCount); + } + } +} From 8edebdaacde0066c1db431ca2c4e0886855c9239 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Tue, 27 Sep 2022 14:56:25 +0530 Subject: [PATCH 41/82] Integration test to execute an aggregation query in a read only transaction --- .../cloud/datastore/it/ITDatastoreTest.java | 54 +++++++++++++++++-- 1 file changed, 49 insertions(+), 5 deletions(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java index ca35057f4..7ae6042b9 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java @@ -71,6 +71,7 @@ import com.google.cloud.datastore.testing.RemoteDatastoreHelper; import com.google.common.base.Preconditions; import com.google.datastore.v1.TransactionOptions; +import com.google.datastore.v1.TransactionOptions.ReadOnly; import java.util.ArrayList; import java.util.Collections; import java.util.HashSet; @@ -625,7 +626,7 @@ public void testRunAggregationQueryInTransactionShouldReturnAConsistentSnapshot( * Source */ @Test - public void testRunAggregationQueryInATransactionShouldLockTheCountedDocuments() throws Exception { + public void testRunAggregationQueryInAReadWriteTransactionShouldLockTheCountedDocuments() throws Exception { ExecutorService executor = Executors.newSingleThreadExecutor(); EntityQuery entityQuery = Query.newEntityQueryBuilder() .setNamespace(NAMESPACE) @@ -637,13 +638,14 @@ public void testRunAggregationQueryInATransactionShouldLockTheCountedDocuments() .addAggregation(count().as("count")) .build(); - Transaction insideTransaction = DATASTORE.newTransaction(); + // read-write transaction + Transaction readWriteTransaction = DATASTORE.newTransaction(); // acquiring lock by executing query in transaction - assertThat(getOnlyElement(insideTransaction.runAggregation(aggregationQuery)).get("count"), + assertThat(getOnlyElement(readWriteTransaction.runAggregation(aggregationQuery)).get("count"), equalTo(2L)); - // Waiting task will be blocked by ongoing transaction. + // Waiting task will be blocked by ongoing transactions. Future addNewEntityTaskOutsideTransaction = executor.submit(() -> { Entity aNewEntity = Entity.newBuilder(ENTITY2) .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build()) @@ -657,11 +659,53 @@ public void testRunAggregationQueryInATransactionShouldLockTheCountedDocuments() assertThrows(TimeoutException.class, () -> addNewEntityTaskOutsideTransaction.get(3, SECONDS)); //cleanup - insideTransaction.commit(); + readWriteTransaction.commit(); addNewEntityTaskOutsideTransaction.cancel(true); executor.shutdownNow(); } + @Test + public void testRunAggregationQueryInAReadOnlyTransactionShouldNotLockTheCountedDocuments() throws Exception { + ExecutorService executor = Executors.newSingleThreadExecutor(); + EntityQuery entityQuery = Query.newEntityQueryBuilder() + .setNamespace(NAMESPACE) + .setFilter(PropertyFilter.hasAncestor(KEY1)) + .build(); + AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() + .setNamespace(NAMESPACE) + .over(entityQuery) + .addAggregation(count().as("count")) + .build(); + + TransactionOptions transactionOptions = TransactionOptions.newBuilder().setReadOnly( + ReadOnly.newBuilder().build()).build(); + Transaction readOnlyTransaction = DATASTORE.newTransaction(transactionOptions); + + // Executing query in transaction + assertThat(getOnlyElement(readOnlyTransaction.runAggregation(aggregationQuery)).get("count"), + equalTo(2L)); + + // Concurrent write task. + Future addNewEntityTaskOutsideTransaction = executor.submit(() -> { + Entity aNewEntity = Entity.newBuilder(ENTITY2) + .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build()) + .set("v_int", 10) + .build(); + DATASTORE.put(aNewEntity); + return null; + }); + + // should not throw exception and complete successfully as the ongoing transaction is read-only. + addNewEntityTaskOutsideTransaction.get(); + + //cleanup + readOnlyTransaction.commit(); + executor.shutdownNow(); + + assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), + equalTo(3L)); + } + @Test public void testRunAggregationQueryWithReadTime() throws InterruptedException { String alias = "total_count"; From 394a642f59f14f98cd7537fb2224878b18d59bfd Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 28 Sep 2022 17:04:34 +0530 Subject: [PATCH 42/82] Count aggregation samples with structuredQuery and gql query --- .../datastore/AggregationQuerySample.java | 149 ++++++++++++++++++ .../AggregationQuerySampleTestIT.java | 47 ++++++ .../test/java/com/rule/SystemsOutRule.java | 52 ++++++ 3 files changed, 248 insertions(+) create mode 100644 samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java create mode 100644 samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java create mode 100644 samples/snippets/src/test/java/com/rule/SystemsOutRule.java diff --git a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java new file mode 100644 index 000000000..1a6cb9580 --- /dev/null +++ b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java @@ -0,0 +1,149 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.example.datastore; + +import com.google.cloud.datastore.AggregationQuery; +import com.google.cloud.datastore.AggregationResult; +import com.google.cloud.datastore.Datastore; +import com.google.cloud.datastore.DatastoreOptions; +import com.google.cloud.datastore.Entity; +import com.google.cloud.datastore.EntityQuery; +import com.google.cloud.datastore.GqlQuery; +import com.google.cloud.datastore.Key; +import com.google.cloud.datastore.Query; +import com.google.cloud.datastore.StructuredQuery.PropertyFilter; +import com.google.cloud.datastore.aggregation.Aggregation; +import com.google.common.collect.Iterables; + +public class AggregationQuerySample { + + public void aggregationQueryAndCountAggregation() { + // [START datastore_count_aggregation_query] + + // Instantiates a client + Datastore datastore = DatastoreOptions.getDefaultInstance().getService(); + + // The kind for the new entity + String kind = "Candidate"; + + Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1"); + Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2"); + Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3"); + + // Save all the candidates + datastore.put( + Entity.newBuilder(candidate1Key).set("qualified", true).build(), + Entity.newBuilder(candidate2Key).set("qualified", false).build(), + Entity.newBuilder(candidate3Key).set("qualified", true).build() + ); + + EntityQuery selectAllCandidates = Query.newEntityQueryBuilder() + .setKind(kind) + .build(); + // Creating an aggregation query to get the count of all candidates + AggregationQuery allCandidatesCountQuery = Query.newAggregationQueryBuilder() + .over(selectAllCandidates) + .addAggregation(Aggregation.count().as("total_count")) + .addAggregation(Aggregation.count().as("count_with_limit").limit(2)) + .build(); + // Executing aggregation query + AggregationResult allCandidatesCountQueryResult = Iterables.getOnlyElement( + datastore.runAggregation(allCandidatesCountQuery)); + + System.out.printf("We have at least %d candidates", allCandidatesCountQueryResult.get("count_with_limit")); // 2 + System.out.printf("Total candidates count is %d", allCandidatesCountQueryResult.get("total_count")); // 3 + + + EntityQuery qualifiedCandidates = Query.newEntityQueryBuilder() + .setKind(kind) + .setFilter(PropertyFilter.eq("qualified", true)) + .build(); + // Creating an aggregation query to get the count of all qualified candidates + AggregationQuery qualifiedCandidatesCountQuery = Query.newAggregationQueryBuilder() + .over(qualifiedCandidates) + .addAggregation(Aggregation.count().as("total_qualified_count")) + .build(); + + // Executing aggregation query + AggregationResult qualifiedCandidatesCountQueryResult = Iterables.getOnlyElement( + datastore.runAggregation(qualifiedCandidatesCountQuery)); + + System.out.printf("Total qualified candidates count is %d", qualifiedCandidatesCountQueryResult.get("total_qualified_count")); // 2 + + // [END datastore_count_aggregation_query] + + datastore.delete(candidate1Key, candidate2Key, candidate3Key); + } + + public void aggregationQueryAndCountAggregationWithGqlQuery() { + // [START datastore_count_aggregation_query_gql] + + // Instantiates a client + Datastore datastore = DatastoreOptions.getDefaultInstance().getService(); + + // The kind for the new entity + String kind = "Candidate"; + + Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1"); + Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2"); + Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3"); + + // Save all the candidates + datastore.put( + Entity.newBuilder(candidate1Key).set("qualified", true).build(), + Entity.newBuilder(candidate2Key).set("qualified", false).build(), + Entity.newBuilder(candidate3Key).set("qualified", true).build() + ); + + GqlQuery selectAllCandidates = Query.newGqlQueryBuilder( + "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(2) AS count_with_limit " + + "OVER (SELECT * FROM Candidate)") + .setAllowLiteral(true) + .build(); + // Creating an aggregation query to get the count of all candidates + AggregationQuery allCandidatesCountQuery = Query.newAggregationQueryBuilder() + .over(selectAllCandidates) + .build(); + // Executing aggregation query + AggregationResult allCandidatesCountQueryResult = Iterables.getOnlyElement( + datastore.runAggregation(allCandidatesCountQuery)); + + System.out.printf("We have at least %d candidates", allCandidatesCountQueryResult.get("count_with_limit")); // 2 + System.out.printf("Total candidates count is %d", allCandidatesCountQueryResult.get("total_count")); // 3 + + + GqlQuery qualifiedCandidates = Query.newGqlQueryBuilder( + "AGGREGATE COUNT(*) AS total_qualified_count " + + "OVER (SELECT * FROM Candidate WHERE qualified = true)") + .setAllowLiteral(true) + .build(); + // Creating an aggregation query to get the count of all qualified candidates + AggregationQuery qualifiedCandidatesCountQuery = Query.newAggregationQueryBuilder() + .over(qualifiedCandidates) + .build(); + + // Executing aggregation query + AggregationResult qualifiedCandidatesCountQueryResult = Iterables.getOnlyElement( + datastore.runAggregation(qualifiedCandidatesCountQuery)); + + System.out.printf("Total qualified candidates count is %d", qualifiedCandidatesCountQueryResult.get("total_qualified_count")); // 2 + + // [END datastore_count_aggregation_query_gql] + + datastore.delete(candidate1Key, candidate2Key, candidate3Key); + } + +} diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java new file mode 100644 index 000000000..f575f8d22 --- /dev/null +++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java @@ -0,0 +1,47 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.example.datastore; + +import com.rule.SystemsOutRule; +import org.junit.Rule; +import org.junit.Test; + +public class AggregationQuerySampleTestIT { + + @Rule + public final SystemsOutRule systemsOutRule = new SystemsOutRule(); + + private final AggregationQuerySample sample = new AggregationQuerySample(); + + + @Test + public void testAggregationQueryAndCountAggregationSample() { + sample.aggregationQueryAndCountAggregation(); + + systemsOutRule.assertContains("We have at least 2 candidates"); + systemsOutRule.assertContains("Total candidates count is 3"); + systemsOutRule.assertContains("Total qualified candidates count is 2"); + } + + @Test + public void testAggregationQueryAndCountAggregationSampleWithGqlQuery() { + sample.aggregationQueryAndCountAggregationWithGqlQuery(); + + systemsOutRule.assertContains("We have at least 2 candidates"); + systemsOutRule.assertContains("Total candidates count is 3"); + systemsOutRule.assertContains("Total qualified candidates count is 2"); + } +} \ No newline at end of file diff --git a/samples/snippets/src/test/java/com/rule/SystemsOutRule.java b/samples/snippets/src/test/java/com/rule/SystemsOutRule.java new file mode 100644 index 000000000..eabc41ce9 --- /dev/null +++ b/samples/snippets/src/test/java/com/rule/SystemsOutRule.java @@ -0,0 +1,52 @@ +/* + * Copyright 2022 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.rule; + +import static com.google.common.truth.Truth.assertThat; + +import java.io.ByteArrayOutputStream; +import java.io.PrintStream; +import org.junit.rules.TestRule; +import org.junit.runner.Description; +import org.junit.runners.model.Statement; + +public class SystemsOutRule implements TestRule { + private ByteArrayOutputStream currentOut; + + @Override + public Statement apply(Statement statement, Description description) { + return new Statement() { + @Override + public void evaluate() throws Throwable { + // Setting up customized PrintStream + PrintStream originalOut = System.out; + currentOut = new ByteArrayOutputStream(); + System.setOut(new PrintStream(currentOut)); + + // Running tests + statement.evaluate(); + + // Restoring original PrintStream + System.setOut(originalOut); + currentOut = null; + } + }; + } + + public void assertContains(String content) { + assertThat(currentOut.toString()).contains(content); + } +} From 41cbef7973291d59e66c489624b8911b1348535b Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Wed, 28 Sep 2022 17:17:29 +0530 Subject: [PATCH 43/82] Count aggregation samples to demonstrate stale read --- .../datastore/AggregationQuerySample.java | 52 +++++++++++++++++++ .../AggregationQuerySampleTestIT.java | 8 +++ 2 files changed, 60 insertions(+) diff --git a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java index 1a6cb9580..38330beb8 100644 --- a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java +++ b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java @@ -15,6 +15,7 @@ */ package com.example.datastore; +import com.google.cloud.Timestamp; import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.AggregationResult; import com.google.cloud.datastore.Datastore; @@ -24,6 +25,7 @@ import com.google.cloud.datastore.GqlQuery; import com.google.cloud.datastore.Key; import com.google.cloud.datastore.Query; +import com.google.cloud.datastore.ReadOption; import com.google.cloud.datastore.StructuredQuery.PropertyFilter; import com.google.cloud.datastore.aggregation.Aggregation; import com.google.common.collect.Iterables; @@ -146,4 +148,54 @@ public void aggregationQueryAndCountAggregationWithGqlQuery() { datastore.delete(candidate1Key, candidate2Key, candidate3Key); } + public void aggregationQueryAndCountAggregationWithStaleRead() throws InterruptedException { + // [START datastore_count_aggregation_query_stale_read] + + // Instantiates a client + Datastore datastore = DatastoreOptions.getDefaultInstance().getService(); + + // The kind for the new entity + String kind = "Candidate"; + + Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1"); + Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2"); + Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3"); + + // Saving only two candidates + datastore.put( + Entity.newBuilder(candidate1Key).set("qualified", true).build(), + Entity.newBuilder(candidate2Key).set("qualified", false).build() + ); + Thread.sleep(1000); + Timestamp pastTimestamp = Timestamp.now(); // we have two candidates in database at this time. + + Thread.sleep(1000); + // Saving third candidates + datastore.put(Entity.newBuilder(candidate3Key).set("qualified", false).build()); + + EntityQuery selectAllCandidates = Query.newEntityQueryBuilder() + .setKind(kind) + .build(); + + // Creating an aggregation query to get the count of all candidates + AggregationQuery allCandidatesCountQuery = Query.newAggregationQueryBuilder() + .over(selectAllCandidates) + .addAggregation(Aggregation.count().as("total_count")) + .build(); + + // Executing aggregation query + AggregationResult candidatesCountLatest = Iterables.getOnlyElement( + datastore.runAggregation(allCandidatesCountQuery)); + System.out.printf("Latest candidates count is %d", candidatesCountLatest.get("total_count")); // 3 + + // Executing aggregation query with past timestamp + AggregationResult candidatesCountInPast = Iterables.getOnlyElement( + datastore.runAggregation(allCandidatesCountQuery, ReadOption.readTime(pastTimestamp))); + System.out.printf("Stale candidates count is %d", candidatesCountInPast.get("total_count")); // 2 + + // [END datastore_count_aggregation_query_stale_read] + + datastore.delete(candidate1Key, candidate2Key, candidate3Key); + } + } diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java index f575f8d22..0a86df3b0 100644 --- a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java +++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java @@ -44,4 +44,12 @@ public void testAggregationQueryAndCountAggregationSampleWithGqlQuery() { systemsOutRule.assertContains("Total candidates count is 3"); systemsOutRule.assertContains("Total qualified candidates count is 2"); } + + @Test + public void testAggregationQueryAndCountWithStaleRead() throws InterruptedException { + sample.aggregationQueryAndCountAggregationWithStaleRead(); + + systemsOutRule.assertContains("Latest candidates count is 3"); + systemsOutRule.assertContains("Stale candidates count is 2"); + } } \ No newline at end of file From 842df2d8a5be9f7b3c5e3d08120eb091043e16ba Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Thu, 29 Sep 2022 11:34:00 +0530 Subject: [PATCH 44/82] Getting rid off limit operation on count aggregation as same behaviour can be achieved by using 'limit' operation on the underlying query --- .../cloud/datastore/AggregationQuery.java | 2 -- .../com/google/cloud/datastore/Datastore.java | 2 -- .../aggregation/CountAggregation.java | 23 +++------------ .../cloud/datastore/AggregationQueryTest.java | 19 ++++++------ .../datastore/AggregationResultTest.java | 4 +-- .../google/cloud/datastore/DatastoreTest.java | 16 +--------- .../aggregation/CountAggregationTest.java | 29 ++----------------- .../AggregationQueryExecutorTest.java | 12 ++++---- ...regationQueryRequestProtoPreparerTest.java | 7 ++--- ...gregationQueryResponseTransformerTest.java | 4 +-- 10 files changed, 29 insertions(+), 89 deletions(-) diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java index 1d7e82428..5edc24d46 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java @@ -41,13 +41,11 @@ * .build(); * AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() * .addAggregation(count().as("total_count")) - * .addAggregation(count().limit(100).as("count_upto_100")) * .over(selectAllQuery) * .build(); * AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery); * for (AggregationResult aggregationResult : aggregationResults) { * System.out.println(aggregationResult.get("total_count")); - * System.out.println(aggregationResult.get("count_upto_100")); * } * } * diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java index 17726da2e..b90889c73 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java @@ -477,13 +477,11 @@ interface TransactionCallable { * .build(); * AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() * .addAggregation(count().as("total_count")) - * .addAggregation(count().limit(100).as("count_upto_100")) * .over(selectAllQuery) * .build(); * AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery); * for (AggregationResult aggregationResult : aggregationResults) { * System.out.println(aggregationResult.get("total_count")); - * System.out.println(aggregationResult.get("count_upto_100")); * } * } * diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java index eef7c16af..b205896fd 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java @@ -20,31 +20,22 @@ import com.google.datastore.v1.AggregationQuery.Aggregation.Count; import com.google.protobuf.Int64Value; import java.util.Objects; -import java.util.Optional; -import java.util.function.Consumer; /** * Represents an {@link Aggregation} which returns count. */ public class CountAggregation extends Aggregation { - private final long limit; - /** * @param alias Alias to used when running this aggregation. - * @param limit Specify the number of item to scan to reduce latency and cost. */ - public CountAggregation(String alias, long limit) { + public CountAggregation(String alias) { super(alias); - this.limit = limit; } @Override public AggregationQuery.Aggregation toPb() { Count.Builder countBuilder = Count.newBuilder(); - if(limit > 0) { - countBuilder.setUpTo(Int64Value.of(limit)); - } AggregationQuery.Aggregation.Builder aggregationBuilder = AggregationQuery.Aggregation.newBuilder() .setCount(countBuilder); @@ -63,12 +54,12 @@ public boolean equals(Object o) { return false; } CountAggregation that = (CountAggregation) o; - return limit == that.limit && getAlias().equals(that.getAlias()); + return getAlias().equals(that.getAlias()); } @Override public int hashCode() { - return Objects.hash(limit, getAlias()); + return Objects.hash(getAlias()); } /** @@ -77,21 +68,15 @@ public int hashCode() { public static class Builder implements AggregationBuilder { private String alias; - private long limit; public Builder as(String alias) { this.alias = alias; return this; } - public Builder limit(long limit) { - this.limit = limit; - return this; - } - @Override public CountAggregation build() { - return new CountAggregation(alias, limit); + return new CountAggregation(alias); } } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java index 53cab0369..66992efa7 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java @@ -26,10 +26,8 @@ import com.google.cloud.datastore.aggregation.CountAggregation; import com.google.common.collect.ImmutableSet; -import org.junit.Assert; import org.junit.Rule; import org.junit.Test; -import org.junit.function.ThrowingRunnable; import org.junit.rules.ExpectedException; public class AggregationQueryTest { @@ -40,6 +38,7 @@ public class AggregationQueryTest { .setNamespace(NAMESPACE) .setKind(KIND) .setFilter(eq("done", true)) + .setLimit(100) .build(); @Rule @@ -49,31 +48,31 @@ public class AggregationQueryTest { public void testAggregations() { AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() .setNamespace(NAMESPACE) - .addAggregation(new CountAggregation("total_upto_100", 100)) + .addAggregation(new CountAggregation("total")) .over(COMPLETED_TASK_QUERY) .build(); assertThat(aggregationQuery.getNamespace(), equalTo(NAMESPACE)); assertThat(aggregationQuery.getAggregations(), equalTo( - ImmutableSet.of(count().limit(100).as("total_upto_100").build()))); + ImmutableSet.of(count().as("total").build()))); assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY)); assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED)); } @Test - public void testAggregationBuilder() { + public void testAggregationBuilderWithMoreThanOneAggregations() { AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() .setNamespace(NAMESPACE) .addAggregation(count().as("total")) - .addAggregation(count().limit(100).as("total_upto_100")) + .addAggregation(count().as("new_total")) .over(COMPLETED_TASK_QUERY) .build(); assertThat(aggregationQuery.getNamespace(), equalTo(NAMESPACE)); assertThat(aggregationQuery.getAggregations(), equalTo(ImmutableSet.of( count().as("total").build(), - count().limit(100).as("total_upto_100").build() + count().as("new_total").build() ))); assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY)); assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED)); @@ -83,14 +82,14 @@ public void testAggregationBuilder() { public void testAggregationBuilderWithDuplicateAggregations() { AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder() .setNamespace(NAMESPACE) - .addAggregation(count().limit(100).as("total_upto_100")) - .addAggregation(count().limit(100).as("total_upto_100")) + .addAggregation(count().as("total")) + .addAggregation(count().as("total")) .over(COMPLETED_TASK_QUERY) .build(); assertThat(aggregationQuery.getNamespace(), equalTo(NAMESPACE)); assertThat(aggregationQuery.getAggregations(), equalTo(ImmutableSet.of( - count().limit(100).as("total_upto_100").build() + count().as("total").build() ))); assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY)); assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED)); diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java index 776fae8ab..78398737e 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java @@ -27,10 +27,10 @@ public class AggregationResultTest { public void shouldGetAggregationResultValueByAlias() { AggregationResult aggregationResult = new AggregationResult(ImmutableMap.of( "count", LongValue.of(45), - "count_upto_30", LongValue.of(30) + "property_2", LongValue.of(30) )); assertThat(aggregationResult.get("count"), equalTo(45L)); - assertThat(aggregationResult.get("count_upto_30"), equalTo(30L)); + assertThat(aggregationResult.get("property_2"), equalTo(30L)); } } \ No newline at end of file diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java index 5256ffd3a..698b2c84b 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java @@ -539,7 +539,7 @@ public void testGqlQueryPagination() throws DatastoreException { public void testRunAggregationQuery() { EntityQuery selectAllQuery = Query.newEntityQueryBuilder().build(); AggregationQuery getCountQuery = Query.newAggregationQueryBuilder() - .addAggregation(count().as("total_count").limit(100)) + .addAggregation(count().as("total_count")) .over(selectAllQuery) .build(); AggregationResult resultBeforeInsert = getOnlyElement( @@ -553,20 +553,6 @@ public void testRunAggregationQuery() { assertThat(resultAfterInsert.get("total_count"), equalTo(3L)); } - @Test - public void testRunAggregationQueryWithUptoOption() { - datastore.put(ENTITY3); - - EntityQuery selectAllQuery = Query.newEntityQueryBuilder().build(); - AggregationQuery getCountQuery = Query.newAggregationQueryBuilder() - .addAggregation(count().as("count_upto_2").limit(2)) - .over(selectAllQuery) - .build(); - AggregationResult resultBeforeInsert = getOnlyElement( - datastoreEmulatorProxy.runAggregation(getCountQuery)); - assertThat(resultBeforeInsert.get("count_upto_2"), equalTo(2L)); - } - @Test public void testRunStructuredQuery() { Query query = diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java index e3275987f..6a35142a4 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java @@ -36,16 +36,6 @@ public void testCountAggregationWithDefaultValues() { assertThat(countAggregationPb.getAlias(), equalTo("")); } - @Test - public void testCountAggregationWithLimit() { - AggregationQuery.Aggregation countAggregationPb = count() - .limit(100) - .build().toPb(); - - assertThat(countAggregationPb.getCount().getUpTo().getValue(), equalTo(100L)); - assertThat(countAggregationPb.getAlias(), equalTo("")); - } - @Test public void testCountAggregationWithAlias() { AggregationQuery.Aggregation countAggregationPb = count() @@ -56,31 +46,18 @@ public void testCountAggregationWithAlias() { assertThat(countAggregationPb.getAlias(), equalTo("column_1")); } - @Test - public void testCountAggregationWithAliasAndLimit() { - AggregationQuery.Aggregation countAggregationPb = count() - .as("column_1") - .limit(100) - .build().toPb(); - - assertThat(countAggregationPb.getCount().getUpTo().getValue(), equalTo(100L)); - assertThat(countAggregationPb.getAlias(), equalTo("column_1")); - } - @Test public void testEquals() { CountAggregation.Builder aggregation1 = count() - .as("total") - .limit(100); + .as("total"); CountAggregation.Builder aggregation2 = count() - .as("total") - .limit(100); + .as("total"); assertEquals(aggregation1.build(), aggregation2.build()); assertEquals(aggregation2.build(), aggregation1.build()); assertNotEquals(aggregation1.as("new-alias").build(), aggregation2.build()); - assertNotEquals(aggregation1.limit(399).build(), aggregation2.build()); + assertNotEquals(aggregation2.build(), aggregation1.as("new-alias").build()); } } \ No newline at end of file diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java index f34257250..c8099e194 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java @@ -91,9 +91,9 @@ public void shouldExecuteAggregationQuery() { verify(mockRpc); assertThat(aggregationResults, equalTo(new AggregationResults(asList( new AggregationResult( - ImmutableMap.of("count", LongValue.of(209), "count_upto_100", LongValue.of(100))), + ImmutableMap.of("count", LongValue.of(209), "property_2", LongValue.of(100))), new AggregationResult( - ImmutableMap.of("count", LongValue.of(509), "count_upto_100", LongValue.of(100))) + ImmutableMap.of("count", LongValue.of(509), "property_2", LongValue.of(100))) ), Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime())))); } @@ -123,21 +123,21 @@ public void shouldExecuteAggregationQueryWithReadOptions() { verify(mockRpc); assertThat(aggregationResults, equalTo(new AggregationResults(asList( new AggregationResult( - ImmutableMap.of("count", LongValue.of(209), "count_upto_100", LongValue.of(100))), + ImmutableMap.of("count", LongValue.of(209), "property_2", LongValue.of(100))), new AggregationResult( - ImmutableMap.of("count", LongValue.of(509), "count_upto_100", LongValue.of(100))) + ImmutableMap.of("count", LongValue.of(509), "property_2", LongValue.of(100))) ), Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime())))); } private RunAggregationQueryResponse dummyAggregationQueryResponse() { Map result1 = new HashMap() {{ put("count", intValue(209)); - put("count_upto_100", intValue(100)); + put("property_2", intValue(100)); }}; Map result2 = new HashMap() {{ put("count", intValue(509)); - put("count_upto_100", intValue(100)); + put("property_2", intValue(100)); }}; AggregationResultBatch resultBatch = AggregationResultBatch.newBuilder() diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java index 0d43f8099..077254858 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java @@ -71,7 +71,6 @@ public class AggregationQueryRequestProtoPreparerTest { private final AggregationQuery AGGREGATION_OVER_STRUCTURED_QUERY = Query.newAggregationQueryBuilder() .setNamespace(NAMESPACE) .addAggregation(count().as("total")) - .addAggregation(count().limit(100).as("total_upto_100")) .over(COMPLETED_TASK_STRUCTURED_QUERY) .build(); @@ -100,10 +99,8 @@ public void shouldPrepareAggregationQueryRequestWithGivenStructuredQuery() { .addKind(kind(KIND)) .setFilter(propertyFilter("done", EQUAL, booleanValue(true))) .build())); - assertThat(aggregationQueryProto.getAggregationsList(), equalTo(asList( - countAggregation("total"), - countAggregation("total_upto_100", 100) - ))); + assertThat(aggregationQueryProto.getAggregationsList(), + equalTo(singletonList(countAggregation("total")))); } @Test diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java index a9570c1c6..dafc5dad2 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java @@ -42,12 +42,12 @@ public class AggregationQueryResponseTransformerTest { public void shouldTransformAggregationQueryResponse() { Map result1 = new HashMap() {{ put("count", intValue(209)); - put("count_upto_100", intValue(100)); + put("property_2", intValue(100)); }}; Map result2 = new HashMap() {{ put("count", intValue(509)); - put("count_upto_100", intValue(100)); + put("property_2", intValue(100)); }}; Timestamp readTime = Timestamp.now(); From 8c81b9f201102b4db79d195263d15fb06e4985d3 Mon Sep 17 00:00:00 2001 From: Prateek Jain Date: Thu, 29 Sep 2022 12:20:19 +0530 Subject: [PATCH 45/82] Removing import statement from javadoc and undo changes in .gitignore file --- .gitignore | 3 +-- .../main/java/com/google/cloud/datastore/AggregationQuery.java | 2 -- .../src/main/java/com/google/cloud/datastore/Datastore.java | 2 -- .../src/main/java/com/google/cloud/datastore/Query.java | 2 -- 4 files changed, 1 insertion(+), 8 deletions(-) diff --git a/.gitignore b/.gitignore index d1e66c922..6f5e4424d 100644 --- a/.gitignore +++ b/.gitignore @@ -3,5 +3,4 @@ target/ *.iml __pycache__/ -.flattened-pom.xml -*.patch \ No newline at end of file +.flattened-pom.xml \ No newline at end of file diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java index 5edc24d46..e6c0c0dc3 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java @@ -34,7 +34,6 @@ *

{@link StructuredQuery} example:

* 
{@code
- * import static com.google.cloud.datastore.aggregation.Aggregation.count;
  *
  * EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
  *    .setKind("Task")
@@ -51,7 +50,6 @@
  *
  * 
{@link GqlQuery} example:

  * 
{@code
- * import static com.google.cloud.datastore.aggregation.Aggregation.count;
  *
  * GqlQuery selectAllGqlQuery = Query.newGqlQueryBuilder(
 *         "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(100) AS count_upto_100 OVER(SELECT * FROM Task)"
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
index b90889c73..d8d2431fc 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
@@ -470,7 +470,6 @@ interface TransactionCallable {
    *
    * 
{@link StructuredQuery} example:

    * 
{@code
-   * import static com.google.cloud.datastore.aggregation.Aggregation.count;
    *
    * EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
    *    .setKind("Task")
@@ -487,7 +486,6 @@ interface TransactionCallable {
    *
    * 
{@link GqlQuery} example:

    * 
{@code
-   * import static com.google.cloud.datastore.aggregation.Aggregation.count;
    *
    * GqlQuery selectAllGqlQuery = Query.newGqlQueryBuilder(
    *         "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(100) AS count_upto_100 OVER(SELECT * FROM Task)"
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java
index f27fca572..5802320a8 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java
@@ -266,7 +266,6 @@ public static ProjectionEntityQuery.Builder newProjectionEntityQueryBuilder() {
    *
    * 
{@link StructuredQuery} example:

    * 
{@code
-   * import static com.google.cloud.datastore.aggregation.Aggregation.count;
    *
    * EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
    *    .setKind("Task")
@@ -281,7 +280,6 @@ public static ProjectionEntityQuery.Builder newProjectionEntityQueryBuilder() {
    *
    * 
{@link GqlQuery} example:

    * 
{@code
-   * import static com.google.cloud.datastore.aggregation.Aggregation.count;
    *
    * GqlQuery selectAllGqlQuery = Query.newGqlQueryBuilder(
    *         "AGGREGATE COUNT(*) AS total_count OVER(SELECT * FROM Task)"

From 1584d9f88ae7af6a87223daaf5d0af6690e5a017 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Thu, 29 Sep 2022 13:00:07 +0530
Subject: [PATCH 46/82] Using Optional instead of returning null from
 ReadOptionsProtoPreparer

---
 .../google/cloud/datastore/DatastoreImpl.java | 26 ++++-----
 .../cloud/datastore/QueryResultsImpl.java     | 10 ++--
 .../datastore/ReadOptionProtoPreparer.java    | 56 ++++++++++---------
 .../cloud/datastore/TransactionImpl.java      |  5 +-
 .../AggregationQueryRequestProtoPreparer.java |  7 +--
 .../ReadOptionProtoPreparerTest.java          | 19 ++++---
 6 files changed, 63 insertions(+), 60 deletions(-)

diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
index 9c2af79b6..78582f196 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
@@ -22,8 +22,6 @@
 import com.google.cloud.RetryHelper;
 import com.google.cloud.RetryHelper.RetryHelperException;
 import com.google.cloud.ServiceOptions;
-import com.google.cloud.datastore.ReadOption.EventualConsistency;
-import com.google.cloud.datastore.ReadOption.ReadTime;
 import com.google.cloud.datastore.execution.AggregationQueryExecutor;
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
 import com.google.common.base.MoreObjects;
@@ -32,7 +30,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
-import com.google.datastore.v1.ReadOptions.ReadConsistency;
+import com.google.datastore.v1.ReadOptions;
 import com.google.datastore.v1.ReserveIdsRequest;
 import com.google.datastore.v1.TransactionOptions;
 import com.google.protobuf.ByteString;
@@ -47,6 +45,7 @@
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Callable;
 
@@ -183,7 +182,7 @@ public  T runInTransaction(
 
   @Override
   public  QueryResults run(Query query) {
-    return run(null, query);
+    return run(Optional.empty(), query);
   }
 
   @Override
@@ -191,9 +190,8 @@ public  QueryResults run(Query query, ReadOption... options) {
     return run(toReadOptionsPb(options), query);
   }
 
-  //TODO- Need to check whether we can handle this typecasting in a better way
   @SuppressWarnings("unchecked")
-   QueryResults run(com.google.datastore.v1.ReadOptions readOptionsPb, Query query) {
+   QueryResults run(Optional readOptionsPb, Query query) {
     return new QueryResultsImpl(this, readOptionsPb, (RecordQuery) query, query.getNamespace());
   }
 
@@ -352,7 +350,7 @@ public Entity get(Key key, ReadOption... options) {
 
   @Override
   public Iterator get(Key... keys) {
-    return get(null, keys);
+    return get(Optional.empty(), keys);
   }
 
   @Override
@@ -360,11 +358,11 @@ public Iterator get(Iterable keys, ReadOption... options) {
     return get(toReadOptionsPb(options), Iterables.toArray(keys, Key.class));
   }
 
-  private com.google.datastore.v1.ReadOptions toReadOptionsPb(ReadOption... options) {
-    if (options != null) {
-      return this.readOptionProtoPreparer.prepare(Arrays.asList(options));
+  private Optional toReadOptionsPb(ReadOption... options) {
+    if (options == null) {
+      return Optional.empty();
     }
-    return null;
+    return this.readOptionProtoPreparer.prepare(Arrays.asList(options));
   }
 
   @Override
@@ -377,15 +375,13 @@ public List fetch(Iterable keys, ReadOption... options) {
     return DatastoreHelper.fetch(this, Iterables.toArray(keys, Key.class), options);
   }
 
-  Iterator get(com.google.datastore.v1.ReadOptions readOptionsPb, final Key... keys) {
+  Iterator get(Optional readOptionsPb, final Key... keys) {
     if (keys.length == 0) {
       return Collections.emptyIterator();
     }
     com.google.datastore.v1.LookupRequest.Builder requestPb =
         com.google.datastore.v1.LookupRequest.newBuilder();
-    if (readOptionsPb != null) {
-      requestPb.setReadOptions(readOptionsPb);
-    }
+    readOptionsPb.ifPresent(requestPb::setReadOptions);
     for (Key k : Sets.newLinkedHashSet(Arrays.asList(keys))) {
       requestPb.addKeys(k.toPb());
     }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
index 09c319aff..d079d5837 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
@@ -20,14 +20,16 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.AbstractIterator;
 import com.google.datastore.v1.QueryResultBatch.MoreResultsType;
+import com.google.datastore.v1.ReadOptions;
 import com.google.protobuf.ByteString;
 import java.util.Iterator;
 import java.util.Objects;
+import java.util.Optional;
 
 class QueryResultsImpl extends AbstractIterator implements QueryResults {
 
   private final DatastoreImpl datastore;
-  private final com.google.datastore.v1.ReadOptions readOptionsPb;
+  private final Optional readOptionsPb;
   private final com.google.datastore.v1.PartitionId partitionIdPb;
   private final ResultType queryResultType;
   private RecordQuery query;
@@ -40,7 +42,7 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults
   private MoreResultsType moreResults;
 
   QueryResultsImpl(
-      DatastoreImpl datastore, com.google.datastore.v1.ReadOptions readOptionsPb, RecordQuery query, String namespace) {
+      DatastoreImpl datastore, Optional readOptionsPb, RecordQuery query, String namespace) {
     this.datastore = datastore;
     this.readOptionsPb = readOptionsPb;
     this.query = query;
@@ -65,9 +67,7 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults
   private void sendRequest() {
     com.google.datastore.v1.RunQueryRequest.Builder requestPb =
         com.google.datastore.v1.RunQueryRequest.newBuilder();
-    if (readOptionsPb != null) {
-      requestPb.setReadOptions(readOptionsPb);
-    }
+    readOptionsPb.ifPresent(requestPb::setReadOptions);
     requestPb.setPartitionId(partitionIdPb);
     query.populatePb(requestPb);
     runQueryResponsePb = datastore.runQuery(requestPb.build());
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java
index d332d3763..c70d2c7b3 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java
@@ -24,44 +24,48 @@
 import com.google.datastore.v1.ReadOptions.ReadConsistency;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
 @InternalApi
-public class ReadOptionProtoPreparer implements ProtoPreparer, ReadOptions> {
+public class ReadOptionProtoPreparer implements
+    ProtoPreparer, Optional> {
 
   @Override
-  public ReadOptions prepare(List options) {
+  public Optional prepare(List options) {
+    if (options == null || options.isEmpty()) {
+      return Optional.empty();
+    }
     com.google.datastore.v1.ReadOptions readOptionsPb = null;
-    if (options != null && !options.isEmpty()) {
-      Map, ReadOption> optionsByType =
-          ReadOption.asImmutableMap(options);
+    Map, ReadOption> optionsByType =
+        ReadOption.asImmutableMap(options);
 
-      boolean moreThanOneReadOption = optionsByType.keySet().size() > 1;
-      if (moreThanOneReadOption) {
-        throw DatastoreException.throwInvalidRequest(
-            String.format("Can not use %s together.", getInvalidOptions(optionsByType)));
-      }
+    boolean moreThanOneReadOption = optionsByType.keySet().size() > 1;
+    if (moreThanOneReadOption) {
+      throw DatastoreException.throwInvalidRequest(
+          String.format("Can not use %s together.", getInvalidOptions(optionsByType)));
+    }
 
-      if (optionsByType.containsKey(EventualConsistency.class)) {
-        readOptionsPb = ReadOptions.newBuilder()
-            .setReadConsistency(ReadConsistency.EVENTUAL)
-            .build();
-      }
+    if (optionsByType.containsKey(EventualConsistency.class)) {
+      readOptionsPb = ReadOptions.newBuilder()
+          .setReadConsistency(ReadConsistency.EVENTUAL)
+          .build();
+    }
 
-      if (optionsByType.containsKey(ReadTime.class)) {
-        readOptionsPb = ReadOptions.newBuilder()
-            .setReadTime(((ReadTime) optionsByType.get(ReadTime.class)).time().toProto())
-            .build();
-      }
+    if (optionsByType.containsKey(ReadTime.class)) {
+      readOptionsPb = ReadOptions.newBuilder()
+          .setReadTime(((ReadTime) optionsByType.get(ReadTime.class)).time().toProto())
+          .build();
+    }
 
-      if (optionsByType.containsKey(TransactionId.class)) {
-        readOptionsPb = ReadOptions.newBuilder()
-            .setTransaction(((TransactionId) optionsByType.get(TransactionId.class)).getTransactionId())
-            .build();
-      }
+    if (optionsByType.containsKey(TransactionId.class)) {
+      readOptionsPb = ReadOptions.newBuilder()
+          .setTransaction(
+              ((TransactionId) optionsByType.get(TransactionId.class)).getTransactionId())
+          .build();
     }
-    return readOptionsPb;
+    return Optional.ofNullable(readOptionsPb);
   }
 
   private String getInvalidOptions(Map, ReadOption> optionsByType) {
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java
index c13439460..5d06e0a1c 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java
@@ -23,6 +23,7 @@
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Optional;
 
 final class TransactionImpl extends BaseDatastoreBatchWriter implements Transaction {
 
@@ -80,7 +81,7 @@ public Iterator get(Key... keys) {
     com.google.datastore.v1.ReadOptions.Builder readOptionsPb =
         com.google.datastore.v1.ReadOptions.newBuilder();
     readOptionsPb.setTransaction(transactionId);
-    return datastore.get(readOptionsPb.build(), keys);
+    return datastore.get(Optional.of(readOptionsPb.build()), keys);
   }
 
   @Override
@@ -95,7 +96,7 @@ public  QueryResults run(Query query) {
     com.google.datastore.v1.ReadOptions.Builder readOptionsPb =
         com.google.datastore.v1.ReadOptions.newBuilder();
     readOptionsPb.setTransaction(transactionId);
-    return datastore.run(readOptionsPb.build(), query);
+    return datastore.run(Optional.of(readOptionsPb.build()), query);
   }
 
   @Override
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java
index 2169cd7bb..7418ac5c6 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java
@@ -32,6 +32,7 @@
 import com.google.datastore.v1.ReadOptions;
 import com.google.datastore.v1.RunAggregationQueryRequest;
 import java.util.List;
+import java.util.Optional;
 
 @InternalApi
 public class AggregationQueryRequestProtoPreparer implements
@@ -65,10 +66,8 @@ public RunAggregationQueryRequest prepare(
       aggregationQueryRequestBuilder.setAggregationQuery(getAggregationQuery(aggregationQuery));
     }
 
-    ReadOptions readOptionsPb = readOptionProtoPreparer.prepare(readOptions);
-    if (readOptionsPb != null) {
-      aggregationQueryRequestBuilder.setReadOptions(readOptionsPb);
-    }
+    Optional readOptionsPb = readOptionProtoPreparer.prepare(readOptions);
+    readOptionsPb.ifPresent(aggregationQueryRequestBuilder::setReadOptions);
     return aggregationQueryRequestBuilder.build();
   }
 
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java
index 90df8885e..6e8ba707c 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java
@@ -22,8 +22,10 @@
 import static java.util.Collections.singletonList;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 
 import com.google.cloud.Timestamp;
 import com.google.common.collect.ImmutableList;
@@ -31,6 +33,7 @@
 import com.google.protobuf.ByteString;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Optional;
 import org.junit.Test;
 
 public class ReadOptionProtoPreparerTest {
@@ -55,35 +58,35 @@ public void shouldThrowErrorWhenUsingMultipleReadOptions() {
 
   @Test
   public void shouldPrepareReadOptionsWithEventualConsistency() {
-    ReadOptions readOptions = protoPreparer.prepare(singletonList(eventualConsistency()));
+    Optional readOptions = protoPreparer.prepare(singletonList(eventualConsistency()));
 
-    assertThat(readOptions.getReadConsistency(), is(EVENTUAL));
+    assertThat(readOptions.get().getReadConsistency(), is(EVENTUAL));
   }
 
   @Test
   public void shouldPrepareReadOptionsWithReadTime() {
     Timestamp timestamp = Timestamp.now();
-    ReadOptions readOptions = protoPreparer.prepare(singletonList(readTime(timestamp)));
+    Optional readOptions = protoPreparer.prepare(singletonList(readTime(timestamp)));
 
-    assertThat(Timestamp.fromProto(readOptions.getReadTime()), is(timestamp));
+    assertThat(Timestamp.fromProto(readOptions.get().getReadTime()), is(timestamp));
   }
 
   @Test
   public void shouldPrepareReadOptionsWithTransactionId() {
     String dummyTransactionId = "transaction-id";
-    ReadOptions readOptions = protoPreparer.prepare(singletonList(transactionId(
+    Optional readOptions = protoPreparer.prepare(singletonList(transactionId(
         dummyTransactionId)));
 
-    assertThat(readOptions.getTransaction().toStringUtf8(), is(dummyTransactionId));
+    assertThat(readOptions.get().getTransaction().toStringUtf8(), is(dummyTransactionId));
   }
 
   @Test
   public void shouldReturnNullWhenReadOptionsIsNull() {
-    assertNull(protoPreparer.prepare(null));
+    assertFalse(protoPreparer.prepare(null).isPresent());
   }
 
   @Test
   public void shouldReturnNullWhenReadOptionsIsAnEmptyList() {
-    assertNull(protoPreparer.prepare(ImmutableList.of()));
+    assertFalse(protoPreparer.prepare(ImmutableList.of()).isPresent());
   }
 }
\ No newline at end of file

From bc425ddc2a887020beb6fd49355ee96ac94a76b4 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Thu, 29 Sep 2022 15:55:20 +0530
Subject: [PATCH 47/82] using assertThat from Truth library

---
 google-cloud-datastore/pom.xml                |  6 ++
 .../cloud/datastore/AggregationQueryTest.java | 46 ++++++------
 .../datastore/AggregationResultTest.java      |  7 +-
 .../google/cloud/datastore/DatastoreTest.java |  9 ++-
 .../datastore/GqlQueryProtoPreparerTest.java  | 23 +++---
 .../ReadOptionProtoPreparerTest.java          | 12 ++--
 ...etryAndTraceDatastoreRpcDecoratorTest.java |  5 +-
 .../StructuredQueryProtoPreparerTest.java     | 31 ++++----
 .../aggregation/CountAggregationTest.java     | 12 ++--
 .../AggregationQueryExecutorTest.java         | 12 ++--
 ...regationQueryRequestProtoPreparerTest.java | 68 +++++++++---------
 ...gregationQueryResponseTransformerTest.java | 14 ++--
 .../cloud/datastore/it/ITDatastoreTest.java   | 70 +++++++++++--------
 13 files changed, 168 insertions(+), 147 deletions(-)

diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml
index f929f07ab..5a7ba1fcf 100644
--- a/google-cloud-datastore/pom.xml
+++ b/google-cloud-datastore/pom.xml
@@ -149,6 +149,12 @@
       4.10.0
       test
     
+    
+      com.google.truth
+      truth
+      1.1.3
+      test
+    
   
 
   
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java
index 66992efa7..b43e55ec5 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java
@@ -19,13 +19,13 @@
 import static com.google.cloud.datastore.AggregationQuery.Mode.STRUCTURED;
 import static com.google.cloud.datastore.StructuredQuery.PropertyFilter.eq;
 import static com.google.cloud.datastore.aggregation.Aggregation.count;
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.MatcherAssert.assertThat;
+import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 
 import com.google.cloud.datastore.aggregation.CountAggregation;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.truth.Truth;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
@@ -52,11 +52,11 @@ public void testAggregations() {
         .over(COMPLETED_TASK_QUERY)
         .build();
 
-    assertThat(aggregationQuery.getNamespace(), equalTo(NAMESPACE));
-    assertThat(aggregationQuery.getAggregations(), equalTo(
-        ImmutableSet.of(count().as("total").build())));
-    assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY));
-    assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED));
+    assertThat(aggregationQuery.getNamespace()).isEqualTo(NAMESPACE);
+    assertThat(aggregationQuery.getAggregations()).isEqualTo(
+        ImmutableSet.of(count().as("total").build()));
+    assertThat(aggregationQuery.getNestedStructuredQuery()).isEqualTo(COMPLETED_TASK_QUERY);
+    assertThat(aggregationQuery.getMode()).isEqualTo(STRUCTURED);
   }
 
 
@@ -69,13 +69,13 @@ public void testAggregationBuilderWithMoreThanOneAggregations() {
         .over(COMPLETED_TASK_QUERY)
         .build();
 
-    assertThat(aggregationQuery.getNamespace(), equalTo(NAMESPACE));
-    assertThat(aggregationQuery.getAggregations(), equalTo(ImmutableSet.of(
+    assertThat(aggregationQuery.getNamespace()).isEqualTo(NAMESPACE);
+    assertThat(aggregationQuery.getAggregations()).isEqualTo(ImmutableSet.of(
         count().as("total").build(),
         count().as("new_total").build()
-    )));
-    assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY));
-    assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED));
+    ));
+    assertThat(aggregationQuery.getNestedStructuredQuery()).isEqualTo(COMPLETED_TASK_QUERY);
+    assertThat(aggregationQuery.getMode()).isEqualTo(STRUCTURED);
   }
 
   @Test
@@ -87,12 +87,12 @@ public void testAggregationBuilderWithDuplicateAggregations() {
         .over(COMPLETED_TASK_QUERY)
         .build();
 
-    assertThat(aggregationQuery.getNamespace(), equalTo(NAMESPACE));
-    assertThat(aggregationQuery.getAggregations(), equalTo(ImmutableSet.of(
+    assertThat(aggregationQuery.getNamespace()).isEqualTo(NAMESPACE);
+    assertThat(aggregationQuery.getAggregations()).isEqualTo(ImmutableSet.of(
         count().as("total").build()
-    )));
-    assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY));
-    assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED));
+    ));
+    assertThat(aggregationQuery.getNestedStructuredQuery()).isEqualTo(COMPLETED_TASK_QUERY);
+    assertThat(aggregationQuery.getMode()).isEqualTo(STRUCTURED);
   }
 
   @Test
@@ -103,11 +103,11 @@ public void testAggregationQueryBuilderWithoutNamespace() {
         .build();
 
     assertNull(aggregationQuery.getNamespace());
-    assertThat(aggregationQuery.getAggregations(), equalTo(ImmutableSet.of(
+    assertThat(aggregationQuery.getAggregations()).isEqualTo(ImmutableSet.of(
         count().as("total").build()
-    )));
-    assertThat(aggregationQuery.getNestedStructuredQuery(), equalTo(COMPLETED_TASK_QUERY));
-    assertThat(aggregationQuery.getMode(), equalTo(STRUCTURED));
+    ));
+    assertThat(aggregationQuery.getNestedStructuredQuery()).isEqualTo(COMPLETED_TASK_QUERY);
+    assertThat(aggregationQuery.getMode()).isEqualTo(STRUCTURED);
   }
 
   @Test
@@ -139,8 +139,8 @@ public void testAggregationQueryBuilderWithGqlQuery() {
         .over(gqlQuery)
         .build();
 
-    assertThat(aggregationQuery.getNestedGqlQuery(), equalTo(gqlQuery));
-    assertThat(aggregationQuery.getMode(), equalTo(GQL));
+    assertThat(aggregationQuery.getNestedGqlQuery()).isEqualTo(gqlQuery);
+    assertThat(aggregationQuery.getMode()).isEqualTo(GQL);
   }
 
   @Test
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java
index 78398737e..4ff694fd2 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java
@@ -15,8 +15,7 @@
  */
 package com.google.cloud.datastore;
 
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.MatcherAssert.assertThat;
+import static com.google.common.truth.Truth.assertThat;
 
 import com.google.common.collect.ImmutableMap;
 import org.junit.Test;
@@ -30,7 +29,7 @@ public void shouldGetAggregationResultValueByAlias() {
         "property_2", LongValue.of(30)
     ));
 
-    assertThat(aggregationResult.get("count"), equalTo(45L));
-    assertThat(aggregationResult.get("property_2"), equalTo(30L));
+    assertThat(aggregationResult.get("count")).isEqualTo(45L);
+    assertThat(aggregationResult.get("property_2")).isEqualTo(30L);
   }
 }
\ No newline at end of file
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
index 698b2c84b..67a4b41cd 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
@@ -18,11 +18,10 @@
 
 import static com.google.cloud.datastore.aggregation.Aggregation.count;
 import static com.google.common.collect.Iterables.getOnlyElement;
+import static com.google.common.truth.Truth.assertThat;
 import static org.easymock.EasyMock.createStrictMock;
 import static org.easymock.EasyMock.replay;
 import static org.easymock.EasyMock.verify;
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
@@ -42,9 +41,9 @@
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
 import com.google.cloud.datastore.testing.LocalDatastoreHelper;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Iterables;
 import com.google.common.collect.Iterators;
 import com.google.common.collect.Lists;
+import com.google.common.truth.Truth;
 import com.google.datastore.v1.BeginTransactionRequest;
 import com.google.datastore.v1.BeginTransactionResponse;
 import com.google.datastore.v1.CommitRequest;
@@ -544,13 +543,13 @@ public void testRunAggregationQuery() {
         .build();
     AggregationResult resultBeforeInsert = getOnlyElement(
         datastoreEmulatorProxy.runAggregation(getCountQuery));
-    assertThat(resultBeforeInsert.get("total_count"), equalTo(2L));
+    assertThat(resultBeforeInsert.get("total_count")).isEqualTo(2L);
 
     datastore.put(ENTITY3);
 
     AggregationResult resultAfterInsert = getOnlyElement(
         datastoreEmulatorProxy.runAggregation(getCountQuery));
-    assertThat(resultAfterInsert.get("total_count"), equalTo(3L));
+    assertThat(resultAfterInsert.get("total_count")).isEqualTo(3L);
   }
 
   @Test
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java
index 1ed707228..1af698306 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java
@@ -19,17 +19,14 @@
 import static com.google.cloud.datastore.ProtoTestData.intValue;
 import static com.google.cloud.datastore.ProtoTestData.stringValue;
 import static com.google.cloud.datastore.Query.newGqlQueryBuilder;
+import static com.google.common.truth.Truth.assertThat;
 import static java.util.Arrays.asList;
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import com.google.common.truth.Truth;
 import com.google.datastore.v1.GqlQueryParameter;
-import java.util.Arrays;
-import java.util.Collections;
 import java.util.HashMap;
-import java.util.Map;
 import org.junit.Test;
 
 public class GqlQueryProtoPreparerTest {
@@ -41,7 +38,7 @@ public class GqlQueryProtoPreparerTest {
   public void testQueryString() {
     com.google.datastore.v1.GqlQuery gqlQuery = protoPreparer.prepare(gqlQueryBuilder.build());
 
-    assertThat(gqlQuery.getQueryString(), equalTo("SELECT * from Character"));
+    assertThat(gqlQuery.getQueryString()).isEqualTo("SELECT * from Character");
   }
 
   @Test
@@ -61,10 +58,11 @@ public void testNamedBinding() {
             .build()
     );
 
-    assertThat(gqlQuery.getNamedBindingsMap(), equalTo(new HashMap() {{
-      put("name", gqlQueryParameter(stringValue("John Doe")));
-      put("age", gqlQueryParameter(intValue(27)));
-    }}));
+    assertThat(gqlQuery.getNamedBindingsMap())
+        .isEqualTo(new HashMap() {{
+          put("name", gqlQueryParameter(stringValue("John Doe")));
+          put("age", gqlQueryParameter(intValue(27)));
+        }});
   }
 
   @Test
@@ -76,10 +74,9 @@ public void testPositionalBinding() {
             .build()
     );
 
-    assertThat(gqlQuery.getPositionalBindingsList(), equalTo(asList(
+    assertThat(gqlQuery.getPositionalBindingsList()).isEqualTo(asList(
         gqlQueryParameter(stringValue("John Doe")),
         gqlQueryParameter(intValue(27))
-    )));
+    ));
   }
-
 }
\ No newline at end of file
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java
index 6e8ba707c..539f9a9e5 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java
@@ -18,10 +18,10 @@
 import static com.google.cloud.datastore.ReadOption.eventualConsistency;
 import static com.google.cloud.datastore.ReadOption.readTime;
 import static com.google.cloud.datastore.ReadOption.transactionId;
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.datastore.v1.ReadOptions.ReadConsistency.EVENTUAL;
 import static java.util.Collections.singletonList;
 import static org.hamcrest.CoreMatchers.is;
-import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
@@ -29,6 +29,7 @@
 
 import com.google.cloud.Timestamp;
 import com.google.common.collect.ImmutableList;
+import com.google.common.truth.Truth;
 import com.google.datastore.v1.ReadOptions;
 import com.google.protobuf.ByteString;
 import java.util.Arrays;
@@ -53,14 +54,15 @@ public void shouldThrowErrorWhenUsingMultipleReadOptions() {
             Arrays.asList(transactionId("transaction-id"), readTime(Timestamp.now()))));
     assertThrows(DatastoreException.class,
         () -> protoPreparer.prepare(
-            Arrays.asList(eventualConsistency(), readTime(Timestamp.now()), transactionId("transaction-id"))));
+            Arrays.asList(eventualConsistency(), readTime(Timestamp.now()),
+                transactionId("transaction-id"))));
   }
 
   @Test
   public void shouldPrepareReadOptionsWithEventualConsistency() {
     Optional readOptions = protoPreparer.prepare(singletonList(eventualConsistency()));
 
-    assertThat(readOptions.get().getReadConsistency(), is(EVENTUAL));
+    assertThat(readOptions.get().getReadConsistency()).isEqualTo(EVENTUAL);
   }
 
   @Test
@@ -68,7 +70,7 @@ public void shouldPrepareReadOptionsWithReadTime() {
     Timestamp timestamp = Timestamp.now();
     Optional readOptions = protoPreparer.prepare(singletonList(readTime(timestamp)));
 
-    assertThat(Timestamp.fromProto(readOptions.get().getReadTime()), is(timestamp));
+    assertThat(Timestamp.fromProto(readOptions.get().getReadTime())).isEqualTo(timestamp);
   }
 
   @Test
@@ -77,7 +79,7 @@ public void shouldPrepareReadOptionsWithTransactionId() {
     Optional readOptions = protoPreparer.prepare(singletonList(transactionId(
         dummyTransactionId)));
 
-    assertThat(readOptions.get().getTransaction().toStringUtf8(), is(dummyTransactionId));
+    assertThat(readOptions.get().getTransaction().toStringUtf8()).isEqualTo(dummyTransactionId);
   }
 
   @Test
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java
index b92233187..1bee0c4e8 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java
@@ -17,14 +17,13 @@
 
 import static com.google.cloud.datastore.TraceUtil.END_SPAN_OPTIONS;
 import static com.google.cloud.datastore.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY;
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.rpc.Code.UNAVAILABLE;
 import static org.easymock.EasyMock.createNiceMock;
 import static org.easymock.EasyMock.createStrictMock;
 import static org.easymock.EasyMock.expect;
 import static org.easymock.EasyMock.replay;
 import static org.easymock.EasyMock.verify;
-import static org.hamcrest.CoreMatchers.sameInstance;
-import static org.hamcrest.MatcherAssert.assertThat;
 
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
@@ -72,7 +71,7 @@ public void testRunAggregationQuery() {
     RunAggregationQueryResponse actualAggregationQueryResponse = datastoreRpcDecorator.runAggregationQuery(
         aggregationQueryRequest);
 
-    assertThat(actualAggregationQueryResponse, sameInstance(aggregationQueryResponse));
+    assertThat(actualAggregationQueryResponse).isSameInstanceAs(aggregationQueryResponse);
     verify(mockDatastoreRpc, mockTraceUtil, mockSpan);
   }
 }
\ No newline at end of file
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java
index 38deee0c5..2479ff912 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java
@@ -21,12 +21,12 @@
 import static com.google.cloud.datastore.ProtoTestData.propertyOrder;
 import static com.google.cloud.datastore.ProtoTestData.propertyReference;
 import static com.google.cloud.datastore.Query.newEntityQueryBuilder;
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.datastore.v1.PropertyFilter.Operator.EQUAL;
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.MatcherAssert.assertThat;
 
 import com.google.cloud.datastore.StructuredQuery.OrderBy;
 import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
+import com.google.common.truth.Truth;
 import com.google.datastore.v1.KindExpression;
 import com.google.datastore.v1.Query;
 import com.google.protobuf.ByteString;
@@ -41,7 +41,8 @@ public class StructuredQueryProtoPreparerTest {
   public void testKind() {
     Query queryProto = protoPreparer.prepare(newEntityQueryBuilder().setKind("kind").build());
 
-    assertThat(queryProto.getKind(0), equalTo(KindExpression.newBuilder().setName("kind").build()));
+    assertThat(queryProto.getKind(0))
+        .isEqualTo(KindExpression.newBuilder().setName("kind").build());
   }
 
   @Test
@@ -50,7 +51,7 @@ public void testStartCursor() {
     Query queryProto = protoPreparer.prepare(newEntityQueryBuilder().setStartCursor(
         Cursor.copyFrom(bytes)).build());
 
-    assertThat(queryProto.getStartCursor(), equalTo(ByteString.copyFrom(bytes)));
+    assertThat(queryProto.getStartCursor()).isEqualTo(ByteString.copyFrom(bytes));
   }
 
   @Test
@@ -59,21 +60,21 @@ public void testEndCursor() {
     Query queryProto = protoPreparer.prepare(newEntityQueryBuilder().setEndCursor(
         Cursor.copyFrom(bytes)).build());
 
-    assertThat(queryProto.getEndCursor(), equalTo(ByteString.copyFrom(bytes)));
+    assertThat(queryProto.getEndCursor()).isEqualTo(ByteString.copyFrom(bytes));
   }
 
   @Test
   public void testOffset() {
     Query queryProto = protoPreparer.prepare(newEntityQueryBuilder().setOffset(5).build());
 
-    assertThat(queryProto.getOffset(), equalTo(5));
+    assertThat(queryProto.getOffset()).isEqualTo(5);
   }
 
   @Test
   public void testLimit() {
     Query queryProto = protoPreparer.prepare(newEntityQueryBuilder().setLimit(5).build());
 
-    assertThat(queryProto.getLimit(), equalTo(Int32Value.of(5)));
+    assertThat(queryProto.getLimit()).isEqualTo(Int32Value.of(5));
   }
 
   @Test
@@ -82,9 +83,9 @@ public void testFilter() {
         .setFilter(PropertyFilter.eq("done", true))
         .build());
 
-    assertThat(queryProto.getFilter(), equalTo(
+    assertThat(queryProto.getFilter()).isEqualTo(
         propertyFilter("done", EQUAL, booleanValue(true))
-    ));
+    );
   }
 
   @Test
@@ -93,8 +94,8 @@ public void testOrderBy() {
         .setOrderBy(OrderBy.asc("dept-id"), OrderBy.asc("rank"))
         .build());
 
-    assertThat(queryProto.getOrder(0), equalTo(propertyOrder("dept-id")));
-    assertThat(queryProto.getOrder(1), equalTo(propertyOrder("rank")));
+    assertThat(queryProto.getOrder(0)).isEqualTo(propertyOrder("dept-id"));
+    assertThat(queryProto.getOrder(1)).isEqualTo(propertyOrder("rank"));
   }
 
   @Test
@@ -103,8 +104,8 @@ public void testDistinctOn() {
         .setDistinctOn("dept-id", "rank")
         .build());
 
-    assertThat(queryProto.getDistinctOn(0), equalTo(propertyReference("dept-id")));
-    assertThat(queryProto.getDistinctOn(1), equalTo(propertyReference("rank")));
+    assertThat(queryProto.getDistinctOn(0)).isEqualTo(propertyReference("dept-id"));
+    assertThat(queryProto.getDistinctOn(1)).isEqualTo(propertyReference("rank"));
   }
 
   @Test
@@ -113,7 +114,7 @@ public void testProjections() {
         .setProjection("dept-id", "rank")
         .build());
 
-    assertThat(queryProto.getProjection(0), equalTo(projection("dept-id")));
-    assertThat(queryProto.getProjection(1), equalTo(projection("rank")));
+    assertThat(queryProto.getProjection(0)).isEqualTo(projection("dept-id"));
+    assertThat(queryProto.getProjection(1)).isEqualTo(projection("rank"));
   }
 }
\ No newline at end of file
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java
index 6a35142a4..859b3c522 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java
@@ -17,11 +17,9 @@
 package com.google.cloud.datastore.aggregation;
 
 import static com.google.cloud.datastore.aggregation.Aggregation.count;
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.MatcherAssert.assertThat;
+import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertTrue;
 
 import com.google.datastore.v1.AggregationQuery;
 import org.junit.Test;
@@ -32,8 +30,8 @@ public class CountAggregationTest {
   public void testCountAggregationWithDefaultValues() {
     AggregationQuery.Aggregation countAggregationPb = count().build().toPb();
 
-    assertThat(countAggregationPb.getCount().getUpTo().getValue(), equalTo(0L));
-    assertThat(countAggregationPb.getAlias(), equalTo(""));
+    assertThat(countAggregationPb.getCount().getUpTo().getValue()).isEqualTo(0L);
+    assertThat(countAggregationPb.getAlias()).isEqualTo("");
   }
 
   @Test
@@ -42,8 +40,8 @@ public void testCountAggregationWithAlias() {
         .as("column_1")
         .build().toPb();
 
-    assertThat(countAggregationPb.getCount().getUpTo().getValue(), equalTo(0L));
-    assertThat(countAggregationPb.getAlias(), equalTo("column_1"));
+    assertThat(countAggregationPb.getCount().getUpTo().getValue()).isEqualTo(0L);
+    assertThat(countAggregationPb.getAlias()).isEqualTo("column_1");
   }
 
   @Test
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java
index c8099e194..c5133a3bf 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java
@@ -20,14 +20,13 @@
 import static com.google.cloud.datastore.StructuredQuery.PropertyFilter.eq;
 import static com.google.cloud.datastore.TestUtils.matches;
 import static com.google.cloud.datastore.aggregation.Aggregation.count;
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.datastore.v1.ReadOptions.ReadConsistency.EVENTUAL;
 import static java.util.Arrays.asList;
 import static org.easymock.EasyMock.anyObject;
 import static org.easymock.EasyMock.expect;
 import static org.easymock.EasyMock.replay;
 import static org.easymock.EasyMock.verify;
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.MatcherAssert.assertThat;
 
 import com.google.cloud.Timestamp;
 import com.google.cloud.datastore.AggregationQuery;
@@ -39,6 +38,7 @@
 import com.google.cloud.datastore.Query;
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.truth.Truth;
 import com.google.datastore.v1.AggregationResultBatch;
 import com.google.datastore.v1.RunAggregationQueryRequest;
 import com.google.datastore.v1.RunAggregationQueryResponse;
@@ -89,12 +89,12 @@ public void shouldExecuteAggregationQuery() {
     AggregationResults aggregationResults = queryExecutor.execute(aggregationQuery);
 
     verify(mockRpc);
-    assertThat(aggregationResults, equalTo(new AggregationResults(asList(
+    assertThat(aggregationResults).isEqualTo(new AggregationResults(asList(
         new AggregationResult(
             ImmutableMap.of("count", LongValue.of(209), "property_2", LongValue.of(100))),
         new AggregationResult(
             ImmutableMap.of("count", LongValue.of(509), "property_2", LongValue.of(100)))
-    ), Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime()))));
+    ), Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime())));
   }
 
   @Test
@@ -121,12 +121,12 @@ public void shouldExecuteAggregationQueryWithReadOptions() {
         eventualConsistency());
 
     verify(mockRpc);
-    assertThat(aggregationResults, equalTo(new AggregationResults(asList(
+    assertThat(aggregationResults).isEqualTo(new AggregationResults(asList(
         new AggregationResult(
             ImmutableMap.of("count", LongValue.of(209), "property_2", LongValue.of(100))),
         new AggregationResult(
             ImmutableMap.of("count", LongValue.of(509), "property_2", LongValue.of(100)))
-    ), Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime()))));
+    ), Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime())));
   }
 
   private RunAggregationQueryResponse dummyAggregationQueryResponse() {
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java
index 077254858..990011d11 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java
@@ -25,14 +25,11 @@
 import static com.google.cloud.datastore.ReadOption.eventualConsistency;
 import static com.google.cloud.datastore.StructuredQuery.PropertyFilter.eq;
 import static com.google.cloud.datastore.aggregation.Aggregation.count;
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.datastore.v1.PropertyFilter.Operator.EQUAL;
 import static com.google.datastore.v1.ReadOptions.ReadConsistency.EVENTUAL;
 import static java.util.Arrays.asList;
 import static java.util.Collections.singletonList;
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.CoreMatchers.is;
-import static org.hamcrest.MatcherAssert.assertThat;
-import static org.junit.Assert.assertNull;
 
 import com.google.cloud.Timestamp;
 import com.google.cloud.datastore.AggregationQuery;
@@ -41,8 +38,8 @@
 import com.google.cloud.datastore.GqlQuery;
 import com.google.cloud.datastore.Query;
 import com.google.cloud.datastore.ReadOption;
-import com.google.cloud.datastore.ReadOption.EventualConsistency;
 import com.google.cloud.datastore.ReadOption.QueryAndReadOptions;
+import com.google.common.truth.Truth;
 import com.google.datastore.v1.GqlQueryParameter;
 import com.google.datastore.v1.RunAggregationQueryRequest;
 import java.util.HashMap;
@@ -88,19 +85,21 @@ public void shouldPrepareAggregationQueryRequestWithGivenStructuredQuery() {
     RunAggregationQueryRequest runAggregationQueryRequest = protoPreparer.prepare(
         QueryAndReadOptions.create(AGGREGATION_OVER_STRUCTURED_QUERY));
 
-    assertThat(runAggregationQueryRequest.getProjectId(), equalTo(PROJECT_ID));
+    assertThat(runAggregationQueryRequest.getProjectId()).isEqualTo(PROJECT_ID);
 
-    assertThat(runAggregationQueryRequest.getPartitionId().getProjectId(), equalTo(PROJECT_ID));
-    assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId(), equalTo(NAMESPACE));
+    assertThat(runAggregationQueryRequest.getPartitionId().getProjectId())
+        .isEqualTo(PROJECT_ID);
+    assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId())
+        .isEqualTo(NAMESPACE);
 
     com.google.datastore.v1.AggregationQuery aggregationQueryProto = runAggregationQueryRequest.getAggregationQuery();
-    assertThat(aggregationQueryProto.getNestedQuery(),
-        equalTo(com.google.datastore.v1.Query.newBuilder()
+    assertThat(aggregationQueryProto.getNestedQuery())
+        .isEqualTo(com.google.datastore.v1.Query.newBuilder()
             .addKind(kind(KIND))
             .setFilter(propertyFilter("done", EQUAL, booleanValue(true)))
-            .build()));
-    assertThat(aggregationQueryProto.getAggregationsList(),
-        equalTo(singletonList(countAggregation("total"))));
+            .build());
+    assertThat(aggregationQueryProto.getAggregationsList())
+        .isEqualTo(singletonList(countAggregation("total")));
   }
 
   @Test
@@ -109,49 +108,52 @@ public void shouldPrepareAggregationQueryRequestWithGivenGqlQuery() {
         QueryAndReadOptions.create(
             AGGREGATION_OVER_GQL_QUERY));
 
-    assertThat(runAggregationQueryRequest.getProjectId(), equalTo(PROJECT_ID));
+    assertThat(runAggregationQueryRequest.getProjectId()).isEqualTo(PROJECT_ID);
 
-    assertThat(runAggregationQueryRequest.getPartitionId().getProjectId(), equalTo(PROJECT_ID));
-    assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId(), equalTo(NAMESPACE));
+    assertThat(runAggregationQueryRequest.getPartitionId().getProjectId())
+        .isEqualTo(PROJECT_ID);
+    assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId())
+        .isEqualTo(NAMESPACE);
 
     com.google.datastore.v1.GqlQuery gqlQueryProto = runAggregationQueryRequest.getGqlQuery();
 
-    assertThat(gqlQueryProto.getQueryString(), equalTo(COMPLETED_TASK_GQL_QUERY.getQueryString()));
-    assertThat(gqlQueryProto.getNamedBindingsMap(),
-        equalTo(new HashMap() {{
+    assertThat(gqlQueryProto.getQueryString())
+        .isEqualTo(COMPLETED_TASK_GQL_QUERY.getQueryString());
+    assertThat(gqlQueryProto.getNamedBindingsMap())
+        .isEqualTo(new HashMap() {{
           put("name", gqlQueryParameter(stringValue("John Doe")));
-        }}));
-    assertThat(gqlQueryProto.getPositionalBindingsList(), equalTo(asList(
+        }});
+    assertThat(gqlQueryProto.getPositionalBindingsList()).isEqualTo(asList(
         gqlQueryParameter(intValue(27))
-    )));
+    ));
   }
 
   @Test
   public void shouldPrepareReadOptionsWithGivenStructuredQuery() {
     RunAggregationQueryRequest eventualConsistencyAggregationRequest = prepareQuery(
         AGGREGATION_OVER_STRUCTURED_QUERY, eventualConsistency());
-    assertThat(eventualConsistencyAggregationRequest.getReadOptions().getReadConsistency(),
-        equalTo(EVENTUAL));
+    assertThat(eventualConsistencyAggregationRequest.getReadOptions().getReadConsistency())
+        .isEqualTo(EVENTUAL);
 
     Timestamp now = Timestamp.now();
     RunAggregationQueryRequest readTimeAggregationRequest = prepareQuery(
         AGGREGATION_OVER_STRUCTURED_QUERY, ReadOption.readTime(now));
-    assertThat(Timestamp.fromProto(readTimeAggregationRequest.getReadOptions().getReadTime()),
-        equalTo(now));
+    assertThat(Timestamp.fromProto(readTimeAggregationRequest.getReadOptions().getReadTime()))
+        .isEqualTo(now);
   }
 
   @Test
   public void shouldPrepareReadOptionsWithGivenGqlQuery() {
     RunAggregationQueryRequest eventualConsistencyAggregationRequest = prepareQuery(
         AGGREGATION_OVER_GQL_QUERY, eventualConsistency());
-    assertThat(eventualConsistencyAggregationRequest.getReadOptions().getReadConsistency(),
-        equalTo(EVENTUAL));
+    assertThat(eventualConsistencyAggregationRequest.getReadOptions().getReadConsistency())
+        .isEqualTo(EVENTUAL);
 
     Timestamp now = Timestamp.now();
     RunAggregationQueryRequest readTimeAggregationRequest = prepareQuery(
         AGGREGATION_OVER_GQL_QUERY, ReadOption.readTime(now));
-    assertThat(Timestamp.fromProto(readTimeAggregationRequest.getReadOptions().getReadTime()),
-        equalTo(now));
+    assertThat(Timestamp.fromProto(readTimeAggregationRequest.getReadOptions().getReadTime()))
+        .isEqualTo(now);
   }
 
   @Test
@@ -169,8 +171,10 @@ public void shouldPrepareAggregationQueryWithoutNamespace() {
     RunAggregationQueryRequest runAggregationQueryFromGqlQuery = protoPreparer.prepare(
         QueryAndReadOptions.create(gqlQueryWithoutNamespace));
 
-    assertThat(runAggregationQueryFromStructuredQuery.getPartitionId().getNamespaceId(), is(""));
-    assertThat(runAggregationQueryFromGqlQuery.getPartitionId().getNamespaceId(), is(""));
+    assertThat(runAggregationQueryFromStructuredQuery.getPartitionId().getNamespaceId())
+        .isEqualTo("");
+    assertThat(runAggregationQueryFromGqlQuery.getPartitionId().getNamespaceId())
+        .isEqualTo("");
 
   }
 
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java
index dafc5dad2..a43ab7f33 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java
@@ -16,13 +16,13 @@
 package com.google.cloud.datastore.execution.response;
 
 import static com.google.cloud.datastore.ProtoTestData.intValue;
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.MatcherAssert.assertThat;
+import static com.google.common.truth.Truth.assertThat;
 
 import com.google.cloud.Timestamp;
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.AggregationResults;
 import com.google.cloud.datastore.LongValue;
+import com.google.common.truth.Truth;
 import com.google.datastore.v1.AggregationResultBatch;
 import com.google.datastore.v1.RunAggregationQueryResponse;
 import com.google.datastore.v1.Value;
@@ -65,10 +65,12 @@ public void shouldTransformAggregationQueryResponse() {
     AggregationResults aggregationResults = responseTransformer.transform(
         runAggregationQueryResponse);
 
-    assertThat(aggregationResults.size(), equalTo(2));
-    assertThat(aggregationResults.get(0), equalTo(new AggregationResult(toDomainValues(result1))));
-    assertThat(aggregationResults.get(1), equalTo(new AggregationResult(toDomainValues(result2))));
-    assertThat(aggregationResults.getReadTime(), equalTo(readTime));
+    assertThat(aggregationResults.size()).isEqualTo(2);
+    assertThat(aggregationResults.get(0))
+        .isEqualTo(new AggregationResult(toDomainValues(result1)));
+    assertThat(aggregationResults.get(1))
+        .isEqualTo(new AggregationResult(toDomainValues(result2)));
+    assertThat(aggregationResults.getReadTime()).isEqualTo(readTime);
 
   }
 
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
index 7ae6042b9..142f94954 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
@@ -18,9 +18,8 @@
 
 import static com.google.cloud.datastore.aggregation.Aggregation.count;
 import static com.google.common.collect.Iterables.getOnlyElement;
+import static com.google.common.truth.Truth.assertThat;
 import static java.util.concurrent.TimeUnit.SECONDS;
-import static org.hamcrest.CoreMatchers.equalTo;
-import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
@@ -568,10 +567,11 @@ public void testRunAggregationQuery() {
 
   /**
    * if an entity is modified or deleted within a transaction, a query or lookup returns the
-   * original version of the entity as of the beginning of the transaction,
-   * or nothing if the entity did not exist then.
+   * original version of the entity as of the beginning of the transaction, or nothing if the entity
+   * did not exist then.
+   *
    * @see 
-   *   Source
+   * Source
    */
   @Test
   public void testRunAggregationQueryInTransactionShouldReturnAConsistentSnapshot() {
@@ -587,46 +587,57 @@ public void testRunAggregationQueryInTransactionShouldReturnAConsistentSnapshot(
         .build();
 
     // original entity count is 2
-    assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(2L));
+    assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+        .isEqualTo(2L);
 
     // FIRST TRANSACTION
     DATASTORE.runInTransaction((TransactionCallable) inFirstTransaction -> {
-        // creating a new entity
+      // creating a new entity
       Entity aNewEntity = Entity.newBuilder(ENTITY2)
           .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build())
           .set("v_int", 10)
           .build();
       inFirstTransaction.put(aNewEntity);
 
-        // count remains 2
-      assertThat(getOnlyElement(inFirstTransaction.runAggregation(aggregationQuery)).get("count"), equalTo(2L));
-      assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(2L));
+      // count remains 2
+      assertThat(
+          getOnlyElement(inFirstTransaction.runAggregation(aggregationQuery)).get("count"))
+          .isEqualTo(2L);
+      assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+          .isEqualTo(2L);
       return null;
     });
     // after first transaction is committed, count is updated to 3 now.
-    assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(3L));
+    assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+        .isEqualTo(3L);
 
     // SECOND TRANSACTION
     DATASTORE.runInTransaction((TransactionCallable) inSecondTransaction -> {
-        // deleting ENTITY2
+      // deleting ENTITY2
       inSecondTransaction.delete(ENTITY2.getKey());
 
-        // count remains 3
-      assertThat(getOnlyElement(inSecondTransaction.runAggregation(aggregationQuery)).get("count"), equalTo(3L));
-      assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(3L));
+      // count remains 3
+      assertThat(
+          getOnlyElement(inSecondTransaction.runAggregation(aggregationQuery)).get("count"))
+          .isEqualTo(3L);
+      assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+          .isEqualTo(3L);
       return null;
     });
     // after second transaction is committed, count is updated to 2 now.
-    assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"), equalTo(2L));
+    assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+        .isEqualTo(2L);
   }
 
   /**
    * Data read or modified by a transaction cannot be concurrently modified.
+   *
    * @see 
-   *   Source
+   * Source
    */
   @Test
-  public void testRunAggregationQueryInAReadWriteTransactionShouldLockTheCountedDocuments() throws Exception {
+  public void testRunAggregationQueryInAReadWriteTransactionShouldLockTheCountedDocuments()
+      throws Exception {
     ExecutorService executor = Executors.newSingleThreadExecutor();
     EntityQuery entityQuery = Query.newEntityQueryBuilder()
         .setNamespace(NAMESPACE)
@@ -642,8 +653,9 @@ public void testRunAggregationQueryInAReadWriteTransactionShouldLockTheCountedDo
     Transaction readWriteTransaction = DATASTORE.newTransaction();
 
     // acquiring lock by executing query in transaction
-    assertThat(getOnlyElement(readWriteTransaction.runAggregation(aggregationQuery)).get("count"),
-        equalTo(2L));
+    assertThat(
+        getOnlyElement(readWriteTransaction.runAggregation(aggregationQuery)).get("count")
+    ).isEqualTo(2L);
 
     // Waiting task will be blocked by ongoing transactions.
     Future addNewEntityTaskOutsideTransaction = executor.submit(() -> {
@@ -665,7 +677,8 @@ public void testRunAggregationQueryInAReadWriteTransactionShouldLockTheCountedDo
   }
 
   @Test
-  public void testRunAggregationQueryInAReadOnlyTransactionShouldNotLockTheCountedDocuments() throws Exception {
+  public void testRunAggregationQueryInAReadOnlyTransactionShouldNotLockTheCountedDocuments()
+      throws Exception {
     ExecutorService executor = Executors.newSingleThreadExecutor();
     EntityQuery entityQuery = Query.newEntityQueryBuilder()
         .setNamespace(NAMESPACE)
@@ -682,8 +695,9 @@ public void testRunAggregationQueryInAReadOnlyTransactionShouldNotLockTheCounted
     Transaction readOnlyTransaction = DATASTORE.newTransaction(transactionOptions);
 
     // Executing query in transaction
-    assertThat(getOnlyElement(readOnlyTransaction.runAggregation(aggregationQuery)).get("count"),
-        equalTo(2L));
+    assertThat(
+        getOnlyElement(readOnlyTransaction.runAggregation(aggregationQuery)).get("count"))
+        .isEqualTo(2L);
 
     // Concurrent write task.
     Future addNewEntityTaskOutsideTransaction = executor.submit(() -> {
@@ -702,8 +716,8 @@ public void testRunAggregationQueryInAReadOnlyTransactionShouldNotLockTheCounted
     readOnlyTransaction.commit();
     executor.shutdownNow();
 
-    assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"),
-        equalTo(3L));
+    assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+        .isEqualTo(3L);
   }
 
   @Test
@@ -1328,7 +1342,7 @@ private void testCountAggregationWith(Consumer configu
     DATASTORE.put(newEntity);
 
     Long countAfterAdd = getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get(alias);
-    assertThat(countAfterAdd, equalTo(expectedCount));
+    assertThat(countAfterAdd).isEqualTo(expectedCount);
 
     DATASTORE.delete(newEntity.getKey());
   }
@@ -1364,12 +1378,12 @@ private void testCountAggregationReadTimeWith(Consumer
 
       Long latestCount = getOnlyElement(DATASTORE.runAggregation(countAggregationQuery))
           .get("total_count");
-      assertThat(latestCount, equalTo(3L));
+      assertThat(latestCount).isEqualTo(3L);
 
       Long oldCount = getOnlyElement(
           DATASTORE.runAggregation(countAggregationQuery, ReadOption.readTime(now))
       ).get("total_count");
-      assertThat(oldCount, equalTo(2L));
+      assertThat(oldCount).isEqualTo(2L);
     } finally {
       DATASTORE.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
     }

From 5801641e67651ef166d6e8cf4b7ca228eec864f6 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Fri, 7 Oct 2022 10:41:16 +0530
Subject: [PATCH 48/82] Removing 'limit' api from count query samples

---
 .../datastore/AggregationQuerySample.java     | 40 +++++++++++++++++--
 .../AggregationQuerySampleTestIT.java         |  8 +++-
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
index 38330beb8..17a5b5e89 100644
--- a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
+++ b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
@@ -59,33 +59,67 @@ public void aggregationQueryAndCountAggregation() {
     AggregationQuery allCandidatesCountQuery = Query.newAggregationQueryBuilder()
         .over(selectAllCandidates)
         .addAggregation(Aggregation.count().as("total_count"))
-        .addAggregation(Aggregation.count().as("count_with_limit").limit(2))
         .build();
     // Executing aggregation query
     AggregationResult allCandidatesCountQueryResult = Iterables.getOnlyElement(
         datastore.runAggregation(allCandidatesCountQuery));
 
-    System.out.printf("We have at least %d candidates", allCandidatesCountQueryResult.get("count_with_limit")); // 2
     System.out.printf("Total candidates count is %d", allCandidatesCountQueryResult.get("total_count")); // 3
 
+    // [END datastore_count_aggregation_query]
+
+    datastore.delete(candidate1Key, candidate2Key, candidate3Key);
+  }
+
+  public void aggregationQueryAndCountAggregationWithPropertyFilter() {
+    // [START datastore_count_aggregation_query_with_filters]
+
+    // Instantiates a client
+    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+
+    // The kind for the new entity
+    String kind = "Candidate";
+
+    Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1");
+    Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2");
+    Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3");
+
+    // Save all the candidates
+    datastore.put(
+        Entity.newBuilder(candidate1Key).set("qualified", true).build(),
+        Entity.newBuilder(candidate2Key).set("qualified", false).build(),
+        Entity.newBuilder(candidate3Key).set("qualified", true).build()
+    );
 
     EntityQuery qualifiedCandidates = Query.newEntityQueryBuilder()
         .setKind(kind)
         .setFilter(PropertyFilter.eq("qualified", true))
         .build();
+    EntityQuery unQualifiedCandidates = Query.newEntityQueryBuilder()
+        .setKind(kind)
+        .setFilter(PropertyFilter.eq("qualified", false))
+        .build();
     // Creating an aggregation query to get the count of all qualified candidates
     AggregationQuery qualifiedCandidatesCountQuery = Query.newAggregationQueryBuilder()
         .over(qualifiedCandidates)
         .addAggregation(Aggregation.count().as("total_qualified_count"))
         .build();
+    // Creating an aggregation query to get the count of all unqualified candidates
+    AggregationQuery unqualifiedCandidatesCountQuery = Query.newAggregationQueryBuilder()
+        .over(unQualifiedCandidates)
+        .addAggregation(Aggregation.count().as("total_unqualified_count"))
+        .build();
 
     // Executing aggregation query
     AggregationResult qualifiedCandidatesCountQueryResult = Iterables.getOnlyElement(
         datastore.runAggregation(qualifiedCandidatesCountQuery));
+    AggregationResult unQualifiedCandidatesCountQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(unqualifiedCandidatesCountQuery));
 
     System.out.printf("Total qualified candidates count is %d", qualifiedCandidatesCountQueryResult.get("total_qualified_count")); // 2
+    System.out.printf("Total unqualified candidates count is %d", unQualifiedCandidatesCountQueryResult.get("total_unqualified_count")); // 1
 
-    // [END datastore_count_aggregation_query]
+    // [END datastore_count_aggregation_query_with_filters]
 
     datastore.delete(candidate1Key, candidate2Key, candidate3Key);
   }
diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
index 0a86df3b0..e15aaeb22 100644
--- a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
@@ -31,9 +31,15 @@ public class AggregationQuerySampleTestIT {
   public void testAggregationQueryAndCountAggregationSample() {
     sample.aggregationQueryAndCountAggregation();
 
-    systemsOutRule.assertContains("We have at least 2 candidates");
     systemsOutRule.assertContains("Total candidates count is 3");
+  }
+
+  @Test
+  public void testAggregationQueryAndCountAggregationWithPropertyFilterSample() {
+    sample.aggregationQueryAndCountAggregationWithPropertyFilter();
+
     systemsOutRule.assertContains("Total qualified candidates count is 2");
+    systemsOutRule.assertContains("Total unqualified candidates count is 1");
   }
 
   @Test

From 41df1bc7633112356e13db118279b96bed85e313 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Fri, 7 Oct 2022 11:12:13 +0530
Subject: [PATCH 49/82] fixing unit test

---
 .../cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java | 2 +-
 .../cloud/datastore/execution/AggregationQueryExecutorTest.java | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java
index 1bee0c4e8..316938de6 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java
@@ -39,7 +39,7 @@ public class RetryAndTraceDatastoreRpcDecoratorTest {
   public static final int MAX_ATTEMPTS = 3;
   private DatastoreRpc mockDatastoreRpc;
   private TraceUtil mockTraceUtil;
-  private DatastoreOptions datastoreOptions = DatastoreOptions.getDefaultInstance();
+  private DatastoreOptions datastoreOptions = DatastoreOptions.newBuilder().setProjectId("project-id").build();
   private RetrySettings retrySettings = RetrySettings.newBuilder()
       .setMaxAttempts(MAX_ATTEMPTS)
       .build();
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java
index c5133a3bf..18c36c6c2 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java
@@ -62,7 +62,7 @@ public class AggregationQueryExecutorTest {
   @Before
   public void setUp() throws Exception {
     mockRpc = EasyMock.createStrictMock(DatastoreRpc.class);
-    datastoreOptions = DatastoreOptions.newBuilder().setNamespace(NAMESPACE).build();
+    datastoreOptions = DatastoreOptions.newBuilder().setProjectId("project-id").setNamespace(NAMESPACE).build();
     queryExecutor = new AggregationQueryExecutor(mockRpc, datastoreOptions);
   }
 

From 16065847b6d041d9ab24a1e03d2290489ec52a74 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Fri, 7 Oct 2022 11:26:26 +0530
Subject: [PATCH 50/82] Getting rid off Double braces initialization syntax

---
 .../datastore/GqlQueryProtoPreparerTest.java  | 11 +++++-----
 .../AggregationQueryExecutorTest.java         | 22 +++++++++----------
 ...regationQueryRequestProtoPreparerTest.java |  7 +++---
 ...gregationQueryResponseTransformerTest.java | 18 +++++++--------
 4 files changed, 29 insertions(+), 29 deletions(-)

diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java
index 1af698306..d37627d62 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java
@@ -24,8 +24,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
-import com.google.common.truth.Truth;
-import com.google.datastore.v1.GqlQueryParameter;
+import com.google.common.collect.ImmutableMap;
 import java.util.HashMap;
 import org.junit.Test;
 
@@ -59,10 +58,10 @@ public void testNamedBinding() {
     );
 
     assertThat(gqlQuery.getNamedBindingsMap())
-        .isEqualTo(new HashMap() {{
-          put("name", gqlQueryParameter(stringValue("John Doe")));
-          put("age", gqlQueryParameter(intValue(27)));
-        }});
+        .isEqualTo(new HashMap<>(ImmutableMap.of(
+            "name", gqlQueryParameter(stringValue("John Doe")),
+            "age", gqlQueryParameter(intValue(27))
+        )));
   }
 
   @Test
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java
index 18c36c6c2..8b6d6d2cf 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java
@@ -38,7 +38,6 @@
 import com.google.cloud.datastore.Query;
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.truth.Truth;
 import com.google.datastore.v1.AggregationResultBatch;
 import com.google.datastore.v1.RunAggregationQueryRequest;
 import com.google.datastore.v1.RunAggregationQueryResponse;
@@ -62,7 +61,8 @@ public class AggregationQueryExecutorTest {
   @Before
   public void setUp() throws Exception {
     mockRpc = EasyMock.createStrictMock(DatastoreRpc.class);
-    datastoreOptions = DatastoreOptions.newBuilder().setProjectId("project-id").setNamespace(NAMESPACE).build();
+    datastoreOptions = DatastoreOptions.newBuilder().setProjectId("project-id")
+        .setNamespace(NAMESPACE).build();
     queryExecutor = new AggregationQueryExecutor(mockRpc, datastoreOptions);
   }
 
@@ -130,15 +130,15 @@ public void shouldExecuteAggregationQueryWithReadOptions() {
   }
 
   private RunAggregationQueryResponse dummyAggregationQueryResponse() {
-    Map result1 = new HashMap() {{
-      put("count", intValue(209));
-      put("property_2", intValue(100));
-    }};
-
-    Map result2 = new HashMap() {{
-      put("count", intValue(509));
-      put("property_2", intValue(100));
-    }};
+    Map result1 = new HashMap<>(ImmutableMap.of(
+        "count", intValue(209),
+        "property_2", intValue(100)
+    ));
+
+    Map result2 = new HashMap<>(ImmutableMap.of(
+        "count", intValue(509),
+        "property_2", intValue(100)
+    ));
 
     AggregationResultBatch resultBatch = AggregationResultBatch.newBuilder()
         .addAggregationResults(com.google.datastore.v1.AggregationResult.newBuilder()
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java
index 990011d11..dec1615fe 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java
@@ -39,6 +39,7 @@
 import com.google.cloud.datastore.Query;
 import com.google.cloud.datastore.ReadOption;
 import com.google.cloud.datastore.ReadOption.QueryAndReadOptions;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.truth.Truth;
 import com.google.datastore.v1.GqlQueryParameter;
 import com.google.datastore.v1.RunAggregationQueryRequest;
@@ -120,9 +121,9 @@ public void shouldPrepareAggregationQueryRequestWithGivenGqlQuery() {
     assertThat(gqlQueryProto.getQueryString())
         .isEqualTo(COMPLETED_TASK_GQL_QUERY.getQueryString());
     assertThat(gqlQueryProto.getNamedBindingsMap())
-        .isEqualTo(new HashMap() {{
-          put("name", gqlQueryParameter(stringValue("John Doe")));
-        }});
+        .isEqualTo(new HashMap<>(
+            ImmutableMap.of("name", gqlQueryParameter(stringValue("John Doe")))
+        ));
     assertThat(gqlQueryProto.getPositionalBindingsList()).isEqualTo(asList(
         gqlQueryParameter(intValue(27))
     ));
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java
index a43ab7f33..e15fc0def 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java
@@ -22,7 +22,7 @@
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.AggregationResults;
 import com.google.cloud.datastore.LongValue;
-import com.google.common.truth.Truth;
+import com.google.common.collect.ImmutableMap;
 import com.google.datastore.v1.AggregationResultBatch;
 import com.google.datastore.v1.RunAggregationQueryResponse;
 import com.google.datastore.v1.Value;
@@ -40,15 +40,15 @@ public class AggregationQueryResponseTransformerTest {
 
   @Test
   public void shouldTransformAggregationQueryResponse() {
-    Map result1 = new HashMap() {{
-      put("count", intValue(209));
-      put("property_2", intValue(100));
-    }};
+    Map result1 = new HashMap<>(ImmutableMap.of(
+        "count", intValue(209),
+        "property_2", intValue(100)
+    ));
 
-    Map result2 = new HashMap() {{
-      put("count", intValue(509));
-      put("property_2", intValue(100));
-    }};
+    Map result2 = new HashMap<>(ImmutableMap.of(
+        "count", intValue(509),
+        "property_2", intValue(100)
+    ));
     Timestamp readTime = Timestamp.now();
 
     AggregationResultBatch resultBatch = AggregationResultBatch.newBuilder()

From ca27e75c50ae7609cb990250561871ced7223296 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Fri, 7 Oct 2022 11:37:33 +0530
Subject: [PATCH 51/82] Fixing lint

---
 .../google/datastore/v1/client/Datastore.java |   3 +-
 .../cloud/datastore/AggregationQuery.java     |  42 +--
 .../cloud/datastore/AggregationResult.java    |   9 +-
 .../cloud/datastore/AggregationResults.java   |  25 +-
 .../com/google/cloud/datastore/Datastore.java |  13 +-
 .../google/cloud/datastore/DatastoreImpl.java |  21 +-
 .../cloud/datastore/DatastoreReader.java      |   3 +-
 .../com/google/cloud/datastore/GqlQuery.java  |   1 -
 .../datastore/GqlQueryProtoPreparer.java      |   4 +-
 .../com/google/cloud/datastore/Query.java     |  14 +-
 .../cloud/datastore/QueryResultsImpl.java     |   5 +-
 .../google/cloud/datastore/ReadOption.java    |  12 +-
 .../datastore/ReadOptionProtoPreparer.java    |  33 +-
 .../google/cloud/datastore/RecordQuery.java   |   6 +-
 .../RetryAndTraceDatastoreRpcDecorator.java   |  13 +-
 .../cloud/datastore/StructuredQuery.java      |   1 -
 .../StructuredQueryProtoPreparer.java         |   8 +-
 .../com/google/cloud/datastore/TraceUtil.java |   3 +-
 .../datastore/aggregation/Aggregation.java    |   8 +-
 .../aggregation/AggregationBuilder.java       |   7 +-
 .../aggregation/CountAggregation.java         |  18 +-
 .../execution/AggregationQueryExecutor.java   |  23 +-
 .../datastore/execution/QueryExecutor.java    |  10 +-
 .../AggregationQueryRequestProtoPreparer.java |  23 +-
 .../AggregationQueryResponseTransformer.java  |  30 +-
 .../cloud/datastore/spi/v1/DatastoreRpc.java  |   2 +-
 .../cloud/datastore/AggregationQueryTest.java | 123 ++++---
 .../datastore/AggregationResultTest.java      |  11 +-
 .../google/cloud/datastore/DatastoreTest.java |  18 +-
 .../datastore/GqlQueryProtoPreparerTest.java  |  43 +--
 .../google/cloud/datastore/ProtoTestData.java |  25 +-
 .../ReadOptionProtoPreparerTest.java          |  47 +--
 ...etryAndTraceDatastoreRpcDecoratorTest.java |  29 +-
 .../StructuredQueryProtoPreparerTest.java     |  40 +--
 .../com/google/cloud/datastore/TestUtils.java |  24 +-
 .../aggregation/CountAggregationTest.java     |  12 +-
 .../datastore/emulator/EmulatorProxy.java     |   3 +-
 .../datastore/emulator/ProxyDispatcher.java   |  31 +-
 .../AggregationQueryExecutorTest.java         | 140 ++++----
 ...regationQueryRequestProtoPreparerTest.java | 161 ++++-----
 ...gregationQueryResponseTransformerTest.java |  69 ++--
 .../cloud/datastore/it/ITDatastoreTest.java   | 334 +++++++++---------
 42 files changed, 726 insertions(+), 721 deletions(-)

diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java
index 4de2cc96e..09101c94b 100644
--- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java
+++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java
@@ -123,7 +123,8 @@ public RunQueryResponse runQuery(RunQueryRequest request) throws DatastoreExcept
     }
   }
 
-  public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) throws DatastoreException {
+  public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request)
+      throws DatastoreException {
     try (InputStream is = remoteRpc.call("runAggregationQuery", request)) {
       return RunAggregationQueryResponse.parseFrom(is);
     } catch (IOException exception) {
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java
index e6c0c0dc3..8e7e7cdc9 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java
@@ -18,23 +18,20 @@
 import static com.google.cloud.datastore.AggregationQuery.Mode.GQL;
 import static com.google.cloud.datastore.AggregationQuery.Mode.STRUCTURED;
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.cloud.datastore.aggregation.Aggregation;
 import com.google.cloud.datastore.aggregation.AggregationBuilder;
-import java.util.ArrayList;
 import java.util.HashSet;
-import java.util.List;
 import java.util.Set;
 
 /**
  * An implementation of a Google Cloud Datastore Query that returns {@link AggregationResults}, It
- * can be constructed by providing a nested query ({@link StructuredQuery} or {@link GqlQuery})
- * to run the aggregations on and a set of {@link Aggregation}.
-
- * 
{@link StructuredQuery} example:

- * 
{@code
+ * can be constructed by providing a nested query ({@link StructuredQuery} or {@link GqlQuery}) to
+ * run the aggregations on and a set of {@link Aggregation}.
+ *
+ * 
{@link StructuredQuery} example:
  *
+ * 
{@code
  * EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
  *    .setKind("Task")
  *    .build();
@@ -49,10 +46,10 @@
  * }

  *
  * 
{@link GqlQuery} example:

- * 
{@code
  *
+ * 
{@code
  * GqlQuery selectAllGqlQuery = Query.newGqlQueryBuilder(
-*         "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(100) AS count_upto_100 OVER(SELECT * FROM Task)"
+ *         "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(100) AS count_upto_100 OVER(SELECT * FROM Task)"
  *     )
  *     .setAllowLiteral(true)
  *     .build();
@@ -76,10 +73,11 @@ public class AggregationQuery extends Query {
   private final Mode mode;
   private GqlQuery nestedGqlQuery;
 
-  AggregationQuery(String namespace, Set aggregations,
-      StructuredQuery nestedQuery) {
+  AggregationQuery(
+      String namespace, Set aggregations, StructuredQuery nestedQuery) {
     super(namespace);
-    checkArgument(!aggregations.isEmpty(),
+    checkArgument(
+        !aggregations.isEmpty(),
         "At least one aggregation is required for an aggregation query to run");
     this.aggregations = aggregations;
     this.nestedStructuredQuery = nestedQuery;
@@ -97,16 +95,18 @@ public Set getAggregations() {
     return aggregations;
   }
 
-  /** Returns the underlying {@link StructuredQuery for this Query}.
-   * Returns null if created with {@link GqlQuery}
-   * */
+  /**
+   * Returns the underlying {@link StructuredQuery for this Query}. Returns null if created with
+   * {@link GqlQuery}
+   */
   public StructuredQuery getNestedStructuredQuery() {
     return nestedStructuredQuery;
   }
 
-  /** Returns the underlying {@link GqlQuery for this Query}.
-   * Returns null if created with {@link StructuredQuery}
-   * */
+  /**
+   * Returns the underlying {@link GqlQuery for this Query}. Returns null if created with {@link
+   * StructuredQuery}
+   */
   public GqlQuery getNestedGqlQuery() {
     return nestedGqlQuery;
   }
@@ -157,8 +157,8 @@ public Builder over(GqlQuery nestedQuery) {
 
     public AggregationQuery build() {
       boolean nestedQueryProvided = nestedGqlQuery != null || nestedStructuredQuery != null;
-      checkArgument(nestedQueryProvided,
-          "Nested query is required for an aggregation query to run");
+      checkArgument(
+          nestedQueryProvided, "Nested query is required for an aggregation query to run");
 
       if (mode == GQL) {
         return new AggregationQuery(namespace, nestedGqlQuery);
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java
index 72421a5de..6e086c30b 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java
@@ -17,14 +17,11 @@
 
 import com.google.common.base.MoreObjects;
 import com.google.common.base.MoreObjects.ToStringHelper;
-import java.util.Iterator;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Objects;
 
-/**
- * Represents a result of an {@link AggregationQuery} query submission.
- */
+/** Represents a result of an {@link AggregationQuery} query submission. */
 public class AggregationResult {
 
   private final Map properties;
@@ -35,7 +32,9 @@ public AggregationResult(Map properties) {
 
   /**
    * Returns a result value for the given alias.
-   * @param alias A custom alias provided in the query or an autogenerated alias in the form of 'property_\d'
+   *
+   * @param alias A custom alias provided in the query or an autogenerated alias in the form of
+   *     'property_\d'
    * @return An aggregation result value for the given alias.
    */
   public Long get(String alias) {
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java
index 9b6148175..420964449 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java
@@ -24,33 +24,30 @@
 import java.util.Objects;
 
 /**
- * The result of an {@link AggregationQuery} query submission. Contains a
- * {@link List} and readTime {@link Timestamp} in it.
+ * The result of an {@link AggregationQuery} query submission. Contains a {@link
+ * List} and readTime {@link Timestamp} in it.
  *
- *  This can be used to iterate over underlying {@link List} directly.
+ * 
This can be used to iterate over underlying {@link List} directly.
  *
- *  Though {@link com.google.cloud.datastore.aggregation.CountAggregation} is guaranteed to return
- *  only one {@link AggregationResult} as part of its execution.
+ * 
Though {@link com.google.cloud.datastore.aggregation.CountAggregation} is guaranteed to return
+ * only one {@link AggregationResult} as part of its execution.
  *
- *  In the future, we might support more complex {@link AggregationQuery} that might result in
- *  multiple {@link AggregationResult}
+ * 
In the future, we might support more complex {@link AggregationQuery} that might result in
+ * multiple {@link AggregationResult}
  */
 public class AggregationResults implements Iterable {
 
   private final List aggregationResults;
   private final Timestamp readTime;
 
-  public AggregationResults(List aggregationResults,
-      Timestamp readTime) {
+  public AggregationResults(List aggregationResults, Timestamp readTime) {
     checkNotNull(aggregationResults, "Aggregation results cannot be null");
     checkNotNull(readTime, "readTime cannot be null");
     this.aggregationResults = aggregationResults;
     this.readTime = readTime;
   }
 
-  /**
-   * Returns {@link Iterator} for underlying {@link List}.
-   */
+  /** Returns {@link Iterator} for underlying {@link List}. */
   @Override
   public Iterator iterator() {
     return this.aggregationResults.iterator();
@@ -65,9 +62,7 @@ public AggregationResult get(int index) {
     return this.aggregationResults.get(index);
   }
 
-  /**
-   * Returns read timestamp this result batch was returned from.
-   */
+  /** Returns read timestamp this result batch was returned from. */
   public Timestamp getReadTime() {
     return this.readTime;
   }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
index d8d2431fc..d58c001b1 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
@@ -463,14 +463,14 @@ interface TransactionCallable {
    QueryResults run(Query query, ReadOption... options);
 
   /**
-   * Submits a {@link AggregationQuery} and returns {@link AggregationResults}.
-   * {@link ReadOption}s can be specified if desired.
+   * Submits a {@link AggregationQuery} and returns {@link AggregationResults}. {@link ReadOption}s
+   * can be specified if desired.
    *
    * 
Example of running an {@link AggregationQuery} to find the count of entities of one kind.
    *
-   * 
{@link StructuredQuery} example:

-   * 
{@code
+   * 
{@link StructuredQuery} example:
    *
+   * 
{@code
    * EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
    *    .setKind("Task")
    *    .build();
@@ -485,8 +485,8 @@ interface TransactionCallable {
    * }

    *
    * 
{@link GqlQuery} example:

-   * 
{@code
    *
+   * 
{@code
    * GqlQuery selectAllGqlQuery = Query.newGqlQueryBuilder(
    *         "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(100) AS count_upto_100 OVER(SELECT * FROM Task)"
    *     )
@@ -501,10 +501,11 @@ interface TransactionCallable {
    *   System.out.println(aggregationResult.get("count_upto_100"));
    * }
    * }

+   *
    * @throws DatastoreException upon failure
    * @return {@link AggregationResults}
    */
-  default AggregationResults runAggregation(AggregationQuery query, ReadOption... options){
+  default AggregationResults runAggregation(AggregationQuery query, ReadOption... options) {
     throw new UnsupportedOperationException("Not implemented.");
   }
 }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
index 78582f196..cf5c6e3ad 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
@@ -69,10 +69,10 @@ final class DatastoreImpl extends BaseService implements Datas
         MoreObjects.firstNonNull(options.getRetrySettings(), ServiceOptions.getNoRetrySettings());
 
     readOptionProtoPreparer = new ReadOptionProtoPreparer();
-    aggregationQueryExecutor = new AggregationQueryExecutor(
-        new RetryAndTraceDatastoreRpcDecorator(datastoreRpc, traceUtil, retrySettings, options),
-        options
-    );
+    aggregationQueryExecutor =
+        new AggregationQueryExecutor(
+            new RetryAndTraceDatastoreRpcDecorator(datastoreRpc, traceUtil, retrySettings, options),
+            options);
   }
 
   @Override
@@ -191,17 +191,19 @@ public  QueryResults run(Query query, ReadOption... options) {
   }
 
   @SuppressWarnings("unchecked")
-   QueryResults run(Optional readOptionsPb, Query query) {
-    return new QueryResultsImpl(this, readOptionsPb, (RecordQuery) query, query.getNamespace());
+   QueryResults run(
+      Optional readOptionsPb, Query query) {
+    return new QueryResultsImpl(
+        this, readOptionsPb, (RecordQuery) query, query.getNamespace());
   }
 
   @Override
-  public AggregationResults runAggregation(AggregationQuery query){
+  public AggregationResults runAggregation(AggregationQuery query) {
     return aggregationQueryExecutor.execute(query);
   }
 
   @Override
-  public AggregationResults runAggregation(AggregationQuery query, ReadOption... options){
+  public AggregationResults runAggregation(AggregationQuery query, ReadOption... options) {
     return aggregationQueryExecutor.execute(query, options);
   }
 
@@ -375,7 +377,8 @@ public List fetch(Iterable keys, ReadOption... options) {
     return DatastoreHelper.fetch(this, Iterables.toArray(keys, Key.class), options);
   }
 
-  Iterator get(Optional readOptionsPb, final Key... keys) {
+  Iterator get(
+      Optional readOptionsPb, final Key... keys) {
     if (keys.length == 0) {
       return Collections.emptyIterator();
     }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java
index 86a18ca3d..2a3071f3c 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java
@@ -54,13 +54,12 @@ public interface DatastoreReader {
    */
    QueryResults run(Query query);
 
-
   /**
    * Submits a {@link AggregationQuery} and returns {@link AggregationResults}.
    *
    * @throws DatastoreException upon failure
    */
-  default AggregationResults runAggregation(AggregationQuery query){
+  default AggregationResults runAggregation(AggregationQuery query) {
     throw new UnsupportedOperationException("Not implemented.");
   }
 }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java
index 814e68ae2..8a8c5e9fc 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java
@@ -83,7 +83,6 @@ public final class GqlQuery extends Query implements RecordQuery {
 
   private final ResultType resultType;
 
-
   static final class Binding implements Serializable {
 
     private static final long serialVersionUID = 2344746877591371548L;
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQueryProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQueryProtoPreparer.java
index 90837d75e..5269740f7 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQueryProtoPreparer.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQueryProtoPreparer.java
@@ -21,8 +21,8 @@
 import java.util.Map;
 
 @InternalApi
-public class GqlQueryProtoPreparer implements
-    ProtoPreparer, com.google.datastore.v1.GqlQuery> {
+public class GqlQueryProtoPreparer
+    implements ProtoPreparer, com.google.datastore.v1.GqlQuery> {
 
   @Override
   public com.google.datastore.v1.GqlQuery prepare(GqlQuery gqlQuery) {
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java
index 5802320a8..8870cf520 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java
@@ -16,18 +16,15 @@
 
 package com.google.cloud.datastore;
 
-import static com.google.common.base.Preconditions.checkNotNull;
-
 import com.google.common.base.MoreObjects;
 import com.google.common.base.MoreObjects.ToStringHelper;
-import com.google.common.base.Preconditions;
 import com.google.common.collect.Maps;
 import java.io.Serializable;
 import java.util.Map;
 
 /**
- * A Google Cloud Datastore query. For usage examples see {@link GqlQuery}, {@link
- * StructuredQuery} and {@link AggregationQuery}.
+ * A Google Cloud Datastore query. For usage examples see {@link GqlQuery}, {@link StructuredQuery}
+ * and {@link AggregationQuery}.
  *
  * 
Note that queries require proper indexing. See Cloud Datastore Index
@@ -168,7 +165,6 @@ ToStringHelper toStringHelper() {
     return MoreObjects.toStringHelper(this).add("namespace", namespace);
   }
 
-
   /**
    * Returns a new {@link GqlQuery} builder.
    *
@@ -264,9 +260,9 @@ public static ProjectionEntityQuery.Builder newProjectionEntityQueryBuilder() {
    *
    * 
Example of creating and running an {@link AggregationQuery}.
    *
-   * 
{@link StructuredQuery} example:

-   * 
{@code
+   * 
{@link StructuredQuery} example:
    *
+   * 
{@code
    * EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
    *    .setKind("Task")
    *    .build();
@@ -279,8 +275,8 @@ public static ProjectionEntityQuery.Builder newProjectionEntityQueryBuilder() {
    * }

    *
    * 
{@link GqlQuery} example:

-   * 
{@code
    *
+   * 
{@code
    * GqlQuery selectAllGqlQuery = Query.newGqlQueryBuilder(
    *         "AGGREGATE COUNT(*) AS total_count OVER(SELECT * FROM Task)"
    *     )
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
index d079d5837..6170c0b8b 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
@@ -42,7 +42,10 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults
   private MoreResultsType moreResults;
 
   QueryResultsImpl(
-      DatastoreImpl datastore, Optional readOptionsPb, RecordQuery query, String namespace) {
+      DatastoreImpl datastore,
+      Optional readOptionsPb,
+      RecordQuery query,
+      String namespace) {
     this.datastore = datastore;
     this.readOptionsPb = readOptionsPb;
     this.query = query;
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java
index 1fad1842b..be5644da0 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java
@@ -72,9 +72,7 @@ public Timestamp time() {
     }
   }
 
-  /**
-   * Specifies transaction to be used when running a {@link Query}.
-   */
+  /** Specifies transaction to be used when running a {@link Query}. */
   @InternalApi
   static class TransactionId extends ReadOption {
 
@@ -89,8 +87,7 @@ public ByteString getTransactionId() {
     }
   }
 
-  private ReadOption() {
-  }
+  private ReadOption() {}
 
   /**
    * Returns a {@code ReadOption} that specifies eventual consistency, allowing Datastore to return
@@ -172,10 +169,9 @@ public static > QueryAndReadOptions create(Q query) {
       return new QueryAndReadOptions<>(query);
     }
 
-    public static > QueryAndReadOptions create(Q query,
-        List readOptions) {
+    public static > QueryAndReadOptions create(
+        Q query, List readOptions) {
       return new QueryAndReadOptions<>(query, readOptions);
     }
   }
-
 }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java
index c70d2c7b3..15713b02f 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOptionProtoPreparer.java
@@ -25,12 +25,11 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 
 @InternalApi
-public class ReadOptionProtoPreparer implements
-    ProtoPreparer, Optional> {
+public class ReadOptionProtoPreparer
+    implements ProtoPreparer, Optional> {
 
   @Override
   public Optional prepare(List options) {
@@ -38,8 +37,7 @@ public Optional prepare(List options) {
       return Optional.empty();
     }
     com.google.datastore.v1.ReadOptions readOptionsPb = null;
-    Map, ReadOption> optionsByType =
-        ReadOption.asImmutableMap(options);
+    Map, ReadOption> optionsByType = ReadOption.asImmutableMap(options);
 
     boolean moreThanOneReadOption = optionsByType.keySet().size() > 1;
     if (moreThanOneReadOption) {
@@ -48,22 +46,22 @@ public Optional prepare(List options) {
     }
 
     if (optionsByType.containsKey(EventualConsistency.class)) {
-      readOptionsPb = ReadOptions.newBuilder()
-          .setReadConsistency(ReadConsistency.EVENTUAL)
-          .build();
+      readOptionsPb = ReadOptions.newBuilder().setReadConsistency(ReadConsistency.EVENTUAL).build();
     }
 
     if (optionsByType.containsKey(ReadTime.class)) {
-      readOptionsPb = ReadOptions.newBuilder()
-          .setReadTime(((ReadTime) optionsByType.get(ReadTime.class)).time().toProto())
-          .build();
+      readOptionsPb =
+          ReadOptions.newBuilder()
+              .setReadTime(((ReadTime) optionsByType.get(ReadTime.class)).time().toProto())
+              .build();
     }
 
     if (optionsByType.containsKey(TransactionId.class)) {
-      readOptionsPb = ReadOptions.newBuilder()
-          .setTransaction(
-              ((TransactionId) optionsByType.get(TransactionId.class)).getTransactionId())
-          .build();
+      readOptionsPb =
+          ReadOptions.newBuilder()
+              .setTransaction(
+                  ((TransactionId) optionsByType.get(TransactionId.class)).getTransactionId())
+              .build();
     }
     return Optional.ofNullable(readOptionsPb);
   }
@@ -71,9 +69,8 @@ public Optional prepare(List options) {
   private String getInvalidOptions(Map, ReadOption> optionsByType) {
     String regex = "([a-z])([A-Z]+)";
     String replacement = "$1 $2";
-    return optionsByType
-        .keySet()
-        .stream().map(Class::getSimpleName)
+    return optionsByType.keySet().stream()
+        .map(Class::getSimpleName)
         .map(s -> s.replaceAll(regex, replacement).toLowerCase())
         .collect(Collectors.joining(", "));
   }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java
index 6b9bc9125..770837602 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java
@@ -18,11 +18,9 @@
 import com.google.api.core.InternalApi;
 import com.google.cloud.datastore.Query.ResultType;
 
-/**
- * An internal marker interface to represent {@link Query} that returns the entity records.
- */
+/** An internal marker interface to represent {@link Query} that returns the entity records. */
 @InternalApi
-public interface RecordQuery{
+public interface RecordQuery {
 
   ResultType getType();
 
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java
index 7d15df188..8d9a7aac4 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java
@@ -54,8 +54,11 @@ public class RetryAndTraceDatastoreRpcDecorator implements DatastoreRpc {
   private final RetrySettings retrySettings;
   private final DatastoreOptions datastoreOptions;
 
-  public RetryAndTraceDatastoreRpcDecorator(DatastoreRpc datastoreRpc, TraceUtil traceUtil,
-      RetrySettings retrySettings, DatastoreOptions datastoreOptions) {
+  public RetryAndTraceDatastoreRpcDecorator(
+      DatastoreRpc datastoreRpc,
+      TraceUtil traceUtil,
+      RetrySettings retrySettings,
+      DatastoreOptions datastoreOptions) {
     this.datastoreRpc = datastoreRpc;
     this.traceUtil = traceUtil;
     this.retrySettings = retrySettings;
@@ -100,13 +103,15 @@ public RunQueryResponse runQuery(RunQueryRequest request) {
 
   @Override
   public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) {
-    return invokeRpc(() -> datastoreRpc.runAggregationQuery(request), SPAN_NAME_RUN_AGGREGATION_QUERY);
+    return invokeRpc(
+        () -> datastoreRpc.runAggregationQuery(request), SPAN_NAME_RUN_AGGREGATION_QUERY);
   }
 
   public  O invokeRpc(Callable block, String startSpan) {
     Span span = traceUtil.startSpan(startSpan);
     try (Scope scope = traceUtil.getTracer().withSpan(span)) {
-      return RetryHelper.runWithRetries(block, this.retrySettings, EXCEPTION_HANDLER, this.datastoreOptions.getClock());
+      return RetryHelper.runWithRetries(
+          block, this.retrySettings, EXCEPTION_HANDLER, this.datastoreOptions.getClock());
     } catch (RetryHelperException e) {
       span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
       throw DatastoreException.translateAndThrow(e);
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java
index 8db6a506d..852627a08 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java
@@ -104,7 +104,6 @@ public abstract class StructuredQuery extends Query implements RecordQuery
 
   private final ResultType resultType;
 
-
   public abstract static class Filter implements Serializable {
 
     private static final long serialVersionUID = -6443285436239990860L;
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQueryProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQueryProtoPreparer.java
index ad07dd838..fda6f8f4a 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQueryProtoPreparer.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQueryProtoPreparer.java
@@ -53,9 +53,11 @@ public Query prepare(StructuredQuery query) {
           com.google.datastore.v1.PropertyReference.newBuilder().setName(value).build());
     }
     for (String value : query.getProjection()) {
-      com.google.datastore.v1.Projection expressionPb = com.google.datastore.v1.Projection.newBuilder()
-          .setProperty(com.google.datastore.v1.PropertyReference.newBuilder().setName(value).build())
-          .build();
+      com.google.datastore.v1.Projection expressionPb =
+          com.google.datastore.v1.Projection.newBuilder()
+              .setProperty(
+                  com.google.datastore.v1.PropertyReference.newBuilder().setName(value).build())
+              .build();
       queryPb.addProjection(expressionPb);
     }
 
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java
index e6909b4df..57525d15d 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java
@@ -39,7 +39,8 @@ public class TraceUtil {
   static final String SPAN_NAME_RESERVEIDS = "CloudDatastoreOperation.reserveIds";
   static final String SPAN_NAME_ROLLBACK = "CloudDatastoreOperation.rollback";
   static final String SPAN_NAME_RUNQUERY = "CloudDatastoreOperation.runQuery";
-  static final String SPAN_NAME_RUN_AGGREGATION_QUERY = "CloudDatastoreOperation.runAggregationQuery";
+  static final String SPAN_NAME_RUN_AGGREGATION_QUERY =
+      "CloudDatastoreOperation.runAggregationQuery";
   static final EndSpanOptions END_SPAN_OPTIONS =
       EndSpanOptions.builder().setSampleToLocalSpanStore(true).build();
 
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java
index b2b48726d..7bd2bbb38 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/Aggregation.java
@@ -30,9 +30,7 @@ public Aggregation(String alias) {
     this.alias = alias;
   }
 
-  /**
-   * Returns the alias for this aggregation.
-   */
+  /** Returns the alias for this aggregation. */
   public String getAlias() {
     return alias;
   }
@@ -40,9 +38,7 @@ public String getAlias() {
   @InternalApi
   public abstract AggregationQuery.Aggregation toPb();
 
-  /**
-   * Returns a {@link CountAggregation} builder.
-   */
+  /** Returns a {@link CountAggregation} builder. */
   public static CountAggregation.Builder count() {
     return new CountAggregation.Builder();
   }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java
index 3a6c826f9..ce23edcf0 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/AggregationBuilder.java
@@ -17,10 +17,11 @@
 package com.google.cloud.datastore.aggregation;
 
 /**
- * An interface to represent the builders which build and customize {@link Aggregation} for
- * {@link com.google.cloud.datastore.AggregationQuery}.
+ * An interface to represent the builders which build and customize {@link Aggregation} for {@link
+ * com.google.cloud.datastore.AggregationQuery}.
  *
- * Used by {@link com.google.cloud.datastore.AggregationQuery.Builder#addAggregation(AggregationBuilder)}.
+ * 
Used by {@link
+ * com.google.cloud.datastore.AggregationQuery.Builder#addAggregation(AggregationBuilder)}.
  */
 public interface AggregationBuilder {
   A build();
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java
index b205896fd..07a68959e 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/aggregation/CountAggregation.java
@@ -18,17 +18,12 @@
 
 import com.google.datastore.v1.AggregationQuery;
 import com.google.datastore.v1.AggregationQuery.Aggregation.Count;
-import com.google.protobuf.Int64Value;
 import java.util.Objects;
 
-/**
- * Represents an {@link Aggregation} which returns count.
- */
+/** Represents an {@link Aggregation} which returns count. */
 public class CountAggregation extends Aggregation {
 
-  /**
-   * @param alias Alias to used when running this aggregation.
-   */
+  /** @param alias Alias to used when running this aggregation. */
   public CountAggregation(String alias) {
     super(alias);
   }
@@ -37,8 +32,8 @@ public CountAggregation(String alias) {
   public AggregationQuery.Aggregation toPb() {
     Count.Builder countBuilder = Count.newBuilder();
 
-    AggregationQuery.Aggregation.Builder aggregationBuilder = AggregationQuery.Aggregation.newBuilder()
-        .setCount(countBuilder);
+    AggregationQuery.Aggregation.Builder aggregationBuilder =
+        AggregationQuery.Aggregation.newBuilder().setCount(countBuilder);
     if (this.getAlias() != null) {
       aggregationBuilder.setAlias(this.getAlias());
     }
@@ -62,9 +57,7 @@ public int hashCode() {
     return Objects.hash(getAlias());
   }
 
-  /**
-   * A builder class to create and customize a {@link CountAggregation}.
-   */
+  /** A builder class to create and customize a {@link CountAggregation}. */
   public static class Builder implements AggregationBuilder {
 
     private String alias;
@@ -79,5 +72,4 @@ public CountAggregation build() {
       return new CountAggregation(alias);
     }
   }
-
 }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java
index 945683dc6..14e425845 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java
@@ -33,8 +33,8 @@
  * {@link AggregationResults}.
  */
 @InternalApi
-public class AggregationQueryExecutor implements
-    QueryExecutor {
+public class AggregationQueryExecutor
+    implements QueryExecutor {
 
   private final DatastoreRpc datastoreRpc;
   private final AggregationQueryRequestProtoPreparer protoPreparer;
@@ -48,18 +48,19 @@ public AggregationQueryExecutor(DatastoreRpc datastoreRpc, DatastoreOptions data
 
   @Override
   public AggregationResults execute(AggregationQuery query, ReadOption... readOptions) {
-    RunAggregationQueryRequest runAggregationQueryRequest = getRunAggregationQueryRequest(
-        query, readOptions);
-    RunAggregationQueryResponse runAggregationQueryResponse = this.datastoreRpc.runAggregationQuery(
-        runAggregationQueryRequest);
+    RunAggregationQueryRequest runAggregationQueryRequest =
+        getRunAggregationQueryRequest(query, readOptions);
+    RunAggregationQueryResponse runAggregationQueryResponse =
+        this.datastoreRpc.runAggregationQuery(runAggregationQueryRequest);
     return this.responseTransformer.transform(runAggregationQueryResponse);
   }
 
-  private RunAggregationQueryRequest getRunAggregationQueryRequest(AggregationQuery query,
-      ReadOption... readOptions) {
-    QueryAndReadOptions queryAndReadOptions = readOptions == null ?
-        QueryAndReadOptions.create(query) :
-        QueryAndReadOptions.create(query, Arrays.asList(readOptions));
+  private RunAggregationQueryRequest getRunAggregationQueryRequest(
+      AggregationQuery query, ReadOption... readOptions) {
+    QueryAndReadOptions queryAndReadOptions =
+        readOptions == null
+            ? QueryAndReadOptions.create(query)
+            : QueryAndReadOptions.create(query, Arrays.asList(readOptions));
     return this.protoPreparer.prepare(queryAndReadOptions);
   }
 }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java
index b737da522..856c64a02 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/QueryExecutor.java
@@ -20,11 +20,11 @@
 import com.google.cloud.datastore.ReadOption;
 
 /**
- * An internal functional interface whose implementation has the responsibility to execute a
- * {@link Query} and returns the result. This class will have the responsibility to orchestrate
- * between {@link com.google.cloud.datastore.execution.request.ProtoPreparer},
- * {@link com.google.cloud.datastore.spi.v1.DatastoreRpc} and
- * {@link com.google.cloud.datastore.execution.response.ResponseTransformer} layers.
+ * An internal functional interface whose implementation has the responsibility to execute a {@link
+ * Query} and returns the result. This class will have the responsibility to orchestrate between
+ * {@link com.google.cloud.datastore.execution.request.ProtoPreparer}, {@link
+ * com.google.cloud.datastore.spi.v1.DatastoreRpc} and {@link
+ * com.google.cloud.datastore.execution.response.ResponseTransformer} layers.
  *
  * @param  A {@link Query} to execute.
  * @param  the type of result produced by Query.
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java
index 7418ac5c6..b5da8d9fe 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java
@@ -35,8 +35,8 @@
 import java.util.Optional;
 
 @InternalApi
-public class AggregationQueryRequestProtoPreparer implements
-    ProtoPreparer, RunAggregationQueryRequest> {
+public class AggregationQueryRequestProtoPreparer
+    implements ProtoPreparer, RunAggregationQueryRequest> {
 
   private final DatastoreOptions datastoreOptions;
   private final StructuredQueryProtoPreparer structuredQueryProtoPreparer;
@@ -56,9 +56,10 @@ public RunAggregationQueryRequest prepare(
     AggregationQuery aggregationQuery = aggregationQueryAndReadOptions.getQuery();
     List readOptions = aggregationQueryAndReadOptions.getReadOptions();
     PartitionId partitionId = getPartitionId(aggregationQuery);
-    RunAggregationQueryRequest.Builder aggregationQueryRequestBuilder = RunAggregationQueryRequest.newBuilder()
-        .setPartitionId(partitionId)
-        .setProjectId(datastoreOptions.getProjectId());
+    RunAggregationQueryRequest.Builder aggregationQueryRequestBuilder =
+        RunAggregationQueryRequest.newBuilder()
+            .setPartitionId(partitionId)
+            .setProjectId(datastoreOptions.getProjectId());
 
     if (aggregationQuery.getMode() == GQL) {
       aggregationQueryRequestBuilder.setGqlQuery(buildGqlQuery(aggregationQuery));
@@ -77,11 +78,11 @@ private GqlQuery buildGqlQuery(AggregationQuery aggregationQuery) {
 
   private com.google.datastore.v1.AggregationQuery getAggregationQuery(
       AggregationQuery aggregationQuery) {
-    Query nestedQueryProto = structuredQueryProtoPreparer.prepare(
-        aggregationQuery.getNestedStructuredQuery());
+    Query nestedQueryProto =
+        structuredQueryProtoPreparer.prepare(aggregationQuery.getNestedStructuredQuery());
 
-    com.google.datastore.v1.AggregationQuery.Builder aggregationQueryProtoBuilder = com.google.datastore.v1.AggregationQuery.newBuilder()
-        .setNestedQuery(nestedQueryProto);
+    com.google.datastore.v1.AggregationQuery.Builder aggregationQueryProtoBuilder =
+        com.google.datastore.v1.AggregationQuery.newBuilder().setNestedQuery(nestedQueryProto);
     for (Aggregation aggregation : aggregationQuery.getAggregations()) {
       aggregationQueryProtoBuilder.addAggregations(aggregation.toPb());
     }
@@ -89,8 +90,8 @@ private com.google.datastore.v1.AggregationQuery getAggregationQuery(
   }
 
   private PartitionId getPartitionId(AggregationQuery aggregationQuery) {
-    PartitionId.Builder builder = PartitionId.newBuilder()
-        .setProjectId(datastoreOptions.getProjectId());
+    PartitionId.Builder builder =
+        PartitionId.newBuilder().setProjectId(datastoreOptions.getProjectId());
     if (aggregationQuery.getNamespace() != null) {
       builder.setNamespaceId(aggregationQuery.getNamespace());
     }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java
index dba180d22..1515a1147 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformer.java
@@ -31,28 +31,28 @@
 import java.util.stream.Collectors;
 
 @InternalApi
-public class AggregationQueryResponseTransformer implements
-    ResponseTransformer {
+public class AggregationQueryResponseTransformer
+    implements ResponseTransformer {
 
   @Override
   public AggregationResults transform(RunAggregationQueryResponse response) {
     Timestamp readTime = Timestamp.fromProto(response.getBatch().getReadTime());
-    List aggregationResults = response
-        .getBatch()
-        .getAggregationResultsList()
-        .stream()
-        .map(aggregationResult -> new AggregationResult(resultWithLongValues(aggregationResult)))
-        .collect(Collectors.toCollection(LinkedList::new));
+    List aggregationResults =
+        response.getBatch().getAggregationResultsList().stream()
+            .map(
+                aggregationResult -> new AggregationResult(resultWithLongValues(aggregationResult)))
+            .collect(Collectors.toCollection(LinkedList::new));
     return new AggregationResults(aggregationResults, readTime);
   }
 
-  private Map resultWithLongValues(com.google.datastore.v1.AggregationResult aggregationResult) {
-    return aggregationResult.getAggregatePropertiesMap()
-        .entrySet()
-        .stream()
-        .map((Function, Entry>) entry ->
-            new SimpleEntry<>(entry.getKey(), (LongValue) LongValue.fromPb(entry.getValue())))
+  private Map resultWithLongValues(
+      com.google.datastore.v1.AggregationResult aggregationResult) {
+    return aggregationResult.getAggregatePropertiesMap().entrySet().stream()
+        .map(
+            (Function, Entry>)
+                entry ->
+                    new SimpleEntry<>(
+                        entry.getKey(), (LongValue) LongValue.fromPb(entry.getValue())))
         .collect(Collectors.toMap(Entry::getKey, Entry::getValue));
   }
 }
-
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java
index cdaab95a3..33b8e11ea 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java
@@ -93,7 +93,7 @@ BeginTransactionResponse beginTransaction(BeginTransactionRequest request)
    *
    * @throws DatastoreException upon failure
    */
-  default RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request){
+  default RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) {
     throw new UnsupportedOperationException("Not implemented.");
   }
 }
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java
index b43e55ec5..840d23bca 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationQueryTest.java
@@ -25,7 +25,6 @@
 
 import com.google.cloud.datastore.aggregation.CountAggregation;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.truth.Truth;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
@@ -34,110 +33,111 @@ public class AggregationQueryTest {
 
   private static final String KIND = "Task";
   private static final String NAMESPACE = "ns";
-  private static final EntityQuery COMPLETED_TASK_QUERY = Query.newEntityQueryBuilder()
-      .setNamespace(NAMESPACE)
-      .setKind(KIND)
-      .setFilter(eq("done", true))
-      .setLimit(100)
-      .build();
+  private static final EntityQuery COMPLETED_TASK_QUERY =
+      Query.newEntityQueryBuilder()
+          .setNamespace(NAMESPACE)
+          .setKind(KIND)
+          .setFilter(eq("done", true))
+          .setLimit(100)
+          .build();
 
-  @Rule
-  public ExpectedException exceptionRule = ExpectedException.none();
+  @Rule public ExpectedException exceptionRule = ExpectedException.none();
 
   @Test
   public void testAggregations() {
-    AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .addAggregation(new CountAggregation("total"))
-        .over(COMPLETED_TASK_QUERY)
-        .build();
+    AggregationQuery aggregationQuery =
+        Query.newAggregationQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .addAggregation(new CountAggregation("total"))
+            .over(COMPLETED_TASK_QUERY)
+            .build();
 
     assertThat(aggregationQuery.getNamespace()).isEqualTo(NAMESPACE);
-    assertThat(aggregationQuery.getAggregations()).isEqualTo(
-        ImmutableSet.of(count().as("total").build()));
+    assertThat(aggregationQuery.getAggregations())
+        .isEqualTo(ImmutableSet.of(count().as("total").build()));
     assertThat(aggregationQuery.getNestedStructuredQuery()).isEqualTo(COMPLETED_TASK_QUERY);
     assertThat(aggregationQuery.getMode()).isEqualTo(STRUCTURED);
   }
 
-
   @Test
   public void testAggregationBuilderWithMoreThanOneAggregations() {
-    AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .addAggregation(count().as("total"))
-        .addAggregation(count().as("new_total"))
-        .over(COMPLETED_TASK_QUERY)
-        .build();
+    AggregationQuery aggregationQuery =
+        Query.newAggregationQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .addAggregation(count().as("total"))
+            .addAggregation(count().as("new_total"))
+            .over(COMPLETED_TASK_QUERY)
+            .build();
 
     assertThat(aggregationQuery.getNamespace()).isEqualTo(NAMESPACE);
-    assertThat(aggregationQuery.getAggregations()).isEqualTo(ImmutableSet.of(
-        count().as("total").build(),
-        count().as("new_total").build()
-    ));
+    assertThat(aggregationQuery.getAggregations())
+        .isEqualTo(ImmutableSet.of(count().as("total").build(), count().as("new_total").build()));
     assertThat(aggregationQuery.getNestedStructuredQuery()).isEqualTo(COMPLETED_TASK_QUERY);
     assertThat(aggregationQuery.getMode()).isEqualTo(STRUCTURED);
   }
 
   @Test
   public void testAggregationBuilderWithDuplicateAggregations() {
-    AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .addAggregation(count().as("total"))
-        .addAggregation(count().as("total"))
-        .over(COMPLETED_TASK_QUERY)
-        .build();
+    AggregationQuery aggregationQuery =
+        Query.newAggregationQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .addAggregation(count().as("total"))
+            .addAggregation(count().as("total"))
+            .over(COMPLETED_TASK_QUERY)
+            .build();
 
     assertThat(aggregationQuery.getNamespace()).isEqualTo(NAMESPACE);
-    assertThat(aggregationQuery.getAggregations()).isEqualTo(ImmutableSet.of(
-        count().as("total").build()
-    ));
+    assertThat(aggregationQuery.getAggregations())
+        .isEqualTo(ImmutableSet.of(count().as("total").build()));
     assertThat(aggregationQuery.getNestedStructuredQuery()).isEqualTo(COMPLETED_TASK_QUERY);
     assertThat(aggregationQuery.getMode()).isEqualTo(STRUCTURED);
   }
 
   @Test
   public void testAggregationQueryBuilderWithoutNamespace() {
-    AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
-        .addAggregation(count().as("total"))
-        .over(COMPLETED_TASK_QUERY)
-        .build();
+    AggregationQuery aggregationQuery =
+        Query.newAggregationQueryBuilder()
+            .addAggregation(count().as("total"))
+            .over(COMPLETED_TASK_QUERY)
+            .build();
 
     assertNull(aggregationQuery.getNamespace());
-    assertThat(aggregationQuery.getAggregations()).isEqualTo(ImmutableSet.of(
-        count().as("total").build()
-    ));
+    assertThat(aggregationQuery.getAggregations())
+        .isEqualTo(ImmutableSet.of(count().as("total").build()));
     assertThat(aggregationQuery.getNestedStructuredQuery()).isEqualTo(COMPLETED_TASK_QUERY);
     assertThat(aggregationQuery.getMode()).isEqualTo(STRUCTURED);
   }
 
   @Test
   public void testAggregationQueryBuilderWithoutNestedQuery() {
-    assertThrows("Nested query is required for an aggregation query to run",
+    assertThrows(
+        "Nested query is required for an aggregation query to run",
         IllegalArgumentException.class,
-        () -> Query.newAggregationQueryBuilder()
-            .setNamespace(NAMESPACE)
-            .addAggregation(count().as("total"))
-            .build());
+        () ->
+            Query.newAggregationQueryBuilder()
+                .setNamespace(NAMESPACE)
+                .addAggregation(count().as("total"))
+                .build());
   }
 
   @Test
   public void testAggregationQueryBuilderWithoutAggregation() {
-    assertThrows("At least one aggregation is required for an aggregation query to run",
+    assertThrows(
+        "At least one aggregation is required for an aggregation query to run",
         IllegalArgumentException.class,
-        () -> Query.newAggregationQueryBuilder()
-            .setNamespace(NAMESPACE)
-            .over(COMPLETED_TASK_QUERY)
-            .build());
+        () ->
+            Query.newAggregationQueryBuilder()
+                .setNamespace(NAMESPACE)
+                .over(COMPLETED_TASK_QUERY)
+                .build());
   }
 
   @Test
   public void testAggregationQueryBuilderWithGqlQuery() {
     GqlQuery gqlQuery = Query.newGqlQueryBuilder("SELECT * FROM Task WHERE done = true").build();
 
-    AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .over(gqlQuery)
-        .build();
+    AggregationQuery aggregationQuery =
+        Query.newAggregationQueryBuilder().setNamespace(NAMESPACE).over(gqlQuery).build();
 
     assertThat(aggregationQuery.getNestedGqlQuery()).isEqualTo(gqlQuery);
     assertThat(aggregationQuery.getMode()).isEqualTo(GQL);
@@ -145,10 +145,9 @@ public void testAggregationQueryBuilderWithGqlQuery() {
 
   @Test
   public void testAggregationQueryBuilderWithoutProvidingAnyNestedQuery() {
-    assertThrows("Nested query is required for an aggregation query to run",
+    assertThrows(
+        "Nested query is required for an aggregation query to run",
         IllegalArgumentException.class,
-        () -> Query.newAggregationQueryBuilder()
-            .setNamespace(NAMESPACE)
-            .build());
+        () -> Query.newAggregationQueryBuilder().setNamespace(NAMESPACE).build());
   }
-}
\ No newline at end of file
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java
index 4ff694fd2..06a5cb5f7 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java
@@ -24,12 +24,13 @@ public class AggregationResultTest {
 
   @Test
   public void shouldGetAggregationResultValueByAlias() {
-    AggregationResult aggregationResult = new AggregationResult(ImmutableMap.of(
-        "count", LongValue.of(45),
-        "property_2", LongValue.of(30)
-    ));
+    AggregationResult aggregationResult =
+        new AggregationResult(
+            ImmutableMap.of(
+                "count", LongValue.of(45),
+                "property_2", LongValue.of(30)));
 
     assertThat(aggregationResult.get("count")).isEqualTo(45L);
     assertThat(aggregationResult.get("property_2")).isEqualTo(30L);
   }
-}
\ No newline at end of file
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
index 67a4b41cd..81e66da55 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
@@ -43,7 +43,6 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterators;
 import com.google.common.collect.Lists;
-import com.google.common.truth.Truth;
 import com.google.datastore.v1.BeginTransactionRequest;
 import com.google.datastore.v1.BeginTransactionResponse;
 import com.google.datastore.v1.CommitRequest;
@@ -537,18 +536,19 @@ public void testGqlQueryPagination() throws DatastoreException {
   @Test
   public void testRunAggregationQuery() {
     EntityQuery selectAllQuery = Query.newEntityQueryBuilder().build();
-    AggregationQuery getCountQuery = Query.newAggregationQueryBuilder()
-        .addAggregation(count().as("total_count"))
-        .over(selectAllQuery)
-        .build();
-    AggregationResult resultBeforeInsert = getOnlyElement(
-        datastoreEmulatorProxy.runAggregation(getCountQuery));
+    AggregationQuery getCountQuery =
+        Query.newAggregationQueryBuilder()
+            .addAggregation(count().as("total_count"))
+            .over(selectAllQuery)
+            .build();
+    AggregationResult resultBeforeInsert =
+        getOnlyElement(datastoreEmulatorProxy.runAggregation(getCountQuery));
     assertThat(resultBeforeInsert.get("total_count")).isEqualTo(2L);
 
     datastore.put(ENTITY3);
 
-    AggregationResult resultAfterInsert = getOnlyElement(
-        datastoreEmulatorProxy.runAggregation(getCountQuery));
+    AggregationResult resultAfterInsert =
+        getOnlyElement(datastoreEmulatorProxy.runAggregation(getCountQuery));
     assertThat(resultAfterInsert.get("total_count")).isEqualTo(3L);
   }
 
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java
index d37627d62..0d2e0ede7 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GqlQueryProtoPreparerTest.java
@@ -42,40 +42,33 @@ public void testQueryString() {
 
   @Test
   public void testAllowLiteral() {
-    assertTrue(protoPreparer.prepare(
-        gqlQueryBuilder.setAllowLiteral(true).build()).getAllowLiterals());
-    assertFalse(protoPreparer.prepare(
-        gqlQueryBuilder.setAllowLiteral(false).build()).getAllowLiterals());
+    assertTrue(
+        protoPreparer.prepare(gqlQueryBuilder.setAllowLiteral(true).build()).getAllowLiterals());
+    assertFalse(
+        protoPreparer.prepare(gqlQueryBuilder.setAllowLiteral(false).build()).getAllowLiterals());
   }
 
   @Test
   public void testNamedBinding() {
-    com.google.datastore.v1.GqlQuery gqlQuery = protoPreparer.prepare(
-        gqlQueryBuilder
-            .setBinding("name", "John Doe")
-            .setBinding("age", 27)
-            .build()
-    );
+    com.google.datastore.v1.GqlQuery gqlQuery =
+        protoPreparer.prepare(
+            gqlQueryBuilder.setBinding("name", "John Doe").setBinding("age", 27).build());
 
     assertThat(gqlQuery.getNamedBindingsMap())
-        .isEqualTo(new HashMap<>(ImmutableMap.of(
-            "name", gqlQueryParameter(stringValue("John Doe")),
-            "age", gqlQueryParameter(intValue(27))
-        )));
+        .isEqualTo(
+            new HashMap<>(
+                ImmutableMap.of(
+                    "name", gqlQueryParameter(stringValue("John Doe")),
+                    "age", gqlQueryParameter(intValue(27)))));
   }
 
   @Test
   public void testPositionalBinding() {
-    com.google.datastore.v1.GqlQuery gqlQuery = protoPreparer.prepare(
-        gqlQueryBuilder
-            .addBinding("John Doe")
-            .addBinding(27)
-            .build()
-    );
+    com.google.datastore.v1.GqlQuery gqlQuery =
+        protoPreparer.prepare(gqlQueryBuilder.addBinding("John Doe").addBinding(27).build());
 
-    assertThat(gqlQuery.getPositionalBindingsList()).isEqualTo(asList(
-        gqlQueryParameter(stringValue("John Doe")),
-        gqlQueryParameter(intValue(27))
-    ));
+    assertThat(gqlQuery.getPositionalBindingsList())
+        .isEqualTo(
+            asList(gqlQueryParameter(stringValue("John Doe")), gqlQueryParameter(intValue(27))));
   }
-}
\ No newline at end of file
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ProtoTestData.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ProtoTestData.java
index 5cff2cf11..e76b3a0a0 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ProtoTestData.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ProtoTestData.java
@@ -52,13 +52,14 @@ public static KindExpression kind(String kind) {
   }
 
   public static Filter propertyFilter(String propertyName, Operator operator, Value value) {
-    return Filter.newBuilder().setPropertyFilter(
-        com.google.datastore.v1.PropertyFilter.newBuilder()
-            .setProperty(propertyReference(propertyName))
-            .setOp(operator)
-            .setValue(value)
-            .build()
-    ).build();
+    return Filter.newBuilder()
+        .setPropertyFilter(
+            com.google.datastore.v1.PropertyFilter.newBuilder()
+                .setProperty(propertyReference(propertyName))
+                .setOp(operator)
+                .setValue(value)
+                .build())
+        .build();
   }
 
   public static PropertyReference propertyReference(String value) {
@@ -66,10 +67,7 @@ public static PropertyReference propertyReference(String value) {
   }
 
   public static Aggregation countAggregation(String alias) {
-    return Aggregation.newBuilder()
-        .setAlias(alias)
-        .setCount(Count.newBuilder().build())
-        .build();
+    return Aggregation.newBuilder().setAlias(alias).setCount(Count.newBuilder().build()).build();
   }
 
   public static Aggregation countAggregation(String alias, long limit) {
@@ -87,9 +85,6 @@ public static PropertyOrder propertyOrder(String value) {
   }
 
   public static Projection projection(String value) {
-    return Projection.newBuilder()
-        .setProperty(propertyReference(value))
-        .build();
+    return Projection.newBuilder().setProperty(propertyReference(value)).build();
   }
-
 }
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java
index 539f9a9e5..e0631348d 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/ReadOptionProtoPreparerTest.java
@@ -21,19 +21,13 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.datastore.v1.ReadOptions.ReadConsistency.EVENTUAL;
 import static java.util.Collections.singletonList;
-import static org.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.assertTrue;
 
 import com.google.cloud.Timestamp;
 import com.google.common.collect.ImmutableList;
-import com.google.common.truth.Truth;
 import com.google.datastore.v1.ReadOptions;
-import com.google.protobuf.ByteString;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.Optional;
 import org.junit.Test;
 
@@ -43,19 +37,28 @@ public class ReadOptionProtoPreparerTest {
 
   @Test
   public void shouldThrowErrorWhenUsingMultipleReadOptions() {
-    assertThrows(DatastoreException.class,
-        () -> protoPreparer.prepare(
-            Arrays.asList(eventualConsistency(), readTime(Timestamp.now()))));
-    assertThrows(DatastoreException.class,
-        () -> protoPreparer.prepare(
-            Arrays.asList(eventualConsistency(), transactionId("transaction-id"))));
-    assertThrows(DatastoreException.class,
-        () -> protoPreparer.prepare(
-            Arrays.asList(transactionId("transaction-id"), readTime(Timestamp.now()))));
-    assertThrows(DatastoreException.class,
-        () -> protoPreparer.prepare(
-            Arrays.asList(eventualConsistency(), readTime(Timestamp.now()),
-                transactionId("transaction-id"))));
+    assertThrows(
+        DatastoreException.class,
+        () ->
+            protoPreparer.prepare(Arrays.asList(eventualConsistency(), readTime(Timestamp.now()))));
+    assertThrows(
+        DatastoreException.class,
+        () ->
+            protoPreparer.prepare(
+                Arrays.asList(eventualConsistency(), transactionId("transaction-id"))));
+    assertThrows(
+        DatastoreException.class,
+        () ->
+            protoPreparer.prepare(
+                Arrays.asList(transactionId("transaction-id"), readTime(Timestamp.now()))));
+    assertThrows(
+        DatastoreException.class,
+        () ->
+            protoPreparer.prepare(
+                Arrays.asList(
+                    eventualConsistency(),
+                    readTime(Timestamp.now()),
+                    transactionId("transaction-id"))));
   }
 
   @Test
@@ -76,8 +79,8 @@ public void shouldPrepareReadOptionsWithReadTime() {
   @Test
   public void shouldPrepareReadOptionsWithTransactionId() {
     String dummyTransactionId = "transaction-id";
-    Optional readOptions = protoPreparer.prepare(singletonList(transactionId(
-        dummyTransactionId)));
+    Optional readOptions =
+        protoPreparer.prepare(singletonList(transactionId(dummyTransactionId)));
 
     assertThat(readOptions.get().getTransaction().toStringUtf8()).isEqualTo(dummyTransactionId);
   }
@@ -91,4 +94,4 @@ public void shouldReturnNullWhenReadOptionsIsNull() {
   public void shouldReturnNullWhenReadOptionsIsAnEmptyList() {
     assertFalse(protoPreparer.prepare(ImmutableList.of()).isPresent());
   }
-}
\ No newline at end of file
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java
index 316938de6..b86355afa 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecoratorTest.java
@@ -39,10 +39,10 @@ public class RetryAndTraceDatastoreRpcDecoratorTest {
   public static final int MAX_ATTEMPTS = 3;
   private DatastoreRpc mockDatastoreRpc;
   private TraceUtil mockTraceUtil;
-  private DatastoreOptions datastoreOptions = DatastoreOptions.newBuilder().setProjectId("project-id").build();
-  private RetrySettings retrySettings = RetrySettings.newBuilder()
-      .setMaxAttempts(MAX_ATTEMPTS)
-      .build();
+  private DatastoreOptions datastoreOptions =
+      DatastoreOptions.newBuilder().setProjectId("project-id").build();
+  private RetrySettings retrySettings =
+      RetrySettings.newBuilder().setMaxAttempts(MAX_ATTEMPTS).build();
 
   private RetryAndTraceDatastoreRpcDecorator datastoreRpcDecorator;
 
@@ -50,17 +50,24 @@ public class RetryAndTraceDatastoreRpcDecoratorTest {
   public void setUp() throws Exception {
     mockDatastoreRpc = createStrictMock(DatastoreRpc.class);
     mockTraceUtil = createStrictMock(TraceUtil.class);
-    datastoreRpcDecorator = new RetryAndTraceDatastoreRpcDecorator(mockDatastoreRpc, mockTraceUtil, retrySettings, datastoreOptions);
+    datastoreRpcDecorator =
+        new RetryAndTraceDatastoreRpcDecorator(
+            mockDatastoreRpc, mockTraceUtil, retrySettings, datastoreOptions);
   }
 
   @Test
   public void testRunAggregationQuery() {
     Span mockSpan = createStrictMock(Span.class);
-    RunAggregationQueryRequest aggregationQueryRequest = RunAggregationQueryRequest.getDefaultInstance();
-    RunAggregationQueryResponse aggregationQueryResponse = RunAggregationQueryResponse.getDefaultInstance();
+    RunAggregationQueryRequest aggregationQueryRequest =
+        RunAggregationQueryRequest.getDefaultInstance();
+    RunAggregationQueryResponse aggregationQueryResponse =
+        RunAggregationQueryResponse.getDefaultInstance();
 
     expect(mockDatastoreRpc.runAggregationQuery(aggregationQueryRequest))
-        .andThrow(new DatastoreException(UNAVAILABLE.getNumber(), "API not accessible currently", UNAVAILABLE.name())).times(2)
+        .andThrow(
+            new DatastoreException(
+                UNAVAILABLE.getNumber(), "API not accessible currently", UNAVAILABLE.name()))
+        .times(2)
         .andReturn(aggregationQueryResponse);
     expect(mockTraceUtil.startSpan(SPAN_NAME_RUN_AGGREGATION_QUERY)).andReturn(mockSpan);
     expect(mockTraceUtil.getTracer()).andReturn(createNiceMock(Tracer.class));
@@ -68,10 +75,10 @@ public void testRunAggregationQuery() {
 
     replay(mockDatastoreRpc, mockTraceUtil, mockSpan);
 
-    RunAggregationQueryResponse actualAggregationQueryResponse = datastoreRpcDecorator.runAggregationQuery(
-        aggregationQueryRequest);
+    RunAggregationQueryResponse actualAggregationQueryResponse =
+        datastoreRpcDecorator.runAggregationQuery(aggregationQueryRequest);
 
     assertThat(actualAggregationQueryResponse).isSameInstanceAs(aggregationQueryResponse);
     verify(mockDatastoreRpc, mockTraceUtil, mockSpan);
   }
-}
\ No newline at end of file
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java
index 2479ff912..60937fc28 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/StructuredQueryProtoPreparerTest.java
@@ -26,7 +26,6 @@
 
 import com.google.cloud.datastore.StructuredQuery.OrderBy;
 import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
-import com.google.common.truth.Truth;
 import com.google.datastore.v1.KindExpression;
 import com.google.datastore.v1.Query;
 import com.google.protobuf.ByteString;
@@ -48,8 +47,9 @@ public void testKind() {
   @Test
   public void testStartCursor() {
     byte[] bytes = {1, 2};
-    Query queryProto = protoPreparer.prepare(newEntityQueryBuilder().setStartCursor(
-        Cursor.copyFrom(bytes)).build());
+    Query queryProto =
+        protoPreparer.prepare(
+            newEntityQueryBuilder().setStartCursor(Cursor.copyFrom(bytes)).build());
 
     assertThat(queryProto.getStartCursor()).isEqualTo(ByteString.copyFrom(bytes));
   }
@@ -57,8 +57,8 @@ public void testStartCursor() {
   @Test
   public void testEndCursor() {
     byte[] bytes = {1, 2};
-    Query queryProto = protoPreparer.prepare(newEntityQueryBuilder().setEndCursor(
-        Cursor.copyFrom(bytes)).build());
+    Query queryProto =
+        protoPreparer.prepare(newEntityQueryBuilder().setEndCursor(Cursor.copyFrom(bytes)).build());
 
     assertThat(queryProto.getEndCursor()).isEqualTo(ByteString.copyFrom(bytes));
   }
@@ -79,20 +79,20 @@ public void testLimit() {
 
   @Test
   public void testFilter() {
-    Query queryProto = protoPreparer.prepare(newEntityQueryBuilder()
-        .setFilter(PropertyFilter.eq("done", true))
-        .build());
+    Query queryProto =
+        protoPreparer.prepare(
+            newEntityQueryBuilder().setFilter(PropertyFilter.eq("done", true)).build());
 
-    assertThat(queryProto.getFilter()).isEqualTo(
-        propertyFilter("done", EQUAL, booleanValue(true))
-    );
+    assertThat(queryProto.getFilter()).isEqualTo(propertyFilter("done", EQUAL, booleanValue(true)));
   }
 
   @Test
   public void testOrderBy() {
-    Query queryProto = protoPreparer.prepare(newEntityQueryBuilder()
-        .setOrderBy(OrderBy.asc("dept-id"), OrderBy.asc("rank"))
-        .build());
+    Query queryProto =
+        protoPreparer.prepare(
+            newEntityQueryBuilder()
+                .setOrderBy(OrderBy.asc("dept-id"), OrderBy.asc("rank"))
+                .build());
 
     assertThat(queryProto.getOrder(0)).isEqualTo(propertyOrder("dept-id"));
     assertThat(queryProto.getOrder(1)).isEqualTo(propertyOrder("rank"));
@@ -100,9 +100,8 @@ public void testOrderBy() {
 
   @Test
   public void testDistinctOn() {
-    Query queryProto = protoPreparer.prepare(newEntityQueryBuilder()
-        .setDistinctOn("dept-id", "rank")
-        .build());
+    Query queryProto =
+        protoPreparer.prepare(newEntityQueryBuilder().setDistinctOn("dept-id", "rank").build());
 
     assertThat(queryProto.getDistinctOn(0)).isEqualTo(propertyReference("dept-id"));
     assertThat(queryProto.getDistinctOn(1)).isEqualTo(propertyReference("rank"));
@@ -110,11 +109,10 @@ public void testDistinctOn() {
 
   @Test
   public void testProjections() {
-    Query queryProto = protoPreparer.prepare(newEntityQueryBuilder()
-        .setProjection("dept-id", "rank")
-        .build());
+    Query queryProto =
+        protoPreparer.prepare(newEntityQueryBuilder().setProjection("dept-id", "rank").build());
 
     assertThat(queryProto.getProjection(0)).isEqualTo(projection("dept-id"));
     assertThat(queryProto.getProjection(1)).isEqualTo(projection("rank"));
   }
-}
\ No newline at end of file
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/TestUtils.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/TestUtils.java
index 9c1ea8ffd..3a3fcfaea 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/TestUtils.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/TestUtils.java
@@ -15,7 +15,6 @@
  */
 package com.google.cloud.datastore;
 
-import com.google.datastore.v1.RunAggregationQueryRequest;
 import java.util.function.Predicate;
 import org.easymock.EasyMock;
 import org.easymock.IArgumentMatcher;
@@ -23,19 +22,18 @@
 public class TestUtils {
 
   public static  T matches(Predicate predicate) {
-    EasyMock.reportMatcher(new IArgumentMatcher() {
-      @Override
-      public boolean matches(Object argument) {
-        return predicate.test(((T) argument));
-      }
+    EasyMock.reportMatcher(
+        new IArgumentMatcher() {
+          @Override
+          public boolean matches(Object argument) {
+            return predicate.test(((T) argument));
+          }
 
-      @Override
-      public void appendTo(StringBuffer buffer) {
-        buffer.append("matches(\"").append(predicate).append("\")");
-      }
-    });
+          @Override
+          public void appendTo(StringBuffer buffer) {
+            buffer.append("matches(\"").append(predicate).append("\")");
+          }
+        });
     return null;
   }
-
-
 }
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java
index 859b3c522..94b16677a 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/aggregation/CountAggregationTest.java
@@ -36,9 +36,7 @@ public void testCountAggregationWithDefaultValues() {
 
   @Test
   public void testCountAggregationWithAlias() {
-    AggregationQuery.Aggregation countAggregationPb = count()
-        .as("column_1")
-        .build().toPb();
+    AggregationQuery.Aggregation countAggregationPb = count().as("column_1").build().toPb();
 
     assertThat(countAggregationPb.getCount().getUpTo().getValue()).isEqualTo(0L);
     assertThat(countAggregationPb.getAlias()).isEqualTo("column_1");
@@ -46,11 +44,9 @@ public void testCountAggregationWithAlias() {
 
   @Test
   public void testEquals() {
-    CountAggregation.Builder aggregation1 = count()
-        .as("total");
+    CountAggregation.Builder aggregation1 = count().as("total");
 
-    CountAggregation.Builder aggregation2 = count()
-        .as("total");
+    CountAggregation.Builder aggregation2 = count().as("total");
 
     assertEquals(aggregation1.build(), aggregation2.build());
     assertEquals(aggregation2.build(), aggregation1.build());
@@ -58,4 +54,4 @@ public void testEquals() {
     assertNotEquals(aggregation1.as("new-alias").build(), aggregation2.build());
     assertNotEquals(aggregation2.build(), aggregation1.as("new-alias").build());
   }
-}
\ No newline at end of file
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/EmulatorProxy.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/EmulatorProxy.java
index 726f3fe11..3c9502ef6 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/EmulatorProxy.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/EmulatorProxy.java
@@ -32,7 +32,8 @@ public EmulatorProxy(DatastoreOptions emulatorDataStoreOptions) {
   }
 
   private void init() {
-    this.mockWebServer.setDispatcher(new ProxyDispatcher(new HttpDatastoreRpc(emulatorDataStoreOptions)));
+    this.mockWebServer.setDispatcher(
+        new ProxyDispatcher(new HttpDatastoreRpc(emulatorDataStoreOptions)));
   }
 
   public void start() throws IOException {
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/ProxyDispatcher.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/ProxyDispatcher.java
index 2a8792e5d..eea3cd30a 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/ProxyDispatcher.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/ProxyDispatcher.java
@@ -52,19 +52,19 @@ public MockResponse dispatch(@NotNull RecordedRequest recordedRequest) {
     String methodName = recordedRequest.getPath().split(":")[1];
     if (!"runAggregationQuery".equals(methodName)) {
       throw new IllegalStateException(
-          String.format("Proxy only supports RunAggregationQuery method, Found %s method.",
-              methodName));
+          String.format(
+              "Proxy only supports RunAggregationQuery method, Found %s method.", methodName));
     }
 
     try {
-      RunAggregationQueryRequest runAggregationQueryRequest = RunAggregationQueryRequest.parseFrom(
-          recordedRequest.getBody().inputStream());
+      RunAggregationQueryRequest runAggregationQueryRequest =
+          RunAggregationQueryRequest.parseFrom(recordedRequest.getBody().inputStream());
 
       RunQueryRequest runQueryRequest = getRunQueryRequest(runAggregationQueryRequest);
       RunQueryResponse runQueryResponse = this.datastoreRpc.runQuery(runQueryRequest);
 
-      RunAggregationQueryResponse runAggregationQueryResponse = getRunAggregationQueryResponse(
-          runAggregationQueryRequest, runQueryResponse);
+      RunAggregationQueryResponse runAggregationQueryResponse =
+          getRunAggregationQueryResponse(runAggregationQueryRequest, runQueryResponse);
       Buffer buffer = new Buffer();
       runAggregationQueryResponse.writeTo(buffer.outputStream());
       return new MockResponse().setBody(buffer);
@@ -77,11 +77,13 @@ public MockResponse dispatch(@NotNull RecordedRequest recordedRequest) {
   @NotNull
   private RunAggregationQueryResponse getRunAggregationQueryResponse(
       RunAggregationQueryRequest runAggregationQueryRequest, RunQueryResponse runQueryResponse) {
-    AggregationResult aggregationResult = AggregationResult.newBuilder()
-        .putAllAggregateProperties(getProperties(
-            runAggregationQueryRequest.getAggregationQuery().getAggregationsList(),
-            runQueryResponse))
-        .build();
+    AggregationResult aggregationResult =
+        AggregationResult.newBuilder()
+            .putAllAggregateProperties(
+                getProperties(
+                    runAggregationQueryRequest.getAggregationQuery().getAggregationsList(),
+                    runQueryResponse))
+            .build();
     return RunAggregationQueryResponse.newBuilder()
         .setBatch(
             AggregationResultBatch.newBuilder().addAggregationResults(aggregationResult).build())
@@ -89,7 +91,8 @@ private RunAggregationQueryResponse getRunAggregationQueryResponse(
   }
 
   @NotNull
-  private RunQueryRequest getRunQueryRequest(RunAggregationQueryRequest runAggregationQueryRequest) {
+  private RunQueryRequest getRunQueryRequest(
+      RunAggregationQueryRequest runAggregationQueryRequest) {
     Query nestedQuery = runAggregationQueryRequest.getAggregationQuery().getNestedQuery();
     RunQueryRequest.Builder builder = RunQueryRequest.newBuilder().setQuery(nestedQuery);
     if (runAggregationQueryRequest.hasReadOptions()) {
@@ -98,8 +101,8 @@ private RunQueryRequest getRunQueryRequest(RunAggregationQueryRequest runAggrega
     return builder.build();
   }
 
-  private Map getProperties(List aggregationsList,
-      RunQueryResponse runQueryResponse) {
+  private Map getProperties(
+      List aggregationsList, RunQueryResponse runQueryResponse) {
     HashMap map = new HashMap<>();
     for (Aggregation aggregation : aggregationsList) {
       map.put(aggregation.getAlias(), getValue(runQueryResponse, aggregation.getCount()));
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java
index 8b6d6d2cf..4b7e86f67 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/AggregationQueryExecutorTest.java
@@ -61,55 +61,64 @@ public class AggregationQueryExecutorTest {
   @Before
   public void setUp() throws Exception {
     mockRpc = EasyMock.createStrictMock(DatastoreRpc.class);
-    datastoreOptions = DatastoreOptions.newBuilder().setProjectId("project-id")
-        .setNamespace(NAMESPACE).build();
+    datastoreOptions =
+        DatastoreOptions.newBuilder().setProjectId("project-id").setNamespace(NAMESPACE).build();
     queryExecutor = new AggregationQueryExecutor(mockRpc, datastoreOptions);
   }
 
   @Test
   public void shouldExecuteAggregationQuery() {
-    EntityQuery nestedQuery = Query.newEntityQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .setKind(KIND)
-        .setFilter(eq("done", true))
-        .build();
-
-    AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .addAggregation(count().as("total"))
-        .over(nestedQuery)
-        .build();
+    EntityQuery nestedQuery =
+        Query.newEntityQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setKind(KIND)
+            .setFilter(eq("done", true))
+            .build();
+
+    AggregationQuery aggregationQuery =
+        Query.newAggregationQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .addAggregation(count().as("total"))
+            .over(nestedQuery)
+            .build();
 
     RunAggregationQueryResponse runAggregationQueryResponse = dummyAggregationQueryResponse();
-    expect(mockRpc.runAggregationQuery(anyObject(RunAggregationQueryRequest.class))).andReturn(
-        runAggregationQueryResponse);
+    expect(mockRpc.runAggregationQuery(anyObject(RunAggregationQueryRequest.class)))
+        .andReturn(runAggregationQueryResponse);
 
     replay(mockRpc);
 
     AggregationResults aggregationResults = queryExecutor.execute(aggregationQuery);
 
     verify(mockRpc);
-    assertThat(aggregationResults).isEqualTo(new AggregationResults(asList(
-        new AggregationResult(
-            ImmutableMap.of("count", LongValue.of(209), "property_2", LongValue.of(100))),
-        new AggregationResult(
-            ImmutableMap.of("count", LongValue.of(509), "property_2", LongValue.of(100)))
-    ), Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime())));
+    assertThat(aggregationResults)
+        .isEqualTo(
+            new AggregationResults(
+                asList(
+                    new AggregationResult(
+                        ImmutableMap.of(
+                            "count", LongValue.of(209), "property_2", LongValue.of(100))),
+                    new AggregationResult(
+                        ImmutableMap.of(
+                            "count", LongValue.of(509), "property_2", LongValue.of(100)))),
+                Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime())));
   }
 
   @Test
   public void shouldExecuteAggregationQueryWithReadOptions() {
-    EntityQuery nestedQuery = Query.newEntityQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .setKind(KIND)
-        .setFilter(eq("done", true))
-        .build();
-
-    AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .addAggregation(count().as("total"))
-        .over(nestedQuery)
-        .build();
+    EntityQuery nestedQuery =
+        Query.newEntityQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setKind(KIND)
+            .setFilter(eq("done", true))
+            .build();
+
+    AggregationQuery aggregationQuery =
+        Query.newAggregationQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .addAggregation(count().as("total"))
+            .over(nestedQuery)
+            .build();
 
     RunAggregationQueryResponse runAggregationQueryResponse = dummyAggregationQueryResponse();
     expect(mockRpc.runAggregationQuery(matches(runAggregationRequestWithEventualConsistency())))
@@ -117,43 +126,52 @@ public void shouldExecuteAggregationQueryWithReadOptions() {
 
     replay(mockRpc);
 
-    AggregationResults aggregationResults = queryExecutor.execute(aggregationQuery,
-        eventualConsistency());
+    AggregationResults aggregationResults =
+        queryExecutor.execute(aggregationQuery, eventualConsistency());
 
     verify(mockRpc);
-    assertThat(aggregationResults).isEqualTo(new AggregationResults(asList(
-        new AggregationResult(
-            ImmutableMap.of("count", LongValue.of(209), "property_2", LongValue.of(100))),
-        new AggregationResult(
-            ImmutableMap.of("count", LongValue.of(509), "property_2", LongValue.of(100)))
-    ), Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime())));
+    assertThat(aggregationResults)
+        .isEqualTo(
+            new AggregationResults(
+                asList(
+                    new AggregationResult(
+                        ImmutableMap.of(
+                            "count", LongValue.of(209), "property_2", LongValue.of(100))),
+                    new AggregationResult(
+                        ImmutableMap.of(
+                            "count", LongValue.of(509), "property_2", LongValue.of(100)))),
+                Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime())));
   }
 
   private RunAggregationQueryResponse dummyAggregationQueryResponse() {
-    Map result1 = new HashMap<>(ImmutableMap.of(
-        "count", intValue(209),
-        "property_2", intValue(100)
-    ));
-
-    Map result2 = new HashMap<>(ImmutableMap.of(
-        "count", intValue(509),
-        "property_2", intValue(100)
-    ));
-
-    AggregationResultBatch resultBatch = AggregationResultBatch.newBuilder()
-        .addAggregationResults(com.google.datastore.v1.AggregationResult.newBuilder()
-            .putAllAggregateProperties(result1).build())
-        .addAggregationResults(com.google.datastore.v1.AggregationResult.newBuilder()
-            .putAllAggregateProperties(result2).build())
-        .build();
-    return RunAggregationQueryResponse.newBuilder()
-        .setBatch(resultBatch)
-        .build();
+    Map result1 =
+        new HashMap<>(
+            ImmutableMap.of(
+                "count", intValue(209),
+                "property_2", intValue(100)));
+
+    Map result2 =
+        new HashMap<>(
+            ImmutableMap.of(
+                "count", intValue(509),
+                "property_2", intValue(100)));
+
+    AggregationResultBatch resultBatch =
+        AggregationResultBatch.newBuilder()
+            .addAggregationResults(
+                com.google.datastore.v1.AggregationResult.newBuilder()
+                    .putAllAggregateProperties(result1)
+                    .build())
+            .addAggregationResults(
+                com.google.datastore.v1.AggregationResult.newBuilder()
+                    .putAllAggregateProperties(result2)
+                    .build())
+            .build();
+    return RunAggregationQueryResponse.newBuilder().setBatch(resultBatch).build();
   }
 
   private Predicate runAggregationRequestWithEventualConsistency() {
     return runAggregationQueryRequest ->
         runAggregationQueryRequest.getReadOptions().getReadConsistency() == EVENTUAL;
   }
-
-}
\ No newline at end of file
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java
index dec1615fe..6301ebeff 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java
@@ -40,8 +40,6 @@
 import com.google.cloud.datastore.ReadOption;
 import com.google.cloud.datastore.ReadOption.QueryAndReadOptions;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.truth.Truth;
-import com.google.datastore.v1.GqlQueryParameter;
 import com.google.datastore.v1.RunAggregationQueryRequest;
 import java.util.HashMap;
 import org.junit.Test;
@@ -51,138 +49,131 @@ public class AggregationQueryRequestProtoPreparerTest {
   private static final String KIND = "Task";
   private static final String NAMESPACE = "ns";
   private static final String PROJECT_ID = "project-id";
-  private static final DatastoreOptions DATASTORE_OPTIONS = DatastoreOptions.newBuilder()
-      .setProjectId(PROJECT_ID)
-      .setNamespace(NAMESPACE)
-      .build();
-  private static final EntityQuery COMPLETED_TASK_STRUCTURED_QUERY = Query.newEntityQueryBuilder()
-      .setNamespace(NAMESPACE).setKind(KIND).setFilter(eq("done", true)).build();
-
-  private static final GqlQuery COMPLETED_TASK_GQL_QUERY = Query.newGqlQueryBuilder(
-          "AGGREGATE COUNT AS total_characters OVER ("
-              + "SELECT * FROM Character WHERE name = @name and age > @1"
-              + ")")
-      .setBinding("name", "John Doe")
-      .addBinding(27)
-      .build();
-
-  private final AggregationQuery AGGREGATION_OVER_STRUCTURED_QUERY = Query.newAggregationQueryBuilder()
-      .setNamespace(NAMESPACE)
-      .addAggregation(count().as("total"))
-      .over(COMPLETED_TASK_STRUCTURED_QUERY)
-      .build();
-
-  private final AggregationQuery AGGREGATION_OVER_GQL_QUERY = Query.newAggregationQueryBuilder()
-      .setNamespace(NAMESPACE)
-      .over(COMPLETED_TASK_GQL_QUERY)
-      .build();
-
-
-  private final AggregationQueryRequestProtoPreparer protoPreparer = new AggregationQueryRequestProtoPreparer(
-      DATASTORE_OPTIONS);
+  private static final DatastoreOptions DATASTORE_OPTIONS =
+      DatastoreOptions.newBuilder().setProjectId(PROJECT_ID).setNamespace(NAMESPACE).build();
+  private static final EntityQuery COMPLETED_TASK_STRUCTURED_QUERY =
+      Query.newEntityQueryBuilder()
+          .setNamespace(NAMESPACE)
+          .setKind(KIND)
+          .setFilter(eq("done", true))
+          .build();
+
+  private static final GqlQuery COMPLETED_TASK_GQL_QUERY =
+      Query.newGqlQueryBuilder(
+              "AGGREGATE COUNT AS total_characters OVER ("
+                  + "SELECT * FROM Character WHERE name = @name and age > @1"
+                  + ")")
+          .setBinding("name", "John Doe")
+          .addBinding(27)
+          .build();
+
+  private final AggregationQuery AGGREGATION_OVER_STRUCTURED_QUERY =
+      Query.newAggregationQueryBuilder()
+          .setNamespace(NAMESPACE)
+          .addAggregation(count().as("total"))
+          .over(COMPLETED_TASK_STRUCTURED_QUERY)
+          .build();
+
+  private final AggregationQuery AGGREGATION_OVER_GQL_QUERY =
+      Query.newAggregationQueryBuilder()
+          .setNamespace(NAMESPACE)
+          .over(COMPLETED_TASK_GQL_QUERY)
+          .build();
+
+  private final AggregationQueryRequestProtoPreparer protoPreparer =
+      new AggregationQueryRequestProtoPreparer(DATASTORE_OPTIONS);
 
   @Test
   public void shouldPrepareAggregationQueryRequestWithGivenStructuredQuery() {
-    RunAggregationQueryRequest runAggregationQueryRequest = protoPreparer.prepare(
-        QueryAndReadOptions.create(AGGREGATION_OVER_STRUCTURED_QUERY));
+    RunAggregationQueryRequest runAggregationQueryRequest =
+        protoPreparer.prepare(QueryAndReadOptions.create(AGGREGATION_OVER_STRUCTURED_QUERY));
 
     assertThat(runAggregationQueryRequest.getProjectId()).isEqualTo(PROJECT_ID);
 
-    assertThat(runAggregationQueryRequest.getPartitionId().getProjectId())
-        .isEqualTo(PROJECT_ID);
-    assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId())
-        .isEqualTo(NAMESPACE);
+    assertThat(runAggregationQueryRequest.getPartitionId().getProjectId()).isEqualTo(PROJECT_ID);
+    assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId()).isEqualTo(NAMESPACE);
 
-    com.google.datastore.v1.AggregationQuery aggregationQueryProto = runAggregationQueryRequest.getAggregationQuery();
+    com.google.datastore.v1.AggregationQuery aggregationQueryProto =
+        runAggregationQueryRequest.getAggregationQuery();
     assertThat(aggregationQueryProto.getNestedQuery())
-        .isEqualTo(com.google.datastore.v1.Query.newBuilder()
-            .addKind(kind(KIND))
-            .setFilter(propertyFilter("done", EQUAL, booleanValue(true)))
-            .build());
+        .isEqualTo(
+            com.google.datastore.v1.Query.newBuilder()
+                .addKind(kind(KIND))
+                .setFilter(propertyFilter("done", EQUAL, booleanValue(true)))
+                .build());
     assertThat(aggregationQueryProto.getAggregationsList())
         .isEqualTo(singletonList(countAggregation("total")));
   }
 
   @Test
   public void shouldPrepareAggregationQueryRequestWithGivenGqlQuery() {
-    RunAggregationQueryRequest runAggregationQueryRequest = protoPreparer.prepare(
-        QueryAndReadOptions.create(
-            AGGREGATION_OVER_GQL_QUERY));
+    RunAggregationQueryRequest runAggregationQueryRequest =
+        protoPreparer.prepare(QueryAndReadOptions.create(AGGREGATION_OVER_GQL_QUERY));
 
     assertThat(runAggregationQueryRequest.getProjectId()).isEqualTo(PROJECT_ID);
 
-    assertThat(runAggregationQueryRequest.getPartitionId().getProjectId())
-        .isEqualTo(PROJECT_ID);
-    assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId())
-        .isEqualTo(NAMESPACE);
+    assertThat(runAggregationQueryRequest.getPartitionId().getProjectId()).isEqualTo(PROJECT_ID);
+    assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId()).isEqualTo(NAMESPACE);
 
     com.google.datastore.v1.GqlQuery gqlQueryProto = runAggregationQueryRequest.getGqlQuery();
 
-    assertThat(gqlQueryProto.getQueryString())
-        .isEqualTo(COMPLETED_TASK_GQL_QUERY.getQueryString());
+    assertThat(gqlQueryProto.getQueryString()).isEqualTo(COMPLETED_TASK_GQL_QUERY.getQueryString());
     assertThat(gqlQueryProto.getNamedBindingsMap())
-        .isEqualTo(new HashMap<>(
-            ImmutableMap.of("name", gqlQueryParameter(stringValue("John Doe")))
-        ));
-    assertThat(gqlQueryProto.getPositionalBindingsList()).isEqualTo(asList(
-        gqlQueryParameter(intValue(27))
-    ));
+        .isEqualTo(
+            new HashMap<>(ImmutableMap.of("name", gqlQueryParameter(stringValue("John Doe")))));
+    assertThat(gqlQueryProto.getPositionalBindingsList())
+        .isEqualTo(asList(gqlQueryParameter(intValue(27))));
   }
 
   @Test
   public void shouldPrepareReadOptionsWithGivenStructuredQuery() {
-    RunAggregationQueryRequest eventualConsistencyAggregationRequest = prepareQuery(
-        AGGREGATION_OVER_STRUCTURED_QUERY, eventualConsistency());
+    RunAggregationQueryRequest eventualConsistencyAggregationRequest =
+        prepareQuery(AGGREGATION_OVER_STRUCTURED_QUERY, eventualConsistency());
     assertThat(eventualConsistencyAggregationRequest.getReadOptions().getReadConsistency())
         .isEqualTo(EVENTUAL);
 
     Timestamp now = Timestamp.now();
-    RunAggregationQueryRequest readTimeAggregationRequest = prepareQuery(
-        AGGREGATION_OVER_STRUCTURED_QUERY, ReadOption.readTime(now));
+    RunAggregationQueryRequest readTimeAggregationRequest =
+        prepareQuery(AGGREGATION_OVER_STRUCTURED_QUERY, ReadOption.readTime(now));
     assertThat(Timestamp.fromProto(readTimeAggregationRequest.getReadOptions().getReadTime()))
         .isEqualTo(now);
   }
 
   @Test
   public void shouldPrepareReadOptionsWithGivenGqlQuery() {
-    RunAggregationQueryRequest eventualConsistencyAggregationRequest = prepareQuery(
-        AGGREGATION_OVER_GQL_QUERY, eventualConsistency());
+    RunAggregationQueryRequest eventualConsistencyAggregationRequest =
+        prepareQuery(AGGREGATION_OVER_GQL_QUERY, eventualConsistency());
     assertThat(eventualConsistencyAggregationRequest.getReadOptions().getReadConsistency())
         .isEqualTo(EVENTUAL);
 
     Timestamp now = Timestamp.now();
-    RunAggregationQueryRequest readTimeAggregationRequest = prepareQuery(
-        AGGREGATION_OVER_GQL_QUERY, ReadOption.readTime(now));
+    RunAggregationQueryRequest readTimeAggregationRequest =
+        prepareQuery(AGGREGATION_OVER_GQL_QUERY, ReadOption.readTime(now));
     assertThat(Timestamp.fromProto(readTimeAggregationRequest.getReadOptions().getReadTime()))
         .isEqualTo(now);
   }
 
   @Test
   public void shouldPrepareAggregationQueryWithoutNamespace() {
-    AggregationQuery structuredQueryWithoutNamespace = Query.newAggregationQueryBuilder()
-        .addAggregation(count().as("total"))
-        .over(COMPLETED_TASK_STRUCTURED_QUERY)
-        .build();
-    AggregationQuery gqlQueryWithoutNamespace = Query.newAggregationQueryBuilder()
-        .over(COMPLETED_TASK_GQL_QUERY)
-        .build();
-
-    RunAggregationQueryRequest runAggregationQueryFromStructuredQuery = protoPreparer.prepare(
-        QueryAndReadOptions.create(structuredQueryWithoutNamespace));
-    RunAggregationQueryRequest runAggregationQueryFromGqlQuery = protoPreparer.prepare(
-        QueryAndReadOptions.create(gqlQueryWithoutNamespace));
+    AggregationQuery structuredQueryWithoutNamespace =
+        Query.newAggregationQueryBuilder()
+            .addAggregation(count().as("total"))
+            .over(COMPLETED_TASK_STRUCTURED_QUERY)
+            .build();
+    AggregationQuery gqlQueryWithoutNamespace =
+        Query.newAggregationQueryBuilder().over(COMPLETED_TASK_GQL_QUERY).build();
+
+    RunAggregationQueryRequest runAggregationQueryFromStructuredQuery =
+        protoPreparer.prepare(QueryAndReadOptions.create(structuredQueryWithoutNamespace));
+    RunAggregationQueryRequest runAggregationQueryFromGqlQuery =
+        protoPreparer.prepare(QueryAndReadOptions.create(gqlQueryWithoutNamespace));
 
     assertThat(runAggregationQueryFromStructuredQuery.getPartitionId().getNamespaceId())
         .isEqualTo("");
-    assertThat(runAggregationQueryFromGqlQuery.getPartitionId().getNamespaceId())
-        .isEqualTo("");
-
+    assertThat(runAggregationQueryFromGqlQuery.getPartitionId().getNamespaceId()).isEqualTo("");
   }
 
   private RunAggregationQueryRequest prepareQuery(AggregationQuery query, ReadOption readOption) {
-    return protoPreparer.prepare(
-        QueryAndReadOptions.create(query,
-            singletonList(readOption)));
+    return protoPreparer.prepare(QueryAndReadOptions.create(query, singletonList(readOption)));
   }
-
-}
\ No newline at end of file
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java
index e15fc0def..8776d4221 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/response/AggregationQueryResponseTransformerTest.java
@@ -36,51 +36,56 @@
 
 public class AggregationQueryResponseTransformerTest {
 
-  private final AggregationQueryResponseTransformer responseTransformer = new AggregationQueryResponseTransformer();
+  private final AggregationQueryResponseTransformer responseTransformer =
+      new AggregationQueryResponseTransformer();
 
   @Test
   public void shouldTransformAggregationQueryResponse() {
-    Map result1 = new HashMap<>(ImmutableMap.of(
-        "count", intValue(209),
-        "property_2", intValue(100)
-    ));
+    Map result1 =
+        new HashMap<>(
+            ImmutableMap.of(
+                "count", intValue(209),
+                "property_2", intValue(100)));
 
-    Map result2 = new HashMap<>(ImmutableMap.of(
-        "count", intValue(509),
-        "property_2", intValue(100)
-    ));
+    Map result2 =
+        new HashMap<>(
+            ImmutableMap.of(
+                "count", intValue(509),
+                "property_2", intValue(100)));
     Timestamp readTime = Timestamp.now();
 
-    AggregationResultBatch resultBatch = AggregationResultBatch.newBuilder()
-        .addAggregationResults(com.google.datastore.v1.AggregationResult.newBuilder()
-            .putAllAggregateProperties(result1).build())
-        .addAggregationResults(com.google.datastore.v1.AggregationResult.newBuilder()
-            .putAllAggregateProperties(result2).build())
-        .setReadTime(readTime.toProto())
-        .build();
-    RunAggregationQueryResponse runAggregationQueryResponse = RunAggregationQueryResponse.newBuilder()
-        .setBatch(resultBatch)
-        .build();
+    AggregationResultBatch resultBatch =
+        AggregationResultBatch.newBuilder()
+            .addAggregationResults(
+                com.google.datastore.v1.AggregationResult.newBuilder()
+                    .putAllAggregateProperties(result1)
+                    .build())
+            .addAggregationResults(
+                com.google.datastore.v1.AggregationResult.newBuilder()
+                    .putAllAggregateProperties(result2)
+                    .build())
+            .setReadTime(readTime.toProto())
+            .build();
+    RunAggregationQueryResponse runAggregationQueryResponse =
+        RunAggregationQueryResponse.newBuilder().setBatch(resultBatch).build();
 
-    AggregationResults aggregationResults = responseTransformer.transform(
-        runAggregationQueryResponse);
+    AggregationResults aggregationResults =
+        responseTransformer.transform(runAggregationQueryResponse);
 
     assertThat(aggregationResults.size()).isEqualTo(2);
-    assertThat(aggregationResults.get(0))
-        .isEqualTo(new AggregationResult(toDomainValues(result1)));
-    assertThat(aggregationResults.get(1))
-        .isEqualTo(new AggregationResult(toDomainValues(result2)));
+    assertThat(aggregationResults.get(0)).isEqualTo(new AggregationResult(toDomainValues(result1)));
+    assertThat(aggregationResults.get(1)).isEqualTo(new AggregationResult(toDomainValues(result2)));
     assertThat(aggregationResults.getReadTime()).isEqualTo(readTime);
-
   }
 
   private Map toDomainValues(Map map) {
 
-    return map
-        .entrySet()
-        .stream()
-        .map((Function, Entry>) entry ->
-            new SimpleEntry<>(entry.getKey(), (LongValue) LongValue.fromPb(entry.getValue())))
+    return map.entrySet().stream()
+        .map(
+            (Function, Entry>)
+                entry ->
+                    new SimpleEntry<>(
+                        entry.getKey(), (LongValue) LongValue.fromPb(entry.getValue())))
         .collect(Collectors.toMap(Entry::getKey, Entry::getValue));
   }
-}
\ No newline at end of file
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
index 142f94954..a505b456e 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
@@ -168,11 +168,9 @@ public class ITDatastoreTest {
           .set("partial2", ENTITY2)
           .build();
 
-  @Rule
-  public Timeout globalTimeout = Timeout.seconds(100);
+  @Rule public Timeout globalTimeout = Timeout.seconds(100);
 
-  @Rule
-  public MultipleAttemptsRule multipleAttemptsRule = new MultipleAttemptsRule(3);
+  @Rule public MultipleAttemptsRule multipleAttemptsRule = new MultipleAttemptsRule(3);
 
   @AfterClass
   public static void afterClass() {
@@ -525,44 +523,40 @@ public void testRunGqlQueryWithCasting() throws InterruptedException {
   @Test
   public void testRunAggregationQuery() {
     // verifying aggregation with an entity query
-    testCountAggregationWith(builder ->
-        builder
-            .addAggregation(count().as("total_count"))
-            .over(Query.newEntityQueryBuilder()
-                .setNamespace(NAMESPACE)
-                .setKind(KIND1)
-                .build()
-            ));
+    testCountAggregationWith(
+        builder ->
+            builder
+                .addAggregation(count().as("total_count"))
+                .over(
+                    Query.newEntityQueryBuilder().setNamespace(NAMESPACE).setKind(KIND1).build()));
 
     // verifying aggregation with a projection query
-    testCountAggregationWith(builder ->
-        builder
-            .addAggregation(count().as("total_count"))
-            .over(Query.newProjectionEntityQueryBuilder()
-                .setProjection("str")
-                .setNamespace(NAMESPACE)
-                .setKind(KIND1)
-                .build()
-            ));
+    testCountAggregationWith(
+        builder ->
+            builder
+                .addAggregation(count().as("total_count"))
+                .over(
+                    Query.newProjectionEntityQueryBuilder()
+                        .setProjection("str")
+                        .setNamespace(NAMESPACE)
+                        .setKind(KIND1)
+                        .build()));
 
     // verifying aggregation with a key query
-    testCountAggregationWith(builder ->
-        builder
-            .addAggregation(count().as("total_count"))
-            .over(Query.newKeyQueryBuilder()
-                .setNamespace(NAMESPACE)
-                .setKind(KIND1)
-                .build()
-            ));
+    testCountAggregationWith(
+        builder ->
+            builder
+                .addAggregation(count().as("total_count"))
+                .over(Query.newKeyQueryBuilder().setNamespace(NAMESPACE).setKind(KIND1).build()));
 
     // verifying aggregation with a GQL query
-    testCountAggregationWith(builder ->
-        builder
-            .over(Query.newGqlQueryBuilder(
-                    "AGGREGATE COUNT(*) AS total_count OVER (SELECT * FROM kind1)")
-                .setNamespace(NAMESPACE)
-                .build()
-            ));
+    testCountAggregationWith(
+        builder ->
+            builder.over(
+                Query.newGqlQueryBuilder(
+                        "AGGREGATE COUNT(*) AS total_count OVER (SELECT * FROM kind1)")
+                    .setNamespace(NAMESPACE)
+                    .build()));
   }
 
   /**
@@ -570,60 +564,70 @@ public void testRunAggregationQuery() {
    * original version of the entity as of the beginning of the transaction, or nothing if the entity
    * did not exist then.
    *
-   * @see 
-   * Source
+   * @see 
+   *     Source
    */
   @Test
   public void testRunAggregationQueryInTransactionShouldReturnAConsistentSnapshot() {
-    EntityQuery entityQuery = Query.newEntityQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .setFilter(PropertyFilter.hasAncestor(KEY1))
-        .build();
+    EntityQuery entityQuery =
+        Query.newEntityQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setFilter(PropertyFilter.hasAncestor(KEY1))
+            .build();
 
-    AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .over(entityQuery)
-        .addAggregation(count().as("count"))
-        .build();
+    AggregationQuery aggregationQuery =
+        Query.newAggregationQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .over(entityQuery)
+            .addAggregation(count().as("count"))
+            .build();
 
     // original entity count is 2
     assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
         .isEqualTo(2L);
 
     // FIRST TRANSACTION
-    DATASTORE.runInTransaction((TransactionCallable) inFirstTransaction -> {
-      // creating a new entity
-      Entity aNewEntity = Entity.newBuilder(ENTITY2)
-          .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build())
-          .set("v_int", 10)
-          .build();
-      inFirstTransaction.put(aNewEntity);
-
-      // count remains 2
-      assertThat(
-          getOnlyElement(inFirstTransaction.runAggregation(aggregationQuery)).get("count"))
-          .isEqualTo(2L);
-      assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
-          .isEqualTo(2L);
-      return null;
-    });
+    DATASTORE.runInTransaction(
+        (TransactionCallable)
+            inFirstTransaction -> {
+              // creating a new entity
+              Entity aNewEntity =
+                  Entity.newBuilder(ENTITY2)
+                      .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build())
+                      .set("v_int", 10)
+                      .build();
+              inFirstTransaction.put(aNewEntity);
+
+              // count remains 2
+              assertThat(
+                      getOnlyElement(inFirstTransaction.runAggregation(aggregationQuery))
+                          .get("count"))
+                  .isEqualTo(2L);
+              assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+                  .isEqualTo(2L);
+              return null;
+            });
     // after first transaction is committed, count is updated to 3 now.
     assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
         .isEqualTo(3L);
 
     // SECOND TRANSACTION
-    DATASTORE.runInTransaction((TransactionCallable) inSecondTransaction -> {
-      // deleting ENTITY2
-      inSecondTransaction.delete(ENTITY2.getKey());
-
-      // count remains 3
-      assertThat(
-          getOnlyElement(inSecondTransaction.runAggregation(aggregationQuery)).get("count"))
-          .isEqualTo(3L);
-      assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
-          .isEqualTo(3L);
-      return null;
-    });
+    DATASTORE.runInTransaction(
+        (TransactionCallable)
+            inSecondTransaction -> {
+              // deleting ENTITY2
+              inSecondTransaction.delete(ENTITY2.getKey());
+
+              // count remains 3
+              assertThat(
+                      getOnlyElement(inSecondTransaction.runAggregation(aggregationQuery))
+                          .get("count"))
+                  .isEqualTo(3L);
+              assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
+                  .isEqualTo(3L);
+              return null;
+            });
     // after second transaction is committed, count is updated to 2 now.
     assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
         .isEqualTo(2L);
@@ -632,45 +636,50 @@ public void testRunAggregationQueryInTransactionShouldReturnAConsistentSnapshot(
   /**
    * Data read or modified by a transaction cannot be concurrently modified.
    *
-   * @see 
-   * Source
+   * @see 
+   *     Source
    */
   @Test
   public void testRunAggregationQueryInAReadWriteTransactionShouldLockTheCountedDocuments()
       throws Exception {
     ExecutorService executor = Executors.newSingleThreadExecutor();
-    EntityQuery entityQuery = Query.newEntityQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .setFilter(PropertyFilter.hasAncestor(KEY1))
-        .build();
-    AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .over(entityQuery)
-        .addAggregation(count().as("count"))
-        .build();
+    EntityQuery entityQuery =
+        Query.newEntityQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setFilter(PropertyFilter.hasAncestor(KEY1))
+            .build();
+    AggregationQuery aggregationQuery =
+        Query.newAggregationQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .over(entityQuery)
+            .addAggregation(count().as("count"))
+            .build();
 
     // read-write transaction
     Transaction readWriteTransaction = DATASTORE.newTransaction();
 
     // acquiring lock by executing query in transaction
-    assertThat(
-        getOnlyElement(readWriteTransaction.runAggregation(aggregationQuery)).get("count")
-    ).isEqualTo(2L);
+    assertThat(getOnlyElement(readWriteTransaction.runAggregation(aggregationQuery)).get("count"))
+        .isEqualTo(2L);
 
     // Waiting task will be blocked by ongoing transactions.
-    Future addNewEntityTaskOutsideTransaction = executor.submit(() -> {
-      Entity aNewEntity = Entity.newBuilder(ENTITY2)
-          .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build())
-          .set("v_int", 10)
-          .build();
-      DATASTORE.put(aNewEntity);
-      return null;
-    });
+    Future addNewEntityTaskOutsideTransaction =
+        executor.submit(
+            () -> {
+              Entity aNewEntity =
+                  Entity.newBuilder(ENTITY2)
+                      .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build())
+                      .set("v_int", 10)
+                      .build();
+              DATASTORE.put(aNewEntity);
+              return null;
+            });
 
     // should throw Timeout exception as we haven't yet committed the transaction
     assertThrows(TimeoutException.class, () -> addNewEntityTaskOutsideTransaction.get(3, SECONDS));
 
-    //cleanup
+    // cleanup
     readWriteTransaction.commit();
     addNewEntityTaskOutsideTransaction.cancel(true);
     executor.shutdownNow();
@@ -680,39 +689,43 @@ public void testRunAggregationQueryInAReadWriteTransactionShouldLockTheCountedDo
   public void testRunAggregationQueryInAReadOnlyTransactionShouldNotLockTheCountedDocuments()
       throws Exception {
     ExecutorService executor = Executors.newSingleThreadExecutor();
-    EntityQuery entityQuery = Query.newEntityQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .setFilter(PropertyFilter.hasAncestor(KEY1))
-        .build();
-    AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
-        .setNamespace(NAMESPACE)
-        .over(entityQuery)
-        .addAggregation(count().as("count"))
-        .build();
-
-    TransactionOptions transactionOptions = TransactionOptions.newBuilder().setReadOnly(
-        ReadOnly.newBuilder().build()).build();
+    EntityQuery entityQuery =
+        Query.newEntityQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .setFilter(PropertyFilter.hasAncestor(KEY1))
+            .build();
+    AggregationQuery aggregationQuery =
+        Query.newAggregationQueryBuilder()
+            .setNamespace(NAMESPACE)
+            .over(entityQuery)
+            .addAggregation(count().as("count"))
+            .build();
+
+    TransactionOptions transactionOptions =
+        TransactionOptions.newBuilder().setReadOnly(ReadOnly.newBuilder().build()).build();
     Transaction readOnlyTransaction = DATASTORE.newTransaction(transactionOptions);
 
     // Executing query in transaction
-    assertThat(
-        getOnlyElement(readOnlyTransaction.runAggregation(aggregationQuery)).get("count"))
+    assertThat(getOnlyElement(readOnlyTransaction.runAggregation(aggregationQuery)).get("count"))
         .isEqualTo(2L);
 
     // Concurrent write task.
-    Future addNewEntityTaskOutsideTransaction = executor.submit(() -> {
-      Entity aNewEntity = Entity.newBuilder(ENTITY2)
-          .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build())
-          .set("v_int", 10)
-          .build();
-      DATASTORE.put(aNewEntity);
-      return null;
-    });
+    Future addNewEntityTaskOutsideTransaction =
+        executor.submit(
+            () -> {
+              Entity aNewEntity =
+                  Entity.newBuilder(ENTITY2)
+                      .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build())
+                      .set("v_int", 10)
+                      .build();
+              DATASTORE.put(aNewEntity);
+              return null;
+            });
 
     // should not throw exception and complete successfully as the ongoing transaction is read-only.
     addNewEntityTaskOutsideTransaction.get();
 
-    //cleanup
+    // cleanup
     readOnlyTransaction.commit();
     executor.shutdownNow();
 
@@ -725,41 +738,39 @@ public void testRunAggregationQueryWithReadTime() throws InterruptedException {
     String alias = "total_count";
 
     // verifying aggregation readTime with an entity query
-    testCountAggregationReadTimeWith(builder ->
-        builder
-            .over(
-                Query.newEntityQueryBuilder().setKind("new_kind").build())
-            .addAggregation(count().as(alias))
-    );
+    testCountAggregationReadTimeWith(
+        builder ->
+            builder
+                .over(Query.newEntityQueryBuilder().setKind("new_kind").build())
+                .addAggregation(count().as(alias)));
 
     // verifying aggregation readTime with a projection query
-    testCountAggregationReadTimeWith(builder ->
-        builder
-            .over(
-                Query.newProjectionEntityQueryBuilder()
-                    .setProjection("name")
-                    .setKind("new_kind").build())
-            .addAggregation(count().as(alias))
-    );
+    testCountAggregationReadTimeWith(
+        builder ->
+            builder
+                .over(
+                    Query.newProjectionEntityQueryBuilder()
+                        .setProjection("name")
+                        .setKind("new_kind")
+                        .build())
+                .addAggregation(count().as(alias)));
 
     // verifying aggregation readTime with a key query
-    testCountAggregationReadTimeWith(builder ->
-        builder
-            .over(
-                Query.newKeyQueryBuilder()
-                    .setKind("new_kind").build())
-            .addAggregation(count().as(alias))
-    );
+    testCountAggregationReadTimeWith(
+        builder ->
+            builder
+                .over(Query.newKeyQueryBuilder().setKind("new_kind").build())
+                .addAggregation(count().as(alias)));
 
     // verifying aggregation readTime with a GQL query
-    testCountAggregationReadTimeWith(builder ->
-        builder
-            .over(
-                Query.newGqlQueryBuilder(
-                        "AGGREGATE COUNT(*) AS total_count OVER (SELECT * FROM new_kind)")
-                    .build())
-            .addAggregation(count().as(alias))
-    );
+    testCountAggregationReadTimeWith(
+        builder ->
+            builder
+                .over(
+                    Query.newGqlQueryBuilder(
+                            "AGGREGATE COUNT(*) AS total_count OVER (SELECT * FROM new_kind)")
+                        .build())
+                .addAggregation(count().as(alias)));
   }
 
   @Test
@@ -1333,12 +1344,13 @@ private void testCountAggregationWith(Consumer configu
     Long countBeforeAdd = getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get(alias);
     long expectedCount = countBeforeAdd + 1;
 
-    Entity newEntity = Entity.newBuilder(ENTITY1)
-        .setKey(Key.newBuilder(KEY3, KIND1, 1).build())
-        .set("null", NULL_VALUE)
-        .set("partial1", PARTIAL_ENTITY2)
-        .set("partial2", ENTITY2)
-        .build();
+    Entity newEntity =
+        Entity.newBuilder(ENTITY1)
+            .setKey(Key.newBuilder(KEY3, KIND1, 1).build())
+            .set("null", NULL_VALUE)
+            .set("partial1", PARTIAL_ENTITY2)
+            .set("partial2", ENTITY2)
+            .build();
     DATASTORE.put(newEntity);
 
     Long countAfterAdd = getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get(alias);
@@ -1376,13 +1388,13 @@ private void testCountAggregationReadTimeWith(Consumer
       configurer.accept(builder);
       AggregationQuery countAggregationQuery = builder.build();
 
-      Long latestCount = getOnlyElement(DATASTORE.runAggregation(countAggregationQuery))
-          .get("total_count");
+      Long latestCount =
+          getOnlyElement(DATASTORE.runAggregation(countAggregationQuery)).get("total_count");
       assertThat(latestCount).isEqualTo(3L);
 
-      Long oldCount = getOnlyElement(
-          DATASTORE.runAggregation(countAggregationQuery, ReadOption.readTime(now))
-      ).get("total_count");
+      Long oldCount =
+          getOnlyElement(DATASTORE.runAggregation(countAggregationQuery, ReadOption.readTime(now)))
+              .get("total_count");
       assertThat(oldCount).isEqualTo(2L);
     } finally {
       DATASTORE.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());

From 9ff91d9e7244af327522ec82c6f4f41a6ae237eb Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Fri, 7 Oct 2022 12:01:00 +0530
Subject: [PATCH 52/82] Getting rid off emulator proxy and using easy mock to
 check the aggregationQuery triggered

---
 google-cloud-datastore/pom.xml                |   6 -
 .../google/cloud/datastore/DatastoreTest.java |  55 ++++++--
 .../datastore/emulator/EmulatorProxy.java     |  52 --------
 .../datastore/emulator/ProxyDispatcher.java   | 121 ------------------
 4 files changed, 43 insertions(+), 191 deletions(-)
 delete mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/EmulatorProxy.java
 delete mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/ProxyDispatcher.java

diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml
index 46e4e5f06..27ca86ab5 100644
--- a/google-cloud-datastore/pom.xml
+++ b/google-cloud-datastore/pom.xml
@@ -143,12 +143,6 @@
       easymock
       test
     
-    
-      com.squareup.okhttp3
-      mockwebserver
-      4.10.0
-      test
-    
     
       com.google.truth
       truth
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
index 81e66da55..cfa5922fc 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
@@ -16,6 +16,8 @@
 
 package com.google.cloud.datastore;
 
+import static com.google.cloud.datastore.ProtoTestData.intValue;
+import static com.google.cloud.datastore.TestUtils.matches;
 import static com.google.cloud.datastore.aggregation.Aggregation.count;
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.truth.Truth.assertThat;
@@ -36,13 +38,14 @@
 import com.google.cloud.datastore.Query.ResultType;
 import com.google.cloud.datastore.StructuredQuery.OrderBy;
 import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
-import com.google.cloud.datastore.emulator.EmulatorProxy;
 import com.google.cloud.datastore.spi.DatastoreRpcFactory;
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
 import com.google.cloud.datastore.testing.LocalDatastoreHelper;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterators;
 import com.google.common.collect.Lists;
+import com.google.datastore.v1.AggregationResultBatch;
 import com.google.datastore.v1.BeginTransactionRequest;
 import com.google.datastore.v1.BeginTransactionResponse;
 import com.google.datastore.v1.CommitRequest;
@@ -58,6 +61,8 @@
 import com.google.datastore.v1.ReserveIdsResponse;
 import com.google.datastore.v1.RollbackRequest;
 import com.google.datastore.v1.RollbackResponse;
+import com.google.datastore.v1.RunAggregationQueryRequest;
+import com.google.datastore.v1.RunAggregationQueryResponse;
 import com.google.datastore.v1.RunQueryRequest;
 import com.google.datastore.v1.RunQueryResponse;
 import com.google.datastore.v1.TransactionOptions;
@@ -66,11 +71,14 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeoutException;
+import java.util.function.Predicate;
 import org.easymock.EasyMock;
 import org.junit.AfterClass;
 import org.junit.Assert;
@@ -87,8 +95,6 @@ public class DatastoreTest {
   private static LocalDatastoreHelper helper = LocalDatastoreHelper.create(1.0);
   private static final DatastoreOptions options = helper.getOptions();
   private static final Datastore datastore = options.getService();
-  private static final EmulatorProxy emulatorProxy = new EmulatorProxy(options);
-  private static final Datastore datastoreEmulatorProxy = emulatorProxy.getOptions().getService();
   private static final String PROJECT_ID = options.getProjectId();
   private static final String KIND1 = "kind1";
   private static final String KIND2 = "kind2";
@@ -186,7 +192,6 @@ public void setUp() {
   @AfterClass
   public static void afterClass() throws IOException, InterruptedException, TimeoutException {
     helper.stop(Duration.ofMinutes(1));
-    emulatorProxy.stop();
   }
 
   @Test
@@ -535,21 +540,23 @@ public void testGqlQueryPagination() throws DatastoreException {
 
   @Test
   public void testRunAggregationQuery() {
+    RunAggregationQueryResponse aggregationQueryResponse = dummyAggregationQueryResponse();
+    EasyMock.expect(rpcMock.runAggregationQuery(matches(aggregationQueryWithAlias("total_count"))))
+        .andReturn(aggregationQueryResponse);
+    EasyMock.replay(rpcFactoryMock, rpcMock);
+
+    Datastore mockDatastore = rpcMockOptions.getService();
+
     EntityQuery selectAllQuery = Query.newEntityQueryBuilder().build();
     AggregationQuery getCountQuery =
         Query.newAggregationQueryBuilder()
             .addAggregation(count().as("total_count"))
             .over(selectAllQuery)
             .build();
-    AggregationResult resultBeforeInsert =
-        getOnlyElement(datastoreEmulatorProxy.runAggregation(getCountQuery));
-    assertThat(resultBeforeInsert.get("total_count")).isEqualTo(2L);
+    AggregationResult result = getOnlyElement(mockDatastore.runAggregation(getCountQuery));
 
-    datastore.put(ENTITY3);
-
-    AggregationResult resultAfterInsert =
-        getOnlyElement(datastoreEmulatorProxy.runAggregation(getCountQuery));
-    assertThat(resultAfterInsert.get("total_count")).isEqualTo(3L);
+    assertThat(result.get("total_count")).isEqualTo(209L);
+    EasyMock.verify(rpcFactoryMock, rpcMock);
   }
 
   @Test
@@ -1337,4 +1344,28 @@ public void testQueryWithStartCursor() {
     assertEquals(cursor2, cursor1);
     datastore.delete(entity1.getKey(), entity2.getKey(), entity3.getKey());
   }
+
+  private RunAggregationQueryResponse dummyAggregationQueryResponse() {
+    Map result1 =
+        new HashMap<>(ImmutableMap.of("total_count", intValue(209)));
+
+    AggregationResultBatch resultBatch =
+        AggregationResultBatch.newBuilder()
+            .addAggregationResults(
+                com.google.datastore.v1.AggregationResult.newBuilder()
+                    .putAllAggregateProperties(result1)
+                    .build())
+            .build();
+    return RunAggregationQueryResponse.newBuilder().setBatch(resultBatch).build();
+  }
+
+  private Predicate aggregationQueryWithAlias(String alias) {
+    return runAggregationQueryRequest ->
+        alias.equals(
+            runAggregationQueryRequest
+                .getAggregationQuery()
+                .getAggregationsList()
+                .get(0)
+                .getAlias());
+  }
 }
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/EmulatorProxy.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/EmulatorProxy.java
deleted file mode 100644
index 3c9502ef6..000000000
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/EmulatorProxy.java
+++ /dev/null
@@ -1,52 +0,0 @@
-/*
- * Copyright 2022 Google LLC
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     https://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.google.cloud.datastore.emulator;
-
-import com.google.cloud.datastore.DatastoreOptions;
-import com.google.cloud.datastore.spi.v1.HttpDatastoreRpc;
-import java.io.IOException;
-import okhttp3.mockwebserver.MockWebServer;
-
-public class EmulatorProxy {
-
-  private final MockWebServer mockWebServer;
-  private final DatastoreOptions emulatorDataStoreOptions;
-
-  public EmulatorProxy(DatastoreOptions emulatorDataStoreOptions) {
-    this.emulatorDataStoreOptions = emulatorDataStoreOptions;
-    this.mockWebServer = new MockWebServer();
-    init();
-  }
-
-  private void init() {
-    this.mockWebServer.setDispatcher(
-        new ProxyDispatcher(new HttpDatastoreRpc(emulatorDataStoreOptions)));
-  }
-
-  public void start() throws IOException {
-    this.mockWebServer.start();
-  }
-
-  public void stop() throws IOException {
-    this.mockWebServer.shutdown();
-  }
-
-  public DatastoreOptions getOptions() {
-    DatastoreOptions.Builder builder = this.emulatorDataStoreOptions.toBuilder();
-    builder.setHost(this.mockWebServer.getHostName() + ":" + this.mockWebServer.getPort());
-    return builder.build();
-  }
-}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/ProxyDispatcher.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/ProxyDispatcher.java
deleted file mode 100644
index eea3cd30a..000000000
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/emulator/ProxyDispatcher.java
+++ /dev/null
@@ -1,121 +0,0 @@
-/*
- * Copyright 2022 Google LLC
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     https://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.google.cloud.datastore.emulator;
-
-import static com.google.cloud.datastore.ProtoTestData.intValue;
-
-import com.google.cloud.datastore.spi.v1.DatastoreRpc;
-import com.google.datastore.v1.AggregationQuery.Aggregation;
-import com.google.datastore.v1.AggregationQuery.Aggregation.Count;
-import com.google.datastore.v1.AggregationResult;
-import com.google.datastore.v1.AggregationResultBatch;
-import com.google.datastore.v1.Query;
-import com.google.datastore.v1.RunAggregationQueryRequest;
-import com.google.datastore.v1.RunAggregationQueryResponse;
-import com.google.datastore.v1.RunQueryRequest;
-import com.google.datastore.v1.RunQueryResponse;
-import com.google.datastore.v1.Value;
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import okhttp3.mockwebserver.Dispatcher;
-import okhttp3.mockwebserver.MockResponse;
-import okhttp3.mockwebserver.RecordedRequest;
-import okio.Buffer;
-import org.jetbrains.annotations.NotNull;
-
-public class ProxyDispatcher extends Dispatcher {
-
-  private final DatastoreRpc datastoreRpc;
-
-  public ProxyDispatcher(DatastoreRpc datastoreRpc) {
-    this.datastoreRpc = datastoreRpc;
-  }
-
-  @NotNull
-  @Override
-  public MockResponse dispatch(@NotNull RecordedRequest recordedRequest) {
-    String methodName = recordedRequest.getPath().split(":")[1];
-    if (!"runAggregationQuery".equals(methodName)) {
-      throw new IllegalStateException(
-          String.format(
-              "Proxy only supports RunAggregationQuery method, Found %s method.", methodName));
-    }
-
-    try {
-      RunAggregationQueryRequest runAggregationQueryRequest =
-          RunAggregationQueryRequest.parseFrom(recordedRequest.getBody().inputStream());
-
-      RunQueryRequest runQueryRequest = getRunQueryRequest(runAggregationQueryRequest);
-      RunQueryResponse runQueryResponse = this.datastoreRpc.runQuery(runQueryRequest);
-
-      RunAggregationQueryResponse runAggregationQueryResponse =
-          getRunAggregationQueryResponse(runAggregationQueryRequest, runQueryResponse);
-      Buffer buffer = new Buffer();
-      runAggregationQueryResponse.writeTo(buffer.outputStream());
-      return new MockResponse().setBody(buffer);
-    } catch (IOException e) {
-      e.printStackTrace();
-    }
-    return new MockResponse();
-  }
-
-  @NotNull
-  private RunAggregationQueryResponse getRunAggregationQueryResponse(
-      RunAggregationQueryRequest runAggregationQueryRequest, RunQueryResponse runQueryResponse) {
-    AggregationResult aggregationResult =
-        AggregationResult.newBuilder()
-            .putAllAggregateProperties(
-                getProperties(
-                    runAggregationQueryRequest.getAggregationQuery().getAggregationsList(),
-                    runQueryResponse))
-            .build();
-    return RunAggregationQueryResponse.newBuilder()
-        .setBatch(
-            AggregationResultBatch.newBuilder().addAggregationResults(aggregationResult).build())
-        .build();
-  }
-
-  @NotNull
-  private RunQueryRequest getRunQueryRequest(
-      RunAggregationQueryRequest runAggregationQueryRequest) {
-    Query nestedQuery = runAggregationQueryRequest.getAggregationQuery().getNestedQuery();
-    RunQueryRequest.Builder builder = RunQueryRequest.newBuilder().setQuery(nestedQuery);
-    if (runAggregationQueryRequest.hasReadOptions()) {
-      builder.setReadOptions(runAggregationQueryRequest.getReadOptions());
-    }
-    return builder.build();
-  }
-
-  private Map getProperties(
-      List aggregationsList, RunQueryResponse runQueryResponse) {
-    HashMap map = new HashMap<>();
-    for (Aggregation aggregation : aggregationsList) {
-      map.put(aggregation.getAlias(), getValue(runQueryResponse, aggregation.getCount()));
-    }
-    return map;
-  }
-
-  private Value getValue(RunQueryResponse runQueryResponse, Count count) {
-    int totalEntityCount = runQueryResponse.getBatch().getEntityResultsCount();
-    if (count.hasUpTo()) {
-      return intValue(Math.min(count.getUpTo().getValue(), totalEntityCount));
-    } else {
-      return intValue(totalEntityCount);
-    }
-  }
-}

From 6953efe9d5c9944d7ea24f3c6aaf1013a12ad646 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Fri, 7 Oct 2022 12:31:14 +0530
Subject: [PATCH 53/82] Deleting a entity created locally in other test which
 is causing failure in other test

---
 .../java/com/google/cloud/datastore/it/ITDatastoreTest.java   | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
index a505b456e..0bd8b8df4 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
@@ -570,6 +570,7 @@ public void testRunAggregationQuery() {
    */
   @Test
   public void testRunAggregationQueryInTransactionShouldReturnAConsistentSnapshot() {
+    Key newEntityKey = Key.newBuilder(KEY1, "newKind", "name-01").build();
     EntityQuery entityQuery =
         Query.newEntityQueryBuilder()
             .setNamespace(NAMESPACE)
@@ -594,7 +595,7 @@ public void testRunAggregationQueryInTransactionShouldReturnAConsistentSnapshot(
               // creating a new entity
               Entity aNewEntity =
                   Entity.newBuilder(ENTITY2)
-                      .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build())
+                      .setKey(newEntityKey)
                       .set("v_int", 10)
                       .build();
               inFirstTransaction.put(aNewEntity);
@@ -631,6 +632,7 @@ public void testRunAggregationQueryInTransactionShouldReturnAConsistentSnapshot(
     // after second transaction is committed, count is updated to 2 now.
     assertThat(getOnlyElement(DATASTORE.runAggregation(aggregationQuery)).get("count"))
         .isEqualTo(2L);
+    DATASTORE.delete(newEntityKey);
   }
 
   /**

From d2e19691007310356bd5cb40eba291836e230b26 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Fri, 7 Oct 2022 13:14:03 +0530
Subject: [PATCH 54/82] Deleting all keys in datastore in integration test so
 that new test can start fresh

---
 .../google/cloud/datastore/it/ITDatastoreTest.java   | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
index 0bd8b8df4..db6d76209 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
@@ -69,6 +69,7 @@
 import com.google.cloud.datastore.ValueType;
 import com.google.cloud.datastore.testing.RemoteDatastoreHelper;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import com.google.datastore.v1.TransactionOptions;
 import com.google.datastore.v1.TransactionOptions.ReadOnly;
 import java.util.ArrayList;
@@ -184,7 +185,11 @@ public void setUp() {
 
   @After
   public void tearDown() {
-    DATASTORE.delete(KEY1, KEY2, KEY3);
+    EntityQuery allEntitiesQuery = Query.newEntityQueryBuilder().build();
+    QueryResults allEntities = DATASTORE.run(allEntitiesQuery);
+    Key[] keysToDelete =
+        ImmutableList.copyOf(allEntities).stream().map(Entity::getKey).toArray(Key[]::new);
+    DATASTORE.delete(keysToDelete);
   }
 
   private  Iterator getStronglyConsistentResults(Query scQuery, Query query)
@@ -594,10 +599,7 @@ public void testRunAggregationQueryInTransactionShouldReturnAConsistentSnapshot(
             inFirstTransaction -> {
               // creating a new entity
               Entity aNewEntity =
-                  Entity.newBuilder(ENTITY2)
-                      .setKey(newEntityKey)
-                      .set("v_int", 10)
-                      .build();
+                  Entity.newBuilder(ENTITY2).setKey(newEntityKey).set("v_int", 10).build();
               inFirstTransaction.put(aNewEntity);
 
               // count remains 2

From 39812d40ff1cacef4a86355f3f2fde1a5ec327a6 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Fri, 7 Oct 2022 15:20:28 +0530
Subject: [PATCH 55/82] Executing two read write transaction simultaneously and
 verifying their behaviour

---
 .../cloud/datastore/it/ITDatastoreTest.java   | 30 +++++++++++--------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
index db6d76209..853aefdc3 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
@@ -661,14 +661,15 @@ public void testRunAggregationQueryInAReadWriteTransactionShouldLockTheCountedDo
             .build();
 
     // read-write transaction
-    Transaction readWriteTransaction = DATASTORE.newTransaction();
+    Transaction readWriteTransaction1 = DATASTORE.newTransaction();
+    Transaction readWriteTransaction2 = DATASTORE.newTransaction();
 
-    // acquiring lock by executing query in transaction
-    assertThat(getOnlyElement(readWriteTransaction.runAggregation(aggregationQuery)).get("count"))
+    // acquiring lock by executing query in transaction 1
+    assertThat(getOnlyElement(readWriteTransaction1.runAggregation(aggregationQuery)).get("count"))
         .isEqualTo(2L);
 
-    // Waiting task will be blocked by ongoing transactions.
-    Future addNewEntityTaskOutsideTransaction =
+    // Transaction 2 will be blocked by ongoing Transaction 1.
+    Future addNewEntityTaskInTransaction2 =
         executor.submit(
             () -> {
               Entity aNewEntity =
@@ -676,17 +677,20 @@ public void testRunAggregationQueryInAReadWriteTransactionShouldLockTheCountedDo
                       .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build())
                       .set("v_int", 10)
                       .build();
-              DATASTORE.put(aNewEntity);
+              readWriteTransaction2.put(aNewEntity);
+              readWriteTransaction2.commit();
               return null;
             });
 
-    // should throw Timeout exception as we haven't yet committed the transaction
-    assertThrows(TimeoutException.class, () -> addNewEntityTaskOutsideTransaction.get(3, SECONDS));
-
-    // cleanup
-    readWriteTransaction.commit();
-    addNewEntityTaskOutsideTransaction.cancel(true);
-    executor.shutdownNow();
+    try {
+      // should throw Timeout exception as we haven't yet committed the Transaction 1
+      assertThrows(TimeoutException.class, () -> addNewEntityTaskInTransaction2.get(3, SECONDS));
+    } finally {
+      // cleanup
+      readWriteTransaction1.commit();
+      addNewEntityTaskInTransaction2.cancel(true);
+      executor.shutdownNow();
+    }
   }
 
   @Test

From 2d9e81a0d37afbfc02d94c6012f749197ca1cf11 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Fri, 7 Oct 2022 19:10:48 +0530
Subject: [PATCH 56/82] Removing tests to verify serializability as it's an
 underlying implementation detail

---
 .../cloud/datastore/it/ITDatastoreTest.java   | 56 -------------------
 1 file changed, 56 deletions(-)

diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
index 853aefdc3..1052487c1 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
@@ -637,62 +637,6 @@ public void testRunAggregationQueryInTransactionShouldReturnAConsistentSnapshot(
     DATASTORE.delete(newEntityKey);
   }
 
-  /**
-   * Data read or modified by a transaction cannot be concurrently modified.
-   *
-   * @see 
-   *     Source
-   */
-  @Test
-  public void testRunAggregationQueryInAReadWriteTransactionShouldLockTheCountedDocuments()
-      throws Exception {
-    ExecutorService executor = Executors.newSingleThreadExecutor();
-    EntityQuery entityQuery =
-        Query.newEntityQueryBuilder()
-            .setNamespace(NAMESPACE)
-            .setFilter(PropertyFilter.hasAncestor(KEY1))
-            .build();
-    AggregationQuery aggregationQuery =
-        Query.newAggregationQueryBuilder()
-            .setNamespace(NAMESPACE)
-            .over(entityQuery)
-            .addAggregation(count().as("count"))
-            .build();
-
-    // read-write transaction
-    Transaction readWriteTransaction1 = DATASTORE.newTransaction();
-    Transaction readWriteTransaction2 = DATASTORE.newTransaction();
-
-    // acquiring lock by executing query in transaction 1
-    assertThat(getOnlyElement(readWriteTransaction1.runAggregation(aggregationQuery)).get("count"))
-        .isEqualTo(2L);
-
-    // Transaction 2 will be blocked by ongoing Transaction 1.
-    Future addNewEntityTaskInTransaction2 =
-        executor.submit(
-            () -> {
-              Entity aNewEntity =
-                  Entity.newBuilder(ENTITY2)
-                      .setKey(Key.newBuilder(KEY1, "newKind", "name-01").build())
-                      .set("v_int", 10)
-                      .build();
-              readWriteTransaction2.put(aNewEntity);
-              readWriteTransaction2.commit();
-              return null;
-            });
-
-    try {
-      // should throw Timeout exception as we haven't yet committed the Transaction 1
-      assertThrows(TimeoutException.class, () -> addNewEntityTaskInTransaction2.get(3, SECONDS));
-    } finally {
-      // cleanup
-      readWriteTransaction1.commit();
-      addNewEntityTaskInTransaction2.cancel(true);
-      executor.shutdownNow();
-    }
-  }
-
   @Test
   public void testRunAggregationQueryInAReadOnlyTransactionShouldNotLockTheCountedDocuments()
       throws Exception {

From 90ab7cb89e4c6a96f4621409dfd79663f8a2913a Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Fri, 7 Oct 2022 19:32:51 +0530
Subject: [PATCH 57/82] Fixing lint

---
 .../java/com/google/cloud/datastore/it/ITDatastoreTest.java    | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
index 1052487c1..0c4de2f5a 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java
@@ -19,14 +19,12 @@
 import static com.google.cloud.datastore.aggregation.Aggregation.count;
 import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.truth.Truth.assertThat;
-import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
-import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -81,7 +79,6 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
-import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
 import org.junit.After;
 import org.junit.AfterClass;

From aae26201ac5a0a0f6951bdfb552bc1f59200a804 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Mon, 10 Oct 2022 11:38:52 +0530
Subject: [PATCH 58/82] Adding runAggregationQuery method to reflect config so
 that it's available and accessible in native image through reflection

---
 .../main/resources/META-INF/native-image/reflect-config.json   | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/datastore-v1-proto-client/src/main/resources/META-INF/native-image/reflect-config.json b/datastore-v1-proto-client/src/main/resources/META-INF/native-image/reflect-config.json
index 32b27f5d9..17876ff43 100644
--- a/datastore-v1-proto-client/src/main/resources/META-INF/native-image/reflect-config.json
+++ b/datastore-v1-proto-client/src/main/resources/META-INF/native-image/reflect-config.json
@@ -8,7 +8,8 @@
       {"name":"lookup","parameterTypes":["com.google.datastore.v1.LookupRequest"] },
       {"name":"reserveIds","parameterTypes":["com.google.datastore.v1.ReserveIdsRequest"] },
       {"name":"rollback","parameterTypes":["com.google.datastore.v1.RollbackRequest"] },
-      {"name":"runQuery","parameterTypes":["com.google.datastore.v1.RunQueryRequest"] }
+      {"name":"runQuery","parameterTypes":["com.google.datastore.v1.RunQueryRequest"] },
+      {"name":"runAggregationQuery","parameterTypes":["com.google.datastore.v1.RunAggregationQueryRequest"] }
     ]
   },
   {

From f9b41b29cc9010937dc7b0894d9735d9c3eb1169 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Mon, 10 Oct 2022 15:38:19 +0530
Subject: [PATCH 59/82] Aggregation query samples with limit and order by
 option

---
 .../datastore/AggregationQuerySample.java     | 89 ++++++++++++++++++-
 .../AggregationQuerySampleTestIT.java         | 17 +++-
 2 files changed, 102 insertions(+), 4 deletions(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
index 17a5b5e89..f06c48046 100644
--- a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
+++ b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
@@ -26,14 +26,15 @@
 import com.google.cloud.datastore.Key;
 import com.google.cloud.datastore.Query;
 import com.google.cloud.datastore.ReadOption;
+import com.google.cloud.datastore.StructuredQuery.OrderBy;
 import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
 import com.google.cloud.datastore.aggregation.Aggregation;
 import com.google.common.collect.Iterables;
 
 public class AggregationQuerySample {
 
-  public void aggregationQueryAndCountAggregation() {
-    // [START datastore_count_aggregation_query]
+  public void aggregationQueryAndCountAggregationOnKind() {
+    // [START datastore_count_aggregation_query_on_kind]
 
     // Instantiates a client
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
@@ -58,6 +59,7 @@ public void aggregationQueryAndCountAggregation() {
     // Creating an aggregation query to get the count of all candidates
     AggregationQuery allCandidatesCountQuery = Query.newAggregationQueryBuilder()
         .over(selectAllCandidates)
+        .addAggregation(Aggregation.count())
         .addAggregation(Aggregation.count().as("total_count"))
         .build();
     // Executing aggregation query
@@ -65,8 +67,89 @@ public void aggregationQueryAndCountAggregation() {
         datastore.runAggregation(allCandidatesCountQuery));
 
     System.out.printf("Total candidates count is %d", allCandidatesCountQueryResult.get("total_count")); // 3
+    System.out.printf("Total candidates from default alias is %d", allCandidatesCountQueryResult.get("property_1")); // 3
 
-    // [END datastore_count_aggregation_query]
+    // [END datastore_count_aggregation_query_on_kind]
+
+    datastore.delete(candidate1Key, candidate2Key, candidate3Key);
+  }
+
+  public void aggregationQueryAndCountAggregationWithLimit() {
+    // [START datastore_count_aggregation_query_with_limit]
+
+    // Instantiates a client
+    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+
+    // The kind for the new entity
+    String kind = "Candidate";
+
+    Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1");
+    Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2");
+    Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3");
+
+    // Save all the candidates
+    datastore.put(
+        Entity.newBuilder(candidate1Key).set("qualified", true).build(),
+        Entity.newBuilder(candidate2Key).set("qualified", false).build(),
+        Entity.newBuilder(candidate3Key).set("qualified", true).build()
+    );
+
+    EntityQuery selectAllCandidates = Query.newEntityQueryBuilder()
+        .setKind(kind)
+        .setLimit(2)
+        .build();
+    // Creating an aggregation query to get the count of all candidates
+    AggregationQuery allCandidatesCountQuery = Query.newAggregationQueryBuilder()
+        .over(selectAllCandidates)
+        .addAggregation(Aggregation.count().as("at_least"))
+        .build();
+    // Executing aggregation query
+    AggregationResult limitQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(allCandidatesCountQuery));
+
+    System.out.printf("We have at least %d candidates", limitQueryResult.get("at_least")); // 2
+
+    // [END datastore_count_aggregation_query_with_limit]
+
+    datastore.delete(candidate1Key, candidate2Key, candidate3Key);
+  }
+
+  public void aggregationQueryAndCountAggregationWithOrderBy() {
+    // [START datastore_count_aggregation_query_with_limit]
+
+    // Instantiates a client
+    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+
+    // The kind for the new entity
+    String kind = "Candidate";
+
+    Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1");
+    Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2");
+    Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3");
+
+    // Save all the candidates
+    datastore.put(
+        Entity.newBuilder(candidate1Key).set("qualified", true).set("rank", 1).build(),
+        Entity.newBuilder(candidate2Key).set("qualified", false).build(),  // no rank specified
+        Entity.newBuilder(candidate3Key).set("qualified", true).set("rank", 2).build()
+    );
+
+    EntityQuery selectAllCandidates = Query.newEntityQueryBuilder()
+        .setKind(kind)
+        .addOrderBy(OrderBy.asc("rank"))    // OrderBy acts as an existence filter
+        .build();
+    // Creating an aggregation query to get the count of all candidates
+    AggregationQuery allCandidatesCountQuery = Query.newAggregationQueryBuilder()
+        .over(selectAllCandidates)
+        .addAggregation(Aggregation.count().as("count"))
+        .build();
+    // Executing aggregation query
+    AggregationResult limitQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(allCandidatesCountQuery));
+
+    System.out.printf("Total %d candidates found with rank field", limitQueryResult.get("count")); // 2
+
+    // [END datastore_count_aggregation_query_with_limit]
 
     datastore.delete(candidate1Key, candidate2Key, candidate3Key);
   }
diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
index e15aaeb22..77cad5105 100644
--- a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
@@ -29,9 +29,24 @@ public class AggregationQuerySampleTestIT {
 
   @Test
   public void testAggregationQueryAndCountAggregationSample() {
-    sample.aggregationQueryAndCountAggregation();
+    sample.aggregationQueryAndCountAggregationOnKind();
 
     systemsOutRule.assertContains("Total candidates count is 3");
+    systemsOutRule.assertContains("Total candidates from default alias is 3");
+  }
+
+  @Test
+  public void testAggregationQueryAndCountAggregationWithLimitSample() {
+    sample.aggregationQueryAndCountAggregationWithLimit();
+
+    systemsOutRule.assertContains("We have at least 2 candidates");
+  }
+
+  @Test
+  public void testAggregationQueryAndCountAggregationWithOrderBySample() {
+    sample.aggregationQueryAndCountAggregationWithOrderBy();
+
+    systemsOutRule.assertContains("Total 2 candidates found with rank field");
   }
 
   @Test

From a4cf3662336bdb4d5d592182f368979656d5b2d2 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Mon, 10 Oct 2022 18:16:33 +0530
Subject: [PATCH 60/82] Aggregation query samples with transactions

---
 .../datastore/AggregationQuerySample.java     | 60 +++++++++++++++----
 .../AggregationQuerySampleTestIT.java         | 30 +++++++++-
 2 files changed, 77 insertions(+), 13 deletions(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
index f06c48046..7d302a099 100644
--- a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
+++ b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
@@ -19,7 +19,9 @@
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.Datastore.TransactionCallable;
 import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.DatastoreReaderWriter;
 import com.google.cloud.datastore.Entity;
 import com.google.cloud.datastore.EntityQuery;
 import com.google.cloud.datastore.GqlQuery;
@@ -67,11 +69,9 @@ public void aggregationQueryAndCountAggregationOnKind() {
         datastore.runAggregation(allCandidatesCountQuery));
 
     System.out.printf("Total candidates count is %d", allCandidatesCountQueryResult.get("total_count")); // 3
-    System.out.printf("Total candidates from default alias is %d", allCandidatesCountQueryResult.get("property_1")); // 3
+    System.out.printf("Total candidates (accessible from default alias) is %d", allCandidatesCountQueryResult.get("property_1")); // 3
 
     // [END datastore_count_aggregation_query_on_kind]
-
-    datastore.delete(candidate1Key, candidate2Key, candidate3Key);
   }
 
   public void aggregationQueryAndCountAggregationWithLimit() {
@@ -110,8 +110,6 @@ public void aggregationQueryAndCountAggregationWithLimit() {
     System.out.printf("We have at least %d candidates", limitQueryResult.get("at_least")); // 2
 
     // [END datastore_count_aggregation_query_with_limit]
-
-    datastore.delete(candidate1Key, candidate2Key, candidate3Key);
   }
 
   public void aggregationQueryAndCountAggregationWithOrderBy() {
@@ -150,8 +148,6 @@ public void aggregationQueryAndCountAggregationWithOrderBy() {
     System.out.printf("Total %d candidates found with rank field", limitQueryResult.get("count")); // 2
 
     // [END datastore_count_aggregation_query_with_limit]
-
-    datastore.delete(candidate1Key, candidate2Key, candidate3Key);
   }
 
   public void aggregationQueryAndCountAggregationWithPropertyFilter() {
@@ -203,8 +199,6 @@ public void aggregationQueryAndCountAggregationWithPropertyFilter() {
     System.out.printf("Total unqualified candidates count is %d", unQualifiedCandidatesCountQueryResult.get("total_unqualified_count")); // 1
 
     // [END datastore_count_aggregation_query_with_filters]
-
-    datastore.delete(candidate1Key, candidate2Key, candidate3Key);
   }
 
   public void aggregationQueryAndCountAggregationWithGqlQuery() {
@@ -261,8 +255,6 @@ public void aggregationQueryAndCountAggregationWithGqlQuery() {
     System.out.printf("Total qualified candidates count is %d", qualifiedCandidatesCountQueryResult.get("total_qualified_count")); // 2
 
     // [END datastore_count_aggregation_query_gql]
-
-    datastore.delete(candidate1Key, candidate2Key, candidate3Key);
   }
 
   public void aggregationQueryAndCountAggregationWithStaleRead() throws InterruptedException {
@@ -311,8 +303,52 @@ public void aggregationQueryAndCountAggregationWithStaleRead() throws Interrupte
     System.out.printf("Stale candidates count is %d", candidatesCountInPast.get("total_count")); // 2
 
     // [END datastore_count_aggregation_query_stale_read]
+  }
+
+  public void aggregationQueryInTransaction() {
+    // [START datastore_count_aggregation_query_in_transaction]
+
+    // Instantiates a client
+    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+
+    // The kind for the new entity
+    String kind = "Operations";
+
+    Key operation1Key = datastore.newKeyFactory().setKind(kind).newKey("operation1");
+    Key operation2Key = datastore.newKeyFactory().setKind(kind).newKey("operation2");
+
+
+    // Save all the candidates
+    datastore.put(
+        Entity.newBuilder(operation1Key).set("owner", "john").build(),
+        Entity.newBuilder(operation2Key).set("owner", "john").build()
+    );
 
-    datastore.delete(candidate1Key, candidate2Key, candidate3Key);
+    // Using transactions to maintain consistent application state.
+    datastore.runInTransaction((TransactionCallable) transaction -> {
+      EntityQuery operationsOfJohn = Query.newEntityQueryBuilder()
+          .setKind(kind)
+          .setFilter(PropertyFilter.eq("owner", "john"))
+          .build();
+      AggregationQuery totalOperationsQuery = Query.newAggregationQueryBuilder()
+          .over(operationsOfJohn)
+          .addAggregation(Aggregation.count().as("operations_count"))
+          .build();
+
+      Long operationsCount = Iterables.getOnlyElement(
+          datastore.runAggregation(totalOperationsQuery)).get("operations_count");
+
+      if (operationsCount < 2) {
+        Key newOperationKey = datastore.newKeyFactory().setKind(kind).newKey("operation3");
+        Entity newOperation = Entity.newBuilder(newOperationKey).set("owner", "john").build();
+        transaction.put(newOperation);
+      } else {
+        System.out.printf("Found existing %d operations, rolling back", operationsCount);
+        throw new Exception("User 'John' cannot have more than 2 operations");
+      }
+      return null;
+    });
+    // [END datastore_count_aggregation_query_in_transaction]
   }
 
 }
diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
index 77cad5105..4c64a4808 100644
--- a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
@@ -15,24 +15,45 @@
  */
 package com.example.datastore;
 
+import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.Entity;
+import com.google.cloud.datastore.EntityQuery;
+import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.Query;
+import com.google.cloud.datastore.QueryResults;
+import com.google.common.collect.ImmutableList;
 import com.rule.SystemsOutRule;
+import org.junit.After;
+import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.function.ThrowingRunnable;
 
 public class AggregationQuerySampleTestIT {
 
+  private Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+
   @Rule
   public final SystemsOutRule systemsOutRule = new SystemsOutRule();
 
   private final AggregationQuerySample sample = new AggregationQuerySample();
 
+  @After
+  public void tearDown() throws Exception {
+    EntityQuery allEntitiesQuery = Query.newEntityQueryBuilder().build();
+    QueryResults allEntities = datastore.run(allEntitiesQuery);
+    Key[] keysToDelete =
+        ImmutableList.copyOf(allEntities).stream().map(Entity::getKey).toArray(Key[]::new);
+    datastore.delete(keysToDelete);
+  }
 
   @Test
   public void testAggregationQueryAndCountAggregationSample() {
     sample.aggregationQueryAndCountAggregationOnKind();
 
     systemsOutRule.assertContains("Total candidates count is 3");
-    systemsOutRule.assertContains("Total candidates from default alias is 3");
+    systemsOutRule.assertContains("Total candidates (accessible from default alias) is 3");
   }
 
   @Test
@@ -73,4 +94,11 @@ public void testAggregationQueryAndCountWithStaleRead() throws InterruptedExcept
     systemsOutRule.assertContains("Latest candidates count is 3");
     systemsOutRule.assertContains("Stale candidates count is 2");
   }
+
+  @Test
+  public void testAggregationQueryAndCountWithTransaction() throws InterruptedException {
+    Assert.assertThrows(Exception.class, sample::aggregationQueryInTransaction);
+
+    systemsOutRule.assertContains("Found existing 2 operations, rolling back");
+  }
 }
\ No newline at end of file

From e265f355f2db44feddc365922b44eb1d26a67c67 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Mon, 10 Oct 2022 18:43:05 +0530
Subject: [PATCH 61/82] fixing lint and reusing sysoutRule

---
 .../datastore/AggregationQuerySample.java     | 44 ++++++++++++-------
 .../AggregationQuerySampleTestIT.java         | 27 ++++++------
 .../example/datastore/QuickstartSampleIT.java | 20 +++------
 .../test/java/com/rule/SystemsOutRule.java    | 29 ++++++------
 4 files changed, 63 insertions(+), 57 deletions(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
index 7d302a099..e9c38dc2b 100644
--- a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
+++ b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
@@ -5,7 +5,7 @@
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
- *     https://www.apache.org/licenses/LICENSE-2.0
+ * http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 package com.example.datastore;
 
 import com.google.cloud.Timestamp;
@@ -68,8 +69,10 @@ public void aggregationQueryAndCountAggregationOnKind() {
     AggregationResult allCandidatesCountQueryResult = Iterables.getOnlyElement(
         datastore.runAggregation(allCandidatesCountQuery));
 
-    System.out.printf("Total candidates count is %d", allCandidatesCountQueryResult.get("total_count")); // 3
-    System.out.printf("Total candidates (accessible from default alias) is %d", allCandidatesCountQueryResult.get("property_1")); // 3
+    System.out.printf("Total candidates count is %d",
+        allCandidatesCountQueryResult.get("total_count")); // 3
+    System.out.printf("Total candidates (accessible from default alias) is %d",
+        allCandidatesCountQueryResult.get("property_1")); // 3
 
     // [END datastore_count_aggregation_query_on_kind]
   }
@@ -145,7 +148,8 @@ public void aggregationQueryAndCountAggregationWithOrderBy() {
     AggregationResult limitQueryResult = Iterables.getOnlyElement(
         datastore.runAggregation(allCandidatesCountQuery));
 
-    System.out.printf("Total %d candidates found with rank field", limitQueryResult.get("count")); // 2
+    System.out.printf("Total %d candidates found with rank field",
+        limitQueryResult.get("count")); // 2
 
     // [END datastore_count_aggregation_query_with_limit]
   }
@@ -195,8 +199,10 @@ public void aggregationQueryAndCountAggregationWithPropertyFilter() {
     AggregationResult unQualifiedCandidatesCountQueryResult = Iterables.getOnlyElement(
         datastore.runAggregation(unqualifiedCandidatesCountQuery));
 
-    System.out.printf("Total qualified candidates count is %d", qualifiedCandidatesCountQueryResult.get("total_qualified_count")); // 2
-    System.out.printf("Total unqualified candidates count is %d", unQualifiedCandidatesCountQueryResult.get("total_unqualified_count")); // 1
+    System.out.printf("Total qualified candidates count is %d",
+        qualifiedCandidatesCountQueryResult.get("total_qualified_count")); // 2
+    System.out.printf("Total unqualified candidates count is %d",
+        unQualifiedCandidatesCountQueryResult.get("total_unqualified_count")); // 1
 
     // [END datastore_count_aggregation_query_with_filters]
   }
@@ -222,8 +228,8 @@ public void aggregationQueryAndCountAggregationWithGqlQuery() {
     );
 
     GqlQuery selectAllCandidates = Query.newGqlQueryBuilder(
-        "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(2) AS count_with_limit "
-            + "OVER (SELECT * FROM Candidate)")
+            "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(2) AS count_with_limit "
+                + "OVER (SELECT * FROM Candidate)")
         .setAllowLiteral(true)
         .build();
     // Creating an aggregation query to get the count of all candidates
@@ -234,9 +240,10 @@ public void aggregationQueryAndCountAggregationWithGqlQuery() {
     AggregationResult allCandidatesCountQueryResult = Iterables.getOnlyElement(
         datastore.runAggregation(allCandidatesCountQuery));
 
-    System.out.printf("We have at least %d candidates", allCandidatesCountQueryResult.get("count_with_limit")); // 2
-    System.out.printf("Total candidates count is %d", allCandidatesCountQueryResult.get("total_count")); // 3
-
+    System.out.printf("We have at least %d candidates",
+        allCandidatesCountQueryResult.get("count_with_limit")); // 2
+    System.out.printf("Total candidates count is %d",
+        allCandidatesCountQueryResult.get("total_count")); // 3
 
     GqlQuery qualifiedCandidates = Query.newGqlQueryBuilder(
             "AGGREGATE COUNT(*) AS total_qualified_count "
@@ -252,7 +259,8 @@ public void aggregationQueryAndCountAggregationWithGqlQuery() {
     AggregationResult qualifiedCandidatesCountQueryResult = Iterables.getOnlyElement(
         datastore.runAggregation(qualifiedCandidatesCountQuery));
 
-    System.out.printf("Total qualified candidates count is %d", qualifiedCandidatesCountQueryResult.get("total_qualified_count")); // 2
+    System.out.printf("Total qualified candidates count is %d",
+        qualifiedCandidatesCountQueryResult.get("total_qualified_count")); // 2
 
     // [END datastore_count_aggregation_query_gql]
   }
@@ -268,7 +276,6 @@ public void aggregationQueryAndCountAggregationWithStaleRead() throws Interrupte
 
     Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1");
     Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2");
-    Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3");
 
     // Saving only two candidates
     datastore.put(
@@ -276,10 +283,12 @@ public void aggregationQueryAndCountAggregationWithStaleRead() throws Interrupte
         Entity.newBuilder(candidate2Key).set("qualified", false).build()
     );
     Thread.sleep(1000);
-    Timestamp pastTimestamp = Timestamp.now(); // we have two candidates in database at this time.
+    final Timestamp pastTimestamp =
+        Timestamp.now(); // we have two candidates in database at this time.
 
     Thread.sleep(1000);
     // Saving third candidates
+    Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3");
     datastore.put(Entity.newBuilder(candidate3Key).set("qualified", false).build());
 
     EntityQuery selectAllCandidates = Query.newEntityQueryBuilder()
@@ -295,12 +304,14 @@ public void aggregationQueryAndCountAggregationWithStaleRead() throws Interrupte
     // Executing aggregation query
     AggregationResult candidatesCountLatest = Iterables.getOnlyElement(
         datastore.runAggregation(allCandidatesCountQuery));
-    System.out.printf("Latest candidates count is %d", candidatesCountLatest.get("total_count")); // 3
+    System.out.printf("Latest candidates count is %d",
+        candidatesCountLatest.get("total_count")); // 3
 
     // Executing aggregation query with past timestamp
     AggregationResult candidatesCountInPast = Iterables.getOnlyElement(
         datastore.runAggregation(allCandidatesCountQuery, ReadOption.readTime(pastTimestamp)));
-    System.out.printf("Stale candidates count is %d", candidatesCountInPast.get("total_count")); // 2
+    System.out.printf("Stale candidates count is %d",
+        candidatesCountInPast.get("total_count")); // 2
 
     // [END datastore_count_aggregation_query_stale_read]
   }
@@ -317,7 +328,6 @@ public void aggregationQueryInTransaction() {
     Key operation1Key = datastore.newKeyFactory().setKind(kind).newKey("operation1");
     Key operation2Key = datastore.newKeyFactory().setKind(kind).newKey("operation2");
 
-
     // Save all the candidates
     datastore.put(
         Entity.newBuilder(operation1Key).set("owner", "john").build(),
diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
index 4c64a4808..093fde14f 100644
--- a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
@@ -1,18 +1,19 @@
 /*
-  * Copyright 2022 Google LLC
-  *
-  * Licensed under the Apache License, Version 2.0 (the "License");
-  * you may not use this file except in compliance with the License.
-  * You may obtain a copy of the License at
-  *
-  *     https://www.apache.org/licenses/LICENSE-2.0
-  *
-  * Unless required by applicable law or agreed to in writing, software
-  * distributed under the License is distributed on an "AS IS" BASIS,
-  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-  * See the License for the specific language governing permissions and
-  * limitations under the License.
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
  */
+
 package com.example.datastore;
 
 import com.google.cloud.datastore.Datastore;
diff --git a/samples/snippets/src/test/java/com/example/datastore/QuickstartSampleIT.java b/samples/snippets/src/test/java/com/example/datastore/QuickstartSampleIT.java
index 1d46f20a4..35084937f 100644
--- a/samples/snippets/src/test/java/com/example/datastore/QuickstartSampleIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/QuickstartSampleIT.java
@@ -16,15 +16,13 @@
 
 package com.example.datastore;
 
-import static com.google.common.truth.Truth.assertThat;
-
 import com.google.cloud.datastore.Datastore;
 import com.google.cloud.datastore.DatastoreOptions;
 import com.google.cloud.datastore.Key;
-import java.io.ByteArrayOutputStream;
-import java.io.PrintStream;
+import com.rule.SystemsOutRule;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -33,8 +31,9 @@
 @RunWith(JUnit4.class)
 @SuppressWarnings("checkstyle:abbreviationaswordinname")
 public class QuickstartSampleIT {
-  private ByteArrayOutputStream bout;
-  private PrintStream out;
+
+  @Rule
+  public final SystemsOutRule systemsOutRule = new SystemsOutRule();
 
   private static final void deleteTestEntity() {
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
@@ -47,10 +46,6 @@ private static final void deleteTestEntity() {
   @Before
   public void setUp() {
     deleteTestEntity();
-
-    bout = new ByteArrayOutputStream();
-    out = new PrintStream(bout);
-    System.setOut(out);
   }
 
   @After
@@ -62,8 +57,7 @@ public void tearDown() {
   @Test
   public void testQuickstart() throws Exception {
     QuickstartSample.main();
-    String got = bout.toString();
-    assertThat(got).contains("Saved sampletask1: Buy milk");
-    assertThat(got).contains("Retrieved sampletask1: Buy milk");
+    systemsOutRule.assertContains("Saved sampletask1: Buy milk");
+    systemsOutRule.assertContains("Retrieved sampletask1: Buy milk");
   }
 }
diff --git a/samples/snippets/src/test/java/com/rule/SystemsOutRule.java b/samples/snippets/src/test/java/com/rule/SystemsOutRule.java
index eabc41ce9..9cdb2744e 100644
--- a/samples/snippets/src/test/java/com/rule/SystemsOutRule.java
+++ b/samples/snippets/src/test/java/com/rule/SystemsOutRule.java
@@ -1,18 +1,19 @@
 /*
-  * Copyright 2022 Google LLC
-  *
-  * Licensed under the Apache License, Version 2.0 (the "License");
-  * you may not use this file except in compliance with the License.
-  * You may obtain a copy of the License at
-  *
-  *     https://www.apache.org/licenses/LICENSE-2.0
-  *
-  * Unless required by applicable law or agreed to in writing, software
-  * distributed under the License is distributed on an "AS IS" BASIS,
-  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-  * See the License for the specific language governing permissions and
-  * limitations under the License.
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
  */
+
 package com.rule;
 
 import static com.google.common.truth.Truth.assertThat;
@@ -32,7 +33,7 @@ public Statement apply(Statement statement, Description description) {
       @Override
       public void evaluate() throws Throwable {
         // Setting up customized PrintStream
-        PrintStream originalOut = System.out;
+        final PrintStream originalOut = System.out;
         currentOut = new ByteArrayOutputStream();
         System.setOut(new PrintStream(currentOut));
 

From 4a9c7baaace44aa617682c5f1e9a6fa0b50d47bf Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Thu, 13 Oct 2022 09:49:43 +0530
Subject: [PATCH 62/82] correcting region tag

---
 .../java/com/example/datastore/AggregationQuerySample.java    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
index e9c38dc2b..acd8906f1 100644
--- a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
+++ b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
@@ -116,7 +116,7 @@ public void aggregationQueryAndCountAggregationWithLimit() {
   }
 
   public void aggregationQueryAndCountAggregationWithOrderBy() {
-    // [START datastore_count_aggregation_query_with_limit]
+    // [START datastore_count_aggregation_query_with_order_by]
 
     // Instantiates a client
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
@@ -151,7 +151,7 @@ public void aggregationQueryAndCountAggregationWithOrderBy() {
     System.out.printf("Total %d candidates found with rank field",
         limitQueryResult.get("count")); // 2
 
-    // [END datastore_count_aggregation_query_with_limit]
+    // [END datastore_count_aggregation_query_with_order_by]
   }
 
   public void aggregationQueryAndCountAggregationWithPropertyFilter() {

From 485212d944719ab533b80d5c12fe4a0f405d9485 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Thu, 13 Oct 2022 09:56:50 +0530
Subject: [PATCH 63/82] simplifying tearDown method

---
 .../datastore/AggregationQuerySampleTestIT.java       | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
index 093fde14f..82972037c 100644
--- a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
@@ -18,9 +18,8 @@
 
 import com.google.cloud.datastore.Datastore;
 import com.google.cloud.datastore.DatastoreOptions;
-import com.google.cloud.datastore.Entity;
-import com.google.cloud.datastore.EntityQuery;
 import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.KeyQuery;
 import com.google.cloud.datastore.Query;
 import com.google.cloud.datastore.QueryResults;
 import com.google.common.collect.ImmutableList;
@@ -29,7 +28,6 @@
 import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
-import org.junit.function.ThrowingRunnable;
 
 public class AggregationQuerySampleTestIT {
 
@@ -42,10 +40,9 @@ public class AggregationQuerySampleTestIT {
 
   @After
   public void tearDown() throws Exception {
-    EntityQuery allEntitiesQuery = Query.newEntityQueryBuilder().build();
-    QueryResults allEntities = datastore.run(allEntitiesQuery);
-    Key[] keysToDelete =
-        ImmutableList.copyOf(allEntities).stream().map(Entity::getKey).toArray(Key[]::new);
+    KeyQuery allKeysQuery = Query.newKeyQueryBuilder().build();
+    QueryResults allKeys = datastore.run(allKeysQuery);
+    Key[] keysToDelete = ImmutableList.copyOf(allKeys).toArray(new Key[0]);
     datastore.delete(keysToDelete);
   }
 

From 7d66848eae0b3c2cf4cbb5b776335859bef134e6 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Tue, 18 Oct 2022 12:28:21 +0530
Subject: [PATCH 64/82] Reverting back to using Task domain for aggregation
 query samples

---
 .../datastore/AggregationQuerySample.java     | 289 +++++++++---------
 .../AggregationQuerySampleTestIT.java         |  24 +-
 2 files changed, 156 insertions(+), 157 deletions(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
index acd8906f1..1bad06ed2 100644
--- a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
+++ b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
@@ -22,7 +22,6 @@
 import com.google.cloud.datastore.Datastore;
 import com.google.cloud.datastore.Datastore.TransactionCallable;
 import com.google.cloud.datastore.DatastoreOptions;
-import com.google.cloud.datastore.DatastoreReaderWriter;
 import com.google.cloud.datastore.Entity;
 import com.google.cloud.datastore.EntityQuery;
 import com.google.cloud.datastore.GqlQuery;
@@ -43,36 +42,36 @@ public void aggregationQueryAndCountAggregationOnKind() {
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
     // The kind for the new entity
-    String kind = "Candidate";
+    String kind = "Task";
 
-    Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1");
-    Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2");
-    Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3");
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
+    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
 
-    // Save all the candidates
+    // Save all the tasks
     datastore.put(
-        Entity.newBuilder(candidate1Key).set("qualified", true).build(),
-        Entity.newBuilder(candidate2Key).set("qualified", false).build(),
-        Entity.newBuilder(candidate3Key).set("qualified", true).build()
+        Entity.newBuilder(task1Key).set("done", true).build(),
+        Entity.newBuilder(task2Key).set("done", false).build(),
+        Entity.newBuilder(task3Key).set("done", true).build()
     );
 
-    EntityQuery selectAllCandidates = Query.newEntityQueryBuilder()
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
         .setKind(kind)
         .build();
-    // Creating an aggregation query to get the count of all candidates
-    AggregationQuery allCandidatesCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllCandidates)
+    // Creating an aggregation query to get the count of all tasks
+    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(selectAllTasks)
         .addAggregation(Aggregation.count())
         .addAggregation(Aggregation.count().as("total_count"))
         .build();
     // Executing aggregation query
-    AggregationResult allCandidatesCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(allCandidatesCountQuery));
+    AggregationResult allTasksCountQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(allTasksCountQuery));
 
-    System.out.printf("Total candidates count is %d",
-        allCandidatesCountQueryResult.get("total_count")); // 3
-    System.out.printf("Total candidates (accessible from default alias) is %d",
-        allCandidatesCountQueryResult.get("property_1")); // 3
+    System.out.printf("Total tasks count is %d",
+        allTasksCountQueryResult.get("total_count")); // 3
+    System.out.printf("Total tasks (accessible from default alias) is %d",
+        allTasksCountQueryResult.get("property_1")); // 3
 
     // [END datastore_count_aggregation_query_on_kind]
   }
@@ -84,33 +83,33 @@ public void aggregationQueryAndCountAggregationWithLimit() {
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
     // The kind for the new entity
-    String kind = "Candidate";
+    String kind = "Task";
 
-    Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1");
-    Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2");
-    Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3");
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
+    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
 
-    // Save all the candidates
+    // Save all the tasks
     datastore.put(
-        Entity.newBuilder(candidate1Key).set("qualified", true).build(),
-        Entity.newBuilder(candidate2Key).set("qualified", false).build(),
-        Entity.newBuilder(candidate3Key).set("qualified", true).build()
+        Entity.newBuilder(task1Key).set("done", true).build(),
+        Entity.newBuilder(task2Key).set("done", false).build(),
+        Entity.newBuilder(task3Key).set("done", true).build()
     );
 
-    EntityQuery selectAllCandidates = Query.newEntityQueryBuilder()
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
         .setKind(kind)
         .setLimit(2)
         .build();
-    // Creating an aggregation query to get the count of all candidates
-    AggregationQuery allCandidatesCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllCandidates)
+    // Creating an aggregation query to get the count of all tasks
+    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(selectAllTasks)
         .addAggregation(Aggregation.count().as("at_least"))
         .build();
     // Executing aggregation query
     AggregationResult limitQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(allCandidatesCountQuery));
+        datastore.runAggregation(allTasksCountQuery));
 
-    System.out.printf("We have at least %d candidates", limitQueryResult.get("at_least")); // 2
+    System.out.printf("We have at least %d tasks", limitQueryResult.get("at_least")); // 2
 
     // [END datastore_count_aggregation_query_with_limit]
   }
@@ -122,33 +121,33 @@ public void aggregationQueryAndCountAggregationWithOrderBy() {
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
     // The kind for the new entity
-    String kind = "Candidate";
+    String kind = "Task";
 
-    Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1");
-    Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2");
-    Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3");
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
+    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
 
-    // Save all the candidates
+    // Save all the tasks
     datastore.put(
-        Entity.newBuilder(candidate1Key).set("qualified", true).set("rank", 1).build(),
-        Entity.newBuilder(candidate2Key).set("qualified", false).build(),  // no rank specified
-        Entity.newBuilder(candidate3Key).set("qualified", true).set("rank", 2).build()
+        Entity.newBuilder(task1Key).set("done", true).set("priority", 1).build(),
+        Entity.newBuilder(task2Key).set("done", false).build(),  // no priority specified
+        Entity.newBuilder(task3Key).set("done", true).set("priority", 2).build()
     );
 
-    EntityQuery selectAllCandidates = Query.newEntityQueryBuilder()
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
         .setKind(kind)
-        .addOrderBy(OrderBy.asc("rank"))    // OrderBy acts as an existence filter
+        .addOrderBy(OrderBy.asc("priority"))    // OrderBy acts as an existence filter
         .build();
-    // Creating an aggregation query to get the count of all candidates
-    AggregationQuery allCandidatesCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllCandidates)
+    // Creating an aggregation query to get the count of all tasks
+    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(selectAllTasks)
         .addAggregation(Aggregation.count().as("count"))
         .build();
     // Executing aggregation query
     AggregationResult limitQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(allCandidatesCountQuery));
+        datastore.runAggregation(allTasksCountQuery));
 
-    System.out.printf("Total %d candidates found with rank field",
+    System.out.printf("Total %d tasks found with priority field",
         limitQueryResult.get("count")); // 2
 
     // [END datastore_count_aggregation_query_with_order_by]
@@ -161,48 +160,48 @@ public void aggregationQueryAndCountAggregationWithPropertyFilter() {
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
     // The kind for the new entity
-    String kind = "Candidate";
+    String kind = "Task";
 
-    Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1");
-    Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2");
-    Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3");
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
+    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
 
-    // Save all the candidates
+    // Save all the tasks
     datastore.put(
-        Entity.newBuilder(candidate1Key).set("qualified", true).build(),
-        Entity.newBuilder(candidate2Key).set("qualified", false).build(),
-        Entity.newBuilder(candidate3Key).set("qualified", true).build()
+        Entity.newBuilder(task1Key).set("done", true).build(),
+        Entity.newBuilder(task2Key).set("done", false).build(),
+        Entity.newBuilder(task3Key).set("done", true).build()
     );
 
-    EntityQuery qualifiedCandidates = Query.newEntityQueryBuilder()
+    EntityQuery completedTasks = Query.newEntityQueryBuilder()
         .setKind(kind)
-        .setFilter(PropertyFilter.eq("qualified", true))
+        .setFilter(PropertyFilter.eq("done", true))
         .build();
-    EntityQuery unQualifiedCandidates = Query.newEntityQueryBuilder()
+    EntityQuery remainingTasks = Query.newEntityQueryBuilder()
         .setKind(kind)
-        .setFilter(PropertyFilter.eq("qualified", false))
+        .setFilter(PropertyFilter.eq("done", false))
         .build();
-    // Creating an aggregation query to get the count of all qualified candidates
-    AggregationQuery qualifiedCandidatesCountQuery = Query.newAggregationQueryBuilder()
-        .over(qualifiedCandidates)
-        .addAggregation(Aggregation.count().as("total_qualified_count"))
+    // Creating an aggregation query to get the count of all completed tasks
+    AggregationQuery completedTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(completedTasks)
+        .addAggregation(Aggregation.count().as("total_completed_count"))
         .build();
-    // Creating an aggregation query to get the count of all unqualified candidates
-    AggregationQuery unqualifiedCandidatesCountQuery = Query.newAggregationQueryBuilder()
-        .over(unQualifiedCandidates)
-        .addAggregation(Aggregation.count().as("total_unqualified_count"))
+    // Creating an aggregation query to get the count of all remaining tasks
+    AggregationQuery remainingTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(remainingTasks)
+        .addAggregation(Aggregation.count().as("total_remaining_count"))
         .build();
 
     // Executing aggregation query
-    AggregationResult qualifiedCandidatesCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(qualifiedCandidatesCountQuery));
-    AggregationResult unQualifiedCandidatesCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(unqualifiedCandidatesCountQuery));
+    AggregationResult completedTasksCountQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(completedTasksCountQuery));
+    AggregationResult remainingTasksCountQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(remainingTasksCountQuery));
 
-    System.out.printf("Total qualified candidates count is %d",
-        qualifiedCandidatesCountQueryResult.get("total_qualified_count")); // 2
-    System.out.printf("Total unqualified candidates count is %d",
-        unQualifiedCandidatesCountQueryResult.get("total_unqualified_count")); // 1
+    System.out.printf("Total completed tasks count is %d",
+        completedTasksCountQueryResult.get("total_completed_count")); // 2
+    System.out.printf("Total remaining tasks count is %d",
+        remainingTasksCountQueryResult.get("total_remaining_count")); // 1
 
     // [END datastore_count_aggregation_query_with_filters]
   }
@@ -214,53 +213,53 @@ public void aggregationQueryAndCountAggregationWithGqlQuery() {
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
     // The kind for the new entity
-    String kind = "Candidate";
+    String kind = "Task";
 
-    Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1");
-    Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2");
-    Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3");
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
+    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
 
-    // Save all the candidates
+    // Save all the tasks
     datastore.put(
-        Entity.newBuilder(candidate1Key).set("qualified", true).build(),
-        Entity.newBuilder(candidate2Key).set("qualified", false).build(),
-        Entity.newBuilder(candidate3Key).set("qualified", true).build()
+        Entity.newBuilder(task1Key).set("done", true).build(),
+        Entity.newBuilder(task2Key).set("done", false).build(),
+        Entity.newBuilder(task3Key).set("done", true).build()
     );
 
-    GqlQuery selectAllCandidates = Query.newGqlQueryBuilder(
+    GqlQuery selectAllTasks = Query.newGqlQueryBuilder(
             "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(2) AS count_with_limit "
-                + "OVER (SELECT * FROM Candidate)")
+                + "OVER (SELECT * FROM Task)")
         .setAllowLiteral(true)
         .build();
-    // Creating an aggregation query to get the count of all candidates
-    AggregationQuery allCandidatesCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllCandidates)
+    // Creating an aggregation query to get the count of all tasks
+    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(selectAllTasks)
         .build();
     // Executing aggregation query
-    AggregationResult allCandidatesCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(allCandidatesCountQuery));
+    AggregationResult allTasksCountQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(allTasksCountQuery));
 
-    System.out.printf("We have at least %d candidates",
-        allCandidatesCountQueryResult.get("count_with_limit")); // 2
-    System.out.printf("Total candidates count is %d",
-        allCandidatesCountQueryResult.get("total_count")); // 3
+    System.out.printf("We have at least %d tasks",
+        allTasksCountQueryResult.get("count_with_limit")); // 2
+    System.out.printf("Total tasks count is %d",
+        allTasksCountQueryResult.get("total_count")); // 3
 
-    GqlQuery qualifiedCandidates = Query.newGqlQueryBuilder(
-            "AGGREGATE COUNT(*) AS total_qualified_count "
-                + "OVER (SELECT * FROM Candidate WHERE qualified = true)")
+    GqlQuery completedTasks = Query.newGqlQueryBuilder(
+            "AGGREGATE COUNT(*) AS total_completed_count "
+                + "OVER (SELECT * FROM Task WHERE done = true)")
         .setAllowLiteral(true)
         .build();
-    // Creating an aggregation query to get the count of all qualified candidates
-    AggregationQuery qualifiedCandidatesCountQuery = Query.newAggregationQueryBuilder()
-        .over(qualifiedCandidates)
+    // Creating an aggregation query to get the count of all completed tasks
+    AggregationQuery completedTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(completedTasks)
         .build();
 
     // Executing aggregation query
-    AggregationResult qualifiedCandidatesCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(qualifiedCandidatesCountQuery));
+    AggregationResult completedTasksCountQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(completedTasksCountQuery));
 
-    System.out.printf("Total qualified candidates count is %d",
-        qualifiedCandidatesCountQueryResult.get("total_qualified_count")); // 2
+    System.out.printf("Total completed tasks count is %d",
+        completedTasksCountQueryResult.get("total_completed_count")); // 2
 
     // [END datastore_count_aggregation_query_gql]
   }
@@ -272,46 +271,46 @@ public void aggregationQueryAndCountAggregationWithStaleRead() throws Interrupte
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
     // The kind for the new entity
-    String kind = "Candidate";
+    String kind = "Task";
 
-    Key candidate1Key = datastore.newKeyFactory().setKind(kind).newKey("candidate1");
-    Key candidate2Key = datastore.newKeyFactory().setKind(kind).newKey("candidate2");
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
 
-    // Saving only two candidates
+    // Saving only two tasks
     datastore.put(
-        Entity.newBuilder(candidate1Key).set("qualified", true).build(),
-        Entity.newBuilder(candidate2Key).set("qualified", false).build()
+        Entity.newBuilder(task1Key).set("done", true).build(),
+        Entity.newBuilder(task2Key).set("done", false).build()
     );
     Thread.sleep(1000);
     final Timestamp pastTimestamp =
-        Timestamp.now(); // we have two candidates in database at this time.
+        Timestamp.now(); // we have two tasks in database at this time.
 
     Thread.sleep(1000);
-    // Saving third candidates
-    Key candidate3Key = datastore.newKeyFactory().setKind(kind).newKey("candidate3");
-    datastore.put(Entity.newBuilder(candidate3Key).set("qualified", false).build());
+    // Saving third tasks
+    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
+    datastore.put(Entity.newBuilder(task3Key).set("done", false).build());
 
-    EntityQuery selectAllCandidates = Query.newEntityQueryBuilder()
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
         .setKind(kind)
         .build();
 
-    // Creating an aggregation query to get the count of all candidates
-    AggregationQuery allCandidatesCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllCandidates)
+    // Creating an aggregation query to get the count of all tasks
+    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(selectAllTasks)
         .addAggregation(Aggregation.count().as("total_count"))
         .build();
 
     // Executing aggregation query
-    AggregationResult candidatesCountLatest = Iterables.getOnlyElement(
-        datastore.runAggregation(allCandidatesCountQuery));
-    System.out.printf("Latest candidates count is %d",
-        candidatesCountLatest.get("total_count")); // 3
+    AggregationResult tasksCountLatest = Iterables.getOnlyElement(
+        datastore.runAggregation(allTasksCountQuery));
+    System.out.printf("Latest tasks count is %d",
+        tasksCountLatest.get("total_count")); // 3
 
     // Executing aggregation query with past timestamp
-    AggregationResult candidatesCountInPast = Iterables.getOnlyElement(
-        datastore.runAggregation(allCandidatesCountQuery, ReadOption.readTime(pastTimestamp)));
-    System.out.printf("Stale candidates count is %d",
-        candidatesCountInPast.get("total_count")); // 2
+    AggregationResult tasksCountInPast = Iterables.getOnlyElement(
+        datastore.runAggregation(allTasksCountQuery, ReadOption.readTime(pastTimestamp)));
+    System.out.printf("Stale tasks count is %d",
+        tasksCountInPast.get("total_count")); // 2
 
     // [END datastore_count_aggregation_query_stale_read]
   }
@@ -323,38 +322,38 @@ public void aggregationQueryInTransaction() {
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
     // The kind for the new entity
-    String kind = "Operations";
+    String kind = "Tasks";
 
-    Key operation1Key = datastore.newKeyFactory().setKind(kind).newKey("operation1");
-    Key operation2Key = datastore.newKeyFactory().setKind(kind).newKey("operation2");
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
 
-    // Save all the candidates
+    // Save all the tasks
     datastore.put(
-        Entity.newBuilder(operation1Key).set("owner", "john").build(),
-        Entity.newBuilder(operation2Key).set("owner", "john").build()
+        Entity.newBuilder(task1Key).set("owner", "john").build(),
+        Entity.newBuilder(task2Key).set("owner", "john").build()
     );
 
     // Using transactions to maintain consistent application state.
     datastore.runInTransaction((TransactionCallable) transaction -> {
-      EntityQuery operationsOfJohn = Query.newEntityQueryBuilder()
+      EntityQuery tasksOfJohn = Query.newEntityQueryBuilder()
           .setKind(kind)
           .setFilter(PropertyFilter.eq("owner", "john"))
           .build();
-      AggregationQuery totalOperationsQuery = Query.newAggregationQueryBuilder()
-          .over(operationsOfJohn)
-          .addAggregation(Aggregation.count().as("operations_count"))
+      AggregationQuery totalTasksQuery = Query.newAggregationQueryBuilder()
+          .over(tasksOfJohn)
+          .addAggregation(Aggregation.count().as("tasks_count"))
           .build();
 
-      Long operationsCount = Iterables.getOnlyElement(
-          datastore.runAggregation(totalOperationsQuery)).get("operations_count");
+      Long tasksCount = Iterables.getOnlyElement(
+          datastore.runAggregation(totalTasksQuery)).get("tasks_count");
 
-      if (operationsCount < 2) {
-        Key newOperationKey = datastore.newKeyFactory().setKind(kind).newKey("operation3");
-        Entity newOperation = Entity.newBuilder(newOperationKey).set("owner", "john").build();
-        transaction.put(newOperation);
+      if (tasksCount < 2) {
+        Key newTaskKey = datastore.newKeyFactory().setKind(kind).newKey("task3");
+        Entity newTask = Entity.newBuilder(newTaskKey).set("owner", "john").build();
+        transaction.put(newTask);
       } else {
-        System.out.printf("Found existing %d operations, rolling back", operationsCount);
-        throw new Exception("User 'John' cannot have more than 2 operations");
+        System.out.printf("Found existing %d tasks, rolling back", tasksCount);
+        throw new Exception("User 'John' cannot have more than 2 tasks");
       }
       return null;
     });
diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
index 82972037c..58fdb0772 100644
--- a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
@@ -50,53 +50,53 @@ public void tearDown() throws Exception {
   public void testAggregationQueryAndCountAggregationSample() {
     sample.aggregationQueryAndCountAggregationOnKind();
 
-    systemsOutRule.assertContains("Total candidates count is 3");
-    systemsOutRule.assertContains("Total candidates (accessible from default alias) is 3");
+    systemsOutRule.assertContains("Total tasks count is 3");
+    systemsOutRule.assertContains("Total tasks (accessible from default alias) is 3");
   }
 
   @Test
   public void testAggregationQueryAndCountAggregationWithLimitSample() {
     sample.aggregationQueryAndCountAggregationWithLimit();
 
-    systemsOutRule.assertContains("We have at least 2 candidates");
+    systemsOutRule.assertContains("We have at least 2 tasks");
   }
 
   @Test
   public void testAggregationQueryAndCountAggregationWithOrderBySample() {
     sample.aggregationQueryAndCountAggregationWithOrderBy();
 
-    systemsOutRule.assertContains("Total 2 candidates found with rank field");
+    systemsOutRule.assertContains("Total 2 tasks found with priority field");
   }
 
   @Test
   public void testAggregationQueryAndCountAggregationWithPropertyFilterSample() {
     sample.aggregationQueryAndCountAggregationWithPropertyFilter();
 
-    systemsOutRule.assertContains("Total qualified candidates count is 2");
-    systemsOutRule.assertContains("Total unqualified candidates count is 1");
+    systemsOutRule.assertContains("Total completed tasks count is 2");
+    systemsOutRule.assertContains("Total remaining tasks count is 1");
   }
 
   @Test
   public void testAggregationQueryAndCountAggregationSampleWithGqlQuery() {
     sample.aggregationQueryAndCountAggregationWithGqlQuery();
 
-    systemsOutRule.assertContains("We have at least 2 candidates");
-    systemsOutRule.assertContains("Total candidates count is 3");
-    systemsOutRule.assertContains("Total qualified candidates count is 2");
+    systemsOutRule.assertContains("We have at least 2 tasks");
+    systemsOutRule.assertContains("Total tasks count is 3");
+    systemsOutRule.assertContains("Total completed tasks count is 2");
   }
 
   @Test
   public void testAggregationQueryAndCountWithStaleRead() throws InterruptedException {
     sample.aggregationQueryAndCountAggregationWithStaleRead();
 
-    systemsOutRule.assertContains("Latest candidates count is 3");
-    systemsOutRule.assertContains("Stale candidates count is 2");
+    systemsOutRule.assertContains("Latest tasks count is 3");
+    systemsOutRule.assertContains("Stale tasks count is 2");
   }
 
   @Test
   public void testAggregationQueryAndCountWithTransaction() throws InterruptedException {
     Assert.assertThrows(Exception.class, sample::aggregationQueryInTransaction);
 
-    systemsOutRule.assertContains("Found existing 2 operations, rolling back");
+    systemsOutRule.assertContains("Found existing 2 tasks, rolling back");
   }
 }
\ No newline at end of file

From 1af69af085330f50660ceb258876d920941ee5c1 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Tue, 18 Oct 2022 13:13:35 +0530
Subject: [PATCH 65/82] Limiting tearDown method to delete only Task entities

---
 .../java/com/example/datastore/AggregationQuerySample.java    | 2 +-
 .../com/example/datastore/AggregationQuerySampleTestIT.java   | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
index 1bad06ed2..58d07f750 100644
--- a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
+++ b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
@@ -322,7 +322,7 @@ public void aggregationQueryInTransaction() {
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
     // The kind for the new entity
-    String kind = "Tasks";
+    String kind = "Task";
 
     Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
     Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
index 58fdb0772..e35d8ad16 100644
--- a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
@@ -40,7 +40,9 @@ public class AggregationQuerySampleTestIT {
 
   @After
   public void tearDown() throws Exception {
-    KeyQuery allKeysQuery = Query.newKeyQueryBuilder().build();
+    KeyQuery allKeysQuery = Query.newKeyQueryBuilder()
+        .setKind("Task")
+        .build();
     QueryResults allKeys = datastore.run(allKeysQuery);
     Key[] keysToDelete = ImmutableList.copyOf(allKeys).toArray(new Key[0]);
     datastore.delete(keysToDelete);

From 9871515a17552ebb620cc6789165ce934823d9ad Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Wed, 19 Oct 2022 11:08:16 +0530
Subject: [PATCH 66/82] Splitting count aggregation query samples into
 individual java files

---
 .../datastore/AggregationQuerySample.java     | 363 ------------------
 .../CountAggregationInTransaction.java        |  62 +++
 .../aggregation/CountAggregationOnKind.java   |  71 ++++
 .../CountAggregationWithGqlQuery.java         |  72 ++++
 .../CountAggregationWithLimit.java            |  52 +++
 .../CountAggregationWithOrderBy.java          |  54 +++
 .../CountAggregationWithPropertyFilter.java   |  69 ++++
 .../CountAggregationWithStaleRead.java        |  69 ++++
 .../AggregationQuerySampleTestIT.java         |  25 +-
 9 files changed, 464 insertions(+), 373 deletions(-)
 delete mode 100644 samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
 create mode 100644 samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
 create mode 100644 samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
 create mode 100644 samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
 create mode 100644 samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
 create mode 100644 samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
 create mode 100644 samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
 create mode 100644 samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java

diff --git a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java b/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
deleted file mode 100644
index 58d07f750..000000000
--- a/samples/snippets/src/main/java/com/example/datastore/AggregationQuerySample.java
+++ /dev/null
@@ -1,363 +0,0 @@
-/*
- * Copyright 2022 Google LLC
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.example.datastore;
-
-import com.google.cloud.Timestamp;
-import com.google.cloud.datastore.AggregationQuery;
-import com.google.cloud.datastore.AggregationResult;
-import com.google.cloud.datastore.Datastore;
-import com.google.cloud.datastore.Datastore.TransactionCallable;
-import com.google.cloud.datastore.DatastoreOptions;
-import com.google.cloud.datastore.Entity;
-import com.google.cloud.datastore.EntityQuery;
-import com.google.cloud.datastore.GqlQuery;
-import com.google.cloud.datastore.Key;
-import com.google.cloud.datastore.Query;
-import com.google.cloud.datastore.ReadOption;
-import com.google.cloud.datastore.StructuredQuery.OrderBy;
-import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
-import com.google.cloud.datastore.aggregation.Aggregation;
-import com.google.common.collect.Iterables;
-
-public class AggregationQuerySample {
-
-  public void aggregationQueryAndCountAggregationOnKind() {
-    // [START datastore_count_aggregation_query_on_kind]
-
-    // Instantiates a client
-    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
-
-    // The kind for the new entity
-    String kind = "Task";
-
-    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
-    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
-    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
-
-    // Save all the tasks
-    datastore.put(
-        Entity.newBuilder(task1Key).set("done", true).build(),
-        Entity.newBuilder(task2Key).set("done", false).build(),
-        Entity.newBuilder(task3Key).set("done", true).build()
-    );
-
-    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .build();
-    // Creating an aggregation query to get the count of all tasks
-    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllTasks)
-        .addAggregation(Aggregation.count())
-        .addAggregation(Aggregation.count().as("total_count"))
-        .build();
-    // Executing aggregation query
-    AggregationResult allTasksCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(allTasksCountQuery));
-
-    System.out.printf("Total tasks count is %d",
-        allTasksCountQueryResult.get("total_count")); // 3
-    System.out.printf("Total tasks (accessible from default alias) is %d",
-        allTasksCountQueryResult.get("property_1")); // 3
-
-    // [END datastore_count_aggregation_query_on_kind]
-  }
-
-  public void aggregationQueryAndCountAggregationWithLimit() {
-    // [START datastore_count_aggregation_query_with_limit]
-
-    // Instantiates a client
-    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
-
-    // The kind for the new entity
-    String kind = "Task";
-
-    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
-    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
-    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
-
-    // Save all the tasks
-    datastore.put(
-        Entity.newBuilder(task1Key).set("done", true).build(),
-        Entity.newBuilder(task2Key).set("done", false).build(),
-        Entity.newBuilder(task3Key).set("done", true).build()
-    );
-
-    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .setLimit(2)
-        .build();
-    // Creating an aggregation query to get the count of all tasks
-    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllTasks)
-        .addAggregation(Aggregation.count().as("at_least"))
-        .build();
-    // Executing aggregation query
-    AggregationResult limitQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(allTasksCountQuery));
-
-    System.out.printf("We have at least %d tasks", limitQueryResult.get("at_least")); // 2
-
-    // [END datastore_count_aggregation_query_with_limit]
-  }
-
-  public void aggregationQueryAndCountAggregationWithOrderBy() {
-    // [START datastore_count_aggregation_query_with_order_by]
-
-    // Instantiates a client
-    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
-
-    // The kind for the new entity
-    String kind = "Task";
-
-    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
-    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
-    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
-
-    // Save all the tasks
-    datastore.put(
-        Entity.newBuilder(task1Key).set("done", true).set("priority", 1).build(),
-        Entity.newBuilder(task2Key).set("done", false).build(),  // no priority specified
-        Entity.newBuilder(task3Key).set("done", true).set("priority", 2).build()
-    );
-
-    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .addOrderBy(OrderBy.asc("priority"))    // OrderBy acts as an existence filter
-        .build();
-    // Creating an aggregation query to get the count of all tasks
-    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllTasks)
-        .addAggregation(Aggregation.count().as("count"))
-        .build();
-    // Executing aggregation query
-    AggregationResult limitQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(allTasksCountQuery));
-
-    System.out.printf("Total %d tasks found with priority field",
-        limitQueryResult.get("count")); // 2
-
-    // [END datastore_count_aggregation_query_with_order_by]
-  }
-
-  public void aggregationQueryAndCountAggregationWithPropertyFilter() {
-    // [START datastore_count_aggregation_query_with_filters]
-
-    // Instantiates a client
-    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
-
-    // The kind for the new entity
-    String kind = "Task";
-
-    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
-    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
-    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
-
-    // Save all the tasks
-    datastore.put(
-        Entity.newBuilder(task1Key).set("done", true).build(),
-        Entity.newBuilder(task2Key).set("done", false).build(),
-        Entity.newBuilder(task3Key).set("done", true).build()
-    );
-
-    EntityQuery completedTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .setFilter(PropertyFilter.eq("done", true))
-        .build();
-    EntityQuery remainingTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .setFilter(PropertyFilter.eq("done", false))
-        .build();
-    // Creating an aggregation query to get the count of all completed tasks
-    AggregationQuery completedTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(completedTasks)
-        .addAggregation(Aggregation.count().as("total_completed_count"))
-        .build();
-    // Creating an aggregation query to get the count of all remaining tasks
-    AggregationQuery remainingTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(remainingTasks)
-        .addAggregation(Aggregation.count().as("total_remaining_count"))
-        .build();
-
-    // Executing aggregation query
-    AggregationResult completedTasksCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(completedTasksCountQuery));
-    AggregationResult remainingTasksCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(remainingTasksCountQuery));
-
-    System.out.printf("Total completed tasks count is %d",
-        completedTasksCountQueryResult.get("total_completed_count")); // 2
-    System.out.printf("Total remaining tasks count is %d",
-        remainingTasksCountQueryResult.get("total_remaining_count")); // 1
-
-    // [END datastore_count_aggregation_query_with_filters]
-  }
-
-  public void aggregationQueryAndCountAggregationWithGqlQuery() {
-    // [START datastore_count_aggregation_query_gql]
-
-    // Instantiates a client
-    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
-
-    // The kind for the new entity
-    String kind = "Task";
-
-    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
-    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
-    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
-
-    // Save all the tasks
-    datastore.put(
-        Entity.newBuilder(task1Key).set("done", true).build(),
-        Entity.newBuilder(task2Key).set("done", false).build(),
-        Entity.newBuilder(task3Key).set("done", true).build()
-    );
-
-    GqlQuery selectAllTasks = Query.newGqlQueryBuilder(
-            "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(2) AS count_with_limit "
-                + "OVER (SELECT * FROM Task)")
-        .setAllowLiteral(true)
-        .build();
-    // Creating an aggregation query to get the count of all tasks
-    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllTasks)
-        .build();
-    // Executing aggregation query
-    AggregationResult allTasksCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(allTasksCountQuery));
-
-    System.out.printf("We have at least %d tasks",
-        allTasksCountQueryResult.get("count_with_limit")); // 2
-    System.out.printf("Total tasks count is %d",
-        allTasksCountQueryResult.get("total_count")); // 3
-
-    GqlQuery completedTasks = Query.newGqlQueryBuilder(
-            "AGGREGATE COUNT(*) AS total_completed_count "
-                + "OVER (SELECT * FROM Task WHERE done = true)")
-        .setAllowLiteral(true)
-        .build();
-    // Creating an aggregation query to get the count of all completed tasks
-    AggregationQuery completedTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(completedTasks)
-        .build();
-
-    // Executing aggregation query
-    AggregationResult completedTasksCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(completedTasksCountQuery));
-
-    System.out.printf("Total completed tasks count is %d",
-        completedTasksCountQueryResult.get("total_completed_count")); // 2
-
-    // [END datastore_count_aggregation_query_gql]
-  }
-
-  public void aggregationQueryAndCountAggregationWithStaleRead() throws InterruptedException {
-    // [START datastore_count_aggregation_query_stale_read]
-
-    // Instantiates a client
-    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
-
-    // The kind for the new entity
-    String kind = "Task";
-
-    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
-    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
-
-    // Saving only two tasks
-    datastore.put(
-        Entity.newBuilder(task1Key).set("done", true).build(),
-        Entity.newBuilder(task2Key).set("done", false).build()
-    );
-    Thread.sleep(1000);
-    final Timestamp pastTimestamp =
-        Timestamp.now(); // we have two tasks in database at this time.
-
-    Thread.sleep(1000);
-    // Saving third tasks
-    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
-    datastore.put(Entity.newBuilder(task3Key).set("done", false).build());
-
-    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .build();
-
-    // Creating an aggregation query to get the count of all tasks
-    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllTasks)
-        .addAggregation(Aggregation.count().as("total_count"))
-        .build();
-
-    // Executing aggregation query
-    AggregationResult tasksCountLatest = Iterables.getOnlyElement(
-        datastore.runAggregation(allTasksCountQuery));
-    System.out.printf("Latest tasks count is %d",
-        tasksCountLatest.get("total_count")); // 3
-
-    // Executing aggregation query with past timestamp
-    AggregationResult tasksCountInPast = Iterables.getOnlyElement(
-        datastore.runAggregation(allTasksCountQuery, ReadOption.readTime(pastTimestamp)));
-    System.out.printf("Stale tasks count is %d",
-        tasksCountInPast.get("total_count")); // 2
-
-    // [END datastore_count_aggregation_query_stale_read]
-  }
-
-  public void aggregationQueryInTransaction() {
-    // [START datastore_count_aggregation_query_in_transaction]
-
-    // Instantiates a client
-    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
-
-    // The kind for the new entity
-    String kind = "Task";
-
-    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
-    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
-
-    // Save all the tasks
-    datastore.put(
-        Entity.newBuilder(task1Key).set("owner", "john").build(),
-        Entity.newBuilder(task2Key).set("owner", "john").build()
-    );
-
-    // Using transactions to maintain consistent application state.
-    datastore.runInTransaction((TransactionCallable) transaction -> {
-      EntityQuery tasksOfJohn = Query.newEntityQueryBuilder()
-          .setKind(kind)
-          .setFilter(PropertyFilter.eq("owner", "john"))
-          .build();
-      AggregationQuery totalTasksQuery = Query.newAggregationQueryBuilder()
-          .over(tasksOfJohn)
-          .addAggregation(Aggregation.count().as("tasks_count"))
-          .build();
-
-      Long tasksCount = Iterables.getOnlyElement(
-          datastore.runAggregation(totalTasksQuery)).get("tasks_count");
-
-      if (tasksCount < 2) {
-        Key newTaskKey = datastore.newKeyFactory().setKind(kind).newKey("task3");
-        Entity newTask = Entity.newBuilder(newTaskKey).set("owner", "john").build();
-        transaction.put(newTask);
-      } else {
-        System.out.printf("Found existing %d tasks, rolling back", tasksCount);
-        throw new Exception("User 'John' cannot have more than 2 tasks");
-      }
-      return null;
-    });
-    // [END datastore_count_aggregation_query_in_transaction]
-  }
-
-}
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
new file mode 100644
index 000000000..cbe3f4dbe
--- /dev/null
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
@@ -0,0 +1,62 @@
+package com.example.datastore.aggregation;
+
+import com.google.cloud.datastore.AggregationQuery;
+import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.Datastore.TransactionCallable;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.Entity;
+import com.google.cloud.datastore.EntityQuery;
+import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.Query;
+import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
+import com.google.cloud.datastore.aggregation.Aggregation;
+import com.google.common.collect.Iterables;
+
+public class CountAggregationInTransaction {
+
+  public static void invoke() {
+    // [START datastore_count_aggregation_query_in_transaction]
+
+    // Instantiates a client
+    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+
+    // The kind for the new entity
+    String kind = "Task";
+
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
+
+    // Save all the tasks
+    datastore.put(
+        Entity.newBuilder(task1Key).set("owner", "john").build(),
+        Entity.newBuilder(task2Key).set("owner", "john").build()
+    );
+
+    // Using transactions to maintain consistent application state.
+    datastore.runInTransaction((TransactionCallable) transaction -> {
+      EntityQuery tasksOfJohn = Query.newEntityQueryBuilder()
+          .setKind(kind)
+          .setFilter(PropertyFilter.eq("owner", "john"))
+          .build();
+      AggregationQuery totalTasksQuery = Query.newAggregationQueryBuilder()
+          .over(tasksOfJohn)
+          .addAggregation(Aggregation.count().as("tasks_count"))
+          .build();
+
+      Long tasksCount = Iterables.getOnlyElement(
+          datastore.runAggregation(totalTasksQuery)).get("tasks_count");
+
+      if (tasksCount < 2) {
+        Key newTaskKey = datastore.newKeyFactory().setKind(kind).newKey("task3");
+        Entity newTask = Entity.newBuilder(newTaskKey).set("owner", "john").build();
+        transaction.put(newTask);
+      } else {
+        System.out.printf("Found existing %d tasks, rolling back", tasksCount);
+        throw new Exception("User 'John' cannot have more than 2 tasks");
+      }
+      return null;
+    });
+    // [END datastore_count_aggregation_query_in_transaction]
+
+  }
+}
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
new file mode 100644
index 000000000..a047c07b1
--- /dev/null
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
@@ -0,0 +1,71 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.example.datastore.aggregation;
+
+import com.google.cloud.datastore.AggregationQuery;
+import com.google.cloud.datastore.AggregationResult;
+import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.Entity;
+import com.google.cloud.datastore.EntityQuery;
+import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.Query;
+import com.google.cloud.datastore.aggregation.Aggregation;
+import com.google.common.collect.Iterables;
+
+public class CountAggregationOnKind {
+
+  public static void invoke() {
+    // [START datastore_count_aggregation_query_on_kind]
+
+    // Instantiates a client
+    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+
+    // The kind for the new entity
+    String kind = "Task";
+
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
+    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
+
+    // Save all the tasks
+    datastore.put(
+        Entity.newBuilder(task1Key).set("done", true).build(),
+        Entity.newBuilder(task2Key).set("done", false).build(),
+        Entity.newBuilder(task3Key).set("done", true).build()
+    );
+
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
+        .setKind(kind)
+        .build();
+    // Creating an aggregation query to get the count of all tasks
+    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(selectAllTasks)
+        .addAggregation(Aggregation.count())
+        .addAggregation(Aggregation.count().as("total_count"))
+        .build();
+    // Executing aggregation query
+    AggregationResult allTasksCountQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(allTasksCountQuery));
+
+    System.out.printf("Total tasks count is %d",
+        allTasksCountQueryResult.get("total_count")); // 3
+    System.out.printf("Total tasks (accessible from default alias) is %d",
+        allTasksCountQueryResult.get("property_1")); // 3
+
+    // [END datastore_count_aggregation_query_on_kind]
+  }
+}
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
new file mode 100644
index 000000000..4eeec810c
--- /dev/null
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
@@ -0,0 +1,72 @@
+package com.example.datastore.aggregation;
+
+import com.google.cloud.datastore.AggregationQuery;
+import com.google.cloud.datastore.AggregationResult;
+import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.Entity;
+import com.google.cloud.datastore.GqlQuery;
+import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.Query;
+import com.google.common.collect.Iterables;
+
+public class CountAggregationWithGqlQuery {
+
+  public static void invoke() {
+    // [START datastore_count_aggregation_query_gql]
+
+    // Instantiates a client
+    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+
+    // The kind for the new entity
+    String kind = "Task";
+
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
+    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
+
+    // Save all the tasks
+    datastore.put(
+        Entity.newBuilder(task1Key).set("done", true).build(),
+        Entity.newBuilder(task2Key).set("done", false).build(),
+        Entity.newBuilder(task3Key).set("done", true).build()
+    );
+
+    GqlQuery selectAllTasks = Query.newGqlQueryBuilder(
+            "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(2) AS count_with_limit "
+                + "OVER (SELECT * FROM Task)")
+        .setAllowLiteral(true)
+        .build();
+    // Creating an aggregation query to get the count of all tasks
+    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(selectAllTasks)
+        .build();
+    // Executing aggregation query
+    AggregationResult allTasksCountQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(allTasksCountQuery));
+
+    System.out.printf("We have at least %d tasks",
+        allTasksCountQueryResult.get("count_with_limit")); // 2
+    System.out.printf("Total tasks count is %d",
+        allTasksCountQueryResult.get("total_count")); // 3
+
+    GqlQuery completedTasks = Query.newGqlQueryBuilder(
+            "AGGREGATE COUNT(*) AS total_completed_count "
+                + "OVER (SELECT * FROM Task WHERE done = true)")
+        .setAllowLiteral(true)
+        .build();
+    // Creating an aggregation query to get the count of all completed tasks
+    AggregationQuery completedTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(completedTasks)
+        .build();
+
+    // Executing aggregation query
+    AggregationResult completedTasksCountQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(completedTasksCountQuery));
+
+    System.out.printf("Total completed tasks count is %d",
+        completedTasksCountQueryResult.get("total_completed_count")); // 2
+
+    // [END datastore_count_aggregation_query_gql]
+  }
+}
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
new file mode 100644
index 000000000..1813f4109
--- /dev/null
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
@@ -0,0 +1,52 @@
+package com.example.datastore.aggregation;
+
+import com.google.cloud.datastore.AggregationQuery;
+import com.google.cloud.datastore.AggregationResult;
+import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.Entity;
+import com.google.cloud.datastore.EntityQuery;
+import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.Query;
+import com.google.cloud.datastore.aggregation.Aggregation;
+import com.google.common.collect.Iterables;
+
+public class CountAggregationWithLimit {
+  public static void invoke() {
+    // [START datastore_count_aggregation_query_with_limit]
+
+    // Instantiates a client
+    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+
+    // The kind for the new entity
+    String kind = "Task";
+
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
+    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
+
+    // Save all the tasks
+    datastore.put(
+        Entity.newBuilder(task1Key).set("done", true).build(),
+        Entity.newBuilder(task2Key).set("done", false).build(),
+        Entity.newBuilder(task3Key).set("done", true).build()
+    );
+
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
+        .setKind(kind)
+        .setLimit(2)
+        .build();
+    // Creating an aggregation query to get the count of all tasks
+    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(selectAllTasks)
+        .addAggregation(Aggregation.count().as("at_least"))
+        .build();
+    // Executing aggregation query
+    AggregationResult limitQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(allTasksCountQuery));
+
+    System.out.printf("We have at least %d tasks", limitQueryResult.get("at_least")); // 2
+
+    // [END datastore_count_aggregation_query_with_limit]
+  }
+}
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
new file mode 100644
index 000000000..4baa74672
--- /dev/null
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
@@ -0,0 +1,54 @@
+package com.example.datastore.aggregation;
+
+import com.google.cloud.datastore.AggregationQuery;
+import com.google.cloud.datastore.AggregationResult;
+import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.Entity;
+import com.google.cloud.datastore.EntityQuery;
+import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.Query;
+import com.google.cloud.datastore.StructuredQuery.OrderBy;
+import com.google.cloud.datastore.aggregation.Aggregation;
+import com.google.common.collect.Iterables;
+
+public class CountAggregationWithOrderBy {
+  public static void invoke() {
+    // [START datastore_count_aggregation_query_with_order_by]
+
+    // Instantiates a client
+    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+
+    // The kind for the new entity
+    String kind = "Task";
+
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
+    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
+
+    // Save all the tasks
+    datastore.put(
+        Entity.newBuilder(task1Key).set("done", true).set("priority", 1).build(),
+        Entity.newBuilder(task2Key).set("done", false).build(),  // no priority specified
+        Entity.newBuilder(task3Key).set("done", true).set("priority", 2).build()
+    );
+
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
+        .setKind(kind)
+        .addOrderBy(OrderBy.asc("priority"))    // OrderBy acts as an existence filter
+        .build();
+    // Creating an aggregation query to get the count of all tasks
+    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(selectAllTasks)
+        .addAggregation(Aggregation.count().as("count"))
+        .build();
+    // Executing aggregation query
+    AggregationResult limitQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(allTasksCountQuery));
+
+    System.out.printf("Total %d tasks found with priority field",
+        limitQueryResult.get("count")); // 2
+
+    // [END datastore_count_aggregation_query_with_order_by]
+  }
+}
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
new file mode 100644
index 000000000..01c4897a3
--- /dev/null
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
@@ -0,0 +1,69 @@
+package com.example.datastore.aggregation;
+
+import com.google.cloud.datastore.AggregationQuery;
+import com.google.cloud.datastore.AggregationResult;
+import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.Entity;
+import com.google.cloud.datastore.EntityQuery;
+import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.Query;
+import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
+import com.google.cloud.datastore.aggregation.Aggregation;
+import com.google.common.collect.Iterables;
+
+public class CountAggregationWithPropertyFilter {
+
+  public static void invoke() {
+    // [START datastore_count_aggregation_query_with_filters]
+
+    // Instantiates a client
+    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+
+    // The kind for the new entity
+    String kind = "Task";
+
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
+    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
+
+    // Save all the tasks
+    datastore.put(
+        Entity.newBuilder(task1Key).set("done", true).build(),
+        Entity.newBuilder(task2Key).set("done", false).build(),
+        Entity.newBuilder(task3Key).set("done", true).build()
+    );
+
+    EntityQuery completedTasks = Query.newEntityQueryBuilder()
+        .setKind(kind)
+        .setFilter(PropertyFilter.eq("done", true))
+        .build();
+    EntityQuery remainingTasks = Query.newEntityQueryBuilder()
+        .setKind(kind)
+        .setFilter(PropertyFilter.eq("done", false))
+        .build();
+    // Creating an aggregation query to get the count of all completed tasks
+    AggregationQuery completedTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(completedTasks)
+        .addAggregation(Aggregation.count().as("total_completed_count"))
+        .build();
+    // Creating an aggregation query to get the count of all remaining tasks
+    AggregationQuery remainingTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(remainingTasks)
+        .addAggregation(Aggregation.count().as("total_remaining_count"))
+        .build();
+
+    // Executing aggregation query
+    AggregationResult completedTasksCountQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(completedTasksCountQuery));
+    AggregationResult remainingTasksCountQueryResult = Iterables.getOnlyElement(
+        datastore.runAggregation(remainingTasksCountQuery));
+
+    System.out.printf("Total completed tasks count is %d",
+        completedTasksCountQueryResult.get("total_completed_count")); // 2
+    System.out.printf("Total remaining tasks count is %d",
+        remainingTasksCountQueryResult.get("total_remaining_count")); // 1
+
+    // [END datastore_count_aggregation_query_with_filters]
+  }
+}
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
new file mode 100644
index 000000000..b558ad4f8
--- /dev/null
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
@@ -0,0 +1,69 @@
+package com.example.datastore.aggregation;
+
+import com.google.cloud.Timestamp;
+import com.google.cloud.datastore.AggregationQuery;
+import com.google.cloud.datastore.AggregationResult;
+import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.Entity;
+import com.google.cloud.datastore.EntityQuery;
+import com.google.cloud.datastore.Key;
+import com.google.cloud.datastore.Query;
+import com.google.cloud.datastore.ReadOption;
+import com.google.cloud.datastore.aggregation.Aggregation;
+import com.google.common.collect.Iterables;
+
+public class CountAggregationWithStaleRead {
+
+  public static void invoke() throws InterruptedException {
+    // [START datastore_count_aggregation_query_stale_read]
+
+    // Instantiates a client
+    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+
+    // The kind for the new entity
+    String kind = "Task";
+
+    Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
+    Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
+
+    // Saving only two tasks
+    datastore.put(
+        Entity.newBuilder(task1Key).set("done", true).build(),
+        Entity.newBuilder(task2Key).set("done", false).build()
+    );
+    Thread.sleep(1000);
+    final Timestamp pastTimestamp =
+        Timestamp.now(); // we have two tasks in database at this time.
+
+    Thread.sleep(1000);
+    // Saving third tasks
+    Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
+    datastore.put(Entity.newBuilder(task3Key).set("done", false).build());
+
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
+        .setKind(kind)
+        .build();
+
+    // Creating an aggregation query to get the count of all tasks
+    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
+        .over(selectAllTasks)
+        .addAggregation(Aggregation.count().as("total_count"))
+        .build();
+
+    // Executing aggregation query
+    AggregationResult tasksCountLatest = Iterables.getOnlyElement(
+        datastore.runAggregation(allTasksCountQuery));
+    System.out.printf("Latest tasks count is %d",
+        tasksCountLatest.get("total_count")); // 3
+
+    // Executing aggregation query with past timestamp
+    AggregationResult tasksCountInPast = Iterables.getOnlyElement(
+        datastore.runAggregation(allTasksCountQuery, ReadOption.readTime(pastTimestamp)));
+    System.out.printf("Stale tasks count is %d",
+        tasksCountInPast.get("total_count")); // 2
+
+    // [END datastore_count_aggregation_query_stale_read]
+  }
+
+}
diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
index e35d8ad16..a2fc6d815 100644
--- a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
@@ -16,6 +16,13 @@
 
 package com.example.datastore;
 
+import com.example.datastore.aggregation.CountAggregationInTransaction;
+import com.example.datastore.aggregation.CountAggregationOnKind;
+import com.example.datastore.aggregation.CountAggregationWithGqlQuery;
+import com.example.datastore.aggregation.CountAggregationWithLimit;
+import com.example.datastore.aggregation.CountAggregationWithOrderBy;
+import com.example.datastore.aggregation.CountAggregationWithPropertyFilter;
+import com.example.datastore.aggregation.CountAggregationWithStaleRead;
 import com.google.cloud.datastore.Datastore;
 import com.google.cloud.datastore.DatastoreOptions;
 import com.google.cloud.datastore.Key;
@@ -31,13 +38,11 @@
 
 public class AggregationQuerySampleTestIT {
 
-  private Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
+  private final Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
   @Rule
   public final SystemsOutRule systemsOutRule = new SystemsOutRule();
 
-  private final AggregationQuerySample sample = new AggregationQuerySample();
-
   @After
   public void tearDown() throws Exception {
     KeyQuery allKeysQuery = Query.newKeyQueryBuilder()
@@ -50,7 +55,7 @@ public void tearDown() throws Exception {
 
   @Test
   public void testAggregationQueryAndCountAggregationSample() {
-    sample.aggregationQueryAndCountAggregationOnKind();
+    CountAggregationOnKind.invoke();
 
     systemsOutRule.assertContains("Total tasks count is 3");
     systemsOutRule.assertContains("Total tasks (accessible from default alias) is 3");
@@ -58,21 +63,21 @@ public void testAggregationQueryAndCountAggregationSample() {
 
   @Test
   public void testAggregationQueryAndCountAggregationWithLimitSample() {
-    sample.aggregationQueryAndCountAggregationWithLimit();
+    CountAggregationWithLimit.invoke();
 
     systemsOutRule.assertContains("We have at least 2 tasks");
   }
 
   @Test
   public void testAggregationQueryAndCountAggregationWithOrderBySample() {
-    sample.aggregationQueryAndCountAggregationWithOrderBy();
+    CountAggregationWithOrderBy.invoke();
 
     systemsOutRule.assertContains("Total 2 tasks found with priority field");
   }
 
   @Test
   public void testAggregationQueryAndCountAggregationWithPropertyFilterSample() {
-    sample.aggregationQueryAndCountAggregationWithPropertyFilter();
+    CountAggregationWithPropertyFilter.invoke();
 
     systemsOutRule.assertContains("Total completed tasks count is 2");
     systemsOutRule.assertContains("Total remaining tasks count is 1");
@@ -80,7 +85,7 @@ public void testAggregationQueryAndCountAggregationWithPropertyFilterSample() {
 
   @Test
   public void testAggregationQueryAndCountAggregationSampleWithGqlQuery() {
-    sample.aggregationQueryAndCountAggregationWithGqlQuery();
+    CountAggregationWithGqlQuery.invoke();
 
     systemsOutRule.assertContains("We have at least 2 tasks");
     systemsOutRule.assertContains("Total tasks count is 3");
@@ -89,7 +94,7 @@ public void testAggregationQueryAndCountAggregationSampleWithGqlQuery() {
 
   @Test
   public void testAggregationQueryAndCountWithStaleRead() throws InterruptedException {
-    sample.aggregationQueryAndCountAggregationWithStaleRead();
+    CountAggregationWithStaleRead.invoke();
 
     systemsOutRule.assertContains("Latest tasks count is 3");
     systemsOutRule.assertContains("Stale tasks count is 2");
@@ -97,7 +102,7 @@ public void testAggregationQueryAndCountWithStaleRead() throws InterruptedExcept
 
   @Test
   public void testAggregationQueryAndCountWithTransaction() throws InterruptedException {
-    Assert.assertThrows(Exception.class, sample::aggregationQueryInTransaction);
+    Assert.assertThrows(Exception.class, CountAggregationInTransaction::invoke);
 
     systemsOutRule.assertContains("Found existing 2 tasks, rolling back");
   }

From 6e6a4230cbfe34974e83428e2d0dbd6f0c3507fa Mon Sep 17 00:00:00 2001
From: Owl Bot 
Date: Wed, 19 Oct 2022 05:45:01 +0000
Subject: [PATCH 67/82] =?UTF-8?q?=F0=9F=A6=89=20Updates=20from=20OwlBot=20?=
 =?UTF-8?q?post-processor?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md
---
 README.md                                     |  7 +++
 .../CountAggregationInTransaction.java        | 54 ++++++++++---------
 .../aggregation/CountAggregationOnKind.java   | 28 +++++-----
 .../CountAggregationWithGqlQuery.java         | 53 +++++++++---------
 .../CountAggregationWithLimit.java            | 21 ++++----
 .../CountAggregationWithOrderBy.java          | 33 ++++++------
 .../CountAggregationWithPropertyFilter.java   | 53 +++++++++---------
 .../CountAggregationWithStaleRead.java        | 35 ++++++------
 .../AggregationQuerySampleTestIT.java         |  9 ++--
 .../example/datastore/QuickstartSampleIT.java |  3 +-
 10 files changed, 149 insertions(+), 147 deletions(-)

diff --git a/README.md b/README.md
index 3516bc787..2ae734023 100644
--- a/README.md
+++ b/README.md
@@ -234,6 +234,13 @@ Samples are in the [`samples/`](https://github.com/googleapis/java-datastore/tre
 | --------------------------- | --------------------------------- | ------ |
 | Native Image Datastore Sample | [source code](https://github.com/googleapis/java-datastore/blob/main/samples/native-image-sample/src/main/java/com/example/datastore/NativeImageDatastoreSample.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-datastore&page=editor&open_in_editor=samples/native-image-sample/src/main/java/com/example/datastore/NativeImageDatastoreSample.java) |
 | Quickstart Sample | [source code](https://github.com/googleapis/java-datastore/blob/main/samples/snippets/src/main/java/com/example/datastore/QuickstartSample.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-datastore&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/datastore/QuickstartSample.java) |
+| Count Aggregation In Transaction | [source code](https://github.com/googleapis/java-datastore/blob/main/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-datastore&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java) |
+| Count Aggregation On Kind | [source code](https://github.com/googleapis/java-datastore/blob/main/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-datastore&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java) |
+| Count Aggregation With Gql Query | [source code](https://github.com/googleapis/java-datastore/blob/main/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-datastore&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java) |
+| Count Aggregation With Limit | [source code](https://github.com/googleapis/java-datastore/blob/main/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-datastore&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java) |
+| Count Aggregation With Order By | [source code](https://github.com/googleapis/java-datastore/blob/main/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-datastore&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java) |
+| Count Aggregation With Property Filter | [source code](https://github.com/googleapis/java-datastore/blob/main/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-datastore&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java) |
+| Count Aggregation With Stale Read | [source code](https://github.com/googleapis/java-datastore/blob/main/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-datastore&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java) |
 | Task List | [source code](https://github.com/googleapis/java-datastore/blob/main/samples/snippets/src/main/java/com/google/datastore/snippets/TaskList.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-datastore&page=editor&open_in_editor=samples/snippets/src/main/java/com/google/datastore/snippets/TaskList.java) |
 
 
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
index cbe3f4dbe..ad9bfb2cf 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
@@ -29,33 +29,37 @@ public static void invoke() {
     // Save all the tasks
     datastore.put(
         Entity.newBuilder(task1Key).set("owner", "john").build(),
-        Entity.newBuilder(task2Key).set("owner", "john").build()
-    );
+        Entity.newBuilder(task2Key).set("owner", "john").build());
 
     // Using transactions to maintain consistent application state.
-    datastore.runInTransaction((TransactionCallable) transaction -> {
-      EntityQuery tasksOfJohn = Query.newEntityQueryBuilder()
-          .setKind(kind)
-          .setFilter(PropertyFilter.eq("owner", "john"))
-          .build();
-      AggregationQuery totalTasksQuery = Query.newAggregationQueryBuilder()
-          .over(tasksOfJohn)
-          .addAggregation(Aggregation.count().as("tasks_count"))
-          .build();
-
-      Long tasksCount = Iterables.getOnlyElement(
-          datastore.runAggregation(totalTasksQuery)).get("tasks_count");
-
-      if (tasksCount < 2) {
-        Key newTaskKey = datastore.newKeyFactory().setKind(kind).newKey("task3");
-        Entity newTask = Entity.newBuilder(newTaskKey).set("owner", "john").build();
-        transaction.put(newTask);
-      } else {
-        System.out.printf("Found existing %d tasks, rolling back", tasksCount);
-        throw new Exception("User 'John' cannot have more than 2 tasks");
-      }
-      return null;
-    });
+    datastore.runInTransaction(
+        (TransactionCallable)
+            transaction -> {
+              EntityQuery tasksOfJohn =
+                  Query.newEntityQueryBuilder()
+                      .setKind(kind)
+                      .setFilter(PropertyFilter.eq("owner", "john"))
+                      .build();
+              AggregationQuery totalTasksQuery =
+                  Query.newAggregationQueryBuilder()
+                      .over(tasksOfJohn)
+                      .addAggregation(Aggregation.count().as("tasks_count"))
+                      .build();
+
+              Long tasksCount =
+                  Iterables.getOnlyElement(datastore.runAggregation(totalTasksQuery))
+                      .get("tasks_count");
+
+              if (tasksCount < 2) {
+                Key newTaskKey = datastore.newKeyFactory().setKind(kind).newKey("task3");
+                Entity newTask = Entity.newBuilder(newTaskKey).set("owner", "john").build();
+                transaction.put(newTask);
+              } else {
+                System.out.printf("Found existing %d tasks, rolling back", tasksCount);
+                throw new Exception("User 'John' cannot have more than 2 tasks");
+              }
+              return null;
+            });
     // [END datastore_count_aggregation_query_in_transaction]
 
   }
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
index a047c07b1..2540ed187 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
@@ -45,25 +45,23 @@ public static void invoke() {
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).build(),
         Entity.newBuilder(task2Key).set("done", false).build(),
-        Entity.newBuilder(task3Key).set("done", true).build()
-    );
+        Entity.newBuilder(task3Key).set("done", true).build());
 
-    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .build();
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder().setKind(kind).build();
     // Creating an aggregation query to get the count of all tasks
-    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllTasks)
-        .addAggregation(Aggregation.count())
-        .addAggregation(Aggregation.count().as("total_count"))
-        .build();
+    AggregationQuery allTasksCountQuery =
+        Query.newAggregationQueryBuilder()
+            .over(selectAllTasks)
+            .addAggregation(Aggregation.count())
+            .addAggregation(Aggregation.count().as("total_count"))
+            .build();
     // Executing aggregation query
-    AggregationResult allTasksCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(allTasksCountQuery));
+    AggregationResult allTasksCountQueryResult =
+        Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
 
-    System.out.printf("Total tasks count is %d",
-        allTasksCountQueryResult.get("total_count")); // 3
-    System.out.printf("Total tasks (accessible from default alias) is %d",
+    System.out.printf("Total tasks count is %d", allTasksCountQueryResult.get("total_count")); // 3
+    System.out.printf(
+        "Total tasks (accessible from default alias) is %d",
         allTasksCountQueryResult.get("property_1")); // 3
 
     // [END datastore_count_aggregation_query_on_kind]
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
index 4eeec810c..9a3c74802 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
@@ -29,42 +29,41 @@ public static void invoke() {
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).build(),
         Entity.newBuilder(task2Key).set("done", false).build(),
-        Entity.newBuilder(task3Key).set("done", true).build()
-    );
+        Entity.newBuilder(task3Key).set("done", true).build());
 
-    GqlQuery selectAllTasks = Query.newGqlQueryBuilder(
-            "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(2) AS count_with_limit "
-                + "OVER (SELECT * FROM Task)")
-        .setAllowLiteral(true)
-        .build();
+    GqlQuery selectAllTasks =
+        Query.newGqlQueryBuilder(
+                "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(2) AS count_with_limit "
+                    + "OVER (SELECT * FROM Task)")
+            .setAllowLiteral(true)
+            .build();
     // Creating an aggregation query to get the count of all tasks
-    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllTasks)
-        .build();
+    AggregationQuery allTasksCountQuery =
+        Query.newAggregationQueryBuilder().over(selectAllTasks).build();
     // Executing aggregation query
-    AggregationResult allTasksCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(allTasksCountQuery));
+    AggregationResult allTasksCountQueryResult =
+        Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
 
-    System.out.printf("We have at least %d tasks",
-        allTasksCountQueryResult.get("count_with_limit")); // 2
-    System.out.printf("Total tasks count is %d",
-        allTasksCountQueryResult.get("total_count")); // 3
+    System.out.printf(
+        "We have at least %d tasks", allTasksCountQueryResult.get("count_with_limit")); // 2
+    System.out.printf("Total tasks count is %d", allTasksCountQueryResult.get("total_count")); // 3
 
-    GqlQuery completedTasks = Query.newGqlQueryBuilder(
-            "AGGREGATE COUNT(*) AS total_completed_count "
-                + "OVER (SELECT * FROM Task WHERE done = true)")
-        .setAllowLiteral(true)
-        .build();
+    GqlQuery completedTasks =
+        Query.newGqlQueryBuilder(
+                "AGGREGATE COUNT(*) AS total_completed_count "
+                    + "OVER (SELECT * FROM Task WHERE done = true)")
+            .setAllowLiteral(true)
+            .build();
     // Creating an aggregation query to get the count of all completed tasks
-    AggregationQuery completedTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(completedTasks)
-        .build();
+    AggregationQuery completedTasksCountQuery =
+        Query.newAggregationQueryBuilder().over(completedTasks).build();
 
     // Executing aggregation query
-    AggregationResult completedTasksCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(completedTasksCountQuery));
+    AggregationResult completedTasksCountQueryResult =
+        Iterables.getOnlyElement(datastore.runAggregation(completedTasksCountQuery));
 
-    System.out.printf("Total completed tasks count is %d",
+    System.out.printf(
+        "Total completed tasks count is %d",
         completedTasksCountQueryResult.get("total_completed_count")); // 2
 
     // [END datastore_count_aggregation_query_gql]
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
index 1813f4109..44cdcb1c7 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
@@ -29,21 +29,18 @@ public static void invoke() {
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).build(),
         Entity.newBuilder(task2Key).set("done", false).build(),
-        Entity.newBuilder(task3Key).set("done", true).build()
-    );
+        Entity.newBuilder(task3Key).set("done", true).build());
 
-    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .setLimit(2)
-        .build();
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder().setKind(kind).setLimit(2).build();
     // Creating an aggregation query to get the count of all tasks
-    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllTasks)
-        .addAggregation(Aggregation.count().as("at_least"))
-        .build();
+    AggregationQuery allTasksCountQuery =
+        Query.newAggregationQueryBuilder()
+            .over(selectAllTasks)
+            .addAggregation(Aggregation.count().as("at_least"))
+            .build();
     // Executing aggregation query
-    AggregationResult limitQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(allTasksCountQuery));
+    AggregationResult limitQueryResult =
+        Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
 
     System.out.printf("We have at least %d tasks", limitQueryResult.get("at_least")); // 2
 
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
index 4baa74672..fdb084d39 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
@@ -29,25 +29,26 @@ public static void invoke() {
     // Save all the tasks
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).set("priority", 1).build(),
-        Entity.newBuilder(task2Key).set("done", false).build(),  // no priority specified
-        Entity.newBuilder(task3Key).set("done", true).set("priority", 2).build()
-    );
-
-    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .addOrderBy(OrderBy.asc("priority"))    // OrderBy acts as an existence filter
-        .build();
+        Entity.newBuilder(task2Key).set("done", false).build(), // no priority specified
+        Entity.newBuilder(task3Key).set("done", true).set("priority", 2).build());
+
+    EntityQuery selectAllTasks =
+        Query.newEntityQueryBuilder()
+            .setKind(kind)
+            .addOrderBy(OrderBy.asc("priority")) // OrderBy acts as an existence filter
+            .build();
     // Creating an aggregation query to get the count of all tasks
-    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllTasks)
-        .addAggregation(Aggregation.count().as("count"))
-        .build();
+    AggregationQuery allTasksCountQuery =
+        Query.newAggregationQueryBuilder()
+            .over(selectAllTasks)
+            .addAggregation(Aggregation.count().as("count"))
+            .build();
     // Executing aggregation query
-    AggregationResult limitQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(allTasksCountQuery));
+    AggregationResult limitQueryResult =
+        Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
 
-    System.out.printf("Total %d tasks found with priority field",
-        limitQueryResult.get("count")); // 2
+    System.out.printf(
+        "Total %d tasks found with priority field", limitQueryResult.get("count")); // 2
 
     // [END datastore_count_aggregation_query_with_order_by]
   }
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
index 01c4897a3..9263557b5 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
@@ -31,37 +31,42 @@ public static void invoke() {
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).build(),
         Entity.newBuilder(task2Key).set("done", false).build(),
-        Entity.newBuilder(task3Key).set("done", true).build()
-    );
+        Entity.newBuilder(task3Key).set("done", true).build());
 
-    EntityQuery completedTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .setFilter(PropertyFilter.eq("done", true))
-        .build();
-    EntityQuery remainingTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .setFilter(PropertyFilter.eq("done", false))
-        .build();
+    EntityQuery completedTasks =
+        Query.newEntityQueryBuilder()
+            .setKind(kind)
+            .setFilter(PropertyFilter.eq("done", true))
+            .build();
+    EntityQuery remainingTasks =
+        Query.newEntityQueryBuilder()
+            .setKind(kind)
+            .setFilter(PropertyFilter.eq("done", false))
+            .build();
     // Creating an aggregation query to get the count of all completed tasks
-    AggregationQuery completedTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(completedTasks)
-        .addAggregation(Aggregation.count().as("total_completed_count"))
-        .build();
+    AggregationQuery completedTasksCountQuery =
+        Query.newAggregationQueryBuilder()
+            .over(completedTasks)
+            .addAggregation(Aggregation.count().as("total_completed_count"))
+            .build();
     // Creating an aggregation query to get the count of all remaining tasks
-    AggregationQuery remainingTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(remainingTasks)
-        .addAggregation(Aggregation.count().as("total_remaining_count"))
-        .build();
+    AggregationQuery remainingTasksCountQuery =
+        Query.newAggregationQueryBuilder()
+            .over(remainingTasks)
+            .addAggregation(Aggregation.count().as("total_remaining_count"))
+            .build();
 
     // Executing aggregation query
-    AggregationResult completedTasksCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(completedTasksCountQuery));
-    AggregationResult remainingTasksCountQueryResult = Iterables.getOnlyElement(
-        datastore.runAggregation(remainingTasksCountQuery));
+    AggregationResult completedTasksCountQueryResult =
+        Iterables.getOnlyElement(datastore.runAggregation(completedTasksCountQuery));
+    AggregationResult remainingTasksCountQueryResult =
+        Iterables.getOnlyElement(datastore.runAggregation(remainingTasksCountQuery));
 
-    System.out.printf("Total completed tasks count is %d",
+    System.out.printf(
+        "Total completed tasks count is %d",
         completedTasksCountQueryResult.get("total_completed_count")); // 2
-    System.out.printf("Total remaining tasks count is %d",
+    System.out.printf(
+        "Total remaining tasks count is %d",
         remainingTasksCountQueryResult.get("total_remaining_count")); // 1
 
     // [END datastore_count_aggregation_query_with_filters]
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
index b558ad4f8..079b360ff 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
@@ -30,40 +30,35 @@ public static void invoke() throws InterruptedException {
     // Saving only two tasks
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).build(),
-        Entity.newBuilder(task2Key).set("done", false).build()
-    );
+        Entity.newBuilder(task2Key).set("done", false).build());
     Thread.sleep(1000);
-    final Timestamp pastTimestamp =
-        Timestamp.now(); // we have two tasks in database at this time.
+    final Timestamp pastTimestamp = Timestamp.now(); // we have two tasks in database at this time.
 
     Thread.sleep(1000);
     // Saving third tasks
     Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
     datastore.put(Entity.newBuilder(task3Key).set("done", false).build());
 
-    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .build();
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder().setKind(kind).build();
 
     // Creating an aggregation query to get the count of all tasks
-    AggregationQuery allTasksCountQuery = Query.newAggregationQueryBuilder()
-        .over(selectAllTasks)
-        .addAggregation(Aggregation.count().as("total_count"))
-        .build();
+    AggregationQuery allTasksCountQuery =
+        Query.newAggregationQueryBuilder()
+            .over(selectAllTasks)
+            .addAggregation(Aggregation.count().as("total_count"))
+            .build();
 
     // Executing aggregation query
-    AggregationResult tasksCountLatest = Iterables.getOnlyElement(
-        datastore.runAggregation(allTasksCountQuery));
-    System.out.printf("Latest tasks count is %d",
-        tasksCountLatest.get("total_count")); // 3
+    AggregationResult tasksCountLatest =
+        Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
+    System.out.printf("Latest tasks count is %d", tasksCountLatest.get("total_count")); // 3
 
     // Executing aggregation query with past timestamp
-    AggregationResult tasksCountInPast = Iterables.getOnlyElement(
-        datastore.runAggregation(allTasksCountQuery, ReadOption.readTime(pastTimestamp)));
-    System.out.printf("Stale tasks count is %d",
-        tasksCountInPast.get("total_count")); // 2
+    AggregationResult tasksCountInPast =
+        Iterables.getOnlyElement(
+            datastore.runAggregation(allTasksCountQuery, ReadOption.readTime(pastTimestamp)));
+    System.out.printf("Stale tasks count is %d", tasksCountInPast.get("total_count")); // 2
 
     // [END datastore_count_aggregation_query_stale_read]
   }
-
 }
diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
index a2fc6d815..0ff3f526a 100644
--- a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
@@ -40,14 +40,11 @@ public class AggregationQuerySampleTestIT {
 
   private final Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
-  @Rule
-  public final SystemsOutRule systemsOutRule = new SystemsOutRule();
+  @Rule public final SystemsOutRule systemsOutRule = new SystemsOutRule();
 
   @After
   public void tearDown() throws Exception {
-    KeyQuery allKeysQuery = Query.newKeyQueryBuilder()
-        .setKind("Task")
-        .build();
+    KeyQuery allKeysQuery = Query.newKeyQueryBuilder().setKind("Task").build();
     QueryResults allKeys = datastore.run(allKeysQuery);
     Key[] keysToDelete = ImmutableList.copyOf(allKeys).toArray(new Key[0]);
     datastore.delete(keysToDelete);
@@ -106,4 +103,4 @@ public void testAggregationQueryAndCountWithTransaction() throws InterruptedExce
 
     systemsOutRule.assertContains("Found existing 2 tasks, rolling back");
   }
-}
\ No newline at end of file
+}
diff --git a/samples/snippets/src/test/java/com/example/datastore/QuickstartSampleIT.java b/samples/snippets/src/test/java/com/example/datastore/QuickstartSampleIT.java
index 35084937f..09823479f 100644
--- a/samples/snippets/src/test/java/com/example/datastore/QuickstartSampleIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/QuickstartSampleIT.java
@@ -32,8 +32,7 @@
 @SuppressWarnings("checkstyle:abbreviationaswordinname")
 public class QuickstartSampleIT {
 
-  @Rule
-  public final SystemsOutRule systemsOutRule = new SystemsOutRule();
+  @Rule public final SystemsOutRule systemsOutRule = new SystemsOutRule();
 
   private static final void deleteTestEntity() {
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();

From 54ace277aefcbc36f8fd802198f014c8e72d6a61 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Wed, 19 Oct 2022 11:29:12 +0530
Subject: [PATCH 68/82] Adding missing header

---
 .../CountAggregationInTransaction.java           | 16 ++++++++++++++++
 .../aggregation/CountAggregationOnKind.java      |  3 ++-
 .../CountAggregationWithGqlQuery.java            | 16 ++++++++++++++++
 .../aggregation/CountAggregationWithLimit.java   | 16 ++++++++++++++++
 .../aggregation/CountAggregationWithOrderBy.java | 16 ++++++++++++++++
 .../CountAggregationWithPropertyFilter.java      | 16 ++++++++++++++++
 .../CountAggregationWithStaleRead.java           | 16 ++++++++++++++++
 7 files changed, 98 insertions(+), 1 deletion(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
index ad9bfb2cf..43cf96232 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.example.datastore.aggregation;
 
 import com.google.cloud.datastore.AggregationQuery;
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
index 2540ed187..364a10917 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
@@ -5,7 +5,7 @@
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
- *     https://www.apache.org/licenses/LICENSE-2.0
+ * http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
 package com.example.datastore.aggregation;
 
 import com.google.cloud.datastore.AggregationQuery;
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
index 9a3c74802..096355e28 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.example.datastore.aggregation;
 
 import com.google.cloud.datastore.AggregationQuery;
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
index 44cdcb1c7..3a4d4026e 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.example.datastore.aggregation;
 
 import com.google.cloud.datastore.AggregationQuery;
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
index fdb084d39..c4fd89668 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.example.datastore.aggregation;
 
 import com.google.cloud.datastore.AggregationQuery;
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
index 9263557b5..d1e6bfccd 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.example.datastore.aggregation;
 
 import com.google.cloud.datastore.AggregationQuery;
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
index 079b360ff..ce4c85a86 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.example.datastore.aggregation;
 
 import com.google.cloud.Timestamp;

From 170590bd44bf012db631ce36399b546bfd0e92ae Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Wed, 19 Oct 2022 12:50:59 +0530
Subject: [PATCH 69/82] Sticking to the practise of deleting the old created
 data in setUp of test

---
 .../com/example/datastore/AggregationQuerySampleTestIT.java | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
index 0ff3f526a..2f7225163 100644
--- a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
@@ -31,8 +31,8 @@
 import com.google.cloud.datastore.QueryResults;
 import com.google.common.collect.ImmutableList;
 import com.rule.SystemsOutRule;
-import org.junit.After;
 import org.junit.Assert;
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 
@@ -42,8 +42,8 @@ public class AggregationQuerySampleTestIT {
 
   @Rule public final SystemsOutRule systemsOutRule = new SystemsOutRule();
 
-  @After
-  public void tearDown() throws Exception {
+  @Before
+  public void setUp() throws Exception {
     KeyQuery allKeysQuery = Query.newKeyQueryBuilder().setKind("Task").build();
     QueryResults allKeys = datastore.run(allKeysQuery);
     Key[] keysToDelete = ImmutableList.copyOf(allKeys).toArray(new Key[0]);

From 670face736f9f647391c10dc71aca6c71f1d2126 Mon Sep 17 00:00:00 2001
From: Prateek 
Date: Wed, 19 Oct 2022 17:05:30 +0530
Subject: [PATCH 70/82] Moving comment to the top of statement.

Co-authored-by: Sita Lakshmi Sangameswaran 
---
 .../datastore/aggregation/CountAggregationWithOrderBy.java     | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
index c4fd89668..6ddae927e 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
@@ -45,7 +45,8 @@ public static void invoke() {
     // Save all the tasks
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).set("priority", 1).build(),
-        Entity.newBuilder(task2Key).set("done", false).build(), // no priority specified
+        // Priority not specified.
+        Entity.newBuilder(task2Key).set("done", false).build(),
         Entity.newBuilder(task3Key).set("done", true).set("priority", 2).build());
 
     EntityQuery selectAllTasks =

From 67b260208591ac559ea76095ae5402c33abad269 Mon Sep 17 00:00:00 2001
From: Prateek 
Date: Wed, 19 Oct 2022 17:06:35 +0530
Subject: [PATCH 71/82] Making builder syntax multiline

Co-authored-by: Sita Lakshmi Sangameswaran 
---
 .../datastore/aggregation/CountAggregationWithLimit.java     | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
index 3a4d4026e..6fd6bb582 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
@@ -47,7 +47,10 @@ public static void invoke() {
         Entity.newBuilder(task2Key).set("done", false).build(),
         Entity.newBuilder(task3Key).set("done", true).build());
 
-    EntityQuery selectAllTasks = Query.newEntityQueryBuilder().setKind(kind).setLimit(2).build();
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
+        .setKind(kind)
+        .setLimit(2)
+        .build();
     // Creating an aggregation query to get the count of all tasks
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder()

From ebc998d9982cd60f7830a6494ba2861740d5f728 Mon Sep 17 00:00:00 2001
From: Prateek 
Date: Wed, 19 Oct 2022 17:07:21 +0530
Subject: [PATCH 72/82] Adding comment describing the purpose of GQL query

Co-authored-by: Sita Lakshmi Sangameswaran 
---
 .../datastore/aggregation/CountAggregationWithGqlQuery.java      | 1 +
 1 file changed, 1 insertion(+)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
index 096355e28..5f19a74df 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
@@ -47,6 +47,7 @@ public static void invoke() {
         Entity.newBuilder(task2Key).set("done", false).build(),
         Entity.newBuilder(task3Key).set("done", true).build());
 
+    // Create a GQL query to get the count of all tasks.
     GqlQuery selectAllTasks =
         Query.newGqlQueryBuilder(
                 "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(2) AS count_with_limit "

From 9988599b3ba8dc180308098752c1437510c4c5ca Mon Sep 17 00:00:00 2001
From: Prateek 
Date: Wed, 19 Oct 2022 17:07:37 +0530
Subject: [PATCH 73/82] Rephrasing the comment

Co-authored-by: Sita Lakshmi Sangameswaran 
---
 .../datastore/aggregation/CountAggregationWithGqlQuery.java     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
index 5f19a74df..b3b9907ee 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
@@ -54,7 +54,7 @@ public static void invoke() {
                     + "OVER (SELECT * FROM Task)")
             .setAllowLiteral(true)
             .build();
-    // Creating an aggregation query to get the count of all tasks
+    // Create the aggregation query builder and set the query.
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder().over(selectAllTasks).build();
     // Executing aggregation query

From c3fa6d7adbf65dd82c54609281ad23dc90cf7a25 Mon Sep 17 00:00:00 2001
From: Prateek 
Date: Wed, 19 Oct 2022 17:08:19 +0530
Subject: [PATCH 74/82] Adding a comment describing the purpose of the GQL
 query

Co-authored-by: Sita Lakshmi Sangameswaran 
---
 .../datastore/aggregation/CountAggregationWithGqlQuery.java      | 1 +
 1 file changed, 1 insertion(+)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
index b3b9907ee..c99c12a0e 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
@@ -65,6 +65,7 @@ public static void invoke() {
         "We have at least %d tasks", allTasksCountQueryResult.get("count_with_limit")); // 2
     System.out.printf("Total tasks count is %d", allTasksCountQueryResult.get("total_count")); // 3
 
+    // Create a query to get the count of all completed tasks.
     GqlQuery completedTasks =
         Query.newGqlQueryBuilder(
                 "AGGREGATE COUNT(*) AS total_completed_count "

From 11a42309cecdd2458eb5f5e5a30974eb022f9dd6 Mon Sep 17 00:00:00 2001
From: Prateek 
Date: Wed, 19 Oct 2022 17:08:31 +0530
Subject: [PATCH 75/82] Rephrasing the comment

Co-authored-by: Sita Lakshmi Sangameswaran 
---
 .../datastore/aggregation/CountAggregationWithGqlQuery.java     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
index c99c12a0e..0d016e2ba 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
@@ -72,7 +72,7 @@ public static void invoke() {
                     + "OVER (SELECT * FROM Task WHERE done = true)")
             .setAllowLiteral(true)
             .build();
-    // Creating an aggregation query to get the count of all completed tasks
+    // Create the aggregation query builder and set the query.
     AggregationQuery completedTasksCountQuery =
         Query.newAggregationQueryBuilder().over(completedTasks).build();
 

From 7189f80936c2aaac76204c3848646cc38872dd28 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Wed, 19 Oct 2022 16:18:17 +0530
Subject: [PATCH 76/82] Relocating region tags to include import statments

---
 .../aggregation/CountAggregationInTransaction.java          | 6 ++----
 .../datastore/aggregation/CountAggregationOnKind.java       | 6 ++----
 .../datastore/aggregation/CountAggregationWithGqlQuery.java | 6 ++----
 .../datastore/aggregation/CountAggregationWithLimit.java    | 6 ++----
 .../datastore/aggregation/CountAggregationWithOrderBy.java  | 6 ++----
 .../aggregation/CountAggregationWithPropertyFilter.java     | 6 ++----
 .../aggregation/CountAggregationWithStaleRead.java          | 6 ++----
 7 files changed, 14 insertions(+), 28 deletions(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
index 43cf96232..4cd484d5f 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
@@ -16,6 +16,7 @@
 
 package com.example.datastore.aggregation;
 
+// [START datastore_count_aggregation_query_in_transaction]
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.Datastore;
 import com.google.cloud.datastore.Datastore.TransactionCallable;
@@ -31,8 +32,6 @@
 public class CountAggregationInTransaction {
 
   public static void invoke() {
-    // [START datastore_count_aggregation_query_in_transaction]
-
     // Instantiates a client
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
@@ -76,7 +75,6 @@ public static void invoke() {
               }
               return null;
             });
-    // [END datastore_count_aggregation_query_in_transaction]
-
   }
 }
+// [END datastore_count_aggregation_query_in_transaction]
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
index 364a10917..5c84d11e3 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
@@ -16,6 +16,7 @@
 
 package com.example.datastore.aggregation;
 
+// [START datastore_count_aggregation_query_on_kind]
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.Datastore;
@@ -30,8 +31,6 @@
 public class CountAggregationOnKind {
 
   public static void invoke() {
-    // [START datastore_count_aggregation_query_on_kind]
-
     // Instantiates a client
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
@@ -64,7 +63,6 @@ public static void invoke() {
     System.out.printf(
         "Total tasks (accessible from default alias) is %d",
         allTasksCountQueryResult.get("property_1")); // 3
-
-    // [END datastore_count_aggregation_query_on_kind]
   }
 }
+// [END datastore_count_aggregation_query_on_kind]
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
index 0d016e2ba..138b56bee 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
@@ -16,6 +16,7 @@
 
 package com.example.datastore.aggregation;
 
+// [START datastore_count_aggregation_query_gql]
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.Datastore;
@@ -29,8 +30,6 @@
 public class CountAggregationWithGqlQuery {
 
   public static void invoke() {
-    // [START datastore_count_aggregation_query_gql]
-
     // Instantiates a client
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
@@ -83,7 +82,6 @@ public static void invoke() {
     System.out.printf(
         "Total completed tasks count is %d",
         completedTasksCountQueryResult.get("total_completed_count")); // 2
-
-    // [END datastore_count_aggregation_query_gql]
   }
 }
+// [END datastore_count_aggregation_query_gql]
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
index 6fd6bb582..9305aacff 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
@@ -16,6 +16,7 @@
 
 package com.example.datastore.aggregation;
 
+// [START datastore_count_aggregation_query_with_limit]
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.Datastore;
@@ -29,8 +30,6 @@
 
 public class CountAggregationWithLimit {
   public static void invoke() {
-    // [START datastore_count_aggregation_query_with_limit]
-
     // Instantiates a client
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
@@ -62,7 +61,6 @@ public static void invoke() {
         Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
 
     System.out.printf("We have at least %d tasks", limitQueryResult.get("at_least")); // 2
-
-    // [END datastore_count_aggregation_query_with_limit]
   }
 }
+// [END datastore_count_aggregation_query_with_limit]
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
index 6ddae927e..ca5ff920e 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
@@ -16,6 +16,7 @@
 
 package com.example.datastore.aggregation;
 
+// [START datastore_count_aggregation_query_with_order_by]
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.Datastore;
@@ -30,8 +31,6 @@
 
 public class CountAggregationWithOrderBy {
   public static void invoke() {
-    // [START datastore_count_aggregation_query_with_order_by]
-
     // Instantiates a client
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
@@ -66,7 +65,6 @@ public static void invoke() {
 
     System.out.printf(
         "Total %d tasks found with priority field", limitQueryResult.get("count")); // 2
-
-    // [END datastore_count_aggregation_query_with_order_by]
   }
 }
+// [END datastore_count_aggregation_query_with_order_by]
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
index d1e6bfccd..d45db430d 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
@@ -16,6 +16,7 @@
 
 package com.example.datastore.aggregation;
 
+// [START datastore_count_aggregation_query_with_filters]
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.Datastore;
@@ -31,8 +32,6 @@
 public class CountAggregationWithPropertyFilter {
 
   public static void invoke() {
-    // [START datastore_count_aggregation_query_with_filters]
-
     // Instantiates a client
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
@@ -84,7 +83,6 @@ public static void invoke() {
     System.out.printf(
         "Total remaining tasks count is %d",
         remainingTasksCountQueryResult.get("total_remaining_count")); // 1
-
-    // [END datastore_count_aggregation_query_with_filters]
   }
 }
+// [END datastore_count_aggregation_query_with_filters]
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
index ce4c85a86..77d3d0efd 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
@@ -16,6 +16,7 @@
 
 package com.example.datastore.aggregation;
 
+// [START datastore_count_aggregation_query_stale_read]
 import com.google.cloud.Timestamp;
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.AggregationResult;
@@ -32,8 +33,6 @@
 public class CountAggregationWithStaleRead {
 
   public static void invoke() throws InterruptedException {
-    // [START datastore_count_aggregation_query_stale_read]
-
     // Instantiates a client
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
@@ -74,7 +73,6 @@ public static void invoke() throws InterruptedException {
         Iterables.getOnlyElement(
             datastore.runAggregation(allTasksCountQuery, ReadOption.readTime(pastTimestamp)));
     System.out.printf("Stale tasks count is %d", tasksCountInPast.get("total_count")); // 2
-
-    // [END datastore_count_aggregation_query_stale_read]
   }
 }
+// [END datastore_count_aggregation_query_stale_read]
\ No newline at end of file

From 350ab615f353a32d3a7a12648f66d9dc91ed0ffa Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Wed, 19 Oct 2022 17:03:39 +0530
Subject: [PATCH 77/82] Few styling and comment fixes

---
 .../aggregation/CountAggregationWithOrderBy.java    |  1 +
 .../datastore/AggregationQuerySampleTestIT.java     | 13 ++++---------
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
index ca5ff920e..ce2a2bfbe 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
@@ -48,6 +48,7 @@ public static void invoke() {
         Entity.newBuilder(task2Key).set("done", false).build(),
         Entity.newBuilder(task3Key).set("done", true).set("priority", 2).build());
 
+    // OrderBy acts as an existence filter.
     EntityQuery selectAllTasks =
         Query.newEntityQueryBuilder()
             .setKind(kind)
diff --git a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
index 2f7225163..662802e51 100644
--- a/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
+++ b/samples/snippets/src/test/java/com/example/datastore/AggregationQuerySampleTestIT.java
@@ -16,6 +16,8 @@
 
 package com.example.datastore;
 
+import static org.junit.Assert.assertThrows;
+
 import com.example.datastore.aggregation.CountAggregationInTransaction;
 import com.example.datastore.aggregation.CountAggregationOnKind;
 import com.example.datastore.aggregation.CountAggregationWithGqlQuery;
@@ -31,7 +33,6 @@
 import com.google.cloud.datastore.QueryResults;
 import com.google.common.collect.ImmutableList;
 import com.rule.SystemsOutRule;
-import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -44,6 +45,7 @@ public class AggregationQuerySampleTestIT {
 
   @Before
   public void setUp() throws Exception {
+    // Retrieving and deleting all the 'Task' entities.
     KeyQuery allKeysQuery = Query.newKeyQueryBuilder().setKind("Task").build();
     QueryResults allKeys = datastore.run(allKeysQuery);
     Key[] keysToDelete = ImmutableList.copyOf(allKeys).toArray(new Key[0]);
@@ -53,7 +55,6 @@ public void setUp() throws Exception {
   @Test
   public void testAggregationQueryAndCountAggregationSample() {
     CountAggregationOnKind.invoke();
-
     systemsOutRule.assertContains("Total tasks count is 3");
     systemsOutRule.assertContains("Total tasks (accessible from default alias) is 3");
   }
@@ -61,21 +62,18 @@ public void testAggregationQueryAndCountAggregationSample() {
   @Test
   public void testAggregationQueryAndCountAggregationWithLimitSample() {
     CountAggregationWithLimit.invoke();
-
     systemsOutRule.assertContains("We have at least 2 tasks");
   }
 
   @Test
   public void testAggregationQueryAndCountAggregationWithOrderBySample() {
     CountAggregationWithOrderBy.invoke();
-
     systemsOutRule.assertContains("Total 2 tasks found with priority field");
   }
 
   @Test
   public void testAggregationQueryAndCountAggregationWithPropertyFilterSample() {
     CountAggregationWithPropertyFilter.invoke();
-
     systemsOutRule.assertContains("Total completed tasks count is 2");
     systemsOutRule.assertContains("Total remaining tasks count is 1");
   }
@@ -83,7 +81,6 @@ public void testAggregationQueryAndCountAggregationWithPropertyFilterSample() {
   @Test
   public void testAggregationQueryAndCountAggregationSampleWithGqlQuery() {
     CountAggregationWithGqlQuery.invoke();
-
     systemsOutRule.assertContains("We have at least 2 tasks");
     systemsOutRule.assertContains("Total tasks count is 3");
     systemsOutRule.assertContains("Total completed tasks count is 2");
@@ -92,15 +89,13 @@ public void testAggregationQueryAndCountAggregationSampleWithGqlQuery() {
   @Test
   public void testAggregationQueryAndCountWithStaleRead() throws InterruptedException {
     CountAggregationWithStaleRead.invoke();
-
     systemsOutRule.assertContains("Latest tasks count is 3");
     systemsOutRule.assertContains("Stale tasks count is 2");
   }
 
   @Test
   public void testAggregationQueryAndCountWithTransaction() throws InterruptedException {
-    Assert.assertThrows(Exception.class, CountAggregationInTransaction::invoke);
-
+    assertThrows(Exception.class, CountAggregationInTransaction::invoke);
     systemsOutRule.assertContains("Found existing 2 tasks, rolling back");
   }
 }

From abd2b953ac6f8cfdbffd659117cb8d04616539bc Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Wed, 19 Oct 2022 17:25:59 +0530
Subject: [PATCH 78/82] Adding comment in the transaction block

---
 .../datastore/aggregation/CountAggregationInTransaction.java | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
index 4cd484d5f..f4215ad98 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
@@ -50,6 +50,7 @@ public static void invoke() {
     datastore.runInTransaction(
         (TransactionCallable)
             transaction -> {
+              // Create a query to get the count of all tasks of owner 'John'.
               EntityQuery tasksOfJohn =
                   Query.newEntityQueryBuilder()
                       .setKind(kind)
@@ -61,13 +62,15 @@ public static void invoke() {
                       .addAggregation(Aggregation.count().as("tasks_count"))
                       .build();
 
+              // Executing aggregation query in the ongoing transaction.
               Long tasksCount =
-                  Iterables.getOnlyElement(datastore.runAggregation(totalTasksQuery))
+                  Iterables.getOnlyElement(transaction.runAggregation(totalTasksQuery))
                       .get("tasks_count");
 
               if (tasksCount < 2) {
                 Key newTaskKey = datastore.newKeyFactory().setKind(kind).newKey("task3");
                 Entity newTask = Entity.newBuilder(newTaskKey).set("owner", "john").build();
+                // Inserting a new entity in the transaction.
                 transaction.put(newTask);
               } else {
                 System.out.printf("Found existing %d tasks, rolling back", tasksCount);

From 1adf105b9ec1d980a9ce453ec9213160e851663d Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Wed, 19 Oct 2022 17:35:16 +0530
Subject: [PATCH 79/82] Adding fullstop to the comment lines

---
 .../aggregation/CountAggregationInTransaction.java |  6 +++---
 .../aggregation/CountAggregationOnKind.java        | 10 +++++-----
 .../aggregation/CountAggregationWithGqlQuery.java  | 10 +++++-----
 .../aggregation/CountAggregationWithLimit.java     | 10 +++++-----
 .../aggregation/CountAggregationWithOrderBy.java   | 12 ++++++------
 .../CountAggregationWithPropertyFilter.java        | 12 ++++++------
 .../aggregation/CountAggregationWithStaleRead.java | 14 +++++++-------
 7 files changed, 37 insertions(+), 37 deletions(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
index f4215ad98..dd8a93775 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
@@ -32,16 +32,16 @@
 public class CountAggregationInTransaction {
 
   public static void invoke() {
-    // Instantiates a client
+    // Instantiates a client.
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
-    // The kind for the new entity
+    // The kind for the new entity.
     String kind = "Task";
 
     Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
     Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
 
-    // Save all the tasks
+    // Save all the tasks.
     datastore.put(
         Entity.newBuilder(task1Key).set("owner", "john").build(),
         Entity.newBuilder(task2Key).set("owner", "john").build());
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
index 5c84d11e3..36a8636b4 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
@@ -31,31 +31,31 @@
 public class CountAggregationOnKind {
 
   public static void invoke() {
-    // Instantiates a client
+    // Instantiates a client.
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
-    // The kind for the new entity
+    // The kind for the new entity.
     String kind = "Task";
 
     Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
     Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
     Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
 
-    // Save all the tasks
+    // Save all the tasks.
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).build(),
         Entity.newBuilder(task2Key).set("done", false).build(),
         Entity.newBuilder(task3Key).set("done", true).build());
 
     EntityQuery selectAllTasks = Query.newEntityQueryBuilder().setKind(kind).build();
-    // Creating an aggregation query to get the count of all tasks
+    // Creating an aggregation query to get the count of all tasks.
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder()
             .over(selectAllTasks)
             .addAggregation(Aggregation.count())
             .addAggregation(Aggregation.count().as("total_count"))
             .build();
-    // Executing aggregation query
+    // Executing aggregation query.
     AggregationResult allTasksCountQueryResult =
         Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
 
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
index 138b56bee..fd2bf8cc4 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithGqlQuery.java
@@ -30,17 +30,17 @@
 public class CountAggregationWithGqlQuery {
 
   public static void invoke() {
-    // Instantiates a client
+    // Instantiates a client.
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
-    // The kind for the new entity
+    // The kind for the new entity.
     String kind = "Task";
 
     Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
     Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
     Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
 
-    // Save all the tasks
+    // Save all the tasks.
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).build(),
         Entity.newBuilder(task2Key).set("done", false).build(),
@@ -56,7 +56,7 @@ public static void invoke() {
     // Create the aggregation query builder and set the query.
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder().over(selectAllTasks).build();
-    // Executing aggregation query
+    // Executing aggregation query.
     AggregationResult allTasksCountQueryResult =
         Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
 
@@ -75,7 +75,7 @@ public static void invoke() {
     AggregationQuery completedTasksCountQuery =
         Query.newAggregationQueryBuilder().over(completedTasks).build();
 
-    // Executing aggregation query
+    // Executing aggregation query.
     AggregationResult completedTasksCountQueryResult =
         Iterables.getOnlyElement(datastore.runAggregation(completedTasksCountQuery));
 
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
index 9305aacff..4401124e4 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
@@ -30,17 +30,17 @@
 
 public class CountAggregationWithLimit {
   public static void invoke() {
-    // Instantiates a client
+    // Instantiates a client.
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
-    // The kind for the new entity
+    // The kind for the new entity.
     String kind = "Task";
 
     Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
     Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
     Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
 
-    // Save all the tasks
+    // Save all the tasks.
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).build(),
         Entity.newBuilder(task2Key).set("done", false).build(),
@@ -50,13 +50,13 @@ public static void invoke() {
         .setKind(kind)
         .setLimit(2)
         .build();
-    // Creating an aggregation query to get the count of all tasks
+    // Creating an aggregation query to get the count of all tasks.
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder()
             .over(selectAllTasks)
             .addAggregation(Aggregation.count().as("at_least"))
             .build();
-    // Executing aggregation query
+    // Executing aggregation query.
     AggregationResult limitQueryResult =
         Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
 
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
index ce2a2bfbe..63956a172 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
@@ -31,17 +31,17 @@
 
 public class CountAggregationWithOrderBy {
   public static void invoke() {
-    // Instantiates a client
+    // Instantiates a client.
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
-    // The kind for the new entity
+    // The kind for the new entity.
     String kind = "Task";
 
     Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
     Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
     Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
 
-    // Save all the tasks
+    // Save all the tasks.
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).set("priority", 1).build(),
         // Priority not specified.
@@ -52,15 +52,15 @@ public static void invoke() {
     EntityQuery selectAllTasks =
         Query.newEntityQueryBuilder()
             .setKind(kind)
-            .addOrderBy(OrderBy.asc("priority")) // OrderBy acts as an existence filter
+            .addOrderBy(OrderBy.asc("priority"))
             .build();
-    // Creating an aggregation query to get the count of all tasks
+    // Creating an aggregation query to get the count of all tasks.
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder()
             .over(selectAllTasks)
             .addAggregation(Aggregation.count().as("count"))
             .build();
-    // Executing aggregation query
+    // Executing aggregation query.
     AggregationResult limitQueryResult =
         Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
 
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
index d45db430d..4e6f34c08 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
@@ -32,17 +32,17 @@
 public class CountAggregationWithPropertyFilter {
 
   public static void invoke() {
-    // Instantiates a client
+    // Instantiates a client.
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
-    // The kind for the new entity
+    // The kind for the new entity.
     String kind = "Task";
 
     Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
     Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
     Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
 
-    // Save all the tasks
+    // Save all the tasks.
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).build(),
         Entity.newBuilder(task2Key).set("done", false).build(),
@@ -58,20 +58,20 @@ public static void invoke() {
             .setKind(kind)
             .setFilter(PropertyFilter.eq("done", false))
             .build();
-    // Creating an aggregation query to get the count of all completed tasks
+    // Creating an aggregation query to get the count of all completed tasks.
     AggregationQuery completedTasksCountQuery =
         Query.newAggregationQueryBuilder()
             .over(completedTasks)
             .addAggregation(Aggregation.count().as("total_completed_count"))
             .build();
-    // Creating an aggregation query to get the count of all remaining tasks
+    // Creating an aggregation query to get the count of all remaining tasks.
     AggregationQuery remainingTasksCountQuery =
         Query.newAggregationQueryBuilder()
             .over(remainingTasks)
             .addAggregation(Aggregation.count().as("total_remaining_count"))
             .build();
 
-    // Executing aggregation query
+    // Executing aggregation query.
     AggregationResult completedTasksCountQueryResult =
         Iterables.getOnlyElement(datastore.runAggregation(completedTasksCountQuery));
     AggregationResult remainingTasksCountQueryResult =
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
index 77d3d0efd..ab18ed649 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
@@ -33,16 +33,16 @@
 public class CountAggregationWithStaleRead {
 
   public static void invoke() throws InterruptedException {
-    // Instantiates a client
+    // Instantiates a client.
     Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
-    // The kind for the new entity
+    // The kind for the new entity.
     String kind = "Task";
 
     Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
     Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
 
-    // Saving only two tasks
+    // Saving only two tasks.
     datastore.put(
         Entity.newBuilder(task1Key).set("done", true).build(),
         Entity.newBuilder(task2Key).set("done", false).build());
@@ -50,25 +50,25 @@ public static void invoke() throws InterruptedException {
     final Timestamp pastTimestamp = Timestamp.now(); // we have two tasks in database at this time.
 
     Thread.sleep(1000);
-    // Saving third tasks
+    // Saving third tasks.
     Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
     datastore.put(Entity.newBuilder(task3Key).set("done", false).build());
 
     EntityQuery selectAllTasks = Query.newEntityQueryBuilder().setKind(kind).build();
 
-    // Creating an aggregation query to get the count of all tasks
+    // Creating an aggregation query to get the count of all tasks.
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder()
             .over(selectAllTasks)
             .addAggregation(Aggregation.count().as("total_count"))
             .build();
 
-    // Executing aggregation query
+    // Executing aggregation query.
     AggregationResult tasksCountLatest =
         Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
     System.out.printf("Latest tasks count is %d", tasksCountLatest.get("total_count")); // 3
 
-    // Executing aggregation query with past timestamp
+    // Executing aggregation query with past timestamp.
     AggregationResult tasksCountInPast =
         Iterables.getOnlyElement(
             datastore.runAggregation(allTasksCountQuery, ReadOption.readTime(pastTimestamp)));

From e89fde6de9956fe905e70a358f035db75d6a4558 Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Thu, 20 Oct 2022 11:44:04 +0530
Subject: [PATCH 80/82] Showcasing with/without alias usage separately

---
 .../aggregation/CountAggregationOnKind.java   | 49 ++++++++++++++-----
 1 file changed, 37 insertions(+), 12 deletions(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
index 36a8636b4..452f3c173 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
@@ -17,6 +17,8 @@
 package com.example.datastore.aggregation;
 
 // [START datastore_count_aggregation_query_on_kind]
+import static com.google.cloud.datastore.aggregation.Aggregation.count;
+
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.Datastore;
@@ -25,18 +27,17 @@
 import com.google.cloud.datastore.EntityQuery;
 import com.google.cloud.datastore.Key;
 import com.google.cloud.datastore.Query;
-import com.google.cloud.datastore.aggregation.Aggregation;
 import com.google.common.collect.Iterables;
 
 public class CountAggregationOnKind {
+  // Instantiates a client.
+  private static final Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
 
-  public static void invoke() {
-    // Instantiates a client.
-    Datastore datastore = DatastoreOptions.getDefaultInstance().getService();
-
-    // The kind for the new entity.
-    String kind = "Task";
+  // The kind for the new entity.
+  private static final String kind = "Task";
 
+  // Setting up Tasks in database
+  private static void setUpTasks() {
     Key task1Key = datastore.newKeyFactory().setKind(kind).newKey("task1");
     Key task2Key = datastore.newKeyFactory().setKind(kind).newKey("task2");
     Key task3Key = datastore.newKeyFactory().setKind(kind).newKey("task3");
@@ -46,23 +47,47 @@ public static void invoke() {
         Entity.newBuilder(task1Key).set("done", true).build(),
         Entity.newBuilder(task2Key).set("done", false).build(),
         Entity.newBuilder(task3Key).set("done", true).build());
+  }
 
+  // Accessing aggregation result by the generated alias.
+  private static void usageWithGeneratedAlias() {
     EntityQuery selectAllTasks = Query.newEntityQueryBuilder().setKind(kind).build();
     // Creating an aggregation query to get the count of all tasks.
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder()
             .over(selectAllTasks)
-            .addAggregation(Aggregation.count())
-            .addAggregation(Aggregation.count().as("total_count"))
+            .addAggregation(count())
             .build();
     // Executing aggregation query.
-    AggregationResult allTasksCountQueryResult =
+    AggregationResult aggregationResult =
         Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
 
-    System.out.printf("Total tasks count is %d", allTasksCountQueryResult.get("total_count")); // 3
     System.out.printf(
         "Total tasks (accessible from default alias) is %d",
-        allTasksCountQueryResult.get("property_1")); // 3
+        aggregationResult.get("property_1")); // 3
+  }
+
+  // Accessing aggregation result by the provided custom alias.
+  private static void usageWithCustomAlias() {
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder().setKind(kind).build();
+    // Creating an aggregation query to get the count of all tasks.
+    AggregationQuery allTasksCountQuery =
+        Query.newAggregationQueryBuilder()
+            .over(selectAllTasks)
+            // passing 'total_count' as alias in the aggregation query.
+            .addAggregation(count().as("total_count"))
+            .build();
+    // Executing aggregation query.
+    AggregationResult aggregationResult =
+        Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
+
+    System.out.printf("Total tasks count is %d", aggregationResult.get("total_count")); // 3
+  }
+
+  public static void invoke() {
+    setUpTasks();
+    usageWithGeneratedAlias();
+    usageWithCustomAlias();
   }
 }
 // [END datastore_count_aggregation_query_on_kind]

From 36f82faa22e5cd3485c4442d9a4925facfb258dc Mon Sep 17 00:00:00 2001
From: Prateek Jain 
Date: Thu, 20 Oct 2022 11:45:59 +0530
Subject: [PATCH 81/82] Importing static methods are we are now including
 import statements in samples

---
 .../aggregation/CountAggregationInTransaction.java     |  6 ++++--
 .../aggregation/CountAggregationWithLimit.java         |  6 ++++--
 .../aggregation/CountAggregationWithOrderBy.java       | 10 ++++++----
 .../CountAggregationWithPropertyFilter.java            |  8 +++++---
 .../aggregation/CountAggregationWithStaleRead.java     |  6 ++++--
 5 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
index dd8a93775..ca62e2cab 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationInTransaction.java
@@ -17,6 +17,9 @@
 package com.example.datastore.aggregation;
 
 // [START datastore_count_aggregation_query_in_transaction]
+
+import static com.google.cloud.datastore.aggregation.Aggregation.count;
+
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.Datastore;
 import com.google.cloud.datastore.Datastore.TransactionCallable;
@@ -26,7 +29,6 @@
 import com.google.cloud.datastore.Key;
 import com.google.cloud.datastore.Query;
 import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
-import com.google.cloud.datastore.aggregation.Aggregation;
 import com.google.common.collect.Iterables;
 
 public class CountAggregationInTransaction {
@@ -59,7 +61,7 @@ public static void invoke() {
               AggregationQuery totalTasksQuery =
                   Query.newAggregationQueryBuilder()
                       .over(tasksOfJohn)
-                      .addAggregation(Aggregation.count().as("tasks_count"))
+                      .addAggregation(count().as("tasks_count"))
                       .build();
 
               // Executing aggregation query in the ongoing transaction.
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
index 4401124e4..a6dcbc725 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
@@ -17,6 +17,9 @@
 package com.example.datastore.aggregation;
 
 // [START datastore_count_aggregation_query_with_limit]
+
+import static com.google.cloud.datastore.aggregation.Aggregation.count;
+
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.Datastore;
@@ -25,7 +28,6 @@
 import com.google.cloud.datastore.EntityQuery;
 import com.google.cloud.datastore.Key;
 import com.google.cloud.datastore.Query;
-import com.google.cloud.datastore.aggregation.Aggregation;
 import com.google.common.collect.Iterables;
 
 public class CountAggregationWithLimit {
@@ -54,7 +56,7 @@ public static void invoke() {
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder()
             .over(selectAllTasks)
-            .addAggregation(Aggregation.count().as("at_least"))
+            .addAggregation(count().as("at_least"))
             .build();
     // Executing aggregation query.
     AggregationResult limitQueryResult =
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
index 63956a172..200f71303 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
@@ -17,6 +17,10 @@
 package com.example.datastore.aggregation;
 
 // [START datastore_count_aggregation_query_with_order_by]
+
+import static com.google.cloud.datastore.StructuredQuery.OrderBy.asc;
+import static com.google.cloud.datastore.aggregation.Aggregation.count;
+
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.Datastore;
@@ -25,8 +29,6 @@
 import com.google.cloud.datastore.EntityQuery;
 import com.google.cloud.datastore.Key;
 import com.google.cloud.datastore.Query;
-import com.google.cloud.datastore.StructuredQuery.OrderBy;
-import com.google.cloud.datastore.aggregation.Aggregation;
 import com.google.common.collect.Iterables;
 
 public class CountAggregationWithOrderBy {
@@ -52,13 +54,13 @@ public static void invoke() {
     EntityQuery selectAllTasks =
         Query.newEntityQueryBuilder()
             .setKind(kind)
-            .addOrderBy(OrderBy.asc("priority"))
+            .addOrderBy(asc("priority"))
             .build();
     // Creating an aggregation query to get the count of all tasks.
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder()
             .over(selectAllTasks)
-            .addAggregation(Aggregation.count().as("count"))
+            .addAggregation(count().as("count"))
             .build();
     // Executing aggregation query.
     AggregationResult limitQueryResult =
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
index 4e6f34c08..ecc53b496 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithPropertyFilter.java
@@ -17,6 +17,9 @@
 package com.example.datastore.aggregation;
 
 // [START datastore_count_aggregation_query_with_filters]
+
+import static com.google.cloud.datastore.aggregation.Aggregation.count;
+
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.AggregationResult;
 import com.google.cloud.datastore.Datastore;
@@ -26,7 +29,6 @@
 import com.google.cloud.datastore.Key;
 import com.google.cloud.datastore.Query;
 import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
-import com.google.cloud.datastore.aggregation.Aggregation;
 import com.google.common.collect.Iterables;
 
 public class CountAggregationWithPropertyFilter {
@@ -62,13 +64,13 @@ public static void invoke() {
     AggregationQuery completedTasksCountQuery =
         Query.newAggregationQueryBuilder()
             .over(completedTasks)
-            .addAggregation(Aggregation.count().as("total_completed_count"))
+            .addAggregation(count().as("total_completed_count"))
             .build();
     // Creating an aggregation query to get the count of all remaining tasks.
     AggregationQuery remainingTasksCountQuery =
         Query.newAggregationQueryBuilder()
             .over(remainingTasks)
-            .addAggregation(Aggregation.count().as("total_remaining_count"))
+            .addAggregation(count().as("total_remaining_count"))
             .build();
 
     // Executing aggregation query.
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
index ab18ed649..c242c6eec 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
@@ -17,6 +17,9 @@
 package com.example.datastore.aggregation;
 
 // [START datastore_count_aggregation_query_stale_read]
+
+import static com.google.cloud.datastore.aggregation.Aggregation.count;
+
 import com.google.cloud.Timestamp;
 import com.google.cloud.datastore.AggregationQuery;
 import com.google.cloud.datastore.AggregationResult;
@@ -27,7 +30,6 @@
 import com.google.cloud.datastore.Key;
 import com.google.cloud.datastore.Query;
 import com.google.cloud.datastore.ReadOption;
-import com.google.cloud.datastore.aggregation.Aggregation;
 import com.google.common.collect.Iterables;
 
 public class CountAggregationWithStaleRead {
@@ -60,7 +62,7 @@ public static void invoke() throws InterruptedException {
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder()
             .over(selectAllTasks)
-            .addAggregation(Aggregation.count().as("total_count"))
+            .addAggregation(count().as("total_count"))
             .build();
 
     // Executing aggregation query.

From de89eb93a2d0a464ddd1d7b3425c8803250a3679 Mon Sep 17 00:00:00 2001
From: Owl Bot 
Date: Fri, 21 Oct 2022 04:28:47 +0000
Subject: [PATCH 82/82] =?UTF-8?q?=F0=9F=A6=89=20Updates=20from=20OwlBot=20?=
 =?UTF-8?q?post-processor?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md
---
 README.md                                                    | 4 ++--
 .../datastore/aggregation/CountAggregationOnKind.java        | 5 +----
 .../datastore/aggregation/CountAggregationWithLimit.java     | 5 +----
 .../datastore/aggregation/CountAggregationWithOrderBy.java   | 5 +----
 .../datastore/aggregation/CountAggregationWithStaleRead.java | 2 +-
 5 files changed, 6 insertions(+), 15 deletions(-)

diff --git a/README.md b/README.md
index 2ae734023..b4c7605b8 100644
--- a/README.md
+++ b/README.md
@@ -56,13 +56,13 @@ implementation 'com.google.cloud:google-cloud-datastore'
 If you are using Gradle without BOM, add this to your dependencies:
 
 ```Groovy
-implementation 'com.google.cloud:google-cloud-datastore:2.12.0'
+implementation 'com.google.cloud:google-cloud-datastore:2.12.1'
 ```
 
 If you are using SBT, add this to your dependencies:
 
 ```Scala
-libraryDependencies += "com.google.cloud" % "google-cloud-datastore" % "2.12.0"
+libraryDependencies += "com.google.cloud" % "google-cloud-datastore" % "2.12.1"
 ```
 
 ## Authentication
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
index 452f3c173..b2673653a 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationOnKind.java
@@ -54,10 +54,7 @@ private static void usageWithGeneratedAlias() {
     EntityQuery selectAllTasks = Query.newEntityQueryBuilder().setKind(kind).build();
     // Creating an aggregation query to get the count of all tasks.
     AggregationQuery allTasksCountQuery =
-        Query.newAggregationQueryBuilder()
-            .over(selectAllTasks)
-            .addAggregation(count())
-            .build();
+        Query.newAggregationQueryBuilder().over(selectAllTasks).addAggregation(count()).build();
     // Executing aggregation query.
     AggregationResult aggregationResult =
         Iterables.getOnlyElement(datastore.runAggregation(allTasksCountQuery));
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
index a6dcbc725..a6a1a4c55 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithLimit.java
@@ -48,10 +48,7 @@ public static void invoke() {
         Entity.newBuilder(task2Key).set("done", false).build(),
         Entity.newBuilder(task3Key).set("done", true).build());
 
-    EntityQuery selectAllTasks = Query.newEntityQueryBuilder()
-        .setKind(kind)
-        .setLimit(2)
-        .build();
+    EntityQuery selectAllTasks = Query.newEntityQueryBuilder().setKind(kind).setLimit(2).build();
     // Creating an aggregation query to get the count of all tasks.
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder()
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
index 200f71303..247bbdefe 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithOrderBy.java
@@ -52,10 +52,7 @@ public static void invoke() {
 
     // OrderBy acts as an existence filter.
     EntityQuery selectAllTasks =
-        Query.newEntityQueryBuilder()
-            .setKind(kind)
-            .addOrderBy(asc("priority"))
-            .build();
+        Query.newEntityQueryBuilder().setKind(kind).addOrderBy(asc("priority")).build();
     // Creating an aggregation query to get the count of all tasks.
     AggregationQuery allTasksCountQuery =
         Query.newAggregationQueryBuilder()
diff --git a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
index c242c6eec..e46670b98 100644
--- a/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
+++ b/samples/snippets/src/main/java/com/example/datastore/aggregation/CountAggregationWithStaleRead.java
@@ -77,4 +77,4 @@ public static void invoke() throws InterruptedException {
     System.out.printf("Stale tasks count is %d", tasksCountInPast.get("total_count")); // 2
   }
 }
-// [END datastore_count_aggregation_query_stale_read]
\ No newline at end of file
+// [END datastore_count_aggregation_query_stale_read]