feat(bigtable): implement async prepare query for DataConnectionImpl (#15667)

* feat(bigtable): implement async prepare query for DataConnectionImpl

Author: mpeddada1

google/cloud/bigtable/internal/data_connection_impl.cc

@@ -645,9 +645,39 @@ StatusOr<bigtable::PreparedQuery> DataConnectionImpl::PrepareQuery(
 }
 
 future<StatusOr<bigtable::PreparedQuery>> DataConnectionImpl::AsyncPrepareQuery(
-    bigtable::PrepareQueryParams const&) {
-  return make_ready_future<StatusOr<bigtable::PreparedQuery>>(
-      Status(StatusCode::kUnimplemented, "not implemented"));
+    bigtable::PrepareQueryParams const& params) {
+  auto current = google::cloud::internal::SaveCurrentOptions();
+  google::bigtable::v2::PrepareQueryRequest request;
+  request.set_instance_name(params.instance.FullName());
+  request.set_app_profile_id(app_profile_id(*current));
+  request.set_query(params.sql_statement.sql());
+  for (auto const& p : params.sql_statement.params()) {
+    (*request.mutable_param_types())[p.first] = p.second.type();
+  }
+  auto retry = retry_policy(*current);
+  auto backoff = backoff_policy(*current);
+  return google::cloud::internal::AsyncRetryLoop(
+             std::move(retry), std::move(backoff), Idempotency::kNonIdempotent,
+             background_->cq(),
+             [this](CompletionQueue& cq,
+                    std::shared_ptr<grpc::ClientContext> context,
+                    google::cloud::internal::ImmutableOptions options,
+                    google::bigtable::v2::PrepareQueryRequest const& request) {
+               return stub_->AsyncPrepareQuery(cq, std::move(context),
+                                               std::move(options), request);
+             },
+             std::move(current), request, __func__)
+      .then([this, params = std::move(params)](
+                future<StatusOr<google::bigtable::v2::PrepareQueryResponse>>
+                    future) -> StatusOr<bigtable::PreparedQuery> {
+        auto response = future.get();
+        if (!response) {
+          return std::move(response).status();
+        }
+        return bigtable::PreparedQuery(background_->cq(), params.instance,
+                                       params.sql_statement,
+                                       *std::move(response));
+      });
 }
 
 StatusOr<bigtable::RowStream> DataConnectionImpl::ExecuteQuery(

google/cloud/bigtable/internal/data_connection_impl_test.cc

@@ -2771,6 +2771,49 @@ TEST_F(DataConnectionTest, PrepareQueryPermanentError) {
   EXPECT_THAT(result, StatusIs(StatusCode::kPermissionDenied));
 }
 
+TEST_F(DataConnectionTest, AsyncPrepareQuerySuccess) {
+  auto mock = std::make_shared<MockBigtableStub>();
+  EXPECT_CALL(*mock, AsyncPrepareQuery)
+      .WillOnce([](CompletionQueue const&, auto, auto,
+                   v2::PrepareQueryRequest const& request) {
+        EXPECT_EQ(kAppProfile, request.app_profile_id());
+        EXPECT_EQ("projects/the-project/instances/the-instance",
+                  request.instance_name());
+        EXPECT_EQ("SELECT * FROM the-table", request.query());
+        return make_ready_future(make_status_or(v2::PrepareQueryResponse{}));
+      });
+
+  auto conn = TestConnection(std::move(mock));
+  internal::OptionsSpan span(CallOptions());
+  auto params = bigtable::PrepareQueryParams{
+      bigtable::InstanceResource(google::cloud::Project("the-project"),
+                                 "the-instance"),
+      bigtable::SqlStatement("SELECT * FROM the-table")};
+  auto future = conn->AsyncPrepareQuery(params);
+  auto result = future.get();
+  ASSERT_STATUS_OK(result);
+}
+
+TEST_F(DataConnectionTest, AsyncPrepareQueryPermanentError) {
+  auto mock = std::make_shared<MockBigtableStub>();
+  EXPECT_CALL(*mock, AsyncPrepareQuery)
+      .WillOnce(
+          [](CompletionQueue&, auto, auto, v2::PrepareQueryRequest const&) {
+            return make_ready_future<StatusOr<v2::PrepareQueryResponse>>(
+                PermanentError());
+          });
+
+  auto conn = TestConnection(std::move(mock));
+  internal::OptionsSpan span(CallOptions());
+  auto params = bigtable::PrepareQueryParams{
+      bigtable::InstanceResource(google::cloud::Project("the-project"),
+                                 "the-instance"),
+      bigtable::SqlStatement("SELECT * FROM the-table")};
+  auto future = conn->AsyncPrepareQuery(params);
+  auto result = future.get();
+  EXPECT_THAT(result, StatusIs(StatusCode::kPermissionDenied));
+}
+
 }  // namespace
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace bigtable_internal