[SQL] Add configurable retry attempts and backoff delay for JDBC connection establishment

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCOptions.scala

@@ -304,6 +304,9 @@ class JdbcOptionsInWrite(
       s"Option '$JDBC_QUERY_STRING' is not applicable while writing.")
 
   val table = parameters(JDBC_TABLE_NAME)
+
+  val connectionRetryAttempts = parameters.get(JDBC_CONNECTION_RETRY_ATTEMPTS).map(_.toInt).getOrElse(0)
+  val connectionRetryDelayMs = parameters.get(JDBC_CONNECTION_RETRY_DELAY_MS).map(_.toLong).getOrElse(1000L)
 }
 
 object JDBCOptions {
@@ -379,4 +382,6 @@ object JDBCOptions {
   val JDBC_PREFER_TIMESTAMP_NTZ = newOption("preferTimestampNTZ")
   val JDBC_PREFER_TIMESTAMP_NANOS = newOption("preferTimestampNanos")
   val JDBC_HINT_STRING = newOption("hint")
+  val JDBC_CONNECTION_RETRY_ATTEMPTS = newOption("connectionRetryAttempts")
+  val JDBC_CONNECTION_RETRY_DELAY_MS = newOption("connectionRetryDelayMs")
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala

@@ -174,7 +174,7 @@ object JDBCRDD extends Logging {
       // these are already quoted in JDBCScanBuilder
       requiredColumns
     }
-    val connectionFactory = dialect.createConnectionFactory(options)
+    val connectionFactory = JdbcUtils.createConnectionFactory(dialect, options)
 
     new JDBCRDD(
       sc,

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcRelationProvider.scala

@@ -57,7 +57,7 @@ class JdbcRelationProvider extends CreatableRelationProvider
     val options = new JdbcOptionsInWrite(parameters)
     val isCaseSensitive = sqlContext.sparkSession.sessionState.conf.caseSensitiveAnalysis
     val dialect = JdbcDialects.get(options.url)
-    val conn = dialect.createConnectionFactory(options)(-1)
+    val conn = JdbcUtils.createConnectionFactory(dialect, options)(-1)
     try {
       val tableExists = JdbcUtils.tableExists(conn, options)
       if (tableExists) {

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala

@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql.execution.datasources.jdbc
 
-import java.sql.{Connection, JDBCType, PreparedStatement, ResultSet, ResultSetMetaData, SQLException}
+import java.sql.{Connection, JDBCType, PreparedStatement, ResultSet, ResultSetMetaData, SQLException, SQLTransientConnectionException}
 import java.time.{Instant, LocalDate}
 import java.util
 
@@ -670,7 +670,7 @@ object JdbcUtils extends Logging with SQLConfHelper {
 
     val outMetrics = TaskContext.get().taskMetrics().outputMetrics
 
-    val conn = dialect.createConnectionFactory(options)(-1)
+    val conn = createConnectionFactory(dialect, options)(-1)
 
     // Close JDBC connection so blocked native reads (e.g. executeBatch) fail instead of
     // ignoring Thread.interrupt(). Listener registered after opening the connection; we don't need
@@ -1250,12 +1250,68 @@ object JdbcUtils extends Logging with SQLConfHelper {
       description = "Failed to connect",
       isRuntime = false
     ) {
-      conn = dialect.createConnectionFactory(options)(-1)
+      conn = createConnectionFactory(dialect, options)(-1)
     }
     try {
       f(conn)
     } finally {
       conn.close()
     }
   }
+
+  /**
+   * Wraps the dialect's connection factory with optional retry logic.
+   * Retry is driven by the connectionRetryAttempts / connectionRetryDelayMs options.
+   */
+  def createConnectionFactory(dialect: JdbcDialect, options: JDBCOptions): Int => Connection = {
+    val rawFactory = dialect.createConnectionFactory(options)
+    (partitionId: Int) => createConnectionWithRetry(rawFactory, partitionId, options)
+  }
+
+  private def isRetryableConnectionException(e: Throwable): Boolean = e match {
+    case _: SQLTransientConnectionException => true
+    case se: SQLException =>
+      val sqlState = Option(se.getSQLState).getOrElse("")
+      sqlState.startsWith("08")
+    case _ => false
+  }
+
+  private def createConnectionWithRetry(
+      rawFactory: Int => Connection,
+      partitionId: Int,
+      options: JDBCOptions): Connection = {
+    val maxRetries = options.connectionRetryAttempts
+    val retryDelayMs = options.connectionRetryDelayMs
+    var attempt = 0
+    var connection: Connection = null
+
+    while (connection == null) {
+      try {
+        connection = rawFactory(partitionId)
+      } catch {
+        case NonFatal(e) if attempt < maxRetries && isRetryableConnectionException(e) =>
+          attempt += 1
+          logWarning(s"JDBC connection attempt failed (attempt $attempt/$maxRetries). Retrying in ${retryDelayMs}ms...", e)
+          val tc = TaskContext.get()
+          if (tc != null && tc.isInterrupted()) {
+            throw e
+          }
+          if (retryDelayMs > 0) {
+            try {
+              Thread.sleep(retryDelayMs)
+            } catch {
+              case _: InterruptedException =>
+                logWarning("Interrupted while sleeping between JDBC connection retries.")
+                if (tc != null && tc.isInterrupted()) {
+                  throw e
+                }
+            }
+          }
+          if (tc != null && tc.isInterrupted()) {
+            throw e
+          }
+      }
+    }
+    connection
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCWriteBuilder.scala

@@ -39,7 +39,7 @@ case class JDBCWriteBuilder(schema: StructType, options: JdbcOptionsInWrite) ext
       // TODO (SPARK-32595): do truncate and append atomically.
       if (isTruncate) {
         val dialect = JdbcDialects.get(options.url)
-        val conn = dialect.createConnectionFactory(options)(-1)
+        val conn = JdbcUtils.createConnectionFactory(dialect, options)(-1)
         JdbcUtils.truncateTable(conn, options)
       }
       JdbcUtils.saveTable(data, Some(schema), SQLConf.get.caseSensitiveAnalysis, options)

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCConnectionRetrySuite.scala

@@ -0,0 +1,160 @@
+/*
+ * 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, SQLException, SQLTransientConnectionException}
+
+import org.mockito.Mockito._
+import org.scalatestplus.mockito.MockitoSugar
+
+import org.apache.spark.{SparkFunSuite, TaskContext, TaskContextImpl}
+import org.apache.spark.sql.jdbc.JdbcDialects
+
+class JDBCConnectionRetrySuite extends SparkFunSuite with MockitoSugar {
+
+  private val url = "jdbc:mock://localhost:1234/test"
+  private val dialect = JdbcDialects.get(url)
+
+  test("No retries by default when connection succeeds on attempt 1") {
+    var attempts = 0
+    val mockConn = mock[Connection]
+    val options = new JDBCOptions(Map("url" -> url, "dbtable" -> "t"))
+    val factory = JdbcUtils.createConnectionFactory(dialect, options)
+
+    // Override raw dialect connection factory for testing
+    val rawFactory: Int => Connection = _ => {
+      attempts += 1
+      mockConn
+    }
+    val conn = rawFactory(-1)
+    assert(attempts === 1)
+    assert(conn === mockConn)
+  }
+
+  test("Retries on SQLTransientConnectionException and succeeds") {
+    var attempts = 0
+    val mockConn = mock[Connection]
+    val options = new JDBCOptions(Map(
+      "url" -> url,
+      "dbtable" -> "t",
+      "connectionRetryAttempts" -> "3",
+      "connectionRetryDelayMs" -> "1"
+    ))
+
+    val rawDialect = new org.apache.spark.sql.jdbc.JdbcDialect {
+      override def canHandle(url: String): Boolean = true
+      override def createConnectionFactory(options: JDBCOptions): Int => Connection = {
+        _ => {
+          attempts += 1
+          if (attempts < 3) {
+            throw new SQLTransientConnectionException("Transient connection failure")
+          }
+          mockConn
+        }
+      }
+    }
+
+    val factory = JdbcUtils.createConnectionFactory(rawDialect, options)
+    val conn = factory(-1)
+    assert(attempts === 3)
+    assert(conn === mockConn)
+  }
+
+  test("Retries on connection SQLState (08001) and succeeds") {
+    var attempts = 0
+    val mockConn = mock[Connection]
+    val options = new JDBCOptions(Map(
+      "url" -> url,
+      "dbtable" -> "t",
+      "connectionRetryAttempts" -> "2",
+      "connectionRetryDelayMs" -> "1"
+    ))
+
+    val rawDialect = new org.apache.spark.sql.jdbc.JdbcDialect {
+      override def canHandle(url: String): Boolean = true
+      override def createConnectionFactory(options: JDBCOptions): Int => Connection = {
+        _ => {
+          attempts += 1
+          if (attempts == 1) {
+            throw new SQLException("Unable to establish connection", "08001")
+          }
+          mockConn
+        }
+      }
+    }
+
+    val factory = JdbcUtils.createConnectionFactory(rawDialect, options)
+    val conn = factory(-1)
+    assert(attempts === 2)
+    assert(conn === mockConn)
+  }
+
+  test("Does NOT retry on non-transient auth error (SQLState 28000)") {
+    var attempts = 0
+    val options = new JDBCOptions(Map(
+      "url" -> url,
+      "dbtable" -> "t",
+      "connectionRetryAttempts" -> "5",
+      "connectionRetryDelayMs" -> "1"
+    ))
+
+    val rawDialect = new org.apache.spark.sql.jdbc.JdbcDialect {
+      override def canHandle(url: String): Boolean = true
+      override def createConnectionFactory(options: JDBCOptions): Int => Connection = {
+        _ => {
+          attempts += 1
+          throw new SQLException("Invalid password", "28000")
+        }
+      }
+    }
+
+    val factory = JdbcUtils.createConnectionFactory(rawDialect, options)
+    val ex = intercept[SQLException] {
+      factory(-1)
+    }
+    assert(attempts === 1)
+    assert(ex.getMessage === "Invalid password")
+  }
+
+  test("Exhausted retries preserves original exception") {
+    var attempts = 0
+    val options = new JDBCOptions(Map(
+      "url" -> url,
+      "dbtable" -> "t",
+      "connectionRetryAttempts" -> "2",
+      "connectionRetryDelayMs" -> "1"
+    ))
+
+    val rawDialect = new org.apache.spark.sql.jdbc.JdbcDialect {
+      override def canHandle(url: String): Boolean = true
+      override def createConnectionFactory(options: JDBCOptions): Int => Connection = {
+        _ => {
+          attempts += 1
+          throw new SQLTransientConnectionException(s"Failure attempt $attempts")
+        }
+      }
+    }
+
+    val factory = JdbcUtils.createConnectionFactory(rawDialect, options)
+    val ex = intercept[SQLTransientConnectionException] {
+      factory(-1)
+    }
+    assert(attempts === 3) // Initial attempt + 2 retries
+    assert(ex.getMessage === "Failure attempt 3")
+  }
+}