nit

apache · Nov 23, 2023 · 60552b5 · 60552b5
1 parent 37ce142
commit 60552b5
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 11 deletions.
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/util/JdbcUtils.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/util/JdbcUtils.scala
@@ -98,6 +98,12 @@ object JdbcUtils extends Logging {
     }
   }
 
+  def mapResultSet[R](rs: ResultSet)(rowMapper: ResultSet => R): Seq[R] = {
+    val builder = Seq.newBuilder[R]
+    while (rs.next()) builder += rowMapper(rs)
+    builder.result
+  }
+
   def redactPassword(password: Option[String]): String = {
     password match {
       case Some(s) if StringUtils.isNotBlank(s) => s"${"*" * s.length}(length:${s.length})"

diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/HiveEngineTests.scala
@@ -18,11 +18,10 @@
 package org.apache.kyuubi
 
 import scala.collection.mutable.ArrayBuffer
-
 import org.apache.commons.lang3.{JavaVersion, SystemUtils}
-
 import org.apache.kyuubi.operation.HiveJDBCTestHelper
 import org.apache.kyuubi.operation.meta.ResultSetSchemaConstant._
+import org.apache.kyuubi.util.JdbcUtils
 
 /**
  * hive tests disabled for JAVA 11
@@ -230,15 +229,10 @@ trait HiveEngineTests extends HiveJDBCTestHelper {
     withJdbcStatement() { statement =>
       val resultSet = statement.getConnection.getMetaData.getTableTypes
       // Hive3 removes support for INDEX_TABLE
-      val expected = Set("TABLE", "VIEW", "MATERIALIZED_VIEW", "INDEX_TABLE")
-      var tableTypes = Set[String]()
-      while (resultSet.next()) {
-        assert(expected.contains(resultSet.getString(TABLE_TYPE)))
-        tableTypes += resultSet.getString(TABLE_TYPE)
-      }
-      assert(!resultSet.next())
-      // Hive3 removes support for INDEX_TABLE
-      assert(expected.dropWhile(_ == "INDEX_TABLE").size === tableTypes.size)
+      val hive2Expected = Set("TABLE", "VIEW", "MATERIALIZED_VIEW", "INDEX_TABLE")
+      val hive3Expected = Set("TABLE", "VIEW", "MATERIALIZED_VIEW")
+      val tableTypes = JdbcUtils.mapResultSet(resultSet) { rs => rs.getString(TABLE_TYPE) }.toSet
+      assert(tableTypes === hive2Expected || tableTypes === hive3Expected)
     }
   }