Merge pull request pandas-dev#552 from srf94/pymongo_kwargs

Added kwargs to Arctic init for extra MongoClient parameters (Issue 493)

Author: dimosped

CHANGES.md

@@ -3,6 +3,7 @@
 ### 1.70
   * Bugfix: #629 - Column kwarg no longer modified
   * Bugfix: #641 - DateRange.intersection open/closed range fix
+  * Feature: #493 Can pass kwargs when calling MongoClient, e.g. for ssl
 
 ### 1.69 (2018-09-12)
   * Docs: VersionStore documentation

arctic/arctic.py

@@ -70,10 +70,13 @@ class Arctic(object):
 
     def __init__(self, mongo_host, app_name=APPLICATION_NAME, allow_secondary=False,
                  socketTimeoutMS=10 * 60 * 1000, connectTimeoutMS=2 * 1000,
-                 serverSelectionTimeoutMS=30 * 1000):
+                 serverSelectionTimeoutMS=30 * 1000, **kwargs):
         """
         Constructs a Arctic Datastore.
 
+        Note: If mongo_host is a pymongo connection and the process is later forked, the
+                new pymongo connection may have different parameters.
+
         Parameters:
         -----------
         mongo_host: A MongoDB hostname, alias or Mongo Connection
@@ -92,6 +95,8 @@ def __init__(self, mongo_host, app_name=APPLICATION_NAME, allow_secondary=False,
             the pymongo driver will spend on MongoDB cluster discovery.  This parameter
             takes precedence over connectTimeoutMS: https://jira.mongodb.org/browse/DRIVERS-222
 
+        kwargs: 'dict' extra keyword arguments to pass when calling pymongo.MongoClient,
+            for example ssl parameters.
         """
         self._application_name = app_name
         self._library_cache = {}
@@ -101,10 +106,13 @@ def __init__(self, mongo_host, app_name=APPLICATION_NAME, allow_secondary=False,
         self._server_selection_timeout = serverSelectionTimeoutMS
         self._lock = threading.RLock()
         self._pid = os.getpid()
+        self._pymongo_kwargs = kwargs
 
         if isinstance(mongo_host, string_types):
+            self._given_instance = False
             self.mongo_host = mongo_host
         else:
+            self._given_instance = True
             self.__conn = mongo_host
             # Workaround for: https://jira.mongodb.org/browse/PYTHON-927
             mongo_host.server_info()
@@ -119,17 +127,21 @@ def _conn(self):
             #    http://api.mongodb.com/python/current/faq.html#using-pymongo-with-multiprocessing
             curr_pid = os.getpid()
             if curr_pid != self._pid:
+                if self._given_instance:
+                    logger.warn("Forking process. Arctic was passed a pymongo connection during init, "
+                                "the new pymongo connection may have different parameters.")
                 self._pid = curr_pid  # this line has to precede reset() otherwise we get to eternal recursion
                 self.reset()  # also triggers re-auth
 
             if self.__conn is None:
                 host = get_mongodb_uri(self.mongo_host)
                 logger.info("Connecting to mongo: {0} ({1})".format(self.mongo_host, host))
                 self.__conn = pymongo.MongoClient(host=host,
-                                                   maxPoolSize=self._MAX_CONNS,
-                                                   socketTimeoutMS=self._socket_timeout,
-                                                   connectTimeoutMS=self._connect_timeout,
-                                                   serverSelectionTimeoutMS=self._server_selection_timeout)
+                                                  maxPoolSize=self._MAX_CONNS,
+                                                  socketTimeoutMS=self._socket_timeout,
+                                                  connectTimeoutMS=self._connect_timeout,
+                                                  serverSelectionTimeoutMS=self._server_selection_timeout,
+                                                  **self._pymongo_kwargs)
                 self._adminDB = self.__conn.admin
 
                 # Authenticate against admin for the user

tests/unit/test_arctic.py

@@ -26,6 +26,38 @@ def test_arctic_lazy_init():
             assert mc.called
 
 
+def test_arctic_lazy_init_ssl_true():
+    with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True) as mc, \
+            patch('arctic.arctic.mongo_retry', side_effect=lambda x: x, autospec=True), \
+            patch('arctic.arctic.get_auth', autospec=True) as ga:
+        store = Arctic('cluster', ssl=True)
+        assert not mc.called
+        # do something to trigger lazy arctic init
+        store.list_libraries()
+        assert mc.called
+        assert len(mc.mock_calls) == 1
+        assert mc.mock_calls[0] == call(connectTimeoutMS=2000,
+                                        host='cluster',
+                                        maxPoolSize=4,
+                                        serverSelectionTimeoutMS=30000,
+                                        socketTimeoutMS=600000,
+                                        ssl=True)
+
+
+def test_connection_passed_warning_raised():
+    with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True), \
+         patch('arctic.arctic.mongo_retry', side_effect=lambda x: x, autospec=True), \
+         patch('arctic.arctic.get_auth', autospec=True), \
+         patch('arctic.arctic.logger') as lg:
+        magic_mock = MagicMock(nodes={("host", "port")})
+        store = Arctic(magic_mock, ssl=True)
+        # Increment _pid to simulate forking the process
+        store._pid += 1
+        _ = store._conn
+        assert lg.mock_calls[0] == call.warn('Forking process. Arctic was passed a pymongo connection during init, '
+                                             'the new pymongo connection may have different parameters.')
+
+
 def test_arctic_auth():
     with patch('pymongo.MongoClient', return_value=MagicMock(), autospec=True), \
         patch('arctic.arctic.mongo_retry', autospec=True), \