Start session timeouts send fatal error and stop further executions

remotespark/kernels/kernelmagics.py

@@ -29,6 +29,8 @@ def __init__(self, shell, data=None):
         self.language = ""
         self.url = None
         self.connection_string = None
+        self.fatal_error = False
+        self.fatal_error_message = ""
 
     @cell_magic
     def help(self, line, cell="", local_ns=None):
@@ -96,7 +98,7 @@ def local(self, line, cell="", local_ns=None):
     def info(self, line, cell="", local_ns=None):
         self.ipython_display.writeln("Endpoint:\n\t{}\n".format(self.url))
 
-        self.ipython_display.writeln("Current session:\n\t{}\n".format(
+        self.ipython_display.writeln("Current session ID number:\n\t{}\n".format(
                 self.spark_controller.get_session_id_for_client(self.session_name)))
 
         self.ipython_display.writeln("Session configs:\n\t{}\n".format(conf.get_session_properties(self.language)))
@@ -145,37 +147,40 @@ def configure(self, line, cell="", local_ns=None):
 
     @cell_magic
     def spark(self, line, cell="", local_ns=None):
-        self._do_not_call_start_session("")
-
-        (success, out) = self.spark_controller.run_cell(cell)
-        if success:
-            self.ipython_display.write(out)
+        if self._do_not_call_start_session(""):
+            (success, out) = self.spark_controller.run_cell(cell)
+            if success:
+                self.ipython_display.write(out)
+            else:
+                self.ipython_display.send_error(out)
         else:
-            self.ipython_display.send_error(out)
+            return None
 
     @magic_arguments()
     @cell_magic
     @needs_local_scope
     @argument("-o", "--output", type=str, default=None, help="If present, query will be stored in variable of this "
                                                              "name.")
     def sql(self, line, cell="", local_ns=None):
-        self._do_not_call_start_session("")
-
-        args = parse_argstring(self.sql, line)
-        return self.execute_against_context_that_returns_df(self.spark_controller.run_cell_sql, cell,
+        if self._do_not_call_start_session(""):
+            args = parse_argstring(self.sql, line)
+            return self.execute_against_context_that_returns_df(self.spark_controller.run_cell_sql, cell,
                                                             None, args.output)
+        else:
+            return None
 
     @magic_arguments()
     @cell_magic
     @needs_local_scope
     @argument("-o", "--output", type=str, default=None, help="If present, query will be stored in variable of this "
                                                              "name.")
     def hive(self, line, cell="", local_ns=None):
-        self._do_not_call_start_session("")
-
-        args = parse_argstring(self.hive, line)
-        return self.execute_against_context_that_returns_df(self.spark_controller.run_cell_hive, cell,
-                                                            None, args.output)
+        if self._do_not_call_start_session(""):
+            args = parse_argstring(self.hive, line)
+            return self.execute_against_context_that_returns_df(self.spark_controller.run_cell_hive, cell,
+                                                                None, args.output)
+        else:
+            return None
 
     @magic_arguments()
     @cell_magic
@@ -190,7 +195,7 @@ def cleanup(self, line, cell="", local_ns=None):
             self.ipython_display.send_error("When you clean up the endpoint, all sessions will be lost, including the "
                                             "one used for this notebook. Include the -f parameter if that's your "
                                             "intention.")
-            return
+            return None
 
     @magic_arguments()
     @cell_magic
@@ -206,23 +211,38 @@ def delete(self, line, cell="", local_ns=None):
                 self.ipython_display.send_error("Cannot delete this kernel's session ({}). Specify a different session,"
                                                 " shutdown the kernel to delete this session, or run %cleanup to "
                                                 "delete all sessions for this endpoint.".format(id))
-                return
+                return None
 
             self.spark_controller.delete_session_by_id(self.connection_string, session)
         else:
             self.ipython_display.send_error("Include the -f parameter if you understand that all statements executed"
                                             "in this session will be lost.")
-            return
+            return None
 
     @cell_magic
     def _do_not_call_start_session(self, line, cell="", local_ns=None):
-        if not self.session_started:
-            self.session_started = True
+        # Starts a session unless session is already created or a fatal error occurred. Returns True when session is
+        # created successfully.
 
+        if self.fatal_error:
+            self.ipython_display.send_error(self.fatal_error_message)
+            return False
+
+        if not self.session_started:
             skip = False
             properties = conf.get_session_properties(self.language)
 
-            self.spark_controller.add_session(self.session_name, self.connection_string, skip, properties)
+            try:
+                self.session_started = True
+                self.spark_controller.add_session(self.session_name, self.connection_string, skip, properties)
+            except Exception as e:
+                self.fatal_error = True
+                self.fatal_error_message = conf.fatal_error_suggestion().format(e)
+                self.logger.error("Error creating session: {}".format(e))
+                self.ipython_display.send_error(self.fatal_error_message)
+                return False
+
+        return self.session_started
 
     @cell_magic
     def _do_not_call_delete_session(self, line, cell="", local_ns=None):
@@ -239,11 +259,11 @@ def _do_not_call_change_language(self, line, cell="", local_ns=None):
 
         if language not in Constants.lang_supported:
             self.ipython_display.send_error("'{}' language not supported in kernel magics.".format(language))
-            return
+            return None
 
         if self.session_started:
             self.ipython_display.send_error("Cannot change the language if a session has been started.")
-            return
+            return None
 
         self.language = language
         self.refresh_configuration()

remotespark/livyclientlib/livysession.py

@@ -158,7 +158,7 @@ def wait_for_idle(self, seconds_to_wait):
             return
 
         if current_status in Constants.final_status:
-            error = "Session {} unexpectedly reached final status {}. See logs:\n{}"\
+            error = "Session {} unexpectedly reached final status '{}'. See logs:\n{}"\
                 .format(self.id, current_status, self.logs)
             self.logger.error(error)
             raise LivyUnexpectedStatusError(error)

tests/test_kernel_magics.py

@@ -2,9 +2,10 @@
 from nose.tools import with_setup, raises, assert_equals
 from IPython.core.magic import magics_class
 
-import remotespark.utils.configuration as conf
 from remotespark.kernels.kernelmagics import KernelMagics
+from remotespark.livyclientlib.livyclienttimeouterror import LivyClientTimeoutError
 from remotespark.utils.constants import Constants
+import remotespark.utils.configuration as conf
 
 magic = None
 spark_controller = None
@@ -53,18 +54,41 @@ def test_start_session():
     line = ""
     assert not magic.session_started
 
-    magic._do_not_call_start_session(line)
+    ret = magic._do_not_call_start_session(line)
 
+    assert ret
     assert magic.session_started
     spark_controller.add_session.assert_called_once_with(magic.session_name, magic.connection_string, False,
                                                          {"kind": Constants.session_kind_pyspark})
 
     # Call a second time
-    magic._do_not_call_start_session(line)
+    ret = magic._do_not_call_start_session(line)
+    assert ret
     assert magic.session_started
     assert spark_controller.add_session.call_count == 1
 
 
+@with_setup(_setup, _teardown)
+def test_start_session_times_out():
+    line = ""
+    spark_controller.add_session = MagicMock(side_effect=LivyClientTimeoutError)
+    assert not magic.session_started
+
+    ret = magic._do_not_call_start_session(line)
+
+    assert not ret
+    assert magic.session_started
+    assert magic.fatal_error
+    assert_equals(ipython_display.send_error.call_count, 1)
+
+    # Call after fatal error
+    ipython_display.send_error.reset_mock()
+    ret = magic._do_not_call_start_session(line)
+    assert not ret
+    assert magic.session_started
+    assert_equals(ipython_display.send_error.call_count, 1)
+
+
 @with_setup(_setup, _teardown)
 def test_delete_session():
     line = ""
@@ -195,7 +219,7 @@ def test_get_session_settings():
     assert magic.get_session_settings("    something", False) == "something"
     assert magic.get_session_settings("-f something", True) == "something"
     assert magic.get_session_settings("something -f", True) == "something"
-    assert magic.get_session_settings("something", True) == None
+    assert magic.get_session_settings("something", True) is None
 
 
 @with_setup(_setup, _teardown)
@@ -225,6 +249,21 @@ def test_spark_error():
                                                          {"kind": Constants.session_kind_pyspark})
     spark_controller.run_cell.assert_called_once_with(cell)
 
+
+@with_setup(_setup, _teardown)
+def test_spark_failed_session_start():
+    line = ""
+    cell = "some spark code"
+    magic._do_not_call_start_session = MagicMock(return_value=False)
+
+    ret = magic.spark(line, cell)
+
+    assert_equals(ret, None)
+    assert_equals(ipython_display.write.call_count, 0)
+    assert_equals(spark_controller.add_session.call_count, 0)
+    assert_equals(spark_controller.run_cell.call_count, 0)
+
+
 @with_setup(_setup, _teardown)
 def test_sql_without_output():
     line = ""
@@ -238,6 +277,7 @@ def test_sql_without_output():
     magic.execute_against_context_that_returns_df.assert_called_once_with(spark_controller.run_cell_sql, cell, None,
                                                                           None)
 
+
 @with_setup(_setup, _teardown)
 def test_sql_with_output():
     line = "-o my_var"
@@ -252,6 +292,19 @@ def test_sql_with_output():
                                                                           "my_var")
 
 
+@with_setup(_setup, _teardown)
+def test_sql_failed_session_start():
+    line = ""
+    cell = "some spark code"
+    magic._do_not_call_start_session = MagicMock(return_value=False)
+
+    ret = magic.sql(line, cell)
+
+    assert_equals(ret, None)
+    assert_equals(spark_controller.add_session.call_count, 0)
+    assert_equals(spark_controller.execute_against_context_that_returns_df.call_count, 0)
+
+
 @with_setup(_setup, _teardown)
 def test_hive_without_output():
     line = ""
@@ -265,6 +318,7 @@ def test_hive_without_output():
     magic.execute_against_context_that_returns_df.assert_called_once_with(spark_controller.run_cell_hive, cell, None,
                                                                           None)
 
+
 @with_setup(_setup, _teardown)
 def test_hive_with_output():
     line = "-o my_var"
@@ -279,6 +333,19 @@ def test_hive_with_output():
                                                                           "my_var")
 
 
+@with_setup(_setup, _teardown)
+def test_hive_failed_session_start():
+    line = ""
+    cell = "some spark code"
+    magic._do_not_call_start_session = MagicMock(return_value=False)
+
+    ret = magic.hive(line, cell)
+
+    assert_equals(ret, None)
+    assert_equals(spark_controller.add_session.call_count, 0)
+    assert_equals(spark_controller.execute_against_context_that_returns_df.call_count, 0)
+
+
 @with_setup(_setup, _teardown)
 def test_cleanup_without_force():
     line = ""