[KYUUBI #3549] Support query id in Flink engine

Flink engine now doesn't support query id. This PR fixes the problem.

- [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible

- [ ] Add screenshots for manual tests if appropriate

- [ ] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request

Closes #3550 from link3280/KYUUBI-3549.

Closes #3549

4f583cb [Paul Lin] Update externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkSQLOperationManager.scala
b8808d7 [Paul Lin] [KYUUBI #3549] Simplify code
a68bc59 [Paul Lin] [KYUUBI #3549] Fix typo in the comments
e6217db [Paul Lin] [KYUUBI #3549] Support query id in Flink engine

Authored-by: Paul Lin <paullin3280@gmail.com>
Signed-off-by: Cheng Pan <chengpan@apache.org>

Author: link3280

externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/ExecuteStatement.scala

@@ -23,6 +23,7 @@ import java.util
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
 
+import org.apache.flink.api.common.JobID
 import org.apache.flink.table.api.ResultKind
 import org.apache.flink.table.client.gateway.TypedResult
 import org.apache.flink.table.data.{GenericArrayData, GenericMapData, RowData}
@@ -54,6 +55,8 @@ class ExecuteStatement(
   private val operationLog: OperationLog =
     OperationLog.createOperationLog(session, getHandle)
 
+  var jobId: Option[JobID] = None
+
   override def getOperationLog: Option[OperationLog] = Option(operationLog)
 
   override protected def beforeRun(): Unit = {
@@ -152,6 +155,7 @@ class ExecuteStatement(
 
   private def runOperation(operation: Operation): Unit = {
     val result = executor.executeOperation(sessionId, operation)
+    jobId = result.getJobClient.asScala.map(_.getJobID)
     result.await()
     resultSet = ResultSet.fromTableResult(result)
   }

externals/kyuubi-flink-sql-engine/src/main/scala/org/apache/kyuubi/engine/flink/operation/FlinkSQLOperationManager.scala

@@ -180,6 +180,13 @@ class FlinkSQLOperationManager extends OperationManager("FlinkSQLOperationManage
   }
 
   override def getQueryId(operation: Operation): String = {
-    throw KyuubiSQLException.featureNotSupported()
+    // return empty string instead of null if there's no query id
+    // otherwise there would be TTransportException
+    operation match {
+      case exec: ExecuteStatement => exec.jobId.map(_.toHexString).getOrElse("")
+      case _: PlanOnlyStatement => ""
+      case _ =>
+        throw new IllegalStateException(s"Unsupported Flink operation class $classOf[operation].")
+    }
   }
 }

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

@@ -19,10 +19,11 @@ package org.apache.kyuubi.engine.flink.operation
 
 import java.nio.file.Files
 import java.sql.DatabaseMetaData
-import java.util.UUID
+import java.util.{Properties, UUID}
 
 import scala.collection.JavaConverters._
 
+import org.apache.flink.api.common.JobID
 import org.apache.flink.table.types.logical.LogicalTypeRoot
 import org.apache.hive.service.rpc.thrift._
 import org.scalatest.concurrent.PatienceConfiguration.Timeout
@@ -34,6 +35,7 @@ import org.apache.kyuubi.engine.flink.FlinkEngineUtils._
 import org.apache.kyuubi.engine.flink.WithFlinkSQLEngine
 import org.apache.kyuubi.engine.flink.result.Constants
 import org.apache.kyuubi.engine.flink.util.TestUserClassLoaderJar
+import org.apache.kyuubi.jdbc.hive.{KyuubiConnection, KyuubiStatement}
 import org.apache.kyuubi.operation.HiveJDBCTestHelper
 import org.apache.kyuubi.operation.meta.ResultSetSchemaConstant._
 import org.apache.kyuubi.service.ServiceState._
@@ -974,4 +976,16 @@ class FlinkOperationSuite extends WithFlinkSQLEngine with HiveJDBCTestHelper {
       assertDefaultDatabase(client, "default_database", true)
     }
   }
+
+  test("get query id") {
+    val conn = new KyuubiConnection(jdbcUrl, new Properties())
+    val stmt = conn.createStatement()
+    stmt.executeQuery("create table tbl_a (a int) with ('connector' = 'blackhole')")
+    assert(stmt.asInstanceOf[KyuubiStatement].getQueryId === null)
+    stmt.executeQuery("insert into tbl_a values (1)")
+    val queryId = stmt.asInstanceOf[KyuubiStatement].getQueryId
+    assert(queryId !== null)
+    // parse the string to check if it's valid Flink job id
+    assert(JobID.fromHexString(queryId) !== null)
+  }
 }