Add tensorstore bazel aliases for cc_grpc_library

PiperOrigin-RevId: 713103236 Change-Id: I11124c491da18e94f7edaca5ffbb3d2676fd9d25
google · Jan 8, 2025 · 424c2c7 · 424c2c7
1 parent 7d7c347
commit 424c2c7
Show file tree

Hide file tree

Showing 8 changed files with 43 additions and 9 deletions.
diff --git a/bazel/tensorstore.bzl b/bazel/tensorstore.bzl
@@ -15,6 +15,10 @@
 """Rule definitions for TensorStore targets."""
 
 load("@bazel_skylib//rules:build_test.bzl", "build_test")
+load(
+    "//bazel:cc_grpc_library.bzl",
+    _cc_grpc_library = "cc_grpc_library",
+)
 
 def _get_tensorstore_copts(copts):
     return (copts or []) + []
@@ -53,3 +57,5 @@ def tensorstore_proto_library(has_services = None, **kwargs):
     native.proto_library(**kwargs)
 
 tensorstore_cc_proto_library = native.cc_proto_library
+
+tensorstore_cc_grpc_library = _cc_grpc_library
diff --git a/tensorstore/internal/thread/BUILD b/tensorstore/internal/thread/BUILD
@@ -62,11 +62,16 @@ tensorstore_cc_test(
     ],
 )
 
+THREAD_POOL_DEFINES = []
+
+THREAD_POOL_DEPS = []
+
 tensorstore_cc_library(
     name = "thread_pool",
     srcs = ["thread_pool.cc"],
     hdrs = ["thread_pool.h"],
-    deps = [
+    local_defines = THREAD_POOL_DEFINES,
+    deps = THREAD_POOL_DEPS + [
         ":pool_impl",
         ":task",
         ":task_group_impl",

diff --git a/tensorstore/internal/tracing/BUILD b/tensorstore/internal/tracing/BUILD
@@ -11,6 +11,10 @@ tensorstore_cc_library(
     ],
 )
 
+TRACING_DEFINES = []
+
+TRACING_DEPS = []
+
 tensorstore_cc_library(
     name = "tracing",
     srcs = [
@@ -22,7 +26,8 @@ tensorstore_cc_library(
         "trace_context.h",
         "trace_span.h",
     ],
-    deps = [
+    defines = TRACING_DEFINES,
+    deps = TRACING_DEPS + [
         ":span_attribute",
         "//tensorstore/internal:source_location",
         "//tensorstore/util:span",

diff --git a/tensorstore/internal/tracing/trace_test.cc b/tensorstore/internal/tracing/trace_test.cc
@@ -24,6 +24,7 @@
 #include "absl/status/status.h"
 #include "tensorstore/internal/tracing/logged_trace_span.h"
 #include "tensorstore/internal/tracing/span_attribute.h"
+#include "tensorstore/internal/tracing/trace_context.h"
 #include "tensorstore/internal/tracing/trace_span.h"
 
 namespace {
@@ -34,6 +35,13 @@ using ::tensorstore::internal_tracing::TraceSpan;
 using ::testing::_;
 using ::testing::HasSubstr;
 
+TEST(TraceTest, SwapContext) {
+  tensorstore::internal_tracing::TraceContext tc(
+      tensorstore::internal_tracing::TraceContext::kThread);
+
+  tensorstore::internal_tracing::SwapCurrentTraceContext(&tc);
+}
+
 TEST(TraceTest, Span) {
   TraceSpan span("TraceSpan",
                  {

diff --git a/tensorstore/kvstore/gcs/BUILD b/tensorstore/kvstore/gcs/BUILD
@@ -109,11 +109,16 @@ tensorstore_cc_library(
     alwayslink = 1,
 )
 
+GCS_CREDENTIALS_DEPS = []
+
+GCS_CREDENTIALS_DEFINES = []
+
 tensorstore_cc_library(
     name = "exp_credentials_spec",
     srcs = ["exp_credentials_spec.cc"],
     hdrs = ["exp_credentials_spec.h"],
-    deps = [
+    defines = GCS_CREDENTIALS_DEFINES,
+    deps = GCS_CREDENTIALS_DEPS + [
         "//tensorstore:json_serialization_options_base",
         "//tensorstore/internal:type_traits",
         "//tensorstore/internal/grpc/clientauth:authentication_strategy",

diff --git a/tensorstore/kvstore/gcs_grpc/BUILD b/tensorstore/kvstore/gcs_grpc/BUILD
@@ -198,11 +198,16 @@ tensorstore_cc_library(
     ],
 )
 
+DEFAULT_STRATEGY_DEPS = []
+
+DEFAULT_STRATEGY_DEFINES = []
+
 tensorstore_cc_library(
     name = "default_strategy",
     srcs = ["default_strategy.cc"],
     hdrs = ["default_strategy.h"],
-    deps = [
+    local_defines = DEFAULT_STRATEGY_DEFINES,
+    deps = DEFAULT_STRATEGY_DEPS + [
         "//tensorstore/internal/grpc/clientauth:authentication_strategy",
         "//tensorstore/internal/grpc/clientauth:channel_authentication",
         "@com_google_absl//absl/log:absl_log",

diff --git a/tensorstore/kvstore/ocdbt/distributed/BUILD b/tensorstore/kvstore/ocdbt/distributed/BUILD
@@ -1,8 +1,8 @@
 # Distributed write support for Optionally-cooperative Distributed B+tree driver
 
-load("//bazel:cc_grpc_library.bzl", "cc_grpc_library")
 load(
     "//bazel:tensorstore.bzl",
+    "tensorstore_cc_grpc_library",
     "tensorstore_cc_library",
     "tensorstore_cc_proto_library",
     "tensorstore_cc_test",
@@ -84,7 +84,7 @@ tensorstore_cc_proto_library(
     deps = [":coordinator_proto"],
 )
 
-cc_grpc_library(
+tensorstore_cc_grpc_library(
     name = "coordinator_cc_grpc",
     srcs = [":coordinator_proto"],
     deps = [
@@ -103,7 +103,7 @@ tensorstore_cc_proto_library(
     deps = [":cooperator_proto"],
 )
 
-cc_grpc_library(
+tensorstore_cc_grpc_library(
     name = "cooperator_cc_grpc",
     srcs = [":cooperator_proto"],
     deps = [

diff --git a/tensorstore/kvstore/tsgrpc/BUILD b/tensorstore/kvstore/tsgrpc/BUILD
@@ -1,7 +1,7 @@
-load("//bazel:cc_grpc_library.bzl", "cc_grpc_library")
 load(
     "//bazel:tensorstore.bzl",
     "tensorstore_cc_binary",
+    "tensorstore_cc_grpc_library",
     "tensorstore_cc_library",
     "tensorstore_cc_proto_library",
     "tensorstore_cc_test",
@@ -48,7 +48,7 @@ tensorstore_cc_proto_library(
     deps = [":kvstore_proto"],
 )
 
-cc_grpc_library(
+tensorstore_cc_grpc_library(
     name = "kvstore_cc_grpc",
     srcs = [":kvstore_proto"],
     service_namespace = "grpc_gen",