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]