You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kyuubi.apache.org by ch...@apache.org on 2023/08/15 08:06:21 UTC
[kyuubi] branch master updated: [KYUUBI #5160] Refactor getNextRowSetInternal to support fetch streaming data
This is an automated email from the ASF dual-hosted git repository.
chengpan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/kyuubi.git
The following commit(s) were added to refs/heads/master by this push:
new 14d0dab69 [KYUUBI #5160] Refactor getNextRowSetInternal to support fetch streaming data
14d0dab69 is described below
commit 14d0dab697c314366750b8b80de057aee0155ffd
Author: Paul Lin <pa...@gmail.com>
AuthorDate: Tue Aug 15 16:06:10 2023 +0800
[KYUUBI #5160] Refactor getNextRowSetInternal to support fetch streaming data
### _Why are the changes needed?_
Currently, `getNextRowSetInternal` returns `TRowSet` which is not friendly to explicit EOS in streaming result fetch.
This PR changes the return type to `TFetchResultsResp` to allow the engines to determine the EOS.
### _How was this patch tested?_
- [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible
- [ ] Add screenshots for manual tests if appropriate
- [x] [Run test](https://kyuubi.readthedocs.io/en/master/contributing/code/testing.html#running-tests) locally before make a pull request
Closes #5160 from link3280/refactor_result.
Closes #5160
09822f2ee [Paul Lin] Fix hasMoreRows missing
c94907e2b [Paul Lin] Explicitly set `resp.setHasMoreRows(false)` for operations
4d193fb1d [Paul Lin] Revert unrelated changes in FlinkOperation
ffd0367b3 [Paul Lin] Refactor getNextRowSetInternal to support fetch streaming data
Authored-by: Paul Lin <pa...@gmail.com>
Signed-off-by: Cheng Pan <ch...@apache.org>
---
.../engine/chat/operation/ChatOperation.scala | 9 ++++++--
.../engine/flink/operation/FlinkOperation.scala | 11 ++++++---
.../engine/hive/operation/HiveOperation.scala | 19 +++++++++++----
.../hive/operation/HiveOperationManager.scala | 4 ++--
.../engine/jdbc/operation/JdbcOperation.scala | 11 ++++++---
.../engine/spark/operation/SparkOperation.scala | 11 ++++++---
.../engine/trino/operation/ExecuteStatement.scala | 11 ++++++---
.../engine/trino/operation/TrinoOperation.scala | 11 ++++++---
.../kyuubi/operation/AbstractOperation.scala | 11 +++++----
.../apache/kyuubi/operation/FetchIterator.scala | 2 +-
.../org/apache/kyuubi/operation/Operation.scala | 4 ++--
.../apache/kyuubi/operation/OperationManager.scala | 10 +++++---
.../kyuubi/service/AbstractBackendService.scala | 2 +-
.../org/apache/kyuubi/service/BackendService.scala | 2 +-
.../apache/kyuubi/service/TFrontendService.scala | 9 +++-----
.../apache/kyuubi/session/AbstractSession.scala | 2 +-
.../scala/org/apache/kyuubi/session/Session.scala | 4 ++--
.../apache/kyuubi/operation/NoopOperation.scala | 11 ++++++---
.../kyuubi/operation/NoopOperationManager.scala | 9 +++++---
.../kyuubi/jdbc/hive/KyuubiQueryResultSet.java | 27 ++++++++++++++--------
.../kyuubi/operation/ExecutedCommandExec.scala | 12 +++++++---
.../operation/KyuubiApplicationOperation.scala | 12 +++++++---
.../apache/kyuubi/operation/KyuubiOperation.scala | 10 ++++++--
.../kyuubi/operation/KyuubiOperationManager.scala | 11 +++++----
.../kyuubi/server/BackendServiceMetric.scala | 7 +++---
.../kyuubi/server/api/v1/OperationsResource.scala | 6 +++--
.../kyuubi/server/mysql/MySQLCommandHandler.scala | 3 ++-
.../org/apache/kyuubi/server/trino/api/Query.scala | 2 +-
.../org/apache/kyuubi/WithKyuubiServerOnYarn.scala | 2 +-
.../server/trino/api/TrinoContextSuite.scala | 4 ++--
30 files changed, 166 insertions(+), 83 deletions(-)
diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ChatOperation.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ChatOperation.scala
index 3ae21aa9f..b0b1806f8 100644
--- a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ChatOperation.scala
+++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/operation/ChatOperation.scala
@@ -31,7 +31,9 @@ abstract class ChatOperation(session: Session) extends AbstractOperation(session
protected lazy val conf: KyuubiConf = session.sessionManager.getConf
- override def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TRowSet = {
+ override def getNextRowSetInternal(
+ order: FetchOrientation,
+ rowSetSize: Int): TFetchResultsResp = {
validateDefaultFetchOrientation(order)
assertState(OperationState.FINISHED)
setHasResultSet(true)
@@ -47,7 +49,10 @@ abstract class ChatOperation(session: Session) extends AbstractOperation(session
val taken = iter.take(rowSetSize)
val resultRowSet = RowSet.toTRowSet(taken.toSeq, 1, getProtocolVersion)
resultRowSet.setStartRowOffset(iter.getPosition)
- resultRowSet
+ val resp = new TFetchResultsResp(OK_STATUS)
+ resp.setResults(resultRowSet)
+ resp.setHasMoreRows(false)
+ resp
}
override def cancel(): Unit = {
diff --git a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperation.scala b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperation.scala
index ef80034d3..5a79d2c0e 100644
--- a/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperation.scala
+++ b/externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkOperation.scala
@@ -25,7 +25,7 @@ import scala.collection.JavaConverters.collectionAsScalaIterableConverter
import org.apache.flink.configuration.Configuration
import org.apache.flink.table.gateway.service.context.SessionContext
import org.apache.flink.table.gateway.service.operation.OperationExecutor
-import org.apache.hive.service.rpc.thrift.{TGetResultSetMetadataResp, TRowSet, TTableSchema}
+import org.apache.hive.service.rpc.thrift.{TFetchResultsResp, TGetResultSetMetadataResp, TTableSchema}
import org.apache.kyuubi.{KyuubiSQLException, Utils}
import org.apache.kyuubi.engine.flink.result.ResultSet
@@ -91,7 +91,9 @@ abstract class FlinkOperation(session: Session) extends AbstractOperation(sessio
resp
}
- override def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TRowSet = {
+ override def getNextRowSetInternal(
+ order: FetchOrientation,
+ rowSetSize: Int): TFetchResultsResp = {
validateDefaultFetchOrientation(order)
assertState(OperationState.FINISHED)
setHasResultSet(true)
@@ -112,7 +114,10 @@ abstract class FlinkOperation(session: Session) extends AbstractOperation(sessio
zoneId,
getProtocolVersion)
resultRowSet.setStartRowOffset(resultSet.getData.getPosition)
- resultRowSet
+ val resp = new TFetchResultsResp(OK_STATUS)
+ resp.setResults(resultRowSet)
+ resp.setHasMoreRows(false)
+ resp
}
override def shouldRunAsync: Boolean = false
diff --git a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/HiveOperation.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/HiveOperation.scala
index c7166356d..9759fa00b 100644
--- a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/HiveOperation.scala
+++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/HiveOperation.scala
@@ -21,7 +21,7 @@ import java.util.concurrent.Future
import org.apache.hive.service.cli.operation.{Operation, OperationManager}
import org.apache.hive.service.cli.session.{HiveSession, SessionManager => HiveSessionManager}
-import org.apache.hive.service.rpc.thrift.{TGetResultSetMetadataResp, TRowSet}
+import org.apache.hive.service.rpc.thrift.{TFetchResultsResp, TGetResultSetMetadataResp}
import org.apache.kyuubi.KyuubiSQLException
import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY
@@ -95,22 +95,31 @@ abstract class HiveOperation(session: Session) extends AbstractOperation(session
resp
}
- override def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TRowSet = {
+ override def getNextRowSetInternal(
+ order: FetchOrientation,
+ rowSetSize: Int): TFetchResultsResp = {
val tOrder = FetchOrientation.toTFetchOrientation(order)
val hiveOrder = org.apache.hive.service.cli.FetchOrientation.getFetchOrientation(tOrder)
val rowSet = internalHiveOperation.getNextRowSet(hiveOrder, rowSetSize)
- rowSet.toTRowSet
+ val resp = new TFetchResultsResp(OK_STATUS)
+ resp.setResults(rowSet.toTRowSet)
+ resp.setHasMoreRows(false)
+ resp
}
- def getOperationLogRowSet(order: FetchOrientation, rowSetSize: Int): TRowSet = {
+ def getOperationLogRowSet(order: FetchOrientation, rowSetSize: Int): TFetchResultsResp = {
val tOrder = FetchOrientation.toTFetchOrientation(order)
val hiveOrder = org.apache.hive.service.cli.FetchOrientation.getFetchOrientation(tOrder)
val handle = internalHiveOperation.getHandle
- delegatedOperationManager.getOperationLogRowSet(
+ val rowSet = delegatedOperationManager.getOperationLogRowSet(
handle,
hiveOrder,
rowSetSize,
hive.getHiveConf).toTRowSet
+ val resp = new TFetchResultsResp(OK_STATUS)
+ resp.setResults(rowSet)
+ resp.setHasMoreRows(false)
+ resp
}
override def isTimedOut: Boolean = internalHiveOperation.isTimedOut(System.currentTimeMillis)
diff --git a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/HiveOperationManager.scala b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/HiveOperationManager.scala
index 0762a2938..4e41e742e 100644
--- a/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/HiveOperationManager.scala
+++ b/externals/kyuubi-hive-sql-engine/src/main/scala/org/apache/kyuubi/engine/hive/operation/HiveOperationManager.scala
@@ -20,7 +20,7 @@ package org.apache.kyuubi.engine.hive.operation
import java.util.List
import org.apache.hadoop.hive.conf.HiveConf.ConfVars
-import org.apache.hive.service.rpc.thrift.TRowSet
+import org.apache.hive.service.rpc.thrift.TFetchResultsResp
import org.apache.kyuubi.config.KyuubiConf._
import org.apache.kyuubi.engine.hive.session.HiveSessionImpl
@@ -154,7 +154,7 @@ class HiveOperationManager() extends OperationManager("HiveOperationManager") {
override def getOperationLogRowSet(
opHandle: OperationHandle,
order: FetchOrientation,
- maxRows: Int): TRowSet = {
+ maxRows: Int): TFetchResultsResp = {
val operation = getOperation(opHandle).asInstanceOf[HiveOperation]
operation.getOperationLogRowSet(order, maxRows)
}
diff --git a/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/operation/JdbcOperation.scala b/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/operation/JdbcOperation.scala
index 1c4e0c57f..2ca173757 100644
--- a/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/operation/JdbcOperation.scala
+++ b/externals/kyuubi-jdbc-engine/src/main/scala/org/apache/kyuubi/engine/jdbc/operation/JdbcOperation.scala
@@ -16,7 +16,7 @@
*/
package org.apache.kyuubi.engine.jdbc.operation
-import org.apache.hive.service.rpc.thrift.{TGetResultSetMetadataResp, TRowSet}
+import org.apache.hive.service.rpc.thrift.{TFetchResultsResp, TGetResultSetMetadataResp, TRowSet}
import org.apache.kyuubi.{KyuubiSQLException, Utils}
import org.apache.kyuubi.config.KyuubiConf
@@ -36,7 +36,9 @@ abstract class JdbcOperation(session: Session) extends AbstractOperation(session
protected lazy val dialect: JdbcDialect = JdbcDialects.get(conf)
- override def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TRowSet = {
+ override def getNextRowSetInternal(
+ order: FetchOrientation,
+ rowSetSize: Int): TFetchResultsResp = {
validateDefaultFetchOrientation(order)
assertState(OperationState.FINISHED)
setHasResultSet(true)
@@ -51,7 +53,10 @@ abstract class JdbcOperation(session: Session) extends AbstractOperation(session
val taken = iter.take(rowSetSize)
val resultRowSet = toTRowSet(taken)
resultRowSet.setStartRowOffset(iter.getPosition)
- resultRowSet
+ val resp = new TFetchResultsResp(OK_STATUS)
+ resp.setResults(resultRowSet)
+ resp.setHasMoreRows(false)
+ resp
}
override def cancel(): Unit = {
diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkOperation.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkOperation.scala
index 5062a4904..782f443df 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkOperation.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkOperation.scala
@@ -20,7 +20,7 @@ package org.apache.kyuubi.engine.spark.operation
import java.io.IOException
import java.time.ZoneId
-import org.apache.hive.service.rpc.thrift.{TGetResultSetMetadataResp, TProgressUpdateResp, TRowSet}
+import org.apache.hive.service.rpc.thrift.{TFetchResultsResp, TGetResultSetMetadataResp, TProgressUpdateResp, TRowSet}
import org.apache.spark.kyuubi.{SparkProgressMonitor, SQLOperationListener}
import org.apache.spark.kyuubi.SparkUtilsHelper.redact
import org.apache.spark.sql.{DataFrame, Row, SparkSession}
@@ -233,7 +233,9 @@ abstract class SparkOperation(session: Session)
resp
}
- override def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TRowSet = {
+ override def getNextRowSetInternal(
+ order: FetchOrientation,
+ rowSetSize: Int): TFetchResultsResp = {
var resultRowSet: TRowSet = null
try {
withLocalProperties {
@@ -264,7 +266,10 @@ abstract class SparkOperation(session: Session)
}
} catch onError(cancel = true)
- resultRowSet
+ val resp = new TFetchResultsResp(OK_STATUS)
+ resp.setResults(resultRowSet)
+ resp.setHasMoreRows(false)
+ resp
}
override def shouldRunAsync: Boolean = false
diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/ExecuteStatement.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/ExecuteStatement.scala
index 405c93b0b..3e7cce80c 100644
--- a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/ExecuteStatement.scala
+++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/ExecuteStatement.scala
@@ -19,7 +19,7 @@ package org.apache.kyuubi.engine.trino.operation
import java.util.concurrent.RejectedExecutionException
-import org.apache.hive.service.rpc.thrift.TRowSet
+import org.apache.hive.service.rpc.thrift.TFetchResultsResp
import org.apache.kyuubi.{KyuubiSQLException, Logging}
import org.apache.kyuubi.engine.trino.TrinoStatement
@@ -82,7 +82,9 @@ class ExecuteStatement(
}
}
- override def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TRowSet = {
+ override def getNextRowSetInternal(
+ order: FetchOrientation,
+ rowSetSize: Int): TFetchResultsResp = {
validateDefaultFetchOrientation(order)
assertState(OperationState.FINISHED)
setHasResultSet(true)
@@ -97,7 +99,10 @@ class ExecuteStatement(
val taken = iter.take(rowSetSize)
val resultRowSet = RowSet.toTRowSet(taken.toList, schema, getProtocolVersion)
resultRowSet.setStartRowOffset(iter.getPosition)
- resultRowSet
+ val fetchResultsResp = new TFetchResultsResp(OK_STATUS)
+ fetchResultsResp.setResults(resultRowSet)
+ fetchResultsResp.setHasMoreRows(false)
+ fetchResultsResp
}
private def executeStatement(trinoStatement: TrinoStatement): Unit = {
diff --git a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperation.scala b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperation.scala
index 20153b84f..11eaa1bc1 100644
--- a/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperation.scala
+++ b/externals/kyuubi-trino-engine/src/main/scala/org/apache/kyuubi/engine/trino/operation/TrinoOperation.scala
@@ -21,7 +21,7 @@ import java.io.IOException
import io.trino.client.Column
import io.trino.client.StatementClient
-import org.apache.hive.service.rpc.thrift.{TGetResultSetMetadataResp, TRowSet}
+import org.apache.hive.service.rpc.thrift.{TFetchResultsResp, TGetResultSetMetadataResp}
import org.apache.kyuubi.KyuubiSQLException
import org.apache.kyuubi.Utils
@@ -54,7 +54,9 @@ abstract class TrinoOperation(session: Session) extends AbstractOperation(sessio
resp
}
- override def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TRowSet = {
+ override def getNextRowSetInternal(
+ order: FetchOrientation,
+ rowSetSize: Int): TFetchResultsResp = {
validateDefaultFetchOrientation(order)
assertState(OperationState.FINISHED)
setHasResultSet(true)
@@ -66,7 +68,10 @@ abstract class TrinoOperation(session: Session) extends AbstractOperation(sessio
val taken = iter.take(rowSetSize)
val resultRowSet = RowSet.toTRowSet(taken.toList, schema, getProtocolVersion)
resultRowSet.setStartRowOffset(iter.getPosition)
- resultRowSet
+ val resp = new TFetchResultsResp(OK_STATUS)
+ resp.setResults(resultRowSet)
+ resp.setHasMoreRows(false)
+ resp
}
override protected def beforeRun(): Unit = {
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/AbstractOperation.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/AbstractOperation.scala
index 94f53111d..0a185b942 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/AbstractOperation.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/AbstractOperation.scala
@@ -23,7 +23,7 @@ import java.util.concurrent.locks.ReentrantLock
import scala.collection.JavaConverters._
import org.apache.commons.lang3.StringUtils
-import org.apache.hive.service.rpc.thrift.{TGetResultSetMetadataResp, TProgressUpdateResp, TProtocolVersion, TRowSet, TStatus, TStatusCode}
+import org.apache.hive.service.rpc.thrift.{TFetchResultsResp, TGetResultSetMetadataResp, TProgressUpdateResp, TProtocolVersion, TStatus, TStatusCode}
import org.apache.kyuubi.{KyuubiSQLException, Logging, Utils}
import org.apache.kyuubi.config.KyuubiConf.OPERATION_IDLE_TIMEOUT
@@ -182,11 +182,12 @@ abstract class AbstractOperation(session: Session) extends Operation with Loggin
override def getResultSetMetadata: TGetResultSetMetadataResp
- def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TRowSet
+ def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TFetchResultsResp
- override def getNextRowSet(order: FetchOrientation, rowSetSize: Int): TRowSet = withLockRequired {
- getNextRowSetInternal(order, rowSetSize)
- }
+ override def getNextRowSet(order: FetchOrientation, rowSetSize: Int): TFetchResultsResp =
+ withLockRequired {
+ getNextRowSetInternal(order, rowSetSize)
+ }
/**
* convert SQL 'like' pattern to a Java regular expression.
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/FetchIterator.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/FetchIterator.scala
index fdada1174..ada155887 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/FetchIterator.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/FetchIterator.scala
@@ -20,7 +20,7 @@ package org.apache.kyuubi.operation
/**
* Borrowed from Apache Spark, see SPARK-33655
*/
-sealed trait FetchIterator[A] extends Iterator[A] {
+trait FetchIterator[A] extends Iterator[A] {
/**
* Begin a fetch block, forward from the current position.
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/Operation.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/Operation.scala
index 6f496c9b8..c20a16f61 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/Operation.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/Operation.scala
@@ -19,7 +19,7 @@ package org.apache.kyuubi.operation
import java.util.concurrent.Future
-import org.apache.hive.service.rpc.thrift.{TGetResultSetMetadataResp, TRowSet}
+import org.apache.hive.service.rpc.thrift.{TFetchResultsResp, TGetResultSetMetadataResp}
import org.apache.kyuubi.operation.FetchOrientation.FetchOrientation
import org.apache.kyuubi.operation.log.OperationLog
@@ -32,7 +32,7 @@ trait Operation {
def close(): Unit
def getResultSetMetadata: TGetResultSetMetadataResp
- def getNextRowSet(order: FetchOrientation, rowSetSize: Int): TRowSet
+ def getNextRowSet(order: FetchOrientation, rowSetSize: Int): TFetchResultsResp
def getSession: Session
def getHandle: OperationHandle
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/OperationManager.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/OperationManager.scala
index 3093bcfbe..38dabcc1a 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/OperationManager.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/operation/OperationManager.scala
@@ -137,18 +137,22 @@ abstract class OperationManager(name: String) extends AbstractService(name) {
final def getOperationNextRowSet(
opHandle: OperationHandle,
order: FetchOrientation,
- maxRows: Int): TRowSet = {
+ maxRows: Int): TFetchResultsResp = {
getOperation(opHandle).getNextRowSet(order, maxRows)
}
def getOperationLogRowSet(
opHandle: OperationHandle,
order: FetchOrientation,
- maxRows: Int): TRowSet = {
+ maxRows: Int): TFetchResultsResp = {
val operationLog = getOperation(opHandle).getOperationLog
- operationLog.map(_.read(order, maxRows)).getOrElse {
+ val rowSet = operationLog.map(_.read(order, maxRows)).getOrElse {
throw KyuubiSQLException(s"$opHandle failed to generate operation log")
}
+ val resp = new TFetchResultsResp(new TStatus(TStatusCode.SUCCESS_STATUS))
+ resp.setResults(rowSet)
+ resp.setHasMoreRows(false)
+ resp
}
final def removeExpiredOperations(handles: Seq[OperationHandle]): Seq[Operation] = synchronized {
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/AbstractBackendService.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/AbstractBackendService.scala
index 171e04901..443b35354 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/AbstractBackendService.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/AbstractBackendService.scala
@@ -201,7 +201,7 @@ abstract class AbstractBackendService(name: String)
operationHandle: OperationHandle,
orientation: FetchOrientation,
maxRows: Int,
- fetchLog: Boolean): TRowSet = {
+ fetchLog: Boolean): TFetchResultsResp = {
maxRowsLimit.foreach(limit =>
if (maxRows > limit) {
throw new IllegalArgumentException(s"Max rows for fetching results " +
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/BackendService.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/BackendService.scala
index 968a94197..85df9024c 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/BackendService.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/BackendService.scala
@@ -101,7 +101,7 @@ trait BackendService {
operationHandle: OperationHandle,
orientation: FetchOrientation,
maxRows: Int,
- fetchLog: Boolean): TRowSet
+ fetchLog: Boolean): TFetchResultsResp
def sessionManager: SessionManager
}
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/TFrontendService.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/TFrontendService.scala
index aa012ab56..7cc23779f 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/service/TFrontendService.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/service/TFrontendService.scala
@@ -520,23 +520,20 @@ abstract class TFrontendService(name: String)
override def FetchResults(req: TFetchResultsReq): TFetchResultsResp = {
debug(req.toString)
- val resp = new TFetchResultsResp
try {
val operationHandle = OperationHandle(req.getOperationHandle)
val orientation = FetchOrientation.getFetchOrientation(req.getOrientation)
// 1 means fetching log
val fetchLog = req.getFetchType == 1
val maxRows = req.getMaxRows.toInt
- val rowSet = be.fetchResults(operationHandle, orientation, maxRows, fetchLog)
- resp.setResults(rowSet)
- resp.setHasMoreRows(false)
- resp.setStatus(OK_STATUS)
+ be.fetchResults(operationHandle, orientation, maxRows, fetchLog)
} catch {
case e: Exception =>
error("Error fetching results: ", e)
+ val resp = new TFetchResultsResp
resp.setStatus(KyuubiSQLException.toTStatus(e))
+ resp
}
- resp
}
protected def notSupportTokenErrorStatus = {
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/session/AbstractSession.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/session/AbstractSession.scala
index 1a8c51ccd..a9e33f5a0 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/session/AbstractSession.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/session/AbstractSession.scala
@@ -233,7 +233,7 @@ abstract class AbstractSession(
operationHandle: OperationHandle,
orientation: FetchOrientation,
maxRows: Int,
- fetchLog: Boolean): TRowSet = {
+ fetchLog: Boolean): TFetchResultsResp = {
if (fetchLog) {
sessionManager.operationManager.getOperationLogRowSet(operationHandle, orientation, maxRows)
} else {
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/session/Session.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/session/Session.scala
index bc9f9a8f6..2cdac9f3a 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/session/Session.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/session/Session.scala
@@ -17,7 +17,7 @@
package org.apache.kyuubi.session
-import org.apache.hive.service.rpc.thrift.{TGetInfoType, TGetInfoValue, TGetResultSetMetadataResp, TProtocolVersion, TRowSet}
+import org.apache.hive.service.rpc.thrift.{TFetchResultsResp, TGetInfoType, TGetInfoValue, TGetResultSetMetadataResp, TProtocolVersion}
import org.apache.kyuubi.operation.FetchOrientation.FetchOrientation
import org.apache.kyuubi.operation.OperationHandle
@@ -91,7 +91,7 @@ trait Session {
operationHandle: OperationHandle,
orientation: FetchOrientation,
maxRows: Int,
- fetchLog: Boolean): TRowSet
+ fetchLog: Boolean): TFetchResultsResp
def closeExpiredOperations(): Unit
}
diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/NoopOperation.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/NoopOperation.scala
index 4093d1fca..c369e00ef 100644
--- a/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/NoopOperation.scala
+++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/NoopOperation.scala
@@ -21,7 +21,7 @@ import java.nio.ByteBuffer
import scala.collection.JavaConverters._
-import org.apache.hive.service.rpc.thrift.{TColumn, TColumnDesc, TGetResultSetMetadataResp, TPrimitiveTypeEntry, TRowSet, TStringColumn, TTableSchema, TTypeDesc, TTypeEntry, TTypeId}
+import org.apache.hive.service.rpc.thrift.{TColumn, TColumnDesc, TFetchResultsResp, TGetResultSetMetadataResp, TPrimitiveTypeEntry, TStringColumn, TTableSchema, TTypeDesc, TTypeEntry, TTypeId}
import org.apache.kyuubi.KyuubiSQLException
import org.apache.kyuubi.operation.FetchOrientation.FetchOrientation
@@ -76,11 +76,16 @@ class NoopOperation(session: Session, shouldFail: Boolean = false)
resp
}
- override def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TRowSet = {
+ override def getNextRowSetInternal(
+ order: FetchOrientation,
+ rowSetSize: Int): TFetchResultsResp = {
val col = TColumn.stringVal(new TStringColumn(Seq(opType).asJava, ByteBuffer.allocate(0)))
val tRowSet = ThriftUtils.newEmptyRowSet
tRowSet.addToColumns(col)
- tRowSet
+ val resp = new TFetchResultsResp(OK_STATUS)
+ resp.setResults(tRowSet)
+ resp.setHasMoreRows(false)
+ resp
}
override def shouldRunAsync: Boolean = false
diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/NoopOperationManager.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/NoopOperationManager.scala
index 455e5d4d2..352aae905 100644
--- a/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/NoopOperationManager.scala
+++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/operation/NoopOperationManager.scala
@@ -20,7 +20,7 @@ package org.apache.kyuubi.operation
import java.nio.ByteBuffer
import java.util
-import org.apache.hive.service.rpc.thrift.{TColumn, TRow, TRowSet, TStringColumn}
+import org.apache.hive.service.rpc.thrift.{TColumn, TFetchResultsResp, TRow, TRowSet, TStatus, TStatusCode, TStringColumn}
import org.apache.kyuubi.operation.FetchOrientation.FetchOrientation
import org.apache.kyuubi.session.Session
@@ -136,13 +136,16 @@ class NoopOperationManager extends OperationManager("noop") {
override def getOperationLogRowSet(
opHandle: OperationHandle,
order: FetchOrientation,
- maxRows: Int): TRowSet = {
+ maxRows: Int): TFetchResultsResp = {
val logs = new util.ArrayList[String]()
logs.add("test")
val tColumn = TColumn.stringVal(new TStringColumn(logs, ByteBuffer.allocate(0)))
val tRow = new TRowSet(0, new util.ArrayList[TRow](logs.size()))
tRow.addToColumns(tColumn)
- tRow
+ val resp = new TFetchResultsResp(new TStatus(TStatusCode.SUCCESS_STATUS))
+ resp.setResults(tRow)
+ resp.setHasMoreRows(false)
+ resp
}
override def getQueryId(operation: Operation): String = {
diff --git a/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiQueryResultSet.java b/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiQueryResultSet.java
index 82ea74a01..242ec7720 100644
--- a/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiQueryResultSet.java
+++ b/kyuubi-hive-jdbc/src/main/java/org/apache/kyuubi/jdbc/hive/KyuubiQueryResultSet.java
@@ -26,6 +26,7 @@ import org.apache.hive.service.rpc.thrift.*;
import org.apache.kyuubi.jdbc.hive.cli.RowSet;
import org.apache.kyuubi.jdbc.hive.cli.RowSetFactory;
import org.apache.kyuubi.jdbc.hive.common.HiveDecimal;
+import org.apache.thrift.TException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -47,6 +48,7 @@ public class KyuubiQueryResultSet extends KyuubiBaseResultSet {
private boolean emptyResultSet = false;
private boolean isScrollable = false;
private boolean fetchFirst = false;
+ private boolean hasMoreToFetch = false;
private final TProtocolVersion protocol;
@@ -317,25 +319,20 @@ public class KyuubiQueryResultSet extends KyuubiBaseResultSet {
try {
TFetchOrientation orientation = TFetchOrientation.FETCH_NEXT;
if (fetchFirst) {
- // If we are asked to start from begining, clear the current fetched resultset
+ // If we are asked to start from beginning, clear the current fetched resultset
orientation = TFetchOrientation.FETCH_FIRST;
fetchedRows = null;
fetchedRowsItr = null;
fetchFirst = false;
}
if (fetchedRows == null || !fetchedRowsItr.hasNext()) {
- TFetchResultsReq fetchReq = new TFetchResultsReq(stmtHandle, orientation, fetchSize);
- TFetchResultsResp fetchResp;
- fetchResp = client.FetchResults(fetchReq);
- Utils.verifySuccessWithInfo(fetchResp.getStatus());
-
- TRowSet results = fetchResp.getResults();
- fetchedRows = RowSetFactory.create(results, protocol);
- fetchedRowsItr = fetchedRows.iterator();
+ fetchResult(orientation);
}
if (fetchedRowsItr.hasNext()) {
row = fetchedRowsItr.next();
+ } else if (hasMoreToFetch) {
+ fetchResult(orientation);
} else {
return false;
}
@@ -350,6 +347,18 @@ public class KyuubiQueryResultSet extends KyuubiBaseResultSet {
return true;
}
+ private void fetchResult(TFetchOrientation orientation) throws SQLException, TException {
+ TFetchResultsReq fetchReq = new TFetchResultsReq(stmtHandle, orientation, fetchSize);
+ TFetchResultsResp fetchResp;
+ fetchResp = client.FetchResults(fetchReq);
+ Utils.verifySuccessWithInfo(fetchResp.getStatus());
+ hasMoreToFetch = fetchResp.isSetHasMoreRows() && fetchResp.isHasMoreRows();
+
+ TRowSet results = fetchResp.getResults();
+ fetchedRows = RowSetFactory.create(results, protocol);
+ fetchedRowsItr = fetchedRows.iterator();
+ }
+
@Override
public ResultSetMetaData getMetaData() throws SQLException {
if (isClosed) {
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecutedCommandExec.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecutedCommandExec.scala
index 93379da4b..70b727e5e 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecutedCommandExec.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecutedCommandExec.scala
@@ -17,7 +17,7 @@
package org.apache.kyuubi.operation
-import org.apache.hive.service.rpc.thrift.{TGetResultSetMetadataResp, TRowSet}
+import org.apache.hive.service.rpc.thrift.{TFetchResultsResp, TGetResultSetMetadataResp}
import org.apache.kyuubi.operation.FetchOrientation.FetchOrientation
import org.apache.kyuubi.operation.log.OperationLog
@@ -67,11 +67,17 @@ class ExecutedCommandExec(
if (!shouldRunAsync) getBackgroundHandle.get()
}
- override def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TRowSet = {
+ override def getNextRowSetInternal(
+ order: FetchOrientation,
+ rowSetSize: Int): TFetchResultsResp = {
validateDefaultFetchOrientation(order)
assertState(OperationState.FINISHED)
setHasResultSet(true)
- command.getNextRowSet(order, rowSetSize, getProtocolVersion)
+ val rowSet = command.getNextRowSet(order, rowSetSize, getProtocolVersion)
+ val resp = new TFetchResultsResp(OK_STATUS)
+ resp.setResults(rowSet)
+ resp.setHasMoreRows(false)
+ resp
}
override def getResultSetMetadata: TGetResultSetMetadataResp = {
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiApplicationOperation.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiApplicationOperation.scala
index 54cfdfe93..93929c59c 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiApplicationOperation.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiApplicationOperation.scala
@@ -22,7 +22,7 @@ import java.util.{ArrayList => JArrayList}
import scala.collection.JavaConverters._
-import org.apache.hive.service.rpc.thrift.{TColumn, TColumnDesc, TGetResultSetMetadataResp, TPrimitiveTypeEntry, TRow, TRowSet, TStringColumn, TTableSchema, TTypeDesc, TTypeEntry, TTypeId}
+import org.apache.hive.service.rpc.thrift.{TColumn, TColumnDesc, TFetchResultsResp, TGetResultSetMetadataResp, TPrimitiveTypeEntry, TRow, TRowSet, TStringColumn, TTableSchema, TTypeDesc, TTypeEntry, TTypeId}
import org.apache.kyuubi.engine.ApplicationInfo
import org.apache.kyuubi.operation.FetchOrientation.FetchOrientation
@@ -54,8 +54,11 @@ abstract class KyuubiApplicationOperation(session: Session) extends KyuubiOperat
resp
}
- override def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TRowSet = {
- applicationInfoMap.map { state =>
+ override def getNextRowSetInternal(
+ order: FetchOrientation,
+ rowSetSize: Int): TFetchResultsResp = {
+ val resp = new TFetchResultsResp(OK_STATUS)
+ val rowSet = applicationInfoMap.map { state =>
val tRow = new TRowSet(0, new JArrayList[TRow](state.size))
Seq(state.keys, state.values.map(Option(_).getOrElse(""))).map(_.toSeq.asJava).foreach {
col =>
@@ -64,5 +67,8 @@ abstract class KyuubiApplicationOperation(session: Session) extends KyuubiOperat
}
tRow
}.getOrElse(ThriftUtils.EMPTY_ROW_SET)
+ resp.setResults(rowSet)
+ resp.setHasMoreRows(false)
+ resp
}
}
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
index e357b7912..b85fa8b02 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala
@@ -179,11 +179,17 @@ abstract class KyuubiOperation(session: Session) extends AbstractOperation(sessi
}
}
- override def getNextRowSetInternal(order: FetchOrientation, rowSetSize: Int): TRowSet = {
+ override def getNextRowSetInternal(
+ order: FetchOrientation,
+ rowSetSize: Int): TFetchResultsResp = {
validateDefaultFetchOrientation(order)
assertState(OperationState.FINISHED)
setHasResultSet(true)
- client.fetchResults(_remoteOpHandle, order, rowSetSize, fetchLog = false)
+ val rowset = client.fetchResults(_remoteOpHandle, order, rowSetSize, fetchLog = false)
+ val resp = new TFetchResultsResp(OK_STATUS)
+ resp.setResults(rowset)
+ resp.setHasMoreRows(false)
+ resp
}
override def shouldRunAsync: Boolean = false
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperationManager.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperationManager.scala
index 3078401b0..8ae9c91f8 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperationManager.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperationManager.scala
@@ -19,7 +19,7 @@ package org.apache.kyuubi.operation
import java.util.concurrent.TimeUnit
-import org.apache.hive.service.rpc.thrift.TRowSet
+import org.apache.hive.service.rpc.thrift.{TFetchResultsResp, TStatus, TStatusCode}
import org.apache.kyuubi.KyuubiSQLException
import org.apache.kyuubi.config.KyuubiConf
@@ -214,11 +214,11 @@ class KyuubiOperationManager private (name: String) extends OperationManager(nam
override def getOperationLogRowSet(
opHandle: OperationHandle,
order: FetchOrientation,
- maxRows: Int): TRowSet = {
-
+ maxRows: Int): TFetchResultsResp = {
+ val resp = new TFetchResultsResp(new TStatus(TStatusCode.SUCCESS_STATUS))
val operation = getOperation(opHandle).asInstanceOf[KyuubiOperation]
val operationLog = operation.getOperationLog
- operationLog match {
+ val rowSet = operationLog match {
case Some(log) => log.read(order, maxRows)
case None =>
val remoteHandle = operation.remoteOpHandle()
@@ -229,6 +229,9 @@ class KyuubiOperationManager private (name: String) extends OperationManager(nam
ThriftUtils.EMPTY_ROW_SET
}
}
+ resp.setResults(rowSet)
+ resp.setHasMoreRows(false)
+ resp
}
override def start(): Unit = synchronized {
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/BackendServiceMetric.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/BackendServiceMetric.scala
index 3a92a5ad0..9da4b78c0 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/BackendServiceMetric.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/BackendServiceMetric.scala
@@ -183,9 +183,10 @@ trait BackendServiceMetric extends BackendService {
operationHandle: OperationHandle,
orientation: FetchOrientation,
maxRows: Int,
- fetchLog: Boolean): TRowSet = {
+ fetchLog: Boolean): TFetchResultsResp = {
MetricsSystem.timerTracing(MetricsConstants.BS_FETCH_RESULTS) {
- val rowSet = super.fetchResults(operationHandle, orientation, maxRows, fetchLog)
+ val fetchResultsResp = super.fetchResults(operationHandle, orientation, maxRows, fetchLog)
+ val rowSet = fetchResultsResp.getResults
// TODO: the statistics are wrong when we enabled the arrow.
val rowsSize =
if (rowSet.getColumnsSize > 0) {
@@ -217,7 +218,7 @@ trait BackendServiceMetric extends BackendService {
operation.increaseFetchResultsCount(rowsSize)
}
- rowSet
+ fetchResultsResp
}
}
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/OperationsResource.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/OperationsResource.scala
index ae5feea52..fdde5bbc5 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/OperationsResource.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/OperationsResource.scala
@@ -145,10 +145,11 @@ private[v1] class OperationsResource extends ApiRequestContext with Logging {
if (fetchOrientation != "FETCH_NEXT" && fetchOrientation != "FETCH_FIRST") {
throw new BadRequestException(s"$fetchOrientation in operation log is not supported")
}
- val rowSet = fe.be.sessionManager.operationManager.getOperationLogRowSet(
+ val fetchResultsResp = fe.be.sessionManager.operationManager.getOperationLogRowSet(
OperationHandle(operationHandleStr),
FetchOrientation.withName(fetchOrientation),
maxRows)
+ val rowSet = fetchResultsResp.getResults
val logRowSet = rowSet.getColumns.get(0).getStringVal.getValues.asScala
new OperationLog(logRowSet.asJava, logRowSet.size)
} catch {
@@ -175,11 +176,12 @@ private[v1] class OperationsResource extends ApiRequestContext with Logging {
@QueryParam("fetchorientation") @DefaultValue("FETCH_NEXT")
fetchOrientation: String): ResultRowSet = {
try {
- val rowSet = fe.be.fetchResults(
+ val fetchResultsResp = fe.be.fetchResults(
OperationHandle(operationHandleStr),
FetchOrientation.withName(fetchOrientation),
maxRows,
fetchLog = false)
+ val rowSet = fetchResultsResp.getResults
val rows = rowSet.getRows.asScala.map(i => {
new Row(i.getColVals.asScala.map(i => {
new Field(
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/mysql/MySQLCommandHandler.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/mysql/MySQLCommandHandler.scala
index 2f574d904..5f7a07f58 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/mysql/MySQLCommandHandler.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/mysql/MySQLCommandHandler.scala
@@ -196,11 +196,12 @@ class MySQLCommandHandler(
.getOrElse(KyuubiSQLException(s"Error operator state ${opStatus.state}"))
}
val resultSetMetadata = be.getResultSetMetadata(opHandle)
- val rowSet = be.fetchResults(
+ val fetchResultResp = be.fetchResults(
opHandle,
FetchOrientation.FETCH_NEXT,
Int.MaxValue,
fetchLog = false)
+ val rowSet = fetchResultResp.getResults
MySQLQueryResult(resultSetMetadata.getSchema, rowSet)
} catch {
case rethrow: Exception =>
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/trino/api/Query.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/trino/api/Query.scala
index 4e768b04a..dc9de4ae2 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/trino/api/Query.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/trino/api/Query.scala
@@ -69,7 +69,7 @@ case class Query(
queryId.operationHandle,
defaultFetchOrientation,
defaultMaxRows,
- false)
+ false).getResults
TrinoContext.createQueryResults(
queryId.getQueryId,
nextUri,
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/WithKyuubiServerOnYarn.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/WithKyuubiServerOnYarn.scala
index 68d95dc80..1826760db 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/WithKyuubiServerOnYarn.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/WithKyuubiServerOnYarn.scala
@@ -151,7 +151,7 @@ class KyuubiOperationYarnClusterSuite extends WithKyuubiServerOnYarn with HiveJD
}
val resultColumns = batchJobSubmissionOp.getNextRowSet(FetchOrientation.FETCH_NEXT, 10)
- .getColumns.asScala
+ .getResults.getColumns.asScala
val keys = resultColumns.head.getStringVal.getValues.asScala
val values = resultColumns.apply(1).getStringVal.getValues.asScala
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/trino/api/TrinoContextSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/trino/api/TrinoContextSuite.scala
index 87c8eda96..6c5a01e46 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/trino/api/TrinoContextSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/trino/api/TrinoContextSuite.scala
@@ -84,7 +84,7 @@ class TrinoContextSuite extends KyuubiFunSuite with RestFrontendTestHelper {
checkOpState(opHandleStr, FINISHED)
val metadataResp = fe.be.getResultSetMetadata(opHandle)
- val tRowSet = fe.be.fetchResults(opHandle, FetchOrientation.FETCH_NEXT, 1000, false)
+ val tRowSet = fe.be.fetchResults(opHandle, FetchOrientation.FETCH_NEXT, 1000, false).getResults
val status = fe.be.getOperationStatus(opHandle, Some(0))
val uri = new URI("sfdsfsdfdsf")
@@ -111,7 +111,7 @@ class TrinoContextSuite extends KyuubiFunSuite with RestFrontendTestHelper {
checkOpState(opHandleStr, FINISHED)
val metadataResp = fe.be.getResultSetMetadata(opHandle)
- val tRowSet = fe.be.fetchResults(opHandle, FetchOrientation.FETCH_NEXT, 1000, false)
+ val tRowSet = fe.be.fetchResults(opHandle, FetchOrientation.FETCH_NEXT, 1000, false).getResults
val status = fe.be.getOperationStatus(opHandle, Some(0))
val uri = new URI("sfdsfsdfdsf")