[FLINK] Skip the max-rows limitation if the statement is not DQL

Author: link3280

externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/result/QueryResultFetchIterator.scala renamed to externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/result/IncrementalResultFetchIterator.scala

@@ -36,8 +36,9 @@ import org.apache.flink.types.Row
 import org.apache.kyuubi.Logging
 import org.apache.kyuubi.engine.flink.shim.FlinkResultSet
 import org.apache.kyuubi.operation.FetchIterator
+import org.apache.kyuubi.util.reflect.ReflectUtils.getField
 
-class QueryResultFetchIterator(
+class IncrementalResultFetchIterator(
     resultFetcher: ResultFetcher,
     maxRows: Int = 1000000,
     resultFetchTimeout: Duration = Duration.Inf) extends FetchIterator[Row] with Logging {
@@ -58,6 +59,10 @@ class QueryResultFetchIterator(
 
   val FETCH_INTERVAL_MS: Long = 1000
 
+  val isQueryResult: Boolean = getField(resultFetcher, "isQueryResult")
+
+  val effectiveMaxRows: Int = if (isQueryResult) maxRows else Int.MaxValue
+
   private val executor = Executors.newSingleThreadScheduledExecutor(
     new ThreadFactoryBuilder().setNameFormat("flink-query-iterator-%d").setDaemon(true).build)
 
@@ -78,7 +83,7 @@ class QueryResultFetchIterator(
       // if no timeout is set, this would block until some rows are fetched
       debug(s"Fetching from result store with timeout $resultFetchTimeout ms")
       while (!fetched && !Thread.interrupted()) {
-        val rs = resultFetcher.fetchResults(token, maxRows - bufferedRows.length)
+        val rs = resultFetcher.fetchResults(token, effectiveMaxRows - bufferedRows.length)
         val flinkRs = new FlinkResultSet(rs)
         // TODO: replace string-based match when Flink 1.16 support is dropped
         flinkRs.getResultType.name() match {
@@ -144,7 +149,7 @@ class QueryResultFetchIterator(
       debug(s"Fetching from buffered rows at pos $pos.")
       val row = bufferedRows(pos.toInt)
       pos += 1
-      if (pos >= maxRows) {
+      if (pos >= effectiveMaxRows) {
         hasNext = false
       }
       row
@@ -154,7 +159,7 @@ class QueryResultFetchIterator(
       if (hasNext) {
         val row = bufferedRows(pos.toInt)
         pos += 1
-        if (pos >= maxRows) {
+        if (pos >= effectiveMaxRows) {
           hasNext = false
         }
         row

externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/result/ResultSet.scala

@@ -53,7 +53,7 @@ case class ResultSet(
 
   def close: Unit = {
     data match {
-      case queryIte: QueryResultFetchIterator => queryIte.close()
+      case incIte: IncrementalResultFetchIterator => incIte.close()
       case _ =>
     }
   }

externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/result/ResultSetUtil.scala

@@ -73,7 +73,7 @@ object ResultSetUtil {
       throw new IllegalArgumentException("maxRows should be positive")
     }
     val schema = resultFetcher.getResultSchema
-    val ite = new QueryResultFetchIterator(resultFetcher, maxRows, resultFetchTimeout)
+    val ite = new IncrementalResultFetchIterator(resultFetcher, maxRows, resultFetchTimeout)
     ResultSet.builder
       .resultKind(ResultKind.SUCCESS_WITH_CONTENT)
       .columns(schema.getColumns)

externals/kyuubi-flink-sql-engine/src/test/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperationSuite.scala

@@ -1148,6 +1148,20 @@ abstract class FlinkOperationSuite extends HiveJDBCTestHelper with WithFlinkTest
         assert(rows === 200)
       }
     }
+    withSessionConf()(Map(ENGINE_FLINK_MAX_ROWS.key -> "10"))(Map.empty) {
+      withJdbcStatement() { statement =>
+        for (i <- 0 to 10) {
+          statement.execute(s"create table tbl_src$i (a bigint) " +
+            s"with ('connector' = 'blackhole')")
+        }
+        val resultSet = statement.executeQuery("show tables")
+        var rows = 0
+        while (resultSet.next()) {
+          rows += 1
+        }
+        assert(rows === 11)
+      }
+    }
   }
 
   test("execute statement - add/show jar") {