From 3afcfe66260c7de2e48296faf7d7a42957a791d2 Mon Sep 17 00:00:00 2001
From: ulysses-you <ulyssesyou18@gmail.com>
Date: Wed, 20 Oct 2021 17:23:18 +0800
Subject: [PATCH] [KYUUBI #1263] Kyuubi auxiliary UDF failed on Spark Yarn mode

<!--
Thanks for sending a pull request!

Here are some tips for you:
  1. If this is your first time, please read our contributor guidelines: https://kyuubi.readthedocs.io/en/latest/community/contributions.html
  2. If the PR is related to an issue in https://github.com/apache/incubator-kyuubi/issues, add '[KYUUBI #XXXX]' in your PR title, e.g., '[KYUUBI #XXXX] Your PR title ...'.
  3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][KYUUBI #XXXX] Your PR title ...'.
-->

### _Why are the changes needed?_
<!--
Please clarify why the changes are needed. For instance,
  1. If you add a feature, you can talk about the use case of it.
  2. If you fix a bug, you can clarify why it is a bug.
-->
Closes https://github.com/apache/incubator-kyuubi/issues/1263

### _How was this patch tested?_
- [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible

- [ ] Add screenshots for manual tests if appropriate

- [x] [Run test](https://kyuubi.readthedocs.io/en/latest/develop_tools/testing.html#running-tests) locally before make a pull request

Closes #1264 from ulysses-you/kyuubi-1263.

Closes #1263

5685a409 [ulysses-you] id
b65cfeff [ulysses-you] test
5a44f0ad [ulysses-you] fix

Authored-by: ulysses-you <ulyssesyou18@gmail.com>
Signed-off-by: Cheng Pan <chengpan@apache.org>
---
 .../org/apache/kyuubi/engine/spark/udf/KDFRegistry.scala  | 4 +++-
 .../operation/KyuubiOperationYarnClusterSuite.scala       | 8 ++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/udf/KDFRegistry.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/udf/KDFRegistry.scala
index 537a77117..11900f77a 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/udf/KDFRegistry.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/udf/KDFRegistry.scala
@@ -19,6 +19,7 @@ package org.apache.kyuubi.engine.spark.udf
 
 import scala.collection.mutable.ArrayBuffer
 
+import org.apache.spark.SparkEnv
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.expressions.UserDefinedFunction
 import org.apache.spark.sql.functions.udf
@@ -27,9 +28,10 @@ import org.apache.kyuubi.KYUUBI_VERSION
 
 object KDFRegistry {
 
+  @transient
   val registeredFunctions = new ArrayBuffer[KyuubiDefinedFunction]()
 
-  val appName = SparkSession.active.sparkContext.appName
+  val appName = SparkEnv.get.conf.get("spark.app.name")
 
   val kyuubi_version: KyuubiDefinedFunction = create(
     "kyuubi_version",
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationYarnClusterSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationYarnClusterSuite.scala
index 122a3dbfa..734406e0f 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationYarnClusterSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/KyuubiOperationYarnClusterSuite.scala
@@ -42,4 +42,12 @@ class KyuubiOperationYarnClusterSuite extends WithKyuubiServerOnYarn with JDBCTe
       assert(resultSet.getString("id").startsWith("application_"))
     }
   }
+
+  test("KYUUBI #1263: Kyuubi auxiliary UDF failed on Spark Yarn mode") {
+    withJdbcStatement() { statement =>
+      val resultSet = statement.executeQuery("""SELECT kyuubi_version() as id""")
+      assert(resultSet.next())
+      assert(resultSet.getString("id") === org.apache.kyuubi.KYUUBI_VERSION)
+    }
+  }
 }