Merge branch 'main' into bigtable_query_snippet_2

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/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

google/cloud/opentelemetry/internal/monitoring_exporter.cc

@@ -31,6 +31,41 @@ std::string FormatProjectFullName(std::string const& project) {
   return absl::StrCat("projects/", project);
 }
 
+otel_internal::ResourceFilterDataFn MakeResourceFilterFn(
+    Options const& options) {
+  if (!options.has<otel_internal::ResourceFilterDataFnOption>()) {
+    return nullptr;
+  }
+
+  // Get the metric labels set to be excluded.
+  auto const& excluded =
+      options.get<otel_internal::ResourceFilterDataFnOption>();
+  if (excluded.empty()) return nullptr;
+
+  // Capture by value to avoid dangling reference in the lambda.
+  return [excluded = std::move(excluded)](std::string const& key) -> bool {
+    return excluded.count(key) > 0;
+  };
+}
+
+otel_internal::MonitoredResourceFromDataFn MakeDynamicResourceFn(
+    Options const& options, absl::optional<Project> const& project,
+    absl::optional<google::api::MonitoredResource> const& mr_proto) {
+  if (!options.has<otel_internal::ResourceFilterDataFnOption>()) {
+    return nullptr;
+  }
+
+  // `resource_filter_fn_` and `dynamic_resource_fn_` are meant to be used as a
+  // pair. Here we have a filter but no dynamic function, create a default one
+  // that returns the same project and monitored resource for all data points.
+  auto project_id = project->project_id();
+  auto monitored_resource = mr_proto.value_or(google::api::MonitoredResource{});
+  return [project_id, monitored_resource](
+             opentelemetry::sdk::metrics::PointDataAttributes const&) {
+    return std::make_pair(project_id, monitored_resource);
+  };
+}
+
 }  // namespace
 
 MonitoringExporter::MonitoringExporter(
@@ -51,6 +86,8 @@ MonitoringExporter::MonitoringExporter(
     Options const& options)
     : MonitoringExporter(std::move(conn), nullptr, nullptr, options) {
   project_ = std::move(project);
+  resource_filter_fn_ = MakeResourceFilterFn(options);
+  dynamic_resource_fn_ = MakeDynamicResourceFn(options, project_, mr_proto_);
 }
 
 opentelemetry::sdk::common::ExportResult MonitoringExporter::Export(

google/cloud/opentelemetry/internal/monitoring_exporter.h

@@ -45,6 +45,12 @@ using MonitoredResourceFromDataFn =
 // of the google::api::Metric proto.
 using ResourceFilterDataFn = std::function<bool(std::string const&)>;
 
+// Filter resource labels. A set of OpenTelemetry resource attribute keys to
+// exclude from metric labels when exporting metrics.
+struct ResourceFilterDataFnOption {
+  using Type = std::set<std::string>;
+};
+
 class MonitoringExporter final
     : public opentelemetry::sdk::metrics::PushMetricExporter {
  public:

google/cloud/opentelemetry/internal/monitoring_exporter_test.cc

@@ -159,6 +159,40 @@ TEST(MonitoringExporter, ExportSuccess) {
   EXPECT_EQ(result, opentelemetry::sdk::common::ExportResult::kSuccess);
 }
 
+TEST(MonitoringExporterTest, MakeFilterNoOption) {
+  auto mock =
+      std::make_shared<monitoring_v3_mocks::MockMetricServiceConnection>();
+  Options options;
+
+  auto exporter = std::make_unique<MonitoringExporter>(Project("test-project"),
+                                                       mock, options);
+  EXPECT_NE(exporter, nullptr);
+}
+
+TEST(MonitoringExporterTest, MakeFilterEmptySet) {
+  auto mock =
+      std::make_shared<monitoring_v3_mocks::MockMetricServiceConnection>();
+  Options options;
+  options.set<otel_internal::ResourceFilterDataFnOption>(
+      std::set<std::string>{});
+
+  auto exporter = std::make_unique<MonitoringExporter>(Project("test-project"),
+                                                       mock, options);
+  EXPECT_NE(exporter, nullptr);
+}
+
+TEST(MonitoringExporterTest, MakeFilterWithExcludedKeys) {
+  auto mock =
+      std::make_shared<monitoring_v3_mocks::MockMetricServiceConnection>();
+  Options options;
+  std::set<std::string> excluded{"service_name", "service_version"};
+  options.set<otel_internal::ResourceFilterDataFnOption>(excluded);
+
+  auto exporter = std::make_unique<MonitoringExporter>(Project("test-project"),
+                                                       mock, options);
+  EXPECT_NE(exporter, nullptr);
+}
+
 }  // namespace
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace otel_internal

google/cloud/storage/grpc_plugin.h

@@ -131,6 +131,27 @@ struct GrpcMetricsExportTimeoutOption {
   using Type = std::chrono::seconds;
 };
 
+/**
+ * gRPC telemetry excluded labels.
+ *
+ * A set of OpenTelemetry resource attribute keys to exclude from metric labels
+ * when exporting gRPC telemetry. For example, to exclude the `service.name`
+ * label, configure the option with `{"service_name"}`.
+ *
+ * @par Example: Exclude specific labels from telemetry
+ * @code
+ * namespace gcs_ex = google::cloud::storage_experimental;
+ * auto client = google::cloud::storage::MakeGrpcClient(
+ *     google::cloud::Options{}
+ *         .set<gcs_ex::EnableGrpcMetricsOption>(true)
+ *         .set<gcs_ex::GrpcMetricsExcludedLabelsOption>(
+ *             std::set<std::string>{"service_name", "service_version"}));
+ * @endcode
+ */
+struct GrpcMetricsExcludedLabelsOption {
+  using Type = std::set<std::string>;
+};
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace storage_experimental
 }  // namespace cloud