From b20ae4e79868375c5e4e3e3a4023653e63263465 Mon Sep 17 00:00:00 2001 From: John Cheng Date: Fri, 11 Jan 2019 12:40:25 +0800 Subject: [PATCH 1/4] [AIRFLOW-3672] Add support for Mongo DB DNS Seedlist Connection Format https://docs.mongodb.com/manual/reference/connection-string/index.html#dns-seedlist-connection-format http://api.mongodb.com/python/current/api/pymongo/mongo_client.html#pymongo.mongo_client.MongoClient --- airflow/contrib/hooks/mongo_hook.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/airflow/contrib/hooks/mongo_hook.py b/airflow/contrib/hooks/mongo_hook.py index 4d72f56566114..eda7ccb21d35d 100644 --- a/airflow/contrib/hooks/mongo_hook.py +++ b/airflow/contrib/hooks/mongo_hook.py @@ -24,8 +24,11 @@ class MongoHook(BaseHook): https://docs.mongodb.com/manual/reference/connection-string/index.html You can specify connection string options in extra field of your connection https://docs.mongodb.com/manual/reference/connection-string/index.html#connection-string-options + + If you want use DNS seedlist, set `srv` to True. + ex. - {replicaSet: test, ssl: True, connectTimeoutMS: 30000} + {"srv": true, "replicaSet": "test", "ssl": true, "connectTimeoutMS": 30000} """ conn_type = 'mongo' @@ -34,8 +37,12 @@ def __init__(self, conn_id='mongo_default', *args, **kwargs): self.mongo_conn_id = conn_id self.connection = self.get_connection(conn_id) - self.extras = self.connection.extra_dejson + self.extras = self.connection.extra_dejson.copy() self.client = None + self.srv = False + + if 'srv' in self.extras: + self.srv = self.extras.pop('srv') def __enter__(self): return self @@ -53,7 +60,10 @@ def get_conn(self): conn = self.connection - uri = 'mongodb://{creds}{host}{port}/{database}'.format( + scheme = 'mongodb+srv' if self.srv else 'mongodb' + + uri = '{scheme}://{creds}{host}{port}/{database}'.format( + scheme=scheme, creds='{}:{}@'.format( conn.login, conn.password ) if conn.login is not None else '', From 37789d12e1752efe821c8206b9f05a9fa367a056 Mon Sep 17 00:00:00 2001 From: John Cheng Date: Sun, 14 Apr 2019 16:50:37 +0800 Subject: [PATCH 2/4] [AIRFLOW-3672] Add unit test for srv uri --- airflow/contrib/hooks/mongo_hook.py | 5 +++-- setup.py | 2 +- tests/contrib/hooks/test_mongo_hook.py | 7 ++++++- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/airflow/contrib/hooks/mongo_hook.py b/airflow/contrib/hooks/mongo_hook.py index 48d7a790f0ee5..34c5c80346f55 100644 --- a/airflow/contrib/hooks/mongo_hook.py +++ b/airflow/contrib/hooks/mongo_hook.py @@ -45,6 +45,7 @@ def __init__(self, conn_id='mongo_default', *args, **kwargs): self.extras = self.connection.extra_dejson.copy() self.client = None self.srv = False + self.uri = None if 'srv' in self.extras: self.srv = self.extras.pop('srv') @@ -67,7 +68,7 @@ def get_conn(self): scheme = 'mongodb+srv' if self.srv else 'mongodb' - uri = '{scheme}://{creds}{host}{port}/{database}'.format( + self.uri = '{scheme}://{creds}{host}{port}/{database}'.format( scheme=scheme, creds='{}:{}@'.format( conn.login, conn.password @@ -85,7 +86,7 @@ def get_conn(self): if options.get('ssl', False): options.update({'ssl_cert_reqs': CERT_NONE}) - self.client = MongoClient(uri, **options) + self.client = MongoClient(self.uri, **options) return self.client diff --git a/setup.py b/setup.py index 063bb04f13fa8..7dbae4845a70f 100644 --- a/setup.py +++ b/setup.py @@ -239,7 +239,7 @@ def write_version(filename=os.path.join(*['airflow', segment = ['analytics-python>=1.2.9'] sendgrid = ['sendgrid>=5.2.0,<6'] slack = ['slackclient>=1.0.0'] -mongo = ['pymongo>=3.6.0'] +mongo = ['pymongo>=3.6.0', 'dnspython>=1.13.0,<2.0.0'] snowflake = ['snowflake-connector-python>=1.5.2', 'snowflake-sqlalchemy>=1.1.0'] ssh = ['paramiko>=2.1.1', 'pysftp>=0.2.9', 'sshtunnel>=0.1.4,<0.2'] diff --git a/tests/contrib/hooks/test_mongo_hook.py b/tests/contrib/hooks/test_mongo_hook.py index 114ba8495911c..9a96335844f5e 100644 --- a/tests/contrib/hooks/test_mongo_hook.py +++ b/tests/contrib/hooks/test_mongo_hook.py @@ -42,13 +42,18 @@ def get_collection(self, mock_collection, mongo_db=None): class TestMongoHook(unittest.TestCase): def setUp(self): configuration.load_test_config() - self.hook = MongoHookTest(conn_id='mongo_default', mongo_db='default') + self.hook = MongoHookTest( + conn_id='mongo_default', + mongo_db='default', + extra='{"srv":true}', + ) self.conn = self.hook.get_conn() @unittest.skipIf(mongomock is None, 'mongomock package not present') def test_get_conn(self): self.assertEqual(self.hook.connection.port, 27017) self.assertIsInstance(self.conn, pymongo.MongoClient) + self.assertTrue(self.hook.uri.startswith('mongodb+srv://')) @unittest.skipIf(mongomock is None, 'mongomock package not present') def test_insert_one(self): From ffd97d612e10f9a70d47f1f2a14cb43c3a2c2f0b Mon Sep 17 00:00:00 2001 From: John Cheng Date: Sun, 14 Apr 2019 18:02:16 +0800 Subject: [PATCH 3/4] [AIRFLOW-3672] Fix unit test for Mongo srv uri --- tests/contrib/hooks/test_mongo_hook.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/tests/contrib/hooks/test_mongo_hook.py b/tests/contrib/hooks/test_mongo_hook.py index 9a96335844f5e..57becea43136b 100644 --- a/tests/contrib/hooks/test_mongo_hook.py +++ b/tests/contrib/hooks/test_mongo_hook.py @@ -42,18 +42,22 @@ def get_collection(self, mock_collection, mongo_db=None): class TestMongoHook(unittest.TestCase): def setUp(self): configuration.load_test_config() - self.hook = MongoHookTest( - conn_id='mongo_default', - mongo_db='default', - extra='{"srv":true}', - ) + self.hook = MongoHookTest(conn_id='mongo_default', mongo_db='default') self.conn = self.hook.get_conn() @unittest.skipIf(mongomock is None, 'mongomock package not present') def test_get_conn(self): self.assertEqual(self.hook.connection.port, 27017) self.assertIsInstance(self.conn, pymongo.MongoClient) - self.assertTrue(self.hook.uri.startswith('mongodb+srv://')) + + @unittest.skipIf(mongomock is None, 'mongomock package not present') + def test_srv(self): + hook = MongoHookTest( + conn_id='mongo_default', + mongo_db='default', + extra='{"srv":true}', + ) + self.assertTrue(hook.uri.startswith('mongodb+srv://')) @unittest.skipIf(mongomock is None, 'mongomock package not present') def test_insert_one(self): From 41e746ec89862dd3ac9056dd16f35ead9b961ac3 Mon Sep 17 00:00:00 2001 From: John Cheng Date: Sun, 14 Apr 2019 19:44:54 +0800 Subject: [PATCH 4/4] [AIRFLOW-3672] Construct MongoDB URI when hook init --- airflow/contrib/hooks/mongo_hook.py | 32 +++++++++++--------------- tests/contrib/hooks/test_mongo_hook.py | 12 ++++++---- 2 files changed, 20 insertions(+), 24 deletions(-) diff --git a/airflow/contrib/hooks/mongo_hook.py b/airflow/contrib/hooks/mongo_hook.py index 34c5c80346f55..e674d9d22adc6 100644 --- a/airflow/contrib/hooks/mongo_hook.py +++ b/airflow/contrib/hooks/mongo_hook.py @@ -44,11 +44,20 @@ def __init__(self, conn_id='mongo_default', *args, **kwargs): self.connection = self.get_connection(conn_id) self.extras = self.connection.extra_dejson.copy() self.client = None - self.srv = False - self.uri = None - if 'srv' in self.extras: - self.srv = self.extras.pop('srv') + srv = self.extras.pop('srv', False) + scheme = 'mongodb+srv' if srv else 'mongodb' + + self.uri = '{scheme}://{creds}{host}{port}/{database}'.format( + scheme=scheme, + creds='{}:{}@'.format( + self.connection.login, self.connection.password + ) if self.connection.login else '', + + host=self.connection.host, + port='' if self.connection.port is None else ':{}'.format(self.connection.port), + database=self.connection.schema + ) def __enter__(self): return self @@ -64,21 +73,6 @@ def get_conn(self): if self.client is not None: return self.client - conn = self.connection - - scheme = 'mongodb+srv' if self.srv else 'mongodb' - - self.uri = '{scheme}://{creds}{host}{port}/{database}'.format( - scheme=scheme, - creds='{}:{}@'.format( - conn.login, conn.password - ) if conn.login else '', - - host=conn.host, - port='' if conn.port is None else ':{}'.format(conn.port), - database=conn.schema - ) - # Mongo Connection Options dict that is unpacked when passed to MongoClient options = self.extras diff --git a/tests/contrib/hooks/test_mongo_hook.py b/tests/contrib/hooks/test_mongo_hook.py index 57becea43136b..4e362c330f9f5 100644 --- a/tests/contrib/hooks/test_mongo_hook.py +++ b/tests/contrib/hooks/test_mongo_hook.py @@ -25,6 +25,8 @@ from airflow import configuration from airflow.contrib.hooks.mongo_hook import MongoHook +from airflow.models import Connection +from airflow.utils import db class MongoHookTest(MongoHook): @@ -44,6 +46,10 @@ def setUp(self): configuration.load_test_config() self.hook = MongoHookTest(conn_id='mongo_default', mongo_db='default') self.conn = self.hook.get_conn() + db.merge_conn( + Connection( + conn_id='mongo_default_with_srv', conn_type='mongo', + host='mongo', port='27017', extra='{"srv": true}')) @unittest.skipIf(mongomock is None, 'mongomock package not present') def test_get_conn(self): @@ -52,11 +58,7 @@ def test_get_conn(self): @unittest.skipIf(mongomock is None, 'mongomock package not present') def test_srv(self): - hook = MongoHookTest( - conn_id='mongo_default', - mongo_db='default', - extra='{"srv":true}', - ) + hook = MongoHook(conn_id='mongo_default_with_srv') self.assertTrue(hook.uri.startswith('mongodb+srv://')) @unittest.skipIf(mongomock is None, 'mongomock package not present')