diff --git a/src/maggma/stores/mongolike.py b/src/maggma/stores/mongolike.py index 25e48f5d7..690421fa0 100644 --- a/src/maggma/stores/mongolike.py +++ b/src/maggma/stores/mongolike.py @@ -5,10 +5,12 @@ """ import warnings +import weakref from collections.abc import Callable, Iterator from itertools import chain, groupby from pathlib import Path from typing import Any, Literal +from uuid import uuid4 import bson import mongomock_ng as mongomock @@ -505,23 +507,142 @@ def connect(self, force_reset: bool = False): class MemoryStore(MongoStore): """ - An in-memory Store that functions similarly - to a MongoStore. + An in-memory Store that functions similarly to a MongoStore. + + By default (``port=None``) the data is held in an in-process ``mongomock`` + database, so the Store works with no MongoDB server installed. As an opt-in, + if a ``port`` is supplied the Store instead uses a real MongoDB server at + ``host:port`` (e.g. a local ``mongod``) for full MongoDB compatibility and + performance. In that case the data is written to an ephemeral database that + is namespaced uniquely per Store instance and dropped automatically when the + Store is garbage collected or the interpreter exits, so the user never has to + create or clean up a database manually. """ - def __init__(self, collection_name: str = "memory_db", **kwargs): + #: Set to True by connect() when a real MongoDB backend is in use. Defined + #: at class level so close() is safe on subclasses (e.g. MontyStore) that + #: do not call MemoryStore.__init__. + _using_real_mongo: bool = False + _finalizer = None + + def __init__( + self, + collection_name: str = "memory_db", + host: str = "localhost", + port: int | None = None, + mongoclient_kwargs: dict | None = None, + server_selection_timeout_ms: int = 500, + **kwargs, + ): """ Initializes the Memory Store. Args: collection_name: name for the collection in memory. + host: hostname of the MongoDB server to use when ``port`` is supplied. + port: TCP port of a MongoDB server to back the Store with. If None + (the default), the Store uses an in-process ``mongomock`` + database and no server is required. If a port is supplied (e.g. + ``port=27017`` for a local ``mongod``), the Store uses that real + MongoDB server, writing to an ephemeral database that is dropped + when the Store is garbage collected or the interpreter exits. + mongoclient_kwargs: Dict of extra kwargs to pass to MongoClient when + a real MongoDB backend is used. + server_selection_timeout_ms: serverSelectionTimeoutMS to use for the + real MongoDB client. Only relevant when ``port`` is supplied. """ self.collection_name = collection_name + self.host = host + self.port = port + self.mongoclient_kwargs = mongoclient_kwargs or {} + self.server_selection_timeout_ms = server_selection_timeout_ms + # unique, ephemeral database name so that multiple MemoryStore instances + # backed by the same real MongoDB server do not clobber one another + self._database = f"maggma_memory_{uuid4().hex}" self.default_sort = None + self._client = None self._coll = None self.kwargs = kwargs super(MongoStore, self).__init__(**kwargs) + def _get_memory_client(self) -> MongoClient: + """ + Return a client backing the in-memory Store. + + Uses a real MongoDB server at ``host:port`` when a ``port`` was supplied + (opt-in, for full MongoDB compatibility and performance), otherwise an + in-process ``mongomock`` client so the Store works without a server. + """ + if self.port is None: + # default: no port supplied -> in-process mongomock backend + self._using_real_mongo = False + return mongomock.MongoClient() # type: ignore + + # a port was supplied -> use a real MongoDB server at host:port + mongoclient_kwargs = dict(self.mongoclient_kwargs) + if self.server_selection_timeout_ms > 0: + mongoclient_kwargs.setdefault("serverSelectionTimeoutMS", self.server_selection_timeout_ms) + client = MongoClient(host=self.host, port=self.port, **mongoclient_kwargs) + self._using_real_mongo = True + # ensure the ephemeral database is dropped when this Store is garbage + # collected or the interpreter exits, so nothing is left behind on the server + if self._finalizer is None: + self._finalizer = weakref.finalize( + self, + self._drop_ephemeral_database, + self.host, + self.port, + self._database, + dict(mongoclient_kwargs), + ) + self.logger.debug(f"{self.name} using real MongoDB backend at {self.host}:{self.port}") + return client + + @staticmethod + def _drop_ephemeral_database(host: str, port: int, database: str, mongoclient_kwargs: dict): + """Drop an ephemeral MongoDB database. + + Used as a weakref finalizer, so it must not hold a reference to the Store. + """ + mongoclient_kwargs.setdefault("serverSelectionTimeoutMS", 500) + try: + client = MongoClient(host=host, port=port, **mongoclient_kwargs) + client.drop_database(database) + client.close() + except Exception: + # Best-effort cleanup: if the server is already gone or unreachable at + # finalization time there is nothing left to drop. Swallow the error so + # a finalizer/atexit hook never raises. + pass + + def _reset_connection(self): + """Tear down the current connection and discard its contents.""" + if getattr(self, "_using_real_mongo", False): + # discard the ephemeral database, even if we are currently + # disconnected (e.g. after close(), when self._client is None) + if self._client is not None: + try: + self._client.drop_database(self._database) + except Exception: + # Best-effort discard; the finalizer will still drop the + # database on garbage collection or interpreter exit. + pass + else: + self._drop_ephemeral_database(self.host, self.port, self._database, dict(self.mongoclient_kwargs)) + if self._client is not None: + self._client.close() + self._client = None + self._coll = None + + def _ensure_connected(self, force_reset: bool = False): + """Establish (or re-establish) the underlying in-memory collection.""" + if force_reset: + self._reset_connection() + if self._coll is None: + if self._client is None: + self._client = self._get_memory_client() + self._coll = self._client[self._database][self.collection_name] # type: ignore + def connect(self, force_reset: bool = False): """ Connect to the source data. @@ -530,12 +651,29 @@ def connect(self, force_reset: bool = False): force_reset: whether to reset the connection or not when the Store is already connected. """ - if self._coll is None or force_reset: - self._coll = mongomock.MongoClient().db[self.name] # type: ignore + self._ensure_connected(force_reset=force_reset) def close(self): - """Close up all collections.""" - self._coll.database.client.close() + """Close the Store's connection. + + After ``close()`` the Store must be reconnected with ``connect()`` before + it can be queried again; querying a closed Store raises a ``StoreError``. + The Store's data is preserved across a ``close()``/``connect()`` cycle: + with a real MongoDB backend it lives in the ephemeral database on the + server (dropped only when the Store is garbage collected or at + interpreter exit, see ``_drop_ephemeral_database``); with the + ``mongomock`` backend it lives in the in-process client, which is kept + alive for this reason. Use ``connect(force_reset=True)`` to discard the + contents and start fresh. + """ + if getattr(self, "_using_real_mongo", False) and self._client is not None: + # data persists in the ephemeral database on the server, so we can + # release the client connection now and reconnect later on demand + self._client.close() + self._client = None + # for the mongomock backend the data lives in the in-process client; + # keep it alive and simply mark the Store as needing a reconnect + self._coll = None @property def name(self): @@ -682,7 +820,7 @@ def connect(self, force_reset: bool = False): 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._ensure_connected(force_reset=force_reset) # create the .json file if it does not exist if not self.read_only and not Path(self.paths[0]).exists(): @@ -877,6 +1015,12 @@ def connect(self, force_reset: bool = False): client = MontyClient(self.database_path, **self.client_kwargs) self._coll = client[self.database_name][self.collection_name] + def close(self): + """Close up the MontyDB client. The Store must be reconnected before use.""" + if self._coll is not None: + self._coll.database.client.close() + self._coll = None + @property def name(self) -> str: """Return a string representing this data source.""" diff --git a/tests/builders/test_copy_builder.py b/tests/builders/test_copy_builder.py index 68b0dd284..7b85e0927 100644 --- a/tests/builders/test_copy_builder.py +++ b/tests/builders/test_copy_builder.py @@ -114,6 +114,7 @@ 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 @@ -129,6 +130,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" diff --git a/tests/builders/test_projection_builder.py b/tests/builders/test_projection_builder.py index f11f9abe2..fc2295bf6 100644 --- a/tests/builders/test_projection_builder.py +++ b/tests/builders/test_projection_builder.py @@ -104,6 +104,7 @@ def test_update_targets(source1, source2, target): def test_run(source1, source2, target): builder = Projection_Builder(source_stores=[source1, source2], target_store=target) builder.run() + target.connect() assert len(list(target.query())) == 15 assert target.query_one(criteria={"k": 0})["a"] == "a" assert target.query_one(criteria={"k": 0})["d"] == "d" @@ -119,4 +120,5 @@ def test_query(source1, source2, target): query_by_key=[0, 1, 2, 3, 4], ) builder.run() + target.connect() assert len(list(target.query())) == 5 diff --git a/tests/stores/test_aws.py b/tests/stores/test_aws.py index 3d86dddca..78a2da19f 100644 --- a/tests/stores/test_aws.py +++ b/tests/stores/test_aws.py @@ -6,6 +6,7 @@ from botocore.exceptions import ClientError from moto import mock_aws +from maggma.core import StoreError from maggma.stores import MemoryStore, MongoStore, S3Store from maggma.stores.ssh_tunnel import SSHTunnel @@ -232,7 +233,9 @@ def objects_in_bucket(key): def test_close(s3store): list(s3store.query()) s3store.close() - with pytest.raises(AttributeError): + # querying a closed store raises: the index (a MemoryStore) is genuinely + # closed, so it raises a StoreError before the S3 access is reached + with pytest.raises(StoreError): list(s3store.query()) diff --git a/tests/stores/test_mongolike.py b/tests/stores/test_mongolike.py index fb264953b..93f5727d2 100644 --- a/tests/stores/test_mongolike.py +++ b/tests/stores/test_mongolike.py @@ -248,13 +248,48 @@ def test_mongostore_newer_in(mongostore): # Memory store tests -def test_memory_store_connect(): +def test_memory_store_connect_default_mongomock(): + # by default (no port supplied) the Store uses the in-process mongomock + # backend. Note this runs in CI with a MongoDB server available, so it also + # confirms the real-mongo backend is strictly opt-in (the default does not + # touch a running server). memorystore = MemoryStore() assert memorystore._coll is None memorystore.connect() + assert memorystore._using_real_mongo is False assert isinstance(memorystore._collection, mongomock_ng.collection.Collection) +def test_memory_store_uses_real_mongo_when_port_supplied(): + # supplying a port opts in to a real MongoDB backend + try: + pymongo.MongoClient(serverSelectionTimeoutMS=500).admin.command("ping") + except Exception: + pytest.skip("no MongoDB server reachable on localhost:27017") + + memorystore = MemoryStore(port=27017) + memorystore.connect() + assert memorystore._using_real_mongo is True + assert isinstance(memorystore._collection, pymongo.collection.Collection) + + # the ephemeral database exists while connected + verify_client = pymongo.MongoClient(serverSelectionTimeoutMS=500) + memorystore.update({"task_id": 1, "val": 2}) + assert memorystore._database in verify_client.list_database_names() + + # data survives close() (the Store remains usable), matching the historical + # in-memory behavior relied upon by builders + memorystore.close() + memorystore.connect() + assert memorystore.count() == 1 + + # the ephemeral database is dropped when the Store is finalized + database_name = memorystore._database + memorystore._finalizer() + assert database_name not in verify_client.list_database_names() + verify_client.close() + + def test_groupby(memorystore): memorystore.update( [ @@ -522,6 +557,10 @@ def test_jsonstore_orjson_options(test_dir): class SubFloat(float): pass + # JSONStore uses the default mongomock backend (no port), which preserves the + # SubFloat subclass so orjson raises. A real MongoDB backend would instead + # coerce it to a plain float, and the serialization_default option this test + # exercises would never be triggered. with ScratchDir("."): jsonstore = JSONStore("d.json", read_only=False) jsonstore.connect()