fix: move OperationType to ApiTracerContext

diegomarquezp · diegomarquezp · commit 1b47eb997dad · 2026-03-10T10:54:50.000-04:00
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/ApiTracerContext.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/ApiTracerContext.java
@@ -32,6 +32,7 @@
 
 import com.google.api.core.InternalApi;
 import com.google.api.gax.rpc.LibraryMetadata;
+import com.google.api.gax.tracing.ApiTracerFactory.OperationType;
 import com.google.auto.value.AutoValue;
 import java.util.HashMap;
 import java.util.Map;
@@ -116,6 +117,14 @@ String rpcSystemName() {
   @Nullable
   abstract Transport transport();
 
+  /**
+   * Returns the type of operation the {@link ApiTracer} is tracing.
+   *
+   * @return the operation type, or {@code null} if not set
+   */
+  @Nullable
+  public abstract OperationType operationType();
+
   /**
    * @return a map of attributes to be included in attempt-level spans
    */
@@ -161,6 +170,9 @@ ApiTracerContext merge(ApiTracerContext other) {
     if (other.transport() != null) {
       builder.setTransport(other.transport());
     }
+    if (other.operationType() != null) {
+      builder.setOperationType(other.operationType());
+    }
     return builder.build();
   }
 
@@ -184,6 +196,8 @@ public abstract static class Builder {
 
     public abstract Builder setTransport(@Nullable Transport transport);
 
+    public abstract Builder setOperationType(@Nullable OperationType operationType);
+
     public abstract ApiTracerContext build();
   }
 }
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/ApiTracerFactory.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/ApiTracerFactory.java
@@ -75,6 +75,16 @@ default ApiTracer newTracer(
     return newTracer(parent, spanName, operationType);
   }
 
+  /**
+   * Create a new {@link ApiTracer} that will be a child of the current context.
+   *
+   * @param parent the parent of this tracer
+   * @param tracerContext the method-definition-specific tracer context
+   */
+  default ApiTracer newTracer(ApiTracer parent, ApiTracerContext tracerContext) {
+    return newTracer(parent, tracerContext, tracerContext.operationType());
+  }
+
   /**
    * @return the {@link ApiTracerContext} for this factory
    */
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedBatchingCallable.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedBatchingCallable.java
@@ -85,8 +85,7 @@ public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext context)
     // So this start a top level tracer.
     ApiTracer tracer;
     if (apiTracerContext != null) {
-      tracer =
-          tracerFactory.newTracer(context.getTracer(), apiTracerContext, OperationType.Batching);
+      tracer = tracerFactory.newTracer(context.getTracer(), apiTracerContext);
     } else {
       tracer = tracerFactory.newTracer(context.getTracer(), spanName, OperationType.Batching);
     }
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedBidiCallable.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedBidiCallable.java
@@ -87,9 +87,7 @@ public ClientStream<RequestT> internalCall(
 
     ApiTracer tracer;
     if (apiTracerContext != null) {
-      tracer =
-          tracerFactory.newTracer(
-              context.getTracer(), apiTracerContext, OperationType.BidiStreaming);
+      tracer = tracerFactory.newTracer(context.getTracer(), apiTracerContext);
     } else {
       tracer = tracerFactory.newTracer(context.getTracer(), spanName, OperationType.BidiStreaming);
     }
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedClientStreamingCallable.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedClientStreamingCallable.java
@@ -85,9 +85,7 @@ public ApiStreamObserver<RequestT> clientStreamingCall(
 
     ApiTracer tracer;
     if (apiTracerContext != null) {
-      tracer =
-          tracerFactory.newTracer(
-              context.getTracer(), apiTracerContext, OperationType.ClientStreaming);
+      tracer = tracerFactory.newTracer(context.getTracer(), apiTracerContext);
     } else {
       tracer =
           tracerFactory.newTracer(context.getTracer(), spanName, OperationType.ClientStreaming);
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedServerStreamingCallable.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedServerStreamingCallable.java
@@ -81,9 +81,7 @@ public void call(
 
     ApiTracer tracer;
     if (apiTracerContext != null) {
-      tracer =
-          tracerFactory.newTracer(
-              context.getTracer(), apiTracerContext, OperationType.ServerStreaming);
+      tracer = tracerFactory.newTracer(context.getTracer(), apiTracerContext);
     } else {
       tracer =
           tracerFactory.newTracer(context.getTracer(), spanName, OperationType.ServerStreaming);
diff --git a/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedUnaryCallable.java b/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedUnaryCallable.java
@@ -82,7 +82,7 @@ public TracedUnaryCallable(
   public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext context) {
     ApiTracer tracer;
     if (apiTracerContext != null) {
-      tracer = tracerFactory.newTracer(context.getTracer(), apiTracerContext, OperationType.Unary);
+      tracer = tracerFactory.newTracer(context.getTracer(), apiTracerContext);
     } else {
       tracer = tracerFactory.newTracer(context.getTracer(), spanName, OperationType.Unary);
     }
diff --git a/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedBatchingCallableTest.java b/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedBatchingCallableTest.java
@@ -60,6 +60,7 @@ class TracedBatchingCallableTest {
           .setFullMethodName("FakeClient/FakeRpc")
           .setTransport(Transport.GRPC)
           .setLibraryMetadata(LibraryMetadata.empty())
+          .setOperationType(OperationType.Batching)
           .build();
 
   @Mock private ApiTracerFactory tracerFactory;
@@ -74,8 +75,7 @@ class TracedBatchingCallableTest {
   void init(boolean useContext) {
     // Wire the mock tracer factory
     if (useContext) {
-      when(tracerFactory.newTracer(
-              any(ApiTracer.class), any(ApiTracerContext.class), eq(OperationType.Batching)))
+      when(tracerFactory.newTracer(any(ApiTracer.class), any(ApiTracerContext.class)))
           .thenReturn(tracer);
       tracedBatchingCallable =
           new TracedBatchingCallable<>(
@@ -102,8 +102,7 @@ void testRootTracerCreated(boolean useContext) {
     init(useContext);
     tracedBatchingCallable.futureCall("test", callContext);
     if (useContext) {
-      verify(tracerFactory, times(1))
-          .newTracer(callContext.getTracer(), TRACER_CONTEXT, OperationType.Batching);
+      verify(tracerFactory, times(1)).newTracer(callContext.getTracer(), TRACER_CONTEXT);
     } else {
       verify(tracerFactory, times(1))
           .newTracer(callContext.getTracer(), SPAN_NAME, OperationType.Batching);
diff --git a/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedBidiCallableTest.java b/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedBidiCallableTest.java
@@ -31,7 +31,6 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -65,6 +64,7 @@ class TracedBidiCallableTest {
           .setFullMethodName("fake-client/fake-method")
           .setTransport(Transport.GRPC)
           .setLibraryMetadata(LibraryMetadata.empty())
+          .setOperationType(OperationType.BidiStreaming)
           .build();
 
   private FakeBidiObserver outerObserver;
@@ -84,8 +84,7 @@ void init(boolean useContext) {
     outerCallContext = FakeCallContext.createDefault();
 
     if (useContext) {
-      when(tracerFactory.newTracer(
-              any(ApiTracer.class), any(ApiTracerContext.class), eq(OperationType.BidiStreaming)))
+      when(tracerFactory.newTracer(any(ApiTracer.class), any(ApiTracerContext.class)))
           .thenReturn(tracer);
     } else {
       when(tracerFactory.newTracer(parentTracer, SPAN_NAME, OperationType.BidiStreaming))
@@ -108,8 +107,7 @@ void testTracerCreated(boolean useContext) {
     tracedCallable.call(outerObserver, outerCallContext);
 
     if (useContext) {
-      verify(tracerFactory, times(1))
-          .newTracer(parentTracer, TRACER_CONTEXT, OperationType.BidiStreaming);
+      verify(tracerFactory, times(1)).newTracer(parentTracer, TRACER_CONTEXT);
     } else {
       verify(tracerFactory, times(1))
           .newTracer(parentTracer, SPAN_NAME, OperationType.BidiStreaming);
diff --git a/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedClientStreamingCallableTest.java b/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedClientStreamingCallableTest.java
@@ -31,7 +31,6 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -65,6 +64,7 @@ class TracedClientStreamingCallableTest {
           .setFullMethodName("fake-client/fake-method")
           .setTransport(Transport.GRPC)
           .setLibraryMetadata(LibraryMetadata.empty())
+          .setOperationType(OperationType.ClientStreaming)
           .build();
 
   @Mock private ApiTracerFactory tracerFactory;
@@ -81,8 +81,7 @@ void init(boolean useContext) {
     callContext = FakeCallContext.createDefault();
     innerCallable = new FakeClientCallable();
     if (useContext) {
-      when(tracerFactory.newTracer(
-              any(ApiTracer.class), any(ApiTracerContext.class), eq(OperationType.ClientStreaming)))
+      when(tracerFactory.newTracer(any(ApiTracer.class), any(ApiTracerContext.class)))
           .thenReturn(tracer);
       tracedCallable =
           new TracedClientStreamingCallable<>(innerCallable, tracerFactory, TRACER_CONTEXT);
@@ -100,8 +99,7 @@ void testTracerCreated(boolean useContext) {
     tracedCallable.clientStreamingCall(outerResponseObsever, callContext);
 
     if (useContext) {
-      verify(tracerFactory, times(1))
-          .newTracer(parentTracer, TRACER_CONTEXT, OperationType.ClientStreaming);
+      verify(tracerFactory, times(1)).newTracer(parentTracer, TRACER_CONTEXT);
     } else {
       verify(tracerFactory, times(1))
           .newTracer(parentTracer, SPAN_NAME, OperationType.ClientStreaming);
diff --git a/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedServerStreamingCallableTest.java b/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedServerStreamingCallableTest.java
@@ -65,6 +65,7 @@ class TracedServerStreamingCallableTest {
           .setFullMethodName("FakeClient/FakeRpc")
           .setTransport(Transport.GRPC)
           .setLibraryMetadata(LibraryMetadata.empty())
+          .setOperationType(OperationType.ServerStreaming)
           .build();
 
   @Mock private ApiTracerFactory tracerFactory;
@@ -80,8 +81,7 @@ void init(boolean useContext) {
     innerCallable = new MockServerStreamingCallable<>();
     // Wire the mock tracer factory
     if (useContext) {
-      when(tracerFactory.newTracer(
-              any(ApiTracer.class), any(ApiTracerContext.class), eq(OperationType.ServerStreaming)))
+      when(tracerFactory.newTracer(any(ApiTracer.class), any(ApiTracerContext.class)))
           .thenReturn(tracer);
       tracedCallable =
           new TracedServerStreamingCallable<>(innerCallable, tracerFactory, TRACER_CONTEXT);
@@ -102,8 +102,7 @@ void testTracerCreated(boolean useContext) {
     init(useContext);
     tracedCallable.call("test", responseObserver, callContext);
     if (useContext) {
-      verify(tracerFactory, times(1))
-          .newTracer(parentTracer, TRACER_CONTEXT, OperationType.ServerStreaming);
+      verify(tracerFactory, times(1)).newTracer(parentTracer, TRACER_CONTEXT);
     } else {
       verify(tracerFactory, times(1))
           .newTracer(parentTracer, SPAN_NAME, OperationType.ServerStreaming);
diff --git a/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedUnaryCallableTest.java b/gax-java/gax/src/test/java/com/google/api/gax/tracing/TracedUnaryCallableTest.java
@@ -58,6 +58,7 @@ class TracedUnaryCallableTest {
           .setFullMethodName("FakeClient/FakeRpc")
           .setTransport(Transport.GRPC)
           .setLibraryMetadata(LibraryMetadata.empty())
+          .setOperationType(OperationType.Unary)
           .build();
 
   @Mock private ApiTracerFactory tracerFactory;
@@ -74,8 +75,7 @@ void init(boolean useContext) {
 
     // Wire the mock tracer factory
     if (useContext) {
-      when(tracerFactory.newTracer(
-              any(ApiTracer.class), any(ApiTracerContext.class), eq(OperationType.Unary)))
+      when(tracerFactory.newTracer(any(ApiTracer.class), any(ApiTracerContext.class)))
           .thenReturn(tracer);
       tracedUnaryCallable = new TracedUnaryCallable<>(innerCallable, tracerFactory, TRACER_CONTEXT);
     } else {
@@ -98,7 +98,7 @@ void testTracerCreated(boolean useContext) {
     init(useContext);
     tracedUnaryCallable.futureCall("test", callContext);
     if (useContext) {
-      verify(tracerFactory, times(1)).newTracer(parentTracer, TRACER_CONTEXT, OperationType.Unary);
+      verify(tracerFactory, times(1)).newTracer(parentTracer, TRACER_CONTEXT);
     } else {
       verify(tracerFactory, times(1)).newTracer(parentTracer, SPAN_NAME, OperationType.Unary);
     }

Original file line number	Diff line number	Diff line change
`@@ -82,7 +82,7 @@ public TracedUnaryCallable(`
`82`	`82`	`public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext context) {`
`83`	`83`	`ApiTracer tracer;`
`84`	`84`	`if (apiTracerContext != null) {`
`85`		`- tracer = tracerFactory.newTracer(context.getTracer(), apiTracerContext, OperationType.Unary);`
	`85`	`+ tracer = tracerFactory.newTracer(context.getTracer(), apiTracerContext);`
`86`	`86`	`} else {`
`87`	`87`	`tracer = tracerFactory.newTracer(context.getTracer(), spanName, OperationType.Unary);`
`88`	`88`	`}`