Automatically parse the dbname from the URI

src/maggma/stores/mongolike.py

@@ -17,8 +17,8 @@
 from monty.json import jsanitize, MSONable
 from monty.serialization import loadfn
 from pydash import get, has, set_
-from pymongo import MongoClient, ReplaceOne
-from pymongo.errors import OperationFailure, DocumentTooLarge
+from pymongo import MongoClient, ReplaceOne, uri_parser
+from pymongo.errors import OperationFailure, DocumentTooLarge, ConfigurationError
 from sshtunnel import SSHTunnelForwarder
 
 
@@ -366,15 +366,26 @@ class MongoURIStore(MongoStore):
     client parameters via TXT records
     """
 
-    def __init__(self, uri: str, database: str, collection_name: str, **kwargs):
+    def __init__(self, uri: str, collection_name: str, database: str = None, **kwargs):
         """
         Args:
             uri: MongoDB+SRV URI
             database: database to connect to
             collection_name: The collection name
         """
         self.uri = uri
-        self.database = database
+
+        # parse the dbname from the uri
+        if database is None:
+            d_uri = uri_parser.parse_uri(uri)
+            if d_uri["database"] is None:
+                raise ConfigurationError(
+                    "If database name is not supplied, a database must be set in the uri"
+                )
+            self.database = d_uri["database"]
+        else:
+            self.database = database
+
         self.collection_name = collection_name
         self.kwargs = kwargs
         self._collection = None

tests/stores/test_mongolike.py

@@ -3,7 +3,7 @@
 
 import mongomock.collection
 import pymongo.collection
-from pymongo.errors import OperationFailure, DocumentTooLarge
+from pymongo.errors import OperationFailure, DocumentTooLarge, ConfigurationError
 import pytest
 
 from maggma.core import StoreError
@@ -286,3 +286,18 @@ def test_mongo_uri():
     is_name = store.name is uri
     # This is try and keep the secret safe
     assert is_name
+
+
+def test_mongo_uri_dbname_parse():
+    # test parsing dbname from uri
+    uri_with_db = "mongodb://uuu:xxxx@host:27017/fake_db"
+    store = MongoURIStore(uri_with_db, "test")
+    assert store.database == "fake_db"
+
+    uri_with_db = "mongodb://uuu:xxxx@host:27017/fake_db"
+    store = MongoURIStore(uri_with_db, "test", database="fake_db2")
+    assert store.database == "fake_db2"
+
+    uri_with_db = "mongodb://uuu:xxxx@host:27017"
+    with pytest.raises(ConfigurationError):
+        MongoURIStore(uri_with_db, "test")