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..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 @@ -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,25 @@ 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. + * + * @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..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 @@ -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 || !datastore.getOptions().getRequestTags().isEmpty()) { + 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..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 @@ -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 || !datastoreOptions.getRequestTags().isEmpty()) { + 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..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 @@ -1399,4 +1399,238 @@ private Predicate aggregationQueryWithAlias(String a .get(0) .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 = + 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"); + } }