diff --git a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/GrpcDatastoreRpc.java b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/GrpcDatastoreRpc.java index f952dfdc9ef0..b3c5d2e8aab0 100644 --- a/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/GrpcDatastoreRpc.java +++ b/java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/GrpcDatastoreRpc.java @@ -26,6 +26,7 @@ import com.google.api.gax.core.GaxProperties; import com.google.api.gax.grpc.GrpcCallContext; import com.google.api.gax.grpc.GrpcTransportChannel; +import com.google.api.gax.rpc.ApiCallContext; import com.google.api.gax.rpc.ClientContext; import com.google.api.gax.rpc.HeaderProvider; import com.google.api.gax.rpc.NoHeaderProvider; @@ -37,7 +38,10 @@ import com.google.cloud.datastore.v1.stub.DatastoreStubSettings; import com.google.cloud.datastore.v1.stub.GrpcDatastoreStub; import com.google.cloud.grpc.GrpcTransportOptions; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Strings; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.datastore.v1.AllocateIdsRequest; import com.google.datastore.v1.AllocateIdsResponse; import com.google.datastore.v1.BeginTransactionRequest; @@ -46,6 +50,7 @@ import com.google.datastore.v1.CommitResponse; import com.google.datastore.v1.LookupRequest; import com.google.datastore.v1.LookupResponse; +import com.google.datastore.v1.ReadOptions; import com.google.datastore.v1.ReserveIdsRequest; import com.google.datastore.v1.ReserveIdsResponse; import com.google.datastore.v1.RollbackRequest; @@ -62,11 +67,24 @@ @InternalApi public class GrpcDatastoreRpc implements DatastoreRpc { + private static final ApiCallContext TRANSACTION_CALL_CONTEXT = + GrpcCallContext.createDefault() + .withExtraHeaders(ImmutableMap.of("x-goog-is-transaction", ImmutableList.of("true"))); + + private static final ApiCallContext NON_TRANSACTION_CALL_CONTEXT = + GrpcCallContext.createDefault() + .withExtraHeaders(ImmutableMap.of("x-goog-is-transaction", ImmutableList.of("false"))); private final GrpcDatastoreStub datastoreStub; private final ClientContext clientContext; private boolean closed; + @VisibleForTesting + GrpcDatastoreRpc(GrpcDatastoreStub datastoreStub) { + this.datastoreStub = datastoreStub; + this.clientContext = null; + } + public GrpcDatastoreRpc(DatastoreOptions datastoreOptions) throws IOException { try { clientContext = @@ -108,43 +126,63 @@ public void close() throws Exception { @Override public AllocateIdsResponse allocateIds(AllocateIdsRequest request) { - return datastoreStub.allocateIdsCallable().call(request); + return datastoreStub.allocateIdsCallable().call(request, NON_TRANSACTION_CALL_CONTEXT); } @Override public BeginTransactionResponse beginTransaction(BeginTransactionRequest request) throws DatastoreException { - return datastoreStub.beginTransactionCallable().call(request); + return datastoreStub.beginTransactionCallable().call(request, TRANSACTION_CALL_CONTEXT); } @Override public CommitResponse commit(CommitRequest request) { - return datastoreStub.commitCallable().call(request); + boolean isTransaction = + request.getMode() == CommitRequest.Mode.TRANSACTIONAL + || !request.getTransaction().isEmpty() + || request.hasSingleUseTransaction(); + if (isTransaction) { + return datastoreStub.commitCallable().call(request, TRANSACTION_CALL_CONTEXT); + } + return datastoreStub.commitCallable().call(request, NON_TRANSACTION_CALL_CONTEXT); } @Override public LookupResponse lookup(LookupRequest request) { - return datastoreStub.lookupCallable().call(request); + if (request.hasReadOptions() && isTransactionalRead(request.getReadOptions())) { + return datastoreStub.lookupCallable().call(request, TRANSACTION_CALL_CONTEXT); + } + return datastoreStub.lookupCallable().call(request, NON_TRANSACTION_CALL_CONTEXT); } @Override public ReserveIdsResponse reserveIds(ReserveIdsRequest request) { - return datastoreStub.reserveIdsCallable().call(request); + return datastoreStub.reserveIdsCallable().call(request, NON_TRANSACTION_CALL_CONTEXT); } @Override public RollbackResponse rollback(RollbackRequest request) { - return datastoreStub.rollbackCallable().call(request); + return datastoreStub.rollbackCallable().call(request, TRANSACTION_CALL_CONTEXT); } @Override public RunQueryResponse runQuery(RunQueryRequest request) { - return datastoreStub.runQueryCallable().call(request); + if (request.hasReadOptions() && isTransactionalRead(request.getReadOptions())) { + return datastoreStub.runQueryCallable().call(request, TRANSACTION_CALL_CONTEXT); + } + return datastoreStub.runQueryCallable().call(request, NON_TRANSACTION_CALL_CONTEXT); } @Override public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) { - return datastoreStub.runAggregationQueryCallable().call(request); + if (request.hasReadOptions() && isTransactionalRead(request.getReadOptions())) { + return datastoreStub.runAggregationQueryCallable().call(request, TRANSACTION_CALL_CONTEXT); + } + return datastoreStub.runAggregationQueryCallable().call(request, NON_TRANSACTION_CALL_CONTEXT); + } + + private static boolean isTransactionalRead(ReadOptions readOptions) { + return !readOptions.getTransaction().isEmpty() || readOptions.hasNewTransaction(); } @Override diff --git a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/spi/v1/GrpcDatastoreRpcTest.java b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/spi/v1/GrpcDatastoreRpcTest.java new file mode 100644 index 000000000000..93ab9bf2caaa --- /dev/null +++ b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/spi/v1/GrpcDatastoreRpcTest.java @@ -0,0 +1,209 @@ +/* + * 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.spi.v1; + +import static com.google.common.truth.Truth.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.UnaryCallable; +import com.google.cloud.datastore.v1.stub.GrpcDatastoreStub; +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.LookupRequest; +import com.google.datastore.v1.LookupResponse; +import com.google.datastore.v1.ReadOptions; +import com.google.datastore.v1.ReserveIdsRequest; +import com.google.datastore.v1.ReserveIdsResponse; +import com.google.datastore.v1.RollbackRequest; +import com.google.datastore.v1.RollbackResponse; +import com.google.datastore.v1.RunAggregationQueryRequest; +import com.google.datastore.v1.RunAggregationQueryResponse; +import com.google.datastore.v1.RunQueryRequest; +import com.google.datastore.v1.RunQueryResponse; +import com.google.datastore.v1.TransactionOptions; +import com.google.protobuf.ByteString; +import org.junit.Test; + +public class GrpcDatastoreRpcTest { + + @Test + public void testRoutingHeaders() throws Exception { + GrpcDatastoreStub mockStub = mock(GrpcDatastoreStub.class); + + RecordingUnaryCallable beginTxCallable = + new RecordingUnaryCallable<>(); + when(mockStub.beginTransactionCallable()).thenReturn(beginTxCallable); + + RecordingUnaryCallable rollbackCallable = + new RecordingUnaryCallable<>(); + when(mockStub.rollbackCallable()).thenReturn(rollbackCallable); + + RecordingUnaryCallable commitCallable = + new RecordingUnaryCallable<>(); + when(mockStub.commitCallable()).thenReturn(commitCallable); + + RecordingUnaryCallable lookupCallable = + new RecordingUnaryCallable<>(); + when(mockStub.lookupCallable()).thenReturn(lookupCallable); + + RecordingUnaryCallable runQueryCallable = + new RecordingUnaryCallable<>(); + when(mockStub.runQueryCallable()).thenReturn(runQueryCallable); + + RecordingUnaryCallable + runAggQueryCallable = new RecordingUnaryCallable<>(); + when(mockStub.runAggregationQueryCallable()).thenReturn(runAggQueryCallable); + + RecordingUnaryCallable allocateIdsCallable = + new RecordingUnaryCallable<>(); + when(mockStub.allocateIdsCallable()).thenReturn(allocateIdsCallable); + + RecordingUnaryCallable reserveIdsCallable = + new RecordingUnaryCallable<>(); + when(mockStub.reserveIdsCallable()).thenReturn(reserveIdsCallable); + + GrpcDatastoreRpc rpc = new GrpcDatastoreRpc(mockStub); + + // beginTransaction -> always true + rpc.beginTransaction(BeginTransactionRequest.getDefaultInstance()); + assertThat(beginTxCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("true"); + + // rollback -> always true + rpc.rollback(RollbackRequest.getDefaultInstance()); + assertThat(rollbackCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("true"); + + // commit (transactional mode) -> true + rpc.commit(CommitRequest.newBuilder().setMode(CommitRequest.Mode.TRANSACTIONAL).build()); + assertThat(commitCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("true"); + + // commit (non-transactional) -> false + rpc.commit(CommitRequest.newBuilder().setMode(CommitRequest.Mode.NON_TRANSACTIONAL).build()); + assertThat(commitCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("false"); + + // commit (has transaction id) -> true + rpc.commit(CommitRequest.newBuilder().setTransaction(ByteString.copyFromUtf8("tx")).build()); + assertThat(commitCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("true"); + + // commit (single-use transaction) -> true + rpc.commit( + CommitRequest.newBuilder() + .setSingleUseTransaction(TransactionOptions.getDefaultInstance()) + .build()); + assertThat(commitCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("true"); + + // lookup (transactional) -> true + rpc.lookup( + LookupRequest.newBuilder() + .setReadOptions(ReadOptions.newBuilder().setTransaction(ByteString.copyFromUtf8("tx"))) + .build()); + assertThat(lookupCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("true"); + + // lookup (non-transactional) -> false + rpc.lookup(LookupRequest.getDefaultInstance()); + assertThat(lookupCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("false"); + + // lookup (newTransaction) -> true + rpc.lookup( + LookupRequest.newBuilder() + .setReadOptions( + ReadOptions.newBuilder().setNewTransaction(TransactionOptions.getDefaultInstance())) + .build()); + assertThat(lookupCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("true"); + + // runQuery (transactional) -> true + rpc.runQuery( + RunQueryRequest.newBuilder() + .setReadOptions(ReadOptions.newBuilder().setTransaction(ByteString.copyFromUtf8("tx"))) + .build()); + assertThat(runQueryCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("true"); + + // runQuery (non-transactional) -> false + rpc.runQuery(RunQueryRequest.getDefaultInstance()); + assertThat(runQueryCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("false"); + + // runQuery (newTransaction) -> true + rpc.runQuery( + RunQueryRequest.newBuilder() + .setReadOptions( + ReadOptions.newBuilder().setNewTransaction(TransactionOptions.getDefaultInstance())) + .build()); + assertThat(runQueryCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("true"); + + // runAggregationQuery (transactional) -> true + rpc.runAggregationQuery( + RunAggregationQueryRequest.newBuilder() + .setReadOptions(ReadOptions.newBuilder().setTransaction(ByteString.copyFromUtf8("tx"))) + .build()); + assertThat(runAggQueryCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("true"); + + // runAggregationQuery (non-transactional) -> false + rpc.runAggregationQuery(RunAggregationQueryRequest.getDefaultInstance()); + assertThat(runAggQueryCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("false"); + + // runAggregationQuery (newTransaction) -> true + rpc.runAggregationQuery( + RunAggregationQueryRequest.newBuilder() + .setReadOptions( + ReadOptions.newBuilder().setNewTransaction(TransactionOptions.getDefaultInstance())) + .build()); + assertThat(runAggQueryCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("true"); + + // allocateIds -> false + rpc.allocateIds(AllocateIdsRequest.getDefaultInstance()); + assertThat(allocateIdsCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("false"); + + // reserveIds -> false + rpc.reserveIds(ReserveIdsRequest.getDefaultInstance()); + assertThat(reserveIdsCallable.lastContext.getExtraHeaders().get("x-goog-is-transaction")) + .containsExactly("false"); + } + + private static class RecordingUnaryCallable + extends UnaryCallable { + ApiCallContext lastContext; + + @Override + public ApiFuture futureCall(RequestT request, ApiCallContext context) { + this.lastContext = context; + return ApiFutures.immediateFuture(null); + } + } +}