From 9d0a267e25b1086cfab04eefc5d2ae112cdb31aa Mon Sep 17 00:00:00 2001 From: kolea2 Date: Wed, 13 Dec 2023 10:37:45 -0500 Subject: [PATCH 1/5] feat: query profiling --- .../clirr-ignored-differences.xml | 56 ++++++ .../cloud/datastore/AggregationResults.java | 23 ++- .../com/google/cloud/datastore/Datastore.java | 53 ++++++ .../google/cloud/datastore/DatastoreImpl.java | 37 +++- .../cloud/datastore/DatastoreReader.java | 12 ++ .../com/google/cloud/datastore/Query.java | 6 + .../google/cloud/datastore/QueryResults.java | 10 + .../cloud/datastore/QueryResultsImpl.java | 23 ++- .../google/cloud/datastore/ReadOption.java | 20 +- .../google/cloud/datastore/Transaction.java | 7 + .../cloud/datastore/TransactionImpl.java | 20 +- .../execution/AggregationQueryExecutor.java | 13 +- .../datastore/execution/QueryExecutor.java | 5 +- .../AggregationQueryRequestProtoPreparer.java | 2 +- .../AggregationQueryResponseTransformer.java | 7 +- .../cloud/datastore/models/QueryPlan.java | 57 ++++++ .../cloud/datastore/models/QueryProfile.java | 59 ++++++ .../datastore/models/ResultSetStats.java | 76 ++++++++ .../AggregationQueryExecutorTest.java | 17 +- ...regationQueryRequestProtoPreparerTest.java | 39 +++- ...gregationQueryResponseTransformerTest.java | 46 +++++ .../cloud/datastore/it/ITDatastoreTest.java | 177 +++++++++++++++++- .../cloud/datastore/models/QueryPlanTest.java | 66 +++++++ .../datastore/models/QueryProfileTest.java | 33 ++++ .../datastore/models/ResultSetStatsTest.java | 97 ++++++++++ 25 files changed, 923 insertions(+), 38 deletions(-) create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/QueryPlan.java create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/QueryProfile.java create mode 100644 google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/ResultSetStats.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/QueryPlanTest.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/QueryProfileTest.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/ResultSetStatsTest.java diff --git a/google-cloud-datastore/clirr-ignored-differences.xml b/google-cloud-datastore/clirr-ignored-differences.xml index 34395cc01..189df195e 100644 --- a/google-cloud-datastore/clirr-ignored-differences.xml +++ b/google-cloud-datastore/clirr-ignored-differences.xml @@ -13,4 +13,60 @@ *QueryConfig* 7005 + + + + com/google/cloud/datastore/Datastore + com.google.cloud.datastore.QueryResults run(com.google.cloud.datastore.Query, com.google.cloud.datastore.models.QueryProfile$QueryMode, com.google.cloud.datastore.ReadOption[]) + 7012 + + + com/google/cloud/datastore/Datastore + com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.models.QueryProfile$QueryMode, com.google.cloud.datastore.ReadOption[]) + 7012 + + + com/google/cloud/datastore/QueryResults + com.google.cloud.datastore.models.ResultSetStats getResultSetStats() + 7012 + + + com/google/cloud/datastore/DatastoreReader + com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.models.QueryProfile$QueryMode) + 7012 + + + com/google/cloud/datastore/DatastoreReader + com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.models.QueryProfile$QueryMode) + 7012 + + + com/google/cloud/datastore/Transaction + com.google.cloud.datastore.QueryResults run(com.google.cloud.datastore.Query, com.google.cloud.datastore.models.QueryProfile$QueryMode) + 7012 + + + + + com/google/cloud/datastore/execution/QueryExecutor + java.lang.Object execute(com.google.cloud.datastore.Query, com.google.cloud.datastore.models.QueryProfile$QueryMode, com.google.cloud.datastore.ReadOption[]) + 7012 + + + + + com/google/cloud/datastore/execution/AggregationQueryExecutor + com.google.cloud.datastore.AggregationResults execute(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.ReadOption[]) + 7004 + + + com/google/cloud/datastore/execution/AggregationQueryExecutor + java.lang.Object execute(com.google.cloud.datastore.Query, com.google.cloud.datastore.ReadOption[]) + 7004 + + + com/google/cloud/datastore/execution/QueryExecutor + java.lang.Object execute(com.google.cloud.datastore.Query, com.google.cloud.datastore.ReadOption[]) + 7004 + 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 d85e8f8ae..53b422889 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 @@ -19,6 +19,7 @@ import com.google.api.core.InternalApi; import com.google.cloud.Timestamp; +import com.google.cloud.datastore.models.ResultSetStats; import java.util.Iterator; import java.util.List; import java.util.Objects; @@ -33,12 +34,21 @@ public class AggregationResults implements Iterable { private final List aggregationResults; private final Timestamp readTime; + private final ResultSetStats resultSetStats; public AggregationResults(List aggregationResults, Timestamp readTime) { + this(aggregationResults, readTime, null); + } + + public AggregationResults( + List aggregationResults, + Timestamp readTime, + ResultSetStats resultSetStats) { checkNotNull(aggregationResults, "Aggregation results cannot be null"); checkNotNull(readTime, "readTime cannot be null"); this.aggregationResults = aggregationResults; this.readTime = readTime; + this.resultSetStats = resultSetStats; } /** Returns {@link Iterator} for underlying List<{@link AggregationResult}>. */ @@ -51,6 +61,13 @@ public int size() { return this.aggregationResults.size(); } + /* + * Returns the ResultSetStats if QueryMode is set to EXPLAIN or EXPLAIN_ANALYZE. Otherwise, returns null. + */ + public ResultSetStats getResultSetStats() { + return this.resultSetStats; + } + @InternalApi public AggregationResult get(int index) { return this.aggregationResults.get(index); @@ -70,11 +87,13 @@ public boolean equals(Object o) { return false; } AggregationResults that = (AggregationResults) o; - return Objects.equals(aggregationResults, that.aggregationResults); + return Objects.equals(aggregationResults, that.aggregationResults) + && Objects.equals(readTime, that.readTime) + && Objects.equals(resultSetStats, that.resultSetStats); } @Override public int hashCode() { - return Objects.hash(aggregationResults); + return Objects.hash(aggregationResults, readTime, resultSetStats); } } 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 d78bea9a2..5dcf55172 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 @@ -16,7 +16,9 @@ package com.google.cloud.datastore; +import com.google.api.core.BetaApi; import com.google.cloud.Service; +import com.google.cloud.datastore.models.QueryProfile; import com.google.datastore.v1.TransactionOptions; import java.util.Iterator; import java.util.List; @@ -462,6 +464,29 @@ interface TransactionCallable { */ QueryResults run(Query query, ReadOption... options); + /** + * Submits a {@link Query} with specified {@link com.google.cloud.datastore.models.QueryProfile.QueryMode} and returns its result. {@link ReadOption}s can be specified if + * desired. + * + *

Example of running a query to find all entities of one kind. + * + * 

{@code
+   * String kind = "my_kind";
+   * StructuredQuery query = Query.newEntityQueryBuilder()
+   *     .setKind(kind)
+   *     .build();
+   * QueryResults results = datastore.run(query, QueryMode.EXPLAIN_ANALYZE);
+   * ResultSetStats resultSetStats = results.getResultSetStats();
+   * }
+ * + * @throws DatastoreException upon failure + */ + @BetaApi + default QueryResults run( + Query query, QueryProfile.QueryMode queryMode, ReadOption... options) { + throw new UnsupportedOperationException("not implemented"); + } + /** * Submits a {@link AggregationQuery} and returns {@link AggregationResults}. {@link ReadOption}s * can be specified if desired. @@ -508,4 +533,32 @@ interface TransactionCallable { default AggregationResults runAggregation(AggregationQuery query, ReadOption... options) { throw new UnsupportedOperationException("Not implemented."); } + + /** + * Submits a {@link AggregationQuery} with specified {@link com.google.cloud.datastore.models.QueryProfile.QueryMode} 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
+   * EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
+   *    .setKind("Task")
+   *    .build();
+   * AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
+   *    .addAggregation(count().as("total_count"))
+   *    .over(selectAllQuery)
+   *    .build();
+   * AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery, QueryMode.EXPLAIN_ANALYZE);
+   * ResultSetStats aggregationStats = aggregationResults.getResultSetStats();
+   * }
+ * @throws DatastoreException upon failure + * @return {@link AggregationResults} + */ + @BetaApi + default AggregationResults runAggregation( + AggregationQuery query, QueryProfile.QueryMode queryMode, 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 a1b337c05..795fe4215 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 @@ -16,6 +16,7 @@ package com.google.cloud.datastore; +import com.google.api.core.BetaApi; import com.google.api.gax.retrying.RetrySettings; import com.google.cloud.BaseService; import com.google.cloud.ExceptionHandler; @@ -23,6 +24,7 @@ import com.google.cloud.RetryHelper.RetryHelperException; import com.google.cloud.ServiceOptions; import com.google.cloud.datastore.execution.AggregationQueryExecutor; +import com.google.cloud.datastore.models.QueryProfile; import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; @@ -30,6 +32,7 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; import com.google.common.collect.Sets; +import com.google.datastore.v1.QueryMode; import com.google.datastore.v1.ReadOptions; import com.google.datastore.v1.ReserveIdsRequest; import com.google.datastore.v1.TransactionOptions; @@ -182,28 +185,50 @@ public T runInTransaction( @Override public QueryResults run(Query query) { - return run(Optional.empty(), query); + return run(Optional.empty(), query, QueryMode.NORMAL); } @Override public QueryResults run(Query query, ReadOption... options) { - return run(toReadOptionsPb(options), query); + return run(toReadOptionsPb(options), query, QueryMode.NORMAL); + } + + @Override + @BetaApi + public QueryResults run( + Query query, QueryProfile.QueryMode queryMode, ReadOption... options) { + return run(toReadOptionsPb(options), query, queryMode.toPb()); } @SuppressWarnings("unchecked") - QueryResults run(Optional readOptionsPb, Query query) { + QueryResults run( + Optional readOptionsPb, Query query, QueryMode queryMode) { return new QueryResultsImpl( - this, readOptionsPb, (RecordQuery) query, query.getNamespace()); + this, readOptionsPb, (RecordQuery) query, query.getNamespace(), queryMode); } @Override public AggregationResults runAggregation(AggregationQuery query) { - return aggregationQueryExecutor.execute(query); + return aggregationQueryExecutor.execute(query, QueryProfile.QueryMode.NORMAL); } @Override public AggregationResults runAggregation(AggregationQuery query, ReadOption... options) { - return aggregationQueryExecutor.execute(query, options); + return aggregationQueryExecutor.execute(query, QueryProfile.QueryMode.NORMAL, options); + } + + @Override + @BetaApi + public AggregationResults runAggregation( + AggregationQuery query, QueryProfile.QueryMode queryMode) { + return aggregationQueryExecutor.execute(query, queryMode); + } + + @Override + @BetaApi + public AggregationResults runAggregation( + AggregationQuery query, QueryProfile.QueryMode queryMode, ReadOption... options) { + return aggregationQueryExecutor.execute(query, queryMode, options); } com.google.datastore.v1.RunQueryResponse runQuery( 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 2a3071f3c..11c270837 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 @@ -16,6 +16,7 @@ package com.google.cloud.datastore; +import com.google.cloud.datastore.models.QueryProfile; import java.util.Iterator; import java.util.List; @@ -62,4 +63,15 @@ public interface DatastoreReader { default AggregationResults runAggregation(AggregationQuery query) { throw new UnsupportedOperationException("Not implemented."); } + + /** + * Submits a {@link AggregationQuery} with a specified {@link com.google.cloud.datastore.models.QueryProfile.QueryMode} and returns {@link + * AggregationResults}. + * + * @throws DatastoreException upon failure + */ + default AggregationResults runAggregation( + AggregationQuery query, QueryProfile.QueryMode queryMode) { + throw new UnsupportedOperationException("Not implemented."); + } } 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 89b8d8807..883051972 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,6 +16,7 @@ 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.common.collect.Maps; @@ -117,6 +118,11 @@ public Class resultClass() { return resultClass; } + @InternalApi + public com.google.datastore.v1.EntityResult.ResultType getQueryType() { + return this.queryType; + } + @Override public int hashCode() { return resultClass.hashCode(); diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java index 3f73824dc..911bcb2a9 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java @@ -16,6 +16,8 @@ package com.google.cloud.datastore; +import com.google.api.core.BetaApi; +import com.google.cloud.datastore.models.ResultSetStats; import com.google.datastore.v1.QueryResultBatch; import java.util.Iterator; @@ -70,4 +72,12 @@ public interface QueryResults extends Iterator { /** Returns MoreResults state of the query after the current batch. */ QueryResultBatch.MoreResultsType getMoreResults(); + + /* + * Returns the {@link ResultSetStats} if {@link QueryMode} was set to EXPLAIN or EXPLAIN_ANALYZE. Otherwise, it will return null. + */ + @BetaApi + default ResultSetStats getResultSetStats() { + throw new UnsupportedOperationException("not implemented"); + } } 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 8fc731ace..608060aab 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 @@ -16,9 +16,12 @@ package com.google.cloud.datastore; +import com.google.api.core.BetaApi; import com.google.cloud.datastore.Query.ResultType; +import com.google.cloud.datastore.models.ResultSetStats; import com.google.common.base.Preconditions; import com.google.common.collect.AbstractIterator; +import com.google.datastore.v1.QueryMode; import com.google.datastore.v1.QueryResultBatch.MoreResultsType; import com.google.datastore.v1.ReadOptions; import com.google.protobuf.ByteString; @@ -40,16 +43,20 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults private Iterator entityResultPbIter; private ByteString cursor; private MoreResultsType moreResults; + private QueryMode queryMode; + private ResultSetStats resultSetStats; QueryResultsImpl( DatastoreImpl datastore, Optional readOptionsPb, RecordQuery query, - String namespace) { + String namespace, + QueryMode queryMode) { this.datastore = datastore; this.readOptionsPb = readOptionsPb; this.query = query; queryResultType = query.getType(); + this.queryMode = queryMode; com.google.datastore.v1.PartitionId.Builder pbBuilder = com.google.datastore.v1.PartitionId.newBuilder(); pbBuilder.setProjectId(datastore.getOptions().getProjectId()); @@ -75,6 +82,7 @@ private void sendRequest() { requestPb.setPartitionId(partitionIdPb); requestPb.setProjectId(datastore.getOptions().getProjectId()); requestPb.setDatabaseId(datastore.getOptions().getDatabaseId()); + requestPb.setMode(queryMode); query.populatePb(requestPb); runQueryResponsePb = datastore.runQuery(requestPb.build()); mostRecentQueryPb = requestPb.getQuery(); @@ -86,9 +94,14 @@ private void sendRequest() { // projection entity can represent all type of results actualResultType = ResultType.PROJECTION_ENTITY; } + boolean isQueryPlan = actualResultType.getQueryType() == null && queryMode == QueryMode.PLAN; Preconditions.checkState( - queryResultType.isAssignableFrom(actualResultType), + queryResultType.isAssignableFrom(actualResultType) || isQueryPlan, "Unexpected result type " + actualResultType + " vs " + queryResultType); + if (runQueryResponsePb.getStats().hasQueryPlan() + || runQueryResponsePb.getStats().hasQueryStats()) { + this.resultSetStats = new ResultSetStats(runQueryResponsePb.getStats()); + } } @Override @@ -127,4 +140,10 @@ public int getSkippedResults() { public MoreResultsType getMoreResults() { return moreResults; } + + @Override + @BetaApi + public ResultSetStats getResultSetStats() { + return this.resultSetStats; + } } 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 c249e45a6..21ccfd06f 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.api.core.InternalApi; import com.google.cloud.Timestamp; import com.google.common.collect.ImmutableMap; +import com.google.datastore.v1.QueryMode; import com.google.protobuf.ByteString; import java.io.Serializable; import java.util.Collections; @@ -146,32 +147,39 @@ public static class QueryConfig> { Q query; List readOptions; + QueryMode queryMode; - private QueryConfig(Q query, List readOptions) { + private QueryConfig(Q query, QueryMode queryMode, List readOptions) { this.query = query; + this.queryMode = queryMode; this.readOptions = readOptions; } - private QueryConfig(Q query) { + private QueryConfig(Q query, QueryMode queryMode) { this.query = query; this.readOptions = Collections.emptyList(); + this.queryMode = queryMode; } public Q getQuery() { return query; } + public QueryMode getQueryMode() { + return this.queryMode; + } + public List getReadOptions() { return readOptions; } - public static > QueryConfig create(Q query) { - return new QueryConfig<>(query); + public static > QueryConfig create(Q query, QueryMode queryMode) { + return new QueryConfig<>(query, queryMode); } public static > QueryConfig create( - Q query, List readOptions) { - return new QueryConfig<>(query, readOptions); + Q query, QueryMode queryMode, List readOptions) { + return new QueryConfig<>(query, queryMode, readOptions); } } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Transaction.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Transaction.java index 69c18d75c..c7b08a4d0 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Transaction.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Transaction.java @@ -16,6 +16,8 @@ package com.google.cloud.datastore; +import com.google.api.core.BetaApi; +import com.google.cloud.datastore.models.QueryProfile.QueryMode; import com.google.protobuf.ByteString; import java.util.Iterator; import java.util.List; @@ -176,6 +178,11 @@ interface Response { @Override QueryResults run(Query query); + @BetaApi + default QueryResults run(Query query, QueryMode queryMode) { + throw new UnsupportedOperationException("not implemented"); + } + /** * Datastore add operation. This method will also allocate id for any entity with an incomplete * key. As opposed to {@link #add(FullEntity)} and {@link #add(FullEntity...)}, this method will 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 3b5e5e4e8..3301b0c0d 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 @@ -18,7 +18,10 @@ import static com.google.cloud.datastore.ReadOption.transactionId; +import com.google.api.core.BetaApi; +import com.google.cloud.datastore.models.QueryProfile; import com.google.common.collect.ImmutableList; +import com.google.datastore.v1.QueryMode; import com.google.datastore.v1.ReadOptions; import com.google.datastore.v1.TransactionOptions; import com.google.protobuf.ByteString; @@ -102,7 +105,16 @@ public QueryResults run(Query query) { validateActive(); Optional readOptions = this.readOptionProtoPreparer.prepare(ImmutableList.of(transactionId(transactionId))); - return datastore.run(readOptions, query); + return datastore.run(readOptions, query, QueryMode.NORMAL); + } + + @Override + @BetaApi + public QueryResults run(Query query, QueryProfile.QueryMode queryMode) { + validateActive(); + Optional readOptions = + this.readOptionProtoPreparer.prepare(ImmutableList.of(transactionId(transactionId))); + return datastore.run(readOptions, query, queryMode.toPb()); } @Override @@ -110,6 +122,12 @@ public AggregationResults runAggregation(AggregationQuery query) { return datastore.runAggregation(query, transactionId(transactionId)); } + @Override + public AggregationResults runAggregation( + AggregationQuery query, QueryProfile.QueryMode queryMode) { + return datastore.runAggregation(query, queryMode, transactionId(transactionId)); + } + @Override public Transaction.Response commit() { validateActive(); 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 5a1fdd2c3..d02b4d842 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 @@ -23,7 +23,9 @@ import com.google.cloud.datastore.ReadOption.QueryConfig; import com.google.cloud.datastore.execution.request.AggregationQueryRequestProtoPreparer; import com.google.cloud.datastore.execution.response.AggregationQueryResponseTransformer; +import com.google.cloud.datastore.models.QueryProfile; import com.google.cloud.datastore.spi.v1.DatastoreRpc; +import com.google.datastore.v1.QueryMode; import com.google.datastore.v1.RunAggregationQueryRequest; import com.google.datastore.v1.RunAggregationQueryResponse; import java.util.Arrays; @@ -47,20 +49,21 @@ public AggregationQueryExecutor(DatastoreRpc datastoreRpc, DatastoreOptions data } @Override - public AggregationResults execute(AggregationQuery query, ReadOption... readOptions) { + public AggregationResults execute( + AggregationQuery query, QueryProfile.QueryMode queryMode, ReadOption... readOptions) { RunAggregationQueryRequest runAggregationQueryRequest = - getRunAggregationQueryRequest(query, readOptions); + getRunAggregationQueryRequest(query, queryMode.toPb(), readOptions); RunAggregationQueryResponse runAggregationQueryResponse = this.datastoreRpc.runAggregationQuery(runAggregationQueryRequest); return this.responseTransformer.transform(runAggregationQueryResponse); } private RunAggregationQueryRequest getRunAggregationQueryRequest( - AggregationQuery query, ReadOption... readOptions) { + AggregationQuery query, QueryMode queryMode, ReadOption... readOptions) { QueryConfig queryConfig = readOptions == null - ? QueryConfig.create(query) - : QueryConfig.create(query, Arrays.asList(readOptions)); + ? QueryConfig.create(query, queryMode) + : QueryConfig.create(query, queryMode, Arrays.asList(readOptions)); return this.protoPreparer.prepare(queryConfig); } } 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 856c64a02..3d5f6b15b 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 @@ -18,6 +18,7 @@ import com.google.api.core.InternalApi; import com.google.cloud.datastore.Query; import com.google.cloud.datastore.ReadOption; +import com.google.cloud.datastore.models.QueryProfile; /** * An internal functional interface whose implementation has the responsibility to execute a {@link @@ -34,7 +35,9 @@ public interface QueryExecutor, OUTPUT> { /** * @param query A {@link Query} to execute. + * @param queryMode {@link com.google.cloud.datastore.models.QueryProfile.QueryMode}s to be used + * when executing {@link Query}. * @param readOptions Optional {@link ReadOption}s to be used when executing {@link Query}. */ - OUTPUT execute(INPUT query, ReadOption... readOptions); + OUTPUT execute(INPUT query, QueryProfile.QueryMode queryMode, 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 475a47b58..2d85682f4 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 @@ -67,7 +67,7 @@ public RunAggregationQueryRequest prepare(QueryConfig queryCon } else { aggregationQueryRequestBuilder.setAggregationQuery(getAggregationQuery(aggregationQuery)); } - + aggregationQueryRequestBuilder.setMode(queryConfig.getQueryMode()); Optional readOptionsPb = readOptionProtoPreparer.prepare(readOptions); readOptionsPb.ifPresent(aggregationQueryRequestBuilder::setReadOptions); return aggregationQueryRequestBuilder.build(); 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 8c99fcd41..1506d5456 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 @@ -19,6 +19,7 @@ import com.google.cloud.Timestamp; import com.google.cloud.datastore.AggregationResult; import com.google.cloud.datastore.AggregationResults; +import com.google.cloud.datastore.models.ResultSetStats; import com.google.datastore.v1.RunAggregationQueryResponse; import com.google.datastore.v1.Value; import java.util.AbstractMap.SimpleEntry; @@ -40,7 +41,11 @@ public AggregationResults transform(RunAggregationQueryResponse response) { response.getBatch().getAggregationResultsList().stream() .map(aggregationResult -> new AggregationResult(transformValues(aggregationResult))) .collect(Collectors.toCollection(LinkedList::new)); - return new AggregationResults(aggregationResults, readTime); + ResultSetStats stats = null; + if (response.getStats().hasQueryStats() || response.getStats().hasQueryPlan()) { + stats = new ResultSetStats(response.getStats()); + } + return new AggregationResults(aggregationResults, readTime, stats); } private Map> transformValues( diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/QueryPlan.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/QueryPlan.java new file mode 100644 index 000000000..6ec98069d --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/QueryPlan.java @@ -0,0 +1,57 @@ +/* + * Copyright 2023 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.models; + +import com.google.api.core.InternalApi; +import com.google.cloud.Structs; +import com.google.common.base.Objects; +import java.util.Map; + +/* + * Class to model the returned plan for the query. This includes plan info - a map containing phase information for the query. + */ +public class QueryPlan { + private final Map planInfo; + + @InternalApi + public QueryPlan(com.google.datastore.v1.QueryPlan proto) { + this.planInfo = Structs.asMap(proto.getPlanInfo()); + } + + /* + * Returns the plan info as a map. + */ + public Map getPlanInfo() { + return this.planInfo; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof QueryPlan)) { + return false; + } + QueryPlan queryPlan = (QueryPlan) o; + return Objects.equal(planInfo, queryPlan.planInfo); + } + + @Override + public int hashCode() { + return Objects.hashCode(planInfo); + } +} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/QueryProfile.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/QueryProfile.java new file mode 100644 index 000000000..4c9262f29 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/QueryProfile.java @@ -0,0 +1,59 @@ +/* + * Copyright 2023 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.models; + +import com.google.api.core.InternalApi; +import com.google.common.annotations.VisibleForTesting; + +public class QueryProfile { + + /* + * Model enums for query mode. + */ + public enum QueryMode { + /* + * The default mode. Only the query results are returned. + */ + NORMAL(com.google.datastore.v1.QueryMode.NORMAL), + + /* + * This mode returns only the query plan, without any results or execution statistics information. + */ + EXPLAIN(com.google.datastore.v1.QueryMode.PLAN), + + /* + * This mode returns both the query plan and the execution statistics along with the results. + */ + EXPLAIN_ANALYZE(com.google.datastore.v1.QueryMode.PROFILE), + + /* + * Unrecognized enum specified; please update your client. + */ + UNRECOGZNIED(com.google.datastore.v1.QueryMode.UNRECOGNIZED); + + private final com.google.datastore.v1.QueryMode proto; + + QueryMode(com.google.datastore.v1.QueryMode proto) { + this.proto = proto; + } + + @InternalApi + @VisibleForTesting + public com.google.datastore.v1.QueryMode toPb() { + return this.proto; + } + } +} diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/ResultSetStats.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/ResultSetStats.java new file mode 100644 index 000000000..7843836d3 --- /dev/null +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/ResultSetStats.java @@ -0,0 +1,76 @@ +/* + * Copyright 2023 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.models; + +import com.google.api.core.InternalApi; +import com.google.cloud.Structs; +import com.google.common.base.Objects; +import java.util.Map; + +/* + * Model class containing the planning and execution stats for the query. If QueryMode.EXPLAIN_ANALYZE was set, + * QueryPlan and QueryStats will be present. If QueryMode.EXPLAIN was set, QueryPlan will be present and QueryStats will return null. + */ +public class ResultSetStats { + + private QueryPlan queryPlan = null; + private Map queryStats = null; + + @InternalApi + public ResultSetStats(com.google.datastore.v1.ResultSetStats proto) { + if (proto.hasQueryPlan()) { + this.queryPlan = new QueryPlan(proto.getQueryPlan()); + } + if (proto.hasQueryStats()) { + this.queryStats = Structs.asMap(proto.getQueryStats()); + } + } + + /* + * Returns the plan for the query, if EXPLAIN or EXPLAIN_ANALYZE was set. Otherwise, returns null. + */ + public QueryPlan getQueryPlan() { + return this.queryPlan; + } + + /* + * Returns the stats for the query if EXPLAIN_ANALYZE was set. Otherwise, returns null. + */ + public Map getQueryStats() { + return this.queryStats; + } + + public boolean hasQueryStats() { + return this.queryStats != null; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof ResultSetStats)) { + return false; + } + ResultSetStats that = (ResultSetStats) o; + return Objects.equal(queryPlan, that.queryPlan) && Objects.equal(queryStats, that.queryStats); + } + + @Override + public int hashCode() { + return Objects.hashCode(queryPlan, queryStats); + } +} 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 f9f23261d..e8dbe6fd5 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 @@ -36,6 +36,7 @@ import com.google.cloud.datastore.EntityQuery; import com.google.cloud.datastore.LongValue; import com.google.cloud.datastore.Query; +import com.google.cloud.datastore.models.QueryProfile.QueryMode; import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.common.collect.ImmutableMap; import com.google.datastore.v1.AggregationResultBatch; @@ -88,7 +89,8 @@ public void shouldExecuteAggregationQuery() { replay(mockRpc); - AggregationResults aggregationResults = queryExecutor.execute(aggregationQuery); + AggregationResults aggregationResults = + queryExecutor.execute(aggregationQuery, QueryMode.NORMAL); verify(mockRpc); assertThat(aggregationResults) @@ -101,7 +103,8 @@ public void shouldExecuteAggregationQuery() { new AggregationResult( ImmutableMap.of( "count", LongValue.of(509), "property_2", LongValue.of(100)))), - Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime()))); + Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime()), + null)); } @Test @@ -127,7 +130,7 @@ public void shouldExecuteAggregationQueryWithReadOptions() { replay(mockRpc); AggregationResults aggregationResults = - queryExecutor.execute(aggregationQuery, eventualConsistency()); + queryExecutor.execute(aggregationQuery, QueryMode.NORMAL, eventualConsistency()); verify(mockRpc); assertThat(aggregationResults) @@ -140,7 +143,8 @@ public void shouldExecuteAggregationQueryWithReadOptions() { new AggregationResult( ImmutableMap.of( "count", LongValue.of(509), "property_2", LongValue.of(100)))), - Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime()))); + Timestamp.fromProto(runAggregationQueryResponse.getBatch().getReadTime()), + null)); } private RunAggregationQueryResponse placeholderAggregationQueryResponse() { @@ -167,7 +171,10 @@ private RunAggregationQueryResponse placeholderAggregationQueryResponse() { .putAllAggregateProperties(result2) .build()) .build(); - return RunAggregationQueryResponse.newBuilder().setBatch(resultBatch).build(); + return RunAggregationQueryResponse.newBuilder() + .setBatch(resultBatch) + .setStats(com.google.datastore.v1.ResultSetStats.newBuilder().build()) + .build(); } private Predicate runAggregationRequestWithEventualConsistency() { 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 0f22828d9..dfc778eb1 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,6 +40,7 @@ import com.google.cloud.datastore.ReadOption; import com.google.cloud.datastore.ReadOption.QueryConfig; import com.google.common.collect.ImmutableMap; +import com.google.datastore.v1.QueryMode; import com.google.datastore.v1.RunAggregationQueryRequest; import java.util.HashMap; import org.junit.Test; @@ -91,13 +92,15 @@ public class AggregationQueryRequestProtoPreparerTest { @Test public void shouldPrepareAggregationQueryRequestWithGivenStructuredQuery() { RunAggregationQueryRequest runAggregationQueryRequest = - protoPreparer.prepare(QueryConfig.create(AGGREGATION_OVER_STRUCTURED_QUERY)); + protoPreparer.prepare( + QueryConfig.create(AGGREGATION_OVER_STRUCTURED_QUERY, QueryMode.NORMAL)); assertThat(runAggregationQueryRequest.getProjectId()).isEqualTo(PROJECT_ID); assertThat(runAggregationQueryRequest.getDatabaseId()).isEqualTo(DATABASE_ID); assertThat(runAggregationQueryRequest.getPartitionId().getProjectId()).isEqualTo(PROJECT_ID); assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId()).isEqualTo(NAMESPACE); + assertThat(runAggregationQueryRequest.getMode()).isEqualTo(QueryMode.NORMAL); com.google.datastore.v1.AggregationQuery aggregationQueryProto = runAggregationQueryRequest.getAggregationQuery(); @@ -114,7 +117,7 @@ public void shouldPrepareAggregationQueryRequestWithGivenStructuredQuery() { @Test public void shouldPrepareAggregationQueryRequestWithGivenGqlQuery() { RunAggregationQueryRequest runAggregationQueryRequest = - protoPreparer.prepare(QueryConfig.create(AGGREGATION_OVER_GQL_QUERY)); + protoPreparer.prepare(QueryConfig.create(AGGREGATION_OVER_GQL_QUERY, QueryMode.NORMAL)); assertThat(runAggregationQueryRequest.getProjectId()).isEqualTo(PROJECT_ID); assertThat(runAggregationQueryRequest.getDatabaseId()).isEqualTo(DATABASE_ID); @@ -122,6 +125,7 @@ public void shouldPrepareAggregationQueryRequestWithGivenGqlQuery() { assertThat(runAggregationQueryRequest.getPartitionId().getProjectId()).isEqualTo(PROJECT_ID); assertThat(runAggregationQueryRequest.getPartitionId().getDatabaseId()).isEqualTo(DATABASE_ID); assertThat(runAggregationQueryRequest.getPartitionId().getNamespaceId()).isEqualTo(NAMESPACE); + assertThat(runAggregationQueryRequest.getMode()).isEqualTo(QueryMode.NORMAL); com.google.datastore.v1.GqlQuery gqlQueryProto = runAggregationQueryRequest.getGqlQuery(); @@ -172,17 +176,42 @@ public void shouldPrepareAggregationQueryWithNamespaceFromDatastoreOptions() { Query.newAggregationQueryBuilder().over(COMPLETED_TASK_GQL_QUERY).build(); RunAggregationQueryRequest runAggregationQueryFromStructuredQuery = - protoPreparer.prepare(QueryConfig.create(structuredQueryWithoutNamespace)); + protoPreparer.prepare( + QueryConfig.create(structuredQueryWithoutNamespace, QueryMode.NORMAL)); RunAggregationQueryRequest runAggregationQueryFromGqlQuery = - protoPreparer.prepare(QueryConfig.create(gqlQueryWithoutNamespace)); + protoPreparer.prepare(QueryConfig.create(gqlQueryWithoutNamespace, QueryMode.NORMAL)); assertThat(runAggregationQueryFromStructuredQuery.getPartitionId().getNamespaceId()) .isEqualTo(NAMESPACE); assertThat(runAggregationQueryFromGqlQuery.getPartitionId().getNamespaceId()) .isEqualTo(NAMESPACE); + assertThat(runAggregationQueryFromStructuredQuery.getMode()).isEqualTo(QueryMode.NORMAL); + } + + @Test + public void shouldPrepareAggregationQueryWithDifferentModes() { + AggregationQuery structuredQueryWithoutNamespace = + Query.newAggregationQueryBuilder() + .addAggregation(count().as("total")) + .over(COMPLETED_TASK_STRUCTURED_QUERY) + .build(); + + RunAggregationQueryRequest runAggregationQueryFromStructuredQuery = + protoPreparer.prepare( + QueryConfig.create(structuredQueryWithoutNamespace, QueryMode.PROFILE)); + assertThat(runAggregationQueryFromStructuredQuery.getPartitionId().getNamespaceId()) + .isEqualTo(NAMESPACE); + assertThat(runAggregationQueryFromStructuredQuery.getMode()).isEqualTo(QueryMode.PROFILE); + + RunAggregationQueryRequest runAggregationQueryFromStructuredQuery2 = + protoPreparer.prepare(QueryConfig.create(structuredQueryWithoutNamespace, QueryMode.PLAN)); + assertThat(runAggregationQueryFromStructuredQuery2.getPartitionId().getNamespaceId()) + .isEqualTo(NAMESPACE); + assertThat(runAggregationQueryFromStructuredQuery2.getMode()).isEqualTo(QueryMode.PLAN); } private RunAggregationQueryRequest prepareQuery(AggregationQuery query, ReadOption readOption) { - return protoPreparer.prepare(QueryConfig.create(query, singletonList(readOption))); + return protoPreparer.prepare( + QueryConfig.create(query, QueryMode.NORMAL, singletonList(readOption))); } } 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 7ba57223f..c06414bad 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 @@ -24,6 +24,8 @@ import com.google.cloud.datastore.AggregationResults; import com.google.common.collect.ImmutableMap; import com.google.datastore.v1.AggregationResultBatch; +import com.google.datastore.v1.QueryPlan; +import com.google.datastore.v1.ResultSetStats; import com.google.datastore.v1.RunAggregationQueryResponse; import com.google.datastore.v1.Value; import java.util.AbstractMap.SimpleEntry; @@ -76,6 +78,50 @@ public void shouldTransformAggregationQueryResponseWithIntValues() { assertThat(aggregationResults.get(0)).isEqualTo(new AggregationResult(toDomainValues(result1))); assertThat(aggregationResults.get(1)).isEqualTo(new AggregationResult(toDomainValues(result2))); assertThat(aggregationResults.getReadTime()).isEqualTo(readTime); + assertThat(aggregationResults.getResultSetStats()).isEqualTo(null); + } + + @Test + public void shouldTransformAggregationQueryResponseWithIntValuesWithStats() { + 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)))); + 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(); + ResultSetStats stats = + ResultSetStats.newBuilder().setQueryPlan(QueryPlan.newBuilder().build()).build(); + RunAggregationQueryResponse runAggregationQueryResponse = + RunAggregationQueryResponse.newBuilder().setBatch(resultBatch).setStats(stats).build(); + + 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.getReadTime()).isEqualTo(readTime); + assertThat(aggregationResults.getResultSetStats()) + .isEqualTo(new com.google.cloud.datastore.models.ResultSetStats(stats)); } @Test 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 7c68ffe32..5ba5d9f9f 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 @@ -31,6 +31,7 @@ import com.google.cloud.Timestamp; import com.google.cloud.Tuple; import com.google.cloud.datastore.AggregationQuery; +import com.google.cloud.datastore.AggregationResults; import com.google.cloud.datastore.Batch; import com.google.cloud.datastore.BooleanValue; import com.google.cloud.datastore.Cursor; @@ -66,9 +67,13 @@ import com.google.cloud.datastore.TimestampValue; import com.google.cloud.datastore.Transaction; import com.google.cloud.datastore.ValueType; +import com.google.cloud.datastore.models.QueryProfile; +import com.google.cloud.datastore.models.QueryProfile.QueryMode; +import com.google.cloud.datastore.models.ResultSetStats; import com.google.cloud.datastore.testing.RemoteDatastoreHelper; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; +import com.google.common.truth.Truth; import com.google.datastore.v1.TransactionOptions; import com.google.datastore.v1.TransactionOptions.ReadOnly; import java.util.ArrayList; @@ -77,6 +82,7 @@ import java.util.HashSet; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.Set; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -358,6 +364,57 @@ public void orQuery() { assertFalse(results4.hasNext()); } + @Test + public void testQueryProfile() { + Key key = Key.newBuilder(KEY1, KIND2, 2).build(); + Entity entity3 = + Entity.newBuilder(ENTITY1) + .setKey(key) + .remove("str") + .set("name", "Dan") + .setNull("null") + .set("age", 19) + .build(); + datastore.put(entity3); + + // age == 19 || age == 20 + CompositeFilter orFilter = + CompositeFilter.or(PropertyFilter.eq("age", 19), PropertyFilter.eq("age", 20)); + StructuredQuery simpleOrQuery = + Query.newEntityQueryBuilder() + .setNamespace(NAMESPACE) + .setKind(KIND2) + .setFilter(orFilter) + .build(); + QueryResults results = + datastore.run(simpleOrQuery, QueryProfile.QueryMode.EXPLAIN_ANALYZE); + Truth.assertThat(results.hasNext()).isTrue(); + ResultSetStats resultSetStats = results.getResultSetStats(); + assertQueryPlan(resultSetStats); + assertQueryStats(resultSetStats); + + QueryResults results2 = datastore.run(simpleOrQuery, QueryProfile.QueryMode.EXPLAIN); + Truth.assertThat(results2.hasNext()).isFalse(); + ResultSetStats resultSetStats2 = results2.getResultSetStats(); + assertQueryPlan(resultSetStats2); + Truth.assertThat(resultSetStats2.getQueryStats()).isNull(); + + QueryResults results3 = datastore.run(simpleOrQuery, QueryProfile.QueryMode.NORMAL); + Truth.assertThat(results3.hasNext()).isTrue(); + Truth.assertThat(results3.getResultSetStats()).isNull(); + + AggregationQuery aggregationQuery = + Query.newAggregationQueryBuilder().over(simpleOrQuery).addAggregation(count()).build(); + AggregationResults resultsAggregation = + datastore.runAggregation(aggregationQuery, QueryProfile.QueryMode.EXPLAIN_ANALYZE); + + Truth.assertThat(resultsAggregation.size() > 0).isTrue(); + + ResultSetStats aggregationStats = resultsAggregation.getResultSetStats(); + assertQueryPlan(aggregationStats); + assertQueryStats(aggregationStats); + } + @Test public void testNewTransactionCommit() { Transaction transaction = datastore.newTransaction(); @@ -468,6 +525,117 @@ public void testTransactionWithQuery() throws Exception { assertThat(onlyOneTransactionIsSuccessful).isTrue(); } + @Test + public void testTransactionQueryModeExplainAnalyze() { + StructuredQuery query = + Query.newEntityQueryBuilder() + .setKind(KIND2) + .setFilter(PropertyFilter.hasAncestor(KEY2)) + .setNamespace(NAMESPACE) + .build(); + Transaction baseTransaction = datastore.newTransaction(); + QueryResults baseResults = + baseTransaction.run(query, QueryProfile.QueryMode.EXPLAIN_ANALYZE); + assertTrue(baseResults.hasNext()); + assertEquals(ENTITY2, baseResults.next()); + assertFalse(baseResults.hasNext()); + + ResultSetStats resultSetStats = baseResults.getResultSetStats(); + assertQueryPlan(resultSetStats); + assertQueryStats(resultSetStats); + + baseTransaction.add(ENTITY3); + baseTransaction.commit(); + assertEquals(ENTITY3, datastore.get(KEY3)); + + AggregationQuery aggregationQuery = + Query.newAggregationQueryBuilder().addAggregation(count()).over(query).build(); + + Transaction aggregationTransaction = datastore.newTransaction(); + AggregationResults results = + aggregationTransaction.runAggregation( + aggregationQuery, QueryProfile.QueryMode.EXPLAIN_ANALYZE); + assertTrue(results.size() > 0); + + ResultSetStats aggregationStats = results.getResultSetStats(); + assertQueryPlan(aggregationStats); + assertQueryStats(aggregationStats); + aggregationTransaction.commit(); + } + + @Test + public void testTransactionQueryModeExplain() { + StructuredQuery query = + Query.newEntityQueryBuilder() + .setKind(KIND2) + .setFilter(PropertyFilter.hasAncestor(KEY2)) + .setNamespace(NAMESPACE) + .build(); + Transaction baseTransaction = datastore.newTransaction(); + QueryResults baseResults = baseTransaction.run(query, QueryProfile.QueryMode.EXPLAIN); + assertFalse(baseResults.hasNext()); + + ResultSetStats resultSetStats = baseResults.getResultSetStats(); + assertQueryPlan(resultSetStats); + Truth.assertThat(resultSetStats.hasQueryStats()).isFalse(); + Truth.assertThat(resultSetStats.getQueryStats()).isNull(); + + AggregationQuery aggregationQuery = + Query.newAggregationQueryBuilder().addAggregation(count()).over(query).build(); + + Transaction aggregationTransaction = datastore.newTransaction(); + AggregationResults results = + aggregationTransaction.runAggregation(aggregationQuery, QueryProfile.QueryMode.EXPLAIN); + assertFalse(results.size() > 0); + + ResultSetStats aggregationStats = results.getResultSetStats(); + assertQueryPlan(aggregationStats); + Truth.assertThat(aggregationStats.hasQueryStats()).isFalse(); + Truth.assertThat(aggregationStats.getQueryStats()).isNull(); + } + + private void assertQueryPlan(ResultSetStats resultSetStats) { + Truth.assertThat(resultSetStats.getQueryPlan().getPlanInfo().keySet()).contains("indexes_used"); + } + + private void assertQueryStats(ResultSetStats resultSetStats) { + Map queryStatsAggregation = resultSetStats.getQueryStats(); + Truth.assertThat(queryStatsAggregation.keySet()) + .containsAtLeast( + "bytes_returned", + "total_execution_time", + "documents_scanned", + "index_entries_scanned", + "results_returned"); + } + + @Test + public void testTransactionQueryModeNormal() { + StructuredQuery query = + Query.newEntityQueryBuilder() + .setKind(KIND2) + .setFilter(PropertyFilter.hasAncestor(KEY2)) + .setNamespace(NAMESPACE) + .build(); + Transaction baseTransaction = datastore.newTransaction(); + QueryResults baseResults = baseTransaction.run(query, QueryProfile.QueryMode.NORMAL); + assertTrue(baseResults.hasNext()); + + Truth.assertThat(baseResults.getResultSetStats()).isNull(); + baseTransaction.commit(); + + AggregationQuery aggregationQuery = + Query.newAggregationQueryBuilder().addAggregation(count()).over(query).build(); + + Transaction aggregationTransaction = datastore.newTransaction(); + AggregationResults results = + aggregationTransaction.runAggregation(aggregationQuery, QueryProfile.QueryMode.NORMAL); + assertTrue(results.size() > 0); + + Truth.assertThat(results.getResultSetStats()).isNull(); + aggregationTransaction.commit(); + } + @Test public void testNewTransactionRollback() { Transaction transaction = datastore.newTransaction(); @@ -1646,10 +1814,13 @@ private void testCountAggregationReadTimeWith(Consumer getOnlyElement(datastore.runAggregation(countAggregationQuery)).getLong("total_count"); assertThat(latestCount).isEqualTo(3L); - Long oldCount = - getOnlyElement(datastore.runAggregation(countAggregationQuery, ReadOption.readTime(now))) - .getLong("total_count"); + AggregationResults results = + datastore.runAggregation( + countAggregationQuery, QueryMode.EXPLAIN_ANALYZE, ReadOption.readTime(now)); + Long oldCount = getOnlyElement(results).getLong("total_count"); assertThat(oldCount).isEqualTo(2L); + assertQueryPlan(results.getResultSetStats()); + assertQueryStats(results.getResultSetStats()); } finally { datastore.delete(entity1.getKey(), entity2.getKey(), entity3.getKey()); } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/QueryPlanTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/QueryPlanTest.java new file mode 100644 index 000000000..62e9903a6 --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/QueryPlanTest.java @@ -0,0 +1,66 @@ +/* + * Copyright 2023 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.models; + +import com.google.cloud.Structs; +import com.google.common.truth.Truth; +import com.google.protobuf.Struct; +import com.google.protobuf.Value; +import java.util.Map; +import org.junit.Test; + +public class QueryPlanTest { + + @Test + public void testQueryPlanModel() { + Struct planInfoStruct = + Struct.newBuilder() + .putFields("indexes_used", Value.newBuilder().setStringValue("val").build()) + .build(); + com.google.datastore.v1.QueryPlan proto = + com.google.datastore.v1.QueryPlan.newBuilder().setPlanInfo(planInfoStruct).build(); + QueryPlan queryPlan = new QueryPlan(proto); + + Map planInfo = queryPlan.getPlanInfo(); + Truth.assertThat(planInfo).isEqualTo(Structs.asMap(planInfoStruct)); + } + + @Test + public void testEqualsAndHashcode() { + Struct planInfoStruct = + Struct.newBuilder() + .putFields("indexes_used", Value.newBuilder().setStringValue("val").build()) + .build(); + com.google.datastore.v1.QueryPlan proto = + com.google.datastore.v1.QueryPlan.newBuilder().setPlanInfo(planInfoStruct).build(); + QueryPlan queryPlan1 = new QueryPlan(proto); + QueryPlan queryPlan2 = new QueryPlan(proto); + + Truth.assertThat(queryPlan1).isEqualTo(queryPlan2); + Truth.assertThat(queryPlan1.hashCode()).isEqualTo(queryPlan2.hashCode()); + + Struct planInfoStruct2 = + Struct.newBuilder() + .putFields("indexes_used", Value.newBuilder().setStringValue("12345").build()) + .build(); + com.google.datastore.v1.QueryPlan proto2 = + com.google.datastore.v1.QueryPlan.newBuilder().setPlanInfo(planInfoStruct2).build(); + QueryPlan queryPlan3 = new QueryPlan(proto2); + + Truth.assertThat(queryPlan3).isNotEqualTo(queryPlan1); + Truth.assertThat(queryPlan3.hashCode()).isNotEqualTo(queryPlan1.hashCode()); + } +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/QueryProfileTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/QueryProfileTest.java new file mode 100644 index 000000000..481d0bd18 --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/QueryProfileTest.java @@ -0,0 +1,33 @@ +/* + * Copyright 2023 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.models; + +import com.google.cloud.datastore.models.QueryProfile.QueryMode; +import com.google.common.truth.Truth; +import org.junit.Test; + +public class QueryProfileTest { + + @Test + public void testToPb() { + Truth.assertThat(QueryMode.NORMAL.toPb()).isEqualTo(com.google.datastore.v1.QueryMode.NORMAL); + Truth.assertThat(QueryMode.EXPLAIN_ANALYZE.toPb()) + .isEqualTo(com.google.datastore.v1.QueryMode.PROFILE); + Truth.assertThat(QueryMode.EXPLAIN.toPb()).isEqualTo(com.google.datastore.v1.QueryMode.PLAN); + Truth.assertThat(QueryMode.UNRECOGZNIED.toPb()) + .isEqualTo(com.google.datastore.v1.QueryMode.UNRECOGNIZED); + } +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/ResultSetStatsTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/ResultSetStatsTest.java new file mode 100644 index 000000000..166a7798c --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/ResultSetStatsTest.java @@ -0,0 +1,97 @@ +/* + * Copyright 2023 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.models; + +import com.google.cloud.Structs; +import com.google.common.truth.Truth; +import com.google.protobuf.Struct; +import com.google.protobuf.Value; +import org.junit.Test; + +public class ResultSetStatsTest { + + @Test + public void testResultSetStatsModel() { + Struct planInfoStruct = + Struct.newBuilder() + .putFields("indexes_used", Value.newBuilder().setStringValue("val").build()) + .build(); + com.google.datastore.v1.QueryPlan queryPlanProto = + com.google.datastore.v1.QueryPlan.newBuilder().setPlanInfo(planInfoStruct).build(); + + Struct statsProto = + Struct.newBuilder() + .putFields("bytes_returned", Value.newBuilder().setStringValue("bytes").build()) + .build(); + com.google.datastore.v1.ResultSetStats proto = + com.google.datastore.v1.ResultSetStats.newBuilder() + .setQueryStats(statsProto) + .setQueryPlan(queryPlanProto) + .build(); + ResultSetStats resultSetStats = new ResultSetStats(proto); + + Truth.assertThat(resultSetStats.getQueryPlan()).isEqualTo(new QueryPlan(queryPlanProto)); + Truth.assertThat(resultSetStats.getQueryStats()).isEqualTo(Structs.asMap(statsProto)); + } + + @Test + public void testEqualsAndHashcode() { + Struct planInfoStruct = + Struct.newBuilder() + .putFields("indexes_used", Value.newBuilder().setStringValue("val").build()) + .build(); + com.google.datastore.v1.QueryPlan queryPlanProto = + com.google.datastore.v1.QueryPlan.newBuilder().setPlanInfo(planInfoStruct).build(); + + Struct statsProto = + Struct.newBuilder() + .putFields("bytes_returned", Value.newBuilder().setStringValue("bytes").build()) + .build(); + com.google.datastore.v1.ResultSetStats proto = + com.google.datastore.v1.ResultSetStats.newBuilder() + .setQueryStats(statsProto) + .setQueryPlan(queryPlanProto) + .build(); + ResultSetStats resultSetStats = new ResultSetStats(proto); + ResultSetStats resultSetStats2 = new ResultSetStats(proto); + + Truth.assertThat(resultSetStats).isEqualTo(resultSetStats2); + Truth.assertThat(resultSetStats.hashCode()).isEqualTo(resultSetStats2.hashCode()); + + Struct planInfoStruct2 = + Struct.newBuilder() + .putFields("indexes_used", Value.newBuilder().setStringValue("val").build()) + .build(); + com.google.datastore.v1.QueryPlan queryPlanProto2 = + com.google.datastore.v1.QueryPlan.newBuilder().setPlanInfo(planInfoStruct2).build(); + + Struct statsProto2 = + Struct.newBuilder() + .putFields("bytes_returned", Value.newBuilder().setStringValue("other bytes").build()) + .build(); + com.google.datastore.v1.ResultSetStats proto2 = + com.google.datastore.v1.ResultSetStats.newBuilder() + .setQueryStats(statsProto2) + .setQueryPlan(queryPlanProto2) + .build(); + ResultSetStats resultSetStats3 = new ResultSetStats(proto2); + + Truth.assertThat(resultSetStats3).isNotEqualTo(resultSetStats); + Truth.assertThat(resultSetStats3).isNotEqualTo(resultSetStats2); + Truth.assertThat(resultSetStats3.hashCode()).isNotEqualTo(resultSetStats.hashCode()); + Truth.assertThat(resultSetStats3.hashCode()).isNotEqualTo(resultSetStats2.hashCode()); + } +} From ca6a85eb57f96621e4414d56df9e9539136553a7 Mon Sep 17 00:00:00 2001 From: kolea2 Date: Thu, 18 Jan 2024 13:32:26 -0500 Subject: [PATCH 2/5] fix clirr, run linter --- .../clirr-ignored-differences.xml | 27 ++++++++++--------- .../com/google/cloud/datastore/Datastore.java | 11 +++++--- .../cloud/datastore/DatastoreReader.java | 3 ++- 3 files changed, 23 insertions(+), 18 deletions(-) diff --git a/google-cloud-datastore/clirr-ignored-differences.xml b/google-cloud-datastore/clirr-ignored-differences.xml index 189df195e..111b828c3 100644 --- a/google-cloud-datastore/clirr-ignored-differences.xml +++ b/google-cloud-datastore/clirr-ignored-differences.xml @@ -1,19 +1,6 @@ - - - com/google/cloud/datastore/ReadOption$QueryAndReadOptions - * - 8001 - - - com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer - *QueryAndReadOptions* - *QueryConfig* - 7005 - - com/google/cloud/datastore/Datastore @@ -69,4 +56,18 @@ java.lang.Object execute(com.google.cloud.datastore.Query, com.google.cloud.datastore.ReadOption[]) 7004 + + + + com/google/cloud/datastore/ReadOption$QueryConfig + *QueryConfig* + *com.google.datastore.v1.QueryMode* + 7005 + + + com/google/cloud/datastore/ReadOption$QueryConfig + com.google.cloud.datastore.ReadOption$QueryConfig create(com.google.cloud.datastore.Query) + 7004 + + 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 5dcf55172..9f8543ac7 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 @@ -465,8 +465,9 @@ interface TransactionCallable { QueryResults run(Query query, ReadOption... options); /** - * Submits a {@link Query} with specified {@link com.google.cloud.datastore.models.QueryProfile.QueryMode} and returns its result. {@link ReadOption}s can be specified if - * desired. + * Submits a {@link Query} with specified {@link + * com.google.cloud.datastore.models.QueryProfile.QueryMode} and returns its result. {@link + * ReadOption}s can be specified if desired. * *

Example of running a query to find all entities of one kind. * @@ -535,8 +536,9 @@ default AggregationResults runAggregation(AggregationQuery query, ReadOption... } /** - * Submits a {@link AggregationQuery} with specified {@link com.google.cloud.datastore.models.QueryProfile.QueryMode} and returns {@link AggregationResults}. {@link ReadOption}s - * can be specified if desired. + * Submits a {@link AggregationQuery} with specified {@link + * com.google.cloud.datastore.models.QueryProfile.QueryMode} 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. * @@ -553,6 +555,7 @@ default AggregationResults runAggregation(AggregationQuery query, ReadOption... * AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery, QueryMode.EXPLAIN_ANALYZE); * ResultSetStats aggregationStats = aggregationResults.getResultSetStats(); * } + * * @throws DatastoreException upon failure * @return {@link AggregationResults} */ 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 11c270837..fbca10b20 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 @@ -65,7 +65,8 @@ default AggregationResults runAggregation(AggregationQuery query) { } /** - * Submits a {@link AggregationQuery} with a specified {@link com.google.cloud.datastore.models.QueryProfile.QueryMode} and returns {@link + * Submits a {@link AggregationQuery} with a specified {@link + * com.google.cloud.datastore.models.QueryProfile.QueryMode} and returns {@link * AggregationResults}. * * @throws DatastoreException upon failure From f23e1ecdf732bd543b683ae915b6a582ed58c275 Mon Sep 17 00:00:00 2001 From: kolea2 Date: Thu, 18 Jan 2024 13:48:26 -0500 Subject: [PATCH 3/5] review feedback --- .../src/main/java/com/google/cloud/datastore/Datastore.java | 2 +- .../main/java/com/google/cloud/datastore/QueryResultsImpl.java | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) 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 9f8543ac7..436f10198 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 @@ -485,7 +485,7 @@ interface TransactionCallable { @BetaApi default QueryResults run( Query query, QueryProfile.QueryMode queryMode, ReadOption... options) { - throw new UnsupportedOperationException("not implemented"); + throw new UnsupportedOperationException("Not implemented."); } /** 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 608060aab..8a4a37423 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 @@ -97,7 +97,8 @@ private void sendRequest() { boolean isQueryPlan = actualResultType.getQueryType() == null && queryMode == QueryMode.PLAN; Preconditions.checkState( queryResultType.isAssignableFrom(actualResultType) || isQueryPlan, - "Unexpected result type " + actualResultType + " vs " + queryResultType); + "Unexpected result type or query mode. Result type: " + actualResultType + " vs " + queryResultType +", query mode: " + queryMode); + if (runQueryResponsePb.getStats().hasQueryPlan() || runQueryResponsePb.getStats().hasQueryStats()) { this.resultSetStats = new ResultSetStats(runQueryResponsePb.getStats()); From 151ce495c2d85569d16c69e058aa903e633889db Mon Sep 17 00:00:00 2001 From: kolea2 Date: Thu, 18 Jan 2024 13:48:42 -0500 Subject: [PATCH 4/5] lint again --- .../java/com/google/cloud/datastore/QueryResultsImpl.java | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) 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 8a4a37423..1d91fac1f 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 @@ -97,7 +97,12 @@ private void sendRequest() { boolean isQueryPlan = actualResultType.getQueryType() == null && queryMode == QueryMode.PLAN; Preconditions.checkState( queryResultType.isAssignableFrom(actualResultType) || isQueryPlan, - "Unexpected result type or query mode. Result type: " + actualResultType + " vs " + queryResultType +", query mode: " + queryMode); + "Unexpected result type or query mode. Result type: " + + actualResultType + + " vs " + + queryResultType + + ", query mode: " + + queryMode); if (runQueryResponsePb.getStats().hasQueryPlan() || runQueryResponsePb.getStats().hasQueryStats()) { From 4a6dc23d055bdd30647778c391a9e202dc220d0c Mon Sep 17 00:00:00 2001 From: kolea2 Date: Thu, 18 Jan 2024 14:18:57 -0500 Subject: [PATCH 5/5] switch return type to optional for ResultSetStats --- .../clirr-ignored-differences.xml | 2 +- .../cloud/datastore/AggregationResults.java | 5 +-- .../google/cloud/datastore/QueryResults.java | 5 +-- .../cloud/datastore/QueryResultsImpl.java | 4 +-- .../datastore/models/ResultSetStats.java | 7 ++-- ...gregationQueryResponseTransformerTest.java | 4 +-- .../cloud/datastore/it/ITDatastoreTest.java | 32 +++++++++---------- .../datastore/models/ResultSetStatsTest.java | 2 +- 8 files changed, 32 insertions(+), 29 deletions(-) diff --git a/google-cloud-datastore/clirr-ignored-differences.xml b/google-cloud-datastore/clirr-ignored-differences.xml index 111b828c3..c115c41b3 100644 --- a/google-cloud-datastore/clirr-ignored-differences.xml +++ b/google-cloud-datastore/clirr-ignored-differences.xml @@ -14,7 +14,7 @@ com/google/cloud/datastore/QueryResults - com.google.cloud.datastore.models.ResultSetStats getResultSetStats() + java.util.Optional getResultSetStats() 7012 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 53b422889..613a1e10a 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,7 @@ import java.util.Iterator; import java.util.List; import java.util.Objects; +import java.util.Optional; /** * The result of an {@link AggregationQuery} query submission. Contains a List<{@link @@ -64,8 +65,8 @@ public int size() { /* * Returns the ResultSetStats if QueryMode is set to EXPLAIN or EXPLAIN_ANALYZE. Otherwise, returns null. */ - public ResultSetStats getResultSetStats() { - return this.resultSetStats; + public Optional getResultSetStats() { + return Optional.ofNullable(this.resultSetStats); } @InternalApi diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java index 911bcb2a9..09271c39d 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java @@ -20,6 +20,7 @@ import com.google.cloud.datastore.models.ResultSetStats; import com.google.datastore.v1.QueryResultBatch; import java.util.Iterator; +import java.util.Optional; /** * The result of a Google Cloud Datastore query submission. When the result is not typed it is @@ -77,7 +78,7 @@ public interface QueryResults extends Iterator { * Returns the {@link ResultSetStats} if {@link QueryMode} was set to EXPLAIN or EXPLAIN_ANALYZE. Otherwise, it will return null. */ @BetaApi - default ResultSetStats getResultSetStats() { - throw new UnsupportedOperationException("not implemented"); + default Optional getResultSetStats() { + throw new UnsupportedOperationException("Not implemented."); } } 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 1d91fac1f..3811022b9 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 @@ -149,7 +149,7 @@ public MoreResultsType getMoreResults() { @Override @BetaApi - public ResultSetStats getResultSetStats() { - return this.resultSetStats; + public Optional getResultSetStats() { + return Optional.ofNullable(this.resultSetStats); } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/ResultSetStats.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/ResultSetStats.java index 7843836d3..ea21e5635 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/ResultSetStats.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/models/ResultSetStats.java @@ -19,6 +19,7 @@ import com.google.cloud.Structs; import com.google.common.base.Objects; import java.util.Map; +import java.util.Optional; /* * Model class containing the planning and execution stats for the query. If QueryMode.EXPLAIN_ANALYZE was set, @@ -40,7 +41,7 @@ public ResultSetStats(com.google.datastore.v1.ResultSetStats proto) { } /* - * Returns the plan for the query, if EXPLAIN or EXPLAIN_ANALYZE was set. Otherwise, returns null. + * Returns the plan for the query. */ public QueryPlan getQueryPlan() { return this.queryPlan; @@ -49,8 +50,8 @@ public QueryPlan getQueryPlan() { /* * Returns the stats for the query if EXPLAIN_ANALYZE was set. Otherwise, returns null. */ - public Map getQueryStats() { - return this.queryStats; + public Optional> getQueryStats() { + return Optional.ofNullable(this.queryStats); } public boolean hasQueryStats() { 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 c06414bad..bb74b912c 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 @@ -78,7 +78,7 @@ public void shouldTransformAggregationQueryResponseWithIntValues() { assertThat(aggregationResults.get(0)).isEqualTo(new AggregationResult(toDomainValues(result1))); assertThat(aggregationResults.get(1)).isEqualTo(new AggregationResult(toDomainValues(result2))); assertThat(aggregationResults.getReadTime()).isEqualTo(readTime); - assertThat(aggregationResults.getResultSetStats()).isEqualTo(null); + assertThat(aggregationResults.getResultSetStats().isPresent()).isFalse(); } @Test @@ -120,7 +120,7 @@ public void shouldTransformAggregationQueryResponseWithIntValuesWithStats() { assertThat(aggregationResults.get(0)).isEqualTo(new AggregationResult(toDomainValues(result1))); assertThat(aggregationResults.get(1)).isEqualTo(new AggregationResult(toDomainValues(result2))); assertThat(aggregationResults.getReadTime()).isEqualTo(readTime); - assertThat(aggregationResults.getResultSetStats()) + assertThat(aggregationResults.getResultSetStats().get()) .isEqualTo(new com.google.cloud.datastore.models.ResultSetStats(stats)); } 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 5ba5d9f9f..00ec57194 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 @@ -389,19 +389,19 @@ public void testQueryProfile() { QueryResults results = datastore.run(simpleOrQuery, QueryProfile.QueryMode.EXPLAIN_ANALYZE); Truth.assertThat(results.hasNext()).isTrue(); - ResultSetStats resultSetStats = results.getResultSetStats(); + ResultSetStats resultSetStats = results.getResultSetStats().get(); assertQueryPlan(resultSetStats); assertQueryStats(resultSetStats); QueryResults results2 = datastore.run(simpleOrQuery, QueryProfile.QueryMode.EXPLAIN); Truth.assertThat(results2.hasNext()).isFalse(); - ResultSetStats resultSetStats2 = results2.getResultSetStats(); + ResultSetStats resultSetStats2 = results2.getResultSetStats().get(); assertQueryPlan(resultSetStats2); - Truth.assertThat(resultSetStats2.getQueryStats()).isNull(); + Truth.assertThat(resultSetStats2.getQueryStats().isPresent()).isFalse(); QueryResults results3 = datastore.run(simpleOrQuery, QueryProfile.QueryMode.NORMAL); Truth.assertThat(results3.hasNext()).isTrue(); - Truth.assertThat(results3.getResultSetStats()).isNull(); + Truth.assertThat(results3.getResultSetStats().isPresent()).isFalse(); AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder().over(simpleOrQuery).addAggregation(count()).build(); @@ -410,7 +410,7 @@ public void testQueryProfile() { Truth.assertThat(resultsAggregation.size() > 0).isTrue(); - ResultSetStats aggregationStats = resultsAggregation.getResultSetStats(); + ResultSetStats aggregationStats = resultsAggregation.getResultSetStats().get(); assertQueryPlan(aggregationStats); assertQueryStats(aggregationStats); } @@ -540,7 +540,7 @@ public void testTransactionQueryModeExplainAnalyze() { assertEquals(ENTITY2, baseResults.next()); assertFalse(baseResults.hasNext()); - ResultSetStats resultSetStats = baseResults.getResultSetStats(); + ResultSetStats resultSetStats = baseResults.getResultSetStats().get(); assertQueryPlan(resultSetStats); assertQueryStats(resultSetStats); @@ -557,7 +557,7 @@ public void testTransactionQueryModeExplainAnalyze() { aggregationQuery, QueryProfile.QueryMode.EXPLAIN_ANALYZE); assertTrue(results.size() > 0); - ResultSetStats aggregationStats = results.getResultSetStats(); + ResultSetStats aggregationStats = results.getResultSetStats().get(); assertQueryPlan(aggregationStats); assertQueryStats(aggregationStats); aggregationTransaction.commit(); @@ -575,10 +575,10 @@ public void testTransactionQueryModeExplain() { QueryResults baseResults = baseTransaction.run(query, QueryProfile.QueryMode.EXPLAIN); assertFalse(baseResults.hasNext()); - ResultSetStats resultSetStats = baseResults.getResultSetStats(); + ResultSetStats resultSetStats = baseResults.getResultSetStats().get(); assertQueryPlan(resultSetStats); Truth.assertThat(resultSetStats.hasQueryStats()).isFalse(); - Truth.assertThat(resultSetStats.getQueryStats()).isNull(); + Truth.assertThat(resultSetStats.getQueryStats().isPresent()).isFalse(); AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder().addAggregation(count()).over(query).build(); @@ -588,10 +588,10 @@ public void testTransactionQueryModeExplain() { aggregationTransaction.runAggregation(aggregationQuery, QueryProfile.QueryMode.EXPLAIN); assertFalse(results.size() > 0); - ResultSetStats aggregationStats = results.getResultSetStats(); + ResultSetStats aggregationStats = results.getResultSetStats().get(); assertQueryPlan(aggregationStats); Truth.assertThat(aggregationStats.hasQueryStats()).isFalse(); - Truth.assertThat(aggregationStats.getQueryStats()).isNull(); + Truth.assertThat(aggregationStats.getQueryStats().isPresent()).isFalse(); } private void assertQueryPlan(ResultSetStats resultSetStats) { @@ -599,7 +599,7 @@ private void assertQueryPlan(ResultSetStats resultSetStats) { } private void assertQueryStats(ResultSetStats resultSetStats) { - Map queryStatsAggregation = resultSetStats.getQueryStats(); + Map queryStatsAggregation = resultSetStats.getQueryStats().get(); Truth.assertThat(queryStatsAggregation.keySet()) .containsAtLeast( "bytes_returned", @@ -621,7 +621,7 @@ public void testTransactionQueryModeNormal() { QueryResults baseResults = baseTransaction.run(query, QueryProfile.QueryMode.NORMAL); assertTrue(baseResults.hasNext()); - Truth.assertThat(baseResults.getResultSetStats()).isNull(); + Truth.assertThat(baseResults.getResultSetStats().isPresent()).isFalse(); baseTransaction.commit(); AggregationQuery aggregationQuery = @@ -632,7 +632,7 @@ public void testTransactionQueryModeNormal() { aggregationTransaction.runAggregation(aggregationQuery, QueryProfile.QueryMode.NORMAL); assertTrue(results.size() > 0); - Truth.assertThat(results.getResultSetStats()).isNull(); + Truth.assertThat(results.getResultSetStats().isPresent()).isFalse(); aggregationTransaction.commit(); } @@ -1819,8 +1819,8 @@ private void testCountAggregationReadTimeWith(Consumer countAggregationQuery, QueryMode.EXPLAIN_ANALYZE, ReadOption.readTime(now)); Long oldCount = getOnlyElement(results).getLong("total_count"); assertThat(oldCount).isEqualTo(2L); - assertQueryPlan(results.getResultSetStats()); - assertQueryStats(results.getResultSetStats()); + assertQueryPlan(results.getResultSetStats().get()); + assertQueryStats(results.getResultSetStats().get()); } finally { datastore.delete(entity1.getKey(), entity2.getKey(), entity3.getKey()); } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/ResultSetStatsTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/ResultSetStatsTest.java index 166a7798c..fea6e4f1b 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/ResultSetStatsTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/models/ResultSetStatsTest.java @@ -44,7 +44,7 @@ public void testResultSetStatsModel() { ResultSetStats resultSetStats = new ResultSetStats(proto); Truth.assertThat(resultSetStats.getQueryPlan()).isEqualTo(new QueryPlan(queryPlanProto)); - Truth.assertThat(resultSetStats.getQueryStats()).isEqualTo(Structs.asMap(statsProto)); + Truth.assertThat(resultSetStats.getQueryStats().get()).isEqualTo(Structs.asMap(statsProto)); } @Test