feat(bigtable): add expiration logic to query plan implementation (#15647)

* feat(bigtable): add expiration logic to query plan implementation

Author: mpeddada1

google/cloud/bigtable/internal/query_plan.cc

@@ -21,11 +21,30 @@ namespace cloud {
 namespace bigtable_internal {
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
-std::string const& QueryPlan::prepared_query() const {
+std::shared_ptr<QueryPlan> QueryPlan::Create(
+    CompletionQueue cq, google::bigtable::v2::PrepareQueryResponse response,
+    RefreshFn fn) {
+  auto plan = std::shared_ptr<QueryPlan>(
+      new QueryPlan(std::move(cq), std::move(response), std::move(fn)));
+  plan->Initialize();
+  return plan;
+}
+
+bool QueryPlan::IsExpired() { return false; }
+
+StatusOr<std::string> QueryPlan::prepared_query() const {
+  std::lock_guard<std::mutex> lock(mu_);
+  if (IsExpired()) {
+    return Status(StatusCode::kUnavailable, "Query plan has expired");
+  }
   return response_.prepared_query();
 }
 
-google::bigtable::v2::ResultSetMetadata const& QueryPlan::metadata() const {
+StatusOr<google::bigtable::v2::ResultSetMetadata> QueryPlan::metadata() const {
+  std::lock_guard<std::mutex> lock(mu_);
+  if (IsExpired()) {
+    return Status(StatusCode::kUnavailable, "Query plan has expired");
+  }
   return response_.metadata();
 }
 

google/cloud/bigtable/internal/query_plan.h

@@ -36,23 +36,21 @@ class QueryPlan : public std::enable_shared_from_this<QueryPlan> {
   // Calls the constructor and then Initialize.
   static std::shared_ptr<QueryPlan> Create(
       CompletionQueue cq, google::bigtable::v2::PrepareQueryResponse response,
-      RefreshFn fn) {
-    return std::shared_ptr<QueryPlan>(
-        new QueryPlan(std::move(cq), std::move(response), std::move(fn)));
-  }
+      RefreshFn fn);
 
   // Accessor for the prepared_query field in response_.
-  std::string const& prepared_query() const;
+  StatusOr<std::string> prepared_query() const;
 
   // Accessor for the metadata field in  response_.
-  google::bigtable::v2::ResultSetMetadata const& metadata() const;
+  StatusOr<google::bigtable::v2::ResultSetMetadata> metadata() const;
 
  private:
   QueryPlan(CompletionQueue cq,
             google::bigtable::v2::PrepareQueryResponse response, RefreshFn fn)
       : cq_(std::move(cq)),
         response_(std::move(response)),
         fn_(std::move(fn)) {}
+  static bool IsExpired();
 
   // Performs the first call to ScheduleRefresh and any other initialization not
   // possible in the constructor.

google/cloud/bigtable/internal/query_plan_test.cc

@@ -13,25 +13,55 @@
 // limitations under the License.
 
 #include "google/cloud/bigtable/internal/query_plan.h"
-#include "google/cloud/completion_queue.h"
+#include "google/cloud/testing_util/is_proto_equal.h"
+#include "google/cloud/testing_util/mock_completion_queue_impl.h"
+#include "google/cloud/testing_util/status_matchers.h"
 #include <google/bigtable/v2/data.pb.h>
+#include <google/protobuf/text_format.h>
+#include <google/protobuf/util/time_util.h>
 #include <gmock/gmock.h>
+#include <gtest/gtest.h>
 
 namespace google {
 namespace cloud {
 namespace bigtable_internal {
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 namespace {
 
-/// @test Basic test to confirm compilation and dummy behavior.
-/// Remove this after implementing QueryPlan.
-TEST(QueryPlanTest, TestDummyMethods) {
-  auto cq = CompletionQueue();
-  auto response = google::bigtable::v2::PrepareQueryResponse();
-  auto fn = QueryPlan::RefreshFn();
-  auto created = QueryPlan::Create(cq, response, fn);
-  ASSERT_EQ("", (*created).prepared_query());
-  ASSERT_EQ(0, created->metadata().ByteSizeLong());
+using ::google::cloud::testing_util::IsProtoEqual;
+using ::google::cloud::testing_util::MockCompletionQueueImpl;
+using ::google::protobuf::util::TimeUtil;
+
+TEST(QueryPlanTest, Accessors) {
+  auto mock_cq = std::make_shared<MockCompletionQueueImpl>();
+  CompletionQueue cq(mock_cq);
+  google::bigtable::v2::PrepareQueryResponse response;
+  response.set_prepared_query("test-query");
+  auto constexpr kResultMetadataText = R"pb(
+    proto_schema {
+      columns {
+        name: "user_id"
+        type { string_type {} }
+      }
+    }
+  )pb";
+  google::bigtable::v2::ResultSetMetadata metadata;
+  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(kResultMetadataText,
+                                                            &metadata));
+  *response.mutable_metadata() = metadata;
+  *response.mutable_valid_until() =
+      TimeUtil::GetCurrentTime() + TimeUtil::SecondsToDuration(300);
+
+  auto plan = QueryPlan::Create(cq, response, [] {
+    return google::bigtable::v2::PrepareQueryResponse{};
+  });
+
+  auto prepared_query = plan->prepared_query();
+  ASSERT_STATUS_OK(prepared_query);
+  EXPECT_EQ(*prepared_query, "test-query");
+  auto actual_metadata = plan->metadata();
+  ASSERT_STATUS_OK(actual_metadata);
+  EXPECT_THAT(*actual_metadata, IsProtoEqual(metadata));
 }
 
 }  // namespace