[GitHub] [incubator-kyuubi] yanghua commented on a change in pull request #1693: [KYUUBI #1681] Extend Flink ExecuteStatement to support more operations

[GitBox <gi...@apache.org>]

yanghua commented on a change in pull request #1693:
URL: https://github.com/apache/incubator-kyuubi/pull/1693#discussion_r780781403



##########
File path: externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/ExecuteStatement.scala
##########
@@ -101,46 +103,58 @@ class ExecuteStatement(
     try {
       setState(OperationState.RUNNING)
 
-      columnInfos = new util.ArrayList[ColumnInfo]
+      columns = new util.ArrayList[Column]
 
       val operation = executor.parseStatement(sessionId, statement)
-      resultDescriptor = executor.executeQuery(sessionId, operation.asInstanceOf[QueryOperation])
-      resultDescriptor.getResultSchema.getColumns.asScala.foreach { column =>
-        columnInfos.add(ColumnInfo.create(column.getName, column.getDataType.getLogicalType))
+      operation match {
+        case queryOperation: QueryOperation => runQueryOperation(queryOperation)
+        case operation: Operation => runOperation(operation)
       }
-
-      val resultID = resultDescriptor.getResultId
-
-      val rows = new ArrayBuffer[Row]()
-      var loop = true
-      while (loop) {
-        Thread.sleep(50) // slow the processing down
-
-        val result = executor.snapshotResult(sessionId, resultID, 2)
-        result.getType match {
-          case TypedResult.ResultType.PAYLOAD =>
-            rows.clear()
-            (1 to result.getPayload).foreach { page =>
-              rows ++= executor.retrieveResultPage(resultID, page).asScala
-            }
-          case TypedResult.ResultType.EOS => loop = false
-          case TypedResult.ResultType.EMPTY =>
-        }
-      }
-
-      resultSet = ResultSet.builder
-        .resultKind(ResultKind.SUCCESS_WITH_CONTENT)
-        .columns(columnInfos)
-        .data(rows.toArray[Row])
-        .build
-      setState(OperationState.FINISHED)
     } catch {
       onError(cancel = true)
     } finally {
       statementTimeoutCleaner.foreach(_.shutdown())
     }
   }
 
+  private def runQueryOperation(operation: QueryOperation): Unit = {
+    resultDescriptor = executor.executeQuery(sessionId, operation)
+
+    val resultID = resultDescriptor.getResultId
+
+    val rows = new ArrayBuffer[Row]()
+    var loop = true
+    while (loop) {
+      Thread.sleep(50) // slow the processing down
+
+      val result = executor.snapshotResult(sessionId, resultID, 2)
+      result.getType match {
+        case TypedResult.ResultType.PAYLOAD =>
+          rows.clear()
+          (1 to result.getPayload).foreach { page =>
+            rows ++= executor.retrieveResultPage(resultID, page).asScala
+          }
+        case TypedResult.ResultType.EOS => loop = false
+        case TypedResult.ResultType.EMPTY =>
+      }
+    }
+
+    resultSet = ResultSet.builder
+      .resultKind(ResultKind.SUCCESS_WITH_CONTENT)
+      .columns(resultDescriptor.getResultSchema.getColumns)
+      .data(rows.toArray[Row])
+      .build
+    setState(OperationState.FINISHED)
+  }
+
+  private def runOperation(operation: Operation): Unit = {
+    val result = executor.executeOperation(sessionId, operation)
+    result.await()
+    columns = result.getResolvedSchema.getColumns

Review comment:
       unnecessary?

##########
File path: externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/ExecuteStatement.scala
##########
@@ -47,7 +49,7 @@ class ExecuteStatement(
 
   private var resultDescriptor: ResultDescriptor = _
 
-  private var columnInfos: util.List[ColumnInfo] = _
+  private var columns: util.List[Column] = _

Review comment:
       unnecessary class level variable?




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@kyuubi.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org