Add requestId to PartitionQuery + PartitionRead

Author: odeke-em

google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java

@@ -250,9 +250,11 @@ private List<Partition> partitionReadUsingIndex(
       }
       builder.setPartitionOptions(pbuilder.build());
 
+      XGoogSpannerRequestId reqId =
+          session.getRequestIdCreator().nextRequestId(1 /*TODO: retrieve channelId*/, 1);
       final PartitionReadRequest request = builder.build();
       try {
-        PartitionResponse response = rpc.partitionRead(request, options);
+        PartitionResponse response = rpc.partitionRead(request, reqId.withOptions(options));
         ImmutableList.Builder<Partition> partitions = ImmutableList.builder();
         for (com.google.spanner.v1.Partition p : response.getPartitionsList()) {
           Partition partition =
@@ -313,9 +315,11 @@ private List<Partition> partitionQuery(
       }
       builder.setPartitionOptions(pbuilder.build());
 
+      XGoogSpannerRequestId reqId =
+          session.getRequestIdCreator().nextRequestId(1 /*TODO: retrieve channelId*/, 1);
       final PartitionQueryRequest request = builder.build();
       try {
-        PartitionResponse response = rpc.partitionQuery(request, options);
+        PartitionResponse response = rpc.partitionQuery(request, reqId.withOptions(options));
         ImmutableList.Builder<Partition> partitions = ImmutableList.builder();
         for (com.google.spanner.v1.Partition p : response.getPartitionsList()) {
           Partition partition =

google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java

@@ -239,9 +239,8 @@ public static void startStaticServer() throws IOException {
                 "google.spanner.v1.Spanner/ExecuteStreamingSql",
                 "google.spanner.v1.Spanner/GetSession",
                 "google.spanner.v1.Spanner/ListSessions",
-                // TODO: Uncomment the methods when the implementation is completed.
-                // "google.spanner.v1.Spanner/PartitionQuery",
-                // "google.spanner.v1.Spanner/PartitionRead",
+                "google.spanner.v1.Spanner/PartitionQuery",
+                "google.spanner.v1.Spanner/PartitionRead",
                 "google.spanner.v1.Spanner/Read",
                 "google.spanner.v1.Spanner/Rollback",
                 "google.spanner.v1.Spanner/StreamingRead"));