Implement #123 arbitrary Mongo URL & collection name (#124)

* Issue #123 arbitrary Mongo URL & collection name * #124 updated unit tests * update doc Co-authored-by: ekzhu <ekzhu@users.noreply.github.com>
ekzhu · Mar 12, 2020 · e5f14ae · e5f14ae
1 parent d65924b
commit e5f14ae
Show file tree

Hide file tree

Showing 8 changed files with 10,663 additions and 7 deletions.
diff --git a/datasketch/experimental/aio/storage.py b/datasketch/experimental/aio/storage.py
@@ -128,10 +128,17 @@ def __init__(self, config, name=None):
             self._mongo_param = self._parse_config(self._config['mongo'])
 
             self._name = name if name else _random_name(11).decode('utf-8')
-            self._collection_name = 'lsh_' + self._name
+            if 'collection_name' in self.mongo_param:
+                self._collection_name = self.mongo_param['collection_name']
+            elif 'collection_prefix' in self.mongo_param:
+                self._collection_name = self.mongo_param['collection_prefix'] + self._name
+            else:
+                self._collection_name = 'lsh_' + self._name
 
             db_lsh = self.mongo_param['db'] if 'db' in self.mongo_param else 'db_0'
-            if 'replica_set' in self.mongo_param:
+            if 'url' in self.mongo_param:
+                dsn = self.mongo_param['url']
+            elif 'replica_set' in self.mongo_param:
                 dsn = 'mongodb://{replica_set_nodes}/?replicaSet={replica_set}'.format(**self.mongo_param)
             elif 'username' in self.mongo_param or 'password' in self.mongo_param:
                 dsn = 'mongodb://{username}:{password}@{host}:{port}'.format(**self.mongo_param)
@@ -142,7 +149,7 @@ def __init__(self, config, name=None):
 
             self._batch_size = 1000
             self._mongo_client = motor.motor_asyncio.AsyncIOMotorClient(dsn, **additional_args)
-            self._collection = self._mongo_client[db_lsh][self._collection_name]
+            self._collection = self._mongo_client.get_default_database(db_lsh).get_collection(self._collection_name)
             self._initialized = True
             self._buffer = AsyncMongoBuffer(self._collection, self._batch_size)
 

diff --git a/docs/_sources/lsh.rst.txt b/docs/_sources/lsh.rst.txt
@@ -291,6 +291,12 @@ To configure Asynchronous MongoDB storage that will connect to a `replica set <h
 
     _storage = {'type': 'aiomongo', 'mongo': {'replica_set': 'rs0', 'replica_set_nodes': 'node1:port1,node2:port2,node3:port3'}}
 
+To connect to a cloud Mongo Atlas cluster (or any other arbitrary ``mongodb`` URI):
+
+.. code:: python
+
+    _storage = {'type': 'aiomongo', 'mongo': {'url': 'mongodb+srv://user:pass@server-ybq4y.mongodb.net/db'}}
+
 If you want to pass additional params to the `Mongo client <http://api.mongodb.com/python/current/api/pymongo/mongo_client.html>` constructor, just put them in the ``mongo.args`` object in the storage config (example usage to configure X509 authentication):
 
 .. code:: python