[incubator-kyuubi] branch master updated: [KYUUBI #1522] Pass the client ip to engine

[ya...@apache.org]

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

yao 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 16add92  [KYUUBI #1522] Pass the client ip to engine
16add92 is described below

commit 16add92299a14ab5bb2e9e3f0a4ddfb51653bc62
Author: Wang Zhen <wa...@qiyi.com>
AuthorDate: Wed Dec 8 17:08:09 2021 +0800

    [KYUUBI #1522] Pass the client ip to engine
    
    <!--
    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.
    -->
    
    Pass the client ip to engine. For details: #1522
    
    ### _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.readthedocs.io/en/latest/develop_tools/testing.html#running-tests) locally before make a pull request
    
    Closes #1524 from wForget/KYUUBI-1522.
    
    Closes #1522
    
    c9d54ed9 [Wang Zhen] fix word error
    a16d914d [Wang Zhen] [KYUUBI-1522] Pass the client ip to engine
    
    Authored-by: Wang Zhen <wa...@qiyi.com>
    Signed-off-by: Kent Yao <ya...@apache.org>
---
 .../apache/kyuubi/engine/spark/session/SparkSQLSessionManager.scala  | 5 +++--
 kyuubi-common/src/main/scala/org/apache/kyuubi/session/package.scala | 2 ++
 .../main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala  | 5 +++--
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/session/SparkSQLSessionManager.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/session/SparkSQLSessionManager.scala
index 76e8c08..693bbc6 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/session/SparkSQLSessionManager.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/session/SparkSQLSessionManager.scala
@@ -57,7 +57,8 @@ class SparkSQLSessionManager private (name: String, spark: SparkSession)
       password: String,
       ipAddress: String,
       conf: Map[String, String]): SessionHandle = {
-    info(s"Opening session for $user@$ipAddress")
+    val clientIp = conf.getOrElse(CLIENT_IP_KEY, ipAddress)
+    info(s"Opening session for $user@$clientIp")
     val sparkSession =
       try {
         if (singleSparkSession) {
@@ -78,7 +79,7 @@ class SparkSQLSessionManager private (name: String, spark: SparkSession)
       protocol,
       user,
       password,
-      ipAddress,
+      clientIp,
       conf,
       this,
       sparkSession)
diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/session/package.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/session/package.scala
index 40abded..a5b874d 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/session/package.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/session/package.scala
@@ -27,4 +27,6 @@ package object session {
   val METACONF_PREFIX = "metaconf:"
 
   val SPARK_PREFIX = "spark."
+
+  val CLIENT_IP_KEY = "kyuubi.client.ip"
 }
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala
index ca5521b..99bd9d6 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala
@@ -50,13 +50,14 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) {
       conf: Map[String, String]): SessionHandle = {
 
     val username = Option(user).filter(_.nonEmpty).getOrElse("anonymous")
-
+    // inject client ip into session conf
+    val newConf = conf + (CLIENT_IP_KEY -> ipAddress)
     val sessionImpl = new KyuubiSessionImpl(
       protocol,
       username,
       password,
       ipAddress,
-      conf,
+      newConf,
       this,
       this.getConf.getUserDefaults(user))
     try {