Enhancement: Supports setting query_timeout hint.

Author: yuanoOo

spark-connector-oceanbase/spark-connector-oceanbase-base/src/main/java/com/oceanbase/spark/config/OceanBaseConfig.java

@@ -241,6 +241,15 @@ public class OceanBaseConfig extends Config implements Serializable {
                     .intConf()
                     .createWithDefault(1);
 
+    // https://www.oceanbase.com/docs/enterprise-oceanbase-database-cn-10000000000881388
+    public static final ConfigEntry<Integer> JDBC_QUERY_TIMEOUT_HINT_DEGREE =
+            new ConfigBuilder("jdbc.query-timeout-hint-degree")
+                    .doc(
+                            "The SQL statements sent by Spark to OB will automatically carry query_timeout Hint. Set this parameter to specify the timeout in microseconds. The default is -1, which means query_timeout hint is not set.")
+                    .version(ConfigConstants.VERSION_1_2_0)
+                    .intConf()
+                    .createWithDefault(-1);
+
     public static final ConfigEntry<Integer> JDBC_STATISTICS_PARALLEL_HINT_DEGREE =
             new ConfigBuilder("jdbc.statistics-parallel-hint-degree")
                     .doc(
@@ -493,6 +502,10 @@ public Integer getJdbcParallelHintDegree() {
         return get(JDBC_PARALLEL_HINT_DEGREE);
     }
 
+    public Integer getQueryTimeoutHintDegree() {
+        return get(JDBC_QUERY_TIMEOUT_HINT_DEGREE);
+    }
+
     public Integer getJdbcStatsParallelHintDegree() {
         return get(JDBC_STATISTICS_PARALLEL_HINT_DEGREE);
     }

spark-connector-oceanbase/spark-connector-oceanbase-base/src/main/scala/com/oceanbase/spark/reader/v2/OBJdbcReader.scala

@@ -126,10 +126,18 @@ class OBJdbcReader(
         ""
     }
 
-    var hint = s"/*+ PARALLEL(${config.getJdbcParallelHintDegree}) */"
-    if (part.useHiddenPKColumn)
-      hint =
-        s"/*+ PARALLEL(${config.getJdbcParallelHintDegree}), opt_param('hidden_column_visible', 'true') */"
+    val useHiddenPKColumnHint = if (part.useHiddenPKColumn) {
+      s", opt_param('hidden_column_visible', 'true') "
+    } else {
+      ""
+    }
+    val queryTimeoutHint = if (config.getQueryTimeoutHintDegree > 0) {
+      s", query_timeout(${config.getQueryTimeoutHintDegree}) "
+    } else {
+      ""
+    }
+    val hint =
+      s"/*+ PARALLEL(${config.getJdbcParallelHintDegree}) $useHiddenPKColumnHint $queryTimeoutHint */"
 
     s"""
        |SELECT $hint $columnStr FROM ${config.getDbTable} ${part.partitionClause}

spark-connector-oceanbase/spark-connector-oceanbase-base/src/main/scala/com/oceanbase/spark/reader/v2/OBMySQLPartition.scala

@@ -253,7 +253,8 @@ object OBMySQLPartition extends Logging {
     val statement = connection.createStatement()
     val tableName = config.getDbTable
     val sql =
-      s"SELECT /*+ PARALLEL(${config.getJdbcStatsParallelHintDegree}) */ count(1) AS cnt FROM $tableName $partName"
+      s"SELECT /*+ PARALLEL(${config.getJdbcStatsParallelHintDegree}) ${queryTimeoutHint(
+          config)} */ count(1) AS cnt FROM $tableName $partName"
     try {
       val rs = statement.executeQuery(sql)
       if (rs.next())
@@ -386,10 +387,12 @@ object OBMySQLPartition extends Logging {
       priKeyColumnName: String) = {
     val statement = connection.createStatement()
     val tableName = config.getDbTable
-    var hint = s"/*+ PARALLEL(${config.getJdbcStatsParallelHintDegree}) */"
-    if (priKeyColumnName.equals(HIDDEN_PK_INCREMENT))
-      hint =
-        s"/*+ PARALLEL(${config.getJdbcParallelHintDegree}), opt_param('hidden_column_visible', 'true') */"
+    val useHiddenPKColHint =
+      if (priKeyColumnName.equals(HIDDEN_PK_INCREMENT))
+        s", opt_param('hidden_column_visible', 'true') "
+      else EMPTY_STRING
+    val hint =
+      s"/*+ PARALLEL(${config.getJdbcStatsParallelHintDegree}) $useHiddenPKColHint ${queryTimeoutHint(config)} */"
 
     val sql =
       s"""
@@ -599,7 +602,8 @@ object OBMySQLPartition extends Logging {
       priKeyColumnName: String,
       config: OceanBaseConfig): Object = {
     val tableName = config.getDbTable
-    val hint = s"/*+ PARALLEL(${config.getJdbcStatsParallelHintDegree}) */"
+    val hint =
+      s"/*+ PARALLEL(${config.getJdbcStatsParallelHintDegree}) ${queryTimeoutHint(config)} */"
     val sql =
       s"""
               SELECT
@@ -629,7 +633,8 @@ object OBMySQLPartition extends Logging {
       priKeyColumnName: String,
       config: OceanBaseConfig): Object = {
     val tableName = config.getDbTable
-    val hint = s"/*+ PARALLEL(${config.getJdbcStatsParallelHintDegree}) */"
+    val hint =
+      s"/*+ PARALLEL(${config.getJdbcStatsParallelHintDegree}) ${queryTimeoutHint(config)} */"
     val sql =
       s"""
               SELECT $hint
@@ -657,7 +662,8 @@ object OBMySQLPartition extends Logging {
       priKeyColumnName: String) = {
     val statement = conn.createStatement()
     val tableName = config.getDbTable
-    val hint = s"/*+ PARALLEL(${config.getJdbcStatsParallelHintDegree}) */"
+    val hint =
+      s"/*+ PARALLEL(${config.getJdbcStatsParallelHintDegree}) ${queryTimeoutHint(config)} */"
     val sql =
       s"""
               SELECT $hint
@@ -682,6 +688,14 @@ object OBMySQLPartition extends Logging {
       obj1.toString.compareTo(obj2.toString)
   }
 
+  def queryTimeoutHint(config: OceanBaseConfig): String = if (
+    config.getQueryTimeoutHintDegree > 0
+  ) {
+    s", query_timeout(${config.getQueryTimeoutHintDegree}) "
+  } else {
+    ""
+  }
+
   private case class UnevenlyPriKeyTableInfo(count: Long, min: Object, max: Object)
 }