capture SQL failures invisible to the listener bus

Author: aboitreaud

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

@@ -1052,6 +1052,31 @@ private synchronized void onStreamingQueryProgressEvent(
     }
   }
 
+  public synchronized void onSqlFailure(String sqlText, Throwable throwable) {
+    if (isRunningOnDatabricks) return;
+    if (applicationEnded) return;
+
+    initApplicationSpanIfNotInitialized();
+
+    AgentSpan sqlSpan =
+        buildSparkSpan("spark.sql", null)
+            .withTag(DDTags.RESOURCE_NAME, sqlText)
+            .withTag("description", sqlText)
+            .asChildOf(applicationSpan.context())
+            .start();
+
+    sqlSpan.setError(true);
+    sqlSpan.setTag(DDTags.ERROR_TYPE, throwable.getClass().getName());
+    sqlSpan.setTag(DDTags.ERROR_MSG, throwable.getMessage());
+
+    java.io.StringWriter sw = new java.io.StringWriter();
+    throwable.printStackTrace(new java.io.PrintWriter(sw));
+    sqlSpan.setTag(DDTags.ERROR_STACK, sw.toString());
+
+    setDataJobsSamplingPriority(sqlSpan);
+    sqlSpan.finish();
+  }
+
   private void setDataJobsSamplingPriority(AgentSpan span) {
     span.setSamplingPriority(PrioritySampling.USER_KEEP, SamplingMechanism.DATA_JOBS);
   }

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

@@ -5,6 +5,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.nameEndsWith;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
@@ -39,7 +40,9 @@ public String[] knownMatchingTypes() {
       "org.apache.spark.util.Utils",
       "org.apache.spark.util.SparkClassUtils",
       "org.apache.spark.scheduler.LiveListenerBus",
-      "org.apache.spark.sql.execution.SparkPlanInfo$"
+      "org.apache.spark.sql.execution.SparkPlanInfo$",
+      "org.apache.spark.sql.SparkSession",
+      "org.apache.spark.sql.DataFrameReader"
     };
   }
 
@@ -74,6 +77,23 @@ public void methodAdvice(MethodTransformer transformer) {
             .and(takesArgument(0, named("org.apache.spark.scheduler.SparkListenerInterface")))
             .and(isDeclaredBy(named("org.apache.spark.scheduler.LiveListenerBus"))),
         AbstractSparkInstrumentation.class.getName() + "$LiveListenerBusAdvice");
+
+    // SparkSession.sql(String) and SparkSession.table(String) — catch AnalysisException failures
+    // that fire before SparkListenerSQLExecutionStart and are invisible to the listener bus
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("sql").or(named("table")))
+            .and(takesArguments(1))
+            .and(takesArgument(0, String.class))
+            .and(isDeclaredBy(named("org.apache.spark.sql.SparkSession"))),
+        AbstractSparkInstrumentation.class.getName() + "$SparkSqlFailureAdvice");
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("table"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, String.class))
+            .and(isDeclaredBy(named("org.apache.spark.sql.DataFrameReader"))),
+        AbstractSparkInstrumentation.class.getName() + "$SparkSqlFailureAdvice");
   }
 
   public static class PrepareSubmitEnvAdvice {
@@ -122,6 +142,21 @@ public static void enter(@Advice.Argument(1) int exitCode, @Advice.Argument(2) S
     }
   }
 
+  public static class SparkSqlFailureAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void enter(
+        @Advice.Argument(0) String identifier, @Advice.Local("sqlIdentifier") String stored) {
+      stored = identifier;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit(
+        @Advice.Thrown Throwable throwable, @Advice.Local("sqlIdentifier") String identifier) {
+      if (throwable == null || AbstractDatadogSparkListener.listener == null) return;
+      AbstractDatadogSparkListener.listener.onSqlFailure(identifier, throwable);
+    }
+  }
+
   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

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

@@ -1639,6 +1639,38 @@ abstract class AbstractSpark32SqlTest extends InstrumentationSpecification {
     return plan
   }
 
+  def "failed spark.sql call creates an error span"() {
+    def sqlText = "SELECT * FROM non_existing_table"
+    def sparkSession = SparkSession.builder()
+    .config("spark.master", "local[2]")
+    .getOrCreate()
+
+    try {
+      sparkSession.sql(sqlText).show()
+    } catch (Exception e) {
+      // expected
+    }
+    sparkSession.stop()
+
+    expect:
+    assertTraces(1) {
+      trace(2) {
+        span {
+          operationName "spark.application"
+          spanType "spark"
+          errored false
+        }
+        span {
+          operationName "spark.sql"
+          spanType "spark"
+          resourceName sqlText
+          childOf(span(0))
+          errored true
+        }
+      }
+    }
+  }
+
   private static Object normalizeColumnRefs(Object plan) {
     if (plan instanceof String) {
       return plan.replaceAll(/#\d+L?/, '#N').replaceAll(/plan_id=\d+/, 'plan_id=N')