feat(bigtable): implement PrepareQuery in DataConnectionImpl (#15661)

* feat(bigtable): implement PrepareQuery in DataConnectionImpl

Author: mpeddada1

google/cloud/bigtable/internal/data_connection_impl.cc

@@ -619,8 +619,29 @@ DataConnectionImpl::AsyncReadRow(std::string const& table_name,
 }
 
 StatusOr<bigtable::PreparedQuery> DataConnectionImpl::PrepareQuery(
-    bigtable::PrepareQueryParams const&) {
-  return 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 response = google::cloud::internal::RetryLoop(
+      retry_policy(*current), backoff_policy(*current),
+      Idempotency::kNonIdempotent,
+      [this](grpc::ClientContext& context, Options const& options,
+             google::bigtable::v2::PrepareQueryRequest const& request) {
+        return stub_->PrepareQuery(context, options, request);
+      },
+      *current, request, __func__);
+  if (!response) {
+    return std::move(response).status();
+  }
+  return bigtable::PreparedQuery(background_->cq(), params.instance,
+                                 std::move(params.sql_statement),
+                                 *std::move(response));
 }
 
 future<StatusOr<bigtable::PreparedQuery>> DataConnectionImpl::AsyncPrepareQuery(

google/cloud/bigtable/internal/data_connection_impl_test.cc

@@ -2730,13 +2730,45 @@ TEST_F(DataConnectionTest, ExecuteQuery) {
               StatusIs(StatusCode::kUnimplemented));
 }
 
-TEST_F(DataConnectionTest, PrepareQuery) {
-  auto conn = TestConnection(std::make_shared<MockBigtableStub>());
+TEST_F(DataConnectionTest, PrepareQuerySuccess) {
+  auto mock = std::make_shared<MockBigtableStub>();
+  EXPECT_CALL(*mock, PrepareQuery)
+      .WillOnce([](grpc::ClientContext&, Options const&,
+                   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());
+        v2::PrepareQueryResponse response;
+        return response;
+      });
+
+  auto conn = TestConnection(std::move(mock));
   internal::OptionsSpan span(CallOptions());
-  EXPECT_THAT(conn->PrepareQuery(bigtable::PrepareQueryParams{
-                  bigtable::InstanceResource(google::cloud::Project(""), ""),
-                  bigtable::SqlStatement("")}),
-              StatusIs(StatusCode::kUnimplemented));
+  auto params = bigtable::PrepareQueryParams{
+      bigtable::InstanceResource(google::cloud::Project("the-project"),
+                                 "the-instance"),
+      bigtable::SqlStatement("SELECT * FROM the-table")};
+  auto prepared_query = conn->PrepareQuery(params);
+  ASSERT_STATUS_OK(prepared_query);
+  EXPECT_EQ(prepared_query->instance(), params.instance);
+  EXPECT_EQ(prepared_query->sql_statement(), params.sql_statement);
+}
+
+TEST_F(DataConnectionTest, PrepareQueryPermanentError) {
+  auto mock = std::make_shared<MockBigtableStub>();
+  EXPECT_CALL(*mock, PrepareQuery)
+      .WillOnce(
+          [](grpc::ClientContext&, Options const&,
+             v2::PrepareQueryRequest const&) { return PermanentError(); });
+
+  auto conn = TestConnection(std::move(mock));
+  internal::OptionsSpan span(CallOptions());
+  auto result = conn->PrepareQuery(bigtable::PrepareQueryParams{
+      bigtable::InstanceResource(google::cloud::Project("the-project"),
+                                 "the-instance"),
+      bigtable::SqlStatement("SELECT * FROM the-table")});
+  EXPECT_THAT(result, StatusIs(StatusCode::kPermissionDenied));
 }
 
 }  // namespace