uros-b commented on code in PR #59962: URL: https://github.com/apache/spark/pull/59962#discussion_r4454152469
########## sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcWriteExceptionSuite.scala: ########## @@ -0,0 +1,397 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.spark.sql.execution.datasources.jdbc + +import java.sql.{Connection, DatabaseMetaData, PreparedStatement, SQLException} + +import org.mockito.ArgumentMatchers._ +import org.mockito.Mockito._ + +import org.apache.spark.{SparkFunSuite, SparkRuntimeException, SparkSQLException, TaskContext} +import org.apache.spark.sql.{AnalysisException, Row} +import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap +import org.apache.spark.sql.jdbc.{JdbcDialect, JdbcDialects, NoLegacyJDBCError} +import org.apache.spark.sql.types.{IntegerType, StructField, StructType} + +/** Tests destination write classification and propagation of upstream JDBC failures. */ +class JdbcWriteExceptionSuite extends SparkFunSuite { + + private val testJdbcUrl = "jdbc:write:test" + + private def newWriteOptions(extra: Map[String, String] = Map.empty): JdbcOptionsInWrite = { + new JdbcOptionsInWrite(CaseInsensitiveMap(Map( + "url" -> testJdbcUrl, + "dbtable" -> "t", + "driver" -> "org.h2.Driver" + ) ++ extra)) + } + + private def registerDialect(dialect: JdbcDialect): Unit = { + JdbcDialects.unregisterDialect(dialect) + JdbcDialects.registerDialect(dialect) + } + + private def savePartitionWithSourceError(iterator: Iterator[Row]): SQLException = { + val conn = mock(classOf[Connection]) + val stmt = mock(classOf[PreparedStatement]) + when(conn.prepareStatement(anyString())).thenReturn(stmt) + + val dialect = new JdbcDialect with NoLegacyJDBCError { + override def canHandle(url: String): Boolean = url == testJdbcUrl + override def createConnectionFactory(options: JDBCOptions): Int => Connection = _ => conn + } + registerDialect(dialect) + TaskContext.setTaskContext(TaskContext.empty()) + try { + intercept[SQLException] { + JdbcUtils.savePartition( + "t", + iterator, + StructType(Seq(StructField("a", IntegerType))), + "INSERT INTO t (a) VALUES (?)", + batchSize = 1, + dialect, + Connection.TRANSACTION_NONE, + newWriteOptions()) + } + } finally { + TaskContext.unset() + JdbcDialects.unregisterDialect(dialect) + } + } + + test("savePartition does not classify a later source hasNext failure as a JDBC write") { + val sourceFailure = new SQLException("source read failed", "08006") + val nextFailure = new SQLException("source connection detail", "08001") + sourceFailure.setNextException(nextFailure) + val iterator = new Iterator[Row] { + private var readFirstRow = false + override def hasNext: Boolean = { + if (readFirstRow) throw sourceFailure + true + } + override def next(): Row = { + readFirstRow = true + Row(1) + } + } + + assert(savePartitionWithSourceError(iterator) eq sourceFailure) + assert(sourceFailure.getSQLState == "08006") + assert(sourceFailure.getNextException eq nextFailure) + assert(sourceFailure.getCause eq nextFailure) + } + + test("savePartition preserves a classified later source next failure") { + val sourceFailure = new SparkSQLException( + "FAILED_JDBC.CONNECTION", + Map("url" -> "jdbc:source"), + Some("08001")) + val iterator = new Iterator[Row] { + private var rowsRead = 0 + override def hasNext: Boolean = rowsRead < 2 + override def next(): Row = { + rowsRead += 1 + if (rowsRead == 1) Row(1) else throw sourceFailure + } + } + + val error = savePartitionWithSourceError(iterator) + assert(error eq sourceFailure) + checkError( + exception = sourceFailure, + condition = "FAILED_JDBC.CONNECTION", + sqlState = Some("08001"), + parameters = Map("url" -> "jdbc:source")) + } + + test("savePartition classifies JDBC row binding failures") { + val conn = mock(classOf[Connection]) + val stmt = mock(classOf[PreparedStatement]) + when(conn.prepareStatement(anyString())).thenReturn(stmt) + val bindingFailure = new SQLException("binding rejected", "22001") + doThrow(bindingFailure).when(stmt).setInt(anyInt(), anyInt()) + + val dialect = new JdbcDialect with NoLegacyJDBCError { + override def canHandle(url: String): Boolean = url == testJdbcUrl + override def createConnectionFactory(options: JDBCOptions): Int => Connection = _ => conn + } + registerDialect(dialect) + TaskContext.setTaskContext(TaskContext.empty()) + try { + val options = newWriteOptions() + val error = intercept[SparkRuntimeException] { + JdbcUtils.savePartition( + "t", + Iterator(Row(1)), + StructType(Seq(StructField("a", IntegerType))), + "INSERT INTO t (a) VALUES (?)", + batchSize = 1, + dialect, + Connection.TRANSACTION_NONE, + options) + } + checkError( + exception = error, + condition = "FAILED_JDBC.WRITE", + sqlState = Some("HV000"), + parameters = Map("url" -> options.getRedactUrl(), "tableName" -> "t")) + assert(error.getCause eq bindingFailure) + verify(stmt, never()).executeBatch() + } finally { + TaskContext.unset() + JdbcDialects.unregisterDialect(dialect) + } + } + + test("savePartition classifies JDBC batch write failures") { + val conn = mock(classOf[Connection]) + val metadata = mock(classOf[DatabaseMetaData]) + val stmt = mock(classOf[PreparedStatement]) + when(conn.getMetaData).thenReturn(metadata) + when(metadata.supportsTransactions()).thenReturn(false) + when(conn.prepareStatement(anyString())).thenReturn(stmt) + + val batchFailure = new SQLException("batch rejected", "23000", 2627) + val nextFailure = new SQLException("duplicate key detail") + batchFailure.setNextException(nextFailure) + when(stmt.executeBatch()).thenThrow(batchFailure) + + val dialect = new JdbcDialect with NoLegacyJDBCError { + override def canHandle(url: String): Boolean = url == testJdbcUrl + override def createConnectionFactory(options: JDBCOptions): Int => Connection = _ => conn + } + registerDialect(dialect) + val context = TaskContext.empty() + TaskContext.setTaskContext(context) + try { + val options = newWriteOptions() + val error = intercept[SparkRuntimeException] { + JdbcUtils.savePartition( + "t", + Iterator(Row(1)), + StructType(Seq(StructField("a", IntegerType))), + "INSERT INTO t (a) VALUES (?)", + batchSize = 1, + dialect, + Connection.TRANSACTION_NONE, + options) + } + checkError( + exception = error, + condition = "FAILED_JDBC.WRITE", + sqlState = Some("HV000"), + parameters = Map("url" -> options.getRedactUrl(), "tableName" -> "t")) + assert(error.getCause eq batchFailure) + assert(batchFailure.getNextException eq nextFailure) + assert(batchFailure.getCause eq nextFailure) + } finally { + TaskContext.unset() + JdbcDialects.unregisterDialect(dialect) + } + } + + test("savePartition classifies JDBC commit failures") { + val conn = mock(classOf[Connection]) + val metadata = mock(classOf[DatabaseMetaData]) + val stmt = mock(classOf[PreparedStatement]) + when(conn.getMetaData).thenReturn(metadata) + when(metadata.supportsTransactions()).thenReturn(true) + when(metadata.supportsTransactionIsolationLevel(Connection.TRANSACTION_READ_COMMITTED)) + .thenReturn(true) + when(conn.prepareStatement(anyString())).thenReturn(stmt) + + val commitFailure = new SQLException("commit rejected", "40001") + doThrow(commitFailure).when(conn).commit() + val dialect = new JdbcDialect with NoLegacyJDBCError { + override def canHandle(url: String): Boolean = url == testJdbcUrl + override def createConnectionFactory(options: JDBCOptions): Int => Connection = _ => conn + } + registerDialect(dialect) + TaskContext.setTaskContext(TaskContext.empty()) + try { + val options = newWriteOptions() + val error = intercept[SparkRuntimeException] { + JdbcUtils.savePartition( + "t", + Iterator(Row(1)), + StructType(Seq(StructField("a", IntegerType))), + "INSERT INTO t (a) VALUES (?)", + batchSize = 1, + dialect, + Connection.TRANSACTION_READ_COMMITTED, + options) + } + checkError( + exception = error, + condition = "FAILED_JDBC.WRITE", + sqlState = Some("HV000"), + parameters = Map("url" -> options.getRedactUrl(), "tableName" -> "t")) + assert(error.getCause eq commitFailure) + verify(conn).rollback() + } finally { + TaskContext.unset() + JdbcDialects.unregisterDialect(dialect) + } + } + + test("savePartition preserves classification from a legacy JDBC dialect") { Review Comment: This test still asserts url -> "jdbc:" and FAILED_JDBC.UNCLASSIFIED as “preserving legacy classification”. These dialects never classified writes. After the wrap in classifyWriteException, this should expect SparkRuntimeException / FAILED_JDBC.WRITE with the redacted URL and tableName, same as the NoLegacyJDBCError cases above. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala: ########## @@ -722,75 +722,132 @@ object JdbcUtils extends Logging with SQLConfHelper { } val supportsTransactions = finalIsolationLevel != Connection.TRANSACTION_NONE var totalRowCount = 0L + def linkNextException(e: SQLException): Unit = { + val cause = e.getNextException + if (cause != null && e.getCause != cause) { + // If there is no cause already, set 'next exception' as cause. If cause is null, + // it *may* be because no cause was set yet + if (e.getCause == null) { + try { + e.initCause(cause) + } catch { + // Or it may be null because the cause *was* explicitly initialized, to *null*, + // in which case this fails. There is no other way to detect it. + // addSuppressed in this case as well. + case _: IllegalStateException => e.addSuppressed(cause) + } + } else { + e.addSuppressed(cause) + } + } + } + + def sourceHasNext: Boolean = { + try { + iterator.hasNext + } catch { + case e: SQLException => + linkNextException(e) + throw e + } + } + + def sourceNext(): Row = { + try { + iterator.next() + } catch { + case e: SQLException => + linkNextException(e) + throw e + } + } + + def classifyWriteException(e: SQLException): Nothing = { + linkNextException(e) + // Closing the connection on task cancellation can cause a SQLException. Do not + // report that cancellation side effect as a JDBC write failure. + if (TaskContext.get().isInterrupted() || e.isInstanceOf[SparkThrowable]) { + throw e + } + throw dialect.classifyException( Review Comment: savePartition already passes FAILED_JDBC.WRITE, the redacted URL, and the table, but the default 5-arg classifier ignores those arguments and delegates to the deprecated 2-arg method, which hardcodes url -> "jdbc:" and always returns AnalysisException (JdbcDialects.scala:867-889). NoLegacyJDBCError is private[sql], so every URL without a built-in dialect (SQLite, Redshift, MariaDB, ClickHouse, Trino, Hive2, …) and every third-party dialect that does not override the 5-arg method still fails as: [FAILED_JDBC.UNCLASSIFIED] Failed JDBC jdbc: on the operation: Write rows to the table t. DAGScheduler.abortStage unwraps that, so it replaces the old job-aborted SQLException. Those dialects never classified writes. Please keep calling the dialect so a real mapping (for example H2 TABLE_OR_VIEW_ALREADY_EXISTS) still wins. If it returns FAILED_JDBC.UNCLASSIFIED, wrap the original SQLException as SparkRuntimeException with FAILED_JDBC.WRITE and {url, tableName}. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
