Skip to content

Commit

Permalink
feat: Count API (#823)
Browse files Browse the repository at this point in the history
* Add method in Datastore client to invoke rpc for aggregation query

* Creating count aggregation and using it to populate Aggregation proto

* Moving aggregation builder method to root level aggregation class

* Introducing RecordQuery to represent queries which returns entity records when executed

* Updating gitignore with patch extension

* Setting up structure of Aggregation query and its builder

* Introducing ProtoPreparer to populate the request protos

* Delegating responsibility of preparing query proto to QueryPreparer

* Populating aggregation query with nested structured query

* Delegating responsibility of preparing query proto in GqlQuery to QueryPreparer

* Removing RecordQuery from the query hierarchy and making it a standalone interface for now

* Populating aggregation query with nested gql query

* Removing deprecation warning by using assertThrows instead of ExpectedException rule

* Making DatastoreRpc call aggregation query method on client

* Creating response transformer to transform aggregation query response into domain objects

* Implementing aggregation query executor to execute AggergationQuery

* Adding missing assertion statements

* Creating RetryExecutor to inject it as a dependency in other components

* Making RetryExecutor accept RetrySettings when creating it

* Revert "Making RetryExecutor accept RetrySettings when creating it"

This reverts commit 1dfafb7.

* Revert "Creating RetryExecutor to inject it as a dependency in other components"

This reverts commit 8872a55.

* Introducing RetryAndTraceDatastoreRpcDecorator to have retry and traceability logic on top of another DatastoreRpc

* Extracting out the responsibility of preparing ReadOption in it's own ProtoPreparer

* Making QueryExecutor to execute query with provided ReadOptions

* Exposing readTime to the user

* Ignoring runAggregationQuery method from clirr check

* Making readTime final

* Allowing namespace to be optional in AggregationQuery

* Add capability to fetch aggregation result by passing alias

* Implementing User facing datastore.runAggrgation method to run aggregation query

* Add integration test for count aggregation

* Add transaction Id support in ReadOptionsProtoPreparer

* Supporting aggregation query with transactions

* Allowing user to create Aggregation directly without involving its builder

* Preventing creating duplicated aggregation when creating an aggregation query

* Marking RecordQuery implemented method as InternalApi

* Writing comments and JavaDoc for aggregation query related class

* Adding a default implementation in the public interfaces to avoid compile time failures

* covering a scenario to maintain consistent snapshot when executing aggregation query in a transaction

* Creating emulator proxy to simulate AggregationQuery response from emulator

* Integration test to execute an aggregation query in a read only transaction

* Getting rid off limit operation on count aggregation as same behaviour can be achieved by using 'limit' operation on the underlying query

* Removing import statement from javadoc and undo changes in .gitignore file

* Using Optional instead of returning null from ReadOptionsProtoPreparer

* using assertThat from Truth library

* fixing unit test

* Getting rid off Double braces initialization syntax

* Fixing lint

* Getting rid off emulator proxy and using easy mock to check the aggregationQuery triggered

* Deleting a entity created locally in other test which is causing failure in other test

* Deleting all keys in datastore in integration test so that new test can start fresh

* Executing two read write transaction simultaneously and verifying their behaviour

* Removing tests to verify serializability as it's an underlying implementation detail

* Fixing lint

* Adding runAggregationQuery method to reflect config so that it's available and accessible in native image through reflection

* Fixing equals of CountAggregation

* Fixing lint

* Adding an integration test of using limit option with aggregation query

* Adding BetaApi annotation to public surface to indicate that aggregation query / count is in preview

* Fixing lint

* Removing unused functiona and fixing javadoc

* fixing variable name
  • Loading branch information
jainsahab authored Oct 17, 2022
1 parent 4641306 commit 8c22e61
Show file tree
Hide file tree
Showing 38 changed files with 2,410 additions and 12 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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"] }
]
},
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
15 changes: 15 additions & 0 deletions google-cloud-datastore/clirr-ignored-differences.xml
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Original file line number Diff line number Diff line change
@@ -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,
}
}
Original file line number Diff line number Diff line change
@@ -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();
}
}
Loading

0 comments on commit 8c22e61

Please sign in to comment.