class BasicTestCase(ClickHouseConnectionEnvVarTestCase):
_hook: ClickHouseHook
def setUp(self):
self._hook = ClickHouseHook()
def test_connection_recreated(self):
temp_table_name = 'test_temp_table'
result = self._hook.run((
f'CREATE TEMPORARY TABLE {temp_table_name} (test_field UInt8)',
f'INSERT INTO {temp_table_name} '
f'SELECT number FROM system.numbers WHERE number < 5 LIMIT 5',
f'SELECT SUM(test_field) FROM {temp_table_name}',
))
self.assertListEqual([(10,)], result)
try:
# a new connection is created and temp table is absent
self._hook.run(f'SELECT * FROM {temp_table_name}')
except ServerException as err:
self.assertEqual(ErrorCodes.UNKNOWN_TABLE, err.code)
else:
raise AssertionError('server did not raise an error')
def test_get_first(self):
self.assertTupleEqual((1,), self._hook.get_first('SELECT 1'))
self.assertIsNone(self._hook.get_first('SELECT 1 WHERE 0'))
def _test(self, sql: str, expected_df):
actual_df = ClickHouseHook().get_pandas_df(sql)
self.assertListEqual(list(actual_df.columns), list(expected_df.columns))
self.assertListEqual(
actual_df.to_dict('records'),
expected_df.to_dict('records'),
)
def execute(self, context: Dict[str, Any]) -> Any:
hook = ClickHouseHook(
clickhouse_conn_id=self._conn_id,
database=self._database,
)
return hook.run(self._sql, self._parameters)
def _get_hook(self) -> ClickHouseHook:
return ClickHouseHook(
clickhouse_conn_id=self.conn_id,
database=self._database,
)
def _test_client_initiated_with(self, *args, **kwargs) -> None:
self._client_mock.reset_mock()
ClickHouseHook().get_conn() # instantiate connection
self._assert_client_initiated_with(*args, **kwargs)
def test_database(self):
database = 'test-db-from-init-args'
ClickHouseHook(database=database).get_conn()
self._assert_client_initiated_with(database=database)
def setUp(self) -> None:
self.hook = ClickHouseHook()
class HookLogQueryTestCase(TestCase):
def setUp(self) -> None:
self.hook = ClickHouseHook()
def test_log_params_dict(self):
self.assertEqual('{}', self.hook._log_params({}))
self.assertEqual('{1: 1}', self.hook._log_params({1: 1}))
self.assertEqual('{1: 1}', self.hook._log_params({1: 1}, limit=1))
self.assertEqual(
'{1: 1 … and 1 more parameters}',
self.hook._log_params({1: 1, 2: 2}, limit=1),
)
self.assertEqual(
'{1: 1, 2: 2}',
self.hook._log_params({1: 1, 2: 2}),
)
self.assertEqual(
'{0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9 …'
' and 1 more parameters}',
self.hook._log_params({k: k for k in range(11)}),
)
self.assertEqual(
'{0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9 …'
' and 10 more parameters}',
self.hook._log_params({k: k for k in range(20)}),
)
self.assertEqual(
'{0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9 …'
' and 10 more parameters}',
self.hook._log_params({k: k for k in range(20)}, limit=10),
)
def test_log_params_generator(self):
def gen():
yield
g = gen()
self.assertEqual(str(g), self.hook._log_params(g))
def test_log_params_tuple(self):
self.assertEqual('()', self.hook._log_params(()))
self.assertEqual('(1,)', self.hook._log_params((1, )))
self.assertEqual('(1,)', self.hook._log_params((1, ), limit=1))
self.assertEqual(
'(1, … and 1 more parameters)',
self.hook._log_params((1, 2), limit=1),
)
self.assertEqual(
'(1, 2)',
self.hook._log_params((1, 2)),
)
self.assertEqual(
'(0, 1, 2, 3, 4, 5, 6, 7, 8, 9 … and 1 more parameters)',
self.hook._log_params(tuple(range(11))),
)
self.assertEqual(
'(0, 1, 2, 3, 4, 5, 6, 7, 8, 9 … and 10 more parameters)',
self.hook._log_params(tuple(range(20))),
)
self.assertEqual(
'(0, 1, 2, 3, 4, 5, 6, 7, 8, 9 … and 10 more parameters)',
self.hook._log_params(tuple(range(20)), limit=10),
)
def test_log_params_list(self):
self.assertEqual('[]', self.hook._log_params([]))
self.assertEqual('[1]', self.hook._log_params([1]))
self.assertEqual('[1]', self.hook._log_params([1], limit=1))
self.assertEqual(
'[1 … and 1 more parameters]',
self.hook._log_params([1, 2], limit=1),
)
self.assertEqual(
'[1, 2]',
self.hook._log_params([1, 2]),
)
self.assertEqual(
'[0, 1, 2, 3, 4, 5, 6, 7, 8, 9 … and 1 more parameters]',
self.hook._log_params(list(range(11))),
)
self.assertEqual(
'[0, 1, 2, 3, 4, 5, 6, 7, 8, 9 … and 10 more parameters]',
self.hook._log_params(list(range(20))),
)
self.assertEqual(
'[0, 1, 2, 3, 4, 5, 6, 7, 8, 9 … and 10 more parameters]',
self.hook._log_params(list(range(20)), limit=10),
)
def test_log_query(self):
_ = self.hook.log # to initialize .log property
with mock.patch.object(self.hook, '_log') as patched:
self.hook._log_query('SELECT 1', {})
patched.info.assert_called_with('%s%s', 'SELECT 1', '')
self.hook._log_query('SELECT 1', {1: 1})
patched.info.assert_called_with('%s%s', 'SELECT 1', ' with {1: 1}')
self.hook._log_query('SELECT 1', [1])
patched.info.assert_called_with('%s%s', 'SELECT 1', ' with [1]')
def setUp(self):
self._hook = ClickHouseHook()