diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/AggregatedDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/AggregatedDialect.scala index 98486701657cb..ea2ec0ed6efcb 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/AggregatedDialect.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/AggregatedDialect.scala @@ -19,6 +19,9 @@ package org.apache.spark.sql.jdbc import java.sql.{Connection, SQLException} +import org.apache.spark.SparkThrowable +import org.apache.spark.sql.connector.catalog.functions.UnboundFunction +import org.apache.spark.sql.connector.expressions.Expression import org.apache.spark.sql.execution.datasources.jdbc.JDBCOptions import org.apache.spark.sql.types.{DataType, MetadataBuilder} @@ -92,4 +95,49 @@ private class AggregatedDialect(dialects: List[JdbcDialect]) override def beforeFetch(connection: Connection, options: JDBCOptions): Unit = { dialects.head.beforeFetch(connection, options) } + + override def compileValue(value: Any): Any = dialects.head.compileValue(value) + + override def compileExpression(expr: Expression): Option[String] = { + dialects.flatMap(_.compileExpression(expr)).headOption + } + + override def isSupportedFunction(funcName: String): Boolean = { + dialects.exists(_.isSupportedFunction(funcName)) + } + + override def getJdbcSQLQueryBuilder(options: JDBCOptions): JdbcSQLQueryBuilder = { + dialects.head.getJdbcSQLQueryBuilder(options) + } + + override def supportsLimit: Boolean = dialects.head.supportsLimit + + override def supportsOffset: Boolean = dialects.head.supportsOffset + + override def supportsHint: Boolean = dialects.head.supportsHint + + override def supportsJoin: Boolean = dialects.head.supportsJoin + + override def getLimitClause(limit: Integer): String = dialects.head.getLimitClause(limit) + + override def getOffsetClause(offset: Integer): String = dialects.head.getOffsetClause(offset) + + override def classifyException( + e: Throwable, + condition: String, + messageParameters: Map[String, String], + description: String, + isRuntime: Boolean): Throwable with SparkThrowable = { + dialects.head.classifyException(e, condition, messageParameters, description, isRuntime) + } + + override def renameTable(oldTable: String, newTable: String): String = { + dialects.head.renameTable(oldTable, newTable) + } + + override def functions: Seq[(String, UnboundFunction)] = dialects.head.functions + + override def createConnectionFactory(options: JDBCOptions): Int => Connection = { + dialects.head.createConnectionFactory(options) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala index 74f803a5a1731..4491e10ced4c5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala @@ -1162,6 +1162,40 @@ class JDBCSuite extends SharedSparkSession { assert(agg.getSchemaQuery ("Dummy") === "My Dummy Schema") } + test("Aggregated dialects: delegation of compileValue and other methods") { + val customDialect = new JdbcDialect { + override def canHandle(url: String): Boolean = true + override def compileValue(value: Any): Any = value match { + case t: Timestamp => s"{ts '$t'}" + case _ => super.compileValue(value) + } + override def isSupportedFunction(funcName: String): Boolean = funcName == "MY_FUNC" + override def supportsLimit: Boolean = true + override def supportsOffset: Boolean = true + override def getLimitClause(limit: Integer): String = s"FETCH FIRST $limit ROWS ONLY" + override def getOffsetClause(offset: Integer): String = s"OFFSET $offset ROWS" + } + val agg = new AggregatedDialect(List(customDialect, new JdbcDialect { + override def canHandle(url: String): Boolean = true + })) + + // compileValue should delegate to first dialect + val ts = Timestamp.valueOf("2024-01-15 10:30:00") + assert(agg.compileValue(ts) === "{ts '2024-01-15 10:30:00.0'}") + + // isSupportedFunction should return true if any dialect supports it + assert(agg.isSupportedFunction("MY_FUNC")) + assert(!agg.isSupportedFunction("OTHER_FUNC")) + + // supportsLimit/Offset should delegate to first dialect + assert(agg.supportsLimit) + assert(agg.supportsOffset) + + // getLimitClause/getOffsetClause should delegate to first dialect + assert(agg.getLimitClause(10) === "FETCH FIRST 10 ROWS ONLY") + assert(agg.getOffsetClause(5) === "OFFSET 5 ROWS") + } + test("Aggregated dialects: isCascadingTruncateTable") { def genDialect(cascadingTruncateTable: Option[Boolean]): JdbcDialect = new JdbcDialect { override def canHandle(url: String): Boolean = true