feat: Count API

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,13 @@ 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/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"] }
     ]
   },
   {

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/src/main/java/com/google/cloud/datastore/AggregationQuery.java

@@ -0,0 +1,176 @@
+/*
+ * 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 com.google.api.core.BetaApi;
+import com.google.cloud.datastore.aggregation.Aggregation;
+import com.google.cloud.datastore.aggregation.AggregationBuilder;
+import java.util.HashSet;
+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:
+ *
+ * <pre>{@code
+ * EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
+ *    .setKind("Task")
+ *    .build();
+ * AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
+ *    .addAggregation(count().as("total_count"))
+ *    .over(selectAllQuery)
+ *    .build();
+ * AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery);
+ * for (AggregationResult aggregationResult : aggregationResults) {
+ *     System.out.println(aggregationResult.get("total_count"));
+ * }
+ * }</pre>
+ *
+ * <h4>{@link GqlQuery} example:</h4>
+ *
+ * <pre>{@code
+ * 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>
+ */
+@BetaApi
+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,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.google.cloud.datastore;
+
+import com.google.api.core.BetaApi;
+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;
+
+/** Represents a result of an {@link AggregationQuery} query submission. */
+@BetaApi
+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();
+  }
+}