feat(bigtable): add support for AsyncPrepareQuery (#15457)

diegomarquezp · web-flow · commit a4d34f33c72f · 2025-09-10T14:14:33.000-04:00
* feat(bigtable): add support for AsyncPrepareQuery

* chore: generate bigtable with new config

* chore: manually introduce async prepare query where expected
diff --git a/generator/generator_config.textproto b/generator/generator_config.textproto
@@ -579,7 +579,8 @@ service {
     "MutateRows",
     "ReadModifyWriteRow",
     "ReadRows",
-    "SampleRowKeys"
+    "SampleRowKeys",
+    "PrepareQuery"
   ]
   omit_repo_metadata: true
 }
diff --git a/google/cloud/bigtable/internal/bigtable_auth_decorator.cc b/google/cloud/bigtable/internal/bigtable_auth_decorator.cc
@@ -242,6 +242,27 @@ BigtableAuth::AsyncReadModifyWriteRow(
       });
 }
 
+future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+BigtableAuth::AsyncPrepareQuery(
+    google::cloud::CompletionQueue& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::PrepareQueryRequest const& request) {
+  return auth_->AsyncConfigureContext(std::move(context))
+      .then([cq, child = child_, options = std::move(options),
+             request](future<StatusOr<std::shared_ptr<grpc::ClientContext>>>
+                          f) mutable {
+        auto context = f.get();
+        if (!context) {
+          return make_ready_future(
+              StatusOr<google::bigtable::v2::PrepareQueryResponse>(
+                  std::move(context).status()));
+        }
+        return child->AsyncPrepareQuery(cq, *std::move(context),
+                                        std::move(options), request);
+      });
+}
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace bigtable_internal
 }  // namespace cloud
diff --git a/google/cloud/bigtable/internal/bigtable_auth_decorator.h b/google/cloud/bigtable/internal/bigtable_auth_decorator.h
@@ -124,6 +124,13 @@ class BigtableAuth : public BigtableStub {
       google::cloud::internal::ImmutableOptions options,
       google::bigtable::v2::ReadModifyWriteRowRequest const& request) override;
 
+  future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+  AsyncPrepareQuery(
+      google::cloud::CompletionQueue& cq,
+      std::shared_ptr<grpc::ClientContext> context,
+      google::cloud::internal::ImmutableOptions options,
+      google::bigtable::v2::PrepareQueryRequest const& request) override;
+
  private:
   std::shared_ptr<google::cloud::internal::GrpcAuthenticationStrategy> auth_;
   std::shared_ptr<BigtableStub> child_;
diff --git a/google/cloud/bigtable/internal/bigtable_channel_refresh.cc b/google/cloud/bigtable/internal/bigtable_channel_refresh.cc
@@ -149,6 +149,16 @@ BigtableChannelRefresh::AsyncReadModifyWriteRow(
                                          std::move(options), request);
 }
 
+future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+BigtableChannelRefresh::AsyncPrepareQuery(
+    google::cloud::CompletionQueue& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::PrepareQueryRequest const& request) {
+  return child_->AsyncPrepareQuery(cq, std::move(context), std::move(options),
+                                   request);
+}
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace bigtable_internal
 }  // namespace cloud
diff --git a/google/cloud/bigtable/internal/bigtable_channel_refresh.h b/google/cloud/bigtable/internal/bigtable_channel_refresh.h
@@ -131,6 +131,13 @@ class BigtableChannelRefresh : public BigtableStub {
       google::cloud::internal::ImmutableOptions,
       google::bigtable::v2::ReadModifyWriteRowRequest const& request) override;
 
+  future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+  AsyncPrepareQuery(
+      google::cloud::CompletionQueue& cq,
+      std::shared_ptr<grpc::ClientContext> context,
+      google::cloud::internal::ImmutableOptions options,
+      google::bigtable::v2::PrepareQueryRequest const& request) override;
+
  private:
   std::shared_ptr<BigtableStub> child_;
   std::shared_ptr<ConnectionRefreshState> refresh_state_;
diff --git a/google/cloud/bigtable/internal/bigtable_logging_decorator.cc b/google/cloud/bigtable/internal/bigtable_logging_decorator.cc
@@ -319,6 +319,24 @@ BigtableLogging::AsyncReadModifyWriteRow(
       tracing_options_);
 }
 
+future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+BigtableLogging::AsyncPrepareQuery(
+    google::cloud::CompletionQueue& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::PrepareQueryRequest const& request) {
+  return google::cloud::internal::LogWrapper(
+      [this](google::cloud::CompletionQueue& cq,
+             std::shared_ptr<grpc::ClientContext> context,
+             google::cloud::internal::ImmutableOptions options,
+             google::bigtable::v2::PrepareQueryRequest const& request) {
+        return child_->AsyncPrepareQuery(cq, std::move(context),
+                                         std::move(options), request);
+      },
+      cq, std::move(context), std::move(options), request, __func__,
+      tracing_options_);
+}
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace bigtable_internal
 }  // namespace cloud
diff --git a/google/cloud/bigtable/internal/bigtable_logging_decorator.h b/google/cloud/bigtable/internal/bigtable_logging_decorator.h
@@ -124,6 +124,13 @@ class BigtableLogging : public BigtableStub {
       google::cloud::internal::ImmutableOptions options,
       google::bigtable::v2::ReadModifyWriteRowRequest const& request) override;
 
+  future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+  AsyncPrepareQuery(
+      google::cloud::CompletionQueue& cq,
+      std::shared_ptr<grpc::ClientContext> context,
+      google::cloud::internal::ImmutableOptions options,
+      google::bigtable::v2::PrepareQueryRequest const& request) override;
+
  private:
   std::shared_ptr<BigtableStub> child_;
   TracingOptions tracing_options_;
diff --git a/google/cloud/bigtable/internal/bigtable_metadata_decorator.cc b/google/cloud/bigtable/internal/bigtable_metadata_decorator.cc
@@ -697,6 +697,42 @@ BigtableMetadata::AsyncReadModifyWriteRow(
                                          std::move(options), request);
 }
 
+future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+BigtableMetadata::AsyncPrepareQuery(
+    google::cloud::CompletionQueue& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::PrepareQueryRequest const& request) {
+  std::vector<std::string> params;
+  params.reserve(2);
+
+  static auto* name_matcher = [] {
+    return new google::cloud::internal::RoutingMatcher<
+        google::bigtable::v2::PrepareQueryRequest>{
+        "name=",
+        {
+            {[](google::bigtable::v2::PrepareQueryRequest const& request)
+                 -> std::string const& { return request.instance_name(); },
+             std::regex{"(projects/[^/]+/instances/[^/]+)",
+                        std::regex::optimize}},
+        }};
+  }();
+  name_matcher->AppendParam(request, params);
+
+  if (!request.app_profile_id().empty()) {
+    params.push_back(absl::StrCat(
+        "app_profile_id=", internal::UrlEncode(request.app_profile_id())));
+  }
+
+  if (params.empty()) {
+    SetMetadata(*context, *options);
+  } else {
+    SetMetadata(*context, *options, absl::StrJoin(params, "&"));
+  }
+  return child_->AsyncPrepareQuery(cq, std::move(context), std::move(options),
+                                   request);
+}
+
 void BigtableMetadata::SetMetadata(grpc::ClientContext& context,
                                    Options const& options,
                                    std::string const& request_params) {
diff --git a/google/cloud/bigtable/internal/bigtable_metadata_decorator.h b/google/cloud/bigtable/internal/bigtable_metadata_decorator.h
@@ -124,6 +124,13 @@ class BigtableMetadata : public BigtableStub {
       google::cloud::internal::ImmutableOptions options,
       google::bigtable::v2::ReadModifyWriteRowRequest const& request) override;
 
+  future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+  AsyncPrepareQuery(
+      google::cloud::CompletionQueue& cq,
+      std::shared_ptr<grpc::ClientContext> context,
+      google::cloud::internal::ImmutableOptions options,
+      google::bigtable::v2::PrepareQueryRequest const& request) override;
+
  private:
   void SetMetadata(grpc::ClientContext& context, Options const& options,
                    std::string const& request_params);
diff --git a/google/cloud/bigtable/internal/bigtable_round_robin_decorator.cc b/google/cloud/bigtable/internal/bigtable_round_robin_decorator.cc
@@ -159,6 +159,16 @@ BigtableRoundRobin::AsyncReadModifyWriteRow(
                                           std::move(options), request);
 }
 
+future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+BigtableRoundRobin::AsyncPrepareQuery(
+    google::cloud::CompletionQueue& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::PrepareQueryRequest const& request) {
+  return Child()->AsyncPrepareQuery(cq, std::move(context), std::move(options),
+                                    request);
+}
+
 std::shared_ptr<BigtableStub> BigtableRoundRobin::Child() {
   std::unique_lock<std::mutex> lk(mu_);
   auto const current = current_;
diff --git a/google/cloud/bigtable/internal/bigtable_round_robin_decorator.h b/google/cloud/bigtable/internal/bigtable_round_robin_decorator.h
@@ -122,6 +122,13 @@ class BigtableRoundRobin : public BigtableStub {
       google::cloud::internal::ImmutableOptions options,
       google::bigtable::v2::ReadModifyWriteRowRequest const& request) override;
 
+  future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+  AsyncPrepareQuery(
+      google::cloud::CompletionQueue& cq,
+      std::shared_ptr<grpc::ClientContext> context,
+      google::cloud::internal::ImmutableOptions options,
+      google::bigtable::v2::PrepareQueryRequest const& request) override;
+
  private:
   std::shared_ptr<BigtableStub> Child();
 
diff --git a/google/cloud/bigtable/internal/bigtable_stub.cc b/google/cloud/bigtable/internal/bigtable_stub.cc
@@ -245,6 +245,24 @@ DefaultBigtableStub::AsyncReadModifyWriteRow(
       request, std::move(context));
 }
 
+future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+DefaultBigtableStub::AsyncPrepareQuery(
+    google::cloud::CompletionQueue& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    // NOLINTNEXTLINE(performance-unnecessary-value-param)
+    google::cloud::internal::ImmutableOptions,
+    google::bigtable::v2::PrepareQueryRequest const& request) {
+  return internal::MakeUnaryRpcImpl<google::bigtable::v2::PrepareQueryRequest,
+                                    google::bigtable::v2::PrepareQueryResponse>(
+      cq,
+      [this](grpc::ClientContext* context,
+             google::bigtable::v2::PrepareQueryRequest const& request,
+             grpc::CompletionQueue* cq) {
+        return grpc_stub_->AsyncPrepareQuery(context, request, cq);
+      },
+      request, std::move(context));
+}
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace bigtable_internal
 }  // namespace cloud
diff --git a/google/cloud/bigtable/internal/bigtable_stub.h b/google/cloud/bigtable/internal/bigtable_stub.h
@@ -125,6 +125,13 @@ class BigtableStub {
       std::shared_ptr<grpc::ClientContext> context,
       google::cloud::internal::ImmutableOptions options,
       google::bigtable::v2::ReadModifyWriteRowRequest const& request) = 0;
+
+  virtual future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+  AsyncPrepareQuery(
+      google::cloud::CompletionQueue& cq,
+      std::shared_ptr<grpc::ClientContext> context,
+      google::cloud::internal::ImmutableOptions options,
+      google::bigtable::v2::PrepareQueryRequest const& request) = 0;
 };
 
 class DefaultBigtableStub : public BigtableStub {
@@ -219,6 +226,13 @@ class DefaultBigtableStub : public BigtableStub {
       google::cloud::internal::ImmutableOptions options,
       google::bigtable::v2::ReadModifyWriteRowRequest const& request) override;
 
+  future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+  AsyncPrepareQuery(
+      google::cloud::CompletionQueue& cq,
+      std::shared_ptr<grpc::ClientContext> context,
+      google::cloud::internal::ImmutableOptions options,
+      google::bigtable::v2::PrepareQueryRequest const& request) override;
+
  private:
   std::unique_ptr<google::bigtable::v2::Bigtable::StubInterface> grpc_stub_;
 };
diff --git a/google/cloud/bigtable/internal/bigtable_tracing_stub.cc b/google/cloud/bigtable/internal/bigtable_tracing_stub.cc
@@ -248,6 +248,20 @@ BigtableTracingStub::AsyncReadModifyWriteRow(
   return internal::EndSpan(std::move(context), std::move(span), std::move(f));
 }
 
+future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+BigtableTracingStub::AsyncPrepareQuery(
+    google::cloud::CompletionQueue& cq,
+    std::shared_ptr<grpc::ClientContext> context,
+    google::cloud::internal::ImmutableOptions options,
+    google::bigtable::v2::PrepareQueryRequest const& request) {
+  auto span =
+      internal::MakeSpanGrpc("google.bigtable.v2.Bigtable", "PrepareQuery");
+  internal::OTelScope scope(span);
+  internal::InjectTraceContext(*context, *propagator_);
+  auto f = child_->AsyncPrepareQuery(cq, context, std::move(options), request);
+  return internal::EndSpan(std::move(context), std::move(span), std::move(f));
+}
+
 #endif  // GOOGLE_CLOUD_CPP_HAVE_OPENTELEMETRY
 
 std::shared_ptr<BigtableStub> MakeBigtableTracingStub(
diff --git a/google/cloud/bigtable/internal/bigtable_tracing_stub.h b/google/cloud/bigtable/internal/bigtable_tracing_stub.h
@@ -124,6 +124,13 @@ class BigtableTracingStub : public BigtableStub {
       google::cloud::internal::ImmutableOptions options,
       google::bigtable::v2::ReadModifyWriteRowRequest const& request) override;
 
+  future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+  AsyncPrepareQuery(
+      google::cloud::CompletionQueue& cq,
+      std::shared_ptr<grpc::ClientContext> context,
+      google::cloud::internal::ImmutableOptions options,
+      google::bigtable::v2::PrepareQueryRequest const& request) override;
+
  private:
   std::shared_ptr<BigtableStub> child_;
   std::shared_ptr<opentelemetry::context::propagation::TextMapPropagator>
diff --git a/google/cloud/bigtable/testing/mock_bigtable_stub.h b/google/cloud/bigtable/testing/mock_bigtable_stub.h
@@ -118,6 +118,13 @@ class MockBigtableStub : public bigtable_internal::BigtableStub {
        google::cloud::internal::ImmutableOptions,
        google::bigtable::v2::ReadModifyWriteRowRequest const&),
       (override));
+  MOCK_METHOD(future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>,
+              AsyncPrepareQuery,
+              (google::cloud::CompletionQueue&,
+               std::shared_ptr<grpc::ClientContext>,
+               google::cloud::internal::ImmutableOptions,
+               google::bigtable::v2::PrepareQueryRequest const&),
+              (override));
 };
 
 class MockMutateRowsStream : public google::cloud::internal::StreamingReadRpc<

Original file line number	Diff line number	Diff line change
`@@ -579,7 +579,8 @@ service {`
`579`	`579`	`"MutateRows",`
`580`	`580`	`"ReadModifyWriteRow",`
`581`	`581`	`"ReadRows",`
`582`		`- "SampleRowKeys"`
	`582`	`+ "SampleRowKeys",`
	`583`	`+ "PrepareQuery"`
`583`	`584`	`]`
`584`	`585`	`omit_repo_metadata: true`
`585`	`586`	`}`