From f1311d9e6791d81d6d44c066237ded4c4ab9d5d5 Mon Sep 17 00:00:00 2001 From: Rahul Mane Date: Mon, 13 Jul 2026 13:07:26 +0000 Subject: [PATCH 1/4] feat(datastore): add support for request tags --- .../com/google/cloud/datastore/Datastore.java | 85 ++++++++ .../google/cloud/datastore/DatastoreImpl.java | 184 +++++++++++------ .../cloud/datastore/DatastoreOptions.java | 31 ++- .../cloud/datastore/QueryResultsImpl.java | 46 +++-- .../google/cloud/datastore/ReadOption.java | 32 ++- .../cloud/datastore/RequestOptionsHelper.java | 46 +++++ .../execution/AggregationQueryExecutor.java | 20 +- .../AggregationQueryRequestProtoPreparer.java | 7 + .../datastore/AbstractDatastoreTest.java | 189 ++++++++++++++++++ ...regationQueryRequestProtoPreparerTest.java | 17 ++ 10 files changed, 573 insertions(+), 84 deletions(-) create mode 100644 java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RequestOptionsHelper.java diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java index 0e769a109372..1c3a836903eb 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java @@ -20,6 +20,8 @@ import com.google.api.core.InternalExtensionOnly; import com.google.cloud.Service; import com.google.cloud.datastore.models.ExplainOptions; +import com.google.common.collect.ImmutableList; +import com.google.datastore.v1.RequestOptions; import com.google.datastore.v1.TransactionOptions; import java.util.Iterator; import java.util.List; @@ -492,6 +494,30 @@ interface TransactionCallable { @BetaApi QueryResults run(Query query, ExplainOptions explainOptions, ReadOption... options); + /** + * Submits a {@link Query} with specified {@link com.google.datastore.v1.RequestOptions} and + * returns its result. + */ + QueryResults run(Query query, RequestOptions requestOptions); + + /** + * Submits a {@link Query} with specified {@link com.google.datastore.v1.RequestOptions} and + * returns its result. {@link ReadOption}s can be specified if desired. + */ + QueryResults run(Query query, RequestOptions requestOptions, ReadOption... options); + + /** + * Submits a {@link Query} with specified {@link com.google.cloud.datastore.models.ExplainOptions} + * and {@link com.google.datastore.v1.RequestOptions} and returns its result. {@link ReadOption}s + * can be specified if desired. + */ + @BetaApi + QueryResults run( + Query query, + ExplainOptions explainOptions, + RequestOptions requestOptions, + ReadOption... options); + /** * Submits a {@link AggregationQuery} and returns {@link AggregationResults}. {@link ReadOption}s * can be specified if desired. @@ -537,6 +563,20 @@ interface TransactionCallable { */ AggregationResults runAggregation(AggregationQuery query, ReadOption... options); + /** + * Submits an {@link AggregationQuery} with specified {@link + * com.google.datastore.v1.RequestOptions} and returns {@link AggregationResults}. + */ + AggregationResults runAggregation(AggregationQuery query, RequestOptions requestOptions); + + /** + * Submits an {@link AggregationQuery} with specified {@link + * com.google.datastore.v1.RequestOptions} and returns {@link AggregationResults}. {@link + * ReadOption}s can be specified if desired. + */ + AggregationResults runAggregation( + AggregationQuery query, RequestOptions requestOptions, ReadOption... options); + /** * Submits a {@link AggregationQuery} with specified {@link * com.google.cloud.datastore.models.ExplainOptions} and returns {@link AggregationResults}. @@ -564,6 +604,28 @@ interface TransactionCallable { AggregationResults runAggregation( AggregationQuery query, ExplainOptions explainOptions, ReadOption... options); + /** + * Submits an {@link AggregationQuery} with specified {@link + * com.google.cloud.datastore.models.ExplainOptions} and {@link + * com.google.datastore.v1.RequestOptions} and returns {@link AggregationResults}. + */ + @BetaApi + AggregationResults runAggregation( + AggregationQuery query, ExplainOptions explainOptions, RequestOptions requestOptions); + + /** + * Submits an {@link AggregationQuery} with specified {@link + * com.google.cloud.datastore.models.ExplainOptions} and {@link + * com.google.datastore.v1.RequestOptions} and returns {@link AggregationResults}. {@link + * ReadOption}s can be specified if desired. + */ + @BetaApi + AggregationResults runAggregation( + AggregationQuery query, + ExplainOptions explainOptions, + RequestOptions requestOptions, + ReadOption... options); + /** * Closes the gRPC channels associated with this instance and frees up their resources. This * method blocks until all channels are closed. Once this method is called, this Datastore client @@ -574,4 +636,27 @@ AggregationResults runAggregation( /** Returns true if this background resource has been shut down. */ boolean isClosed(); + + /** + * Returns a new Datastore client with the specified request tags added. + * + * @param requestTags the request tags to append to existing ones + */ + default Datastore withRequestTags(String... requestTags) { + ImmutableList.Builder builder = ImmutableList.builder(); + builder.addAll(getOptions().getRequestTags()).add(requestTags); + return getOptions().toBuilder().setTags(builder.build()).build().getService(); + } + + /** + * Returns a new Datastore client with the specified request tags added. + * + * @param requestTags the request tags to append to existing ones + */ + default Datastore withRequestTags(List requestTags) { + ImmutableList.Builder builder = ImmutableList.builder(); + builder.addAll(getOptions().getRequestTags()); + builder.addAll(requestTags); + return getOptions().toBuilder().setTags(builder.build()).build().getService(); + } } diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index 700b5e4faf4a..5dd9799f2734 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -61,11 +61,24 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; import com.google.common.collect.Sets; +import com.google.datastore.v1.AllocateIdsRequest; +import com.google.datastore.v1.AllocateIdsResponse; +import com.google.datastore.v1.BeginTransactionRequest; +import com.google.datastore.v1.BeginTransactionResponse; +import com.google.datastore.v1.CommitRequest; import com.google.datastore.v1.CommitResponse; +import com.google.datastore.v1.EntityResult; import com.google.datastore.v1.ExplainOptions; +import com.google.datastore.v1.LookupRequest; +import com.google.datastore.v1.LookupResponse; +import com.google.datastore.v1.Mutation; +import com.google.datastore.v1.MutationResult; import com.google.datastore.v1.ReadOptions; +import com.google.datastore.v1.RequestOptions; import com.google.datastore.v1.ReserveIdsRequest; +import com.google.datastore.v1.ReserveIdsResponse; import com.google.datastore.v1.RollbackRequest; +import com.google.datastore.v1.RunQueryRequest; import com.google.datastore.v1.RunQueryResponse; import com.google.datastore.v1.TransactionOptions; import com.google.protobuf.ByteString; @@ -73,7 +86,6 @@ import io.opentelemetry.api.OpenTelemetry; import io.opentelemetry.context.Context; import io.opentelemetry.sdk.OpenTelemetrySdk; -import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.Iterator; @@ -333,11 +345,49 @@ public QueryResults run( return run(toReadOptionsPb(options), query, explainOptions.toPb()); } + /** Implements run query with request options. */ + @Override + public QueryResults run(Query query, RequestOptions requestOptions) { + return run(Optional.empty(), query, null, requestOptions); + } + + /** Implements run query with request options and read options. */ + @Override + public QueryResults run( + Query query, RequestOptions requestOptions, ReadOption... options) { + return run(toReadOptionsPb(options), query, null, requestOptions); + } + + /** Implements run query with explain options, request options and read options. */ + @Override + @BetaApi + public QueryResults run( + Query query, + com.google.cloud.datastore.models.ExplainOptions explainOptions, + RequestOptions requestOptions, + ReadOption... options) { + return run(toReadOptionsPb(options), query, explainOptions.toPb(), requestOptions); + } + @SuppressWarnings("unchecked") QueryResults run( Optional readOptionsPb, Query query, ExplainOptions explainOptions) { + return run(readOptionsPb, query, explainOptions, null); + } + + @SuppressWarnings("unchecked") + QueryResults run( + Optional readOptionsPb, + Query query, + ExplainOptions explainOptions, + RequestOptions requestOptions) { return new QueryResultsImpl( - this, readOptionsPb, (RecordQuery) query, query.getNamespace(), explainOptions); + this, + readOptionsPb, + (RecordQuery) query, + query.getNamespace(), + explainOptions, + requestOptions); } @Override @@ -345,11 +395,24 @@ public AggregationResults runAggregation(AggregationQuery query) { return aggregationQueryExecutor.execute(query, null); } + /** Implements run aggregation query with request options. */ + @Override + public AggregationResults runAggregation(AggregationQuery query, RequestOptions requestOptions) { + return aggregationQueryExecutor.execute(query, null, requestOptions); + } + @Override public AggregationResults runAggregation(AggregationQuery query, ReadOption... options) { return aggregationQueryExecutor.execute(query, null, options); } + /** Implements run aggregation query with request options and read options. */ + @Override + public AggregationResults runAggregation( + AggregationQuery query, RequestOptions requestOptions, ReadOption... options) { + return aggregationQueryExecutor.execute(query, null, requestOptions, options); + } + @Override @BetaApi public AggregationResults runAggregation( @@ -357,6 +420,16 @@ public AggregationResults runAggregation( return aggregationQueryExecutor.execute(query, explainOptions); } + /** Implements run aggregation query with explain options and request options. */ + @Override + @BetaApi + public AggregationResults runAggregation( + AggregationQuery query, + com.google.cloud.datastore.models.ExplainOptions explainOptions, + RequestOptions requestOptions) { + return aggregationQueryExecutor.execute(query, explainOptions, requestOptions); + } + @Override @BetaApi public AggregationResults runAggregation( @@ -366,8 +439,18 @@ public AggregationResults runAggregation( return aggregationQueryExecutor.execute(query, explainOptions, options); } - com.google.datastore.v1.RunQueryResponse runQuery( - final com.google.datastore.v1.RunQueryRequest requestPb) { + /** Implements run aggregation query with explain options, request options and read options. */ + @Override + @BetaApi + public AggregationResults runAggregation( + AggregationQuery query, + com.google.cloud.datastore.models.ExplainOptions explainOptions, + RequestOptions requestOptions, + ReadOption... options) { + return aggregationQueryExecutor.execute(query, explainOptions, requestOptions, options); + } + + RunQueryResponse runQuery(final RunQueryRequest requestPb) { ReadOptions readOptions = requestPb.getReadOptions(); boolean isTransactional = readOptions.hasTransaction() || readOptions.hasNewTransaction(); String spanName = (isTransactional ? SPAN_NAME_TRANSACTION_RUN_QUERY : SPAN_NAME_RUN_QUERY); @@ -425,14 +508,13 @@ public List allocateId(IncompleteKey... keys) { if (keys.length == 0) { return Collections.emptyList(); } - com.google.datastore.v1.AllocateIdsRequest.Builder requestPb = - com.google.datastore.v1.AllocateIdsRequest.newBuilder(); + AllocateIdsRequest.Builder requestPb = AllocateIdsRequest.newBuilder(); for (IncompleteKey key : keys) { requestPb.addKeys(trimNameOrId(key).toPb()); } requestPb.setProjectId(getOptions().getProjectId()); requestPb.setDatabaseId(getOptions().getDatabaseId()); - com.google.datastore.v1.AllocateIdsResponse responsePb = allocateIds(requestPb.build()); + AllocateIdsResponse responsePb = allocateIds(requestPb.build()); ImmutableList.Builder keyList = ImmutableList.builder(); for (com.google.datastore.v1.Key keyPb : responsePb.getKeysList()) { keyList.add(Key.fromPb(keyPb)); @@ -440,8 +522,7 @@ public List allocateId(IncompleteKey... keys) { return keyList.build(); } - private com.google.datastore.v1.AllocateIdsResponse allocateIds( - final com.google.datastore.v1.AllocateIdsRequest requestPb) { + private AllocateIdsResponse allocateIds(final AllocateIdsRequest requestPb) { return runWithObservability( () -> datastoreRpc.allocateIds(requestPb), TelemetryConstants.METHOD_ALLOCATE_IDS, @@ -467,7 +548,7 @@ public List add(FullEntity... entities) { if (entities.length == 0) { return Collections.emptyList(); } - List mutationsPb = new ArrayList<>(); + ImmutableList.Builder mutationsPb = ImmutableList.builder(); Map completeEntities = new LinkedHashMap<>(); for (FullEntity entity : entities) { Entity completeEntity = null; @@ -482,12 +563,10 @@ public List add(FullEntity... entities) { } else { Preconditions.checkArgument(entity.hasKey(), "Entity %s is missing a key", entity); } - mutationsPb.add( - com.google.datastore.v1.Mutation.newBuilder().setInsert(entity.toPb()).build()); + mutationsPb.add(Mutation.newBuilder().setInsert(entity.toPb()).build()); } - com.google.datastore.v1.CommitResponse commitResponse = commitMutation(mutationsPb); - Iterator mutationResults = - commitResponse.getMutationResultsList().iterator(); + CommitResponse commitResponse = commitMutation(mutationsPb.build()); + Iterator mutationResults = commitResponse.getMutationResultsList().iterator(); ImmutableList.Builder responseBuilder = ImmutableList.builder(); for (FullEntity entity : entities) { Entity completeEntity = completeEntities.get(entity.getKey()); @@ -543,8 +622,7 @@ Iterator get(Optional readOptionsPb, final Key... keys) { if (keys.length == 0) { return Collections.emptyIterator(); } - com.google.datastore.v1.LookupRequest.Builder requestPb = - com.google.datastore.v1.LookupRequest.newBuilder(); + LookupRequest.Builder requestPb = LookupRequest.newBuilder(); readOptionsPb.ifPresent(requestPb::setReadOptions); for (Key k : Sets.newLinkedHashSet(Arrays.asList(keys))) { requestPb.addKeys(k.toPb()); @@ -556,16 +634,16 @@ Iterator get(Optional readOptionsPb, final Key... keys) { final class ResultsIterator extends AbstractIterator { - private final com.google.datastore.v1.LookupRequest.Builder requestPb; - Iterator iter; + private final LookupRequest.Builder requestPb; + Iterator iter; - ResultsIterator(com.google.datastore.v1.LookupRequest.Builder requestPb) { + ResultsIterator(LookupRequest.Builder requestPb) { this.requestPb = requestPb; loadResults(); } private void loadResults() { - com.google.datastore.v1.LookupResponse responsePb = lookup(requestPb.build()); + LookupResponse responsePb = lookup(requestPb.build()); iter = responsePb.getFoundList().iterator(); requestPb.clearKeys(); if (responsePb.getDeferredCount() > 0) { @@ -586,15 +664,14 @@ protected Entity computeNext() { } } - com.google.datastore.v1.LookupResponse lookup( - final com.google.datastore.v1.LookupRequest requestPb) { + LookupResponse lookup(final LookupRequest requestPb) { ReadOptions readOptions = requestPb.getReadOptions(); boolean isTransactional = readOptions.hasTransaction() || readOptions.hasNewTransaction(); String spanName = (isTransactional ? SPAN_NAME_TRANSACTION_LOOKUP : SPAN_NAME_LOOKUP); return runWithObservability( () -> { - com.google.datastore.v1.LookupResponse response = datastoreRpc.lookup(requestPb); + LookupResponse response = datastoreRpc.lookup(requestPb); com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.getCurrentSpan(); if (span != null) { span.addEvent( @@ -626,7 +703,7 @@ public List reserveIds(Key... keys) { } requestPb.setProjectId(getOptions().getProjectId()); requestPb.setDatabaseId(getOptions().getDatabaseId()); - com.google.datastore.v1.ReserveIdsResponse responsePb = reserveIds(requestPb.build()); + ReserveIdsResponse responsePb = reserveIds(requestPb.build()); ImmutableList.Builder keyList = ImmutableList.builder(); if (responsePb.isInitialized()) { for (Key key : keys) { @@ -636,8 +713,7 @@ public List reserveIds(Key... keys) { return keyList.build(); } - com.google.datastore.v1.ReserveIdsResponse reserveIds( - final com.google.datastore.v1.ReserveIdsRequest requestPb) { + ReserveIdsResponse reserveIds(final ReserveIdsRequest requestPb) { return runWithObservability( () -> datastoreRpc.reserveIds(requestPb), TelemetryConstants.METHOD_RESERVE_IDS, @@ -648,16 +724,15 @@ com.google.datastore.v1.ReserveIdsResponse reserveIds( @Override public void update(Entity... entities) { if (entities.length > 0) { - List mutationsPb = new ArrayList<>(); + ImmutableList.Builder mutationsPb = ImmutableList.builder(); Map dedupEntities = new LinkedHashMap<>(); for (Entity entity : entities) { dedupEntities.put(entity.getKey(), entity); } for (Entity entity : dedupEntities.values()) { - mutationsPb.add( - com.google.datastore.v1.Mutation.newBuilder().setUpdate(entity.toPb()).build()); + mutationsPb.add(Mutation.newBuilder().setUpdate(entity.toPb()).build()); } - commitMutation(mutationsPb); + commitMutation(mutationsPb.build()); } } @@ -672,7 +747,7 @@ public List put(FullEntity... entities) { if (entities.length == 0) { return Collections.emptyList(); } - List mutationsPb = new ArrayList<>(); + ImmutableList.Builder mutationsPb = ImmutableList.builder(); Map dedupEntities = new LinkedHashMap<>(); for (FullEntity entity : entities) { Preconditions.checkArgument(entity.hasKey(), "Entity %s is missing a key", entity); @@ -680,17 +755,14 @@ public List put(FullEntity... entities) { Entity completeEntity = Entity.convert((FullEntity) entity); dedupEntities.put(completeEntity.getKey(), completeEntity); } else { - mutationsPb.add( - com.google.datastore.v1.Mutation.newBuilder().setUpsert(entity.toPb()).build()); + mutationsPb.add(Mutation.newBuilder().setUpsert(entity.toPb()).build()); } } for (Entity entity : dedupEntities.values()) { - mutationsPb.add( - com.google.datastore.v1.Mutation.newBuilder().setUpsert(entity.toPb()).build()); + mutationsPb.add(Mutation.newBuilder().setUpsert(entity.toPb()).build()); } - com.google.datastore.v1.CommitResponse commitResponse = commitMutation(mutationsPb); - Iterator mutationResults = - commitResponse.getMutationResultsList().iterator(); + CommitResponse commitResponse = commitMutation(mutationsPb.build()); + Iterator mutationResults = commitResponse.getMutationResultsList().iterator(); ImmutableList.Builder responseBuilder = ImmutableList.builder(); for (FullEntity entity : entities) { Entity completeEntity = dedupEntities.get(entity.getKey()); @@ -707,13 +779,12 @@ public List put(FullEntity... entities) { @Override public void delete(Key... keys) { if (keys.length > 0) { - List mutationsPb = new ArrayList<>(); + ImmutableList.Builder mutationsPb = ImmutableList.builder(); Set dedupKeys = new LinkedHashSet<>(Arrays.asList(keys)); for (Key key : dedupKeys) { - mutationsPb.add( - com.google.datastore.v1.Mutation.newBuilder().setDelete(key.toPb()).build()); + mutationsPb.add(Mutation.newBuilder().setDelete(key.toPb()).build()); } - commitMutation(mutationsPb); + commitMutation(mutationsPb.build()); } } @@ -722,19 +793,17 @@ public KeyFactory newKeyFactory() { return DatastoreHelper.newKeyFactory(getOptions()); } - private com.google.datastore.v1.CommitResponse commitMutation( - List mutationsPb) { - com.google.datastore.v1.CommitRequest.Builder requestPb = - com.google.datastore.v1.CommitRequest.newBuilder(); - requestPb.setMode(com.google.datastore.v1.CommitRequest.Mode.NON_TRANSACTIONAL); - requestPb.setProjectId(getOptions().getProjectId()); - requestPb.setDatabaseId(getOptions().getDatabaseId()); - requestPb.addAllMutations(mutationsPb); + private CommitResponse commitMutation(ImmutableList mutationsPb) { + CommitRequest.Builder requestPb = + CommitRequest.newBuilder() + .setMode(CommitRequest.Mode.NON_TRANSACTIONAL) + .setProjectId(getOptions().getProjectId()) + .setDatabaseId(getOptions().getDatabaseId()) + .addAllMutations(mutationsPb); return commit(requestPb.build()); } - com.google.datastore.v1.CommitResponse commit( - final com.google.datastore.v1.CommitRequest requestPb) { + CommitResponse commit(final CommitRequest requestPb) { final boolean isTransactional = requestPb.hasTransaction() || requestPb.hasSingleUseTransaction(); final String spanName = isTransactional ? SPAN_NAME_TRANSACTION_COMMIT : SPAN_NAME_COMMIT; @@ -763,13 +832,11 @@ com.google.datastore.v1.CommitResponse commit( : TRANSACTION_OPERATION_EXCEPTION_HANDLER); } - ByteString requestTransactionId( - com.google.datastore.v1.BeginTransactionRequest.Builder requestPb) { + ByteString requestTransactionId(BeginTransactionRequest.Builder requestPb) { return beginTransaction(requestPb.build()).getTransaction(); } - com.google.datastore.v1.BeginTransactionResponse beginTransaction( - final com.google.datastore.v1.BeginTransactionRequest requestPb) { + BeginTransactionResponse beginTransaction(final BeginTransactionRequest requestPb) { return runWithObservability( () -> datastoreRpc.beginTransaction(requestPb), TelemetryConstants.METHOD_BEGIN_TRANSACTION, @@ -778,8 +845,7 @@ com.google.datastore.v1.BeginTransactionResponse beginTransaction( } void rollbackTransaction(ByteString transaction) { - com.google.datastore.v1.RollbackRequest.Builder requestPb = - com.google.datastore.v1.RollbackRequest.newBuilder(); + RollbackRequest.Builder requestPb = RollbackRequest.newBuilder(); requestPb.setTransaction(transaction); requestPb.setProjectId(getOptions().getProjectId()); requestPb.setDatabaseId(getOptions().getDatabaseId()); diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java index 48d4adbf065f..c66969543e73 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java @@ -30,11 +30,13 @@ import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.cloud.datastore.spi.v1.GrpcDatastoreRpc; import com.google.cloud.datastore.spi.v1.HttpDatastoreRpc; +import com.google.cloud.datastore.telemetry.TraceUtil; import com.google.cloud.datastore.v1.DatastoreSettings; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.http.HttpTransportOptions; import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import java.io.IOException; import java.lang.reflect.Method; @@ -80,9 +82,10 @@ public class DatastoreOptions extends ServiceOptions requestTags; private final transient @Nonnull DatastoreOpenTelemetryOptions openTelemetryOptions; - private final transient @Nonnull com.google.cloud.datastore.telemetry.TraceUtil traceUtil; + @Nonnull private final transient TraceUtil traceUtil; public static class DefaultDatastoreFactory implements DatastoreFactory { @@ -113,7 +116,7 @@ public ServiceRpc create(DatastoreOptions options) { } @Nonnull - com.google.cloud.datastore.telemetry.TraceUtil getTraceUtil() { + TraceUtil getTraceUtil() { return traceUtil; } @@ -127,6 +130,7 @@ public static class Builder extends ServiceOptions.Builder requestTags = ImmutableList.of(); private TransportChannelProvider channelProvider = null; private String host; @@ -142,6 +146,7 @@ private Builder(DatastoreOptions options) { super(options); this.namespace = options.namespace; this.databaseId = options.databaseId; + this.requestTags = options.requestTags; this.openTelemetryOptions = options.openTelemetryOptions; this.channelProvider = validateChannelProvider(options.channelProvider); this.host = options.getHost(); @@ -251,6 +256,16 @@ public Builder setDatabaseId(String databaseId) { return this; } + public Builder setTags(ImmutableList requestTags) { + this.requestTags = requestTags; + return this; + } + + public Builder setTags(String... requestTags) { + this.requestTags = ImmutableList.copyOf(requestTags); + return this; + } + /** * Sets the {@link DatastoreOpenTelemetryOptions} to be used for this Datastore instance. * @@ -272,10 +287,11 @@ private DatastoreOptions(Builder builder) { builder.openTelemetryOptions != null ? builder.openTelemetryOptions : DatastoreOpenTelemetryOptions.newBuilder().build(); - this.traceUtil = com.google.cloud.datastore.telemetry.TraceUtil.getInstance(this); + this.traceUtil = TraceUtil.getInstance(this); namespace = MoreObjects.firstNonNull(builder.namespace, defaultNamespace()); databaseId = MoreObjects.firstNonNull(builder.databaseId, DEFAULT_DATABASE_ID); + requestTags = builder.requestTags; // ChannelProvider is used by GAX but HttpJson does not use it so we safely ignore it. if (getTransportOptions() instanceof HttpTransportOptions && builder.channelProvider != null) { @@ -370,6 +386,10 @@ public String getDatabaseId() { return this.databaseId; } + public ImmutableList getRequestTags() { + return requestTags; + } + /** Returns a default {@code DatastoreOptions} instance. */ public static DatastoreOptions getDefaultInstance() { return newBuilder().build(); @@ -404,7 +424,7 @@ public Builder toBuilder() { @Override public int hashCode() { - return Objects.hash(baseHashCode(), namespace, databaseId); + return Objects.hash(baseHashCode(), namespace, databaseId, requestTags); } @Override @@ -415,7 +435,8 @@ public boolean equals(Object obj) { DatastoreOptions other = (DatastoreOptions) obj; return baseEquals(other) && Objects.equals(namespace, other.namespace) - && Objects.equals(databaseId, other.databaseId); + && Objects.equals(databaseId, other.databaseId) + && Objects.equals(requestTags, other.requestTags); } public static Builder newBuilder() { diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java index ad51129891b5..37cbd6ce4d6f 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java @@ -21,9 +21,14 @@ import com.google.cloud.datastore.models.ExplainMetrics; import com.google.common.base.Preconditions; import com.google.common.collect.AbstractIterator; +import com.google.datastore.v1.EntityResult; import com.google.datastore.v1.ExplainOptions; +import com.google.datastore.v1.PartitionId; import com.google.datastore.v1.QueryResultBatch.MoreResultsType; import com.google.datastore.v1.ReadOptions; +import com.google.datastore.v1.RequestOptions; +import com.google.datastore.v1.RunQueryRequest; +import com.google.datastore.v1.RunQueryResponse; import com.google.protobuf.ByteString; import java.util.Iterator; import java.util.Objects; @@ -33,32 +38,45 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults private final DatastoreImpl datastore; private final Optional readOptionsPb; - private final com.google.datastore.v1.PartitionId partitionIdPb; + private final PartitionId partitionIdPb; private final ResultType queryResultType; private RecordQuery query; private ResultType actualResultType; - private com.google.datastore.v1.RunQueryResponse runQueryResponsePb; + private RunQueryResponse runQueryResponsePb; private com.google.datastore.v1.Query mostRecentQueryPb; private boolean lastBatch; - private Iterator entityResultPbIter; + private Iterator entityResultPbIter; private ByteString cursor; private MoreResultsType moreResults; private final ExplainOptions explainOptions; private ExplainMetrics explainMetrics; + private final RequestOptions requestOptions; + /** Creates a QueryResultsImpl. */ QueryResultsImpl( DatastoreImpl datastore, Optional readOptionsPb, RecordQuery query, String namespace, ExplainOptions explainOptions) { + this(datastore, readOptionsPb, query, namespace, explainOptions, null); + } + + /** Creates a QueryResultsImpl with RequestOptions. */ + QueryResultsImpl( + DatastoreImpl datastore, + Optional readOptionsPb, + RecordQuery query, + String namespace, + ExplainOptions explainOptions, + RequestOptions requestOptions) { this.datastore = datastore; this.readOptionsPb = readOptionsPb; this.query = query; queryResultType = query.getType(); this.explainOptions = explainOptions; - com.google.datastore.v1.PartitionId.Builder pbBuilder = - com.google.datastore.v1.PartitionId.newBuilder(); + this.requestOptions = requestOptions; + PartitionId.Builder pbBuilder = PartitionId.newBuilder(); pbBuilder.setProjectId(datastore.getOptions().getProjectId()); pbBuilder.setDatabaseId(datastore.getOptions().getDatabaseId()); if (namespace != null) { @@ -76,8 +94,7 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults } private void sendRequest() { - com.google.datastore.v1.RunQueryRequest.Builder requestPb = - com.google.datastore.v1.RunQueryRequest.newBuilder(); + RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder(); readOptionsPb.ifPresent(requestPb::setReadOptions); requestPb.setPartitionId(partitionIdPb); requestPb.setProjectId(datastore.getOptions().getProjectId()); @@ -85,6 +102,10 @@ private void sendRequest() { if (explainOptions != null) { requestPb.setExplainOptions(explainOptions); } + if (requestOptions != null) { + requestPb.setRequestOptions( + RequestOptionsHelper.createRequestOptions(datastore.getOptions(), requestOptions)); + } query.populatePb(requestPb); runQueryResponsePb = datastore.runQuery(requestPb.build()); mostRecentQueryPb = requestPb.getQuery(); @@ -102,11 +123,9 @@ private void sendRequest() { && !explainOptions.getAnalyze(); Preconditions.checkState( queryResultType.isAssignableFrom(actualResultType) || isExplain, - "Unexpected result type or explain options set " - + actualResultType - + " vs " - + queryResultType - + ", explain options = false"); + "Unexpected result type or explain options set %s vs %s, explain options = false", + actualResultType, + queryResultType); if (runQueryResponsePb.hasExplainMetrics()) { this.explainMetrics = new ExplainMetrics(runQueryResponsePb.getExplainMetrics()); } @@ -122,8 +141,9 @@ protected T computeNext() { cursor = runQueryResponsePb.getBatch().getEndCursor(); return endOfData(); } - com.google.datastore.v1.EntityResult entityResultPb = entityResultPbIter.next(); + EntityResult entityResultPb = entityResultPbIter.next(); cursor = entityResultPb.getCursor(); + // Safe because we verify actualResultType compatibility in sendRequest(). @SuppressWarnings("unchecked") T result = (T) actualResultType.convert(entityResultPb.getEntity()); return result; diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java index f7d6d82c76a2..793a48a39db1 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java @@ -21,6 +21,7 @@ import com.google.cloud.Timestamp; import com.google.common.collect.ImmutableMap; import com.google.datastore.v1.ExplainOptions; +import com.google.datastore.v1.RequestOptions; import com.google.protobuf.ByteString; import java.io.Serializable; import java.util.Collections; @@ -142,21 +143,32 @@ static Map, ReadOption> asImmutableMap(List> { Q query; List readOptions; ExplainOptions explainOptions; + RequestOptions requestOptions; private QueryConfig(Q query, ExplainOptions explainOptions, List readOptions) { - this.query = query; - this.explainOptions = explainOptions; - this.readOptions = readOptions; + this(query, explainOptions, readOptions, null); } private QueryConfig(Q query, ExplainOptions explainOptions) { - this(query, explainOptions, Collections.emptyList()); + this(query, explainOptions, Collections.emptyList(), null); + } + + private QueryConfig( + Q query, + ExplainOptions explainOptions, + List readOptions, + RequestOptions requestOptions) { + this.query = query; + this.explainOptions = explainOptions; + this.readOptions = readOptions; + this.requestOptions = requestOptions; } public Q getQuery() { @@ -171,6 +183,10 @@ public List getReadOptions() { return readOptions; } + public RequestOptions getRequestOptions() { + return requestOptions; + } + public static > QueryConfig create( Q query, ExplainOptions explainOptions) { return new QueryConfig<>(query, explainOptions); @@ -180,5 +196,13 @@ public static > QueryConfig create( Q query, ExplainOptions explainOptions, List readOptions) { return new QueryConfig<>(query, explainOptions, readOptions); } + + public static > QueryConfig create( + Q query, + ExplainOptions explainOptions, + List readOptions, + RequestOptions requestOptions) { + return new QueryConfig<>(query, explainOptions, readOptions, requestOptions); + } } } diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RequestOptionsHelper.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RequestOptionsHelper.java new file mode 100644 index 000000000000..7771c23a4290 --- /dev/null +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RequestOptionsHelper.java @@ -0,0 +1,46 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.datastore; + +import com.google.datastore.v1.RequestOptions; + +/** Helper class for building and merging Datastore request options. */ +public final class RequestOptionsHelper { + + private RequestOptionsHelper() {} + + /** + * Creates a merged {@link RequestOptions} instance containing both the instance-level tags and + * the request-level options. + * + * @param datastoreOptions the global Datastore options containing instance-level tags + * @param requestOptions the individual request-level options, or {@code null} + * @return the merged RequestOptions proto message + */ + public static RequestOptions createRequestOptions( + DatastoreOptions datastoreOptions, RequestOptions requestOptions) { + RequestOptions.Builder builder = RequestOptions.newBuilder(); + if (requestOptions != null) { + // copy existing options. This copies the request level tags as well. + builder.mergeFrom(requestOptions); + } + + // add instance-level tags. + builder.addAllRequestTags(datastoreOptions.getRequestTags()); + return builder.build(); + } +} diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java index bedcf34e4da1..f5a1c30933d8 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java @@ -25,6 +25,7 @@ import com.google.cloud.datastore.execution.response.AggregationQueryResponseTransformer; import com.google.cloud.datastore.models.ExplainOptions; import com.google.cloud.datastore.spi.v1.DatastoreRpc; +import com.google.datastore.v1.RequestOptions; import com.google.datastore.v1.RunAggregationQueryRequest; import com.google.datastore.v1.RunAggregationQueryResponse; import java.util.Arrays; @@ -50,9 +51,20 @@ public AggregationQueryExecutor(DatastoreRpc datastoreRpc, DatastoreOptions data @Override public AggregationResults execute( AggregationQuery query, ExplainOptions explainOptions, ReadOption... readOptions) { + return execute(query, explainOptions, null, readOptions); + } + + public AggregationResults execute( + AggregationQuery query, + ExplainOptions explainOptions, + RequestOptions requestOptions, + ReadOption... readOptions) { RunAggregationQueryRequest runAggregationQueryRequest = getRunAggregationQueryRequest( - query, explainOptions == null ? null : explainOptions.toPb(), readOptions); + query, + explainOptions == null ? null : explainOptions.toPb(), + requestOptions, + readOptions); RunAggregationQueryResponse runAggregationQueryResponse = this.datastoreRpc.runAggregationQuery(runAggregationQueryRequest); return this.responseTransformer.transform(runAggregationQueryResponse); @@ -61,11 +73,13 @@ public AggregationResults execute( private RunAggregationQueryRequest getRunAggregationQueryRequest( AggregationQuery query, com.google.datastore.v1.ExplainOptions explainOptions, + RequestOptions requestOptions, ReadOption... readOptions) { QueryConfig queryConfig = readOptions == null - ? QueryConfig.create(query, explainOptions) - : QueryConfig.create(query, explainOptions, Arrays.asList(readOptions)); + ? QueryConfig.create( + query, explainOptions, java.util.Collections.emptyList(), requestOptions) + : QueryConfig.create(query, explainOptions, Arrays.asList(readOptions), requestOptions); return this.protoPreparer.prepare(queryConfig); } } diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java index a89ebb11a18c..2e0db97a3bb6 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java @@ -16,6 +16,7 @@ package com.google.cloud.datastore.execution.request; import static com.google.cloud.datastore.AggregationQuery.Mode.GQL; +import static com.google.cloud.datastore.RequestOptionsHelper.createRequestOptions; import com.google.api.core.InternalApi; import com.google.cloud.datastore.AggregationQuery; @@ -72,6 +73,12 @@ public RunAggregationQueryRequest prepare(QueryConfig queryCon } Optional readOptionsPb = readOptionProtoPreparer.prepare(readOptions); readOptionsPb.ifPresent(aggregationQueryRequestBuilder::setReadOptions); + + // add request options to the request. + if (queryConfig.getRequestOptions() != null) { + aggregationQueryRequestBuilder.setRequestOptions( + createRequestOptions(datastoreOptions, queryConfig.getRequestOptions())); + } return aggregationQueryRequestBuilder.build(); } diff --git a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java index 9d4ecac57a00..e4cc6e1e23b3 100644 --- a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java +++ b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java @@ -1399,4 +1399,193 @@ private Predicate aggregationQueryWithAlias(String a .get(0) .getAlias()); } + + @Test + public void testRunQueryWithRequestOptions() { + com.google.datastore.v1.RequestOptions requestOptions = + com.google.datastore.v1.RequestOptions.newBuilder().addRequestTags("test-tag").build(); + + PartitionId partitionId = + PartitionId.newBuilder() + .setProjectId(PROJECT_ID) + .setDatabaseId(options.getDatabaseId()) + .setNamespaceId(options.getNamespace()) + .build(); + + com.google.datastore.v1.Query queryPb = + com.google.datastore.v1.Query.newBuilder() + .addKind(com.google.datastore.v1.KindExpression.newBuilder().setName(KIND1)) + .build(); + + // 1. RunQueryRequest for run(query, requestOptions) + RunQueryRequest expectedRequest1 = + RunQueryRequest.newBuilder() + .setProjectId(PROJECT_ID) + .setPartitionId(partitionId) + .setQuery(queryPb) + .setRequestOptions(requestOptions) + .build(); + + // 2. RunQueryRequest for run(query, requestOptions, ReadOption) + RunQueryRequest expectedRequest2 = + RunQueryRequest.newBuilder() + .setProjectId(PROJECT_ID) + .setPartitionId(partitionId) + .setQuery(queryPb) + .setReadOptions( + ReadOptions.newBuilder().setReadConsistency(ReadConsistency.EVENTUAL).build()) + .setRequestOptions(requestOptions) + .build(); + + // 3. RunQueryRequest for run(query, explainOptions, requestOptions, ReadOption) + RunQueryRequest expectedRequest3 = + RunQueryRequest.newBuilder() + .setProjectId(PROJECT_ID) + .setPartitionId(partitionId) + .setQuery(queryPb) + .setReadOptions( + ReadOptions.newBuilder().setReadConsistency(ReadConsistency.EVENTUAL).build()) + .setExplainOptions( + com.google.datastore.v1.ExplainOptions.newBuilder().setAnalyze(true).build()) + .setRequestOptions(requestOptions) + .build(); + + RunQueryResponse response = + RunQueryResponse.newBuilder() + .setBatch( + com.google.datastore.v1.QueryResultBatch.newBuilder() + .addEntityResults(EntityResult.newBuilder().setEntity(ENTITY1.toPb())) + .build()) + .build(); + + EasyMock.expect(rpcMock.runQuery(expectedRequest1)).andReturn(response); + EasyMock.expect(rpcMock.runQuery(expectedRequest2)).andReturn(response); + EasyMock.expect(rpcMock.runQuery(expectedRequest3)).andReturn(response); + EasyMock.replay(rpcFactoryMock, rpcMock); + + Datastore targetDatastore = rpcMockOptions.getService(); + Query query = Query.newEntityQueryBuilder().setKind(KIND1).build(); + + // Test 1: run(query, requestOptions) + QueryResults results1 = targetDatastore.run(query, requestOptions); + assertTrue(results1.hasNext()); + assertEquals(ENTITY1, results1.next()); + + // Test 2: run(query, requestOptions, readOptions) + QueryResults results2 = + targetDatastore.run(query, requestOptions, ReadOption.eventualConsistency()); + assertTrue(results2.hasNext()); + assertEquals(ENTITY1, results2.next()); + + // Test 3: run(query, explainOptions, requestOptions, readOptions) + QueryResults results3 = + targetDatastore.run( + query, + com.google.cloud.datastore.models.ExplainOptions.newBuilder().setAnalyze(true).build(), + requestOptions, + ReadOption.eventualConsistency()); + assertTrue(results3.hasNext()); + assertEquals(ENTITY1, results3.next()); + + EasyMock.verify(rpcFactoryMock, rpcMock); + } + + @Test + public void testRunAggregationQueryWithRequestOptions() { + com.google.datastore.v1.RequestOptions requestOptions = + com.google.datastore.v1.RequestOptions.newBuilder().addRequestTags("agg-test-tag").build(); + + RunAggregationQueryResponse aggregationQueryResponse = placeholderAggregationQueryResponse(); + + // 1. runAggregation(query, requestOptions) + EasyMock.expect( + rpcMock.runAggregationQuery( + matches( + aggregationQueryWithAliasAndRequestOptions( + "total_count", requestOptions, false, false)))) + .andReturn(aggregationQueryResponse); + + // 2. runAggregation(query, requestOptions, options) + EasyMock.expect( + rpcMock.runAggregationQuery( + matches( + aggregationQueryWithAliasAndRequestOptions( + "total_count", requestOptions, true, false)))) + .andReturn(aggregationQueryResponse); + + // 3. runAggregation(query, explainOptions, requestOptions) + EasyMock.expect( + rpcMock.runAggregationQuery( + matches( + aggregationQueryWithAliasAndRequestOptions( + "total_count", requestOptions, false, true)))) + .andReturn(aggregationQueryResponse); + + // 4. runAggregation(query, explainOptions, requestOptions, options) + EasyMock.expect( + rpcMock.runAggregationQuery( + matches( + aggregationQueryWithAliasAndRequestOptions( + "total_count", requestOptions, true, true)))) + .andReturn(aggregationQueryResponse); + + EasyMock.replay(rpcFactoryMock, rpcMock); + + Datastore mockDatastore = rpcMockOptions.getService(); + + EntityQuery selectAllQuery = Query.newEntityQueryBuilder().build(); + AggregationQuery getCountQuery = + Query.newAggregationQueryBuilder() + .addAggregation(count().as("total_count")) + .over(selectAllQuery) + .build(); + + com.google.cloud.datastore.models.ExplainOptions explainOptions = + com.google.cloud.datastore.models.ExplainOptions.newBuilder().setAnalyze(true).build(); + + // Test 1 + AggregationResult result1 = + getOnlyElement(mockDatastore.runAggregation(getCountQuery, requestOptions)); + assertThat(result1.getLong("total_count")).isEqualTo(209L); + + // Test 2 + AggregationResult result2 = + getOnlyElement( + mockDatastore.runAggregation( + getCountQuery, requestOptions, ReadOption.eventualConsistency())); + assertThat(result2.getLong("total_count")).isEqualTo(209L); + + // Test 3 + AggregationResult result3 = + getOnlyElement(mockDatastore.runAggregation(getCountQuery, explainOptions, requestOptions)); + assertThat(result3.getLong("total_count")).isEqualTo(209L); + + // Test 4 + AggregationResult result4 = + getOnlyElement( + mockDatastore.runAggregation( + getCountQuery, explainOptions, requestOptions, ReadOption.eventualConsistency())); + assertThat(result4.getLong("total_count")).isEqualTo(209L); + + EasyMock.verify(rpcFactoryMock, rpcMock); + } + + private Predicate aggregationQueryWithAliasAndRequestOptions( + String alias, + com.google.datastore.v1.RequestOptions requestOptions, + boolean checkEventual, + boolean checkExplain) { + return req -> { + boolean match = + alias.equals(req.getAggregationQuery().getAggregationsList().get(0).getAlias()) + && requestOptions.equals(req.getRequestOptions()); + if (checkEventual) { + match = match && req.getReadOptions().getReadConsistency() == ReadConsistency.EVENTUAL; + } + if (checkExplain) { + match = match && req.hasExplainOptions(); + } + return match; + }; + } } diff --git a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java index 4aec8f7bba43..ec34197ef7cb 100644 --- a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java +++ b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparerTest.java @@ -41,6 +41,7 @@ import com.google.cloud.datastore.ReadOption.QueryConfig; import com.google.common.collect.ImmutableMap; import com.google.datastore.v1.ExplainOptions; +import com.google.datastore.v1.RequestOptions; import com.google.datastore.v1.RunAggregationQueryRequest; import java.util.HashMap; import org.junit.Test; @@ -213,4 +214,20 @@ public void shouldPrepareAggregationQueryWithDifferentModes() { private RunAggregationQueryRequest prepareQuery(AggregationQuery query, ReadOption readOption) { return protoPreparer.prepare(QueryConfig.create(query, null, singletonList(readOption))); } + + @Test + public void shouldPrepareAggregationQueryRequestWithRequestOptions() { + RequestOptions requestOptions = + RequestOptions.newBuilder().addRequestTags("aggregation-tag").build(); + + RunAggregationQueryRequest request = + protoPreparer.prepare( + QueryConfig.create( + AGGREGATION_OVER_STRUCTURED_QUERY, + /* explainOptions= */ null, + /* readOptions= */ java.util.Collections.emptyList(), + requestOptions)); + + assertThat(request.getRequestOptions().getRequestTagsList()).containsExactly("aggregation-tag"); + } } From 5f377f8b6b7efcce0c1a99d541193c04523e6162 Mon Sep 17 00:00:00 2001 From: Rahul Mane Date: Mon, 13 Jul 2026 13:38:18 +0000 Subject: [PATCH 2/4] fix: address code review comments regarding instance-level tags and duplication --- .../com/google/cloud/datastore/Datastore.java | 4 +- .../cloud/datastore/QueryResultsImpl.java | 2 +- .../AggregationQueryRequestProtoPreparer.java | 2 +- .../datastore/AbstractDatastoreTest.java | 45 +++++++++++++++++++ 4 files changed, 48 insertions(+), 5 deletions(-) diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java index 1c3a836903eb..c00eda6efc1b 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java @@ -643,9 +643,7 @@ AggregationResults runAggregation( * @param requestTags the request tags to append to existing ones */ default Datastore withRequestTags(String... requestTags) { - ImmutableList.Builder builder = ImmutableList.builder(); - builder.addAll(getOptions().getRequestTags()).add(requestTags); - return getOptions().toBuilder().setTags(builder.build()).build().getService(); + return withRequestTags(java.util.Arrays.asList(requestTags)); } /** diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java index 37cbd6ce4d6f..8312129d7d50 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java @@ -102,7 +102,7 @@ private void sendRequest() { if (explainOptions != null) { requestPb.setExplainOptions(explainOptions); } - if (requestOptions != null) { + if (requestOptions != null || !datastore.getOptions().getRequestTags().isEmpty()) { requestPb.setRequestOptions( RequestOptionsHelper.createRequestOptions(datastore.getOptions(), requestOptions)); } diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java index 2e0db97a3bb6..a07ce14c172b 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer.java @@ -75,7 +75,7 @@ public RunAggregationQueryRequest prepare(QueryConfig queryCon readOptionsPb.ifPresent(aggregationQueryRequestBuilder::setReadOptions); // add request options to the request. - if (queryConfig.getRequestOptions() != null) { + if (queryConfig.getRequestOptions() != null || !datastoreOptions.getRequestTags().isEmpty()) { aggregationQueryRequestBuilder.setRequestOptions( createRequestOptions(datastoreOptions, queryConfig.getRequestOptions())); } diff --git a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java index e4cc6e1e23b3..b93022c76f6f 100644 --- a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java +++ b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java @@ -1400,6 +1400,51 @@ private Predicate aggregationQueryWithAlias(String a .getAlias()); } + @Test + public void testRunQueryWithInstanceLevelRequestTags() { + DatastoreOptions optionsWithTags = options.toBuilder().setTags("instance-tag").build(); + Datastore datastoreWithTags = optionsWithTags.getService(); + + PartitionId partitionId = + PartitionId.newBuilder() + .setProjectId(PROJECT_ID) + .setDatabaseId(options.getDatabaseId()) + .setNamespaceId(options.getNamespace()) + .build(); + + com.google.datastore.v1.Query queryPb = + com.google.datastore.v1.Query.newBuilder() + .addKind(com.google.datastore.v1.KindExpression.newBuilder().setName(KIND1)) + .build(); + + RunQueryRequest expectedRequest = + RunQueryRequest.newBuilder() + .setProjectId(PROJECT_ID) + .setPartitionId(partitionId) + .setQuery(queryPb) + .setRequestOptions( + com.google.datastore.v1.RequestOptions.newBuilder().addRequestTags("instance-tag").build()) + .build(); + + RunQueryResponse response = + RunQueryResponse.newBuilder() + .setBatch( + com.google.datastore.v1.QueryResultBatch.newBuilder() + .addEntityResults(EntityResult.newBuilder().setEntity(ENTITY1.toPb())) + .build()) + .build(); + + EasyMock.expect(rpcMock.runQuery(expectedRequest)).andReturn(response); + EasyMock.replay(rpcFactoryMock, rpcMock); + + Query query = Query.newEntityQueryBuilder().setKind(KIND1).build(); + QueryResults results = datastoreWithTags.run(query); + assertTrue(results.hasNext()); + assertEquals(ENTITY1, results.next()); + + EasyMock.verify(rpcFactoryMock, rpcMock); + } + @Test public void testRunQueryWithRequestOptions() { com.google.datastore.v1.RequestOptions requestOptions = From e5396c080fa04b9e6cf1379425f65d2d89ce1b33 Mon Sep 17 00:00:00 2001 From: Rahul Mane Date: Tue, 14 Jul 2026 06:24:41 +0000 Subject: [PATCH 3/4] Fixed review comments --- .../com/google/cloud/datastore/Datastore.java | 11 +-- .../google/cloud/datastore/DatastoreImpl.java | 15 ++-- .../cloud/datastore/DatastoreOptions.java | 23 +++++- .../cloud/datastore/QueryResultsImpl.java | 80 +++++++++++++------ .../google/cloud/datastore/ReadOption.java | 64 ++++++++++++--- .../cloud/datastore/RequestOptionsHelper.java | 2 + .../datastore/AbstractDatastoreTest.java | 38 ++++++++- .../cloud/datastore/DatastoreOptionsTest.java | 14 ++++ 8 files changed, 189 insertions(+), 58 deletions(-) diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java index c00eda6efc1b..9f67182a9c47 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java @@ -637,15 +637,6 @@ AggregationResults runAggregation( /** Returns true if this background resource has been shut down. */ boolean isClosed(); - /** - * Returns a new Datastore client with the specified request tags added. - * - * @param requestTags the request tags to append to existing ones - */ - default Datastore withRequestTags(String... requestTags) { - return withRequestTags(java.util.Arrays.asList(requestTags)); - } - /** * Returns a new Datastore client with the specified request tags added. * @@ -655,6 +646,6 @@ default Datastore withRequestTags(List requestTags) { ImmutableList.Builder builder = ImmutableList.builder(); builder.addAll(getOptions().getRequestTags()); builder.addAll(requestTags); - return getOptions().toBuilder().setTags(builder.build()).build().getService(); + return getOptions().toBuilder().setRequestTags(builder.build()).build().getService(); } } diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index 5dd9799f2734..923a09182f5f 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -381,13 +381,14 @@ QueryResults run( Query query, ExplainOptions explainOptions, RequestOptions requestOptions) { - return new QueryResultsImpl( - this, - readOptionsPb, - (RecordQuery) query, - query.getNamespace(), - explainOptions, - requestOptions); + return new QueryResultsImpl.Builder() + .setDatastore(this) + .setReadOptionsPb(readOptionsPb) + .setQuery((RecordQuery) query) + .setNamespace(query.getNamespace()) + .setExplainOptions(explainOptions) + .setRequestOptions(requestOptions) + .build(); } @Override diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java index c66969543e73..6a1e092d373e 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java @@ -256,12 +256,26 @@ public Builder setDatabaseId(String databaseId) { return this; } - public Builder setTags(ImmutableList requestTags) { + /** + * Sets the request tags to be associated with all requests sent by this client. + * + * @param requestTags the list of request tags to set + * @return the builder object + */ + public Builder setRequestTags(ImmutableList requestTags) { + Preconditions.checkNotNull(requestTags, "Request tags cannot be null"); this.requestTags = requestTags; return this; } - public Builder setTags(String... requestTags) { + /** + * Sets the request tags to be associated with all requests sent by this client. + * + * @param requestTags the request tags to set + * @return the builder object + */ + public Builder setRequestTags(String... requestTags) { + Preconditions.checkNotNull(requestTags, "Request tags cannot be null"); this.requestTags = ImmutableList.copyOf(requestTags); return this; } @@ -386,6 +400,11 @@ public String getDatabaseId() { return this.databaseId; } + /** + * Returns the request tags to be associated with all requests sent by this client. + * + * @return the request tags + */ public ImmutableList getRequestTags() { return requestTags; } diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java index 8312129d7d50..504fc66b33ca 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java @@ -52,35 +52,18 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults private ExplainMetrics explainMetrics; private final RequestOptions requestOptions; - /** Creates a QueryResultsImpl. */ - QueryResultsImpl( - DatastoreImpl datastore, - Optional readOptionsPb, - RecordQuery query, - String namespace, - ExplainOptions explainOptions) { - this(datastore, readOptionsPb, query, namespace, explainOptions, null); - } - - /** Creates a QueryResultsImpl with RequestOptions. */ - QueryResultsImpl( - DatastoreImpl datastore, - Optional readOptionsPb, - RecordQuery query, - String namespace, - ExplainOptions explainOptions, - RequestOptions requestOptions) { - this.datastore = datastore; - this.readOptionsPb = readOptionsPb; - this.query = query; - queryResultType = query.getType(); - this.explainOptions = explainOptions; - this.requestOptions = requestOptions; + private QueryResultsImpl(Builder builder) { + this.datastore = builder.datastore; + this.readOptionsPb = builder.readOptionsPb; + this.query = builder.query; + queryResultType = builder.query.getType(); + this.explainOptions = builder.explainOptions; + this.requestOptions = builder.requestOptions; PartitionId.Builder pbBuilder = PartitionId.newBuilder(); pbBuilder.setProjectId(datastore.getOptions().getProjectId()); pbBuilder.setDatabaseId(datastore.getOptions().getDatabaseId()); - if (namespace != null) { - pbBuilder.setNamespaceId(namespace); + if (builder.namespace != null) { + pbBuilder.setNamespaceId(builder.namespace); } else if (datastore.getOptions().getNamespace() != null) { pbBuilder.setNamespaceId(datastore.getOptions().getNamespace()); } @@ -93,6 +76,51 @@ class QueryResultsImpl extends AbstractIterator implements QueryResults } } + static class Builder { + private DatastoreImpl datastore; + private Optional readOptionsPb = Optional.empty(); + private RecordQuery query; + private String namespace; + private ExplainOptions explainOptions; + private RequestOptions requestOptions; + + Builder setDatastore(DatastoreImpl datastore) { + this.datastore = datastore; + return this; + } + + Builder setReadOptionsPb(Optional readOptionsPb) { + this.readOptionsPb = readOptionsPb; + return this; + } + + Builder setQuery(RecordQuery query) { + this.query = query; + return this; + } + + Builder setNamespace(String namespace) { + this.namespace = namespace; + return this; + } + + Builder setExplainOptions(ExplainOptions explainOptions) { + this.explainOptions = explainOptions; + return this; + } + + Builder setRequestOptions(RequestOptions requestOptions) { + this.requestOptions = requestOptions; + return this; + } + + QueryResultsImpl build() { + Preconditions.checkNotNull(datastore, "datastore cannot be null"); + Preconditions.checkNotNull(query, "query cannot be null"); + return new QueryResultsImpl<>(this); + } + } + private void sendRequest() { RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder(); readOptionsPb.ifPresent(requestPb::setReadOptions); diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java index 793a48a39db1..7d01730b566d 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java @@ -19,6 +19,7 @@ import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; import com.google.cloud.Timestamp; +import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; import com.google.datastore.v1.ExplainOptions; import com.google.datastore.v1.RequestOptions; @@ -152,14 +153,6 @@ public static class QueryConfig> { ExplainOptions explainOptions; RequestOptions requestOptions; - private QueryConfig(Q query, ExplainOptions explainOptions, List readOptions) { - this(query, explainOptions, readOptions, null); - } - - private QueryConfig(Q query, ExplainOptions explainOptions) { - this(query, explainOptions, Collections.emptyList(), null); - } - private QueryConfig( Q query, ExplainOptions explainOptions, @@ -189,12 +182,16 @@ public RequestOptions getRequestOptions() { public static > QueryConfig create( Q query, ExplainOptions explainOptions) { - return new QueryConfig<>(query, explainOptions); + return QueryConfig.newBuilder().setQuery(query).setExplainOptions(explainOptions).build(); } public static > QueryConfig create( Q query, ExplainOptions explainOptions, List readOptions) { - return new QueryConfig<>(query, explainOptions, readOptions); + return QueryConfig.newBuilder() + .setQuery(query) + .setExplainOptions(explainOptions) + .setReadOptions(readOptions) + .build(); } public static > QueryConfig create( @@ -202,7 +199,52 @@ public static > QueryConfig create( ExplainOptions explainOptions, List readOptions, RequestOptions requestOptions) { - return new QueryConfig<>(query, explainOptions, readOptions, requestOptions); + return QueryConfig.newBuilder() + .setQuery(query) + .setExplainOptions(explainOptions) + .setReadOptions(readOptions) + .setRequestOptions(requestOptions) + .build(); + } + + /** Creates a new builder for {@link QueryConfig}. */ + public static > Builder newBuilder() { + return new Builder<>(); + } + + /** Builder for {@link QueryConfig}. */ + public static class Builder> { + private Q query; + private List readOptions = Collections.emptyList(); + private ExplainOptions explainOptions; + private RequestOptions requestOptions; + + private Builder() {} + + public Builder setQuery(Q query) { + this.query = query; + return this; + } + + public Builder setReadOptions(List readOptions) { + this.readOptions = readOptions; + return this; + } + + public Builder setExplainOptions(ExplainOptions explainOptions) { + this.explainOptions = explainOptions; + return this; + } + + public Builder setRequestOptions(RequestOptions requestOptions) { + this.requestOptions = requestOptions; + return this; + } + + public QueryConfig build() { + Preconditions.checkNotNull(query, "query cannot be null"); + return new QueryConfig<>(query, explainOptions, readOptions, requestOptions); + } } } } diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RequestOptionsHelper.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RequestOptionsHelper.java index 7771c23a4290..7f7ec09c513b 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RequestOptionsHelper.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RequestOptionsHelper.java @@ -16,9 +16,11 @@ package com.google.cloud.datastore; +import com.google.api.core.InternalApi; import com.google.datastore.v1.RequestOptions; /** Helper class for building and merging Datastore request options. */ +@InternalApi public final class RequestOptionsHelper { private RequestOptionsHelper() {} diff --git a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java index b93022c76f6f..c3466b156fc7 100644 --- a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java +++ b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java @@ -1402,7 +1402,7 @@ private Predicate aggregationQueryWithAlias(String a @Test public void testRunQueryWithInstanceLevelRequestTags() { - DatastoreOptions optionsWithTags = options.toBuilder().setTags("instance-tag").build(); + DatastoreOptions optionsWithTags = options.toBuilder().setRequestTags("instance-tag").build(); Datastore datastoreWithTags = optionsWithTags.getService(); PartitionId partitionId = @@ -1423,7 +1423,9 @@ public void testRunQueryWithInstanceLevelRequestTags() { .setPartitionId(partitionId) .setQuery(queryPb) .setRequestOptions( - com.google.datastore.v1.RequestOptions.newBuilder().addRequestTags("instance-tag").build()) + com.google.datastore.v1.RequestOptions.newBuilder() + .addRequestTags("instance-tag") + .build()) .build(); RunQueryResponse response = @@ -1535,6 +1537,38 @@ public void testRunQueryWithRequestOptions() { EasyMock.verify(rpcFactoryMock, rpcMock); } + @Test + public void testRunAggregationQueryWithInstanceLevelRequestTags() { + DatastoreOptions optionsWithTags = options.toBuilder().setRequestTags("instance-tag").build(); + Datastore datastoreWithTags = optionsWithTags.getService(); + + com.google.datastore.v1.RequestOptions requestOptions = + com.google.datastore.v1.RequestOptions.newBuilder().addRequestTags("test-tag").build(); + + RunAggregationQueryResponse aggregationQueryResponse = placeholderAggregationQueryResponse(); + + EasyMock.expect( + rpcMock.runAggregationQuery( + matches( + aggregationQueryWithAliasAndRequestOptions( + "total_count", requestOptions, false, false)))) + .andReturn(aggregationQueryResponse); + + EasyMock.replay(rpcFactoryMock, rpcMock); + + EntityQuery selectAllQuery = Query.newEntityQueryBuilder().build(); + AggregationQuery getCountQuery = + Query.newAggregationQueryBuilder() + .addAggregation(count().as("total_count")) + .over(selectAllQuery) + .build(); + + AggregationResult result = getOnlyElement(datastoreWithTags.runAggregation(getCountQuery)); + assertThat(result.getLong("total_count")).isEqualTo(209L); + + EasyMock.verify(rpcFactoryMock, rpcMock); + } + @Test public void testRunAggregationQueryWithRequestOptions() { com.google.datastore.v1.RequestOptions requestOptions = diff --git a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java index f5b793c29774..43ab5696749c 100644 --- a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java +++ b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java @@ -204,6 +204,20 @@ public void testNamespace() { assertEquals("ns1", options.setNamespace("ns1").build().getNamespace()); } + @Test + public void testRequestTags() { + assertTrue(options.build().getRequestTags().isEmpty()); + assertThat(options.setRequestTags("tag1", "tag2").build().getRequestTags()) + .containsExactly("tag1", "tag2") + .inOrder(); + assertThat( + options + .setRequestTags(com.google.common.collect.ImmutableList.of("tag3")) + .build() + .getRequestTags()) + .containsExactly("tag3"); + } + @Test public void testDatastore() { assertSame(datastoreRpc, options.build().getRpc()); From ae79a99644b799dd27813507eceaff6d612ed22d Mon Sep 17 00:00:00 2001 From: Rahul Mane Date: Thu, 16 Jul 2026 08:38:36 +0000 Subject: [PATCH 4/4] refactor(datastore): extract DatastoreExecutionOptions for query execution options --- .../com/google/cloud/datastore/Datastore.java | 56 +------- .../datastore/DatastoreExecutionOptions.java | 134 ++++++++++++++++++ .../google/cloud/datastore/DatastoreImpl.java | 61 ++------ .../cloud/datastore/QueryResultsImpl.java | 2 +- .../google/cloud/datastore/ReadOption.java | 14 ++ .../execution/AggregationQueryExecutor.java | 34 ++++- .../datastore/AbstractDatastoreTest.java | 74 ++++++---- .../DatastoreExecutionOptionsTest.java | 91 ++++++++++++ 8 files changed, 336 insertions(+), 130 deletions(-) create mode 100644 java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreExecutionOptions.java create mode 100644 java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreExecutionOptionsTest.java diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java index 9f67182a9c47..bb71aa0880a0 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java @@ -495,28 +495,10 @@ interface TransactionCallable { QueryResults run(Query query, ExplainOptions explainOptions, ReadOption... options); /** - * Submits a {@link Query} with specified {@link com.google.datastore.v1.RequestOptions} and - * returns its result. - */ - QueryResults run(Query query, RequestOptions requestOptions); - - /** - * Submits a {@link Query} with specified {@link com.google.datastore.v1.RequestOptions} and - * returns its result. {@link ReadOption}s can be specified if desired. - */ - QueryResults run(Query query, RequestOptions requestOptions, ReadOption... options); - - /** - * Submits a {@link Query} with specified {@link com.google.cloud.datastore.models.ExplainOptions} - * and {@link com.google.datastore.v1.RequestOptions} and returns its result. {@link ReadOption}s - * can be specified if desired. + * Submits a {@link Query} with specified {@link DatastoreExecutionOptions} and returns its result. */ @BetaApi - QueryResults run( - Query query, - ExplainOptions explainOptions, - RequestOptions requestOptions, - ReadOption... options); + QueryResults run(Query query, DatastoreExecutionOptions executionOptions); /** * Submits a {@link AggregationQuery} and returns {@link AggregationResults}. {@link ReadOption}s @@ -563,20 +545,6 @@ QueryResults run( */ AggregationResults runAggregation(AggregationQuery query, ReadOption... options); - /** - * Submits an {@link AggregationQuery} with specified {@link - * com.google.datastore.v1.RequestOptions} and returns {@link AggregationResults}. - */ - AggregationResults runAggregation(AggregationQuery query, RequestOptions requestOptions); - - /** - * Submits an {@link AggregationQuery} with specified {@link - * com.google.datastore.v1.RequestOptions} and returns {@link AggregationResults}. {@link - * ReadOption}s can be specified if desired. - */ - AggregationResults runAggregation( - AggregationQuery query, RequestOptions requestOptions, ReadOption... options); - /** * Submits a {@link AggregationQuery} with specified {@link * com.google.cloud.datastore.models.ExplainOptions} and returns {@link AggregationResults}. @@ -605,26 +573,12 @@ AggregationResults runAggregation( AggregationQuery query, ExplainOptions explainOptions, ReadOption... options); /** - * Submits an {@link AggregationQuery} with specified {@link - * com.google.cloud.datastore.models.ExplainOptions} and {@link - * com.google.datastore.v1.RequestOptions} and returns {@link AggregationResults}. - */ - @BetaApi - AggregationResults runAggregation( - AggregationQuery query, ExplainOptions explainOptions, RequestOptions requestOptions); - - /** - * Submits an {@link AggregationQuery} with specified {@link - * com.google.cloud.datastore.models.ExplainOptions} and {@link - * com.google.datastore.v1.RequestOptions} and returns {@link AggregationResults}. {@link - * ReadOption}s can be specified if desired. + * Submits an {@link AggregationQuery} with specified {@link DatastoreExecutionOptions} and returns + * {@link AggregationResults}. */ @BetaApi AggregationResults runAggregation( - AggregationQuery query, - ExplainOptions explainOptions, - RequestOptions requestOptions, - ReadOption... options); + AggregationQuery query, DatastoreExecutionOptions executionOptions); /** * Closes the gRPC channels associated with this instance and frees up their resources. This diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreExecutionOptions.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreExecutionOptions.java new file mode 100644 index 000000000000..ab42ca75dad6 --- /dev/null +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreExecutionOptions.java @@ -0,0 +1,134 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.datastore; + +import com.google.api.core.BetaApi; +import com.google.cloud.datastore.models.ExplainOptions; +import com.google.common.base.Objects; +import com.google.common.collect.ImmutableList; +import com.google.datastore.v1.RequestOptions; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +/** + * Class representing options for query execution in Google Cloud Datastore. Combines {@link + * ExplainOptions}, {@link RequestOptions}, and {@link ReadOption}s. + */ +@BetaApi +public class DatastoreExecutionOptions { + + private final ExplainOptions explainOptions; + private final RequestOptions requestOptions; + private final List readOptions; + + private DatastoreExecutionOptions(Builder builder) { + this.explainOptions = builder.explainOptions; + this.requestOptions = builder.requestOptions; + this.readOptions = ImmutableList.copyOf(builder.readOptions); + } + + public ExplainOptions getExplainOptions() { + return explainOptions; + } + + public RequestOptions getRequestOptions() { + return requestOptions; + } + + public List getReadOptions() { + return readOptions; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (!(o instanceof DatastoreExecutionOptions)) return false; + DatastoreExecutionOptions that = (DatastoreExecutionOptions) o; + return Objects.equal(explainOptions, that.explainOptions) + && Objects.equal(requestOptions, that.requestOptions) + && Objects.equal(readOptions, that.readOptions); + } + + @Override + public int hashCode() { + return Objects.hashCode(explainOptions, requestOptions, readOptions); + } + + public static Builder newBuilder() { + return new Builder(); + } + + public static DatastoreExecutionOptions of(RequestOptions requestOptions) { + return newBuilder().setRequestOptions(requestOptions).build(); + } + + public static DatastoreExecutionOptions of(ExplainOptions explainOptions) { + return newBuilder().setExplainOptions(explainOptions).build(); + } + + public static DatastoreExecutionOptions of(ReadOption... readOptions) { + return newBuilder().setReadOptions(readOptions).build(); + } + + public static DatastoreExecutionOptions of( + RequestOptions requestOptions, ReadOption... readOptions) { + return newBuilder().setRequestOptions(requestOptions).setReadOptions(readOptions).build(); + } + + public static DatastoreExecutionOptions of( + ExplainOptions explainOptions, RequestOptions requestOptions, ReadOption... readOptions) { + return newBuilder() + .setExplainOptions(explainOptions) + .setRequestOptions(requestOptions) + .setReadOptions(readOptions) + .build(); + } + + /** Builder for {@link DatastoreExecutionOptions}. */ + public static class Builder { + private ExplainOptions explainOptions; + private RequestOptions requestOptions; + private List readOptions = Collections.emptyList(); + + private Builder() {} + + public Builder setExplainOptions(ExplainOptions explainOptions) { + this.explainOptions = explainOptions; + return this; + } + + public Builder setRequestOptions(RequestOptions requestOptions) { + this.requestOptions = requestOptions; + return this; + } + + public Builder setReadOptions(ReadOption... readOptions) { + this.readOptions = readOptions != null ? Arrays.asList(readOptions) : Collections.emptyList(); + return this; + } + + public Builder setReadOptions(List readOptions) { + this.readOptions = readOptions != null ? readOptions : Collections.emptyList(); + return this; + } + + public DatastoreExecutionOptions build() { + return new DatastoreExecutionOptions(this); + } + } +} diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index 923a09182f5f..cbd55278529f 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -345,28 +345,17 @@ public QueryResults run( return run(toReadOptionsPb(options), query, explainOptions.toPb()); } - /** Implements run query with request options. */ - @Override - public QueryResults run(Query query, RequestOptions requestOptions) { - return run(Optional.empty(), query, null, requestOptions); - } - - /** Implements run query with request options and read options. */ - @Override - public QueryResults run( - Query query, RequestOptions requestOptions, ReadOption... options) { - return run(toReadOptionsPb(options), query, null, requestOptions); - } - - /** Implements run query with explain options, request options and read options. */ + /** Implements run query with execution options. */ @Override @BetaApi - public QueryResults run( - Query query, - com.google.cloud.datastore.models.ExplainOptions explainOptions, - RequestOptions requestOptions, - ReadOption... options) { - return run(toReadOptionsPb(options), query, explainOptions.toPb(), requestOptions); + public QueryResults run(Query query, DatastoreExecutionOptions executionOptions) { + com.google.cloud.datastore.models.ExplainOptions explainOptions = + executionOptions.getExplainOptions(); + return run( + toReadOptionsPb(executionOptions.getReadOptions().toArray(new ReadOption[0])), + query, + explainOptions != null ? explainOptions.toPb() : null, + executionOptions.getRequestOptions()); } @SuppressWarnings("unchecked") @@ -396,24 +385,11 @@ public AggregationResults runAggregation(AggregationQuery query) { return aggregationQueryExecutor.execute(query, null); } - /** Implements run aggregation query with request options. */ - @Override - public AggregationResults runAggregation(AggregationQuery query, RequestOptions requestOptions) { - return aggregationQueryExecutor.execute(query, null, requestOptions); - } - @Override public AggregationResults runAggregation(AggregationQuery query, ReadOption... options) { return aggregationQueryExecutor.execute(query, null, options); } - /** Implements run aggregation query with request options and read options. */ - @Override - public AggregationResults runAggregation( - AggregationQuery query, RequestOptions requestOptions, ReadOption... options) { - return aggregationQueryExecutor.execute(query, null, requestOptions, options); - } - @Override @BetaApi public AggregationResults runAggregation( @@ -421,16 +397,6 @@ public AggregationResults runAggregation( return aggregationQueryExecutor.execute(query, explainOptions); } - /** Implements run aggregation query with explain options and request options. */ - @Override - @BetaApi - public AggregationResults runAggregation( - AggregationQuery query, - com.google.cloud.datastore.models.ExplainOptions explainOptions, - RequestOptions requestOptions) { - return aggregationQueryExecutor.execute(query, explainOptions, requestOptions); - } - @Override @BetaApi public AggregationResults runAggregation( @@ -440,15 +406,12 @@ public AggregationResults runAggregation( return aggregationQueryExecutor.execute(query, explainOptions, options); } - /** Implements run aggregation query with explain options, request options and read options. */ + /** Implements run aggregation query with execution options. */ @Override @BetaApi public AggregationResults runAggregation( - AggregationQuery query, - com.google.cloud.datastore.models.ExplainOptions explainOptions, - RequestOptions requestOptions, - ReadOption... options) { - return aggregationQueryExecutor.execute(query, explainOptions, requestOptions, options); + AggregationQuery query, DatastoreExecutionOptions executionOptions) { + return aggregationQueryExecutor.execute(query, executionOptions); } RunQueryResponse runQuery(final RunQueryRequest requestPb) { diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java index 504fc66b33ca..afa33be8d5b9 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java @@ -130,7 +130,7 @@ private void sendRequest() { if (explainOptions != null) { requestPb.setExplainOptions(explainOptions); } - if (requestOptions != null || !datastore.getOptions().getRequestTags().isEmpty()) { + if (requestOptions != null) { requestPb.setRequestOptions( RequestOptionsHelper.createRequestOptions(datastore.getOptions(), requestOptions)); } diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java index 7d01730b566d..437816f614d3 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java @@ -207,6 +207,20 @@ public static > QueryConfig create( .build(); } + public static > QueryConfig createWithExecutionOptions( + Q query, DatastoreExecutionOptions executionOptions) { + if (executionOptions == null) { + return create(query, (com.google.datastore.v1.ExplainOptions) null); + } + return create( + query, + executionOptions.getExplainOptions() != null + ? executionOptions.getExplainOptions().toPb() + : null, + executionOptions.getReadOptions(), + executionOptions.getRequestOptions()); + } + /** Creates a new builder for {@link QueryConfig}. */ public static > Builder newBuilder() { return new Builder<>(); diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java index f5a1c30933d8..4d56b30bcd1d 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/execution/AggregationQueryExecutor.java @@ -18,6 +18,7 @@ import com.google.api.core.InternalApi; import com.google.cloud.datastore.AggregationQuery; import com.google.cloud.datastore.AggregationResults; +import com.google.cloud.datastore.DatastoreExecutionOptions; import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.ReadOption; import com.google.cloud.datastore.ReadOption.QueryConfig; @@ -29,6 +30,8 @@ import com.google.datastore.v1.RunAggregationQueryRequest; import com.google.datastore.v1.RunAggregationQueryResponse; import java.util.Arrays; +import java.util.Collections; +import java.util.List; /** * An implementation of {@link QueryExecutor} which executes {@link AggregationQuery} and returns @@ -54,11 +57,35 @@ public AggregationResults execute( return execute(query, explainOptions, null, readOptions); } + public AggregationResults execute( + AggregationQuery query, DatastoreExecutionOptions executionOptions) { + if (executionOptions == null) { + return execute(query, null, null, Collections.emptyList()); + } + return execute( + query, + executionOptions.getExplainOptions(), + executionOptions.getRequestOptions(), + executionOptions.getReadOptions()); + } + public AggregationResults execute( AggregationQuery query, ExplainOptions explainOptions, RequestOptions requestOptions, ReadOption... readOptions) { + return execute( + query, + explainOptions, + requestOptions, + readOptions != null ? Arrays.asList(readOptions) : Collections.emptyList()); + } + + private AggregationResults execute( + AggregationQuery query, + ExplainOptions explainOptions, + RequestOptions requestOptions, + List readOptions) { RunAggregationQueryRequest runAggregationQueryRequest = getRunAggregationQueryRequest( query, @@ -74,12 +101,9 @@ private RunAggregationQueryRequest getRunAggregationQueryRequest( AggregationQuery query, com.google.datastore.v1.ExplainOptions explainOptions, RequestOptions requestOptions, - ReadOption... readOptions) { + List readOptions) { QueryConfig queryConfig = - readOptions == null - ? QueryConfig.create( - query, explainOptions, java.util.Collections.emptyList(), requestOptions) - : QueryConfig.create(query, explainOptions, Arrays.asList(readOptions), requestOptions); + QueryConfig.create(query, explainOptions, readOptions, requestOptions); return this.protoPreparer.prepare(queryConfig); } } diff --git a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java index c3466b156fc7..c003c9b852f6 100644 --- a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java +++ b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java @@ -1402,8 +1402,11 @@ private Predicate aggregationQueryWithAlias(String a @Test public void testRunQueryWithInstanceLevelRequestTags() { - DatastoreOptions optionsWithTags = options.toBuilder().setRequestTags("instance-tag").build(); - Datastore datastoreWithTags = optionsWithTags.getService(); + DatastoreOptions optionsWithTags = + rpcMockOptions.toBuilder() + .setRequestTags("instance-tag") + .setServiceRpcFactory(rpcFactoryMock) + .build(); PartitionId partitionId = PartitionId.newBuilder() @@ -1432,13 +1435,17 @@ public void testRunQueryWithInstanceLevelRequestTags() { RunQueryResponse.newBuilder() .setBatch( com.google.datastore.v1.QueryResultBatch.newBuilder() + .setEntityResultType(com.google.datastore.v1.EntityResult.ResultType.FULL) .addEntityResults(EntityResult.newBuilder().setEntity(ENTITY1.toPb())) .build()) .build(); + EasyMock.reset(rpcFactoryMock); + EasyMock.expect(rpcFactoryMock.create(optionsWithTags)).andReturn(rpcMock); EasyMock.expect(rpcMock.runQuery(expectedRequest)).andReturn(response); EasyMock.replay(rpcFactoryMock, rpcMock); + Datastore datastoreWithTags = optionsWithTags.getService(); Query query = Query.newEntityQueryBuilder().setKind(KIND1).build(); QueryResults results = datastoreWithTags.run(query); assertTrue(results.hasNext()); @@ -1464,7 +1471,7 @@ public void testRunQueryWithRequestOptions() { .addKind(com.google.datastore.v1.KindExpression.newBuilder().setName(KIND1)) .build(); - // 1. RunQueryRequest for run(query, requestOptions) + // 1. RunQueryRequest for run(query, DatastoreExecutionOptions.of(requestOptions)) RunQueryRequest expectedRequest1 = RunQueryRequest.newBuilder() .setProjectId(PROJECT_ID) @@ -1473,7 +1480,7 @@ public void testRunQueryWithRequestOptions() { .setRequestOptions(requestOptions) .build(); - // 2. RunQueryRequest for run(query, requestOptions, ReadOption) + // 2. RunQueryRequest for run(query, DatastoreExecutionOptions.of(requestOptions, ReadOption)) RunQueryRequest expectedRequest2 = RunQueryRequest.newBuilder() .setProjectId(PROJECT_ID) @@ -1484,7 +1491,7 @@ public void testRunQueryWithRequestOptions() { .setRequestOptions(requestOptions) .build(); - // 3. RunQueryRequest for run(query, explainOptions, requestOptions, ReadOption) + // 3. RunQueryRequest for run(query, DatastoreExecutionOptions.of(explainOptions, requestOptions, ReadOption)) RunQueryRequest expectedRequest3 = RunQueryRequest.newBuilder() .setProjectId(PROJECT_ID) @@ -1501,6 +1508,7 @@ public void testRunQueryWithRequestOptions() { RunQueryResponse.newBuilder() .setBatch( com.google.datastore.v1.QueryResultBatch.newBuilder() + .setEntityResultType(com.google.datastore.v1.EntityResult.ResultType.FULL) .addEntityResults(EntityResult.newBuilder().setEntity(ENTITY1.toPb())) .build()) .build(); @@ -1513,24 +1521,28 @@ public void testRunQueryWithRequestOptions() { Datastore targetDatastore = rpcMockOptions.getService(); Query query = Query.newEntityQueryBuilder().setKind(KIND1).build(); - // Test 1: run(query, requestOptions) - QueryResults results1 = targetDatastore.run(query, requestOptions); + // Test 1: run(query, DatastoreExecutionOptions.of(requestOptions)) + QueryResults results1 = + targetDatastore.run(query, DatastoreExecutionOptions.of(requestOptions)); assertTrue(results1.hasNext()); assertEquals(ENTITY1, results1.next()); - // Test 2: run(query, requestOptions, readOptions) + // Test 2: run(query, DatastoreExecutionOptions.of(requestOptions, readOptions)) QueryResults results2 = - targetDatastore.run(query, requestOptions, ReadOption.eventualConsistency()); + targetDatastore.run( + query, + DatastoreExecutionOptions.of(requestOptions, ReadOption.eventualConsistency())); assertTrue(results2.hasNext()); assertEquals(ENTITY1, results2.next()); - // Test 3: run(query, explainOptions, requestOptions, readOptions) + // Test 3: run(query, DatastoreExecutionOptions.of(explainOptions, requestOptions, readOptions)) QueryResults results3 = targetDatastore.run( query, - com.google.cloud.datastore.models.ExplainOptions.newBuilder().setAnalyze(true).build(), - requestOptions, - ReadOption.eventualConsistency()); + DatastoreExecutionOptions.of( + com.google.cloud.datastore.models.ExplainOptions.newBuilder().setAnalyze(true).build(), + requestOptions, + ReadOption.eventualConsistency())); assertTrue(results3.hasNext()); assertEquals(ENTITY1, results3.next()); @@ -1539,14 +1551,19 @@ public void testRunQueryWithRequestOptions() { @Test public void testRunAggregationQueryWithInstanceLevelRequestTags() { - DatastoreOptions optionsWithTags = options.toBuilder().setRequestTags("instance-tag").build(); - Datastore datastoreWithTags = optionsWithTags.getService(); + DatastoreOptions optionsWithTags = + rpcMockOptions.toBuilder() + .setRequestTags("instance-tag") + .setServiceRpcFactory(rpcFactoryMock) + .build(); com.google.datastore.v1.RequestOptions requestOptions = - com.google.datastore.v1.RequestOptions.newBuilder().addRequestTags("test-tag").build(); + com.google.datastore.v1.RequestOptions.newBuilder().addRequestTags("instance-tag").build(); RunAggregationQueryResponse aggregationQueryResponse = placeholderAggregationQueryResponse(); + EasyMock.reset(rpcFactoryMock); + EasyMock.expect(rpcFactoryMock.create(optionsWithTags)).andReturn(rpcMock); EasyMock.expect( rpcMock.runAggregationQuery( matches( @@ -1556,6 +1573,8 @@ public void testRunAggregationQueryWithInstanceLevelRequestTags() { EasyMock.replay(rpcFactoryMock, rpcMock); + Datastore datastoreWithTags = optionsWithTags.getService(); + EntityQuery selectAllQuery = Query.newEntityQueryBuilder().build(); AggregationQuery getCountQuery = Query.newAggregationQueryBuilder() @@ -1576,7 +1595,7 @@ public void testRunAggregationQueryWithRequestOptions() { RunAggregationQueryResponse aggregationQueryResponse = placeholderAggregationQueryResponse(); - // 1. runAggregation(query, requestOptions) + // 1. runAggregation(query, DatastoreExecutionOptions.of(requestOptions)) EasyMock.expect( rpcMock.runAggregationQuery( matches( @@ -1584,7 +1603,7 @@ public void testRunAggregationQueryWithRequestOptions() { "total_count", requestOptions, false, false)))) .andReturn(aggregationQueryResponse); - // 2. runAggregation(query, requestOptions, options) + // 2. runAggregation(query, DatastoreExecutionOptions.of(requestOptions, options)) EasyMock.expect( rpcMock.runAggregationQuery( matches( @@ -1592,7 +1611,7 @@ public void testRunAggregationQueryWithRequestOptions() { "total_count", requestOptions, true, false)))) .andReturn(aggregationQueryResponse); - // 3. runAggregation(query, explainOptions, requestOptions) + // 3. runAggregation(query, DatastoreExecutionOptions.of(explainOptions, requestOptions)) EasyMock.expect( rpcMock.runAggregationQuery( matches( @@ -1600,7 +1619,7 @@ public void testRunAggregationQueryWithRequestOptions() { "total_count", requestOptions, false, true)))) .andReturn(aggregationQueryResponse); - // 4. runAggregation(query, explainOptions, requestOptions, options) + // 4. runAggregation(query, DatastoreExecutionOptions.of(explainOptions, requestOptions, options)) EasyMock.expect( rpcMock.runAggregationQuery( matches( @@ -1624,26 +1643,33 @@ public void testRunAggregationQueryWithRequestOptions() { // Test 1 AggregationResult result1 = - getOnlyElement(mockDatastore.runAggregation(getCountQuery, requestOptions)); + getOnlyElement( + mockDatastore.runAggregation( + getCountQuery, DatastoreExecutionOptions.of(requestOptions))); assertThat(result1.getLong("total_count")).isEqualTo(209L); // Test 2 AggregationResult result2 = getOnlyElement( mockDatastore.runAggregation( - getCountQuery, requestOptions, ReadOption.eventualConsistency())); + getCountQuery, + DatastoreExecutionOptions.of(requestOptions, ReadOption.eventualConsistency()))); assertThat(result2.getLong("total_count")).isEqualTo(209L); // Test 3 AggregationResult result3 = - getOnlyElement(mockDatastore.runAggregation(getCountQuery, explainOptions, requestOptions)); + getOnlyElement( + mockDatastore.runAggregation( + getCountQuery, DatastoreExecutionOptions.of(explainOptions, requestOptions))); assertThat(result3.getLong("total_count")).isEqualTo(209L); // Test 4 AggregationResult result4 = getOnlyElement( mockDatastore.runAggregation( - getCountQuery, explainOptions, requestOptions, ReadOption.eventualConsistency())); + getCountQuery, + DatastoreExecutionOptions.of( + explainOptions, requestOptions, ReadOption.eventualConsistency()))); assertThat(result4.getLong("total_count")).isEqualTo(209L); EasyMock.verify(rpcFactoryMock, rpcMock); diff --git a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreExecutionOptionsTest.java b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreExecutionOptionsTest.java new file mode 100644 index 000000000000..6f7f1fb3d0d8 --- /dev/null +++ b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreExecutionOptionsTest.java @@ -0,0 +1,91 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.datastore; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.cloud.datastore.models.ExplainOptions; +import com.google.datastore.v1.RequestOptions; +import org.junit.Test; + +public class DatastoreExecutionOptionsTest { + + @Test + public void testBuilderDefaults() { + DatastoreExecutionOptions option = DatastoreExecutionOptions.newBuilder().build(); + + assertThat(option.getExplainOptions()).isNull(); + assertThat(option.getRequestOptions()).isNull(); + assertThat(option.getReadOptions()).isEmpty(); + } + + @Test + public void testBuilderFull() { + ExplainOptions explainOptions = ExplainOptions.newBuilder().setAnalyze(true).build(); + RequestOptions requestOptions = RequestOptions.newBuilder().addRequestTags("tag1").build(); + ReadOption readOption = ReadOption.eventualConsistency(); + + DatastoreExecutionOptions option = + DatastoreExecutionOptions.newBuilder() + .setExplainOptions(explainOptions) + .setRequestOptions(requestOptions) + .setReadOptions(readOption) + .build(); + + assertThat(option.getExplainOptions()).isEqualTo(explainOptions); + assertThat(option.getRequestOptions()).isEqualTo(requestOptions); + assertThat(option.getReadOptions()).containsExactly(readOption); + } + + @Test + public void testFactoryMethods() { + RequestOptions requestOptions = RequestOptions.newBuilder().addRequestTags("tag1").build(); + ExplainOptions explainOptions = ExplainOptions.newBuilder().setAnalyze(true).build(); + ReadOption readOption = ReadOption.eventualConsistency(); + + DatastoreExecutionOptions opt1 = DatastoreExecutionOptions.of(requestOptions); + assertThat(opt1.getRequestOptions()).isEqualTo(requestOptions); + + DatastoreExecutionOptions opt2 = DatastoreExecutionOptions.of(explainOptions); + assertThat(opt2.getExplainOptions()).isEqualTo(explainOptions); + + DatastoreExecutionOptions opt3 = DatastoreExecutionOptions.of(readOption); + assertThat(opt3.getReadOptions()).containsExactly(readOption); + + DatastoreExecutionOptions opt4 = DatastoreExecutionOptions.of(requestOptions, readOption); + assertThat(opt4.getRequestOptions()).isEqualTo(requestOptions); + assertThat(opt4.getReadOptions()).containsExactly(readOption); + + DatastoreExecutionOptions opt5 = + DatastoreExecutionOptions.of(explainOptions, requestOptions, readOption); + assertThat(opt5.getExplainOptions()).isEqualTo(explainOptions); + assertThat(opt5.getRequestOptions()).isEqualTo(requestOptions); + assertThat(opt5.getReadOptions()).containsExactly(readOption); + } + + @Test + public void testEqualsAndHashCode() { + RequestOptions reqOpts = RequestOptions.newBuilder().addRequestTags("tag1").build(); + ExplainOptions expOpts = ExplainOptions.newBuilder().setAnalyze(true).build(); + + DatastoreExecutionOptions opt1 = DatastoreExecutionOptions.of(expOpts, reqOpts); + DatastoreExecutionOptions opt2 = DatastoreExecutionOptions.of(expOpts, reqOpts); + + assertThat(opt1).isEqualTo(opt2); + assertThat(opt1.hashCode()).isEqualTo(opt2.hashCode()); + } +}