Add the linkis task id to trino in the first line of tsql.

WeDataSphere · Oct 7, 2023 · 5950d61 · 5950d61
1 parent 42f5837
commit 5950d61
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/...rc/main/scala/org/apache/linkis/engineplugin/trino/executor/TrinoEngineConnExecutor.scala b/...rc/main/scala/org/apache/linkis/engineplugin/trino/executor/TrinoEngineConnExecutor.scala
@@ -80,6 +80,7 @@ import scala.collection.JavaConverters._
 import com.google.common.cache.{Cache, CacheBuilder}
 import io.trino.client._
 import okhttp3.OkHttpClient
+import org.apache.linkis.governance.common.utils.JobUtils
 
 class TrinoEngineConnExecutor(override val outputPrintLimit: Int, val id: Int)
     extends ConcurrentComputationExecutor(outputPrintLimit) {
@@ -151,7 +152,7 @@ class TrinoEngineConnExecutor(override val outputPrintLimit: Int, val id: Int)
       code: String
   ): ExecuteResponse = {
     val enableSqlHook = TRINO_SQL_HOOK_ENABLED.getValue
-    val realCode = if (StringUtils.isBlank(code)) {
+    var realCode = if (StringUtils.isBlank(code)) {
       "SELECT 1"
     } else if (enableSqlHook) {
       TrinoSQLHook.preExecuteHook(code.trim)
@@ -161,7 +162,9 @@ class TrinoEngineConnExecutor(override val outputPrintLimit: Int, val id: Int)
 
     TrinoCode.checkCode(realCode)
     logger.info(s"trino client begins to run psql code:\n $realCode")
-
+    val jobId = JobUtils.getJobIdFromMap(engineExecutorContext.getProperties)
+    // 在第一行加taskid，trino接收后做定制化处理
+    realCode = s"--linkis_task_id=$jobId" + "\n" + realCode
     val currentUser = getCurrentUser(engineExecutorContext.getLabels)
     val trinoUser = Optional
       .ofNullable(TRINO_DEFAULT_USER.getValue)