feat: Count API

.gitignore

@@ -3,4 +3,5 @@ target/
 *.iml
 __pycache__/
 
-.flattened-pom.xml
+.flattened-pom.xml
+*.patch

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);
+    }
+  }
 }

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;

google-cloud-datastore/clirr-ignored-differences.xml

@@ -11,4 +11,19 @@
         <method>com.google.datastore.v1.ReserveIdsResponse reserveIds(com.google.datastore.v1.ReserveIdsRequest)</method>
         <differenceType>7012</differenceType>
     </difference>
+    <difference>
+        <className>com/google/cloud/datastore/spi/v1/DatastoreRpc</className>
+        <method>com.google.datastore.v1.RunAggregationQueryResponse runAggregationQuery(com.google.datastore.v1.RunAggregationQueryRequest)</method>
+        <differenceType>7012</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/Datastore</className>
+        <method>com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.ReadOption[])</method>
+        <differenceType>7012</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/DatastoreReader</className>
+        <method>com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery)</method>
+        <differenceType>7012</differenceType>
+    </difference>
 </differences>

google-cloud-datastore/pom.xml

@@ -143,6 +143,12 @@
       <artifactId>easymock</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.squareup.okhttp3</groupId>
+      <artifactId>mockwebserver</artifactId>
+      <version>4.10.0</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java

@@ -0,0 +1,178 @@
+/*
+ * 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.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}.
+
+ * <p>{@link StructuredQuery} example:</p>
+ * <pre>{@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"));
+ * }
+ * }</pre>
+ *
+ * <h4>{@link GqlQuery} example:</h4>
+ * <pre>{@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"));
+ * }
+ * }</pre>
+ *
+ * @see <a href="https://cloud.google.com/appengine/docs/java/datastore/queries">Datastore
+ *     queries</a>
+ */
+public class AggregationQuery extends Query<AggregationResults> {
+
+  private Set<Aggregation> aggregations;
+  private StructuredQuery<?> nestedStructuredQuery;
+  private final Mode mode;
+  private GqlQuery<?> nestedGqlQuery;
+
+  AggregationQuery(String namespace, Set<Aggregation> aggregations,
+      StructuredQuery<?> nestedQuery) {
+    super(namespace);
+    checkArgument(!aggregations.isEmpty(),
+        "At least one aggregation is required for an aggregation query to run");
+    this.aggregations = aggregations;
+    this.nestedStructuredQuery = nestedQuery;
+    this.mode = STRUCTURED;
+  }
+
+  AggregationQuery(String namespace, GqlQuery<?> gqlQuery) {
+    super(namespace);
+    this.nestedGqlQuery = gqlQuery;
+    this.mode = GQL;
+  }
+
+  /** Returns the {@link Aggregation}(s) for this Query. */
+  public Set<Aggregation> 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;
+  }
+
+  public static class Builder {
+
+    private String namespace;
+    private Mode mode;
+    private final Set<Aggregation> aggregations;
+    private StructuredQuery<?> nestedStructuredQuery;
+    private GqlQuery<?> nestedGqlQuery;
+
+    public Builder() {
+      this.aggregations = new HashSet<>();
+    }
+
+    public Builder setNamespace(String namespace) {
+      this.namespace = namespace;
+      return this;
+    }
+
+    public Builder addAggregation(AggregationBuilder<?> aggregationBuilder) {
+      this.aggregations.add(aggregationBuilder.build());
+      return this;
+    }
+
+    public Builder addAggregation(Aggregation aggregation) {
+      this.aggregations.add(aggregation);
+      return this;
+    }
+
+    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() {
+      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,
+  }
+}

google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java

@@ -0,0 +1,70 @@
+/*
+ * 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.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<String, LongValue> properties;
+
+  public AggregationResult(Map<String, LongValue> 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();
+  }
+
+  @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);
+  }
+
+  @Override
+  public String toString() {
+    ToStringHelper toStringHelper = MoreObjects.toStringHelper(this);
+    for (Entry<String, LongValue> entry : properties.entrySet()) {
+      toStringHelper.add(entry.getKey(), entry.getValue().get());
+    }
+    return toStringHelper.toString();
+  }
+}