WIP: Add Connection Info Line Magic

sparkmagic/sparkmagic/kernels/kernelmagics.py

@@ -84,6 +84,17 @@ def help(self, line, cell="", local_ns=None):
     <td>%%info</td>
     <td>Outputs session information for the current Livy endpoint.</td>
   </tr>
+  <tr>
+    <td>connection_info</td>
+    <td>%connection_info -o connection</td>
+    <td>Outputs or returns the url and session id for the current session.</td>
+    Parameters:
+      <ul>
+        <li>-o VAR_NAME: The variable of name VAR_NAME will be available in the %%local Python context as a
+          dict: {"endpoint": "http:<ip_address>:<port>", "session": "<session_number>"}
+      </ul>
+    </td>
+  </tr>
   <tr>
     <td>cleanup</td>
     <td>%%cleanup -f</td>
@@ -211,6 +222,33 @@ def info(self, line, cell=u"", local_ns=None):
         info_sessions = self.spark_controller.get_all_sessions_endpoint(self.endpoint)
         self._print_endpoint_info(info_sessions, current_session_id)
 
+    @magic_arguments()
+    @argument(
+        "-o",
+        "--output",
+        type=str,
+        default=None,
+        help="If present, indicated variable will be stored in variable"
+        "of this name in user's local context.",
+    )
+    @line_magic
+    @wrap_unexpected_exceptions
+    @handle_expected_exceptions
+    @_event
+    def connection_info(self, line, cell=u"", local_ns=None):
+        args = parse_argstring_or_throw(self.info, line)
+        connection_info = {"endpoint": None, "session": None}
+        if self.session_started:
+            current_session_id = self.spark_controller.get_session_id_for_client(
+                self.session_name
+            )
+            connection_info["endpoint"] = self.endpoint.url
+            connection_info["session"] = current_session_id
+        if args.output:
+            self.shell.user_ns[args.output] = connection_info
+        else:
+            self._print_current_connection_info(connection_info, current_session_id)
+
     @magic_arguments()
     @cell_magic
     @wrap_unexpected_exceptions

sparkmagic/sparkmagic/magics/sparkmagicsbase.py

@@ -15,7 +15,7 @@
 import sparkmagic.utils.configuration as conf
 from sparkmagic.utils.sparklogger import SparkLog
 from sparkmagic.utils.sparkevents import SparkEvents
-from sparkmagic.utils.utils import get_sessions_info_html
+from sparkmagic.utils.utils import get_connection_info_html, get_sessions_info_html
 from sparkmagic.utils.constants import MIMETYPE_TEXT_HTML 
 from sparkmagic.livyclientlib.sparkcontroller import SparkController
 from sparkmagic.livyclientlib.sqlquery import SQLQuery
@@ -131,3 +131,10 @@ def _print_endpoint_info(self, info_sessions, current_session_id):
             self.ipython_display.html(html)
         else:
             self.ipython_display.html(u'No active sessions.')
+
+    def _print_current_connection_info(self, connection_info):
+        if connection_info.get("endpoint"):
+            html = get_connection_info_html(connection_info)
+            self.ipython_display.html(html)
+        else:
+            self.ipython_display.html(u'No active sessions.')

sparkmagic/sparkmagic/utils/utils.py

@@ -94,6 +94,15 @@ def get_sessions_info_html(info_sessions, current_session_id):
 
     return html
 
+
+def get_connection_info_html(connection_info):
+    html = u"""<table>
+<tr><th>endpoint</th><th>session</th></tr>""" + \
+    u"""<tr><td>{0}</td><td>{1}</td></tr>""".format(connection_info["endpoint"], connection_info["session"]) + \
+    u"</table>"
+
+    return html
+
 def initialize_auth(args):
     """Creates an authenticatior class instance for the given auth type