Enhancement: Supports setting query_timeout hint. (#60)

Author: yuanoOo

docs/spark-catalog-oceanbase.md

@@ -332,6 +332,13 @@ select * from spark_catalog.default.orders;
                 <td>Int</td>
                 <td>Controls the parallelism level for statistical queries (e.g., COUNT, MIN, MAX) by adding /*+ PARALLEL(N) */ hint to generated SQL.</td>
             </tr>
+            <tr>
+                <td>spark.sql.catalog.your_catalog_name.jdbc.query-timeout-hint-degree</td>
+                <td>否</td>
+                <td style="word-wrap: break-word;">-1</td>
+                <td>Int</td>
+                <td>Control the query timeout by adding /*+ query_timeout(N) */ hint to the generated SQL. This parameter can be used to specify the timeout in microseconds. The default value is -1, which means that the hint is not added.</td>
+            </tr>
             <tr>
                 <td>spark.sql.catalog.your_catalog_name.jdbc.disable-pk-table-use-where-partition</td>
                 <td>No</td>

docs/spark-catalog-oceanbase_cn.md

@@ -331,6 +331,13 @@ select * from spark_catalog.default.orders;
                 <td>Int</td>
                 <td>通过向生成的 SQL 添加 /*+ PARALLEL(N) */ hint 来控制统计查询（例如 COUNT、MIN、MAX）的并行级别。</td>
             </tr>
+            <tr>
+                <td>spark.sql.catalog.your_catalog_name.jdbc.query-timeout-hint-degree</td>
+                <td>否</td>
+                <td style="word-wrap: break-word;">-1</td>
+                <td>Int</td>
+                <td>通过向生成的 SQL 添加 /*+ query_timeout(n) */ hint 来控制查询超时时间。通过该参数可以设置超时时间，单位为微妙。默认为-1，表示不添加该Hint。</td>
+            </tr>
             <tr>
                 <td>spark.sql.catalog.your_catalog_name.jdbc.disable-pk-table-use-where-partition</td>
                 <td>否</td>

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)
 }