feat(bigtable): implement execute query in data connection impl (#15694)

* feat(bigtable): implement execute query in data connection impl

Author: mpeddada1

google/cloud/bigtable/bound_query.cc

@@ -33,7 +33,7 @@ std::unordered_map<std::string, Value> const& BoundQuery::parameters() const {
 
 InstanceResource const& BoundQuery::instance() const { return instance_; }
 
-google::bigtable::v2::ExecuteQueryRequest BoundQuery::ToRequestProto() {
+google::bigtable::v2::ExecuteQueryRequest BoundQuery::ToRequestProto() const {
   google::bigtable::v2::ExecuteQueryRequest result;
   *result.mutable_instance_name() = instance_.FullName();
   auto prepared_query = query_plan_->prepared_query();

google/cloud/bigtable/bound_query.h

@@ -47,7 +47,7 @@ class BoundQuery {
   std::unordered_map<std::string, Value> const& parameters() const;
   InstanceResource const& instance() const;
 
-  google::bigtable::v2::ExecuteQueryRequest ToRequestProto();
+  google::bigtable::v2::ExecuteQueryRequest ToRequestProto() const;
 
  private:
   friend class PreparedQuery;

google/cloud/bigtable/client_test.cc

@@ -13,9 +13,14 @@
 // limitations under the License.
 
 #include "google/cloud/bigtable/client.h"
+#include "google/cloud/bigtable/mocks/mock_data_connection.h"
+#include "google/cloud/bigtable/mocks/mock_query_row.h"
+#include "google/cloud/bigtable/mocks/mock_row_reader.h"
+#include "google/cloud/bigtable/query_row.h"
+#include "google/cloud/stream_range.h"
 #include "google/cloud/testing_util/fake_completion_queue_impl.h"
 #include "google/cloud/testing_util/status_matchers.h"
-#include "mocks/mock_data_connection.h"
+#include <google/protobuf/text_format.h>
 
 namespace google {
 namespace cloud {
@@ -24,7 +29,10 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 namespace {
 
 using ::google::bigtable::v2::PrepareQueryResponse;
+using ::google::cloud::testing_util::FakeCompletionQueueImpl;
+using ::google::cloud::testing_util::IsOkAndHolds;
 using ::google::cloud::testing_util::StatusIs;
+using ::testing::Return;
 
 TEST(Client, PrepareQuery) {
   auto fake_cq_impl = std::make_shared<testing_util::FakeCompletionQueueImpl>();
@@ -71,23 +79,91 @@ TEST(Client, AsyncPrepareQuery) {
   fake_cq_impl->SimulateCompletion(false);
 }
 
-TEST(Client, ExecuteQuery) {
-  auto fake_cq_impl = std::make_shared<testing_util::FakeCompletionQueueImpl>();
-  auto conn = MakeDataConnection();
-  Client client(conn);
+class MockQueryRowSource : public ResultSourceInterface {
+ public:
+  MOCK_METHOD(StatusOr<QueryRow>, NextRow, (), (override));
+  MOCK_METHOD(absl::optional<google::bigtable::v2::ResultSetMetadata>, Metadata,
+              (), (override));
+};
+
+TEST(ClientTest, ExecuteQuery) {
+  auto conn_mock = std::make_shared<bigtable_mocks::MockDataConnection>();
+  auto fake_cq_impl = std::make_shared<FakeCompletionQueueImpl>();
+  auto refresh_fn = []() {
+    return make_ready_future(
+        StatusOr<google::bigtable::v2::PrepareQueryResponse>(
+            Status{StatusCode::kUnimplemented, "not implemented"}));
+  };
+  PrepareQueryResponse pq_response;
+  pq_response.set_prepared_query("test-pq-id-54321");
+  auto constexpr kResultMetadataText = R"pb(
+    proto_schema {
+      columns {
+        name: "key"
+        type { string_type {} }
+      }
+      columns {
+        name: "val"
+        type { string_type {} }
+      }
+    }
+  )pb";
+  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(
+      kResultMetadataText, pq_response.mutable_metadata()));
+  EXPECT_CALL(*conn_mock, ExecuteQuery)
+      .WillOnce(
+          [&](bigtable::ExecuteQueryParams const&) -> StatusOr<RowStream> {
+            auto mock_source = std::make_unique<MockQueryRowSource>();
+            EXPECT_CALL(*mock_source, Metadata)
+                .WillRepeatedly(Return(pq_response.metadata()));
+
+            testing::InSequence s;
+            EXPECT_CALL(*mock_source, NextRow)
+                .WillOnce(Return(bigtable_mocks::MakeQueryRow(
+                    {{"key", bigtable::Value("r1")},
+                     {"val", bigtable::Value("v1")}})));
+            EXPECT_CALL(*mock_source, NextRow)
+                .WillOnce(Return(bigtable_mocks::MakeQueryRow(
+                    {{"key", bigtable::Value("r2")},
+                     {"val", bigtable::Value("v2")}})));
+            EXPECT_CALL(*mock_source, NextRow)
+                // Signal end of stream
+                .WillOnce(
+                    Return(Status(StatusCode::kOutOfRange, "End of stream")));
+
+            // Create RowStream with the mock result source
+            RowStream row_stream(std::move(mock_source));
+            return StatusOr<RowStream>(std::move(row_stream));
+          });
+
+  Client client(conn_mock);
   InstanceResource instance(Project("test-project"), "test-instance");
   SqlStatement sql("SELECT * FROM `test-table`");
-  auto prepared_query = PreparedQuery(CompletionQueue{fake_cq_impl}, instance,
-                                      sql, PrepareQueryResponse{});
+  auto query_plan = bigtable_internal::QueryPlan::Create(
+      CompletionQueue(fake_cq_impl), std::move(pq_response),
+      std::move(refresh_fn));
+  auto prepared_query = PreparedQuery(instance, sql, std::move(query_plan));
   auto bound_query = prepared_query.BindParameters({});
-  auto row_stream = client.ExecuteQuery(std::move(bound_query));
-  // We expect a row stream with a single unimplemented status row while
-  // this is not implemented.
-  for (auto const& row : row_stream) {
-    EXPECT_THAT(row.status(),
-                StatusIs(StatusCode::kUnimplemented, "not implemented"));
+  RowStream row_stream = client.ExecuteQuery(std::move(bound_query));
+  std::vector<StatusOr<bigtable::QueryRow>> rows;
+  for (auto const& row : std::move(row_stream)) {
+    rows.push_back(row);
   }
-  EXPECT_EQ(1, std::distance(row_stream.begin(), row_stream.end()));
+
+  ASSERT_EQ(rows.size(), 3);
+  ASSERT_STATUS_OK(rows[0]);
+  auto const& row1 = *rows[0];
+  EXPECT_EQ(row1.columns().size(), 2);
+  EXPECT_THAT(row1.values()[0].get<std::string>(), IsOkAndHolds("r1"));
+  EXPECT_THAT(row1.values()[1].get<std::string>(), IsOkAndHolds("v1"));
+
+  ASSERT_STATUS_OK(rows[1]);
+  auto const& row2 = *rows[1];
+  EXPECT_EQ(row2.columns().size(), 2);
+  EXPECT_THAT(row2.values()[0].get<std::string>(), IsOkAndHolds("r2"));
+  EXPECT_THAT(row2.values()[1].get<std::string>(), IsOkAndHolds("v2"));
+
+  EXPECT_THAT(rows[2], StatusIs(StatusCode::kOutOfRange, "End of stream"));
 
   // Cancel all pending operations, satisfying any remaining futures.
   fake_cq_impl->SimulateCompletion(false);

google/cloud/bigtable/internal/data_connection_impl.cc

@@ -20,6 +20,9 @@
 #include "google/cloud/bigtable/internal/default_row_reader.h"
 #include "google/cloud/bigtable/internal/defaults.h"
 #include "google/cloud/bigtable/internal/operation_context.h"
+#include "google/cloud/bigtable/internal/partial_result_set_reader.h"
+#include "google/cloud/bigtable/internal/partial_result_set_resume.h"
+#include "google/cloud/bigtable/internal/partial_result_set_source.h"
 #include "google/cloud/bigtable/options.h"
 #include "google/cloud/bigtable/results.h"
 #include "google/cloud/background_threads.h"
@@ -29,6 +32,8 @@
 #include "google/cloud/internal/make_status.h"
 #include "google/cloud/internal/random.h"
 #include "google/cloud/internal/retry_loop.h"
+#include "google/cloud/internal/streaming_read_rpc.h"
+#include "absl/functional/bind_front.h"
 #include <memory>
 #include <string>
 
@@ -77,6 +82,83 @@ bigtable::RowReader ReadRowsHelper(
   return MakeRowReader(std::move(impl));
 }
 
+class DefaultPartialResultSetReader
+    : public bigtable_internal::PartialResultSetReader {
+ public:
+  DefaultPartialResultSetReader(std::shared_ptr<grpc::ClientContext> context,
+                                std::unique_ptr<internal::StreamingReadRpc<
+                                    google::bigtable::v2::ExecuteQueryResponse>>
+                                    reader)
+      : context_(std::move(context)), reader_(std::move(reader)) {}
+
+  ~DefaultPartialResultSetReader() override = default;
+
+  void TryCancel() override { context_->TryCancel(); }
+
+  bool Read(absl::optional<std::string> const&,
+            bigtable_internal::UnownedPartialResultSet& result_set) override {
+    while (true) {
+      google::bigtable::v2::ExecuteQueryResponse response;
+      absl::optional<google::cloud::Status> status = reader_->Read(&response);
+
+      if (status.has_value()) {
+        // Stream has ended or an error occurred.
+        final_status_ = *std::move(status);
+        return false;
+      }
+
+      // Message successfully read into response.
+      if (response.has_results()) {
+        result_set.result = std::move(*response.mutable_results());
+        result_set.resumption = false;
+        return true;
+      }
+
+      // Ignore metadata from the stream because PartialResultSetSource already
+      // has it set (in ExecuteQuery).
+      // TODO(#15701): Investigate expected behavior for processing metadata.
+      if (response.has_metadata()) {
+        continue;
+      }
+
+      final_status_ = google::cloud::Status(
+          google::cloud::StatusCode::kInternal,
+          "Empty ExecuteQueryResponse received from stream");
+      return false;
+    }
+  }
+
+  Status Finish() override { return final_status_; }
+
+ private:
+  std::shared_ptr<grpc::ClientContext> context_;
+  std::unique_ptr<
+      internal::StreamingReadRpc<google::bigtable::v2::ExecuteQueryResponse>>
+      reader_;
+  Status final_status_;
+};
+
+class StatusOnlyResultSetSource : public bigtable::ResultSourceInterface {
+ public:
+  explicit StatusOnlyResultSetSource(google::cloud::Status status)
+      : status_(std::move(status)) {}
+  ~StatusOnlyResultSetSource() override = default;
+
+  StatusOr<bigtable::QueryRow> NextRow() override { return status_; }
+  absl::optional<google::bigtable::v2::ResultSetMetadata> Metadata() override {
+    return {};
+  }
+
+ private:
+  google::cloud::Status status_;
+};
+
+template <typename ResultType>
+ResultType MakeStatusOnlyResult(Status status) {
+  return ResultType(
+      std::make_unique<StatusOnlyResultSetSource>(std::move(status)));
+}
+
 }  // namespace
 
 bigtable::Row TransformReadModifyWriteRowResponse(
@@ -737,9 +819,60 @@ future<StatusOr<bigtable::PreparedQuery>> DataConnectionImpl::AsyncPrepareQuery(
       });
 }
 
+std::unique_ptr<PartialResultSetReader>
+DataConnectionImpl::CreateResumableReader(
+    google::bigtable::v2::ExecuteQueryRequest request,
+    std::string const& resume_token) {
+  if (!resume_token.empty()) {
+    request.set_resume_token(resume_token);
+  }
+  auto context = std::make_shared<grpc::ClientContext>();
+  auto const& options = google::cloud::internal::CurrentOptions();
+
+  google::cloud::internal::ConfigureContext(*context, options);
+  auto stream = stub_->ExecuteQuery(context, options, request);
+
+  return std::make_unique<DefaultPartialResultSetReader>(std::move(context),
+                                                         std::move(stream));
+}
+
 StatusOr<bigtable::RowStream> DataConnectionImpl::ExecuteQuery(
-    bigtable::ExecuteQueryParams const&) {
-  return Status(StatusCode::kUnimplemented, "not implemented");
+    bigtable::ExecuteQueryParams const& params) {
+  auto current = google::cloud::internal::SaveCurrentOptions();
+  StatusOr<google::bigtable::v2::ResultSetMetadata> status_or_metadata =
+      params.bound_query.metadata();
+  google::bigtable::v2::ExecuteQueryRequest request =
+      params.bound_query.ToRequestProto();
+  request.set_app_profile_id(app_profile_id(*current));
+  if (!status_or_metadata) {
+    return status_or_metadata.status();
+  }
+  google::bigtable::v2::ResultSetMetadata metadata =
+      *std::move(status_or_metadata);
+
+  auto retry_resume_fn =
+      [this, retry_policy_prototype = retry_policy(*current),
+       backoff_policy_prototype = backoff_policy(*current)](
+          google::bigtable::v2::ResultSetMetadata metadata,
+          google::bigtable::v2::ExecuteQueryRequest const& initial_request)
+      -> StatusOr<std::unique_ptr<bigtable::ResultSourceInterface>> {
+    auto factory = absl::bind_front(&DataConnectionImpl::CreateResumableReader,
+                                    this, initial_request);
+
+    auto resume_reader = std::make_unique<PartialResultSetResume>(
+        std::move(factory), Idempotency::kIdempotent,
+        retry_policy_prototype->clone(), backoff_policy_prototype->clone());
+
+    return PartialResultSetSource::Create(std::move(metadata),
+                                          std::move(resume_reader));
+  };
+
+  auto response = retry_resume_fn(metadata, request);
+  if (!response) {
+    auto status = std::move(response).status();
+    return MakeStatusOnlyResult<bigtable::RowStream>(std::move(status));
+  }
+  return bigtable::RowStream(*std::move(response));
 };
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

google/cloud/bigtable/internal/data_connection_impl.h

@@ -19,6 +19,7 @@
 #include "google/cloud/bigtable/internal/bigtable_stub.h"
 #include "google/cloud/bigtable/internal/mutate_rows_limiter.h"
 #include "google/cloud/bigtable/internal/operation_context_factory.h"
+#include "google/cloud/bigtable/internal/partial_result_set_reader.h"
 #include "google/cloud/bigtable/prepared_query.h"
 #include "google/cloud/bigtable/results.h"
 #include "google/cloud/background_threads.h"
@@ -118,6 +119,9 @@ class DataConnectionImpl : public bigtable::DataConnection {
                            bigtable::Filter filter,
                            internal::ImmutableOptions const& current,
                            std::shared_ptr<OperationContext> operation_context);
+  std::unique_ptr<PartialResultSetReader> CreateResumableReader(
+      google::bigtable::v2::ExecuteQueryRequest request,
+      std::string const& resume_token);
 
   std::unique_ptr<BackgroundThreads> background_;
   std::shared_ptr<BigtableStub> stub_;