Repository: incubator-airflow Updated Branches: refs/heads/master d704dc357 -> 68854e99b
[AIRFLOW-2785] Add context manager entry points to mongoHook Closes #3628 from andscoop/Add-connection-close- to-mongo-hook Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/68854e99 Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/68854e99 Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/68854e99 Branch: refs/heads/master Commit: 68854e99b5f1412a1f8168246a40c5d5d35c0157 Parents: d704dc3 Author: Andy Cooper <[email protected]> Authored: Tue Jul 24 01:01:38 2018 +0100 Committer: Kaxil Naik <[email protected]> Committed: Tue Jul 24 01:01:38 2018 +0100 ---------------------------------------------------------------------- airflow/contrib/hooks/mongo_hook.py | 13 +++++++++++++ tests/contrib/hooks/test_mongo_hook.py | 9 +++++++++ 2 files changed, 22 insertions(+) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/68854e99/airflow/contrib/hooks/mongo_hook.py ---------------------------------------------------------------------- diff --git a/airflow/contrib/hooks/mongo_hook.py b/airflow/contrib/hooks/mongo_hook.py index 6ae71a8..80cedde 100644 --- a/airflow/contrib/hooks/mongo_hook.py +++ b/airflow/contrib/hooks/mongo_hook.py @@ -37,6 +37,13 @@ class MongoHook(BaseHook): self.extras = self.connection.extra_dejson self.client = None + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if self.client is not None: + self.close_conn() + def get_conn(self): """ Fetches PyMongo Client @@ -67,6 +74,12 @@ class MongoHook(BaseHook): return self.client + def close_conn(self): + client = self.client + if client is not None: + client.close() + self.client = None + def get_collection(self, mongo_collection, mongo_db=None): """ Fetches a mongo collection object for querying. http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/68854e99/tests/contrib/hooks/test_mongo_hook.py ---------------------------------------------------------------------- diff --git a/tests/contrib/hooks/test_mongo_hook.py b/tests/contrib/hooks/test_mongo_hook.py index 7696ede..00fe0f0 100644 --- a/tests/contrib/hooks/test_mongo_hook.py +++ b/tests/contrib/hooks/test_mongo_hook.py @@ -118,6 +118,15 @@ class TestMongoHook(unittest.TestCase): results = [result for result in results] self.assertEqual(len(results), 2) + def test_context_manager(self): + with MongoHook(conn_id='mongo_default', mongo_db='default') as ctxHook: + ctxHook.get_conn() + + self.assertIsInstance(ctxHook, MongoHook) + self.assertIsNotNone(ctxHook.client) + + self.assertIsNone(ctxHook.client) + if __name__ == '__main__': unittest.main()
