impl(v3): add generator support for connection cq access (#15912)

scotthart · web-flow · commit 2b9130f6b284 · 2026-01-27T20:03:52.000Z
diff --git a/generator/generator_config.proto b/generator/generator_config.proto
@@ -167,6 +167,12 @@ message ServiceConfiguration {
   // RPCs. If set to false (the default), a no-op resumption function will be
   // generated.
   bool omit_streaming_updater = 29;
+
+  // In rare cases, specifically bigtable::WaitForConsistency, the
+  // CompletionQueue from the BackgroundThreads owned by the Connection is
+  // needed elsewhere. This emits a protected accessor and friend function for
+  // that purpose.
+  bool emit_completion_queue_accessor = 30;
 }
 
 message DiscoveryDocumentDefinedProduct {
diff --git a/generator/generator_config.textproto b/generator/generator_config.textproto
@@ -588,6 +588,7 @@ service {
     {rpc_name: "BigtableTableAdmin.CheckConsistency", idempotency: IDEMPOTENT}
   ]
   omit_repo_metadata: true
+  emit_completion_queue_accessor: true
 }
 
 # Billing
diff --git a/generator/internal/connection_generator.cc b/generator/internal/connection_generator.cc
@@ -57,6 +57,8 @@ Status ConnectionGenerator::GenerateHeader() {
       {vars("idempotency_policy_header_path"), vars("retry_traits_header_path"),
        HasLongrunningMethod() ? "google/cloud/no_await_tag.h" : "",
        IsExperimental() ? "google/cloud/experimental_tag.h" : "",
+       HasEmitCompletionQueueAccessor() ? "google/cloud/completion_queue.h"
+                                        : "",
        "google/cloud/backoff_policy.h",
        HasLongrunningMethod() || HasAsyncMethod() ? "google/cloud/future.h"
                                                   : "",
@@ -91,7 +93,21 @@ Status ConnectionGenerator::GenerateHeader() {
     }
   }
 
-  auto result = HeaderOpenNamespaces();
+  Status result;
+  if (HasEmitCompletionQueueAccessor()) {
+    result = HeaderOpenNamespaces();
+    if (!result.ok()) return result;
+    HeaderPrint(R"""(class $connection_class_name$;)""");
+    HeaderCloseNamespaces();
+
+    result = HeaderOpenNamespaces(NamespaceType::kInternal);
+    if (!result.ok()) return result;
+    HeaderPrint(
+        R"""(StatusOr<CompletionQueue> completion_queue($product_namespace$::$connection_class_name$ const& conn);)""");
+    HeaderCloseNamespaces();
+  }
+
+  result = HeaderOpenNamespaces();
   if (!result.ok()) return result;
 
   HeaderPrint(R"""(
@@ -315,6 +331,13 @@ class $connection_class_name$ {
         __FILE__, __LINE__);
   }
 
+  if (HasEmitCompletionQueueAccessor()) {
+    HeaderPrint(R"""( protected:
+  friend StatusOr<CompletionQueue> $product_internal_namespace$::completion_queue(
+      $connection_class_name$ const& conn);
+  virtual StatusOr<CompletionQueue> completion_queue() const;
+)""");
+  }
   // close abstract interface Connection base class
   HeaderPrint("};\n");
 
@@ -490,6 +513,15 @@ future<StatusOr<$response_type$>>
         __FILE__, __LINE__);
   }
 
+  if (HasEmitCompletionQueueAccessor()) {
+    CcPrint(
+        R"""(
+StatusOr<CompletionQueue> $connection_class_name$::completion_queue() const {
+  return Status(StatusCode::kUnimplemented, "not implemented");
+}
+)""");
+  }
+
   if (HasGenerateGrpcTransport()) {
     EmitFactoryFunctionDefinition(EndpointLocationStyle());
   }
diff --git a/generator/internal/connection_impl_generator.cc b/generator/internal/connection_impl_generator.cc
@@ -133,6 +133,13 @@ class $connection_class_name$Impl
           std::make_shared<google::cloud::internal::InvocationIdGenerator>();)""");
   }
 
+  if (HasEmitCompletionQueueAccessor()) {
+    HeaderPrint(
+        R"""(
+  StatusOr<CompletionQueue> completion_queue() const override;
+)""");
+  }
+
   // This closes the *ConnectionImpl class definition.
   HeaderPrint("\n};\n");
 
@@ -208,6 +215,15 @@ std::unique_ptr<PollingPolicy> polling_policy(Options const& options) {
 } // namespace
 )""");
 
+  if (HasEmitCompletionQueueAccessor()) {
+    CcPrint(R"""(
+StatusOr<CompletionQueue> completion_queue(
+    $product_namespace$::$connection_class_name$ const& conn) {
+  return conn.completion_queue();
+}
+)""");
+  }
+
   // streaming updater functions
   if (!OmitStreamingUpdater(vars())) {
     for (auto const& method : methods()) {
@@ -242,6 +258,14 @@ void $service_name$$method_name$StreamingUpdater(
     CcPrintMethod(method, __FILE__, __LINE__, AsyncMethodDefinition(method));
   }
 
+  if (HasEmitCompletionQueueAccessor()) {
+    CcPrint(R"""(
+StatusOr<CompletionQueue> $connection_class_name$Impl::completion_queue() const {
+  return background_->cq();
+}
+)""");
+  }
+
   CcCloseNamespaces();
   return {};
 }
diff --git a/generator/internal/service_code_generator.cc b/generator/internal/service_code_generator.cc
@@ -511,6 +511,11 @@ bool ServiceCodeGenerator::IsDeprecated() const {
   return service_descriptor_->options().deprecated();
 }
 
+bool ServiceCodeGenerator::HasEmitCompletionQueueAccessor() const {
+  return vars().find("emit_completion_queue_accessor") != vars().end() &&
+         vars().at("emit_completion_queue_accessor") == "true";
+}
+
 }  // namespace generator_internal
 }  // namespace cloud
 }  // namespace google
diff --git a/generator/internal/service_code_generator.h b/generator/internal/service_code_generator.h
@@ -246,6 +246,8 @@ class ServiceCodeGenerator : public GeneratorInterface {
    */
   bool IsDeprecated() const;
 
+  bool HasEmitCompletionQueueAccessor() const;
+
  private:
   void SetMethods();
 
diff --git a/generator/standalone_main.cc b/generator/standalone_main.cc
@@ -296,6 +296,11 @@ std::vector<std::future<google::cloud::Status>> GenerateCodeFromProtos(
     if (service.omit_streaming_updater()) {
       args.emplace_back("--cpp_codegen_opt=omit_streaming_updater=true");
     }
+    if (service.emit_completion_queue_accessor()) {
+      args.emplace_back(
+          "--cpp_codegen_opt=emit_completion_queue_accessor=true");
+    }
+
     if (service.generate_round_robin_decorator()) {
       args.emplace_back(
           "--cpp_codegen_opt=generate_round_robin_decorator=true");
diff --git a/google/cloud/bigtable/admin/bigtable_table_admin_connection.cc b/google/cloud/bigtable/admin/bigtable_table_admin_connection.cc
@@ -382,6 +382,11 @@ BigtableTableAdminConnection::AsyncCheckConsistency(
       Status(StatusCode::kUnimplemented, "not implemented"));
 }
 
+StatusOr<CompletionQueue> BigtableTableAdminConnection::completion_queue()
+    const {
+  return Status(StatusCode::kUnimplemented, "not implemented");
+}
+
 std::shared_ptr<BigtableTableAdminConnection> MakeBigtableTableAdminConnection(
     Options options) {
   internal::CheckExpectedOptions<CommonOptionList, GrpcOptionList,
diff --git a/google/cloud/bigtable/admin/bigtable_table_admin_connection.h b/google/cloud/bigtable/admin/bigtable_table_admin_connection.h
@@ -22,6 +22,7 @@
 #include "google/cloud/bigtable/admin/bigtable_table_admin_connection_idempotency_policy.h"
 #include "google/cloud/bigtable/admin/internal/bigtable_table_admin_retry_traits.h"
 #include "google/cloud/backoff_policy.h"
+#include "google/cloud/completion_queue.h"
 #include "google/cloud/future.h"
 #include "google/cloud/internal/retry_policy_impl.h"
 #include "google/cloud/no_await_tag.h"
@@ -34,6 +35,27 @@
 #include "google/longrunning/operations.grpc.pb.h"
 #include <memory>
 
+namespace google {
+namespace cloud {
+namespace bigtable_admin {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+class BigtableTableAdminConnection;
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace bigtable_admin
+}  // namespace cloud
+}  // namespace google
+
+namespace google {
+namespace cloud {
+namespace bigtable_admin_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+StatusOr<CompletionQueue> completion_queue(
+    bigtable_admin::BigtableTableAdminConnection const& conn);
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace bigtable_admin_internal
+}  // namespace cloud
+}  // namespace google
+
 namespace google {
 namespace cloud {
 namespace bigtable_admin {
@@ -359,6 +381,11 @@ class BigtableTableAdminConnection {
       StatusOr<google::bigtable::admin::v2::CheckConsistencyResponse>>
   AsyncCheckConsistency(
       google::bigtable::admin::v2::CheckConsistencyRequest const& request);
+
+ protected:
+  friend StatusOr<CompletionQueue> bigtable_admin_internal::completion_queue(
+      BigtableTableAdminConnection const& conn);
+  virtual StatusOr<CompletionQueue> completion_queue() const;
 };
 
 /**
diff --git a/google/cloud/bigtable/admin/internal/bigtable_table_admin_connection_impl.cc b/google/cloud/bigtable/admin/internal/bigtable_table_admin_connection_impl.cc
@@ -60,6 +60,11 @@ std::unique_ptr<PollingPolicy> polling_policy(Options const& options) {
 
 }  // namespace
 
+StatusOr<CompletionQueue> completion_queue(
+    bigtable_admin::BigtableTableAdminConnection const& conn) {
+  return conn.completion_queue();
+}
+
 BigtableTableAdminConnectionImpl::BigtableTableAdminConnectionImpl(
     std::unique_ptr<google::cloud::BackgroundThreads> background,
     std::shared_ptr<bigtable_admin_internal::BigtableTableAdminStub> stub,
@@ -1293,6 +1298,11 @@ BigtableTableAdminConnectionImpl::AsyncCheckConsistency(
       std::move(current), std::move(request_copy), __func__);
 }
 
+StatusOr<CompletionQueue> BigtableTableAdminConnectionImpl::completion_queue()
+    const {
+  return background_->cq();
+}
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace bigtable_admin_internal
 }  // namespace cloud
diff --git a/google/cloud/bigtable/admin/internal/bigtable_table_admin_connection_impl.h b/google/cloud/bigtable/admin/internal/bigtable_table_admin_connection_impl.h
@@ -238,6 +238,7 @@ class BigtableTableAdminConnectionImpl
   std::unique_ptr<google::cloud::BackgroundThreads> background_;
   std::shared_ptr<bigtable_admin_internal::BigtableTableAdminStub> stub_;
   Options options_;
+  StatusOr<CompletionQueue> completion_queue() const override;
 };
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

Original file line number	Diff line number	Diff line change
`@@ -588,6 +588,7 @@ service {`
`588`	`588`	`{rpc_name: "BigtableTableAdmin.CheckConsistency", idempotency: IDEMPOTENT}`
`589`	`589`	`]`
`590`	`590`	`omit_repo_metadata: true`
	`591`	`+ emit_completion_queue_accessor: true`
`591`	`592`	`}`
`592`	`593`
`593`	`594`	`# Billing`