diff --git a/sparkmagic/sparkmagic/kernels/kernelmagics.py b/sparkmagic/sparkmagic/kernels/kernelmagics.py
index b22771994..a9b86bd65 100644
--- a/sparkmagic/sparkmagic/kernels/kernelmagics.py
+++ b/sparkmagic/sparkmagic/kernels/kernelmagics.py
@@ -84,6 +84,17 @@ def help(self, line, cell="", local_ns=None):
%%info |
Outputs session information for the current Livy endpoint. |
+
+ connection_info |
+ %connection_info -o connection |
+ Outputs or returns the url and session id for the current session. |
+ Parameters:
+
+ - -o VAR_NAME: The variable of name VAR_NAME will be available in the %%local Python context as a
+ dict: {"endpoint": "http::", "session": ""}
+
+
+
cleanup |
%%cleanup -f |
@@ -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
diff --git a/sparkmagic/sparkmagic/magics/sparkmagicsbase.py b/sparkmagic/sparkmagic/magics/sparkmagicsbase.py
index 5c5f62045..31eb63cc3 100644
--- a/sparkmagic/sparkmagic/magics/sparkmagicsbase.py
+++ b/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.')
diff --git a/sparkmagic/sparkmagic/utils/utils.py b/sparkmagic/sparkmagic/utils/utils.py
index 5af205182..35fd48d9f 100644
--- a/sparkmagic/sparkmagic/utils/utils.py
+++ b/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"""
+endpoint | session |
""" + \
+ u"""{0} | {1} |
""".format(connection_info["endpoint"], connection_info["session"]) + \
+ u"
"
+
+ return html
+
def initialize_auth(args):
"""Creates an authenticatior class instance for the given auth type