Merge pull request #169 from azat-archive/Exception-picke-fix

Make Exception picklable
mymarilyn · Nov 2, 2020 · ab7dde0 · ab7dde0
2 parents d7971c0 + b89f364
commit ab7dde0
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/clickhouse_driver/errors.py b/clickhouse_driver/errors.py
@@ -382,16 +382,18 @@ class Error(Exception):
 
         def __init__(self, message=None):
             self.message = message
+            super().__init__(message)
 
         def __str__(self):
             message = ' ' + self.message if self.message is not None else ''
             return 'Code: {}.{}'.format(self.code, message)
 
     class ServerException(Error):
-        def __init__(self, message, code, nested=None):
+        def __init__(self, message, code=None, nested=None):
             self.message = message
             self.code = code
             self.nested = nested
+            super().__init__(message)
 
         def __str__(self):
             nested = '\nNested: {}'.format(self.nested) if self.nested else ''
@@ -403,6 +405,7 @@ class Error(Exception):
 
         def __init__(self, message=None):
             self.message = message
+            super(Error, self).__init__(message)
 
         def __unicode__(self):
             message = ' ' + self.message if self.message is not None else ''
@@ -412,10 +415,11 @@ def __str__(self):
             return compat.text_type(self).encode('utf-8')
 
     class ServerException(Error):
-        def __init__(self, message, code, nested=None):
+        def __init__(self, message, code=None, nested=None):
             self.message = message
             self.code = code
             self.nested = nested
+            super(ServerException, self).__init__(message)
 
         def __unicode__(self):
             nested = '\nNested: {}'.format(self.nested) if self.nested else ''

diff --git a/tests/test_errors.py b/tests/test_errors.py
@@ -0,0 +1,14 @@
+import pickle
+import clickhouse_driver.errors as err
+
+def picklable(o):
+    picked = pickle.loads(pickle.dumps(o))
+    assert repr(o) == repr(picked)
+    assert str(o) == str(picked)
+
+def test_exception_picklable():
+    picklable(err.Error('foo'))
+    picklable(err.Error(message='foo'))
+
+    picklable(err.ServerException('foo', 0, Exception()))
+    picklable(err.ServerException(message='foo', code=0, nested=Exception()))