impl(v3): add generator support for bespoke WaitForConsistency method

Author: scotthart

generator/generator_config.proto

@@ -167,6 +167,21 @@ message ServiceConfiguration {
   // RPCs. If set to false (the default), a no-op resumption function will be
   // generated.
   bool omit_streaming_updater = 29;
+
+  message BespokeMethod {
+    string name = 1;
+    string return_type = 2;
+    string parameters = 3;
+  }
+
+  // Only added to maintain feature parity when migrating from the handwritten
+  // Bigtable Table Admin class to the generated class. While some attempts were
+  // made to generalize this feature, it currently only supports the
+  // WaitForConsistency method. This functionality can be enhanced later if we
+  // ever need to use this again.
+  // 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;
 }
 
 message DiscoveryDocumentDefinedProduct {

generator/integration_tests/golden_config.textproto

@@ -59,6 +59,13 @@ service {
   ]
   override_service_config_yaml_name: "generator/integration_tests/test2.yaml"
   endpoint_location_style: LOCATION_OPTIONALLY_DEPENDENT
+  bespoke_methods : [
+    {
+    name: "WaitForConsistency",
+    return_type: "StatusOr<google::protobuf::Empty>",
+    parameters: "(google::protobuf::Empty const& request, Options opts = {})"
+    }
+  ]
 }
 
 service {

generator/internal/client_generator.cc

@@ -21,12 +21,51 @@
 #include "generator/internal/predicate_utils.h"
 #include "generator/internal/printer.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_replace.h"
 #include "google/api/client.pb.h"
 #include <google/protobuf/descriptor.h>
 
 namespace google {
 namespace cloud {
 namespace generator_internal {
+namespace {
+std::string FormatBespokeMethodComments(std::string const& method_name) {
+  if (method_name == "WaitForConsistency") {
+    return R"""(
+  // clang-format off
+  ///
+  /// Polls a table until it is consistent or the RetryPolicy is exhausted based
+  /// on a consistency token, that is, if replication has caught up based on the
+  /// provided conditions specified in the token and the check request.
+  ///
+  /// @param request Unary RPCs, such as the one wrapped by this
+  ///     function, receive a single `request` proto message which includes all
+  ///     the inputs for the RPC. In this case, the proto message is a
+  ///     [google.bigtable.admin.v2.CheckConsistencyRequest].
+  ///     Proto messages are converted to C++ classes by Protobuf, using the
+  ///     [Protobuf mapping rules].
+  /// @param opts Optional. Override the class-level options, such as retry and
+  ///     backoff policies.
+  /// @return the result of the RPC. The response message type
+  ///     ([google.bigtable.admin.v2.CheckConsistencyResponse])
+  ///     is mapped to a C++ class using the [Protobuf mapping rules].
+  ///     If the request fails, the [`StatusOr`] contains the error details.
+  ///
+  /// [Protobuf mapping rules]: https://protobuf.dev/reference/cpp/cpp-generated/
+  /// [input iterator requirements]: https://en.cppreference.com/w/cpp/named_req/InputIterator
+  /// [`std::string`]: https://en.cppreference.com/w/cpp/string/basic_string
+  /// [`future`]: @ref google::cloud::future
+  /// [`StatusOr`]: @ref google::cloud::StatusOr
+  /// [`Status`]: @ref google::cloud::Status
+  /// [google.bigtable.admin.v2.CheckConsistencyRequest]: @googleapis_reference_link{google/bigtable/admin/v2/bigtable_table_admin.proto#L909}
+  /// [google.bigtable.admin.v2.CheckConsistencyResponse]: @googleapis_reference_link{google/bigtable/admin/v2/bigtable_table_admin.proto#L948}
+  ///
+  // clang-format on
+)""";
+  }
+  return "";
+}
+}  // namespace
 
 ClientGenerator::ClientGenerator(
     google::protobuf::ServiceDescriptor const* service_descriptor,
@@ -380,6 +419,13 @@ R"""(  std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc<
         __FILE__, __LINE__);
   }
 
+  for (auto const& method : bespoke_methods()) {
+    HeaderPrint("\n");
+    HeaderPrint(FormatBespokeMethodComments(method.name()));
+    HeaderPrint(absl::StrCat(method.return_type(), " ", method.name(),
+                             method.parameters(), ";"));
+  }
+
   HeaderPrint(  // clang-format off
     "\n"
     " private:\n"
@@ -716,6 +762,19 @@ std::unique_ptr<::google::cloud::AsyncStreamingReadWriteRpc<
         __FILE__, __LINE__);
   }
 
+  for (auto const& method : bespoke_methods()) {
+    CcPrint("\n");
+    CcPrint(absl::StrCat(
+        method.return_type(), R"""( $client_class_name$::)""", method.name(),
+        absl::StrReplaceAll(method.parameters(), {{" = {}", ""}}),
+        absl::StrFormat(R"""( {
+  internal::OptionsSpan span(internal::MergeOptions(std::move(opts), options_));
+  return connection_->%s(request);
+}
+)""",
+                        method.name())));
+  }
+
   CcCloseNamespaces();
   return {};
 }

generator/internal/codegen_utils.cc

@@ -111,6 +111,11 @@ void ProcessArgOmitRpc(
   ProcessRepeated("omit_rpc", "omitted_rpcs", command_line_args);
 }
 
+void ProcessArgBespokeMethod(
+    std::vector<std::pair<std::string, std::string>>& command_line_args) {
+  ProcessRepeated("bespoke_method", "bespoke_methods", command_line_args);
+}
+
 void ProcessArgServiceEndpointEnvVar(
     std::vector<std::pair<std::string, std::string>>& command_line_args) {
   auto service_endpoint_env_var =
@@ -269,6 +274,7 @@ ProcessCommandLineArgs(std::string const& parameters) {
   ProcessArgCopyrightYear(command_line_args);
   ProcessArgOmitService(command_line_args);
   ProcessArgOmitRpc(command_line_args);
+  ProcessArgBespokeMethod(command_line_args);
   ProcessArgServiceEndpointEnvVar(command_line_args);
   ProcessArgEmulatorEndpointEnvVar(command_line_args);
   ProcessArgEndpointLocationStyle(command_line_args);

generator/internal/connection_generator.cc

@@ -19,6 +19,7 @@
 #include "generator/internal/pagination.h"
 #include "generator/internal/predicate_utils.h"
 #include "generator/internal/printer.h"
+#include "absl/strings/str_replace.h"
 #include "absl/strings/str_split.h"
 #include <google/protobuf/descriptor.h>
 
@@ -315,6 +316,14 @@ class $connection_class_name$ {
         __FILE__, __LINE__);
   }
 
+  for (auto const& method : bespoke_methods()) {
+    HeaderPrint("\n");
+    HeaderPrint(absl::StrCat(
+        "virtual ", method.return_type(), " ", method.name(),
+        absl::StrReplaceAll(method.parameters(), {{", Options opts = {}", ""}}),
+        ";"));
+  }
+
   // close abstract interface Connection base class
   HeaderPrint("};\n");
 
@@ -490,6 +499,26 @@ future<StatusOr<$response_type$>>
         __FILE__, __LINE__);
   }
 
+  for (auto const& method : bespoke_methods()) {
+    CcPrint("\n");
+    std::string make_return =
+        absl::StrContains(method.return_type(), "future")
+            ? absl::StrCat("google::cloud::make_ready_", method.return_type())
+            : method.return_type();
+
+    CcPrint(
+        absl::StrCat(method.return_type(), R"""( $connection_class_name$::)""",
+                     method.name(),
+                     absl::StrReplaceAll(method.parameters(),
+                                         {{" request, Options opts = {}", ""}}),
+                     " {\n",
+                     absl::StrFormat(R"""(  return %s(
+      Status(StatusCode::kUnimplemented, "not implemented"));
+}
+)""",
+                                     make_return)));
+  }
+
   if (HasGenerateGrpcTransport()) {
     EmitFactoryFunctionDefinition(EndpointLocationStyle());
   }

generator/internal/connection_impl_generator.cc

@@ -19,6 +19,7 @@
 #include "generator/internal/predicate_utils.h"
 #include "generator/internal/printer.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_replace.h"
 #include <google/protobuf/descriptor.h>
 
 namespace google {
@@ -120,6 +121,14 @@ class $connection_class_name$Impl
     HeaderPrintMethod(method, __FILE__, __LINE__, AsyncMethodDeclaration());
   }
 
+  for (auto const& method : bespoke_methods()) {
+    HeaderPrint("\n");
+    HeaderPrint(absl::StrCat(
+        method.return_type(), " ", method.name(),
+        absl::StrReplaceAll(method.parameters(), {{", Options opts = {}", ""}}),
+        " override;"));
+  }
+
   HeaderPrint(R"""(
  private:
   std::unique_ptr<google::cloud::BackgroundThreads> background_;

generator/internal/service_code_generator.cc

@@ -482,6 +482,20 @@ void ServiceCodeGenerator::SetMethods() {
   for (auto const& mixin_method : mixin_methods_) {
     methods_.emplace_back(mixin_method.method.get());
   }
+
+  auto bespoke_methods_var = service_vars_.find("bespoke_methods");
+  if (bespoke_methods_var != service_vars_.end()) {
+    auto methods = absl::StrSplit(bespoke_methods_var->second, ',');
+    for (auto const& method : methods) {
+      std::vector<std::string> pieces = absl::StrSplit(method, "@@");
+      assert(pieces.size() == 3);
+      cpp::generator::ServiceConfiguration::BespokeMethod bespoke_method;
+      bespoke_method.set_name(SafeReplaceAll(pieces[0], "@", ","));
+      bespoke_method.set_return_type(SafeReplaceAll(pieces[1], "@", ","));
+      bespoke_method.set_parameters(SafeReplaceAll(pieces[2], "@", ","));
+      bespoke_methods_.emplace_back(std::move(bespoke_method));
+    }
+  }
 }
 
 std::string ServiceCodeGenerator::GetPbIncludeByTransport() const {

generator/internal/service_code_generator.h

@@ -68,6 +68,10 @@ class ServiceCodeGenerator : public GeneratorInterface {
   std::string vars(std::string const& key) const;
   MethodDescriptorList const& methods() const { return methods_; }
   MethodDescriptorList const& async_methods() const { return async_methods_; }
+  std::vector<cpp::generator::ServiceConfiguration::BespokeMethod> const&
+  bespoke_methods() const {
+    return bespoke_methods_;
+  }
   void SetVars(absl::string_view header_path);
   VarsDictionary MergeServiceAndMethodVars(
       google::protobuf::MethodDescriptor const& method) const;
@@ -268,6 +272,8 @@ class ServiceCodeGenerator : public GeneratorInterface {
   bool pb_h_system_includes_ = false;
   MethodDescriptorList methods_;
   MethodDescriptorList async_methods_;
+  std::vector<cpp::generator::ServiceConfiguration::BespokeMethod>
+      bespoke_methods_;
   Printer header_;
   Printer cc_;
   std::vector<MixinMethod> mixin_methods_;

generator/standalone_main.cc

@@ -281,6 +281,14 @@ std::vector<std::future<google::cloud::Status>> GenerateCodeFromProtos(
       args.emplace_back(absl::StrCat("--cpp_codegen_opt=omit_rpc=",
                                      SafeReplaceAll(omit_rpc, ",", "@")));
     }
+    for (auto const& bespoke_method : service.bespoke_methods()) {
+      args.emplace_back(absl::StrCat(
+          "--cpp_codegen_opt=bespoke_method=",
+          absl::StrJoin({SafeReplaceAll(bespoke_method.name(), ",", "@"),
+                         SafeReplaceAll(bespoke_method.return_type(), ",", "@"),
+                         SafeReplaceAll(bespoke_method.parameters(), ",", "@")},
+                        "@@")));
+    }
     for (auto const& retry_code : service.retryable_status_codes()) {
       args.emplace_back("--cpp_codegen_opt=retry_status_code=" + retry_code);
     }