Repository: googleapis/java-datastore
Pull request #940 (2022-12-16 17:43:59)

feat: add x-goog-request-params to header

blakeli0: I know we probably only need `databaseId` and `projectId` at this moment, but in case we need more info that we cannot easily get from the request directly in the future, do we want to implement extract and match logic based on the proto? So that we don't need to change the code every time the routing annotation changes. 
We implemented it for pure gRPC gapic client libraries, and [this](https://github.com/googleapis/gax-java/blob/a1d62706624539a699745a22241c9ebc14b3495d/gax-grpc/src/main/java/com/google/api/gax/grpc/CallOptionsUtil.java#L48) is how the autogenerated client for Bigtable look like(I know Bigtable does not really use it, but you get the idea).
kolea2: > I know we probably only need `databaseId` and `projectId` at this moment, but in case we need more info that we cannot easily get from the request directly in the future, do we want to implement extract and match logic based on the proto? So that we don't need to change the code every time the routing annotation changes. We implemented it for pure gRPC gapic client libraries, and [this](https://github.com/googleapis/gax-java/blob/a1d62706624539a699745a22241c9ebc14b3495d/gax-grpc/src/main/java/com/google/api/gax/grpc/CallOptionsUtil.java#L48) is how the autogenerated client for Bigtable look like(I know Bigtable does not really use it, but you get the idea).

@blakeli0 this is an interesting idea - I will file an issue to look into this as a refactor for the future

Patch:
diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java
index 09101c94..6b886aef 100644
--- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java
+++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java
@@ -67,7 +67,8 @@ public class Datastore {
   }
 
   public AllocateIdsResponse allocateIds(AllocateIdsRequest request) throws DatastoreException {
-    try (InputStream is = remoteRpc.call("allocateIds", request)) {
+    try (InputStream is =
+        remoteRpc.call("allocateIds", request, request.getProjectId(), request.getDatabaseId())) {
       return AllocateIdsResponse.parseFrom(is);
     } catch (IOException exception) {
       throw invalidResponseException("allocateIds", exception);
@@ -76,7 +77,9 @@ public class Datastore {
 
   public BeginTransactionResponse beginTransaction(BeginTransactionRequest request)
       throws DatastoreException {
-    try (InputStream is = remoteRpc.call("beginTransaction", request)) {
+    try (InputStream is =
+        remoteRpc.call(
+            "beginTransaction", request, request.getProjectId(), request.getDatabaseId())) {
       return BeginTransactionResponse.parseFrom(is);
     } catch (IOException exception) {
       throw invalidResponseException("beginTransaction", exception);
@@ -84,7 +87,8 @@ public class Datastore {
   }
 
   public CommitResponse commit(CommitRequest request) throws DatastoreException {
-    try (InputStream is = remoteRpc.call("commit", request)) {
+    try (InputStream is =
+        remoteRpc.call("commit", request, request.getProjectId(), request.getDatabaseId())) {
       return CommitResponse.parseFrom(is);
     } catch (IOException exception) {
       throw invalidResponseException("commit", exception);
@@ -92,7 +96,8 @@ public class Datastore {
   }
 
   public LookupResponse lookup(LookupRequest request) throws DatastoreException {
-    try (InputStream is = remoteRpc.call("lookup", request)) {
+    try (InputStream is =
+        remoteRpc.call("lookup", request, request.getProjectId(), request.getDatabaseId())) {
       return LookupResponse.parseFrom(is);
     } catch (IOException exception) {
       throw invalidResponseException("lookup", exception);
@@ -100,7 +105,8 @@ public class Datastore {
   }
 
   public ReserveIdsResponse reserveIds(ReserveIdsRequest request) throws DatastoreException {
-    try (InputStream is = remoteRpc.call("reserveIds", request)) {
+    try (InputStream is =
+        remoteRpc.call("reserveIds", request, request.getProjectId(), request.getDatabaseId())) {
       return ReserveIdsResponse.parseFrom(is);
     } catch (IOException exception) {
       throw invalidResponseException("reserveIds", exception);
@@ -108,7 +114,8 @@ public class Datastore {
   }
 
   public RollbackResponse rollback(RollbackRequest request) throws DatastoreException {
-    try (InputStream is = remoteRpc.call("rollback", request)) {
+    try (InputStream is =
+        remoteRpc.call("rollback", request, request.getProjectId(), request.getDatabaseId())) {
       return RollbackResponse.parseFrom(is);
     } catch (IOException exception) {
       throw invalidResponseException("rollback", exception);
@@ -116,7 +123,8 @@ public class Datastore {
   }
 
   public RunQueryResponse runQuery(RunQueryRequest request) throws DatastoreException {
-    try (InputStream is = remoteRpc.call("runQuery", request)) {
+    try (InputStream is =
+        remoteRpc.call("runQuery", request, request.getProjectId(), request.getDatabaseId())) {
       return RunQueryResponse.parseFrom(is);
     } catch (IOException exception) {
       throw invalidResponseException("runQuery", exception);
@@ -125,7 +133,9 @@ public class Datastore {
 
   public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request)
       throws DatastoreException {
-    try (InputStream is = remoteRpc.call("runAggregationQuery", request)) {
+    try (InputStream is =
+        remoteRpc.call(
+            "runAggregationQuery", request, request.getProjectId(), request.getDatabaseId())) {
       return RunAggregationQueryResponse.parseFrom(is);
     } catch (IOException exception) {
       throw invalidResponseException("runAggregationQuery", exception);
diff --git a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/RemoteRpc.java b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/RemoteRpc.java
index 321eea72..b0b47c50 100644
--- a/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/RemoteRpc.java
+++ b/datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/RemoteRpc.java
@@ -24,6 +24,7 @@ import com.google.api.client.http.HttpStatusCodes;
 import com.google.api.client.http.protobuf.ProtoHttpContent;
 import com.google.api.client.util.IOUtils;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Strings;
 import com.google.protobuf.MessageLite;
 import com.google.rpc.Code;
 import com.google.rpc.Status;
@@ -46,6 +47,8 @@ class RemoteRpc {
   @VisibleForTesting static final String API_FORMAT_VERSION_HEADER = "X-Goog-Api-Format-Version";
   private static final String API_FORMAT_VERSION = "2";
 
+  @VisibleForTesting static final String X_GOOG_REQUEST_PARAMS_HEADER = "x-goog-request-params";
+
   private final HttpRequestFactory client;
   private final HttpRequestInitializer initializer;
   private final String url;
@@ -74,7 +77,9 @@ class RemoteRpc {
    *
    * @throws DatastoreException if the RPC fails.
    */
-  public InputStream call(String methodName, MessageLite request) throws DatastoreException {
+  public InputStream call(
+      String methodName, MessageLite request, String projectId, String databaseId)
+      throws DatastoreException {
     logger.fine("remote datastore call " + methodName);
 
     long startTime = System.currentTimeMillis();
@@ -84,7 +89,7 @@ class RemoteRpc {
         rpcCount.incrementAndGet();
         ProtoHttpContent payload = new ProtoHttpContent(request);
         HttpRequest httpRequest = client.buildPostRequest(resolveURL(methodName), payload);
-        setHeaders(request, httpRequest);
+        setHeaders(request, httpRequest, projectId, databaseId);
         // Don't throw an HTTPResponseException on error. It converts the response to a String and
         // throws away the original, whereas we need the raw bytes to parse it as a proto.
         httpRequest.setThrowExceptionOnExecuteError(false);
@@ -123,8 +128,16 @@ class RemoteRpc {
   }
 
   @VisibleForTesting
-  void setHeaders(MessageLite request, HttpRequest httpRequest) {
+  void setHeaders(
+      MessageLite request, HttpRequest httpRequest, String projectId, String databaseId) {
     httpRequest.getHeaders().put(API_FORMAT_VERSION_HEADER, API_FORMAT_VERSION);
+    StringBuilder builder = new StringBuilder("project_id=");
+    builder.append(projectId);
+    if (!Strings.isNullOrEmpty(databaseId)) {
+      builder.append("&database_id=");
+      builder.append(databaseId);
+    }
+    httpRequest.getHeaders().put(X_GOOG_REQUEST_PARAMS_HEADER, builder.toString());
     if (enableE2EChecksum && request != null) {
       String checksum = EndToEndChecksumHandler.computeChecksum(request.toByteArray());
       if (checksum != null) {


Test patch:
diff --git a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/RemoteRpcTest.java b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/RemoteRpcTest.java
index ebcb1239..281e92f0 100644
--- a/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/RemoteRpcTest.java
+++ b/datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/RemoteRpcTest.java
@@ -146,7 +146,8 @@ public class RemoteRpcTest {
         new InjectedTestValues(gzip(response), new byte[1], true);
     RemoteRpc rpc = newRemoteRpc(injectedTestValues);
 
-    InputStream is = rpc.call("beginTransaction", BeginTransactionResponse.getDefaultInstance());
+    InputStream is =
+        rpc.call("beginTransaction", BeginTransactionResponse.getDefaultInstance(), "", "");
     BeginTransactionResponse parsedResponse = BeginTransactionResponse.parseFrom(is);
     is.close();
 
@@ -159,14 +160,15 @@ public class RemoteRpcTest {
   public void testHttpHeaders_expectE2eChecksumHeader() throws IOException {
     // Enable E2E-Checksum system env variable
     RemoteRpc.setSystemEnvE2EChecksum(true);
+    String projectId = "project-id";
     MessageLite request =
-        RollbackRequest.newBuilder().setTransaction(ByteString.copyFromUtf8("project-id")).build();
+        RollbackRequest.newBuilder().setTransaction(ByteString.copyFromUtf8(projectId)).build();
     RemoteRpc rpc =
         newRemoteRpc(
             new InjectedTestValues(gzip(newBeginTransactionResponse()), new byte[1], true));
     HttpRequest httpRequest =
         rpc.getClient().buildPostRequest(rpc.resolveURL("blah"), new ProtoHttpContent(request));
-    rpc.setHeaders(request, httpRequest);
+    rpc.setHeaders(request, httpRequest, projectId, "");
     assertNotNull(
         httpRequest.getHeaders().getFirstHeaderStringValue(RemoteRpc.API_FORMAT_VERSION_HEADER));
     // Expect to find e2e-checksum header
@@ -181,14 +183,15 @@ public class RemoteRpcTest {
   public void testHttpHeaders_doNotExpectE2eChecksumHeader() throws IOException {
     // disable E2E-Checksum system env variable
     RemoteRpc.setSystemEnvE2EChecksum(false);
+    String projectId = "project-id";
     MessageLite request =
-        RollbackRequest.newBuilder().setTransaction(ByteString.copyFromUtf8("project-id")).build();
+        RollbackRequest.newBuilder().setTransaction(ByteString.copyFromUtf8(projectId)).build();
     RemoteRpc rpc =
         newRemoteRpc(
             new InjectedTestValues(gzip(newBeginTransactionResponse()), new byte[1], true));
     HttpRequest httpRequest =
         rpc.getClient().buildPostRequest(rpc.resolveURL("blah"), new ProtoHttpContent(request));
-    rpc.setHeaders(request, httpRequest);
+    rpc.setHeaders(request, httpRequest, projectId, "");
     assertNotNull(
         httpRequest.getHeaders().getFirstHeaderStringValue(RemoteRpc.API_FORMAT_VERSION_HEADER));
     // Do not expect to find e2e-checksum header
@@ -198,6 +201,38 @@ public class RemoteRpcTest {
             .getFirstHeaderStringValue(EndToEndChecksumHandler.HTTP_REQUEST_CHECKSUM_HEADER));
   }
 
+  @Test
+  public void testHttpHeaders_prefixHeader() throws IOException {
+    String projectId = "my-project";
+    String databaseId = "my-db";
+    MessageLite request =
+        RollbackRequest.newBuilder()
+            .setTransaction(ByteString.copyFromUtf8(projectId))
+            .setDatabaseId(databaseId)
+            .build();
+    RemoteRpc rpc =
+        newRemoteRpc(
+            new InjectedTestValues(gzip(newBeginTransactionResponse()), new byte[1], true));
+    HttpRequest httpRequest =
+        rpc.getClient().buildPostRequest(rpc.resolveURL("blah"), new ProtoHttpContent(request));
+    rpc.setHeaders(request, httpRequest, projectId, databaseId);
+    assertEquals(
+        "project_id=my-project&database_id=my-db",
+        httpRequest.getHeaders().get(RemoteRpc.X_GOOG_REQUEST_PARAMS_HEADER));
+
+    MessageLite request2 =
+        RollbackRequest.newBuilder().setTransaction(ByteString.copyFromUtf8(projectId)).build();
+    RemoteRpc rpc2 =
+        newRemoteRpc(
+            new InjectedTestValues(gzip(newBeginTransactionResponse()), new byte[1], true));
+    HttpRequest httpRequest2 =
+        rpc2.getClient().buildPostRequest(rpc2.resolveURL("blah"), new ProtoHttpContent(request2));
+    rpc2.setHeaders(request, httpRequest2, projectId, "");
+    assertEquals(
+        "project_id=my-project",
+        httpRequest2.getHeaders().get(RemoteRpc.X_GOOG_REQUEST_PARAMS_HEADER));
+  }
+
   private static BeginTransactionResponse newBeginTransactionResponse() {
     return BeginTransactionResponse.newBuilder()
         .setTransaction(ByteString.copyFromUtf8("blah-blah-blah"))