mirror of
https://github.com/hanzo-ds/driver.git
synced 2026-10-30 20:32:40 +00:00
zeekay
eb74c996a6
Imports read `datastore_driver`, distribution `datastore-driver`, so nothing has to import another product's name to talk to Hanzo Datastore. Cut from upstream 0.2.9 — the version this stack already ran — rather than from upstream HEAD, so it is a PURE RENAME with no behaviour change. Forking from HEAD would have dragged versions along with the name (connect alone jumped 0.8.17 -> 1.6.0, a major), which is a different and much riskier change wearing a rename's clothes. X-ClickHouse-* headers are untouched: they are wire protocol the server sends and expects, and a different token shape from the module name.
248 lines
9.0 KiB
Python
248 lines
9.0 KiB
Python
import types
|
|
from unittest.mock import patch
|
|
|
|
from datastore_driver.errors import ServerException
|
|
from tests.testcase import BaseTestCase, file_config
|
|
from tests.util import capture_logging
|
|
from datastore_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'
|
|
|
|
inserted = self.client.execute(query, columnar=True)
|
|
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')
|
|
self.assertEqual(
|
|
list(progress),
|
|
[(1, 0), (1, 0)] if self.server_version > (20,) else [(1, 0)]
|
|
)
|
|
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)
|
|
self.assertEqual(progress.progress_totals.total_rows, 0)
|
|
|
|
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('datastore_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('datastore_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('datastore_driver.log', 'INFO') as buffer:
|
|
settings = {'send_logs_level': 'debug'}
|
|
query = 'SELECT 1'
|
|
client.execute(query, settings=settings)
|
|
self.assertIn(query, buffer.getvalue())
|