[AIRFLOW-3892] Create Redis pub sub sensor

airflow/contrib/sensors/redis_pub_sub_sensor.py

@@ -0,0 +1,73 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from airflow.sensors.base_sensor_operator import BaseSensorOperator
+from airflow.utils.decorators import apply_defaults
+from airflow.contrib.hooks.redis_hook import RedisHook
+
+
+class RedisPubSubSensor(BaseSensorOperator):
+
+    """
+    Redis sensor for reading a message from pub sub channels
+    """
+    template_fields = ('channels',)
+    ui_color = '#f0eee4'
+
+    @apply_defaults
+    def __init__(self, channels, redis_conn_id, *args, **kwargs):
+        """
+        Create a new RedisPubSubSensor and subscribe to the channels
+
+        :param channels: The channels to be subscribed to (templated)
+        :type channels: str or list of str
+        :param redis_conn_id: the redis connection id
+        :type redis_conn_id: str
+        """
+
+        super(RedisPubSubSensor, self).__init__(*args, **kwargs)
+        self.channels = channels
+        self.redis_conn_id = redis_conn_id
+        self.pubsub = RedisHook(redis_conn_id=self.redis_conn_id).get_conn().pubsub()
+        self.pubsub.subscribe(self.channels)
+
+    def poke(self, context):
+        """
+        Check for message on subscribed channels and write to xcom the message with key ``message``
+
+        An example of message ``{'type': 'message', 'pattern': None, 'channel': b'test', 'data': b'hello'}``
+
+        :param context: the context object
+        :type context: dict
+        :return: ``True`` if message (with type 'message') is available or ``False`` if not
+        """
+        self.log.info('RedisPubSubSensor checking for message on channels: %s', self.channels)
+
+        message = self.pubsub.get_message()
+        self.log.info('Message %s from channel %s', message, self.channels)
+
+        # Process only message types
+        if message and message['type'] == 'message':
+
+            context['ti'].xcom_push(key='message', value=message)
+            self.pubsub.unsubscribe(self.channels)
+
+            return True
+
+        return False

tests/contrib/sensors/test_redis_pub_sub_sensor.py

@@ -0,0 +1,125 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+
+import unittest
+from airflow import DAG, configuration
+from airflow.contrib.sensors.redis_pub_sub_sensor import RedisPubSubSensor
+from airflow.utils import timezone
+from airflow.contrib.hooks.redis_hook import RedisHook
+from mock import patch, call, MagicMock
+
+DEFAULT_DATE = timezone.datetime(2017, 1, 1)
+
+
+class TestRedisPubSubSensor(unittest.TestCase):
+
+    def setUp(self):
+        configuration.load_test_config()
+
+        args = {
+            'owner': 'airflow',
+            'start_date': DEFAULT_DATE
+        }
+
+        self.dag = DAG('test_dag_id', default_args=args)
+
+        self.mock_context = MagicMock()
+
+    @patch('airflow.contrib.hooks.redis_hook.RedisHook.get_conn')
+    def test_poke_mock_true(self, mock_redis_conn):
+        sensor = RedisPubSubSensor(
+            task_id='test_task',
+            dag=self.dag,
+            channels='test',
+            redis_conn_id='redis_default'
+        )
+
+        self.mock_redis_conn = mock_redis_conn
+        self.mock_redis_conn().pubsub().get_message.return_value = \
+            {'type': 'message', 'channel': b'test', 'data': b'd1'}
+
+        result = sensor.poke(self.mock_context)
+        self.assertTrue(result)
+
+        context_calls = [call.xcom_push(key='message',
+                                        value={'type': 'message', 'channel': b'test', 'data': b'd1'})]
+
+        self.assertTrue(self.mock_context['ti'].method_calls == context_calls, "context call  should be same")
+
+    @patch('airflow.contrib.hooks.redis_hook.RedisHook.get_conn')
+    def test_poke_mock_false(self, mock_redis_conn):
+        sensor = RedisPubSubSensor(
+            task_id='test_task',
+            dag=self.dag,
+            channels='test',
+            redis_conn_id='redis_default'
+        )
+        self.mock_redis_conn = mock_redis_conn
+        self.mock_redis_conn().pubsub().get_message.return_value = \
+            {'type': 'subscribe', 'channel': b'test', 'data': b'd1'}
+
+        result = sensor.poke(self.mock_context)
+        self.assertFalse(result)
+
+        context_calls = []
+        self.assertTrue(self.mock_context['ti'].method_calls == context_calls, "context calls should be same")
+
+    def test_poke_true(self):
+        sensor = RedisPubSubSensor(
+            task_id='test_task',
+            dag=self.dag,
+            channels='test',
+            redis_conn_id='redis_default'
+        )
+
+        hook = RedisHook(redis_conn_id='redis_default')
+        redis = hook.get_conn()
+        redis.publish('test', 'message')
+
+        result = sensor.poke(self.mock_context)
+        self.assertFalse(result)
+        result = sensor.poke(self.mock_context)
+        self.assertTrue(result)
+        context_calls = [
+            call.xcom_push(
+                key='message',
+                value={'type': 'message', 'pattern': None, 'channel': b'test', 'data': b'message'})]
+        self.assertTrue(self.mock_context['ti'].method_calls == context_calls, "context calls should be same")
+        result = sensor.poke(self.mock_context)
+        self.assertFalse(result)
+
+    def test_poke_false(self):
+        sensor = RedisPubSubSensor(
+            task_id='test_task',
+            dag=self.dag,
+            channels='test',
+            redis_conn_id='redis_default'
+        )
+
+        result = sensor.poke(self.mock_context)
+        self.assertFalse(result)
+        self.assertTrue(self.mock_context['ti'].method_calls == [], "context calls should be same")
+        result = sensor.poke(self.mock_context)
+        self.assertFalse(result)
+        self.assertTrue(self.mock_context['ti'].method_calls == [], "context calls should be same")
+
+
+if __name__ == '__main__':
+    unittest.main()