mirror of
https://github.com/hanzo-ds/driver.git
synced 2026-10-30 20:32:40 +00:00
Kim GustyrandGitHub
4a8514c68e
* Make columnar block-extension test deterministic test_columnar_block_extend relied on its two single-row inserts reaching the client as two blocks. A background merge between the inserts and the select lets the server send one block, silently skipping the column-extension path in QueryResult.store. Force one row per block with max_block_size so the path is always exercised. beep boop * Cover sparse serialization without relying on the server Whether a block arrives in sparse form is the server's choice and varies by version and run, so the integration tests cover SparseSerialization only when the server cooperates. Feed read_sparse and apply_sparse crafted buffers directly so every line is exercised deterministically on every job. beep boop
260 lines
9.6 KiB
Python
260 lines
9.6 KiB
Python
import types
|
|
from unittest.mock import patch
|
|
|
|
from clickhouse_driver.errors import ServerException
|
|
from tests.testcase import BaseTestCase, file_config
|
|
from tests.util import capture_logging
|
|
from clickhouse_driver.util.helpers import chunks
|
|
|
|
|
|
class BlocksTestCase(BaseTestCase):
|
|
|
|
def test_return_totals_extremes(self):
|
|
rv = self.client.execute(
|
|
'SELECT a, sum(b + a) FROM ('
|
|
'SELECT arrayJoin(range(3)) - 1 AS a,'
|
|
'arrayJoin(range(4)) AS b'
|
|
') AS t '
|
|
'GROUP BY a WITH TOTALS '
|
|
'ORDER BY a',
|
|
settings={'extremes': 1}
|
|
)
|
|
self.assertEqual(rv, [
|
|
(-1, 2),
|
|
(0, 6),
|
|
(1, 10),
|
|
|
|
# TOTALS
|
|
(0, 18),
|
|
|
|
# EXTREMES
|
|
(-1, 2),
|
|
(1, 10)
|
|
])
|
|
|
|
def test_columnar_result(self):
|
|
rv = self.client.execute(
|
|
'SELECT a, sum(b + a) FROM ('
|
|
'SELECT arrayJoin(range(3)) - 1 AS a,'
|
|
'arrayJoin(range(4)) AS b'
|
|
') AS t '
|
|
'GROUP BY a '
|
|
'ORDER BY a',
|
|
columnar=True
|
|
)
|
|
self.assertEqual(rv, [
|
|
(-1, 0, 1),
|
|
(2, 6, 10)
|
|
])
|
|
|
|
def test_columnar_block_extend(self):
|
|
with self.create_table('a Int32'):
|
|
self.client.execute('INSERT INTO test (a) VALUES', [(1, )])
|
|
self.client.execute('INSERT INTO test (a) VALUES', [(2, )])
|
|
|
|
query = 'SELECT * FROM test ORDER BY a'
|
|
|
|
# One row per block even if the two parts get merged.
|
|
inserted = self.client.execute(
|
|
query, columnar=True, settings={'max_block_size': 1}
|
|
)
|
|
self.assertEqual(inserted, [(1, 2)])
|
|
|
|
def test_select_with_column_types(self):
|
|
rv = self.client.execute(
|
|
'SELECT CAST(1 AS Int32) AS x', with_column_types=True
|
|
)
|
|
self.assertEqual(rv, ([(1,)], [('x', 'Int32')]))
|
|
|
|
def test_select_with_columnar_with_column_types(self):
|
|
progress = self.client.execute_with_progress(
|
|
'SELECT arrayJoin(A) -1 as j,'
|
|
'arrayJoin(A)+1 as k FROM('
|
|
'SELECT range(3) as A)',
|
|
columnar=True, with_column_types=True)
|
|
rv = ([(-1, 0, 1), (1, 2, 3)], [('j', 'Int16'), ('k', 'UInt16')])
|
|
self.assertEqual(progress.get_result(), rv)
|
|
|
|
def test_close_connection_on_keyboard_interrupt(self):
|
|
connection = self.client.connection
|
|
with self.assertRaises(KeyboardInterrupt):
|
|
with patch.object(connection, 'send_query') as mocked_send_query:
|
|
mocked_send_query.side_effect = KeyboardInterrupt
|
|
self.client.execute('SELECT 1')
|
|
|
|
self.assertFalse(self.client.connection.connected)
|
|
|
|
|
|
class ProgressTestCase(BaseTestCase):
|
|
def test_select_with_progress(self):
|
|
progress = self.client.execute_with_progress('SELECT 2')
|
|
# CH 25 reports 1 byte per progress chunk for trivial selects;
|
|
# earlier versions reported 0.
|
|
if self.server_version >= (25,):
|
|
expected = [(1, 1), (1, 1)]
|
|
elif self.server_version > (20,):
|
|
expected = [(1, 0), (1, 0)]
|
|
else:
|
|
expected = [(1, 0)]
|
|
self.assertEqual(list(progress), expected)
|
|
self.assertEqual(progress.get_result(), [(2,)])
|
|
self.assertTrue(self.client.connection.connected)
|
|
|
|
def test_progress_totals(self):
|
|
progress = self.client.execute_with_progress('SELECT 2')
|
|
self.assertEqual(progress.progress_totals.rows, 0)
|
|
self.assertEqual(progress.progress_totals.bytes, 0)
|
|
self.assertEqual(progress.progress_totals.total_rows, 0)
|
|
|
|
self.assertEqual(progress.get_result(), [(2,)])
|
|
|
|
self.assertEqual(progress.progress_totals.rows, 1)
|
|
self.assertEqual(progress.progress_totals.bytes, 1)
|
|
# CH 25 sets total_rows for trivial selects to 1; earlier
|
|
# versions left it at 0.
|
|
expected_total_rows = 1 if self.server_version >= (25,) else 0
|
|
self.assertEqual(
|
|
progress.progress_totals.total_rows, expected_total_rows)
|
|
|
|
def test_select_with_progress_error(self):
|
|
with self.assertRaises(ServerException):
|
|
progress = self.client.execute_with_progress('SELECT error')
|
|
list(progress)
|
|
self.assertFalse(self.client.connection.connected)
|
|
|
|
def test_select_with_progress_no_progress_unwind(self):
|
|
progress = self.client.execute_with_progress('SELECT 2')
|
|
self.assertEqual(progress.get_result(), [(2,)])
|
|
self.assertTrue(self.client.connection.connected)
|
|
|
|
def test_select_with_progress_cancel(self):
|
|
self.client.execute_with_progress('SELECT 2')
|
|
rv = self.client.cancel()
|
|
self.assertEqual(rv, [(2,)])
|
|
self.assertTrue(self.client.connection.connected)
|
|
|
|
def test_select_with_progress_cancel_with_column_types(self):
|
|
self.client.execute_with_progress('SELECT CAST(2 AS Int32) as x')
|
|
rv = self.client.cancel(with_column_types=True)
|
|
self.assertEqual(rv, ([(2,)], [('x', 'Int32')]))
|
|
self.assertTrue(self.client.connection.connected)
|
|
|
|
def test_select_with_progress_with_params(self):
|
|
progress = self.client.execute_with_progress(
|
|
'SELECT %(x)s', params={'x': 2}
|
|
)
|
|
self.assertEqual(progress.get_result(), [(2,)])
|
|
self.assertTrue(self.client.connection.connected)
|
|
|
|
def test_close_connection_on_keyboard_interrupt(self):
|
|
connection = self.client.connection
|
|
with self.assertRaises(KeyboardInterrupt):
|
|
with patch.object(connection, 'send_query') as mocked_send_query:
|
|
mocked_send_query.side_effect = KeyboardInterrupt
|
|
self.client.execute_with_progress('SELECT 1')
|
|
|
|
self.assertFalse(self.client.connection.connected)
|
|
|
|
|
|
class IteratorTestCase(BaseTestCase):
|
|
def test_select_with_iter(self):
|
|
result = self.client.execute_iter(
|
|
'SELECT number FROM system.numbers LIMIT 10'
|
|
)
|
|
self.assertIsInstance(result, types.GeneratorType)
|
|
self.assertEqual(list(result), list(zip(range(10))))
|
|
self.assertEqual(list(result), [])
|
|
|
|
def test_select_with_chunk_one_iter(self):
|
|
result = self.client.execute_iter(
|
|
'SELECT number FROM system.numbers LIMIT 10',
|
|
chunk_size=1
|
|
)
|
|
self.assertIsInstance(result, types.GeneratorType)
|
|
self.assertEqual(list(result), list(zip(range(10))))
|
|
self.assertEqual(list(result), [])
|
|
|
|
def test_select_with_chunk_some_iter(self):
|
|
result = self.client.execute_iter(
|
|
'SELECT number FROM system.numbers LIMIT 10',
|
|
chunk_size=3
|
|
)
|
|
self.assertIsInstance(result, types.GeneratorType)
|
|
self.assertEqual(list(result), list(chunks(zip(range(10)), 3)))
|
|
self.assertEqual(list(result), [])
|
|
|
|
def test_select_with_iter_with_column_types(self):
|
|
result = self.client.execute_iter(
|
|
'SELECT CAST(number AS UInt32) as number '
|
|
'FROM system.numbers LIMIT 10',
|
|
with_column_types=True
|
|
)
|
|
self.assertIsInstance(result, types.GeneratorType)
|
|
|
|
self.assertEqual(
|
|
list(result),
|
|
[[('number', 'UInt32')]] + list(zip(range(10)))
|
|
)
|
|
self.assertEqual(list(result), [])
|
|
|
|
def test_select_with_iter_error(self):
|
|
with self.assertRaises(ServerException):
|
|
result = self.client.execute_iter('SELECT error')
|
|
|
|
self.assertIsInstance(result, types.GeneratorType)
|
|
list(result)
|
|
|
|
self.assertFalse(self.client.connection.connected)
|
|
|
|
def test_close_connection_on_keyboard_interrupt(self):
|
|
connection = self.client.connection
|
|
with self.assertRaises(KeyboardInterrupt):
|
|
with patch.object(connection, 'send_query') as mocked_send_query:
|
|
mocked_send_query.side_effect = KeyboardInterrupt
|
|
self.client.execute_iter('SELECT 1')
|
|
|
|
self.assertFalse(self.client.connection.connected)
|
|
|
|
|
|
class LogTestCase(BaseTestCase):
|
|
required_server_version = (18, 12, 13)
|
|
|
|
def test_logs(self):
|
|
with capture_logging('clickhouse_driver.log', 'INFO') as buffer:
|
|
settings = {'send_logs_level': 'debug'}
|
|
query = 'SELECT 1'
|
|
self.client.execute(query, settings=settings)
|
|
self.assertIn(query, buffer.getvalue())
|
|
|
|
def test_logs_insert(self):
|
|
with capture_logging('clickhouse_driver.log', 'INFO') as buffer:
|
|
with self.create_table('a Int32'):
|
|
settings = {'send_logs_level': 'debug'}
|
|
|
|
query = 'INSERT INTO test (a) VALUES'
|
|
self.client.execute(query, [(1, )], settings=settings)
|
|
logs = buffer.getvalue()
|
|
self.assertIn(query, logs)
|
|
|
|
# Test all packets of INSERT query are consumed.
|
|
rv = self.client.execute('SELECT 1', settings=settings)
|
|
self.assertEqual(rv, [(1, )])
|
|
|
|
def test_logs_with_compression(self):
|
|
compression = 'lz4'
|
|
supported_compressions = (
|
|
file_config.get('db', 'compression').split(',')
|
|
)
|
|
|
|
if compression not in supported_compressions:
|
|
self.skipTest(
|
|
'Compression {} is not supported'.format(compression)
|
|
)
|
|
|
|
with self.created_client(compression='lz4') as client:
|
|
with capture_logging('clickhouse_driver.log', 'INFO') as buffer:
|
|
settings = {'send_logs_level': 'debug'}
|
|
query = 'SELECT 1'
|
|
client.execute(query, settings=settings)
|
|
self.assertIn(query, buffer.getvalue())
|