databricks support changes

Author: pawel-big-lebowski

dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java

@@ -142,7 +142,9 @@ public AbstractDatadogSparkListener(SparkConf sparkConf, String appId, String sp
     this.appId = appId;
     this.sparkVersion = sparkVersion;
 
-    isRunningOnDatabricks = sparkConf.contains("spark.databricks.sparkContextId");
+    // isRunningOnDatabricks = sparkConf.contains("spark.databricks.sparkContextId");
+    isRunningOnDatabricks = false;
+    log.info("Overriding is Running on Databricks. Setting to false");
     databricksClusterName = sparkConf.get("spark.databricks.clusterUsageTags.clusterName", null);
     databricksServiceName = getDatabricksServiceName(sparkConf, databricksClusterName);
     sparkServiceName = getSparkServiceName(sparkConf, isRunningOnDatabricks);
@@ -167,7 +169,7 @@ public AbstractDatadogSparkListener(SparkConf sparkConf, String appId, String sp
   }
 
   public void setupOpenLineage(DDTraceId traceId) {
-    log.debug("Setting up OpenLineage configuration with trace id {}", traceId);
+    log.info("Setting up OpenLineage configuration with trace id {}", traceId);
     if (openLineageSparkListener != null) {
       openLineageSparkConf.set("spark.openlineage.transport.type", "composite");
       openLineageSparkConf.set("spark.openlineage.transport.continueOnFailure", "true");
@@ -190,10 +192,10 @@ public void setupOpenLineage(DDTraceId traceId) {
       setupOpenLineageCircuitBreaker();
       return;
     }
-    log.debug(
+    log.info(
         "There is no OpenLineage Spark Listener in the context. Skipping setting tags. {}",
         openLineageSparkListener);
-    log.debug(
+    log.info(
         "There is no OpenLineage SparkConf in the context. Skipping setting tags. {}",
         openLineageSparkConf);
   }
@@ -231,6 +233,8 @@ public synchronized void onApplicationStart(SparkListenerApplicationStart applic
           OpenlineageParentContext.from(sparkConf)
               .map(context -> context.getTraceId())
               .orElse(predeterminedTraceIdContext.getTraceId()));
+    } else {
+      log.info("Skipping openlineage setup. What a shame...");
     }
     notifyOl(x -> openLineageSparkListener.onApplicationStart(x), applicationStart);
   }
@@ -1301,9 +1305,9 @@ private static String getSparkServiceName(SparkConf conf, boolean isRunningOnDat
 
   private static String getServiceForOpenLineage(SparkConf conf, boolean isRunningOnDatabricks) {
     // Service for OpenLineage in Databricks is not supported yet
-    if (isRunningOnDatabricks) {
-      return null;
-    }
+    //    if (isRunningOnDatabricks) {
+    //      return null;
+    //    }
 
     // Keep service set by user, except if it is only "spark" or "hadoop" that can be set by USM
     String serviceName = Config.get().getServiceName();

dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractSparkInstrumentation.java

@@ -10,7 +10,9 @@
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.InstanceStore;
+import java.lang.reflect.InvocationTargetException;
 import net.bytebuddy.asm.Advice;
+import org.apache.spark.SparkConf;
 import org.apache.spark.deploy.SparkSubmitArguments;
 import org.apache.spark.scheduler.SparkListenerInterface;
 import org.slf4j.Logger;
@@ -127,6 +129,17 @@ public static boolean enter(@Advice.Argument(0) SparkListenerInterface listener)
           && "io.openlineage.spark.agent.OpenLineageSparkListener"
               .equals(listener.getClass().getCanonicalName())) {
         log.debug("Detected OpenLineage listener, skipping adding it to ListenerBus");
+
+        // for some reason on databricks env, the config was not captured
+        try {
+          log.info("Getting openlineage conf from the listener");
+          Object openLineageConf = listener.getClass().getMethod("getConf").invoke(listener);
+          InstanceStore.of(SparkConf.class)
+              .put("openLineageSparkConf", (SparkConf) openLineageConf);
+        } catch (IllegalAccessException | NoSuchMethodException | InvocationTargetException e) {
+          log.error("Something got wrong with getting openlineage conf: {}", e.getMessage());
+        }
+
         InstanceStore.of(SparkListenerInterface.class).put("openLineageListener", listener);
         return true;
       }

dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/OpenLineageInstrumentation.java

@@ -64,6 +64,7 @@ public static class OpenLineageSparkListenerAdvice {
     public static void exit(@Advice.This Object self, @Advice.FieldValue("conf") SparkConf conf)
         throws IllegalAccessException {
       Logger log = LoggerFactory.getLogger("OpenLineageSparkListenerAdvice");
+      log.info("Applying OpenLineageSparkListenerAdvice");
       if (!Config.get().isDataJobsOpenLineageEnabled()) {
         log.debug(
             "OpenLineage - Data Jobs integration disabled. Not manipulating OpenLineageSparkListener");