From 71d71a465bec2144c0aede618f9ac6ecb19a2501 Mon Sep 17 00:00:00 2001 From: Augustin Date: Thu, 15 Dec 2022 14:18:58 +0100 Subject: [PATCH] airbyte-metrics/server: add tags to root span from OAuth handler (#20468) --- .../io/airbyte/metrics/lib/ApmTraceUtils.java | 21 ++++++- .../metrics/lib/ApmTraceUtilsTest.java | 59 +++++++++++++++++++ .../airbyte/server/handlers/OAuthHandler.java | 6 +- 3 files changed, 82 insertions(+), 4 deletions(-) diff --git a/airbyte-metrics/metrics-lib/src/main/java/io/airbyte/metrics/lib/ApmTraceUtils.java b/airbyte-metrics/metrics-lib/src/main/java/io/airbyte/metrics/lib/ApmTraceUtils.java index d38c29bbe04d..bcd273a973c2 100644 --- a/airbyte-metrics/metrics-lib/src/main/java/io/airbyte/metrics/lib/ApmTraceUtils.java +++ b/airbyte-metrics/metrics-lib/src/main/java/io/airbyte/metrics/lib/ApmTraceUtils.java @@ -87,6 +87,21 @@ 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 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. * @@ -94,8 +109,10 @@ public static void addExceptionToTrace(final Span span, final Throwable t) { */ 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); diff --git a/airbyte-metrics/metrics-lib/src/test/java/io/airbyte/metrics/lib/ApmTraceUtilsTest.java b/airbyte-metrics/metrics-lib/src/test/java/io/airbyte/metrics/lib/ApmTraceUtilsTest.java index c45eb92df676..d6b95df15e31 100644 --- a/airbyte-metrics/metrics-lib/src/test/java/io/airbyte/metrics/lib/ApmTraceUtilsTest.java +++ b/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)); + } + } diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/OAuthHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/OAuthHandler.java index 1addffd1d58f..1d1517fd444d 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/OAuthHandler.java +++ b/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 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);