[zeppelin] branch master updated: [ZEPPELIN-5962] LivyInterpreter support more parameter when create session (#4657)

[jongyoul]

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

jongyoul pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/zeppelin.git


The following commit(s) were added to refs/heads/master by this push:
     new 5de2d035fa [ZEPPELIN-5962] LivyInterpreter support more parameter when 
create session (#4657)
5de2d035fa is described below

commit 5de2d035fadff6aa43e463f33ac030f831ee1d58
Author: mrzhao <imrz...@qq.com>
AuthorDate: Tue Sep 19 05:55:11 2023 +0800

    [ZEPPELIN-5962] LivyInterpreter support more parameter when create session 
(#4657)
    
    * LivyInterpreter support more parameter when create session
    
    * Add test for ZEPPELIN-5962
    
    ---------
    
    Co-authored-by: mrzhao <mrz...@iflytek.com>
---
 .../apache/zeppelin/livy/BaseLivyInterpreter.java  | 19 ++++++++++--
 .../apache/zeppelin/livy/LivyInterpreterIT.java    | 36 ++++++++++++++++++++++
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git 
a/livy/src/main/java/org/apache/zeppelin/livy/BaseLivyInterpreter.java 
b/livy/src/main/java/org/apache/zeppelin/livy/BaseLivyInterpreter.java
index f7b576d818..b897537ec2 100644
--- a/livy/src/main/java/org/apache/zeppelin/livy/BaseLivyInterpreter.java
+++ b/livy/src/main/java/org/apache/zeppelin/livy/BaseLivyInterpreter.java
@@ -19,6 +19,7 @@ package org.apache.zeppelin.livy;
 
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
+import com.google.gson.JsonObject;
 import com.google.gson.annotations.SerializedName;
 
 import org.apache.commons.lang3.StringUtils;
@@ -302,15 +303,19 @@ public abstract class BaseLivyInterpreter extends 
Interpreter {
       throws LivyException {
     try {
       Map<String, String> conf = new HashMap<>();
+      Map<String, String> params = new HashMap<>();
       for (Map.Entry<Object, Object> entry : getProperties().entrySet()) {
         if (entry.getKey().toString().startsWith("livy.spark.") &&
             !entry.getValue().toString().isEmpty()) {
           conf.put(entry.getKey().toString().substring(5), 
entry.getValue().toString());
+        } else if (entry.getKey().toString().startsWith("livy.") &&
+                !entry.getValue().toString().isEmpty()) {
+          params.put(entry.getKey().toString().substring(5), 
entry.getValue().toString());
         }
       }
 
       CreateSessionRequest request = new CreateSessionRequest(kind,
-          user == null || user.equals("anonymous") ? null : user, conf);
+          user == null || user.equals("anonymous") ? null : user, conf, 
params);
       SessionInfo sessionInfo = SessionInfo.fromJson(
           callRestAPI("/sessions", "POST", request.toJson()));
       long start = System.currentTimeMillis();
@@ -776,15 +781,23 @@ public abstract class BaseLivyInterpreter extends 
Interpreter {
     @SerializedName("proxyUser")
     public final String user;
     public final Map<String, String> conf;
+    public final Map<String, String> params;
 
-    CreateSessionRequest(String kind, String user, Map<String, String> conf) {
+    CreateSessionRequest(String kind, String user, Map<String, String> conf,
+                         Map<String, String> params) {
       this.kind = kind;
       this.user = user;
       this.conf = conf;
+      this.params = params;
     }
 
     public String toJson() {
-      return gson.toJson(this);
+      JsonObject jsonObject = new JsonObject();
+      jsonObject.add("conf", gson.toJsonTree(conf));
+      params.forEach(jsonObject::addProperty);
+      jsonObject.addProperty("kind", kind);
+      jsonObject.addProperty("proxyUser", user);
+      return gson.toJson(jsonObject);
     }
   }
 
diff --git a/livy/src/test/java/org/apache/zeppelin/livy/LivyInterpreterIT.java 
b/livy/src/test/java/org/apache/zeppelin/livy/LivyInterpreterIT.java
index a649085897..3a3a17c5af 100644
--- a/livy/src/test/java/org/apache/zeppelin/livy/LivyInterpreterIT.java
+++ b/livy/src/test/java/org/apache/zeppelin/livy/LivyInterpreterIT.java
@@ -644,6 +644,42 @@ public class LivyInterpreterIT {
     }
   }
 
+  @Test
+  void testLivyParams() throws InterpreterException {
+    if (!checkPreCondition()) {
+      return;
+    }
+    InterpreterGroup interpreterGroup = new InterpreterGroup("group_1");
+    interpreterGroup.put("session_1", new ArrayList<Interpreter>());
+    Properties props = new Properties(properties);
+    props.setProperty("livy.spark.executor.cores", "4");
+    props.setProperty("livy.name", "zeppelin-livy");
+    LivySparkInterpreter sparkInterpreter = new LivySparkInterpreter(props);
+    sparkInterpreter.setInterpreterGroup(interpreterGroup);
+    interpreterGroup.get("session_1").add(sparkInterpreter);
+    AuthenticationInfo authInfo = new AuthenticationInfo("user1");
+    MyInterpreterOutputListener outputListener = new 
MyInterpreterOutputListener();
+    InterpreterOutput output = new InterpreterOutput(outputListener);
+    InterpreterContext context = InterpreterContext.builder()
+            .setNoteId("noteId")
+            .setParagraphId("paragraphId")
+            .setAuthenticationInfo(authInfo)
+            .setInterpreterOut(output)
+            .build();
+    sparkInterpreter.open();
+
+    try {
+      InterpreterResult result = sparkInterpreter.interpret("sc.version\n" +
+              "assert(sc.getConf.get(\"spark.executor.cores\") == \"4\" && " +
+                      "sc.getConf.get(\"spark.app.name\") == 
\"zeppelin-livy\")"
+              , context);
+      assertEquals(InterpreterResult.Code.SUCCESS, result.code(), 
result.toString());
+      assertEquals(1, result.message().size());
+    } finally {
+      sparkInterpreter.close();
+    }
+  }
+
   @Test
   void testLivyTutorialNote() throws IOException, InterpreterException {
     if (!checkPreCondition()) {