diff --git a/src/infi/clickhouse_orm/database.py b/src/infi/clickhouse_orm/database.py index 2be8534..3afd1c3 100644 --- a/src/infi/clickhouse_orm/database.py +++ b/src/infi/clickhouse_orm/database.py @@ -96,7 +96,9 @@ class Database(object): field_types = parse_tsv(next(lines)) model_class = model_class or ModelBase.create_ad_hoc_model(zip(field_names, field_types)) for line in lines: - yield model_class.from_tsv(line, field_names, self.server_timezone, self) + # skip blank line left by WITH TOTALS modifier + if line: + yield model_class.from_tsv(line, field_names, self.server_timezone, self) def raw(self, query, settings=None, stream=False): """ diff --git a/src/infi/clickhouse_orm/utils.py b/src/infi/clickhouse_orm/utils.py index c24a93d..f5b5b22 100644 --- a/src/infi/clickhouse_orm/utils.py +++ b/src/infi/clickhouse_orm/utils.py @@ -38,7 +38,7 @@ def unescape(value): def parse_tsv(line): if PY3 and isinstance(line, binary_type): line = line.decode() - if line[-1] == '\n': + if line and line[-1] == '\n': line = line[:-1] return [unescape(value) for value in line.split('\t')] diff --git a/tests/test_database.py b/tests/test_database.py index 62b340b..1897d8f 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -89,6 +89,15 @@ class DatabaseTestCase(unittest.TestCase): self.assertEqual(results[0].get_database(), self.database) self.assertEqual(results[1].get_database(), self.database) + def test_select_with_totals(self): + self._insert_and_check(self._sample_data(), len(data)) + query = "SELECT last_name, sum(height) as height FROM `test-db`.person GROUP BY last_name WITH TOTALS" + results = list(self.database.select(query)) + total = sum(r.height for r in results[:-1]) + # Last line has an empty last name, and total of all heights + self.assertFalse(results[-1].last_name) + self.assertEquals(total, results[-1].height) + def test_pagination(self): self._insert_and_check(self._sample_data(), len(data)) # Try different page sizes