[LIVY-758] Document how to attach to an existing session from Java client and add a property to the builder

Author: agilelab-tmnd1991

api/src/main/java/org/apache/livy/LivyClientBuilder.java

@@ -36,9 +36,10 @@
 public final class LivyClientBuilder {
 
   public static final String LIVY_URI_KEY = "livy.uri";
+  public static final String LIVY_SESSION_ID_KEY = "livy.sessionId";
 
   private static final ServiceLoader<LivyClientFactory> CLIENT_FACTORY_LOADER =
-    ServiceLoader.load(LivyClientFactory.class, classLoader());
+      ServiceLoader.load(LivyClientFactory.class, classLoader());
 
   private static List<LivyClientFactory> getLivyClientFactories() {
     List<LivyClientFactory> factories = new ArrayList<>();
@@ -96,11 +97,30 @@ public LivyClientBuilder(boolean loadDefaults) throws IOException {
     }
   }
 
+  /**
+   *
+   * @param uri the uri of the livy server,
+   *            if the uri contains <pre>sessions/{sessionId}</pre>
+   *            the client will connect to the existing session,
+   *            otherwise it will create a new session.
+   * @return the builder itself.
+   */
   public LivyClientBuilder setURI(URI uri) {
     config.setProperty(LIVY_URI_KEY, uri.toString());
     return this;
   }
 
+  /**
+   *
+   * @param sessionId the id of the session to attach to,
+   *                  if not set a new session will be created when the client is built.
+   * @return the builder itself.
+   */
+  public LivyClientBuilder setSessionId(int sessionId) {
+    config.setProperty(LIVY_SESSION_ID_KEY, "" + sessionId);
+    return this;
+  }
+
   public LivyClientBuilder setConf(String key, String value) {
     if (value != null) {
       config.setProperty(key, value);

client-http/src/main/java/org/apache/livy/client/http/HttpClient.java

@@ -33,6 +33,7 @@
 import org.apache.livy.Job;
 import org.apache.livy.JobHandle;
 import org.apache.livy.LivyClient;
+import org.apache.livy.LivyClientBuilder;
 import org.apache.livy.client.common.Serializer;
 import static org.apache.livy.client.common.HttpMessages.*;
 
@@ -54,17 +55,18 @@ class HttpClient implements LivyClient {
     this.config = httpConf;
     this.stopped = false;
 
-    // If the given URI looks like it refers to an existing session, then try to connect to
-    // an existing session. Note this means that any Spark configuration in httpConf will be
-    // unused.
-    Matcher m = Pattern.compile("(.*)" + LivyConnection.SESSIONS_URI + "/([0-9]+)")
-      .matcher(uri.getPath());
-
+    Matcher m = Pattern.compile("(.*)" + LivyConnection.SESSIONS_URI + "/([0-9]+)").matcher(uri.getPath());
+    String sessionIdFromConf = httpConf.get(LivyClientBuilder.LIVY_SESSION_ID_KEY);
     try {
-      if (m.matches()) {
+      if (sessionIdFromConf != null && m.matches()) {
+        throw new IllegalArgumentException("Cannot set existing session both from URI and configuration");
+      } else if (sessionIdFromConf != null) {
+        this.conn = new LivyConnection(uri, httpConf);
+        this.sessionId = Integer.parseInt(sessionIdFromConf);
+        conn.post(null, SessionInfo.class, "/%d/connect", sessionId);
+      } else if (m.matches()) {
         URI base = new URI(uri.getScheme(), uri.getUserInfo(), uri.getHost(), uri.getPort(),
-          m.group(1), uri.getQuery(), uri.getFragment());
-
+            m.group(1), uri.getQuery(), uri.getFragment());
         this.conn = new LivyConnection(base, httpConf);
         this.sessionId = Integer.parseInt(m.group(2));
         conn.post(null, SessionInfo.class, "/%d/connect", sessionId);

client-http/src/test/scala/org/apache/livy/client/http/HttpClientSpec.scala

@@ -181,15 +181,34 @@ class HttpClientSpec extends FunSpecLike with BeforeAndAfterAll with LivyBaseUni
       testJob(false, response = Some(null))
     }
 
-    withClient("should connect to existing sessions") {
-      var sid = client.asInstanceOf[HttpClient].getSessionId()
+    withClient("should connect to existing sessions using the URI") {
+      val sid = client.asInstanceOf[HttpClient].getSessionId()
       val uri = s"http://${InetAddress.getLocalHost.getHostAddress}:${server.port}" +
         s"${LivyConnection.SESSIONS_URI}/$sid"
       val newClient = new LivyClientBuilder(false).setURI(new URI(uri)).build()
       newClient.stop(false)
       verify(session, never()).stop()
     }
 
+    withClient("should connect to existing sessions using the conf") {
+      val sid = client.asInstanceOf[HttpClient].getSessionId()
+      val uri = s"http://${InetAddress.getLocalHost.getHostAddress}:${server.port}"
+      val newClient = new LivyClientBuilder(false).setURI(new URI(uri)).setSessionId(sid).build()
+      newClient.stop(false)
+      verify(session, never()).stop()
+    }
+
+    withClient("should throw an exception if both the sessionId conf is set and the URI contains the session") {
+      val sid = client.asInstanceOf[HttpClient].getSessionId()
+      val uri = s"http://${InetAddress.getLocalHost.getHostAddress}:${server.port}" +
+        s"${LivyConnection.SESSIONS_URI}/$sid"
+      intercept[IllegalArgumentException]{
+        val newClient = new LivyClientBuilder(false).setURI(new URI(uri)).setSessionId(sid).build()
+        newClient.stop(false)
+        verify(session, never()).stop()
+      }
+    }
+
     withClient("should tear down clients") {
       client.stop(true)
       verify(session, times(1)).stop()