code review fixes, add version tag

Signed-off-by: Maciej Obuchowski <[email protected]>

Author: mobuchowski

dd-java-agent/instrumentation/spark/spark_2.12/src/test/groovy/SparkListenerTest.groovy

@@ -1,12 +1,16 @@
 import datadog.trace.instrumentation.spark.AbstractSparkListenerTest
 import datadog.trace.instrumentation.spark.DatadogSpark212Listener
 import org.apache.spark.SparkConf
-import org.apache.spark.scheduler.SparkListener
 
 class SparkListenerTest extends AbstractSparkListenerTest {
   @Override
-  protected SparkListener getTestDatadogSparkListener() {
+  protected DatadogSpark212Listener getTestDatadogSparkListener() {
     def conf = new SparkConf()
     return new DatadogSpark212Listener(conf, "some_app_id", "some_version")
   }
+
+  @Override
+  protected DatadogSpark212Listener getTestDatadogSparkListener(SparkConf conf) {
+    return new DatadogSpark212Listener(conf, "some_app_id", "some_version")
+  }
 }

dd-java-agent/instrumentation/spark/spark_2.13/src/test/groovy/SparkListenerTest.groovy

@@ -1,12 +1,16 @@
 import datadog.trace.instrumentation.spark.AbstractSparkListenerTest
 import datadog.trace.instrumentation.spark.DatadogSpark213Listener
 import org.apache.spark.SparkConf
-import org.apache.spark.scheduler.SparkListener
 
 class SparkListenerTest extends AbstractSparkListenerTest {
   @Override
-  protected SparkListener getTestDatadogSparkListener() {
+  protected DatadogSpark213Listener getTestDatadogSparkListener() {
     def conf = new SparkConf()
     return new DatadogSpark213Listener(conf, "some_app_id", "some_version")
   }
+
+  @Override
+  protected DatadogSpark213Listener getTestDatadogSparkListener(SparkConf conf) {
+    return new DatadogSpark213Listener(conf, "some_app_id", "some_version")
+  }
 }

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

@@ -168,7 +168,7 @@ public AbstractDatadogSparkListener(SparkConf sparkConf, String appId, String sp
   }
 
   public void setupOpenLineage(DDTraceId traceId) {
-    log.error("Setting up OpenLineage configuration with trace id {}", traceId);
+    log.debug("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");

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

@@ -50,7 +50,7 @@ public static Optional<OpenlineageParentContext> from(SparkConf sparkConf) {
     }
 
     if (!sparkConf.contains(OPENLINEAGE_ROOT_PARENT_RUN_ID)) {
-      log.error("Have parent info, but not root parent info. Can't construct valid trace id.");
+      log.debug("Found parent info, but not root parent info. Can't construct valid trace id.");
       return Optional.empty();
     }
 
@@ -59,7 +59,7 @@ public static Optional<OpenlineageParentContext> from(SparkConf sparkConf) {
     String parentRunId = sparkConf.get(OPENLINEAGE_PARENT_RUN_ID);
 
     if (!UUID.matcher(parentRunId).matches()) {
-      log.error("OpenLineage parent run id is not a valid UUID: {}", parentRunId);
+      log.debug("OpenLineage parent run id is not a valid UUID: {}", parentRunId);
       return Optional.empty();
     }
 
@@ -68,7 +68,7 @@ public static Optional<OpenlineageParentContext> from(SparkConf sparkConf) {
     String rootParentRunId = sparkConf.get(OPENLINEAGE_ROOT_PARENT_RUN_ID);
 
     if (!UUID.matcher(rootParentRunId).matches()) {
-      log.error("OpenLineage root parent run id is not a valid UUID: {}", parentRunId);
+      log.debug("OpenLineage root parent run id is not a valid UUID: {}", parentRunId);
       return Optional.empty();
     }
 

dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkListenerTest.groovy

@@ -4,10 +4,10 @@ import com.datadoghq.sketch.ddsketch.DDSketchProtoBinding
 import com.datadoghq.sketch.ddsketch.proto.DDSketch
 import com.datadoghq.sketch.ddsketch.store.CollapsingLowestDenseStore
 import datadog.trace.agent.test.AgentTestRunner
+import org.apache.spark.SparkConf
 import org.apache.spark.Success$
 import org.apache.spark.executor.TaskMetrics
 import org.apache.spark.scheduler.JobSucceeded$
-import org.apache.spark.scheduler.SparkListener
 import org.apache.spark.scheduler.SparkListenerApplicationEnd
 import org.apache.spark.scheduler.SparkListenerApplicationStart
 import org.apache.spark.scheduler.SparkListenerExecutorAdded
@@ -30,7 +30,8 @@ import scala.collection.JavaConverters
 
 abstract class AbstractSparkListenerTest extends AgentTestRunner {
 
-  protected abstract SparkListener getTestDatadogSparkListener()
+  protected abstract AbstractDatadogSparkListener getTestDatadogSparkListener()
+  protected abstract AbstractDatadogSparkListener getTestDatadogSparkListener(SparkConf conf)
 
   protected applicationStartEvent(time=0L) {
     // Constructor of SparkListenerApplicationStart changed starting spark 3.0
@@ -463,6 +464,22 @@ abstract class AbstractSparkListenerTest extends AgentTestRunner {
     }
   }
 
+  def "sets up OpenLineage trace id properly"() {
+    setup:
+    def conf = new SparkConf()
+    conf.set("spark.openlineage.parentRunId", "ad3b6baa-8d88-3b38-8dbe-f06232249a84")
+    conf.set("spark.openlineage.parentJobNamespace", "default")
+    conf.set("spark.openlineage.parentJobName", "dag-push-to-s3-spark.upload_to_s3")
+    conf.set("spark.openlineage.rootParentRunId", "01964820-5280-7674-b04e-82fbed085f39")
+    conf.set("spark.openlineage.rootParentJobNamespace", "default")
+    conf.set("spark.openlineage.rootParentJobName", "dag-push-to-s3-spark")
+    def listener = getTestDatadogSparkListener(conf)
+
+    expect:
+    listener.onApplicationStart(applicationStartEvent(1000L))
+    assert listener.openLineageSparkConf.get("spark.openlineage.run.tags").contains("13959090542865903119")
+  }
+
   def "test lastJobFailed is not set when job is cancelled"() {
     setup:
     def listener = getTestDatadogSparkListener()