From 57112f9de600c820dcdb1fa4297f188e61bd179e Mon Sep 17 00:00:00 2001 From: Ivan Ladelshchikov Date: Mon, 6 Nov 2017 21:58:10 +0500 Subject: [PATCH] add ServerError exception --- src/infi/clickhouse_orm/database.py | 47 +++++++++++++++++++++++++++-- tests/test_database.py | 14 +++++++-- tests/test_readonly.py | 21 +++++++++---- 3 files changed, 70 insertions(+), 12 deletions(-) diff --git a/src/infi/clickhouse_orm/database.py b/src/infi/clickhouse_orm/database.py index cb47d02..92bcb09 100644 --- a/src/infi/clickhouse_orm/database.py +++ b/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\d+), + \ e\.displayText\(\)\ =\ (?P[^ \n]+):\ (?P.+?), + \ e.what\(\)\ =\ (?P[^ \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): diff --git a/tests/test_database.py b/tests/test_database.py index 0214f36..8e7c5c4 100644 --- a/tests/test_database.py +++ b/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) diff --git a/tests/test_readonly.py b/tests/test_readonly.py index facbaa0..b1aa595 100644 --- a/tests/test_readonly.py +++ b/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')