diff --git a/src/airflow_clickhouse_plugin/hooks/clickhouse.py b/src/airflow_clickhouse_plugin/hooks/clickhouse.py index 698352e..63a2e61 100644 --- a/src/airflow_clickhouse_plugin/hooks/clickhouse.py +++ b/src/airflow_clickhouse_plugin/hooks/clickhouse.py @@ -90,6 +90,15 @@ def execute( return last_result +def strtobool(str_value: str) -> bool: + str_value = str_value.lower() + if str_value == 'true': + return True + elif str_value == 'false': + return False + raise ValueError(f'unsupported string value: {str_value!r}') + + def conn_to_kwargs(conn: Connection, database: t.Optional[str]) -> t.Dict[str, t.Any]: """ Translate Airflow Connection to clickhouse-driver Connection kwargs. """ connection_kwargs = conn.extra_dejson.copy() @@ -105,6 +114,13 @@ def conn_to_kwargs(conn: Connection, database: t.Optional[str]) -> t.Dict[str, t connection_kwargs.update(database=database) elif conn.schema: connection_kwargs.update(database=conn.schema) + # converting types for flags + for key in ('secure', 'verify'): + if key in connection_kwargs: + current_value = connection_kwargs[key] + if isinstance(current_value, (bool, int)): + continue + connection_kwargs[key] = strtobool(current_value) return connection_kwargs diff --git a/tests/unit/hooks/test_clickhouse.py b/tests/unit/hooks/test_clickhouse.py index ee08355..e419276 100644 --- a/tests/unit/hooks/test_clickhouse.py +++ b/tests/unit/hooks/test_clickhouse.py @@ -19,7 +19,7 @@ def test_arguments(self): login='test-login', password='test-pass', schema='test-schema', - extra='{"test_extra": "test-extra-value"}', + extra='{"test_extra": "test-extra-value", "secure": "true", "verify": "false"}', ) return_value = ClickHouseHook( @@ -47,6 +47,8 @@ def test_arguments(self): password='test-pass', database='test-database', test_extra='test-extra-value', + secure=True, + verify=False, ) for query, mock_call \ diff --git a/tests/unit/hooks/test_strtobool.py b/tests/unit/hooks/test_strtobool.py new file mode 100644 index 0000000..e9273bc --- /dev/null +++ b/tests/unit/hooks/test_strtobool.py @@ -0,0 +1,23 @@ +import unittest + +from airflow_clickhouse_plugin.hooks.clickhouse import strtobool + + +class StrToBoolTestCase(unittest.TestCase): + def test_correct_true(self): + for v in ['true', 'True']: + with self.subTest('Testing %s value' % v, v=v): + self.assertTrue(strtobool(v)) + + def test_correct_false(self): + for v in ['false', 'False']: + with self.subTest('Testing %s value' % v, v=v): + self.assertFalse(strtobool(v)) + + def test_unknown_throws(self): + with self.assertRaises(ValueError): + strtobool('unknown') + + +if __name__ == '__main__': + unittest.main()