Repository: bryzgaloff/airflow-clickhouse-plugin (language: python)

Issue:
ClickHouse sql sensor implementation



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 5129448..385d67b 100644
--- a/README.md
+++ b/README.md
@@ -15,10 +15,15 @@ Provides `ClickHouseHook` and `ClickHouseOperator` for [Apache Airflow][airflow]
     as various timeouts, `compression`, `secure`, etc through Airflow
     [Connection.extra][airflow-conn-extra] property.
 
-# Installation
+# Installation and dependencies
 
 `pip install -U airflow-clickhouse-plugin`
 
+Requires `apache-airflow` and `clickhouse-driver`. Primarily supports Airflow
+    1.10.6 since it is the latest version
+    [supported by Google Cloud Composer][cloud-composer-versions]: later
+    versions are expected to work properly but may be not fully tested.
+
 # Usage
 
 See [examples](#examples) below.
@@ -236,3 +241,4 @@ From the root project directory: `python -m unittest discover -s tests`
 [airflow-conn-dejson]: https://airflow.apache.org/docs/1.10.6/_api/airflow/models/index.html?highlight=connection#airflow.models.Connection.extra_dejson
 [airflow-conn-env]: https://airflow.apache.org/docs/stable/howto/connection/index.html#storing-a-connection-in-environment-variables
 [python-dbapi2-fetchone]: https://www.python.org/dev/peps/pep-0249/#fetchone
+[cloud-composer-versions]: https://cloud.google.com/composer/docs/concepts/versioning/composer-versions#supported_versions
diff --git a/airflow_clickhouse_plugin/__init__.py b/airflow_clickhouse_plugin/__init__.py
index e997e8a..f65d639 100644
--- a/airflow_clickhouse_plugin/__init__.py
+++ b/airflow_clickhouse_plugin/__init__.py
@@ -1,6 +1,7 @@
 from airflow.plugins_manager import AirflowPlugin
 from .hooks.clickhouse_hook import ClickHouseHook
 from .operators.clickhouse_operator import ClickHouseOperator
+from .sensors.clickhouse_sql_sensor import ClickHouseSqlSensor
 
 
 class ClickHouseOperatorPlugin(AirflowPlugin):
@@ -11,3 +12,8 @@ class ClickHouseOperatorPlugin(AirflowPlugin):
 class ClickHouseHookPlugin(AirflowPlugin):
     name = 'clickhouse_hook'
     hooks = [ClickHouseHook]
+
+
+class ClickHouseSqlSensorPlugin(AirflowPlugin):
+    name = 'clickhouse_sql_sensor'
+    sensors = [ClickHouseSqlSensor]
diff --git a/airflow_clickhouse_plugin/__version__.py b/airflow_clickhouse_plugin/__version__.py
index 073e645..1c1d7b9 100644
--- a/airflow_clickhouse_plugin/__version__.py
+++ b/airflow_clickhouse_plugin/__version__.py
@@ -9,7 +9,7 @@
 __description__ = \
     'airflow-clickhouse-plugin - ' \
     'Airflow plugin to execute ClickHouse commands and queries'
-__version__ = '0.5.6'
+__version__ = '0.5.7'
 __url__ = 'https://github.com/whisklabs/airflow-clickhouse-plugin'
 __license__ = 'MIT License'
 __author__ = 'Viktor Taranenko'
diff --git a/airflow_clickhouse_plugin/sensors/__init__.py b/airflow_clickhouse_plugin/sensors/__init__.py
new file mode 100644
index 0000000..89c10dc
--- /dev/null
+++ b/airflow_clickhouse_plugin/sensors/__init__.py
@@ -0,0 +1,1 @@
+from .clickhouse_sql_sensor import ClickHouseSqlSensor
diff --git a/airflow_clickhouse_plugin/sensors/clickhouse_sql_sensor.py b/airflow_clickhouse_plugin/sensors/clickhouse_sql_sensor.py
new file mode 100644
index 0000000..f2b1d57
--- /dev/null
+++ b/airflow_clickhouse_plugin/sensors/clickhouse_sql_sensor.py
@@ -0,0 +1,85 @@
+from typing import Dict, Callable, Any, Optional
+
+from airflow.exceptions import AirflowException
+from airflow.sensors.sql_sensor import SqlSensor
+from airflow_clickhouse_plugin.hooks.clickhouse_hook import ClickHouseHook
+
+
+class ClickHouseSqlSensor(SqlSensor):
+    _DEFAULT_CONN_ID = 'clickhouse_default'
+
+    def __init__(
+        self,
+        sql: str = None,
+        clickhouse_conn_id: str = _DEFAULT_CONN_ID,
+        parameters: Dict[str, Any] = None,
+        database: str = None,
+        success: Callable[[Any], bool] = None,
+        failure: Callable[[Any], bool] = None,
+        fail_on_empty: bool = False,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(
+            conn_id=clickhouse_conn_id,
+            sql=sql,
+            parameters=parameters,
+            success=success,
+            failure=failure,
+            fail_on_empty=fail_on_empty,
+            *args,
+            **kwargs,
+        )
+        self._database = database
+
+    def _get_hook(self) -> ClickHouseHook:
+        return ClickHouseHook(
+            clickhouse_conn_id=self.conn_id,
+            database=self._database,
+        )
+
+    def poke(self, context: Optional[Dict[str, Any]]) -> bool:
+        if hasattr(super(), '_get_hook'):
+            # .poke in Airflow≥1.10.7 will call overridden _get_hook method
+            return super().poke(context)
+        else:
+            # hook used by .poke in Airflow≤1.10.6 can not be overridden
+            return self.__poke_clickhouse()
+
+    def __poke_clickhouse(self) -> bool:
+        # github.com/apache/airflow/blob/1.10.6/airflow/sensors/sql_sensor.py#L71
+        hook = self._get_hook()
+        self.log.info(
+            'Poking: %s (with parameters %s)',
+            self.sql, self.parameters,
+        )
+        records = hook.get_records(self.sql, self.parameters)
+        if not records:
+            if self.fail_on_empty:
+                raise AirflowException(
+                    'No rows returned, raising as per fail_on_empty flag',
+                )
+            else:
+                return False
+        first_cell = records[0][0]
+        if self.failure is not None:
+            if callable(self.failure):
+                if self.failure(first_cell):
+                    raise AirflowException(
+                        'Failure criteria met. self.failure({}) returned True'
+                            .format(first_cell),
+                    )
+            else:
+                raise AirflowException(
+                    'self.failure is present, but not callable -> {}'
+                        .format(self.success),
+                )
+        if self.success is not None:
+            if callable(self.success):
+                return self.success(first_cell)
+            else:
+                raise AirflowException(
+                    'self.success is present, but not callable -> {}'
+                        .format(self.success),
+                )
+        return str(first_cell) not in ('0', '')
diff --git a/setup.py b/setup.py
index 062f643..f549c40 100755
--- a/setup.py
+++ b/setup.py
@@ -37,6 +37,8 @@
                 '= airflow_clickhouse_plugin:ClickHouseHookPlugin',
             'airflow_clickhouse_operator '
                 '= airflow_clickhouse_plugin:ClickHouseOperatorPlugin',
+            'airflow_clickhouse_sql_sensor '
+                '= airflow_clickhouse_plugin:ClickHouseSqlSensorPlugin',
         ]
     },
     classifiers=[


Test patch:
diff --git a/tests/integration/test_sql_sensor.py b/tests/integration/test_sql_sensor.py
new file mode 100644
index 0000000..f7650d2
--- /dev/null
+++ b/tests/integration/test_sql_sensor.py
@@ -0,0 +1,23 @@
+import unittest
+
+from airflow_clickhouse_plugin import ClickHouseSqlSensor
+
+
+class BasicTestCase(unittest.TestCase):
+    def test_sql_sensor(self):
+        self.assertFalse(
+            ClickHouseSqlSensor(task_id='_', sql='SELECT 0').poke(None),
+        )
+        self.assertTrue(
+            ClickHouseSqlSensor(task_id='_', sql='SELECT 1').poke(None),
+        )
+        self.assertTrue(
+            ClickHouseSqlSensor(task_id='_', sql='SELECT NULL').poke(None),
+        )
+        self.assertFalse(
+            ClickHouseSqlSensor(task_id='_', sql='SELECT 1 WHERE 0').poke(None),
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/unit/test_sql_sensor.py b/tests/unit/test_sql_sensor.py
new file mode 100644
index 0000000..7576ede
--- /dev/null
+++ b/tests/unit/test_sql_sensor.py
@@ -0,0 +1,154 @@
+import unittest
+from unittest import mock
+
+from airflow.exceptions import AirflowException
+from airflow.sensors.sql_sensor import SqlSensor
+
+from airflow_clickhouse_plugin import ClickHouseSqlSensor
+
+
+class ClickHouseSqlSensorTestCase(unittest.TestCase):
+    def test_clickhouse_sql_sensor_poke(self):
+        op = ClickHouseSqlSensor(task_id='_', sql='',)
+        for return_value, expected_result in (
+                ([], False),
+                ([[None]], True),
+                ([['None']], True),
+                ([[0.0]], True),
+                ([[0]], False),
+                ([['0']], False),
+                ([['1']], True),
+        ):
+            with self.subTest(return_value):
+                self._get_records_mock.return_value = return_value
+                self.assertEqual(expected_result, op.poke(context=None))
+
+    def test_clickhouse_sql_sensor_poke_fail_on_empty(self):
+        op = ClickHouseSqlSensor(task_id='_', sql='', fail_on_empty=True)
+        self._get_records_mock.return_value = []
+        with self.assertRaises(AirflowException):
+            op.poke(context=None)
+
+    def test_clickhouse_sql_sensor_poke_success(self):
+        op = ClickHouseSqlSensor(task_id='_', sql='', success=[1].__contains__)
+        for return_value, expected_result in (
+                ([], False),
+                ([[1]], True),
+                ([['1']], False),
+        ):
+            with self.subTest(return_value):
+                self._get_records_mock.return_value = return_value
+                self.assertEqual(expected_result, op.poke(context=None))
+
+    def test_clickhouse_sql_sensor_poke_failure(self):
+        op = ClickHouseSqlSensor(task_id='_', sql='', failure=[1].__contains__)
+
+        self._get_records_mock.return_value = []
+        with self.subTest(self._get_records_mock.return_value):
+            self.assertFalse(op.poke(None))
+
+        self._get_records_mock.return_value = [[1]]
+        with self.subTest(self._get_records_mock.return_value):
+            with self.assertRaises(AirflowException):
+                op.poke(context=None)
+
+    def test_clickhouse_sql_sensor_poke_failure_success(self):
+        op = ClickHouseSqlSensor(
+            task_id='_', sql='',
+            failure=[1].__contains__, success=[2].__contains__,
+        )
+        for return_value, expected_result in (
+                ([], False),
+                ([[2]], True),
+        ):
+            with self.subTest(return_value):
+                self._get_records_mock.return_value = return_value
+                self.assertEqual(expected_result, op.poke(context=None))
+        self._get_records_mock.return_value = [[1]]
+        with self.subTest(self._get_records_mock.return_value):
+            with self.assertRaises(AirflowException):
+                op.poke(context=None)
+
+    def test_clickhouse_sql_sensor_poke_failure_success_same(self):
+        fn = [2].__contains__
+        op = ClickHouseSqlSensor(task_id='_', sql='', failure=fn, success=fn)
+
+        self._get_records_mock.return_value = []
+        with self.subTest(self._get_records_mock.return_value):
+            self.assertFalse(op.poke(None))
+
+        self._get_records_mock.return_value = [[2]]
+        with self.subTest(self._get_records_mock.return_value):
+            with self.assertRaises(AirflowException):
+                op.poke(context=None)
+
+    def test_clickhouse_sql_sensor_poke_invalid_failure(self):
+        op = ClickHouseSqlSensor(task_id='_', sql='', failure=[1])
+        self._get_records_mock.return_value = [[1]]
+        with self.assertRaises(AirflowException):
+            op.poke(context=None)
+
+    def test_clickhouse_sql_sensor_poke_invalid_success(self):
+        op = ClickHouseSqlSensor(task_id='_', sql='', success=[1])
+        self._get_records_mock.return_value = [[1]]
+        with self.assertRaises(AirflowException):
+            op.poke(context=None)
+
+    _get_records_patch: mock._patch
+    _get_records_mock: mock.MagicMock
+
+    @classmethod
+    def setUpClass(cls):
+        cls._get_records_patch = \
+            mock.patch('airflow_clickhouse_plugin.ClickHouseHook.get_records')
+        cls._get_records_mock = cls._get_records_patch.__enter__()
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        cls._get_records_patch.__exit__(None, None, None)
+
+
+class ClickHouseLegacySqlSensorTestCase(unittest.TestCase):
+    @mock.patch(
+        'airflow_clickhouse_plugin.ClickHouseSqlSensor'
+            '._ClickHouseSqlSensor__poke_clickhouse',
+    )
+    @mock.patch('airflow.sensors.sql_sensor.SqlSensor.poke')
+    @mock.patch('airflow.sensors.sql_sensor.SqlSensor._get_hook', create=True)
+    def test_get_hook_defined(
+            self,
+            _: mock.MagicMock,  # force creation of _get_hook
+            sql_sensor_poke_mock: mock.MagicMock,
+            clickhouse_sensor_poke_mock: mock.MagicMock,
+    ):
+        """ Test for case when ``SqlSensor._get_hook`` method is present. """
+        op = ClickHouseSqlSensor(task_id='_', sql='')
+        context = dict(some_test_context=True)
+        op.poke(context)
+        sql_sensor_poke_mock.assert_called_once_with(context)
+        clickhouse_sensor_poke_mock.assert_not_called()
+
+    @mock.patch(
+        'airflow_clickhouse_plugin.ClickHouseSqlSensor'
+            '._ClickHouseSqlSensor__poke_clickhouse',
+    )
+    @mock.patch('airflow.sensors.sql_sensor.SqlSensor.poke')
+    @mock.patch('airflow.sensors.sql_sensor.SqlSensor._get_hook', create=True)
+    def test_get_hook_missing(
+            self,
+            sql_sensor_get_hook_mock: mock.MagicMock,
+            sql_sensor_poke_mock: mock.MagicMock,
+            clickhouse_sensor_poke_mock: mock.MagicMock,
+    ):
+        """ Test for case when ``SqlSensor._get_hook`` method is present. """
+        op = ClickHouseSqlSensor(task_id='_', sql='')
+        del SqlSensor._get_hook
+        op.poke(context=None)
+        sql_sensor_poke_mock.assert_not_called()
+        clickhouse_sensor_poke_mock.assert_called_once_with()
+        # line below prevents AttributeError on patch exit
+        SqlSensor._get_hook = sql_sensor_get_hook_mock
+
+
+if __name__ == '__main__':
+    unittest.main()