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: + + + 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""" +""" + \ + u"""""".format(connection_info["endpoint"], connection_info["session"]) + \ + u"
endpointsession
{0}{1}
" + + return html + def initialize_auth(args): """Creates an authenticatior class instance for the given auth type