Merge branch 'main' into gemini_styleguide_updates

Author: scotthart

ci/cloudbuild/dockerfiles/checkers.Dockerfile

@@ -23,7 +23,6 @@ ARG ARCH=amd64
 
 RUN dnf makecache && \
     dnf install -y \
-        cargo \
         cmake \
         clang-tools-extra \
         diffutils \
@@ -35,7 +34,11 @@ RUN dnf makecache && \
         python-pip \
         ShellCheck
 
-RUN cargo install typos-cli --version 1.24.1 --root /usr/local
+RUN dnf makecache && \
+    dnf install -y \
+    cargo
+
+RUN cargo install typos-cli --locked --version 1.24.1 --root /usr/local
 
 RUN curl -L -o /usr/bin/buildifier https://github.com/bazelbuild/buildtools/releases/download/v6.4.0/buildifier-linux-amd64 && \
     chmod 755 /usr/bin/buildifier

google/cloud/bigtable/client.h

@@ -49,36 +49,7 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
  * For this reason, it is recommended to reuse `Client` objects when possible.
  *
  * @par Example
- * @code
- * #include "google/cloud/bigtable/client.h"
- * #include "google/cloud/project.h"
- * #include <iostream>
- *
- * int main() {
- *   namespace cbt = google::cloud::bigtable;
- *   cbt::Client client(cbt::MakeDataConnection());
- *   cbt::InstanceResource instance(google::cloud::Project("my-project"),
- * "my-instance");
- *
- *   // Declare a parameter with a type, but no value.
- *   cbt::SqlStatement statement(
- *       "SELECT _key, CAST(family['qual'] AS STRING) AS value "
- *       "FROM my-table WHERE _key = @key",
- *       {{"key", cbt::Value(cbt::Bytes())}});
- *
- *   google::cloud::StatusOr<cbt::PreparedQuery> prepared_query =
- *       client.PrepareQuery(instance, statement);
- *   if (!prepared_query) throw std::move(prepared_query).status();
- *
- *   auto bound_query = prepared_query->BindParameters(
- *       {{"key", cbt::Value("row-key-2")}});
- *
- *   RowStream results =
- *       client.ExecuteQuery(std::move(bound_query));
- *
- *   ... // process rows
- * }
- * @endcode
+ * @snippet data_snippets.cc prepare-and-execute-query
  */
 class Client {
  public:

google/cloud/bigtable/examples/bigtable_examples_common.cc

@@ -129,6 +129,24 @@ Commands::value_type MakeCommandEntry(std::string const& name,
   return Commands::value_type{name, std::move(adapter)};
 }
 
+Commands::value_type MakeCommandEntry(std::string const& name,
+                                      std::vector<std::string> const& args,
+                                      ClientCommandType const& function) {
+  auto command = [=](std::vector<std::string> argv) {
+    auto constexpr kFixedArguments = 0;
+    if ((argv.size() == 1 && argv[0] == "--help") ||
+        argv.size() != args.size() + kFixedArguments) {
+      std::ostringstream os;
+      os << name;
+      if (!args.empty()) os << " " << absl::StrJoin(args, " ");
+      throw Usage{std::move(os).str()};
+    }
+    auto client = bigtable::Client(bigtable::MakeDataConnection());
+    function(client, argv);
+  };
+  return {name, command};
+}
+
 }  // namespace examples
 }  // namespace bigtable
 }  // namespace cloud

google/cloud/bigtable/examples/bigtable_examples_common.h

@@ -17,6 +17,7 @@
 
 #include "google/cloud/bigtable/admin/bigtable_instance_admin_client.h"
 #include "google/cloud/bigtable/admin/bigtable_table_admin_client.h"
+#include "google/cloud/bigtable/client.h"
 #include "google/cloud/bigtable/table.h"
 #include "google/cloud/internal/random.h"
 #include "google/cloud/testing_util/example_driver.h"
@@ -85,6 +86,13 @@ Commands::value_type MakeCommandEntry(std::string const& name,
                                       std::vector<std::string> const& args,
                                       TableAsyncCommandType const& command);
 
+using ClientCommandType = std::function<void(google::cloud::bigtable::Client,
+                                             std::vector<std::string>)>;
+
+Commands::value_type MakeCommandEntry(std::string const& name,
+                                      std::vector<std::string> const& args,
+                                      ClientCommandType const& function);
+
 }  // namespace examples
 }  // namespace bigtable
 }  // namespace cloud

google/cloud/bigtable/examples/data_snippets.cc

@@ -18,6 +18,7 @@
 #include "google/cloud/bigtable/testing/cleanup_stale_resources.h"
 #include "google/cloud/bigtable/testing/random_names.h"
 //! [bigtable includes]
+#include "google/cloud/bigtable/client.h"
 #include "google/cloud/bigtable/table.h"
 //! [bigtable includes]
 #include "google/cloud/internal/getenv.h"
@@ -778,6 +779,38 @@ void RunWriteExamples(
   admin.DeleteTable(schema->name());
 }
 
+void PrepareAndExecuteQuery(google::cloud::bigtable::Client client,
+                            std::vector<std::string> const& args) {
+  // [prepare-and-execute-query] [START bigtable_api_execute_query]
+  namespace cbt = ::google::cloud::bigtable;
+  [](cbt::Client client, std::string const& project_id,
+     std::string const& instance_id, std::string const& table_id) {
+    cbt::InstanceResource instance(google::cloud::Project(project_id),
+                                   instance_id);
+    cbt::SqlStatement statement(
+        "SELECT _key, CAST(fam['column0'] AS STRING) as value0 FROM `" +
+            table_id + "` WHERE _key=@key",
+        {{"key", cbt::Parameter(cbt::Bytes())}});
+
+    auto prepared_query = client.PrepareQuery(instance, statement);
+    if (!prepared_query) throw std::move(prepared_query).status();
+
+    auto bound_query = prepared_query->BindParameters(
+        {{"key", cbt::Value(cbt::Bytes("test-key-for-apply"))}});
+
+    auto results = client.ExecuteQuery(std::move(bound_query));
+
+    using RowType = std::tuple<cbt::Bytes, absl::optional<std::string>>;
+    for (auto& row : cbt::StreamOf<RowType>(results)) {
+      if (!row.ok()) throw std::move(row.status());
+      auto v = std::get<1>(*row);
+      std::cout << std::get<0>(*row) << "; " << (v ? *v : "null") << std::endl;
+    }
+  }
+  // [prepare-and-execute-query] [END bigtable_api_execute_query]
+  (std::move(client), args.at(0), args.at(1), args.at(2));
+}
+
 void RunDataExamples(
     google::cloud::bigtable_admin::BigtableTableAdminClient admin,
     google::cloud::internal::DefaultPRNG& generator,
@@ -879,6 +912,11 @@ void RunDataExamples(
   std::cout << "Running ReadModifyWrite() example [3]" << std::endl;
   ReadModifyWrite(table, {"read-modify-write"});
 
+  if (!google::cloud::bigtable::examples::UsingEmulator()) {
+    auto client = cbt::Client(cbt::MakeDataConnection());
+    std::cout << "Running PrepareAndExecuteQuery() example" << std::endl;
+    PrepareAndExecuteQuery(client, {project_id, instance_id, table_id});
+  }
   admin.DeleteTable(schema->name());
 }
 
@@ -946,6 +984,9 @@ int main(int argc, char* argv[]) try {
       MakeCommandEntry("write-batch", {}, WriteBatch),
       MakeCommandEntry("write-increment", {}, WriteIncrement),
       MakeCommandEntry("write-conditional", {}, WriteConditionally),
+      MakeCommandEntry("prepare-and-execute-query",
+                       {"<project_id>", "<instance_id>", "<table_id>"},
+                       PrepareAndExecuteQuery),
       {"auto", RunAll},
   };
 

google/cloud/bigtable/internal/data_connection_impl.cc

@@ -119,10 +119,12 @@ class DefaultPartialResultSetReader
       std::shared_ptr<OperationContext> operation_context,
       std::unique_ptr<internal::StreamingReadRpc<
           google::bigtable::v2::ExecuteQueryResponse>>
-          reader)
+          reader,
+      google::bigtable::v2::ResultSetMetadata initial_metadata)
       : context_(std::move(context)),
         operation_context_(std::move(operation_context)),
-        reader_(std::move(reader)) {}
+        reader_(std::move(reader)),
+        initial_metadata_(std::move(initial_metadata)) {}
 
   ~DefaultPartialResultSetReader() override = default;
 
@@ -148,10 +150,22 @@ class DefaultPartialResultSetReader
         return true;
       }
 
-      // Ignore metadata from the stream because PartialResultSetSource already
-      // has it set (in ExecuteQuery).
-      // TODO(#15701): Investigate expected behavior for processing metadata.
+      // Throw an error when there is a schema difference between
+      // ExecuteQueryResponse and PrepareQueryResponse.
       if (response.has_metadata()) {
+        std::string initial_metadata_str;
+        std::string response_metadata_str;
+        bool metadata_matched =
+            initial_metadata_.SerializeToString(&initial_metadata_str) &&
+            response.metadata().SerializeToString(&response_metadata_str) &&
+            initial_metadata_str == response_metadata_str;
+        if (response.metadata().ByteSizeLong() > 0 && !metadata_matched) {
+          final_status_ = google::cloud::Status(
+              google::cloud::StatusCode::kAborted,
+              "Schema changed during ExecuteQuery operation");
+          operation_context_->PostCall(*context_, final_status_);
+          return false;
+        }
         continue;
       }
 
@@ -173,6 +187,7 @@ class DefaultPartialResultSetReader
   std::unique_ptr<
       internal::StreamingReadRpc<google::bigtable::v2::ExecuteQueryResponse>>
       reader_;
+  google::bigtable::v2::ResultSetMetadata initial_metadata_;
   Status final_status_;
 };
 
@@ -896,8 +911,8 @@ bigtable::RowStream DataConnectionImpl::ExecuteQuery(
           std::shared_ptr<OperationContext> const& operation_context) mutable
       -> StatusOr<std::unique_ptr<bigtable::ResultSourceInterface>> {
     auto factory =
-        [stub, request, tracing_enabled, tracing_options,
-         operation_context](std::string const& resume_token) mutable {
+        [stub, request, tracing_enabled, tracing_options, operation_context,
+         initial_metadata = metadata](std::string const& resume_token) mutable {
           if (!resume_token.empty()) request.set_resume_token(resume_token);
           auto context = std::make_shared<grpc::ClientContext>();
           auto const& options = internal::CurrentOptions();
@@ -906,7 +921,8 @@ bigtable::RowStream DataConnectionImpl::ExecuteQuery(
           auto stream = stub->ExecuteQuery(context, options, request);
           std::unique_ptr<PartialResultSetReader> reader =
               std::make_unique<DefaultPartialResultSetReader>(
-                  std::move(context), operation_context, std::move(stream));
+                  std::move(context), operation_context, std::move(stream),
+                  std::move(initial_metadata));
           if (tracing_enabled) {
             reader = std::make_unique<LoggingResultSetReader>(std::move(reader),
                                                               tracing_options);

google/cloud/bigtable/internal/data_connection_impl_test.cc

@@ -77,6 +77,7 @@ using ::testing::Contains;
 using ::testing::ElementsAre;
 using ::testing::ElementsAreArray;
 using ::testing::Eq;
+using ::testing::HasSubstr;
 using ::testing::Matcher;
 using ::testing::MockFunction;
 using ::testing::Pair;
@@ -3695,6 +3696,92 @@ TEST_F(DataConnectionTest,
   EXPECT_THAT(row2.values().at(1).get<std::string>(), IsOkAndHolds("v2"));
 }
 
+TEST_F(DataConnectionTest, ExecuteQueryFailureWithSchemaChange) {
+  auto factory = std::make_unique<SimpleOperationContextFactory>();
+
+  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};
+  });
+  auto constexpr kResultMetadataText = R"pb(
+    proto_schema {
+      columns {
+        name: "row_key"
+        type { string_type {} }
+      }
+      columns {
+        name: "value"
+        type { string_type {} }
+      }
+    }
+  )pb";
+  v2::PrepareQueryResponse pq_response;
+  pq_response.set_prepared_query("test-pq-id-54321");
+  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));
+
+  auto constexpr kExecuteQueryResultMetadataText = R"pb(
+    proto_schema {
+      columns {
+        name: "row_key"
+        type { string_type {} }
+      }
+      columns {
+        name: "different_value"
+        type { string_type {} }
+      }
+    }
+  )pb";
+  v2::ExecuteQueryResponse eq_response;
+  ASSERT_TRUE(google::protobuf::TextFormat::ParseFromString(
+      kExecuteQueryResultMetadataText, eq_response.mutable_metadata()));
+
+  auto refresh_fn = []() {
+    return make_ready_future(
+        StatusOr<google::bigtable::v2::PrepareQueryResponse>(
+            Status{StatusCode::kUnimplemented, "not implemented"}));
+  };
+  EXPECT_CALL(*mock, ExecuteQuery)
+      .Times(3)
+      .WillRepeatedly(
+          [&](auto, auto const&,
+              google::bigtable::v2::ExecuteQueryRequest const& request) {
+            EXPECT_EQ(request.app_profile_id(), kAppProfile);
+            EXPECT_EQ(request.instance_name(),
+                      "projects/test-project/instances/test-instance");
+            auto stream = std::make_unique<MockExecuteQueryStream>();
+            EXPECT_CALL(*stream, Read)
+                .WillOnce([&](google::bigtable::v2::ExecuteQueryResponse* r) {
+                  *r = eq_response;
+                  return absl::nullopt;
+                });
+            return stream;
+          });
+
+  auto conn = TestConnection(std::move(mock), std::move(factory));
+  internal::OptionsSpan span(CallOptions());
+  Project p("test-project");
+  bigtable::SqlStatement statement("SELECT * FROM the-table");
+  bigtable::InstanceResource instance(p, "test-instance");
+  auto query_plan =
+      QueryPlan::Create(CompletionQueue(fake_cq_impl), std::move(pq_response),
+                        std::move(refresh_fn));
+  auto prepared_query =
+      bigtable::PreparedQuery(instance, statement, std::move(query_plan));
+  auto bq = prepared_query.BindParameters({});
+  bigtable::ExecuteQueryParams params{std::move(bq)};
+  auto row_stream = conn->ExecuteQuery(std::move(params));
+  for (auto const& row : row_stream) {
+    EXPECT_THAT(row,
+                StatusIs(StatusCode::kAborted, HasSubstr("Schema changed")));
+  }
+  fake_cq_impl->SimulateCompletion(false);
+}
+
 }  // namespace
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace bigtable_internal

google/cloud/bigtable/internal/query_plan.cc

@@ -35,6 +35,10 @@ std::shared_ptr<QueryPlan> QueryPlan::Create(
   return plan;
 }
 
+QueryPlan::~QueryPlan() {
+  if (refresh_timer_.valid()) refresh_timer_.cancel();
+}
+
 void QueryPlan::Initialize() {
   std::unique_lock<std::mutex> lock(mu_);
   if (state_ == RefreshState::kDone) ScheduleRefresh(lock);

google/cloud/bigtable/internal/query_plan.h

@@ -43,6 +43,8 @@ class QueryPlan : public std::enable_shared_from_this<QueryPlan> {
       StatusOr<google::bigtable::v2::PrepareQueryResponse> response,
       RefreshFn fn, std::shared_ptr<Clock> clock = std::make_shared<Clock>());
 
+  ~QueryPlan();
+
   // Invalidates the current QueryPlan and triggers a refresh.
   void Invalidate(Status status, std::string const& invalid_query_plan_id);
 

google/cloud/bigtable/internal/query_plan_test.cc

@@ -229,10 +229,11 @@ TEST(QueryPlanTest, CreateFailedPlanAndRefresh) {
 }
 
 // TODO(#15695): For reasons not yet understood, the fedora m32 CI build has
-//  failures not seen in m64 builds when the number of threads is "too" high.
+//   failures not seen in m64 builds when the number of threads is "too" high.
+//  These failures occur while trying to create more than 500 threads.
 constexpr int LimitNumThreadsOn32Bit(int num_threads) {
 #if INTPTR_MAX == INT32_MAX
-  return std::min(num_threads, 500);
+  return std::min(num_threads, 200);
 #else
   return num_threads;
 #endif