Repository: nteract/papermill (language: python)

Issue:
PapermillExecutionError cannot be loaded from picklefile
## 🐛 Bug

As part of my team's workflow, we extensively use Papermill to run notebooks as part of our Apache Airflow pipelines. When upgrading Airflow to version 2.X.X we found that our default error handling tools were no longer providing detailed errors from our PapermIll tasks. The cause for this is that Airflow changed the way it handles errors in it's tasks. It temporarily pickles the error message, and then at a later point loads the pickled exception. 

Because PapermillExecutionError is called with multiple arguments, but then only passes up a message argument to its superclass, pickle is unable to load it. See [this stack overflow question/answer](https://stackoverflow.com/questions/41808912/cannot-unpickle-exception-subclass) for a general case of this issue.

This error can be recreated with this simple test:
```{python}
import tempfile
import pickle

from papermill.exceptions import PapermillExecutionError

def test_exceptions_are_unpickleable():
    """Ensure exceptions can be unpickled"""
    temp_file = NamedTemporaryFile()
    err = PapermillExecutionError(1,2, "TestSource", "Exception", Exception(), ["Traceback", "message"])

    with open(temp_file.name, 'wb') as fd:
        pickle.dump(err, fd)
    # Read the Pickled File
    temp_file.seek(0)
    data = temp_file.read()
    pickled_err = pickle.loads(data)
    assert str(pickled_err) == str(err)
```


Discussion / hints:

Fix patch:
diff --git a/papermill/exceptions.py b/papermill/exceptions.py
index ac3c49f1..6a45a8ce 100644
--- a/papermill/exceptions.py
+++ b/papermill/exceptions.py
@@ -23,18 +23,26 @@ class PapermillExecutionError(PapermillException):
     """Raised when an exception is encountered in a notebook."""
 
     def __init__(self, cell_index, exec_count, source, ename, evalue, traceback):
+        args = cell_index, exec_count, source, ename, evalue, traceback
         self.cell_index = cell_index
         self.exec_count = exec_count
         self.source = source
         self.ename = ename
         self.evalue = evalue
         self.traceback = traceback
+
+        super(PapermillExecutionError, self).__init__(*args)
+
+    def __str__(self):
+        # Standard Behavior of an exception is to produce a string representation of its arguments
+        # when called with str(). In order to maintain compatability with previous versions which
+        # passed only the message to the superclass constructor, __str__ method is implemented to
+        # provide the same result as was produced in the past.
         message = "\n" + 75 * "-" + "\n"
-        message += 'Exception encountered at "In [%s]":\n' % str(exec_count)
-        message += strip_color("\n".join(traceback))
+        message += 'Exception encountered at "In [%s]":\n' % str(self.exec_count)
+        message += strip_color("\n".join(self.traceback))
         message += "\n"
-
-        super(PapermillExecutionError, self).__init__(message)
+        return message
 
 
 class PapermillRateLimitException(PapermillException):


Test patch:
diff --git a/papermill/tests/test_exceptions.py b/papermill/tests/test_exceptions.py
index 45027c7e..2c6968c9 100644
--- a/papermill/tests/test_exceptions.py
+++ b/papermill/tests/test_exceptions.py
@@ -1,1 +1,43 @@
+import pickle
+import tempfile
+
 import pytest  # noqa
+
+from .. import exceptions
+
+
+@pytest.fixture
+def temp_file():
+    f = tempfile.NamedTemporaryFile()
+    yield f
+    f.close()
+
+
+@pytest.mark.parametrize(
+    "exc,args",
+    [
+        (
+            exceptions.PapermillExecutionError,
+            (1, 2, "TestSource", "Exception", Exception(), ["Traceback", "Message"]),
+        ),
+        (exceptions.PapermillMissingParameterException, ("PapermillMissingParameterException",)),
+        (exceptions.AwsError, ("AwsError",)),
+        (exceptions.FileExistsError, ("FileExistsError",)),
+        (exceptions.PapermillException, ("PapermillException",)),
+        (exceptions.PapermillRateLimitException, ("PapermillRateLimitException",)),
+        (
+            exceptions.PapermillOptionalDependencyException,
+            ("PapermillOptionalDependencyException",),
+        ),
+    ],
+)
+def test_exceptions_are_unpickleable(temp_file, exc, args):
+    """Ensure exceptions can be unpickled"""
+    err = exc(*args)
+    with open(temp_file.name, 'wb') as fd:
+        pickle.dump(err, fd)
+    # Read the Pickled File
+    temp_file.seek(0)
+    data = temp_file.read()
+    pickled_err = pickle.loads(data)
+    assert str(pickled_err) == str(err)