linhongliu-db commented on code in PR #59962:
URL: https://github.com/apache/spark/pull/59962#discussion_r4446666376


##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcTaskInterruptSuite.scala:
##########
@@ -57,6 +57,252 @@ class JdbcTaskInterruptSuite extends SharedSparkSession {
     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") {
+    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)
+    when(stmt.executeBatch()).thenThrow(batchFailure)
+
+    val dialect = new JdbcDialect {
+      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 error = intercept[AnalysisException] {
+        JdbcUtils.savePartition(
+          "t",
+          Iterator(Row(1)),
+          StructType(Seq(StructField("a", IntegerType))),
+          "INSERT INTO t (a) VALUES (?)",
+          batchSize = 1,
+          dialect,
+          Connection.TRANSACTION_NONE,
+          newWriteOptions())
+      }
+      checkError(
+        exception = error,
+        condition = "FAILED_JDBC.UNCLASSIFIED",
+        sqlState = Some("HV000"),
+        parameters = Map("url" -> "jdbc:", "message" -> "Write rows to the 
table."))
+      assert(error.getCause eq batchFailure)
+    } finally {
+      TaskContext.unset()
+      JdbcDialects.unregisterDialect(dialect)
+    }
+  }
+

Review Comment:
   Included the table in the legacy description in bdf667fcb5a, with an 
assertion for the table-bearing FAILED_JDBC.UNCLASSIFIED parameters. I kept 
dispatch through the dialect so existing custom mappings remain authoritative. 
The migration guide and PR description now explicitly describe the legacy 
AnalysisException behavior and where the original vendor diagnostics are 
retained.



##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala:
##########
@@ -46,6 +46,21 @@ import org.apache.spark.util.Utils
 
 private[connect] object ErrorUtils extends Logging {
 
+  private def compatibleErrorClass(error: SparkThrowable): (String, 
Map[String, String]) = {
+    val parameters = Option(error.getMessageParameters)
+      .map(_.asScala.toMap)
+      .getOrElse(Map.empty[String, String])
+    if (error.getCondition == "FAILED_JDBC.WRITE") {
+      // Connect clients have no negotiated error-catalog capability. Keep the 
URL supplied
+      // by the WRITE error; it has already been redacted by JDBCOptions.
+      ("FAILED_JDBC.UNCLASSIFIED", Map(

Review Comment:
   Fixed in bdf667fcb5a: both the gRPC description and FetchErrorDetails 
message are rendered from the projected FAILED_JDBC.UNCLASSIFIED condition and 
parameters. Tests assert both exact messages, plus a control for an unchanged 
JDBC condition and its cause chain.



##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcTaskInterruptSuite.scala:
##########
@@ -57,6 +57,252 @@ class JdbcTaskInterruptSuite extends SharedSparkSession {
     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") {

Review Comment:
   Added the plain SQLException next() regression in bdf667fcb5a. It asserts 
exception identity, original SQLSTATE, getNextException, and the linked cause; 
the existing SparkSQLException case remains as separate coverage.



##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala:
##########
@@ -718,75 +718,127 @@ 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(
+        e,
+        condition = "FAILED_JDBC.WRITE",
+        messageParameters = Map("url" -> options.getRedactUrl(), "tableName" 
-> table),
+        description = "Write rows to the table.",

Review Comment:
   Fixed in bdf667fcb5a: the legacy description includes the target table, and 
the legacy-dialect test asserts it. The condition-based WRITE parameters remain 
the redacted URL and table name.



##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcTaskInterruptSuite.scala:
##########
@@ -57,6 +57,252 @@ class JdbcTaskInterruptSuite extends SharedSparkSession {
     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") {

Review Comment:
   Moved the six classification/source-propagation tests and their helpers into 
JdbcWriteExceptionSuite in bdf667fcb5a. JdbcTaskInterruptSuite now contains 
only interruption tests. The new suite also covers plain source next() failures 
and statement cleanup, and both suites pass.



-- 
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]

Reply via email to