Merge pull request #84 from cuiy0006/add-tracing-connection

impl: Add TracingConnection decorator

Author: cuiy0006

.bazelrc

@@ -55,7 +55,9 @@ 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
+# Enable the config when the bool flag is visible.
+# https://github.com/googleapis/google-cloud-cpp-bigquery/issues/98
+# 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.35.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-bigquery/issues/98
+    # 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,29 @@ 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
+                "Instrumentation with OpenTelemetry was requested "
+                "for google-cloud-cpp-bigquery, but it was not enabled "
+                "in google-cloud-cpp. Please build and install "
+                "google-cloud-cpp with GOOGLE_CLOUD_CPP_ENABLE=opentelemetry "
+                "if you want to use this feature.")
+    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 +256,12 @@ 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
+        connection_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

@@ -18,6 +18,9 @@
 
 bigquery_unified_client_unit_tests = [
     "client_test.cc",
+    "connection_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/connection_test.cc

@@ -0,0 +1,66 @@
+// 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/connection.h"
+#include "google/cloud/bigquery_unified/testing_util/opentelemetry_matchers.h"
+#include "google/cloud/bigquery_unified/version.h"
+
+namespace google::cloud::bigquery_unified {
+GOOGLE_CLOUD_CPP_BIGQUERY_INLINE_NAMESPACE_BEGIN
+namespace {
+
+#ifdef GOOGLE_CLOUD_CPP_BIGQUERY_HAVE_OPENTELEMETRY
+using ::google::cloud::bigquery_unified::testing_util::DisableTracing;
+using ::google::cloud::bigquery_unified::testing_util::EnableTracing;
+using ::google::cloud::bigquery_unified::testing_util::SpanNamed;
+using ::testing::Contains;
+using ::testing::Not;
+
+TEST(BigQueryUnifiedConnectionTest, TracingEnabled) {
+  auto span_catcher = testing_util::InstallSpanCatcher();
+
+  auto options = EnableTracing(Options{});
+  auto conn = MakeConnection(std::move(options));
+  google::cloud::internal::OptionsSpan span(
+      google::cloud::internal::MergeOptions(Options{}, conn->options()));
+  // Make a call, which should fail fast. The error itself is not important.
+  google::cloud::bigquery::v2::GetJobRequest request;
+  (void)conn->GetJob(request, Options{});
+
+  auto spans = span_catcher->GetSpans();
+  EXPECT_THAT(spans,
+              Contains(SpanNamed("bigquery_unified::Connection::GetJob")));
+}
+
+TEST(BigQueryUnifiedConnectionTest, TracingDisabled) {
+  auto span_catcher = testing_util::InstallSpanCatcher();
+
+  auto options = DisableTracing(Options{});
+  auto conn = MakeConnection(std::move(options));
+  google::cloud::internal::OptionsSpan span(
+      google::cloud::internal::MergeOptions(Options{}, conn->options()));
+  // Make a call, which should fail fast. The error itself is not important.
+  google::cloud::bigquery::v2::GetJobRequest request;
+  (void)conn->GetJob(request, Options{});
+
+  auto spans = span_catcher->GetSpans();
+  EXPECT_THAT(spans,
+              Not(Contains(SpanNamed("bigquery_unified::Connection::GetJob"))));
+}
+
+#endif  // GOOGLE_CLOUD_CPP_BIGQUERY_HAVE_OPENTELEMETRY
+
+}  // namespace
+GOOGLE_CLOUD_CPP_BIGQUERY_INLINE_NAMESPACE_END
+}  // namespace google::cloud::bigquery_unified

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",
 ]