diff --git a/airflow/providers/mongo/sensors/mongo.py b/airflow/providers/mongo/sensors/mongo.py index cdd48c76986d6..9e96944da0ef0 100644 --- a/airflow/providers/mongo/sensors/mongo.py +++ b/airflow/providers/mongo/sensors/mongo.py @@ -41,16 +41,17 @@ class MongoSensor(BaseSensorOperator): template_fields = ('collection', 'query') def __init__( - self, *, collection: str, query: dict, mongo_conn_id: str = "mongo_default", **kwargs + self, *, collection: str, query: dict, mongo_conn_id: str = "mongo_default", mongo_db=None, **kwargs ) -> None: super().__init__(**kwargs) self.mongo_conn_id = mongo_conn_id self.collection = collection self.query = query + self.mongo_db = mongo_db def poke(self, context: dict) -> bool: 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 + return hook.find(self.collection, self.query, mongo_db=self.mongo_db, find_one=True) is not None diff --git a/tests/providers/mongo/sensors/test_mongo.py b/tests/providers/mongo/sensors/test_mongo.py index 6623631550855..e5504e747941e 100644 --- a/tests/providers/mongo/sensors/test_mongo.py +++ b/tests/providers/mongo/sensors/test_mongo.py @@ -53,3 +53,17 @@ def setUp(self): def test_poke(self): assert self.sensor.poke(None) + + def test_sensor_with_db(self): + hook = MongoHook('mongo_test') + hook.insert_one('nontest', {'1': '2'}, mongo_db='nontest') + + sensor = MongoSensor( + task_id='test_task2', + mongo_conn_id='mongo_test', + dag=self.dag, + collection='nontest', + query={'1': '2'}, + mongo_db="nontest", + ) + assert sensor.poke(None)