impl: Add TracingConnection decorator

.bazelrc

@@ -55,7 +55,7 @@ common:macos --repo_env=BAZEL_NO_APPLE_CPP_TOOLCHAIN=1
 test --test_env=GTEST_SHUFFLE --test_env=GTEST_RANDOM_SEED
 
 # By default, build the library with OpenTelemetry
-#build --//:enable_opentelemetry
+# build --@com_google_googleapis_google_cloud_cpp//:enable_opentelemetry
 
 # Don't show warnings when building external dependencies. This still shows
 # warnings when using these dependencies (say in headers).

CMakeLists.txt

@@ -26,6 +26,10 @@ project(
     LANGUAGES CXX)
 set(PROJECT_VERSION_PRE_RELEASE "rc")
 
+set(GOOGLE_CLOUD_CPP_BIGQUERY_ENABLE_OPENTELEMETRY
+    OFF
+    CACHE BOOL "Enable OpenTelemetry in google-cloud-cpp-bigquery.")
+
 list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake)
 
 include(GoogleCloudCppFeatures)

ci/cloudbuild/builds/cmake-install.sh

@@ -39,7 +39,8 @@ io::run cmake "${cmake_args[@]}" \
   -DGOOGLE_CLOUD_CPP_ENABLE_CLANG_ABI_COMPAT_17=ON \
   -DBUILD_TESTING=OFF \
   -DGOOGLE_CLOUD_CPP_ENABLE_EXAMPLES=OFF \
-  -DGOOGLE_CLOUD_CPP_ENABLE="${ENABLED_FEATURES}"
+  -DGOOGLE_CLOUD_CPP_ENABLE="${ENABLED_FEATURES}" \
+  -DGOOGLE_CLOUD_CPP_BIGQUERY_ENABLE_OPENTELEMETRY=ON
 io::run cmake --build cmake-out
 io::run cmake --install cmake-out --component google_cloud_cpp_bigquery_development
 

ci/cloudbuild/dockerfiles/fedora-latest-cmake.Dockerfile

@@ -198,7 +198,7 @@ RUN curl -fsSL https://github.com/googleapis/google-cloud-cpp/archive/v2.33.0.ta
         -DCMAKE_CXX_STANDARD=17 \
         -DGOOGLE_CLOUD_CPP_ENABLE_CLANG_ABI_COMPAT_17=ON \
         -DBUILD_SHARED_LIBS=yes \
-        -DGOOGLE_CLOUD_CPP_ENABLE="bigquerycontrol,bigquery" \
+        -DGOOGLE_CLOUD_CPP_ENABLE="bigquerycontrol,bigquery,opentelemetry" \
       -GNinja -S . -B cmake-out && \
     cmake --build cmake-out --target install && \
     ldconfig && cd /var/tmp && rm -fr build

google/cloud/bigquery_unified/BUILD.bazel

@@ -21,6 +21,13 @@ package(default_visibility = ["//visibility:private"])
 
 licenses(["notice"])  # Apache 2.0
 
+config_setting(
+    name = "enable_opentelemetry",
+    flag_values = {
+        "@google_cloud_cpp//:enable_opentelemetry": "true",
+    },
+)
+
 filegroup(
     name = "common_public_hdrs",
     srcs = [h for h in google_cloud_cpp_bigquery_common_hdrs if not h.startswith("internal/")],
@@ -53,6 +60,15 @@ cc_library(
     name = "google_cloud_cpp_bigquery_bigquery_unified",
     srcs = google_cloud_cpp_bigquery_bigquery_unified_srcs,
     hdrs = google_cloud_cpp_bigquery_bigquery_unified_hdrs,
+    # Enable the config when the bool flag is visible.
+    # https://github.com/googleapis/google-cloud-cpp/pull/14961
+    # defines = select({
+    #     ":enable_opentelemetry": [
+    #         # Enable OpenTelemetry features in google-cloud-cpp
+    #         "GOOGLE_CLOUD_CPP_BIGQUERY_HAVE_OPENTELEMETRY",
+    #     ],
+    #     "//conditions:default": [],
+    # }),
     visibility = ["//:__pkg__"],
     deps = [
         ":google_cloud_cpp_bigquery_common",

google/cloud/bigquery_unified/CMakeLists.txt

@@ -147,6 +147,8 @@ set(bigquery_unified_library_files
     internal/default_options.cc
     internal/default_options.h
     internal/retry_traits.h
+    internal/tracing_connection.cc
+    internal/tracing_connection.h
     job_options.h
     read_arrow_response.h
     read_options.h
@@ -178,11 +180,27 @@ bigquery_unified_library(
     PKGCONFIG_DEPS
     ${bigquery_unified_pkgconfig_deps})
 
+if (GOOGLE_CLOUD_CPP_BIGQUERY_ENABLE_OPENTELEMETRY)
+    find_package(google_cloud_cpp_opentelemetry CONFIG)
+    if (TARGET google-cloud-cpp::opentelemetry)
+        target_compile_definitions(
+            google_cloud_cpp_bigquery_bigquery_unified
+            PUBLIC # Enable OpenTelemetry features in google-cloud-cpp-bigquery
+                   GOOGLE_CLOUD_CPP_BIGQUERY_HAVE_OPENTELEMETRY)
+    else ()
+        message(
+            FATAL_ERROR
+                "OpenTelemetry is not enabled in the package "
+                "google-cloud-cpp, the target "
+                "google-cloud-cpp::opentelemetry is not available.")
+    endif ()
+endif ()
+
 if (GOOGLE_CLOUD_CPP_BIGQUERY_WITH_MOCKS)
     add_library(google_cloud_cpp_bigquery_bigquery_unified_mocks INTERFACE)
     set(google_cloud_cpp_bigquery_bigquery_unified_mocks_hdrs
         # cmake-format: sort
-        mocks/mock_connection.h)
+        mocks/mock_connection.h mocks/mock_stream_range.h)
     export_list_to_bazel(
         "google_cloud_cpp_bigquery_bigquery_unified_mocks.bzl"
         "google_cloud_cpp_bigquery_bigquery_unified_mocks_hdrs" YEAR "2024")
@@ -236,8 +254,11 @@ function (bigquery_unified_client_define_tests)
     find_package(GTest CONFIG REQUIRED)
     set(bigquery_unified_client_unit_tests
         # cmake-format: sort
-        client_test.cc internal/connection_impl_test.cc
-        internal/default_options_test.cc)
+        client_test.cc
+        internal/connection_impl_test.cc
+        internal/default_options_test.cc
+        internal/tracing_connection_test.cc
+        mocks/mock_stream_range_test.cc)
 
     # Export the list of unit tests to a .bzl file so we do not need to maintain
     # the list in two places.

google/cloud/bigquery_unified/bigquery_unified_client_unit_tests.bzl

@@ -20,4 +20,6 @@ bigquery_unified_client_unit_tests = [
     "client_test.cc",
     "internal/connection_impl_test.cc",
     "internal/default_options_test.cc",
+    "internal/tracing_connection_test.cc",
+    "mocks/mock_stream_range_test.cc",
 ]

google/cloud/bigquery_unified/google_cloud_cpp_bigquery_bigquery_unified.bzl

@@ -25,6 +25,7 @@ google_cloud_cpp_bigquery_bigquery_unified_hdrs = [
     "internal/connection_impl.h",
     "internal/default_options.h",
     "internal/retry_traits.h",
+    "internal/tracing_connection.h",
     "job_options.h",
     "read_arrow_response.h",
     "read_options.h",
@@ -38,4 +39,5 @@ google_cloud_cpp_bigquery_bigquery_unified_srcs = [
     "internal/arrow_reader.cc",
     "internal/connection_impl.cc",
     "internal/default_options.cc",
+    "internal/tracing_connection.cc",
 ]

google/cloud/bigquery_unified/google_cloud_cpp_bigquery_bigquery_unified_mocks.bzl

@@ -18,4 +18,5 @@
 
 google_cloud_cpp_bigquery_bigquery_unified_mocks_hdrs = [
     "mocks/mock_connection.h",
+    "mocks/mock_stream_range.h",
 ]

google/cloud/bigquery_unified/internal/connection_impl.cc

@@ -21,6 +21,7 @@
 #include "google/cloud/bigquery_unified/internal/arrow_reader.h"
 #include "google/cloud/bigquery_unified/internal/async_rest_long_running_operation_custom.h"
 #include "google/cloud/bigquery_unified/internal/default_options.h"
+#include "google/cloud/bigquery_unified/internal/tracing_connection.h"
 #include "google/cloud/bigquery_unified/job_options.h"
 #include "google/cloud/bigquery_unified/read_options.h"
 #include "google/cloud/bigquery_unified/retry_policy.h"
@@ -446,10 +447,11 @@ std::shared_ptr<bigquery_unified::Connection> MakeDefaultConnectionImpl(
   // TODO: We should probably add a TracingConnection decorator in order to
   //  associate all the various rpcs that are called as part of these
   //  operations.
-  return std::make_shared<bigquery_unified_internal::ConnectionImpl>(
-      std::move(read_connection), std::move(job_connection),
-      std::move(read_options), std::move(job_options), std::move(job_stub),
-      std::move(background), std::move(options));
+  return MakeTracingConnection(
+      std::make_shared<bigquery_unified_internal::ConnectionImpl>(
+          std::move(read_connection), std::move(job_connection),
+          std::move(read_options), std::move(job_options), std::move(job_stub),
+          std::move(background), std::move(options)));
 }
 
 GOOGLE_CLOUD_CPP_BIGQUERY_INLINE_NAMESPACE_END

google/cloud/bigquery_unified/internal/tracing_connection.cc

@@ -0,0 +1,127 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "google/cloud/bigquery_unified/internal/tracing_connection.h"
+#include "google/cloud/internal/opentelemetry.h"
+#include "google/cloud/internal/traced_stream_range.h"
+
+namespace google::cloud::bigquery_unified_internal {
+GOOGLE_CLOUD_CPP_BIGQUERY_INLINE_NAMESPACE_BEGIN
+
+#ifdef GOOGLE_CLOUD_CPP_BIGQUERY_HAVE_OPENTELEMETRY
+
+TracingConnection::TracingConnection(
+    std::shared_ptr<bigquery_unified::Connection> child)
+    : child_(std::move(child)) {}
+
+future<StatusOr<google::cloud::bigquery::v2::Job>> TracingConnection::CancelJob(
+    google::cloud::bigquery::v2::CancelJobRequest const& request,
+    Options opts) {
+  auto span = internal::MakeSpan("bigquery_unified::Connection::CancelJob");
+  internal::OTelScope scope(span);
+  return internal::EndSpan(std::move(span), child_->CancelJob(request, opts));
+}
+
+StatusOr<google::cloud::bigquery::v2::JobReference>
+TracingConnection::CancelJob(
+    google::cloud::NoAwaitTag,
+    google::cloud::bigquery::v2::CancelJobRequest const& request,
+    Options opts) {
+  auto span = internal::MakeSpan("bigquery_unified::Connection::CancelJob");
+  auto scope = opentelemetry::trace::Scope(span);
+  return internal::EndSpan(*span,
+                           child_->CancelJob(NoAwaitTag{}, request, opts));
+}
+
+future<StatusOr<google::cloud::bigquery::v2::Job>> TracingConnection::CancelJob(
+    google::cloud::bigquery::v2::JobReference const& job_reference,
+    Options opts) {
+  auto span = internal::MakeSpan("bigquery_unified::Connection::CancelJob");
+  internal::OTelScope scope(span);
+  return internal::EndSpan(std::move(span),
+                           child_->CancelJob(job_reference, opts));
+}
+
+StatusOr<google::cloud::bigquery::v2::Job> TracingConnection::GetJob(
+    google::cloud::bigquery::v2::GetJobRequest const& request, Options opts) {
+  auto span = internal::MakeSpan("bigquery_unified::Connection::GetJob");
+  auto scope = opentelemetry::trace::Scope(span);
+  return internal::EndSpan(*span, child_->GetJob(request, opts));
+}
+
+future<StatusOr<google::cloud::bigquery::v2::Job>> TracingConnection::InsertJob(
+    google::cloud::bigquery::v2::Job const& job, Options opts) {
+  auto span = internal::MakeSpan("bigquery_unified::Connection::InsertJob");
+  internal::OTelScope scope(span);
+  return internal::EndSpan(std::move(span), child_->InsertJob(job, opts));
+}
+
+StatusOr<google::cloud::bigquery::v2::JobReference>
+TracingConnection::InsertJob(google::cloud::NoAwaitTag,
+                             google::cloud::bigquery::v2::Job const& job,
+                             Options opts) {
+  auto span = internal::MakeSpan("bigquery_unified::Connection::InsertJob");
+  auto scope = opentelemetry::trace::Scope(span);
+  return internal::EndSpan(*span, child_->InsertJob(NoAwaitTag{}, job, opts));
+}
+
+future<StatusOr<google::cloud::bigquery::v2::Job>> TracingConnection::InsertJob(
+    google::cloud::bigquery::v2::JobReference const& job_reference,
+    Options opts) {
+  auto span = internal::MakeSpan("bigquery_unified::Connection::InsertJob");
+  internal::OTelScope scope(span);
+  return internal::EndSpan(std::move(span),
+                           child_->InsertJob(job_reference, opts));
+}
+
+Status TracingConnection::DeleteJob(
+    google::cloud::bigquery::v2::DeleteJobRequest const& request,
+    Options opts) {
+  auto span = internal::MakeSpan("bigquery_unified::Connection::DeleteJob");
+  auto scope = opentelemetry::trace::Scope(span);
+  return internal::EndSpan(*span, child_->DeleteJob(request, opts));
+}
+
+StreamRange<google::cloud::bigquery::v2::ListFormatJob>
+TracingConnection::ListJobs(
+    google::cloud::bigquery::v2::ListJobsRequest request, Options opts) {
+  auto span = internal::MakeSpan("bigquery_unified::Connection::ListJobs");
+  internal::OTelScope scope(span);
+  auto sr = child_->ListJobs(std::move(request), opts);
+  return internal::MakeTracedStreamRange<
+      google::cloud::bigquery::v2::ListFormatJob>(std::move(span),
+                                                  std::move(sr));
+}
+
+StatusOr<bigquery_unified::ReadArrowResponse> TracingConnection::ReadArrow(
+    google::cloud::bigquery::storage::v1::CreateReadSessionRequest const&
+        read_session,
+    Options opts) {
+  // Not add span tracing for now, will add it after discussion.
+  return child_->ReadArrow(read_session, opts);
+}
+#endif  // GOOGLE_CLOUD_CPP_BIGQUERY_HAVE_OPENTELEMETRY
+
+std::shared_ptr<bigquery_unified::Connection> MakeTracingConnection(
+    std::shared_ptr<bigquery_unified::Connection> conn) {
+#ifdef GOOGLE_CLOUD_CPP_BIGQUERY_HAVE_OPENTELEMETRY
+  if (internal::TracingEnabled(conn->options())) {
+    conn = std::make_shared<TracingConnection>(std::move(conn));
+  }
+#endif  // GOOGLE_CLOUD_CPP_BIGQUERY_HAVE_OPENTELEMETRY
+  return conn;
+}
+
+GOOGLE_CLOUD_CPP_BIGQUERY_INLINE_NAMESPACE_END
+}  // namespace google::cloud::bigquery_unified_internal

google/cloud/bigquery_unified/internal/tracing_connection.h

@@ -0,0 +1,92 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef GOOGLE_CLOUD_CPP_BIGQUERY_GOOGLE_CLOUD_BIGQUERY_UNIFIED_INTERNAL_TRACING_CONNECTION_H
+#define GOOGLE_CLOUD_CPP_BIGQUERY_GOOGLE_CLOUD_BIGQUERY_UNIFIED_INTERNAL_TRACING_CONNECTION_H
+
+#include "google/cloud/bigquery/storage/v1/bigquery_read_connection.h"
+#include "google/cloud/bigquery_unified/connection.h"
+#include "google/cloud/bigquery_unified/version.h"
+#include "google/cloud/bigquerycontrol/v2/job_connection.h"
+#include <memory>
+
+namespace google::cloud::bigquery_unified_internal {
+GOOGLE_CLOUD_CPP_BIGQUERY_INLINE_NAMESPACE_BEGIN
+
+class TracingConnection : public bigquery_unified::Connection {
+ public:
+  ~TracingConnection() override = default;
+
+  explicit TracingConnection(
+      std::shared_ptr<bigquery_unified::Connection> child);
+
+  Options options() override { return child_->options(); }
+
+  future<StatusOr<google::cloud::bigquery::v2::Job>> CancelJob(
+      google::cloud::bigquery::v2::CancelJobRequest const& request,
+      Options opts) override;
+
+  StatusOr<google::cloud::bigquery::v2::JobReference> CancelJob(
+      google::cloud::NoAwaitTag,
+      google::cloud::bigquery::v2::CancelJobRequest const& request,
+      Options opts) override;
+
+  future<StatusOr<google::cloud::bigquery::v2::Job>> CancelJob(
+      google::cloud::bigquery::v2::JobReference const& job_reference,
+      Options opts) override;
+
+  StatusOr<google::cloud::bigquery::v2::Job> GetJob(
+      google::cloud::bigquery::v2::GetJobRequest const& request,
+      Options opts) override;
+
+  future<StatusOr<google::cloud::bigquery::v2::Job>> InsertJob(
+      google::cloud::bigquery::v2::Job const& job, Options opts) override;
+
+  StatusOr<google::cloud::bigquery::v2::JobReference> InsertJob(
+      google::cloud::NoAwaitTag, google::cloud::bigquery::v2::Job const& job,
+      Options opts) override;
+
+  future<StatusOr<google::cloud::bigquery::v2::Job>> InsertJob(
+      google::cloud::bigquery::v2::JobReference const& job_reference,
+      Options opts) override;
+
+  Status DeleteJob(google::cloud::bigquery::v2::DeleteJobRequest const& request,
+                   Options opts) override;
+
+  StreamRange<google::cloud::bigquery::v2::ListFormatJob> ListJobs(
+      google::cloud::bigquery::v2::ListJobsRequest request,
+      Options opts) override;
+
+  StatusOr<bigquery_unified::ReadArrowResponse> ReadArrow(
+      google::cloud::bigquery::storage::v1::CreateReadSessionRequest const&
+          read_session,
+      Options opts) override;
+
+ private:
+  std::shared_ptr<bigquery_unified::Connection> child_;
+};
+
+/**
+ * Conditionally applies the tracing decorator to the given connection.
+ *
+ * The connection is only decorated if tracing is enabled (as determined by the
+ * connection's options).
+ */
+std::shared_ptr<bigquery_unified::Connection> MakeTracingConnection(
+    std::shared_ptr<bigquery_unified::Connection> conn);
+
+GOOGLE_CLOUD_CPP_BIGQUERY_INLINE_NAMESPACE_END
+}  // namespace google::cloud::bigquery_unified_internal
+
+#endif  // GOOGLE_CLOUD_CPP_BIGQUERY_GOOGLE_CLOUD_BIGQUERY_UNIFIED_INTERNAL_TRACING_CONNECTION_H