[AIRFLOW-2758] Add a sensor for MongoDB

.travis.yml

@@ -23,6 +23,7 @@ jdk:
   - oraclejdk8
 services:
   - cassandra
+  - mongodb
   - mysql
   - postgresql
   - rabbitmq

airflow/contrib/sensors/mongo_sensor.py

@@ -0,0 +1,58 @@
+# -*- 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.contrib.hooks.mongo_hook import MongoHook
+from airflow.sensors.base_sensor_operator import BaseSensorOperator
+from airflow.utils.decorators import apply_defaults
+
+
+class MongoSensor(BaseSensorOperator):
+    """
+    Checks for the existence of a document which
+    matches the given query in MongoDB. Example:
+
+    >>> mongo_sensor = MongoSensor(collection="coll",
+    ...                            query={"key": "value"},
+    ...                            mongo_conn_id="mongo_default",
+    ...                            task_id="mongo_sensor")
+    """
+    template_fields = ('collection', 'query')
+
+    @apply_defaults
+    def __init__(self, collection, query, mongo_conn_id="mongo_default", *args, **kwargs):
+        """
+        Create a new MongoSensor
+
+        :param collection: Target MongoDB collection.
+        :type collection: string
+        :param query: The query to find the target document.
+        :type query: dict
+        :param mongo_conn_id: The connection ID to use
+                              when connecting to MongoDB.
+        :type mongo_conn_id: string
+        """
+        super(MongoSensor, self).__init__(*args, **kwargs)
+        self.mongo_conn_id = mongo_conn_id
+        self.collection = collection
+        self.query = query
+
+    def poke(self, context):
+        self.log.info("Sensor check existence of the document "
+                      "that matches the following query: %s", self.query)
+        hook = MongoHook(self.mongo_conn_id)
+        return hook.find(self.collection, self.query, find_one=True) is not None

tests/contrib/sensors/test_mongo_sensor.py

@@ -0,0 +1,65 @@
+# -*- 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
+from airflow import configuration
+from airflow.contrib.hooks.mongo_hook import MongoHook
+from airflow.contrib.sensors.mongo_sensor import MongoSensor
+from airflow.models import Connection
+from airflow.utils import db, timezone
+
+
+DEFAULT_DATE = timezone.datetime(2017, 1, 1)
+
+
+class TestMongoSensor(unittest.TestCase):
+
+    def setUp(self):
+        configuration.load_test_config()
+        db.merge_conn(
+            Connection(
+                conn_id='mongo_test', conn_type='mongo',
+                host='localhost', port='27017', schema='test'))
+
+        args = {
+            'owner': 'airflow',
+            'start_date': DEFAULT_DATE
+        }
+        self.dag = DAG('test_dag_id', default_args=args)
+
+        hook = MongoHook('mongo_test')
+        hook.insert_one('foo', {'bar': 'baz'})
+
+        self.sensor = MongoSensor(
+            task_id='test_task',
+            mongo_conn_id='mongo_test',
+            dag=self.dag,
+            collection='foo',
+            query={'bar': 'baz'}
+        )
+
+    def test_poke(self):
+        self.assertTrue(self.sensor.poke(None))
+
+
+if __name__ == '__main__':
+    unittest.main()