impl(generator): add OperationContext to Bigtable stub and decorators

generator/generator_config.proto

@@ -182,6 +182,10 @@ message ServiceConfiguration {
   // The implementation for this method in the ConnectionImpl class is not
   // generated and must be handwritten in a separate .cc file.
   repeated BespokeMethod bespoke_methods = 30;
+
+  // If set to `true`, the stub class methods will include an additional
+  // parameter of type `google::cloud::bigtable_internal::OperationContext&`.
+  bool experimental_bigtable_operation_context = 31;
 }
 
 message DiscoveryDocumentDefinedProduct {

generator/generator_config.textproto

@@ -576,6 +576,7 @@ service {
   omit_connection: true
   omit_stub_factory: true
   generate_round_robin_decorator: true
+  experimental_bigtable_operation_context: true
   omitted_rpcs: [
     "GenerateInitialChangeStreamPartitions",
     "ReadChangeStream"

generator/internal/auth_decorator_generator.cc

@@ -17,6 +17,7 @@
 #include "generator/internal/longrunning.h"
 #include "generator/internal/predicate_utils.h"
 #include "generator/internal/printer.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/str_split.h"
 #include <google/protobuf/descriptor.h>
 
@@ -130,37 +131,39 @@ Status AuthDecoratorGenerator::GenerateCc() {
   // Auth decorator class member methods
   for (auto const& method : methods()) {
     if (IsStreamingWrite(method)) {
-      CcPrintMethod(method, __FILE__, __LINE__, R"""(
+      CcPrintMethod(method, __FILE__, __LINE__,
+                    R"""(
 std::unique_ptr<::google::cloud::internal::StreamingWriteRpc<
     $request_type$,
     $response_type$>>
 $auth_class_name$::$method_name$(
     std::shared_ptr<grpc::ClientContext> context,
-    Options const& options) {
+    Options const& options$op_ctx_shared_decl$) {
   using ErrorStream = ::google::cloud::internal::StreamingWriteRpcError<
       $request_type$, $response_type$>;
   auto status = auth_->ConfigureContext(*context);
   if (!status.ok()) return std::make_unique<ErrorStream>(std::move(status));
-  return child_->$method_name$(std::move(context), options);
+  return child_->$method_name$(std::move(context), options$op_ctx_shared_arg$);
 }
 )""");
       continue;
     }
     if (IsBidirStreaming(method)) {
-      CcPrintMethod(method, __FILE__, __LINE__, R"""(
+      CcPrintMethod(method, __FILE__, __LINE__,
+                    R"""(
 std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc<
     $request_type$,
     $response_type$>>
 $auth_class_name$::Async$method_name$(
     google::cloud::CompletionQueue const& cq,
     std::shared_ptr<grpc::ClientContext> context,
-    google::cloud::internal::ImmutableOptions options) {
+    google::cloud::internal::ImmutableOptions options$op_ctx_shared_decl$) {
   using StreamAuth = google::cloud::internal::AsyncStreamingReadWriteRpcAuth<
     $request_type$, $response_type$>;
 
-  auto call = [child = child_, cq, options = std::move(options)](
+  auto call = [child = child_, cq, options = std::move(options)$op_ctx_shared_cap$](
                   std::shared_ptr<grpc::ClientContext> ctx) {
-    return child->Async$method_name$(cq, std::move(ctx), options);
+    return child->Async$method_name$(cq, std::move(ctx), options$op_ctx_shared_arg$);
   };
   return std::make_unique<StreamAuth>(
     std::move(context), auth_, StreamAuth::StreamFactory(std::move(call)));
@@ -169,53 +172,56 @@ std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc<
       continue;
     }
     if (IsLongrunningOperation(method)) {
-      CcPrintMethod(method, __FILE__, __LINE__, R"""(
+      CcPrintMethod(method, __FILE__, __LINE__,
+                    R"""(
 future<StatusOr<google::longrunning::Operation>>
 $auth_class_name$::Async$method_name$(
       google::cloud::CompletionQueue& cq,
       std::shared_ptr<grpc::ClientContext> context,
       google::cloud::internal::ImmutableOptions options,
-      $request_type$ const& request) {
+      $request_type$ const& request$op_ctx_shared_decl$) {
   using ReturnType = StatusOr<google::longrunning::Operation>;
   return auth_->AsyncConfigureContext(std::move(context)).then(
-      [cq, child = child_, options = std::move(options), request](
+      [cq, child = child_, options = std::move(options), request$op_ctx_shared_cap$](
           future<StatusOr<std::shared_ptr<grpc::ClientContext>>> f) mutable {
         auto context = f.get();
         if (!context) {
           return make_ready_future(ReturnType(std::move(context).status()));
         }
         return child->Async$method_name$(
-            cq, *std::move(context), std::move(options), request);
+            cq, *std::move(context), std::move(options), request$op_ctx_shared_arg$);
       });
 }
 )""");
 
-      CcPrintMethod(method, __FILE__, __LINE__, R"""(
+      CcPrintMethod(method, __FILE__, __LINE__,
+                    R"""(
 StatusOr<google::longrunning::Operation>
 $auth_class_name$::$method_name$(
       grpc::ClientContext& context,
       Options options,
-      $request_type$ const& request) {
+      $request_type$ const& request$op_ctx_decl$) {
   auto status = auth_->ConfigureContext(context);
   if (!status.ok()) return status;
-  return child_->$method_name$(context, options, request);
+  return child_->$method_name$(context, options, request$op_ctx_arg$);
 }
 )""");
 
       continue;
     }
     if (IsStreamingRead(method)) {
-      CcPrintMethod(method, __FILE__, __LINE__, R"""(
+      CcPrintMethod(method, __FILE__, __LINE__,
+                    R"""(
 std::unique_ptr<google::cloud::internal::StreamingReadRpc<$response_type$>>
 $auth_class_name$::$method_name$(
    std::shared_ptr<grpc::ClientContext> context,
    Options const& options,
-   $request_type$ const& request) {
+   $request_type$ const& request$op_ctx_shared_decl$) {
   using ErrorStream = ::google::cloud::internal::StreamingReadRpcError<
       $response_type$>;
   auto status = auth_->ConfigureContext(*context);
   if (!status.ok()) return std::make_unique<ErrorStream>(std::move(status));
-  return child_->$method_name$(std::move(context), options, request);
+  return child_->$method_name$(std::move(context), options, request$op_ctx_shared_arg$);
 }
 )""");
       continue;
@@ -225,10 +231,10 @@ std::unique_ptr<google::cloud::internal::StreamingReadRpc<$response_type$>>
                   R"""( $auth_class_name$::$method_name$(
     grpc::ClientContext& context,
     Options const& options,
-    $request_type$ const& request) {
+    $request_type$ const& request$op_ctx_decl$) {
   auto status = auth_->ConfigureContext(context);
   if (!status.ok()) return status;
-  return child_->$method_name$(context, options, request);
+  return child_->$method_name$(context, options, request$op_ctx_arg$);
 }
 )""");
   }
@@ -237,126 +243,129 @@ std::unique_ptr<google::cloud::internal::StreamingReadRpc<$response_type$>>
     // Nothing to do, these are always asynchronous.
     if (IsBidirStreaming(method) || IsLongrunningOperation(method)) continue;
     if (IsStreamingRead(method)) {
-      auto constexpr kDefinition = R"""(
+      CcPrintMethod(method, __FILE__, __LINE__,
+                    R"""(
 std::unique_ptr<::google::cloud::internal::AsyncStreamingReadRpc<
     $response_type$>>
 $auth_class_name$::Async$method_name$(
     google::cloud::CompletionQueue const& cq,
     std::shared_ptr<grpc::ClientContext> context,
     google::cloud::internal::ImmutableOptions options,
-    $request_type$ const& request) {
+    $request_type$ const& request$op_ctx_shared_decl$) {
   using StreamAuth = google::cloud::internal::AsyncStreamingReadRpcAuth<
     $response_type$>;
 
   auto& child = child_;
-  auto call = [child, cq, opts = std::move(options), request](std::shared_ptr<grpc::ClientContext> ctx) {
-    return child->Async$method_name$(cq, std::move(ctx), opts, request);
+  auto call = [child, cq, opts = std::move(options), request$op_ctx_shared_cap$](std::shared_ptr<grpc::ClientContext> ctx) {
+    return child->Async$method_name$(cq, std::move(ctx), opts, request$op_ctx_shared_arg$);
   };
   return std::make_unique<StreamAuth>(
     std::move(context), auth_, StreamAuth::StreamFactory(std::move(call)));
 }
-)""";
-      CcPrintMethod(method, __FILE__, __LINE__, kDefinition);
+)""");
       continue;
     }
     if (IsStreamingWrite(method)) {
-      auto constexpr kDefinition = R"""(
+      CcPrintMethod(method, __FILE__, __LINE__,
+                    R"""(
 std::unique_ptr<::google::cloud::internal::AsyncStreamingWriteRpc<
     $request_type$, $response_type$>>
 $auth_class_name$::Async$method_name$(
     google::cloud::CompletionQueue const& cq,
     std::shared_ptr<grpc::ClientContext> context,
-    google::cloud::internal::ImmutableOptions options) {
+    google::cloud::internal::ImmutableOptions options$op_ctx_shared_decl$) {
   using StreamAuth = google::cloud::internal::AsyncStreamingWriteRpcAuth<
     $request_type$, $response_type$>;
 
-  auto call = [child = child_, cq, options = std::move(options)](
+  auto call = [child = child_, cq, options = std::move(options)$op_ctx_shared_cap$](
                   std::shared_ptr<grpc::ClientContext> ctx) {
-    return child->Async$method_name$(cq, std::move(ctx), options);
+    return child->Async$method_name$(cq, std::move(ctx), options$op_ctx_shared_arg$);
   };
   return std::make_unique<StreamAuth>(
     std::move(context), auth_, StreamAuth::StreamFactory(std::move(call)));
 }
-)""";
-      CcPrintMethod(method, __FILE__, __LINE__, kDefinition);
+)""");
       continue;
     }
     if (IsResponseTypeEmpty(method)) {
-      CcPrintMethod(method, __FILE__, __LINE__, R"""(
+      CcPrintMethod(method, __FILE__, __LINE__,
+                    R"""(
 future<Status>
 $auth_class_name$::Async$method_name$(
       google::cloud::CompletionQueue& cq,
       std::shared_ptr<grpc::ClientContext> context,
       google::cloud::internal::ImmutableOptions options,
-      $request_type$ const& request) {
+      $request_type$ const& request$op_ctx_shared_decl$) {
   return auth_->AsyncConfigureContext(std::move(context)).then(
-      [cq, child = child_, options = std::move(options), request](
+      [cq, child = child_, options = std::move(options), request$op_ctx_shared_cap$](
           future<StatusOr<std::shared_ptr<grpc::ClientContext>>> f) mutable {
         auto context = f.get();
         if (!context) return make_ready_future(std::move(context).status());
         return child->Async$method_name$(
-            cq, *std::move(context), std::move(options), request);
+            cq, *std::move(context), std::move(options), request$op_ctx_shared_arg$);
       });
 }
 )""");
       continue;
     }
-    CcPrintMethod(method, __FILE__, __LINE__, R"""(
+    CcPrintMethod(method, __FILE__, __LINE__,
+                  R"""(
 future<StatusOr<$response_type$>>
 $auth_class_name$::Async$method_name$(
       google::cloud::CompletionQueue& cq,
       std::shared_ptr<grpc::ClientContext> context,
       google::cloud::internal::ImmutableOptions options,
-      $request_type$ const& request) {
+      $request_type$ const& request$op_ctx_shared_decl$) {
   return auth_->AsyncConfigureContext(std::move(context)).then(
-      [cq, child = child_, options = std::move(options), request](
+      [cq, child = child_, options = std::move(options), request$op_ctx_shared_cap$](
           future<StatusOr<std::shared_ptr<grpc::ClientContext>>> f) mutable {
         auto context = f.get();
         if (!context) {
           return make_ready_future(StatusOr<$response_type$>(
               std::move(context).status()));
         }
         return child->Async$method_name$(
-            cq, *std::move(context), std::move(options), request);
+            cq, *std::move(context), std::move(options), request$op_ctx_shared_arg$);
       });
 }
 )""");
   }
 
   // long running operation support methods
   if (HasLongrunningMethod()) {
-    CcPrint(R"""(
+    CcPrint(
+        R"""(
 future<StatusOr<google::longrunning::Operation>>
 $auth_class_name$::AsyncGetOperation(
     google::cloud::CompletionQueue& cq,
     std::shared_ptr<grpc::ClientContext> context,
     google::cloud::internal::ImmutableOptions options,
-    google::longrunning::GetOperationRequest const& request) {
+    google::longrunning::GetOperationRequest const& request$op_ctx_shared_decl$) {
   using ReturnType = StatusOr<google::longrunning::Operation>;
   return auth_->AsyncConfigureContext(std::move(context)).then(
-      [cq, child = child_, options = std::move(options), request](
+      [cq, child = child_, options = std::move(options), request$op_ctx_shared_cap$](
           future<StatusOr<std::shared_ptr<grpc::ClientContext>>> f) mutable {
         auto context = f.get();
         if (!context) {
           return make_ready_future(ReturnType(std::move(context).status()));
         }
         return child->AsyncGetOperation(
-            cq, *std::move(context), std::move(options), request);
+            cq, *std::move(context), std::move(options), request$op_ctx_shared_arg$);
       });
 }
 
 future<Status> $auth_class_name$::AsyncCancelOperation(
     google::cloud::CompletionQueue& cq,
     std::shared_ptr<grpc::ClientContext> context,
     google::cloud::internal::ImmutableOptions options,
-    google::longrunning::CancelOperationRequest const& request) {
+    google::longrunning::CancelOperationRequest const& request$op_ctx_shared_decl$) {
   return auth_->AsyncConfigureContext(std::move(context)).then(
-      [cq, child = child_, options = std::move(options), request](
+      [cq, child = child_, options = std::move(options), request$op_ctx_shared_cap$](
           future<StatusOr<std::shared_ptr<grpc::ClientContext>>> f) mutable {
         auto context = f.get();
         if (!context) return make_ready_future(std::move(context).status());
         return child->AsyncCancelOperation(
-            cq, *std::move(context), std::move(options), request);
+            cq, *std::move(context), std::move(options), request$op_ctx_shared_arg$);
       });
 }
 )""");

generator/internal/codegen_utils_test.cc

@@ -323,6 +323,15 @@ TEST(ProcessCommandLineArgs, ProcessExperimental) {
   EXPECT_THAT(*result, Contains(Pair("experimental", "true")));
 }
 
+TEST(ProcessCommandLineArgs, ProcessExperimentalBigtableOperationContext) {
+  auto result = ProcessCommandLineArgs(
+      "product_path=google/cloud/bigtable/"
+      ",experimental_bigtable_operation_context=true");
+  ASSERT_THAT(result, IsOk());
+  EXPECT_THAT(*result, Contains(Pair("experimental_bigtable_operation_context",
+                                     "true")));
+}
+
 TEST(ProcessCommandLineArgs, ProcessServiceNameMapping) {
   auto result = ProcessCommandLineArgs(
       "product_path=google/cloud/pubsub/"

generator/internal/descriptor_utils.cc

@@ -851,6 +851,37 @@ VarsDictionary CreateServiceVars(
   SetRetryStatusCodeExpression(vars);
   vars["transient_errors_comment"] = TransientErrorsComment(vars);
   SetLongrunningOperationServiceVars(descriptor, vars);
+  auto const experimental_bigtable_operation_context =
+      vars.find("experimental_bigtable_operation_context");
+  if (experimental_bigtable_operation_context != vars.end() &&
+      experimental_bigtable_operation_context->second == "true") {
+    vars["op_ctx_decl"] =
+        ",\n      google::cloud::bigtable_internal::OperationContext& "
+        "operation_context";
+    vars["op_ctx_arg"] = ", operation_context";
+    vars["op_ctx_cap"] = ", &operation_context";
+    vars["op_ctx_stub_decl"] =
+        ",\n    google::cloud::bigtable_internal::OperationContext&";
+    vars["op_ctx_shared_decl"] =
+        ",\n      "
+        "std::shared_ptr<google::cloud::bigtable_internal::OperationContext> "
+        "operation_context";
+    vars["op_ctx_shared_arg"] = ", std::move(operation_context)";
+    vars["op_ctx_shared_cap"] =
+        ", operation_context = std::move(operation_context)";
+    vars["op_ctx_shared_stub_decl"] =
+        ",\n    "
+        "std::shared_ptr<google::cloud::bigtable_internal::OperationContext>";
+  } else {
+    vars["op_ctx_decl"] = "";
+    vars["op_ctx_arg"] = "";
+    vars["op_ctx_cap"] = "";
+    vars["op_ctx_stub_decl"] = "";
+    vars["op_ctx_shared_decl"] = "";
+    vars["op_ctx_shared_arg"] = "";
+    vars["op_ctx_shared_cap"] = "";
+    vars["op_ctx_shared_stub_decl"] = "";
+  }
   return vars;
 }