fix GridFSURIStore and increase testing

Author: gpetretto

src/maggma/stores/gridfs.py

@@ -87,6 +87,7 @@ def __init__(
         self.searchable_fields = [] if searchable_fields is None else searchable_fields
         self.kwargs = kwargs
         self.ssh_tunnel = ssh_tunnel
+        self._fs = None
 
         if auth_source is None:
             auth_source = self.database
@@ -157,9 +158,9 @@ def connect(self, force_reset: bool = False):
             db = conn[self.database]
             self._coll = gridfs.GridFS(db, self.collection_name)
             self._files_collection = db[f"{self.collection_name}.files"]
-            self._files_store = MongoStore.from_collection(self._files_collection)
-            self._files_store.last_updated_field = f"metadata.{self.last_updated_field}"
-            self._files_store.key = self.key
+            self._fs = MongoStore.from_collection(self._files_collection)
+            self._fs.last_updated_field = f"metadata.{self.last_updated_field}"
+            self._fs.key = self.key
             self._chunks_collection = db[f"{self.collection_name}.chunks"]
 
     @property
@@ -169,6 +170,13 @@ def _collection(self):
             raise StoreError("Must connect Mongo-like store before attempting to use it")
         return self._coll
 
+    @property
+    def _files_store(self):
+        """Property referring to MongoStore associated to the files_collection."""
+        if self._fs is None:
+            raise StoreError("Must connect Mongo-like store before attempting to use it")
+        return self._fs
+
     @property
     def last_updated(self) -> datetime:
         """
@@ -448,6 +456,7 @@ def __init__(
         ensure_metadata: bool = False,
         searchable_fields: Optional[list[str]] = None,
         mongoclient_kwargs: Optional[dict] = None,
+        ssh_tunnel: Optional[SSHTunnel] = None,
         **kwargs,
     ):
         """
@@ -463,6 +472,10 @@ def __init__(
         """
         self.uri = uri
 
+        if ssh_tunnel:
+            raise ValueError(f"At the moment ssh_tunnel is not supported for {self.__class__.__name__}")
+        self.ssh_tunnel = None
+
         # parse the dbname from the uri
         if database is None:
             d_uri = uri_parser.parse_uri(uri)
@@ -479,6 +492,7 @@ def __init__(
         self.searchable_fields = [] if searchable_fields is None else searchable_fields
         self.kwargs = kwargs
         self.mongoclient_kwargs = mongoclient_kwargs or {}
+        self._fs = None
 
         if "key" not in kwargs:
             kwargs["key"] = "_id"
@@ -497,7 +511,26 @@ def connect(self, force_reset: bool = False):
             db = conn[self.database]
             self._coll = gridfs.GridFS(db, self.collection_name)
             self._files_collection = db[f"{self.collection_name}.files"]
-            self._files_store = MongoStore.from_collection(self._files_collection)
-            self._files_store.last_updated_field = f"metadata.{self.last_updated_field}"
-            self._files_store.key = self.key
+            self._fs = MongoStore.from_collection(self._files_collection)
+            self._fs.last_updated_field = f"metadata.{self.last_updated_field}"
+            self._fs.key = self.key
             self._chunks_collection = db[f"{self.collection_name}.chunks"]
+
+    @property
+    def name(self) -> str:
+        """
+        Return a string representing this data source.
+        """
+        # TODO: This is not very safe since it exposes the username/password info
+        return self.uri
+
+    def __eq__(self, other: object) -> bool:
+        """
+        Check equality for GridFSURIStore
+        other: other GridFSURIStore to compare with.
+        """
+        if not isinstance(other, GridFSStore):
+            return False
+
+        fields = ["uri", "database", "collection_name"]
+        return all(getattr(self, f) == getattr(other, f) for f in fields)

src/maggma/stores/mongolike.py

@@ -460,7 +460,9 @@ def __init__(
                 ensure determinacy in query results.
         """
         self.uri = uri
-        self.ssh_tunnel = ssh_tunnel
+        if ssh_tunnel:
+            raise ValueError(f"At the moment ssh_tunnel is not supported for {self.__class__.__name__}")
+        self.ssh_tunnel = None
         self.default_sort = default_sort
         self.safe_update = safe_update
         self.mongoclient_kwargs = mongoclient_kwargs or {}

tests/stores/test_gridfs.py

@@ -20,9 +20,20 @@ def mongostore():
     store._collection.drop()
 
 
-@pytest.fixture()
-def gridfsstore():
-    store = GridFSStore("maggma_test", "test", key="task_id")
+@pytest.fixture(params=["std", "uri"])
+def gridfsstore(request):
+    """
+    Fixture providing both a standard GridFSStore and a GridFSURIStore.
+    """
+    store_type = request.param
+    if store_type == "std":
+        store = GridFSStore("maggma_test", "test", key="task_id")
+    elif store_type == "uri":
+        store = GridFSURIStore(
+            uri="mongodb://localhost:27017", database="maggma_test", collection_name="test", key="task_id"
+        )
+    else:
+        raise ValueError(f"Unknown store_type {store_type}")
     store.connect()
     yield store
     store._files_collection.drop()
@@ -243,6 +254,7 @@ def test_gridfs_uri():
     is_name = store.name is uri
     # This is try and keep the secret safe
     assert is_name
+    store.close()
 
 
 def test_gridfs_uri_dbname_parse():