[Feature][udf] Load specific udfs as we needed

linkis-computation-governance/linkis-engineconn-manager/linkis-engineconn-manager-server/src/main/scala/org/apache/linkis/ecm/server/hook/JarUDFLoadECMHook.scala

@@ -56,6 +56,24 @@ class JarUDFLoadECMHook extends ECMHook with Logging {
         logger.info("start loading UDFs")
         val user = pel.user
         val ticketId = pel.ticketId
+
+        val udfAllLoad =
+          pel
+            .creationDesc
+            .properties
+            .getOrDefault("linkis.user.udf.all.load", "true")
+            .toBoolean
+        val udfIdStr =
+          pel
+            .creationDesc
+            .properties
+            .getOrDefault("linkis.user.udf.custom.ids", "")
+        val udfIds =
+          udfIdStr
+            .split(",")
+            .filter(StringUtils.isNotBlank)
+            .map(s => s.toLong)
+
         val engineType = LabelUtil.getEngineType(pel.labels)
         val workDir = localDirsHandleService.getEngineConnWorkDir(user, ticketId, engineType)
         val pubDir = localDirsHandleService.getEngineConnPublicDir
@@ -66,7 +84,11 @@ class JarUDFLoadECMHook extends ECMHook with Logging {
           fs.setPermission(fsPath, "rwxrwxrwx")
         }
 
-        val udfInfos = UDFClient.getJarUdf(pel.user)
+        val udfInfos =
+          if (udfAllLoad)
+            UDFClient.getJarUdf(pel.user)
+          else
+            UDFClient.getJarUdfByIds(pel.user, udfIds)
         val fileNameSet: mutable.HashSet[String] = new mutable.HashSet[String]()
         import util.control.Breaks._
         udfInfos.foreach { udfInfo =>

linkis-computation-governance/linkis-engineconn/linkis-computation-engineconn/src/main/scala/org/apache/linkis/engineconn/computation/executor/hook/UDFLoadEngineConnHook.scala

@@ -37,7 +37,6 @@ import org.apache.linkis.manager.label.entity.engine.RunType.RunType
 import org.apache.linkis.udf.UDFClient
 import org.apache.linkis.udf.utils.ConstantVar
 import org.apache.linkis.udf.vo.UDFInfoVo
-
 import org.apache.commons.io.{FileUtils, IOUtils}
 import org.apache.commons.lang3.StringUtils
 
@@ -58,11 +57,11 @@ abstract class UDFLoad extends Logging {
 
   protected def constructCode(udfInfo: UDFInfoVo): String
 
-  protected def generateCode(user: String): Array[String] = {
+  protected def generateCode(user: String, udfAllLoad: Boolean, udfIds: Array[Long]): Array[String] = {
     val codeBuffer = new ArrayBuffer[String]
     val statementBuffer = new ArrayBuffer[String]
     var accept = true
-    getLoadUdfCode(user).split("\n").foreach {
+    getLoadUdfCode(user, udfAllLoad, udfIds).split("\n").foreach {
       case "" =>
       case l if l.startsWith("%") =>
         if (acceptCodeType(l)) {
@@ -83,12 +82,17 @@ abstract class UDFLoad extends Logging {
     line.startsWith("%" + runType.toString)
   }
 
-  protected def getLoadUdfCode(user: String): String = {
-    logger.info("start loading UDFs")
-    val udfInfos = UDFClient.getUdfInfosByUdfType(user, category, udfType)
+  protected def getLoadUdfCode(user: String, udfAllLoad: Boolean, udfIds: Array[Long]): String = {
+    logger.info(s"start loading UDFs, load all: $udfAllLoad, udfIds: ${udfIds.mkString("Array(", ", ", ")")}")
+    val udfInfos = if (udfAllLoad) {
+      UDFClient.getUdfInfosByUdfType(user, category, udfType)
+    } else {
+      UDFClient.getUdfInfosByUdfIds(user, udfIds, category, udfType)
+    }
     logger.info("all udfs: ")
+
     udfInfos.foreach { l =>
-      logger.info("udfName:" + l.getUdfName + " bml_resource_id:" + l.getBmlResourceId + "\n")
+      logger.info(s"udfName:${l.getUdfName}, bml_resource_id:${l.getBmlResourceId}, bml_id:${l.getId}\n")
     }
     udfInfos
       .filter { info => StringUtils.isNotEmpty(info.getBmlResourceId) }
@@ -126,19 +130,20 @@ abstract class UDFLoad extends Logging {
     }
   }
 
-  private def getFunctionCode(): Array[String] = {
+  private def getFunctionCode(udfAllLoad: Boolean, udfIds: Array[Long]): Array[String] = {
     val engineCreationContext =
       EngineConnManager.getEngineConnManager.getEngineConn.getEngineCreationContext
     val user = engineCreationContext.getUser
-    Utils.tryCatch(generateCode(user)) { case t: Throwable =>
-      if (!ComputationExecutorConf.UDF_LOAD_FAILED_IGNORE.getValue) {
-        logger.error("Failed to load function, executor close ")
-        throw t
-      } else {
-        logger.error("Failed to load function", t)
-        Array.empty[String]
+    Utils.tryCatch(generateCode(user, udfAllLoad, udfIds)) {
+      case t: Throwable =>
+        if (!ComputationExecutorConf.UDF_LOAD_FAILED_IGNORE.getValue) {
+          logger.error("Failed to load function, executor close ")
+          throw t
+        } else {
+          logger.error("Failed to load function", t)
+          Array.empty[String]
+        }
       }
-    }
   }
 
   private def executeFunctionCode(codes: Array[String], executor: ComputationExecutor): Unit = {
@@ -163,7 +168,7 @@ abstract class UDFLoad extends Logging {
 
   protected def loadFunctions(executor: Executor): Unit = {
 
-    val codes = getFunctionCode()
+    val codes = getFunctionCode(true, null)
     if (null != codes && codes.nonEmpty) {
       executor match {
         case computationExecutor: ComputationExecutor =>
@@ -174,9 +179,9 @@ abstract class UDFLoad extends Logging {
     logger.info(s"Successful to execute function code ${runType}, type : ${udfType}")
   }
 
-  protected def loadUDF(labels: Array[Label[_]]): Unit = {
+  protected def loadUDF(labels: Array[Label[_]], udfAllLoad: Boolean, udfIds: Array[Long]): Unit = {
 
-    val codes = getFunctionCode()
+    val codes = getFunctionCode(udfAllLoad, udfIds)
     if (null != codes && codes.nonEmpty) {
       val executor = ExecutorManager.getInstance.getExecutorByLabels(labels)
       executor match {
@@ -206,8 +211,24 @@ abstract class UDFLoadEngineConnHook extends UDFLoad with EngineConnHook with Lo
         codeLanguageLabel.setCodeType(runType.toString)
         logger.warn("no EngineTypeLabel found, use default runType")
     }
+
+    val udfAllLoad =
+      engineCreationContext
+        .getOptions
+        .getOrDefault("linkis.user.udf.all.load", "true")
+        .toBoolean
+    val udfIdStr =
+      engineCreationContext
+        .getOptions
+        .getOrDefault("linkis.user.udf.custom.ids", "")
+    val udfIds =
+      udfIdStr
+        .split(",")
+        .filter(StringUtils.isNotBlank)
+        .map(s => s.toLong)
+
     val labels = Array[Label[_]](codeLanguageLabel)
-    loadUDF(labels)
+    loadUDF(labels, udfAllLoad, udfIds)
   }
 
   override def afterEngineServerStartFailed(

linkis-public-enhancements/linkis-udf/linkis-udf-client/src/main/scala/org/apache/linkis/udf/UDFClient.scala

@@ -18,13 +18,19 @@
 package org.apache.linkis.udf
 
 import org.apache.linkis.rpc.Sender
-import org.apache.linkis.udf.api.rpc.{RequestUdfTree, ResponseUdfTree}
-import org.apache.linkis.udf.entity.{UDFInfo, UDFTree}
+import org.apache.linkis.udf.api.rpc.{
+  RequestUdfIds,
+  RequestUdfTree,
+  ResponseUdfTree,
+  ResponseUdfs
+}
+import org.apache.linkis.udf.entity.UDFTree
 import org.apache.linkis.udf.utils.ConstantVar
 import org.apache.linkis.udf.vo.UDFInfoVo
 
 import org.apache.commons.collections.CollectionUtils
 
+import scala.collection.JavaConversions._
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
 
@@ -51,10 +57,33 @@ object UDFClient {
     udfInfoBuilder
   }
 
+  def getUdfInfosByUdfIds(
+      userName: String,
+      udfIds: Array[Long],
+      category: String,
+      udfType: BigInt
+  ): ArrayBuffer[UDFInfoVo] = {
+    val udfInfoBuilder = new ArrayBuffer[UDFInfoVo]
+
+    val udfTree = Sender.getSender(UDFClientConfiguration.UDF_SERVICE_NAME.getValue)
+      .ask(RequestUdfIds(userName, udfIds, category))
+      .asInstanceOf[ResponseUdfs]
+
+    if (CollectionUtils.isNotEmpty(udfTree.udfInfos)) {
+      udfTree.udfInfos.filter(infoVo => infoVo.getUdfType == udfType)
+        .foreach(infoVo => udfInfoBuilder.append(infoVo))
+    }
+    udfInfoBuilder
+  }
+
   def getJarUdf(userName: String): ArrayBuffer[UDFInfoVo] = {
     getUdfInfosByUdfType(userName, ConstantVar.UDF, ConstantVar.UDF_JAR)
   }
 
+  def getJarUdfByIds(userName: String, udfIds: Array[Long]): ArrayBuffer[UDFInfoVo] = {
+    getUdfInfosByUdfIds(userName, udfIds, ConstantVar.UDF, ConstantVar.UDF_JAR)
+  }
+
   def getPyUdf(userName: String): ArrayBuffer[UDFInfoVo] = {
     getUdfInfosByUdfType(userName, ConstantVar.UDF, ConstantVar.UDF_PY)
   }

linkis-public-enhancements/linkis-udf/linkis-udf-common/src/main/scala/org/apache/linkis/udf/api/rpc/RequestUdfIds.scala

@@ -0,0 +1,23 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.linkis.udf.api.rpc
+
+import org.apache.linkis.protocol.{CacheableProtocol, RetryableProtocol}
+
+case class RequestUdfIds(userName: String, udfIds: Array[Long], treeCategory: String)
+  extends RetryableProtocol with CacheableProtocol with UdfProtocol

linkis-public-enhancements/linkis-udf/linkis-udf-common/src/main/scala/org/apache/linkis/udf/api/rpc/ResponseUdfs.scala

@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.linkis.udf.api.rpc
+
+import org.apache.linkis.udf.vo.UDFInfoVo
+
+class ResponseUdfs(val udfInfos: java.util.List[UDFInfoVo]) extends UdfProtocol

linkis-public-enhancements/linkis-udf/linkis-udf-service/src/main/java/org/apache/linkis/udf/dao/UDFDao.java

@@ -61,6 +61,10 @@ List<UDFInfoVo> getUDFSByTreeIdAndUser(
   List<UDFInfoVo> getUDFInfoByTreeId(
       Long treeId, String userName, Collection<Integer> categoryCodes);
 
+  List<UDFInfoVo> getUDFInfoByIds(
+          @Param("ids") Long[] ids,
+          @Param("categoryCodes") Collection<Integer> categoryCodes);
+
   List<UDFInfo> getLoadedUDFs(String userName);
 
   List<Long> getLoadedUDFIds(String userName);

linkis-public-enhancements/linkis-udf/linkis-udf-service/src/main/java/org/apache/linkis/udf/dao/impl/UDFDao.xml

@@ -164,6 +164,20 @@
         on a.id=b.udf_id
     </select>
 
+    <select id="getUDFInfoByIds" resultType="org.apache.linkis.udf.vo.UDFInfoVo">
+        select pub.*, puv.path, puv.register_format, puv.use_format, puv.bml_resource_id, puv.bml_resource_version, puv.description
+        from
+            linkis_ps_udf_baseinfo pub
+            join (
+                select udf_id, max(bml_resource_version) as max_version
+                from linkis_ps_udf_version
+                where udf_id in <foreach collection="ids" open="(" separator="," close=")" item="item">#{item}</foreach>
+                group by udf_id
+            ) puvg on pub.id=puvg.udf_id
+            join linkis_ps_udf_version puv on puvg.udf_id=puv.udf_id and puvg.max_version=puv.bml_resource_version
+        where udf_type in <foreach collection="categoryCodes" open="(" separator="," close=")" item="item">#{item}</foreach>
+    </select>
+
     <select id="getLoadedUDFs" resultMap="UDFInfoMap">
         select b.*,if(a.udf_id IS NULL,0,1) as is_load
         from

linkis-public-enhancements/linkis-udf/linkis-udf-service/src/main/java/org/apache/linkis/udf/service/UDFService.java

@@ -52,6 +52,8 @@ List<UDFInfoVo> getUDFSByTreeIdAndUser(Long treeId, String userName, String cate
   List<UDFInfoVo> getUDFInfoByTreeId(Long treeId, String userName, String category)
       throws UDFException;
 
+  List<UDFInfoVo> getUDFInfoByIds(Long[] ids, String category) throws UDFException;
+
   Map<String, List<String>> generateInitSql(String userName) throws UDFException;
 
   Iterator<String> getAllLoadJars(String userName) throws UDFException;

linkis-public-enhancements/linkis-udf/linkis-udf-service/src/main/java/org/apache/linkis/udf/service/impl/UDFServiceImpl.java

@@ -806,6 +806,14 @@ public List<UDFInfoVo> getUDFInfoByTreeId(Long treeId, String userName, String c
     return udfDao.getUDFInfoByTreeId(treeId, userName, categoryToCodes.get(category));
   }
 
+  @Override
+  public List<UDFInfoVo> getUDFInfoByIds(Long[] ids, String category) {
+    if (ids == null || ids.length == 0) {
+      return new ArrayList<>(0);
+    }
+    return udfDao.getUDFInfoByIds(ids, categoryToCodes.get(category));
+  }
+
   /**
    * Generate sql needs content: divided into jar, python, scala Save Path and registration syntax
    * separately 生成sql需要内容： 分为jar，python，scala 分别保存Path和注册语法

linkis-public-enhancements/linkis-udf/linkis-udf-service/src/main/scala/org/apache/linkis/udf/api/rpc/UdfReceiver.scala

@@ -18,17 +18,21 @@
 package org.apache.linkis.udf.api.rpc
 
 import org.apache.linkis.rpc.{Receiver, Sender}
-import org.apache.linkis.udf.service.UDFTreeService
+import org.apache.linkis.udf.service.{UDFService, UDFTreeService}
 
+import java.lang
 import scala.concurrent.duration.Duration
 
 class UdfReceiver extends Receiver {
 
   private var udfTreeService: UDFTreeService = _
 
-  def this(udfTreeService: UDFTreeService) = {
+  private var udfService: UDFService = _
+
+  def this(udfTreeService: UDFTreeService, udfService: UDFService) = {
     this()
     this.udfTreeService = udfTreeService
+    this.udfService = udfService
   }
 
   override def receive(message: Any, sender: Sender): Unit = {}
@@ -38,6 +42,9 @@ class UdfReceiver extends Receiver {
       case RequestUdfTree(userName, treeType, treeId, treeCategory) =>
         val udfTree = udfTreeService.getTreeById(treeId, userName, treeType, treeCategory)
         new ResponseUdfTree(udfTree)
+      case RequestUdfIds(userName, udfIds, treeCategory) =>
+        val udfs = udfService.getUDFInfoByIds(udfIds.map(id => new lang.Long(id)), treeCategory)
+        new ResponseUdfs(udfs)
       case _ =>
     }
   }

linkis-public-enhancements/linkis-udf/linkis-udf-service/src/main/scala/org/apache/linkis/udf/api/rpc/UdfReceiverChooser.scala

@@ -18,7 +18,7 @@
 package org.apache.linkis.udf.api.rpc
 
 import org.apache.linkis.rpc.{Receiver, ReceiverChooser, RPCMessageEvent}
-import org.apache.linkis.udf.service.UDFTreeService
+import org.apache.linkis.udf.service.{UDFService, UDFTreeService}
 
 import org.springframework.beans.factory.annotation.Autowired
 import org.springframework.stereotype.Component
@@ -31,10 +31,13 @@ class UdfReceiverChooser extends ReceiverChooser {
   @Autowired
   private var udfTreeService: UDFTreeService = _
 
+  @Autowired
+  private var udfService: UDFService = _
+
   private var udfReceiver: Option[UdfReceiver] = None
 
   @PostConstruct
-  def init(): Unit = udfReceiver = Some(new UdfReceiver(udfTreeService))
+  def init(): Unit = udfReceiver = Some(new UdfReceiver(udfTreeService, udfService))
 
   override def chooseReceiver(event: RPCMessageEvent): Option[Receiver] = event.message match {
     case _: UdfProtocol => udfReceiver