airbyte-server: set tags on span in OAuthHandler (#20229)

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

@@ -52,6 +52,12 @@ public static final class Tags {
      */
     public static final String CONNECTOR_VERSION_KEY = "connector_version";
 
+    /**
+     * Name of the APM trace tag that holds the destination definition ID value associated with the
+     * trace.
+     */
+    public static final String DESTINATION_DEFINITION_ID_KEY = "destination.definition_id";
+
     /**
      * Name of the APM trace tag that holds the destination Docker image value associated with the
      * trace.
@@ -100,6 +106,11 @@ public static final class Tags {
      */
     public static final String REPLICATION_STATUS_KEY = "replication.status";
 
+    /**
+     * Name of the APM trace tag that holds the source definition ID value associated with the trace.
+     */
+    public static final String SOURCE_DEFINITION_ID_KEY = "source.definition_id";
+
     /**
      * Name of the APM trace tag that holds the source Docker image value associated with the trace.
      */
@@ -115,6 +126,11 @@ public static final class Tags {
      */
     public static final String WEBHOOK_CONFIG_ID_KEY = "webhook.config_id";
 
+    /**
+     * Name of the APM trace tag that holds the workspace ID value associated with the trace.
+     */
+    public static final String WORKSPACE_ID_KEY = "workspace.id";
+
     private Tags() {}
 
   }

airbyte-server/build.gradle

@@ -17,6 +17,7 @@ dependencies {
     implementation project(':airbyte-config:config-models')
     implementation project(':airbyte-config:config-persistence')
     implementation project(':airbyte-config:specs')
+    implementation project(':airbyte-metrics:metrics-lib')
     implementation project(':airbyte-db:db-lib')
     implementation project(":airbyte-json-validation")
     implementation project(':airbyte-notification')

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

@@ -4,6 +4,10 @@
 
 package io.airbyte.server.handlers;
 
+import static io.airbyte.metrics.lib.ApmTraceConstants.Tags.DESTINATION_DEFINITION_ID_KEY;
+import static io.airbyte.metrics.lib.ApmTraceConstants.Tags.SOURCE_DEFINITION_ID_KEY;
+import static io.airbyte.metrics.lib.ApmTraceConstants.Tags.WORKSPACE_ID_KEY;
+
 import com.fasterxml.jackson.databind.JsonNode;
 import com.google.common.annotations.VisibleForTesting;
 import io.airbyte.analytics.TrackingClient;
@@ -26,6 +30,7 @@
 import io.airbyte.config.persistence.ConfigNotFoundException;
 import io.airbyte.config.persistence.ConfigRepository;
 import io.airbyte.config.persistence.SecretsRepositoryReader;
+import io.airbyte.metrics.lib.ApmTraceUtils;
 import io.airbyte.oauth.OAuthFlowImplementation;
 import io.airbyte.oauth.OAuthImplementationFactory;
 import io.airbyte.persistence.job.factory.OAuthConfigSupplier;
@@ -65,6 +70,8 @@ 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 StandardSourceDefinition sourceDefinition =
         configRepository.getStandardSourceDefinition(sourceOauthConsentRequest.getSourceDefinitionId());
     final OAuthFlowImplementation oAuthFlowImplementation = oAuthImplementationFactory.create(sourceDefinition);
@@ -107,6 +114,9 @@ public OAuthConsentRead getSourceOAuthConsent(final SourceOauthConsentRequest so
 
   public OAuthConsentRead getDestinationOAuthConsent(final DestinationOauthConsentRequest destinationOauthConsentRequest)
       throws JsonValidationException, ConfigNotFoundException, IOException {
+    ApmTraceUtils.addTagsToTrace(Map.of(WORKSPACE_ID_KEY, destinationOauthConsentRequest.getWorkspaceId(), DESTINATION_DEFINITION_ID_KEY,
+        destinationOauthConsentRequest.getDestinationDefinitionId()));
+
     final StandardDestinationDefinition destinationDefinition =
         configRepository.getStandardDestinationDefinition(destinationOauthConsentRequest.getDestinationDefinitionId());
     final OAuthFlowImplementation oAuthFlowImplementation = oAuthImplementationFactory.create(destinationDefinition);
@@ -150,6 +160,9 @@ public OAuthConsentRead getDestinationOAuthConsent(final DestinationOauthConsent
 
   public Map<String, Object> completeSourceOAuth(final CompleteSourceOauthRequest completeSourceOauthRequest)
       throws JsonValidationException, ConfigNotFoundException, IOException {
+    ApmTraceUtils.addTagsToTrace(Map.of(WORKSPACE_ID_KEY, completeSourceOauthRequest.getWorkspaceId(), SOURCE_DEFINITION_ID_KEY,
+        completeSourceOauthRequest.getSourceDefinitionId()));
+
     final StandardSourceDefinition sourceDefinition =
         configRepository.getStandardSourceDefinition(completeSourceOauthRequest.getSourceDefinitionId());
     final OAuthFlowImplementation oAuthFlowImplementation = oAuthImplementationFactory.create(sourceDefinition);
@@ -195,6 +208,9 @@ public Map<String, Object> completeSourceOAuth(final CompleteSourceOauthRequest
 
   public Map<String, Object> completeDestinationOAuth(final CompleteDestinationOAuthRequest completeDestinationOAuthRequest)
       throws JsonValidationException, ConfigNotFoundException, IOException {
+    ApmTraceUtils.addTagsToTrace(Map.of(WORKSPACE_ID_KEY, completeDestinationOAuthRequest.getWorkspaceId(), DESTINATION_DEFINITION_ID_KEY,
+        completeDestinationOAuthRequest.getDestinationDefinitionId()));
+
     final StandardDestinationDefinition destinationDefinition =
         configRepository.getStandardDestinationDefinition(completeDestinationOAuthRequest.getDestinationDefinitionId());
     final OAuthFlowImplementation oAuthFlowImplementation = oAuthImplementationFactory.create(destinationDefinition);