impl(bigtable): BulkMutator can use Stub (#9114)

Author: dbolduc

google/cloud/bigtable/CMakeLists.txt

@@ -340,6 +340,7 @@ if (BUILD_TESTING)
         internal/default_row_reader_test.cc
         internal/defaults_test.cc
         internal/google_bytes_traits_test.cc
+        internal/legacy_bulk_mutator_test.cc
         internal/legacy_row_reader_test.cc
         internal/logging_data_client_test.cc
         internal/prefix_range_end_test.cc

google/cloud/bigtable/bigtable_client_unit_tests.bzl

@@ -50,6 +50,7 @@ bigtable_client_unit_tests = [
     "internal/default_row_reader_test.cc",
     "internal/defaults_test.cc",
     "internal/google_bytes_traits_test.cc",
+    "internal/legacy_bulk_mutator_test.cc",
     "internal/legacy_row_reader_test.cc",
     "internal/logging_data_client_test.cc",
     "internal/prefix_range_end_test.cc",

google/cloud/bigtable/internal/async_bulk_apply.cc

@@ -76,7 +76,7 @@ void AsyncRetryBulkApply::StartIteration(CompletionQueue cq) {
 
 void AsyncRetryBulkApply::OnRead(
     google::bigtable::v2::MutateRowsResponse response) {
-  state_.OnRead(response);
+  state_.OnRead(std::move(response));
 }
 
 void AsyncRetryBulkApply::OnFinish(CompletionQueue cq, Status const& status) {

google/cloud/bigtable/internal/bulk_mutator.cc

@@ -79,7 +79,7 @@ google::bigtable::v2::MutateRowsRequest const& BulkMutatorState::BeforeStart() {
 }
 
 std::vector<int> BulkMutatorState::OnRead(
-    google::bigtable::v2::MutateRowsResponse& response) {
+    google::bigtable::v2::MutateRowsResponse response) {
   std::vector<int> res;
   for (auto& entry : *response.mutable_entries()) {
     // The type of `entry.index()` is a 64-bit int. But we can never create more
@@ -197,14 +197,42 @@ grpc::Status BulkMutator::MakeOneRequest(bigtable::DataClient& client,
   // Read the stream of responses.
   btproto::MutateRowsResponse response;
   while (stream->Read(&response)) {
-    state_.OnRead(response);
+    state_.OnRead(std::move(response));
   }
   // Handle any errors in the stream.
   auto grpc_status = stream->Finish();
   state_.OnFinish(MakeStatusFromRpcError(grpc_status));
   return grpc_status;
 }
 
+Status BulkMutator::MakeOneRequest(bigtable_internal::BigtableStub& stub) {
+  // Send the request to the server.
+  auto const& mutations = state_.BeforeStart();
+
+  // Configure the context
+  auto const& options = google::cloud::internal::CurrentOptions();
+  auto context = absl::make_unique<grpc::ClientContext>();
+  google::cloud::internal::ConfigureContext(*context, options);
+
+  struct UnpackVariant {
+    BulkMutatorState& state;
+    bool operator()(btproto::MutateRowsResponse r) {
+      state.OnRead(std::move(r));
+      return true;
+    }
+    bool operator()(Status s) {
+      state.OnFinish(std::move(s));
+      return false;
+    }
+  };
+
+  // Read the stream of responses.
+  auto stream = stub.MutateRows(std::move(context), mutations);
+  while (absl::visit(UnpackVariant{state_}, stream->Read())) {
+  }
+  return state_.last_status();
+}
+
 std::vector<FailedMutation> BulkMutator::OnRetryDone() && {
   return std::move(state_).OnRetryDone();
 }

google/cloud/bigtable/internal/bulk_mutator.h

@@ -18,6 +18,7 @@
 #include "google/cloud/bigtable/completion_queue.h"
 #include "google/cloud/bigtable/data_client.h"
 #include "google/cloud/bigtable/idempotent_mutation_policy.h"
+#include "google/cloud/bigtable/internal/bigtable_stub.h"
 #include "google/cloud/bigtable/version.h"
 #include "google/cloud/internal/invoke_result.h"
 #include "google/cloud/internal/retry_policy.h"
@@ -49,14 +50,17 @@ class BulkMutatorState {
    *
    * Returns the original index of any successful operations.
    */
-  std::vector<int> OnRead(google::bigtable::v2::MutateRowsResponse& response);
+  std::vector<int> OnRead(google::bigtable::v2::MutateRowsResponse response);
 
   /// Handle the result of a `Finish()` operation on the MutateRows() RPC.
   void OnFinish(google::cloud::Status finish_status);
 
   /// Terminate the retry loop and return all the failures.
   std::vector<FailedMutation> OnRetryDone() &&;
 
+  /// The status of the most recent stream.
+  Status last_status() const { return last_status_; };
+
  private:
   /// The current request proto.
   google::bigtable::v2::MutateRowsRequest mutations_;
@@ -116,6 +120,9 @@ class BulkMutator {
   grpc::Status MakeOneRequest(bigtable::DataClient& client,
                               grpc::ClientContext& client_context);
 
+  /// Synchronously send one batch request to the given stub.
+  Status MakeOneRequest(bigtable_internal::BigtableStub& stub);
+
   /// Give up on any pending mutations, move them to the failures array.
   std::vector<FailedMutation> OnRetryDone() &&;