This is an automated email from the ASF dual-hosted git repository.

uros-b pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new f6953f77e613 [SPARK-51264][SQL][JDBC][FOLLOWUP] Close JDBC connection 
after DSv2 truncate
f6953f77e613 is described below

commit f6953f77e61359bffca4e7b11818d05cd724194b
Author: Xinyao Zhang <[email protected]>
AuthorDate: Mon Oct 5 13:20:40 2026 +0000

    [SPARK-51264][SQL][JDBC][FOLLOWUP] Close JDBC connection after DSv2 truncate
    
    ### What changes were proposed in this pull request?
    
    Close the JDBC connection created for a DSv2 truncate operation in 
`JDBCWriteBuilder`.
    Add a regression test that verifies every JDBC connection created during 
the write is closed.
    
    ### Why are the changes needed?
    
    `JDBCWriteBuilder` creates a driver-side connection for truncate, while
    `JdbcUtils.truncateTable` closes only its statement. The connection 
otherwise remains open.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Added a `JDBCV2Suite` regression test for `DataFrameWriterV2` truncate 
connection closure.
    
    Attempted:
    `build/sbt 'sql/testOnly org.apache.spark.sql.jdbc.JDBCV2Suite -- -z 
"DataFrameWriterV2: truncate closes JDBC connections"'`
    
    The local build was blocked while compiling unrelated Parquet test sources 
because
    `org.apache.parquet.column.values.Utils` and
    `org.apache.parquet.column.page.mem.MemPageStore` were unavailable. The new 
test's missing
    `JDBCOptions` import was fixed before that unrelated failure.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Codex
    
    Closes #59075 from 
zhangxinyao88/codex/SPARK-59317-close-dsv2-jdbc-truncate-connection.
    
    Authored-by: Xinyao Zhang <[email protected]>
    Signed-off-by: Uros Bojanic <[email protected]>
---
 .../datasources/v2/jdbc/JDBCWriteBuilder.scala     |  6 +++-
 .../org/apache/spark/sql/jdbc/JDBCV2Suite.scala    | 42 +++++++++++++++++++++-
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git 
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCWriteBuilder.scala
 
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCWriteBuilder.scala
index 7449f66ee020..7b87b7f668ea 100644
--- 
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCWriteBuilder.scala
+++ 
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCWriteBuilder.scala
@@ -40,7 +40,11 @@ case class JDBCWriteBuilder(schema: StructType, options: 
JdbcOptionsInWrite) ext
       if (isTruncate) {
         val dialect = JdbcDialects.get(options.url)
         val conn = dialect.createConnectionFactory(options)(-1)
-        JdbcUtils.truncateTable(conn, options)
+        try {
+          JdbcUtils.truncateTable(conn, options)
+        } finally {
+          conn.close()
+        }
       }
       JdbcUtils.saveTable(data, Some(schema), 
SQLConf.get.caseSensitiveAnalysis, options)
     }
diff --git 
a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala 
b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala
index 8f456c3e57ea..f40188bf422c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala
@@ -19,7 +19,9 @@ package org.apache.spark.sql.jdbc
 
 import java.sql.{Connection, DriverManager}
 import java.util.{HexFormat, Properties}
+import java.util.concurrent.ConcurrentLinkedQueue
 
+import scala.jdk.CollectionConverters._
 import scala.util.control.NonFatal
 
 import 
test.org.apache.spark.sql.connector.catalog.functions.JavaStrLen.JavaStrLenStaticMagic
@@ -35,7 +37,7 @@ import 
org.apache.spark.sql.connector.catalog.functions.{ScalarFunction, Unbound
 import org.apache.spark.sql.connector.catalog.index.SupportsIndex
 import org.apache.spark.sql.connector.expressions.Expression
 import org.apache.spark.sql.execution.{FormattedMode, RowDataSourceScanExec}
-import org.apache.spark.sql.execution.datasources.jdbc.{JDBCDatabaseMetadata, 
JDBCRDD}
+import org.apache.spark.sql.execution.datasources.jdbc.{JDBCDatabaseMetadata, 
JDBCOptions, JDBCRDD}
 import 
org.apache.spark.sql.execution.datasources.v2.{DataSourceV2ScanRelation, 
V1ScanWrapper}
 import org.apache.spark.sql.execution.datasources.v2.jdbc.JDBCTableCatalog
 import org.apache.spark.sql.functions.{abs, acos, asin, atan, atan2, avg, 
ceil, coalesce, cos, cosh, cot, count, count_distinct, degrees, exp, floor, 
lit, log => logarithm, log10, not, pow, radians, round, signum, sin, sinh, 
sqrt, sum, tan, tanh, udf, when}
@@ -1899,6 +1901,23 @@ class JDBCV2Suite extends SharedSparkSession with 
ExplainSuiteHelper {
     }
   }
 
+  test("DataFrameWriterV2: truncate closes JDBC connections") {
+    JdbcDialects.unregisterDialect(h2Dialect)
+    try {
+      ConnectionTrackingH2Dialect.clearConnections()
+      JdbcDialects.registerDialect(ConnectionTrackingH2Dialect)
+      withTable("h2.test.abc") {
+        sql("CREATE TABLE h2.test.abc AS SELECT * FROM h2.test.people")
+        sql("SELECT 'bob' AS NAME, 4 AS 
ID").writeTo("h2.test.abc").overwrite(lit(true))
+        assert(ConnectionTrackingH2Dialect.hasTrackedConnections)
+        assert(ConnectionTrackingH2Dialect.allConnectionsClosed)
+      }
+    } finally {
+      JdbcDialects.unregisterDialect(ConnectionTrackingH2Dialect)
+      JdbcDialects.registerDialect(h2Dialect)
+    }
+  }
+
   private def checkAggregateRemoved(df: DataFrame, removed: Boolean = true): 
Unit = {
     val aggregates = df.queryExecution.optimizedPlan.collectFirst {
       case agg: Aggregate => agg
@@ -3195,3 +3214,24 @@ class JDBCV2Suite extends SharedSparkSession with 
ExplainSuiteHelper {
   }
 
 }
+
+private object ConnectionTrackingH2Dialect extends JdbcDialect {
+  private val connections = new ConcurrentLinkedQueue[Connection]()
+
+  def clearConnections(): Unit = connections.clear()
+
+  def hasTrackedConnections: Boolean = !connections.isEmpty
+
+  def allConnectionsClosed: Boolean = connections.asScala.forall(_.isClosed)
+
+  override def canHandle(url: String): Boolean = H2Dialect().canHandle(url)
+
+  override def createConnectionFactory(options: JDBCOptions): Int => 
Connection = {
+    val connectionFactory = H2Dialect().createConnectionFactory(options)
+    (partitionId: Int) => {
+      val connection = connectionFactory(partitionId)
+      connections.add(connection)
+      connection
+    }
+  }
+}


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to