add config variable to determine whether enable ol injection

mobuchowski · mobuchowski · commit 1c43c96a2da7 · 2025-03-26T16:02:18.000+01:00
Signed-off-by: Maciej Obuchowski &lt;maciej.obuchowski@datadoghq.com&gt;
diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java
@@ -242,6 +242,11 @@ public static void start(
       setSystemPropertyDefault(
           propertyNameToSystemPropertyName("integration.kafka.enabled"), "true");
 
+      if (Config.get().isDataJobsOpenLineageEnabled()) {
+        setSystemPropertyDefault(
+            propertyNameToSystemPropertyName("integration.openlineage-spark.enabled"), "true");
+      }
+
       String javaCommand = System.getProperty("sun.java.command");
       String dataJobsCommandPattern = Config.get().getDataJobsCommandPattern();
       if (!isDataJobsSupported(javaCommand, dataJobsCommandPattern)) {
diff --git a/dd-java-agent/instrumentation/spark/spark_2.12/src/main/java/datadog/trace/instrumentation/spark/Spark212Instrumentation.java b/dd-java-agent/instrumentation/spark/spark_2.12/src/main/java/datadog/trace/instrumentation/spark/Spark212Instrumentation.java
@@ -5,8 +5,10 @@
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Config;
 import net.bytebuddy.asm.Advice;
 import org.apache.spark.SparkContext;
+import org.apache.spark.util.Utils;
 
 @AutoService(InstrumenterModule.class)
 public class Spark212Instrumentation extends AbstractSparkInstrumentation {
@@ -44,27 +46,25 @@ public static class InjectListener {
     public static void enter(@Advice.This SparkContext sparkContext) {
       // checking whether OpenLineage integration is available and that it supports tags
       // Disabling this mechanism for this PR. Will be enabled with provided with Config option.
-      //      if (Utils.classIsLoadable("io.openlineage.spark.agent.OpenLineageSparkListener")
-      //          && Utils.classIsLoadable(
-      //              "io.openlineage.spark.agent.facets.builder.TagsRunFacetBuilder")) {
-      //        if (!sparkContext.conf().contains("spark.extraListeners")) {
-      //          sparkContext
-      //              .conf()
-      //              .set("spark.extraListeners",
-      // "io.openlineage.spark.agent.OpenLineageSparkListener");
-      //        } else {
-      //          String extraListeners = sparkContext.conf().get("spark.extraListeners");
-      //          if
-      // (!extraListeners.contains("io.openlineage.spark.agent.OpenLineageSparkListener")) {
-      //            sparkContext
-      //                .conf()
-      //                .set(
-      //                    "spark.extraListeners",
-      //                    extraListeners +
-      // ",io.openlineage.spark.agent.OpenLineageSparkListener");
-      //          }
-      //        }
-      //      }
+      if (Config.get().isDataJobsOpenLineageEnabled()
+          && Utils.classIsLoadable("io.openlineage.spark.agent.OpenLineageSparkListener")
+          && Utils.classIsLoadable(
+              "io.openlineage.spark.agent.facets.builder.TagsRunFacetBuilder")) {
+        if (!sparkContext.conf().contains("spark.extraListeners")) {
+          sparkContext
+              .conf()
+              .set("spark.extraListeners", "io.openlineage.spark.agent.OpenLineageSparkListener");
+        } else {
+          String extraListeners = sparkContext.conf().get("spark.extraListeners");
+          if (!extraListeners.contains("io.openlineage.spark.agent.OpenLineageSparkListener")) {
+            sparkContext
+                .conf()
+                .set(
+                    "spark.extraListeners",
+                    extraListeners + ",io.openlineage.spark.agent.OpenLineageSparkListener");
+          }
+        }
+      }
 
       // We want to add the Datadog listener as the first listener
       AbstractDatadogSparkListener.listener =
diff --git a/dd-java-agent/instrumentation/spark/spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213Instrumentation.java b/dd-java-agent/instrumentation/spark/spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213Instrumentation.java
@@ -5,8 +5,10 @@
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Config;
 import net.bytebuddy.asm.Advice;
 import org.apache.spark.SparkContext;
+import org.apache.spark.util.Utils;
 
 @AutoService(InstrumenterModule.class)
 public class Spark213Instrumentation extends AbstractSparkInstrumentation {
@@ -42,29 +44,29 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class InjectListener {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter(@Advice.This SparkContext sparkContext) {
-      // checking whether OpenLineage integration is available and that it supports tags
-      // Disabling this mechanism for this PR. Will be enabled with provided with Config option.
-      //      if (Utils.classIsLoadable("io.openlineage.spark.agent.OpenLineageSparkListener")
-      //          && Utils.classIsLoadable(
-      //              "io.openlineage.spark.agent.facets.builder.TagsRunFacetBuilder")) {
-      //        if (!sparkContext.conf().contains("spark.extraListeners")) {
-      //          sparkContext
-      //              .conf()
-      //              .set("spark.extraListeners",
-      // "io.openlineage.spark.agent.OpenLineageSparkListener");
-      //        } else {
-      //          String extraListeners = sparkContext.conf().get("spark.extraListeners");
-      //          if
-      // (!extraListeners.contains("io.openlineage.spark.agent.OpenLineageSparkListener")) {
-      //            sparkContext
-      //                .conf()
-      //                .set(
-      //                    "spark.extraListeners",
-      //                    extraListeners +
-      // ",io.openlineage.spark.agent.OpenLineageSparkListener");
-      //          }
-      //        }
-      //      }
+      // Checking whether OpenLineage integration is available and that it supports tags
+      // Disabling this mechanism for this PR. Will be enabled with provided with Config
+      // option.
+
+      if (Config.get().isDataJobsOpenLineageEnabled()
+          && Utils.classIsLoadable("io.openlineage.spark.agent.OpenLineageSparkListener")
+          && Utils.classIsLoadable(
+              "io.openlineage.spark.agent.facets.builder.TagsRunFacetBuilder")) {
+        if (!sparkContext.conf().contains("spark.extraListeners")) {
+          sparkContext
+              .conf()
+              .set("spark.extraListeners", "io.openlineage.spark.agent.OpenLineageSparkListener");
+        } else {
+          String extraListeners = sparkContext.conf().get("spark.extraListeners");
+          if (!extraListeners.contains("io.openlineage.spark.agent.OpenLineageSparkListener")) {
+            sparkContext
+                .conf()
+                .set(
+                    "spark.extraListeners",
+                    extraListeners + ",io.openlineage.spark.agent.OpenLineageSparkListener");
+          }
+        }
+      }
 
       // We want to add the Datadog listener as the first listener
       AbstractDatadogSparkListener.listener =
diff --git a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
@@ -206,10 +206,12 @@ public void setupOpenLineage(DDTraceId traceId) {
   public synchronized void onApplicationStart(SparkListenerApplicationStart applicationStart) {
     this.applicationStart = applicationStart;
 
-    setupOpenLineage(
-        OpenlineageParentContext.from(sparkConf)
-            .map(context -> context.getTraceId())
-            .orElse(predeterminedTraceIdContext.getTraceId()));
+    if (Config.get().isDataJobsOpenLineageEnabled()) {
+      setupOpenLineage(
+          OpenlineageParentContext.from(sparkConf)
+              .map(context -> context.getTraceId())
+              .orElse(predeterminedTraceIdContext.getTraceId()));
+    }
     notifyOl(x -> openLineageSparkListener.onApplicationStart(x), applicationStart);
   }
 
@@ -750,6 +752,11 @@ public void onOtherEvent(SparkListenerEvent event) {
   }
 
   private <T extends SparkListenerEvent> void notifyOl(Consumer<T> ol, T event) {
+    if (!Config.get().isDataJobsOpenLineageEnabled()) {
+      log.trace("Ignoring event {} - OpenLineage not enabled", event);
+      return;
+    }
+
     if (isRunningOnDatabricks || isStreamingJob) {
       log.debug("Not emitting event when running on databricks or on streaming jobs");
       return;
diff --git a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractSparkInstrumentation.java b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractSparkInstrumentation.java
@@ -117,8 +117,11 @@ public static void enter(@Advice.Argument(1) int exitCode, @Advice.Argument(2) S
 
   public static class LiveListenerBusAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class, skipOn = Advice.OnNonDefaultValue.class)
+    // If OL is disabled in tracer config but user set it up manually don't interfere
     public static boolean enter(@Advice.Argument(0) SparkListenerInterface listener) {
-      if (listener == null || listener.getClass().getCanonicalName() == null) {
+      if (!Config.get().isDataJobsOpenLineageEnabled()
+          || listener == null
+          || listener.getClass().getCanonicalName() == null) {
         return false;
       }
       if (listener
diff --git a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/OpenLineageInstrumentation.java b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/OpenLineageInstrumentation.java
@@ -41,7 +41,7 @@ public String[] helperClassNames() {
 
   @Override
   public boolean defaultEnabled() {
-    return true;
+    return false;
   }
 
   @Override
diff --git a/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSpark24SqlTest.groovy b/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSpark24SqlTest.groovy
@@ -14,6 +14,7 @@ abstract class AbstractSpark24SqlTest extends AgentTestRunner {
   void configurePreAgent() {
     super.configurePreAgent()
     injectSysConfig("dd.integration.spark.enabled", "true")
+    injectSysConfig("dd.integration.openlineage-spark.enabled", "true")
   }
 
   static Dataset<Row> generateSampleDataframe(SparkSession spark) {
diff --git a/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSpark32SqlTest.groovy b/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSpark32SqlTest.groovy
@@ -10,6 +10,7 @@ abstract class AbstractSpark32SqlTest extends AgentTestRunner {
   void configurePreAgent() {
     super.configurePreAgent()
     injectSysConfig("dd.integration.spark.enabled", "true")
+    injectSysConfig("dd.integration.openlineage-spark.enabled", "true")
   }
 
   def "compute a GROUP BY sql query plan"() {
diff --git a/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkStructuredStreamingTest.groovy b/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkStructuredStreamingTest.groovy
@@ -25,6 +25,7 @@ class AbstractSparkStructuredStreamingTest extends AgentTestRunner {
   void configurePreAgent() {
     super.configurePreAgent()
     injectSysConfig("dd.integration.spark.enabled", "true")
+    injectSysConfig("dd.integration.openlineage-spark.enabled", "true")
   }
 
   private SparkSession createSparkSession(String appName) {
diff --git a/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy b/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy
@@ -34,6 +34,7 @@ abstract class AbstractSparkTest extends AgentTestRunner {
   void configurePreAgent() {
     super.configurePreAgent()
     injectSysConfig("dd.integration.spark.enabled", "true")
+    injectSysConfig("dd.integration.openlineage-spark.enabled", "true")
   }
 
   def "generate application span with child job and stages"() {
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java
@@ -221,6 +221,7 @@ public final class ConfigDefaults {
   static final int DEFAULT_CWS_TLS_REFRESH = 5000;
 
   static final boolean DEFAULT_DATA_JOBS_ENABLED = false;
+  static final boolean DEFAULT_DATA_JOBS_OPENLINEAGE_ENABLED = false;
 
   static final boolean DEFAULT_DATA_STREAMS_ENABLED = false;
   static final int DEFAULT_DATA_STREAMS_BUCKET_DURATION = 10; // seconds
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java
@@ -69,6 +69,7 @@ public final class GeneralConfig {
 
   public static final String DATA_JOBS_ENABLED = "data.jobs.enabled";
   public static final String DATA_JOBS_COMMAND_PATTERN = "data.jobs.command.pattern";
+  public static final String DATA_JOBS_OPENLINEAGE_ENABLED = "data.jobs.openlineage.enabled";
 
   public static final String DATA_STREAMS_ENABLED = "data.streams.enabled";
   public static final String DATA_STREAMS_BUCKET_DURATION_SECONDS =
diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java
@@ -498,6 +498,7 @@ public static String getHostName() {
 
   private final boolean dataJobsEnabled;
   private final String dataJobsCommandPattern;
+  private final boolean dataJobsOpenLineageEnabled;
 
   private final boolean dataStreamsEnabled;
   private final float dataStreamsBucketDurationSeconds;
@@ -1827,6 +1828,9 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment())
     cwsTlsRefresh = configProvider.getInteger(CWS_TLS_REFRESH, DEFAULT_CWS_TLS_REFRESH);
 
     dataJobsEnabled = configProvider.getBoolean(DATA_JOBS_ENABLED, DEFAULT_DATA_JOBS_ENABLED);
+    dataJobsOpenLineageEnabled =
+        configProvider.getBoolean(
+            DATA_JOBS_OPENLINEAGE_ENABLED, DEFAULT_DATA_JOBS_OPENLINEAGE_ENABLED);
     dataJobsCommandPattern = configProvider.getString(DATA_JOBS_COMMAND_PATTERN);
 
     dataStreamsEnabled =
@@ -3587,6 +3591,10 @@ public boolean isDataJobsEnabled() {
     return dataJobsEnabled;
   }
 
+  public boolean isDataJobsOpenLineageEnabled() {
+    return dataJobsOpenLineageEnabled;
+  }
+
   public String getDataJobsCommandPattern() {
     return dataJobsCommandPattern;
   }

Original file line number	Diff line number	Diff line change
`@@ -41,7 +41,7 @@ public String[] helperClassNames() {`
`41`	`41`
`42`	`42`	`@Override`
`43`	`43`	`public boolean defaultEnabled() {`
`44`		`- return true;`
	`44`	`+ return false;`
`45`	`45`	`}`
`46`	`46`
`47`	`47`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -14,6 +14,7 @@ abstract class AbstractSpark24SqlTest extends AgentTestRunner {`
`14`	`14`	`void configurePreAgent() {`
`15`	`15`	`super.configurePreAgent()`
`16`	`16`	`injectSysConfig("dd.integration.spark.enabled", "true")`
	`17`	`+ injectSysConfig("dd.integration.openlineage-spark.enabled", "true")`
`17`	`18`	`}`
`18`	`19`
`19`	`20`	`static Dataset<Row> generateSampleDataframe(SparkSession spark) {`
Original file line number	Diff line number	Diff line change
`@@ -10,6 +10,7 @@ abstract class AbstractSpark32SqlTest extends AgentTestRunner {`
`10`	`10`	`void configurePreAgent() {`
`11`	`11`	`super.configurePreAgent()`
`12`	`12`	`injectSysConfig("dd.integration.spark.enabled", "true")`
	`13`	`+ injectSysConfig("dd.integration.openlineage-spark.enabled", "true")`
`13`	`14`	`}`
`14`	`15`
`15`	`16`	`def "compute a GROUP BY sql query plan"() {`
Original file line number	Diff line number	Diff line change
`@@ -25,6 +25,7 @@ class AbstractSparkStructuredStreamingTest extends AgentTestRunner {`
`25`	`25`	`void configurePreAgent() {`
`26`	`26`	`super.configurePreAgent()`
`27`	`27`	`injectSysConfig("dd.integration.spark.enabled", "true")`
	`28`	`+ injectSysConfig("dd.integration.openlineage-spark.enabled", "true")`
`28`	`29`	`}`
`29`	`30`
`30`	`31`	`private SparkSession createSparkSession(String appName) {`
Original file line number	Diff line number	Diff line change
`@@ -34,6 +34,7 @@ abstract class AbstractSparkTest extends AgentTestRunner {`
`34`	`34`	`void configurePreAgent() {`
`35`	`35`	`super.configurePreAgent()`
`36`	`36`	`injectSysConfig("dd.integration.spark.enabled", "true")`
	`37`	`+ injectSysConfig("dd.integration.openlineage-spark.enabled", "true")`
`37`	`38`	`}`
`38`	`39`
`39`	`40`	`def "generate application span with child job and stages"() {`