diff --git a/src/main/scala/za/co/absa/commons/spark/NonFatalQueryExecutionListenerAdapter.scala b/src/main/scala/za/co/absa/commons/spark/NonFatalQueryExecutionListenerAdapter.scala new file mode 100644 index 0000000..a78e170 --- /dev/null +++ b/src/main/scala/za/co/absa/commons/spark/NonFatalQueryExecutionListenerAdapter.scala @@ -0,0 +1,38 @@ +/* + * Copyright 2021 ABSA Group Limited + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package za.co.absa.commons.spark + +import org.apache.spark.sql.execution.QueryExecution +import org.apache.spark.sql.util.QueryExecutionListener + +import scala.util.control.NonFatal + +/** + * Since Spark 2.4.6 a passed exception can be a wrapper on top of an Error. + * This is a controversial decision that this trait is trying to revert + * + * @see https://issues.apache.org/jira/browse/SPARK-31144 + */ +trait NonFatalQueryExecutionListenerAdapter extends QueryExecutionListener { + + abstract override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = { + val causeOrItself = Option(exception.getCause).getOrElse(exception) + + if (NonFatal(causeOrItself)) + super.onFailure(funcName, qe, exception) + } +} diff --git a/src/test/scala/za/co/absa/commons/spark/NonFatalQueryExecutionListenerAdapterSpec.scala b/src/test/scala/za/co/absa/commons/spark/NonFatalQueryExecutionListenerAdapterSpec.scala new file mode 100644 index 0000000..02ee801 --- /dev/null +++ b/src/test/scala/za/co/absa/commons/spark/NonFatalQueryExecutionListenerAdapterSpec.scala @@ -0,0 +1,70 @@ +/* + * Copyright 2021 ABSA Group Limited + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package za.co.absa.commons.spark + +import org.apache.spark.sql.execution.QueryExecution +import org.apache.spark.sql.util.QueryExecutionListener +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers +import za.co.absa.commons.spark.NonFatalQueryExecutionListenerAdapterSpec.ExceptionCollectingQueryExecutionListener + +class NonFatalQueryExecutionListenerAdapterSpec + extends AnyFlatSpec + with Matchers { + + it should "only allow non-fatal exceptions to pass through" in { + val lsn = new ExceptionCollectingQueryExecutionListener with NonFatalQueryExecutionListenerAdapter + + val ex1 = new Exception("one") + val ex2 = new Exception("two", new Exception("I'm non-fatal")) + + val threadDeath = new ThreadDeath + val threadInterruption = new InterruptedException + val oom = new OutOfMemoryError + + val evilEx1 = new Exception(threadDeath) + val evilEx2 = new Exception(threadInterruption) + val evilEx3 = new Exception(oom) + + lsn.onFailure(null, null, ex1) + lsn.onFailure(null, null, evilEx1) + lsn.onFailure(null, null, evilEx2) + lsn.onFailure(null, null, evilEx3) + lsn.onFailure(null, null, threadInterruption) + lsn.onFailure(null, null, ex2) + + lsn.collectedExecutions should contain theSameElementsInOrderAs Seq(ex1, ex2) + } + +} + +object NonFatalQueryExecutionListenerAdapterSpec { + + private class ExceptionCollectingQueryExecutionListener extends QueryExecutionListener { + private var _exceptions: Vector[Exception] = Vector.empty + + def collectedExecutions: Seq[Exception] = _exceptions + + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = { + this._exceptions :+= exception + } + + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = + throw new UnsupportedOperationException + } + +}