[incubator-kyuubi] branch master updated: [KYUUBI #3818] [Subtask] [PySpark] set local properties to spark context for python scripts

[ch...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

chengpan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-kyuubi.git


The following commit(s) were added to refs/heads/master by this push:
     new a54f214bc [KYUUBI #3818] [Subtask] [PySpark] set local properties to spark context for python scripts
a54f214bc is described below

commit a54f214bc583aabc41a72f73915f7ee1330e3c3f
Author: liangbowen <li...@gf.com.cn>
AuthorDate: Sun Nov 20 19:08:07 2022 +0800

    [KYUUBI #3818] [Subtask] [PySpark] set local properties to spark context for python scripts
    
    ### _Why are the changes needed?_
    
    to close #3818 .
    - add `withLocalProperties` in `ExecutePython`
    - set Spark local properties `kyuubi.session.user`, `kyuubi.statement.id`, `spark.scheduler.pool`, `spark.scheduler.pool` in withLocalProperties for other feature relied on it , e.g. Authz plugin.
    
    ### _How was this patch tested?_
    - [ ] 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.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request
    
    Closes #3817 from bowenliang123/pyspark-sessionuser.
    
    Closes #3818
    
    86876391 [liangbowen] add local props `kyuubi.statement.id` , `spark.scheduler.pool`, `spark.scheduler.pool`. and setting job group
    1c2bec4e [liangbowen] nit
    276ad43b [liangbowen] generalize setSparkLocalProperties
    00da08eb [liangbowen] nit
    9f2cdaa5 [liangbowen] set session user to spark local property before execute python
    
    Authored-by: liangbowen <li...@gf.com.cn>
    Signed-off-by: Cheng Pan <ch...@apache.org>
---
 .../engine/spark/operation/ExecutePython.scala     | 27 +++++++++++++++++++++-
 .../engine/spark/operation/SparkOperation.scala    |  4 ++--
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala
index cb3d8f842..3254f0e2c 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala
@@ -32,6 +32,8 @@ import org.apache.spark.sql.{Row, RuntimeConfig}
 import org.apache.spark.sql.types.StructType
 
 import org.apache.kyuubi.Logging
+import org.apache.kyuubi.config.KyuubiReservedKeys.{KYUUBI_SESSION_USER_KEY, KYUUBI_STATEMENT_ID_KEY}
+import org.apache.kyuubi.engine.spark.KyuubiSparkUtil.SPARK_SCHEDULER_POOL_KEY
 import org.apache.kyuubi.operation.ArrayFetchIterator
 import org.apache.kyuubi.session.Session
 
@@ -52,7 +54,7 @@ class ExecutePython(
     }
   }
 
-  override protected def runInternal(): Unit = {
+  override protected def runInternal(): Unit = withLocalProperties {
     val response = worker.runCode(statement)
     val output = response.map(_.content.getOutput()).getOrElse("")
     val status = response.map(_.content.status).getOrElse("UNKNOWN_STATUS")
@@ -63,6 +65,29 @@ class ExecutePython(
       new ArrayFetchIterator[Row](Array(Row(output, status, ename, evalue, Row(traceback: _*))))
   }
 
+  override protected def withLocalProperties[T](f: => T): T = {
+    val setSparkLocalProperties = (key: String, value: String) => {
+      worker.runCode(s"spark.sparkContext.setLocalProperty('$key', '$value')")
+    }
+    try {
+      worker.runCode("spark.sparkContext.setJobGroup" +
+        s"($statementId, $redactedStatement, $forceCancel)")
+      setSparkLocalProperties(KYUUBI_SESSION_USER_KEY, session.user)
+      setSparkLocalProperties(KYUUBI_STATEMENT_ID_KEY, statementId)
+      schedulerPool match {
+        case Some(pool) =>
+          setSparkLocalProperties(SPARK_SCHEDULER_POOL_KEY, pool)
+        case None =>
+      }
+
+      f
+    } finally {
+      setSparkLocalProperties(KYUUBI_SESSION_USER_KEY, "")
+      setSparkLocalProperties(KYUUBI_STATEMENT_ID_KEY, "")
+      setSparkLocalProperties(SPARK_SCHEDULER_POOL_KEY, "")
+      worker.runCode("spark.sparkContext.clearJobGroup()")
+    }
+  }
 }
 
 case class SessionPythonWorker(
diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkOperation.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkOperation.scala
index dfa2f5c43..1d08f540d 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkOperation.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkOperation.scala
@@ -66,10 +66,10 @@ abstract class SparkOperation(session: Session)
     }
   }
 
-  private val forceCancel =
+  protected val forceCancel =
     session.sessionManager.getConf.get(KyuubiConf.OPERATION_FORCE_CANCEL)
 
-  private val schedulerPool =
+  protected val schedulerPool =
     spark.conf.getOption(KyuubiConf.OPERATION_SCHEDULER_POOL.key).orElse(
       session.sessionManager.getConf.get(KyuubiConf.OPERATION_SCHEDULER_POOL))