Merge pull request #11 from ods/refactor_exceptions

Richer exception interface to simplify debugging huge inserts

Author: ods

aiochsa/client.py

@@ -4,10 +4,10 @@
 
 import aiohttp
 
-from .compiler import Compiler
+from .compiler import Compiler, Statement
 from .dialect import ClickhouseSaDialect
-from .exc import DBException, ProtocolError
-from .parser import parse_json_compact
+from .exc import DBException, ProtocolError, exc_message_re
+from .parser import parse_json_compact, JSONDecodeError
 from .record import Record
 from .types import TypeRegistry
 
@@ -43,21 +43,26 @@ def __init__(
         self._types = types
         self._compiler = Compiler(dialect=dialect, escape=types.escape)
 
-    async def _execute(self, statement: str, *args) -> Iterable[Record]:
-        query, json_each_row_parameters = self._compiler.compile_statement(
+    async def _execute(self, statement: Statement, *args) -> Iterable[Record]:
+        compiled, json_each_row_parameters = self._compiler.compile_statement(
             statement, args,
         )
+        sql_logger.debug(compiled)
+        compiled_with_params = compiled
+        rows = None
         if json_each_row_parameters:
             to_json = self._types.to_json # lookup optimization
-            query += '\n'
-            query += '\n'.join(
+            rows = [
                 json.dumps(
                     {name: to_json(value) for name, value in row.items()},
                     use_decimal=True,
                 )
                 for row in json_each_row_parameters
-            )
-        sql_logger.debug(query)
+            ]
+            if sql_logger.isEnabledFor(logging.DEBUG):
+                for idx, row in enumerate(rows):
+                    sql_logger.debug(f'{idx}: {row}')
+            compiled_with_params += '\n' + '\n'.join(rows)
 
         # First attempt may fail due to broken state of aiohttp session
         # (aiohttp doesn't handle connection closing properly?)
@@ -66,18 +71,27 @@ async def _execute(self, statement: str, *args) -> Iterable[Record]:
                 async with self._session.post(
                     self.url,
                     params = {'default_format': 'JSONCompact', **self.params},
-                    data = query.encode(),
+                    data = compiled_with_params.encode(),
                 ) as response:
+                    body = await response.read()
                     if response.status != 200:
-                        body = await response.read()
                         raise DBException.from_message(
-                            query, body.decode(errors='replace'),
+                            body.decode(errors='replace'),
+                            statement=compiled, rows=rows,
                         )
 
-                    if response.content_type == 'application/json':
-                        return await parse_json_compact(
-                            self._types, response.content,
-                        )
+                    elif response.content_type == 'application/json':
+                        try:
+                            return parse_json_compact(self._types, body)
+                        except JSONDecodeError:
+                            body_str = body.decode(errors='replace')
+                            m = exc_message_re.search(body_str)
+                            if not m:
+                                raise
+                            raise DBException.from_message(
+                                body_str[m.start():],
+                                statement=compiled, rows=rows,
+                            )
                     else:
                         return ()
             except aiohttp.ClientError as exc:
@@ -86,22 +100,22 @@ async def _execute(self, statement: str, *args) -> Iterable[Record]:
                 logger.debug(f'First attempt failed, retrying (error: {exc})')
 
     async def iterate(
-        self, statement: str, *args,
+        self, statement: Statement, *args,
     ) -> AsyncGenerator[Record, None]:
         for row in await self._execute(statement, *args):
             yield row
 
-    async def execute(self, statement: str, *args) -> None:
+    async def execute(self, statement: Statement, *args) -> None:
         await self._execute(statement, *args)
 
-    async def fetch(self, statement: str, *args) -> List[Record]:
+    async def fetch(self, statement: Statement, *args) -> List[Record]:
         return list(await self._execute(statement, *args))
 
-    async def fetchrow(self, statement: str, *args) -> Optional[Record]:
+    async def fetchrow(self, statement: Statement, *args) -> Optional[Record]:
         gen = await self._execute(statement, *args)
         return next(iter(gen), None)
 
-    async def fetchval(self, statement: str, *args) -> Any:
+    async def fetchval(self, statement: Statement, *args) -> Any:
         row = await self.fetchrow(statement, *args)
         if row is not None:
             return row[0]

aiochsa/compiler.py

@@ -1,4 +1,5 @@
 from types import SimpleNamespace
+from typing import Union
 
 from sqlalchemy.engine.util import _distill_params
 from sqlalchemy.sql import func, ClauseElement
@@ -7,6 +8,9 @@
 from sqlalchemy.sql.functions import FunctionElement
 
 
+Statement = Union[str, ClauseElement]
+
+
 class Compiler:
 
     def __init__(self, dialect, escape):
@@ -91,7 +95,7 @@ def _execute_context(self, dialect, constructor, statement, parameters, *args):
             return context.statement % escaped, ()
 
 
-    def compile_statement(self, statement, args):
+    def compile_statement(self, statement: Statement, args):
         if isinstance(statement, str):
             assert not args
             return statement, args

aiochsa/exc.py

@@ -1,15 +1,17 @@
+from collections import namedtuple
 import re
 
 
 # Based on `getExceptionMessage()`:
-# https://github.com/yandex/ClickHouse/blob/master/dbms/src/Common/Exception.cpp#L261
-_match_exc_message = re.compile(
-    r'^Code: (?P<code>\d+), '
+# https://github.com/ClickHouse/ClickHouse/blob/master/src/Common/Exception.cpp#L370-L373
+exc_message_re = re.compile(
+    r'Code: (?P<code>\d+), '
         r'e\.displayText\(\) = (?P<display_text>.+?)'
         r'(?:, Stack trace[^:]*:\s+(?P<stack_trace>.+))?$',
     re.M,
-).match
+)
 
+at_row_re = re.compile(r'[(]at row (?P<num>\d+)[)]')
 
 class AiochsaException(Exception):
     """ Base class for aiochsa exceptions """
@@ -19,33 +21,56 @@ class ProtocolError(AiochsaException):
     """ Error communicating to Clickhouse server """
 
 
+RowInfo = namedtuple('RowInfo', ['num', 'content'])
+
+
 class DBException(AiochsaException):
     """ Error returned from Clickhouse database """
 
-    def __init__(self, statement, code, display_text, stack_trace=None):
-        super().__init__(statement, code, display_text, stack_trace)
-        self.statement = statement
+    def __init__(
+        self, code, display_text, stack_trace=None, statement=None, row=None,
+    ):
+        super().__init__(code, display_text, stack_trace)
         self.code = code
         self.display_text = display_text
         self.stack_trace = stack_trace
+        self.statement = statement
+        self.row = row
 
     def __str__(self):
-        statement = self.statement
-        if len(statement) > 200:
-            statement = statement[:200] + '...'
-        return f'[Code={self.code}] {self.display_text}: {statement}'
+        message = f'[Code={self.code}] {self.display_text}'
+        if self.statement:
+            statement = self.statement
+            if len(statement) > 200:
+                statement = statement[:200] + '...'
+            message += f'\n{statement}'
+        if self.row:
+            message += f'\n{self.row.num}: {self.row.content}'
+        return message
 
     @classmethod
-    def from_message(cls, statement, exc_message):
-        m = _match_exc_message(exc_message)
+    def from_message(cls, exc_message, *, statement=None, rows=None):
+        m = exc_message_re.match(exc_message)
         if m:
+            display_text = m.group('display_text')
+
+            row = None
+            if rows:
+                at_row_m = at_row_re.search(display_text)
+                if at_row_m:
+                    # It's 1-based
+                    row_num = int(at_row_m.group('num'))
+                    if len(rows) >= row_num:
+                        row = RowInfo(row_num, rows[row_num - 1])
+
             return cls(
-                statement=statement,
                 code=int(m.group('code')),
-                display_text=m.group('display_text'),
+                display_text=display_text,
                 stack_trace=m.group('stack_trace'),
+                statement=statement,
+                row=row,
             )
         else: # pragma: nocover
             return cls(
-                statement=statement, code=None, display_text=exc_message,
+                code=None, display_text=exc_message, statement=statement,
             )

aiochsa/parser.py

@@ -1,16 +1,18 @@
-import asyncio
 import pkgutil
 import simplejson as json
-from typing import AsyncGenerator, Iterable, Union
+from typing import Iterable
 
-import aiohttp
 from lark import Lark, Transformer, v_args
 
 from .record import Record
 from .types import TypeRegistry
 
 
-__all__ = ['parse_type']
+__all__ = ['parse_type', 'parse_json_compact', 'JSONDecodeError']
+
+
+# Re-export
+JSONDecodeError = json.JSONDecodeError
 
 
 type_parser = Lark(
@@ -43,19 +45,18 @@ def parse_type(types: TypeRegistry, type_str):
     return TypeTransformer(types).transform(tree)
 
 
-async def parse_json_compact(
-    types: TypeRegistry,
-    content: Union[asyncio.StreamReader, aiohttp.StreamReader],
+def parse_json_compact(
+    types: TypeRegistry, content: bytes,
 ) -> Iterable[Record]:
     # The method is split into three phases:
-    #   1. Read whole response.  It's done immediately after await-ing, so we
-    #      can close response after it.
+    #   1. Parse JSON.  It's done immediately, so that we can fall back to
+    #      parsing exception when it breaks normal response.
     #   2. Parse type information from meta.  It's done at first iteration of
-    #      returned async-generator.
+    #      returned generator.
     #   3. Convert each row one-by-one.  It's done on demand at each iteration.
     # This way we can cleanup resources even when result is not used.
 
-    json_data = json.loads(await content.read(), parse_float=str)
+    json_data = json.loads(content, parse_float=str)
     return convert_json_compact(types, json_data)
 
 

tests/test_exc.py

@@ -1,3 +1,5 @@
+from decimal import Decimal
+
 import pytest
 
 import aiochsa
@@ -33,3 +35,13 @@ async def test_exc_stacktrace(dsn, statement):
     assert str(error_codes.SYNTAX_ERROR) in str(exc_info.value)
     assert 'Syntax error' in str(exc_info.value)
     assert len(str(exc_info.value)) < 2000
+
+
+async def test_exc_row(conn, table_test1):
+    with pytest.raises(aiochsa.DBException) as exc_info:
+        await conn.execute(
+            table_test1.insert(),
+            {'amount': Decimal('1234567890.1234567890')},
+        )
+    assert exc_info.value.code == error_codes.ARGUMENT_OUT_OF_BOUND
+    assert exc_info.value.row == (1, '{"amount": 1234567890.1234567890}')

tests/test_parser.py

@@ -74,7 +74,7 @@ def test_parse_type(type_str, type_obj):
 
 async def test_parse_json_compact_decimal():
     """ Insure Decimal is parsed without loss of precision """
-    json_data = b'''\
+    content = b'''\
         {
             "meta": [
                 {"name": "column", "type": "Decimal(18, 13)"}
@@ -84,10 +84,6 @@ async def test_parse_json_compact_decimal():
             ]
         }
     '''
-    content = asyncio.StreamReader()
-    content.feed_data(json_data)
-    content.feed_eof()
-
-    [[value]] = list(await parse_json_compact(t.TypeRegistry(), content))
+    [[value]] = list(parse_json_compact(t.TypeRegistry(), content))
     assert isinstance(value, Decimal)
     assert str(value) == '1.2345678901230'

tests/test_pool.py

@@ -1,7 +1,7 @@
 import asyncio
 import pytest
 
-from aiochsa import DBException
+from aiochsa import DBException, error_codes
 from aiochsa.pool import dsn_to_params, create_pool
 
 
@@ -43,20 +43,25 @@ async def test_pool_acquire_release(pool):
     await pool.release(conn)
 
 
+LONG_QUERY = 'SELECT sleep(1) FROM numbers(10) SETTINGS max_block_size=1'
+
+
 async def test_pool_params(dsn):
     async with create_pool(dsn, max_execution_time=1) as conn:
-        with pytest.raises(DBException):
-            await conn.execute('SELECT sleep(3)')
+        with pytest.raises(DBException) as exc_info:
+            await conn.execute(LONG_QUERY)
+    assert exc_info.value.code == error_codes.TIMEOUT_EXCEEDED
 
 
 async def test_pool_dsn_params(dsn):
     dsn += '?max_execution_time=1'
     async with create_pool(dsn) as conn:
-        with pytest.raises(DBException):
-            await conn.execute('SELECT sleep(3)')
+        with pytest.raises(DBException) as exc_info:
+            await conn.execute(LONG_QUERY)
+    assert exc_info.value.code == error_codes.TIMEOUT_EXCEEDED
 
 
 async def test_session_timeout(dsn):
     async with create_pool(dsn, session_timeout=0.1) as conn:
         with pytest.raises(asyncio.TimeoutError):
-            await conn.execute('SELECT sleep(1)')
+            await conn.execute(LONG_QUERY)