airbyte-metrics/server: add tags to root span from OAuth handler (#20468

)

airbyte-metrics/metrics-lib/src/main/java/io/airbyte/metrics/lib/ApmTraceUtils.java

@@ -87,15 +87,32 @@ public static void addExceptionToTrace(final Span span, final Throwable t) {
     }
   }
 
+  /**
+   * Adds all the provided tags to the root span.
+   *
+   * @param tags A map of tags to be added to the root span.
+   */
+  public static void addTagsToRootSpan(final Map<String, Object> tags) {
+    final Span activeSpan = GlobalTracer.get().activeSpan();
+    if (activeSpan instanceof MutableSpan) {
+      final MutableSpan localRootSpan = ((MutableSpan) activeSpan).getLocalRootSpan();
+      tags.entrySet().forEach(entry -> {
+        localRootSpan.setTag(formatTag(entry.getKey(), TAG_PREFIX), entry.getValue().toString());
+      });
+    }
+  }
+
   /**
    * Adds an exception to the root span, if an active one exists.
    *
    * @param t The {@link Throwable} to be added to the provided span.
    */
   public static void recordErrorOnRootSpan(final Throwable t) {
     final Span activeSpan = GlobalTracer.get().activeSpan();
-    activeSpan.setTag(Tags.ERROR, true);
-    activeSpan.log(Map.of(Fields.ERROR_OBJECT, t));
+    if (activeSpan != null) {
+      activeSpan.setTag(Tags.ERROR, true);
+      activeSpan.log(Map.of(Fields.ERROR_OBJECT, t));
+    }
     if (activeSpan instanceof MutableSpan) {
       final MutableSpan localRootSpan = ((MutableSpan) activeSpan).getLocalRootSpan();
       localRootSpan.setError(true);

airbyte-metrics/metrics-lib/src/test/java/io/airbyte/metrics/lib/ApmTraceUtilsTest.java

@@ -11,10 +11,17 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.withSettings;
 
+import datadog.trace.api.DDTags;
+import datadog.trace.api.interceptor.MutableSpan;
 import io.opentracing.Span;
 import io.opentracing.Tracer;
+import io.opentracing.log.Fields;
+import io.opentracing.tag.Tags;
 import io.opentracing.util.GlobalTracerTestUtil;
+import java.io.PrintWriter;
+import java.io.StringWriter;
 import java.util.Map;
 import org.junit.After;
 import org.junit.Before;
@@ -89,4 +96,56 @@ void testFormattingTagKeys() {
     assertEquals("airbyte." + tagPrefix1 + "." + tagKey1, result2);
   }
 
+  @Test
+  void testAddingTagsToRootSpan() {
+    final Span activeSpan = mock(Span.class, withSettings().extraInterfaces(MutableSpan.class));
+    final Tracer tracer = mock(Tracer.class);
+    final MutableSpan localRootSpan = mock(MutableSpan.class);
+    when(tracer.activeSpan()).thenReturn(activeSpan);
+    when(((MutableSpan) activeSpan).getLocalRootSpan()).thenReturn(localRootSpan);
+    GlobalTracerTestUtil.setGlobalTracerUnconditionally(tracer);
+    ApmTraceUtils.addTagsToRootSpan(TAGS);
+    verify(localRootSpan, times(1)).setTag(String.format(TAG_FORMAT, TAG_PREFIX, TAG_1), VALUE_1);
+    verify(localRootSpan, times(1)).setTag(String.format(TAG_FORMAT, TAG_PREFIX, TAG_2), VALUE_2);
+  }
+
+  @Test
+  void testAddingTagsToRootSpanWhenActiveSpanIsNull() {
+    final Tracer tracer = mock(Tracer.class);
+    when(tracer.activeSpan()).thenReturn(null);
+    GlobalTracerTestUtil.setGlobalTracerUnconditionally(tracer);
+    Assertions.assertDoesNotThrow(() -> ApmTraceUtils.addTagsToRootSpan(TAGS));
+  }
+
+  @Test
+  void testRecordErrorOnRootSpan() {
+    final Span activeSpan = mock(Span.class, withSettings().extraInterfaces(MutableSpan.class));
+    final Tracer tracer = mock(Tracer.class);
+    final MutableSpan localRootSpan = mock(MutableSpan.class);
+    final Throwable exception = mock(Throwable.class);
+    when(tracer.activeSpan()).thenReturn(activeSpan);
+    when(((MutableSpan) activeSpan).getLocalRootSpan()).thenReturn(localRootSpan);
+    GlobalTracerTestUtil.setGlobalTracerUnconditionally(tracer);
+
+    ApmTraceUtils.recordErrorOnRootSpan(exception);
+    verify(activeSpan, times(1)).setTag(Tags.ERROR, true);
+    verify(activeSpan, times(1)).log(Map.of(Fields.ERROR_OBJECT, exception));
+
+    verify(localRootSpan, times(1)).setError(true);
+    verify(localRootSpan, times(1)).setTag(DDTags.ERROR_MSG, exception.getMessage());
+    verify(localRootSpan, times(1)).setTag(DDTags.ERROR_TYPE, exception.getClass().getName());
+    final StringWriter expectedErrorString = new StringWriter();
+    exception.printStackTrace(new PrintWriter(expectedErrorString));
+    verify(localRootSpan, times(1)).setTag(DDTags.ERROR_STACK, expectedErrorString.toString());
+  }
+
+  @Test
+  void testRecordErrorOnRootSpanWhenActiveSpanIsNull() {
+    final Throwable exception = mock(Throwable.class);
+    final Tracer tracer = mock(Tracer.class);
+    when(tracer.activeSpan()).thenReturn(null);
+    GlobalTracerTestUtil.setGlobalTracerUnconditionally(tracer);
+    Assertions.assertDoesNotThrow(() -> ApmTraceUtils.recordErrorOnRootSpan(exception));
+  }
+
 }

airbyte-server/src/main/java/io/airbyte/server/handlers/OAuthHandler.java

@@ -70,8 +70,10 @@ public OAuthHandler(final ConfigRepository configRepository,
 
   public OAuthConsentRead getSourceOAuthConsent(final SourceOauthConsentRequest sourceOauthConsentRequest)
       throws JsonValidationException, ConfigNotFoundException, IOException {
-    ApmTraceUtils.addTagsToTrace(Map.of(WORKSPACE_ID_KEY, sourceOauthConsentRequest.getWorkspaceId(), SOURCE_DEFINITION_ID_KEY,
-        sourceOauthConsentRequest.getSourceDefinitionId()));
+    final Map<String, Object> traceTags = Map.of(WORKSPACE_ID_KEY, sourceOauthConsentRequest.getWorkspaceId(), SOURCE_DEFINITION_ID_KEY,
+        sourceOauthConsentRequest.getSourceDefinitionId());
+    ApmTraceUtils.addTagsToTrace(traceTags);
+    ApmTraceUtils.addTagsToRootSpan(traceTags);
     final StandardSourceDefinition sourceDefinition =
         configRepository.getStandardSourceDefinition(sourceOauthConsentRequest.getSourceDefinitionId());
     final OAuthFlowImplementation oAuthFlowImplementation = oAuthImplementationFactory.create(sourceDefinition);