fix invalid type

Author: mpeddada1

google/cloud/bigtable/internal/data_connection_impl.cc

@@ -757,6 +757,13 @@ StatusOr<bigtable::PreparedQuery> DataConnectionImpl::PrepareQuery(
   if (!response) {
     return std::move(response).status();
   }
+  for (auto const& column : response->metadata().proto_schema().columns()) {
+    if (!column.has_type() ||
+        column.type().kind_case() == google::bigtable::v2::Type::KIND_NOT_SET) {
+      return internal::InternalError("Column type cannot be empty",
+                                     GCP_ERROR_INFO());
+    }
+  }
   auto const* func = __func__;
   auto refresh_fn = [this, request, func]() mutable {
     auto current = google::cloud::internal::SaveCurrentOptions();
@@ -840,7 +847,15 @@ future<StatusOr<bigtable::PreparedQuery>> DataConnectionImpl::AsyncPrepareQuery(
         if (!response) {
           return std::move(response).status();
         }
-
+        for (auto const& column :
+             response->metadata().proto_schema().columns()) {
+          if (!column.has_type() ||
+              column.type().kind_case() ==
+                  google::bigtable::v2::Type::KIND_NOT_SET) {
+            return StatusOr<bigtable::PreparedQuery>(internal::InternalError(
+                "Column type cannot be empty", GCP_ERROR_INFO()));
+          }
+        }
         auto refresh_fn = [this, request, func]() mutable {
           auto current = google::cloud::internal::SaveCurrentOptions();
           auto retry = retry_policy(*current);
@@ -870,6 +885,9 @@ future<StatusOr<bigtable::PreparedQuery>> DataConnectionImpl::AsyncPrepareQuery(
               .then([operation_context](auto f) mutable {
                 StatusOr<google::bigtable::v2::PrepareQueryResponse> response =
                     f.get();
+                if (!response.ok()) {
+                  return response;
+                }
                 operation_context->OnDone(response.status());
                 return response;
               });

google/cloud/bigtable/internal/data_connection_impl_test.cc

@@ -3757,6 +3757,96 @@ TEST_F(DataConnectionTest, ExecuteQueryFailureWithSchemaChange) {
   fake_cq_impl->SimulateCompletion(false);
 }
 
+TEST_F(DataConnectionTest, PrepareQueryFailsOnInvalidType) {
+  auto mock = std::make_shared<MockBigtableStub>();
+  auto fake_cq_impl = std::make_shared<FakeCompletionQueueImpl>();
+  auto mock_bg = std::make_unique<MockBackgroundThreads>();
+  EXPECT_CALL(*mock_bg, cq).WillRepeatedly([&]() {
+    return CompletionQueue{fake_cq_impl};
+  });
+
+  v2::PrepareQueryResponse pq_response;
+  pq_response.set_prepared_query("test-pq-id-54321");
+  auto constexpr kResultMetadataText = R"pb(
+    proto_schema {
+      columns {
+        name: "test"
+        type { string_type {} }
+      }
+      columns {
+        name: "invalid"
+        type {}
+      }
+    }
+  )pb";
+  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(
+      kResultMetadataText, pq_response.mutable_metadata()));
+  *pq_response.mutable_valid_until() = internal::ToProtoTimestamp(
+      std::chrono::system_clock::now() + std::chrono::seconds(3600));
+
+  EXPECT_CALL(*mock, PrepareQuery)
+      .WillOnce([&](grpc::ClientContext&, Options const&,
+                    v2::PrepareQueryRequest const&) { return pq_response; });
+
+  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 prepared_query = conn->PrepareQuery(params);
+  EXPECT_THAT(prepared_query,
+              StatusIs(StatusCode::kInternal, "Column type cannot be empty"));
+
+  fake_cq_impl->SimulateCompletion(false);
+}
+
+TEST_F(DataConnectionTest, AsyncPrepareQueryFailsOnInvalidType) {
+  auto mock = std::make_shared<MockBigtableStub>();
+  auto fake_cq_impl = std::make_shared<FakeCompletionQueueImpl>();
+  auto mock_bg = std::make_unique<MockBackgroundThreads>();
+  EXPECT_CALL(*mock_bg, cq).WillRepeatedly([&]() {
+    return CompletionQueue{fake_cq_impl};
+  });
+
+  v2::PrepareQueryResponse pq_response;
+  pq_response.set_prepared_query("test-pq-id-54321");
+  auto constexpr kResultMetadataText = R"pb(
+    proto_schema {
+      columns {
+        name: "test"
+        type { string_type {} }
+      }
+      columns {
+        name: "invalid"
+        type {}
+      }
+    }
+  )pb";
+  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(
+      kResultMetadataText, pq_response.mutable_metadata()));
+  *pq_response.mutable_valid_until() = internal::ToProtoTimestamp(
+      std::chrono::system_clock::now() + std::chrono::seconds(3600));
+
+  EXPECT_CALL(*mock, AsyncPrepareQuery)
+      .WillOnce([&](CompletionQueue const&, auto, auto,
+                    v2::PrepareQueryRequest const&) {
+        return make_ready_future(make_status_or(pq_response));
+      });
+
+  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 prepared_query = conn->AsyncPrepareQuery(params).get();
+  EXPECT_THAT(prepared_query,
+              StatusIs(StatusCode::kInternal, "Column type cannot be empty"));
+
+  fake_cq_impl->SimulateCompletion(false);
+}
+
 }  // namespace
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace bigtable_internal