MemoryStore: replace mongomock with pymongo-inmemory

.github/workflows/testing.yml

@@ -51,7 +51,9 @@ jobs:
       max-parallel: 6
       matrix:
         os: [ubuntu-latest]
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
+        # python-version: ["3.8", "3.9", "3.10", "3.11"]
+        # temporary workaround since pymongo-inmemory dropped 3.8 support
+        python-version: ["3.9", "3.10", "3.11"]
 
     runs-on: ${{ matrix.os }}
 
@@ -72,6 +74,8 @@ jobs:
       env:
         CONTINUOUS_INTEGRATION: True
         MONGODB_SRV_URI: ${{ secrets.MONGODB_SRV_URI }}
+        PYMONGOIM__OPERATING_SYSTEM: ubuntu
+        PYMONGOIM__MONGO_VERSION: 6.0
       run: |
         pip install -e .
         pytest --cov=maggma --cov-report=xml

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
         "pymongo>=4.2.0",
         "monty>=2023.9.25",
         "mongomock>=3.10.0",
+        "pymongo-inmemory>=0.4.1",
         "pydash>=4.1.0",
         "jsonschema>=3.1.1",
         "tqdm>=4.19.6",

setup.cfg

@@ -0,0 +1,5 @@
+# TODO - this entire file can be removed once pymongo-inmemory supports pyproject.toml
+# see https://github.com/kaizendorks/pymongo_inmemory/issues/81
+[pymongo_inmemory]
+mongod_port = 27019
+mongo_version = 6.0

src/maggma/stores/file_store.py

@@ -102,7 +102,6 @@ def __init__(
         self.json_name = json_name
         file_filters = file_filters if file_filters else ["*"]
         self.file_filters = re.compile("|".join(fnmatch.translate(p) for p in file_filters))
-        self.collection_name = "file_store"
         self.key = "file_id"
         self.include_orphans = include_orphans
         self.read_only = read_only
@@ -112,14 +111,12 @@ def __init__(
         self.metadata_store = JSONStore(
             paths=[str(self.path / self.json_name)],
             read_only=self.read_only,
-            collection_name=self.collection_name,
             key=self.key,
         )
 
         self.kwargs = kwargs
 
         super().__init__(
-            collection_name=self.collection_name,
             key=self.key,
             **self.kwargs,
         )

src/maggma/stores/mongolike.py

@@ -5,6 +5,7 @@
 """
 
 import warnings
+from datetime import datetime
 from itertools import chain, groupby
 from pathlib import Path
 
@@ -20,7 +21,6 @@
     from typing_extensions import Literal
 
 import bson
-import mongomock
 import orjson
 from monty.dev import requires
 from monty.io import zopen
@@ -29,6 +29,7 @@
 from pydash import get, has, set_
 from pymongo import MongoClient, ReplaceOne, uri_parser
 from pymongo.errors import ConfigurationError, DocumentTooLarge, OperationFailure
+from pymongo_inmemory import MongoClient as MemoryClient
 
 from maggma.core import Sort, Store, StoreError
 from maggma.utils import confirm_field_index, to_dt
@@ -67,10 +68,12 @@
             port: TCP port to connect to
             username: Username for the collection
             password: Password to connect with
+            ssh_tunnel: SSHTunnel instance to use for connection.
             safe_update: fail gracefully on DocumentTooLarge errors on update
             auth_source: The database to authenticate on. Defaults to the database name.
             default_sort: Default sort field and direction to use when querying. Can be used to
                 ensure determinacy in query results.
+            mongoclient_kwargs: Dict of extra kwargs to pass to MongoClient()
         """
         self.database = database
         self.collection_name = collection_name
@@ -437,7 +440,7 @@
         Check equality for MongoStore
         other: other mongostore to compare with.
         """
-        if not isinstance(other, MongoStore):
+        if not isinstance(other, self.__class__):
             return False
 
         fields = ["database", "collection_name", "host", "port", "last_updated_field"]
@@ -516,98 +519,67 @@
     to a MongoStore.
     """
 
-    def __init__(self, collection_name: str = "memory_db", **kwargs):
+    def __init__(
+        self,
+        database: str = "mem",
+        collection_name: Optional[str] = None,
+        host: str = "localhost",
+        port: int = 27019,  # to avoid conflicts with localhost
+        safe_update: bool = False,
+        mongoclient_kwargs: Optional[Dict] = None,
+        default_sort: Optional[Dict[str, Union[Sort, int]]] = None,
+        **kwargs,
+    ):
         """
-        Initializes the Memory Store.
-
         Args:
-            collection_name: name for the collection in memory.
-        """
-        self.collection_name = collection_name
-        self.default_sort = None
-        self._coll = None
-        self.kwargs = kwargs
-        super(MongoStore, self).__init__(**kwargs)
+            database: The database name
+            collection_name: The collection name. If None (default) a unique collection name is set based
+                on the current date and time. This ensures that multiple Store instances can coexist without
+                overwriting one another.
+            host: Hostname for the database
+            port: TCP port to connect to
+            safe_update: fail gracefully on DocumentTooLarge errors on update
+            default_sort: Default sort field and direction to use when querying.
+                Can be used to ensure determinacy in query results.
+            mongoclient_kwargs: Dict of extra kwargs to pass to MongoClient()
+        """
+        if not collection_name:
+            collection_name = str(datetime.utcnow())
+        super().__init__(
+            database=database,
+            collection_name=collection_name,
+            host=host,
+            port=port,
+            safe_update=safe_update,
+            mongoclient_kwargs=mongoclient_kwargs,
+            default_sort=default_sort,
+            **kwargs,
+        )
 
     def connect(self, force_reset: bool = False):
         """
-        Connect to the source data.
-
-        Args:
-            force_reset: whether to reset the connection or not when the Store is
-                already connected.
+        Connect to the source data
         """
         if self._coll is None or force_reset:
-            self._coll = mongomock.MongoClient().db[self.name]  # type: ignore
-
-    def close(self):
-        """Close up all collections."""
-        self._coll.database.client.close()
+            conn: MemoryClient = MemoryClient(
+                host=self.host,
+                port=self.port,
+                **self.mongoclient_kwargs,
+            )
+            db = conn[self.database]
+            self._coll = db[self.collection_name]  # type: ignore
 
     @property
     def name(self):
-        """Name for the store."""
-        return f"mem://{self.collection_name}"
+        """Name for the store"""
+        return f"mem://{self.database}/{self.collection_name}"
 
-    def __hash__(self):
-        """Hash for the store."""
-        return hash((self.name, self.last_updated_field))
-
-    def groupby(
-        self,
-        keys: Union[List[str], str],
-        criteria: Optional[Dict] = None,
-        properties: Union[Dict, List, None] = None,
-        sort: Optional[Dict[str, Union[Sort, int]]] = None,
-        skip: int = 0,
-        limit: int = 0,
-    ) -> Iterator[Tuple[Dict, List[Dict]]]:
-        """
-        Simple grouping function that will group documents
-        by keys.
-
-        Args:
-            keys: fields to group documents
-            criteria: PyMongo filter for documents to search in
-            properties: properties to return in grouped documents
-            sort: Dictionary of sort order for fields. Keys are field names and
-                values are 1 for ascending or -1 for descending.
-            skip: number documents to skip
-            limit: limit on total number of documents returned
-
-        Returns:
-            generator returning tuples of (key, list of elements)
-        """
-        keys = keys if isinstance(keys, list) else [keys]
-
-        if properties is None:
-            properties = []
-        if isinstance(properties, dict):
-            properties = list(properties.keys())
-
-        data = [
-            doc for doc in self.query(properties=keys + properties, criteria=criteria) if all(has(doc, k) for k in keys)
-        ]
-
-        def grouping_keys(doc):
-            return tuple(get(doc, k) for k in keys)
-
-        for vals, group in groupby(sorted(data, key=grouping_keys), key=grouping_keys):
-            doc = {}  # type: ignore
-            for k, v in zip(keys, vals):
-                set_(doc, k, v)
-            yield doc, list(group)
-
-    def __eq__(self, other: object) -> bool:
-        """
-        Check equality for MemoryStore
-        other: other MemoryStore to compare with.
-        """
-        if not isinstance(other, MemoryStore):
-            return False
-
-        fields = ["collection_name", "last_updated_field"]
-        return all(getattr(self, f) == getattr(other, f) for f in fields)
+    # def __del__(self):
+    #     """
+    #     Ensure collection is dropped from memory on object destruction, even if .close() has not been called.
+    #     """
+    #     if self._coll is not None:
+    #         self._collection.drop()
 
 
 class JSONStore(MemoryStore):
@@ -689,7 +661,8 @@
             on systems with slow storage when multiple connect / disconnects are performed.
         """
         if self._coll is None or force_reset:
-            self._coll = mongomock.MongoClient().db[self.name]  # type: ignore
+            # self._coll = mongomock.MongoClient().db[self.name]  # type: ignore
+            self._coll = MemoryClient().db[self.name]  # type: ignore
 
             # create the .json file if it does not exist
             if not self.read_only and not Path(self.paths[0]).exists():
@@ -791,7 +764,7 @@
         Args:
             other: other JSONStore to compare with
         """
-        if not isinstance(other, JSONStore):
+        if not isinstance(other, self.__class__):
             return False
 
         fields = ["paths", "last_updated_field"]
@@ -803,7 +776,7 @@
     "MontyStore requires MontyDB to be installed. See the MontyDB repository for more "
     "information: https://github.com/davidlatwe/montydb",
 )
-class MontyStore(MemoryStore):
+class MontyStore(MongoStore):
     """
     A MongoDB compatible store that uses on disk files for storage.
 
@@ -944,3 +917,59 @@
                 search_doc = {k: d[k] for k in key} if isinstance(key, list) else {key: d[key]}
 
                 self._collection.replace_one(search_doc, d, upsert=True)
+
+    def groupby(
+        self,
+        keys: Union[List[str], str],
+        criteria: Optional[Dict] = None,
+        properties: Union[Dict, List, None] = None,
+        sort: Optional[Dict[str, Union[Sort, int]]] = None,
+        skip: int = 0,
+        limit: int = 0,
+    ) -> Iterator[Tuple[Dict, List[Dict]]]:
+        """
+        Simple grouping function that will group documents
+        by keys.
+
+        Args:
+            keys: fields to group documents
+            criteria: PyMongo filter for documents to search in
+            properties: properties to return in grouped documents
+            sort: Dictionary of sort order for fields. Keys are field names and
+                values are 1 for ascending or -1 for descending.
+            skip: number documents to skip
+            limit: limit on total number of documents returned
+
+        Returns:
+            generator returning tuples of (key, list of elements)
+        """
+        keys = keys if isinstance(keys, list) else [keys]
+
+        if properties is None:
+            properties = []
+        if isinstance(properties, dict):
+            properties = list(properties.keys())
+
+        data = [
+            doc for doc in self.query(properties=keys + properties, criteria=criteria) if all(has(doc, k) for k in keys)
+        ]
+
+        def grouping_keys(doc):
+            return tuple(get(doc, k) for k in keys)
+
+        for vals, group in groupby(sorted(data, key=grouping_keys), key=grouping_keys):
+            doc = {}  # type: ignore
+            for k, v in zip(keys, vals):
+                set_(doc, k, v)
+            yield doc, list(group)
+
+    def __eq__(self, other: object) -> bool:
+        """
+        Check equality for MontyStore
+        other: other Store to compare with
+        """
+        if not isinstance(other, self.__class__):
+            return False
+
+        fields = ["database_name", "collection_name", "last_updated_field"]
+        return all(getattr(self, f) == getattr(other, f) for f in fields)

tests/api/test_submission_resource.py

@@ -1,4 +1,3 @@
-import json
 from datetime import datetime
 from random import randint
 
@@ -113,10 +112,10 @@ def test_submission_patch(owner_store, post_query_op, patch_query_op):
     app.include_router(endpoint.router)
 
     client = TestClient(app)
-    update = json.dumps({"last_updated": "2023-06-22T17:32:11.645713"})
+    update = {"last_updated": "2023-06-22T17:32:11.645713"}
 
     assert client.get("/").status_code == 200
-    assert client.patch(f"/?name=PersonAge9&update={update}").status_code == 200
+    assert client.patch("/?name=PersonAge9", json=update).status_code == 200
 
 
 def test_key_fields(owner_store, post_query_op):

tests/builders/test_copy_builder.py

@@ -102,7 +102,8 @@ def test_run(source, target, old_docs, new_docs):
 
     builder = CopyBuilder(source, target)
     builder.run()
-    builder.target.connect()
+
+    target.connect()
     assert builder.target.query_one(criteria={"k": 0})["v"] == "new"
     assert builder.target.query_one(criteria={"k": 10})["v"] == "old"
 
@@ -113,6 +114,8 @@ def test_query(source, target, old_docs, new_docs):
     source.update(old_docs)
     source.update(new_docs)
     builder.run()
+
+    target.connect()
     all_docs = list(target.query(criteria={}))
     assert len(all_docs) == 14
     assert min([d["k"] for d in all_docs]) == 6
@@ -128,6 +131,7 @@ def test_delete_orphans(source, target, old_docs, new_docs):
     source._collection.delete_many(deletion_criteria)
     builder.run()
 
+    target.connect()
     assert target._collection.count_documents(deletion_criteria) == 0
     assert target.query_one(criteria={"k": 5})["v"] == "new"
     assert target.query_one(criteria={"k": 10})["v"] == "old"