#AI commit# feat: AISQL支持starrocks切换 (#5363)

Co-authored-by: v-kkhuang <420895376@qq.com>

Author: aiceflower

linkis-computation-governance/linkis-entrance/src/main/java/org/apache/linkis/entrance/persistence/QueryPersistenceEngine.java

@@ -19,7 +19,6 @@
 
 import org.apache.linkis.common.conf.Configuration$;
 import org.apache.linkis.common.exception.ErrorException;
-import org.apache.linkis.common.utils.JsonUtils;
 import org.apache.linkis.entrance.conf.EntranceConfiguration;
 import org.apache.linkis.entrance.exception.EntranceIllegalParamException;
 import org.apache.linkis.entrance.exception.EntranceRPCException;
@@ -47,7 +46,6 @@
 import java.util.List;
 import java.util.Map;
 
-import com.fasterxml.jackson.core.JsonProcessingException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -146,11 +144,13 @@ public void persist(JobRequest jobReq) throws ErrorException {
           JOBREQUEST_NOT_NULL.getErrorCode(), JOBREQUEST_NOT_NULL.getErrorDesc());
     }
     if (logger.isDebugEnabled()) {
-      try {
-        logger.debug("jobReq:" + JsonUtils.jackson().writeValueAsString(jobReq));
-      } catch (JsonProcessingException e) {
-        logger.debug("convert jobReq to string with error:" + e.getMessage());
-      }
+      logger.debug(
+          "Persisting job request - id: {}, reqId: {}, submitUser: {}, executeUser: {}, status: {}",
+          jobReq.getId(),
+          jobReq.getReqId(),
+          jobReq.getSubmitUser(),
+          jobReq.getExecuteUser(),
+          jobReq.getStatus());
     }
 
     JobReqInsert jobReqInsert = new JobReqInsert(jobReq);

linkis-computation-governance/linkis-entrance/src/main/scala/org/apache/linkis/entrance/conf/EntranceConfiguration.scala

@@ -320,7 +320,7 @@ object EntranceConfiguration {
     CommonVars[String]("linkis.ai.sql.hive.template.keys", "hive,mapreduce").getValue
 
   val AI_SQL_CREATORS: String =
-    CommonVars[String]("linkis.ai.sql.support.creators", "IDE,MCP").getValue
+    CommonVars[String]("linkis.ai.sql.support.creators", "IDE,MCP,nodeexecution").getValue
 
   val AI_SQL_KEY: CommonVars[String] =
     CommonVars[String]("linkis.ai.sql.enable", "true")
@@ -369,6 +369,24 @@ object EntranceConfiguration {
 
   val DOCTOR_HTTP_MAX_CONNECT = CommonVars("linkis.aisql.doctor.http.max.connect", 20).getValue
 
+  // StarRocks engine switch configurations
+  val AISQL_STARROCKS_SWITCH = CommonVars("linkis.aisql.starrocks.switch", false)
+
+  val AISQL_DEFAULT_STARROCKS_ENGINE_TYPE =
+    CommonVars("linkis.aisql.default.starrocks.engine.type", "jdbc-4").getValue
+
+  val AISQL_STARROCKS_TEMPLATE_KEYS =
+    CommonVars("linkis.aisql.starrocks.template.keys", "starrocks")
+
+  val AISQL_STARROCKS_DATASOURCE_PREFIX =
+    CommonVars("linkis.aisql.starrocks.datasource.prefix", "starrocks_")
+
+  val AISQL_STARROCKS_WHITELIST_USERS =
+    CommonVars("linkis.aisql.starrocks.whitelist.users", "")
+
+  val AISQL_STARROCKS_WHITELIST_DEPARTMENTS =
+    CommonVars("linkis.aisql.starrocks.whitelist.departments", "")
+
   val SPARK_EXECUTOR_CORES = CommonVars.apply("spark.executor.cores", "2");
 
   var SPARK_EXECUTOR_MEMORY = CommonVars.apply("spark.executor.memory", "6G");

linkis-computation-governance/linkis-entrance/src/main/scala/org/apache/linkis/entrance/interceptor/impl/AISQLTransformInterceptor.scala

@@ -21,6 +21,7 @@ import org.apache.linkis.common.conf.Configuration
 import org.apache.linkis.common.log.LogUtils
 import org.apache.linkis.common.utils.{Logging, Utils}
 import org.apache.linkis.common.utils.CodeAndRunTypeUtils.LANGUAGE_TYPE_AI_SQL
+import org.apache.linkis.datasourcemanager.common.domain.DataSource
 import org.apache.linkis.entrance.conf.EntranceConfiguration
 import org.apache.linkis.entrance.conf.EntranceConfiguration._
 import org.apache.linkis.entrance.interceptor.EntranceInterceptor
@@ -52,6 +53,7 @@ class AISQLTransformInterceptor extends EntranceInterceptor with Logging {
     val sqlLanguage: String = LANGUAGE_TYPE_AI_SQL
     val sparkEngineType: String = AI_SQL_DEFAULT_SPARK_ENGINE_TYPE
     val hiveEngineType: String = AI_SQL_DEFAULT_HIVE_ENGINE_TYPE
+    val starrocksEngineType: String = AISQL_DEFAULT_STARROCKS_ENGINE_TYPE
     val labels: util.List[Label[_]] = jobRequest.getLabels
     val codeType: String = LabelUtil.getCodeType(labels)
     // engineType and creator have been verified in LabelCheckInterceptor.
@@ -105,24 +107,110 @@ class AISQLTransformInterceptor extends EntranceInterceptor with Logging {
         )
       } else {
         logger.info(s"start intelligent selection execution engine for ${jobRequest.getId}")
-        val engineType: String =
-          EntranceUtils.getDynamicEngineType(jobRequest.getExecutionCode, logAppender)
+
+        /**
+         * Check for StarRocks engine switch if feature is enabled Priority: runtime parameters >
+         * script comment > template configuration
+         */
+        var forceEngineType: String = null
+        if (AISQL_STARROCKS_SWITCH.getValue) {
+          // 1. Check runtime parameters
+          val runtimeEngineType = getRuntimeEngineType(jobRequest.getParams)
+
+          // 2. Check script comment
+          val scriptEngineType = if (runtimeEngineType == null) {
+            TemplateConfUtils.getCustomEngineType(
+              jobRequest.getExecutionCode,
+              org.apache.linkis.common.utils.CodeAndRunTypeUtils.getLanguageTypeByCodeType(codeType)
+            )
+          } else null
+
+          // Determine final engine type
+          val targetEngineType = Option(runtimeEngineType)
+            .orElse(Option(scriptEngineType))
+            .orNull
+
+          // If StarRocks engine is specified
+          if ("starrocks".equalsIgnoreCase(targetEngineType)) {
+            // Check whitelist
+            val (user, creator) = LabelUtil.getUserCreator(jobRequest.getLabels)
+            if (!isUserInStarRocksWhitelist(user)) {
+              logger.warn(
+                s"User $user is not in StarRocks whitelist for task ${jobRequest.getId}, using default engine selection"
+              )
+              logAppender.append(
+                LogUtils.generateWarn(
+                  s"User $user is not in StarRocks whitelist, using default engine selection\n"
+                )
+              )
+            } else {
+              forceEngineType = "starrocks"
+            }
+          }
+        }
+        val engineType: String = {
+          EntranceUtils.getDynamicEngineType(
+            jobRequest.getExecutionCode,
+            logAppender,
+            forceEngineType
+          )
+        }
         if ("hive".equals(engineType)) {
           changeEngineLabel(hiveEngineType, labels)
-          logAppender.append(
-            LogUtils.generateInfo(s"use $hiveEngineType by intelligent selection.\n")
-          )
           currentEngineType = hiveEngineType
+        } else if ("starrocks".equals(engineType)) {
+
+          Utils.tryCatch {
+            val dataSource: DataSource = EntranceUtils.getDatasourceByDatasourceTypeAndUser(
+              "starrocks",
+              jobRequest.getSubmitUser,
+              jobRequest.getExecuteUser
+            )
+            if (dataSource != null) {
+              val dsParams: util.Map[String, AnyRef] = new util.HashMap[String, AnyRef]()
+              dsParams.put("wds.linkis.engine.runtime.datasource", dataSource.getDataSourceName)
+              TaskUtils.addRuntimeMap(jobRequest.getParams, dsParams)
+              changeEngineLabel(starrocksEngineType, labels)
+              currentEngineType = starrocksEngineType
+              logAppender.append(
+                LogUtils.generateInfo(
+                  s"use starrocks engine with datasource name : ${dataSource.getDataSourceName} \n"
+                )
+              )
+            } else {
+              // use hive for datasource not exists
+              changeEngineLabel(hiveEngineType, labels)
+              currentEngineType = hiveEngineType
+              logger.warn(
+                s"Failed to select starrocks engine, for user ${jobRequest.getExecuteUser} datasource not exists."
+              )
+              logAppender.append(
+                LogUtils.generateInfo(
+                  s"Failed to select starrocks engine, ${jobRequest.getExecuteUser} datasource does not exist. now use $currentEngineType \n"
+                )
+              )
+            }
+          } { t =>
+            // use hive for exception
+            changeEngineLabel(hiveEngineType, labels)
+            currentEngineType = hiveEngineType
+            logger.warn("Failed to select starrocks engine: ", t)
+            logAppender.append(
+              LogUtils.generateInfo(
+                s"Failed to select starrocks engine, service exception. now use $currentEngineType \n"
+              )
+            )
+          }
         } else {
           changeEngineLabel(sparkEngineType, labels)
-          logAppender.append(
-            LogUtils.generateInfo(s"use $sparkEngineType by intelligent selection.\n")
-          )
           currentEngineType = sparkEngineType
         }
         logger.info(
           s"end intelligent selection execution engine, and engineType is ${currentEngineType} for ${jobRequest.getId}."
         )
+        logAppender.append(
+          LogUtils.generateInfo(s"use $currentEngineType by intelligent selection.\n")
+        )
         EntranceUtils.dealsparkDynamicConf(jobRequest, logAppender, jobRequest.getParams)
       }
 
@@ -179,4 +267,64 @@ class AISQLTransformInterceptor extends EntranceInterceptor with Logging {
     labels.add(newEngineTypeLabel)
   }
 
+  /**
+   * Get engine type from runtime parameters
+   * @param params
+   *   job request parameters
+   * @return
+   *   engine type, such as "starrocks", null if not found
+   */
+  private def getRuntimeEngineType(params: util.Map[String, AnyRef]): String = {
+    if (params == null) return null
+
+    val runtimeParams = TaskUtils.getRuntimeMap(params)
+    if (runtimeParams == null) return null
+
+    val engineType = runtimeParams.get(TemplateConfUtils.confEngineTypeKey)
+    if (engineType != null) engineType.toString else null
+  }
+
+  /**
+   * Check if user is in StarRocks whitelist
+   * @param submitUser
+   *   the user who submits the task
+   * @return
+   *   true if user is in whitelist or whitelist is empty (allow all users), false otherwise
+   */
+  private def isUserInStarRocksWhitelist(submitUser: String): Boolean = {
+    val whitelistUsers = AISQL_STARROCKS_WHITELIST_USERS.getValue
+    val whitelistDepartments = AISQL_STARROCKS_WHITELIST_DEPARTMENTS.getValue
+
+    // If both whitelists are empty, allow all users
+    if (StringUtils.isBlank(whitelistUsers) && StringUtils.isBlank(whitelistDepartments)) {
+      return true
+    }
+
+    // Check user whitelist
+    if (StringUtils.isNotBlank(whitelistUsers)) {
+      val users = whitelistUsers.split(",").map(_.trim)
+      if (users.contains(submitUser)) {
+        logger.info(s"User $submitUser is in StarRocks whitelist (user)")
+        return true
+      }
+    }
+
+    // Check department whitelist
+    if (StringUtils.isNotBlank(whitelistDepartments)) {
+      val userDepartmentId = EntranceUtils.getUserDepartmentId(submitUser)
+      if (StringUtils.isNotBlank(userDepartmentId)) {
+        val departments = whitelistDepartments.split(",").map(_.trim)
+        if (departments.contains(userDepartmentId)) {
+          logger.info(
+            s"User $submitUser (department: $userDepartmentId) is in StarRocks whitelist (department)"
+          )
+          return true
+        }
+      }
+    }
+
+    logger.warn(s"User $submitUser is not in StarRocks whitelist")
+    false
+  }
+
 }

linkis-computation-governance/linkis-entrance/src/main/scala/org/apache/linkis/entrance/interceptor/impl/Explain.scala

@@ -219,7 +219,7 @@ object SQLExplain extends Explain {
         EntranceErrorCode.EXECUTION_CODE_ISNULL.getDesc
       )
     }
-    logger.debug(s"after sql limit code is ${requestPersistTask.getExecutionCode}")
+    logger.debug(s"SQL limit applied.")
   }
 
   private def findRealSemicolonIndex(tempCode: String): Array[Int] = {

linkis-computation-governance/linkis-entrance/src/main/scala/org/apache/linkis/entrance/interceptor/impl/TemplateConfUtils.scala

@@ -55,6 +55,7 @@ object TemplateConfUtils extends Logging {
 
   val confTemplateNameKey = "ec.resource.name"
   val confFixedEngineConnLabelKey = "ec.fixed.sessionId"
+  val confEngineTypeKey = "ec.engine.type"
 
   /**
    * 按模板uuid缓存模板配置
@@ -233,6 +234,40 @@ object TemplateConfUtils extends Logging {
     templateConfName
   }
 
+  /**
+   * Get user-defined engine type from script code
+   *
+   * @param code
+   *   :code
+   * @param languageType
+   *   :sql,hql,python,scala
+   * @return
+   *   String engine type, such as "starrocks", "spark", "hive", null if not found
+   */
+  def getCustomEngineType(code: String, languageType: String): String = {
+    if (StringUtils.isBlank(code) || StringUtils.isBlank(languageType)) {
+      return null
+    }
+
+    val confPattern = languageType.toLowerCase match {
+      case x if x.contains("python") || x.contains("shell") =>
+        s"##@set\\s+${confEngineTypeKey}\\s*=\\s*([^\\s#]+)".r
+      case x if x.contains("scala") =>
+        s"///@set\\s+${confEngineTypeKey}\\s*=\\s*([^\\s/]+)".r
+      case _ =>
+        s"---@set\\s+${confEngineTypeKey}\\s*=\\s*([^\\s-]+)".r
+    }
+
+    val codeRes = code.replaceAll("\r\n", "\n")
+    confPattern.findFirstMatchIn(codeRes) match {
+      case Some(m) =>
+        val engineType = m.group(1).trim
+        logger.info(s"Get custom engine type: $engineType from script")
+        engineType
+      case None => null
+    }
+  }
+
   def dealWithTemplateConf(jobRequest: JobRequest, logAppender: lang.StringBuilder): JobRequest = {
     jobRequest match {
       case requestPersistTask: JobRequest =>
@@ -245,7 +280,7 @@ object TemplateConfUtils extends Logging {
         // only for Creator:IDE, try to get template conf name from code string. eg:---@set ec.resource.name=xxxx
         val codeType = LabelUtil.getCodeType(jobRequest.getLabels)
         val (user, creator) = LabelUtil.getUserCreator(jobRequest.getLabels)
-        if (EntranceConfiguration.DEFAULT_REQUEST_APPLICATION_NAME.getValue.equals(creator)) {
+        if (EntranceConfiguration.DEFAULT_REQUEST_APPLICATION_NAME.getValue.contains(creator)) {
           templateName = getCustomTemplateConfName(jobRequest, codeType, logAppender)
           if (StringUtils.isNotBlank(templateName)) {
             logAppender.append(