add ServerError exception

src/infi/clickhouse_orm/database.py

@@ -1,5 +1,6 @@
 from __future__ import unicode_literals
 
+import re
 import requests
 from collections import namedtuple
 from .models import ModelBase
@@ -24,6 +25,46 @@ class DatabaseException(Exception):
     pass
 
 
+class ServerError(DatabaseException):
+    """
+    Raised when a server returns an error.
+    """
+    def __init__(self, message):
+        self.code = None
+        processed = self.get_error_code_msg(message)
+        if processed:
+            self.code, self.message = processed
+        else:
+            # just skip custom init
+            # if non-standard message format
+            super(ServerError, self).__init__(message)
+
+    ERROR_PATTERN = re.compile(r'''
+        Code:\ (?P<code>\d+),
+        \ e\.displayText\(\)\ =\ (?P<type1>[^ \n]+):\ (?P<msg>.+?),
+        \ e.what\(\)\ =\ (?P<type2>[^ \n]+)
+    ''', re.VERBOSE | re.DOTALL)
+
+    @classmethod
+    def get_error_code_msg(cls, full_error_message):
+        """
+        Extract the code and message of the exception that clickhouse-server generated.
+
+        See the list of error codes here:
+        https://github.com/yandex/ClickHouse/blob/master/dbms/src/Common/ErrorCodes.cpp
+        """
+        match = cls.ERROR_PATTERN.match(full_error_message)
+        if match:
+            # assert match.group('type1') == match.group('type2')
+            return int(match.group('code')), match.group('msg')
+
+        return 0, full_error_message
+
+    def __str__(self):
+        if self.code is not None:
+            return "{} ({})".format(self.message, self.code)
+
+
 class Database(object):
     '''
     Database instances connect to a specific ClickHouse database for running queries,
@@ -250,7 +291,7 @@ class Database(object):
         params = self._build_params(settings)
         r = requests.post(self.db_url, params=params, data=data, stream=stream)
         if r.status_code != 200:
-            raise DatabaseException(r.text)
+            raise ServerError(r.text)
         return r
 
     def _build_params(self, settings):
@@ -281,8 +322,8 @@ class Database(object):
         try:
             r = self._send('SELECT timezone()')
             return pytz.timezone(r.text.strip())
-        except DatabaseException:
-            logger.exception('Cannot determine server timezone, assuming UTC')
+        except ServerError as e:
+            logger.exception('Cannot determine server timezone (%s), assuming UTC', e)
             return pytz.utc
 
     def _is_connection_readonly(self):

tests/test_database.py

@@ -2,7 +2,7 @@
 from __future__ import unicode_literals
 import unittest
 
-from infi.clickhouse_orm.database import Database, DatabaseException
+from infi.clickhouse_orm.database import ServerError
 from .base_test_with_data import *
 
 
@@ -131,14 +131,22 @@ class DatabaseTestCase(TestCaseWithData):
         self.assertEqual(results, "Whitney\tDurham\t1977-09-15\t1.72\nWhitney\tScott\t1971-07-04\t1.7\n")
 
     def test_invalid_user(self):
-        with self.assertRaises(DatabaseException):
+        with self.assertRaises(ServerError) as cm:
             Database(self.database.db_name, username='default', password='wrong')
 
+        exc = cm.exception
+        self.assertEqual(exc.code, 193)
+        self.assertEqual(exc.message, 'Wrong password for user default')
+
     def test_nonexisting_db(self):
         db = Database('db_not_here', autocreate=False)
-        with self.assertRaises(DatabaseException):
+        with self.assertRaises(ServerError) as cm:
             db.create_table(Person)
 
+        exc = cm.exception
+        self.assertEqual(exc.code, 81)
+        self.assertEqual(exc.message, "Database db_not_here doesn't exist")
+
     def test_preexisting_db(self):
         db = Database(self.database.db_name, autocreate=False)
         db.count(Person)

tests/test_readonly.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 from __future__ import unicode_literals
 
-from infi.clickhouse_orm.database import DatabaseException
+from infi.clickhouse_orm.database import DatabaseException, ServerError
 from .base_test_with_data import *
 
 
@@ -12,22 +12,31 @@ class ReadonlyTestCase(TestCaseWithData):
         orig_database = self.database
         try:
             self.database = Database(orig_database.db_name, username=username, readonly=True)
-            with self.assertRaises(DatabaseException):
+            with self.assertRaises(ServerError) as cm:
                 self._insert_and_check(self._sample_data(), len(data))
+            self._check_db_readonly_err(cm.exception)
+
             self.assertEquals(self.database.count(Person), 100)
             list(self.database.select('SELECT * from $table', Person))
-            with self.assertRaises(DatabaseException):
+            with self.assertRaises(ServerError) as cm:
                 self.database.drop_table(Person)
-            with self.assertRaises(DatabaseException):
+            self._check_db_readonly_err(cm.exception)
+
+            with self.assertRaises(ServerError) as cm:
                 self.database.drop_database()
-        except DatabaseException as e:
-            if 'Unknown user' in six.text_type(e):
+            self._check_db_readonly_err(cm.exception)
+        except ServerError as e:
+            if e.code == 192 and e.message.startswith('Unknown user'):
                 raise unittest.SkipTest('Database user "%s" is not defined' % username)
             else:
                 raise
         finally:
             self.database = orig_database
 
+    def _check_db_readonly_err(self, exc):
+        self.assertEqual(exc.code, 164)
+        self.assertEqual(exc.message, 'Cannot execute query in readonly mode')
+
     def test_readonly_db_with_default_user(self):
         self._test_readonly_db('default')