From 24e95fa4aaa1cc680ea7fc7b52418805755a838b Mon Sep 17 00:00:00 2001 From: Ryan Kingsbury Date: Wed, 8 Jul 2026 11:10:59 -0400 Subject: [PATCH 1/5] Back MemoryStore with a real MongoDB when available, fall back to mongomock-ng MemoryStore now probes host:port (default localhost:27017) on connect and, if a MongoDB server answers, stores data in a real, ephemeral MongoDB database for full MongoDB compatibility and performance. Each instance uses a unique maggma_memory_ database that is dropped automatically (via weakref finalizer) when the store is garbage collected or the interpreter exits, so the user never has to set up or clean up a database. If no server is reachable the store transparently falls back to an in-process mongomock-ng database. - close() is now a no-op that leaves the store usable, preserving the historical mongomock behavior relied upon by the builders (pymongo raises on use-after-close). Use connect(force_reset=True) to discard contents. - JSONStore/FileStore inherit the behavior via a shared _connect_collection; MontyStore gets its own close() so its disk-backed client still closes. - Adds backward-compatible host/port/mongoclient_kwargs/server_selection_timeout_ms kwargs; server_selection_timeout_ms=0 forces the mongomock backend. Co-Authored-By: Claude Opus 4.8 --- src/maggma/stores/mongolike.py | 133 +++++++++++++++++++++++++++++++-- tests/stores/test_mongolike.py | 50 ++++++++++++- 2 files changed, 173 insertions(+), 10 deletions(-) diff --git a/src/maggma/stores/mongolike.py b/src/maggma/stores/mongolike.py index 25e48f5d7..0e5486c28 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,126 @@ 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. + + If a MongoDB server is reachable (by default on ``localhost:27017``), the + data is stored in a real, ephemeral MongoDB database for full MongoDB + compatibility and performance. That database is namespaced uniquely per + Store instance and is 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. If no server is reachable, the Store + transparently falls back to an in-process ``mongomock`` database, so it + works even with no MongoDB installed. """ - 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 = 27017, + 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 to probe for a running MongoDB server to back the + Store with. The data is written to an ephemeral database that + is dropped when the Store is closed. + port: TCP port to probe for a running MongoDB server. + mongoclient_kwargs: Dict of extra kwargs to pass to MongoClient when + a real MongoDB backend is used. + server_selection_timeout_ms: how long, in milliseconds, to wait when + probing ``host:port`` for a MongoDB server before falling back + to an in-process ``mongomock`` database. Set to 0 (or less) to + skip the probe entirely and always use ``mongomock``. """ 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._coll = None self.kwargs = kwargs super(MongoStore, self).__init__(**kwargs) + def _get_memory_client(self) -> MongoClient: + """ + Return a client backing the in-memory Store. + + Attempts to connect to a real MongoDB server at ``host:port`` for full + MongoDB compatibility and performance. If none is reachable within + ``server_selection_timeout_ms``, falls back to an in-process + ``mongomock`` client so the Store works without a MongoDB server. + """ + if self.server_selection_timeout_ms > 0: + mongoclient_kwargs = dict(self.mongoclient_kwargs) + mongoclient_kwargs.setdefault("serverSelectionTimeoutMS", self.server_selection_timeout_ms) + try: + client = MongoClient(host=self.host, port=self.port, **mongoclient_kwargs) + # force server selection to confirm a server is actually reachable + client.admin.command("ping") + 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 + except Exception: + self.logger.debug(f"{self.name}: no MongoDB server reachable, falling back to mongomock") + + self._using_real_mongo = False + return mongomock.MongoClient() # type: ignore + + @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: + pass + + def _connect_collection(self, force_reset: bool = False): + """Establish (or re-establish) the underlying in-memory collection.""" + if force_reset and self._coll is not None: + old_client = self._coll.database.client + # on a forced reset, discard the previous contents (matching the + # historical mongomock behavior of starting from a fresh client) + if getattr(self, "_using_real_mongo", False): + try: + old_client.drop_database(self._database) + except Exception: + pass + old_client.close() + client = self._get_memory_client() + self._coll = client[self._database][self.collection_name] # type: ignore + def connect(self, force_reset: bool = False): """ Connect to the source data. @@ -531,11 +636,20 @@ def connect(self, force_reset: bool = False): already connected. """ if self._coll is None or force_reset: - self._coll = mongomock.MongoClient().db[self.name] # type: ignore + self._connect_collection(force_reset=force_reset) def close(self): - """Close up all collections.""" - self._coll.database.client.close() + """Close up all collections. + + For an in-memory Store this is intentionally a no-op that leaves the + Store usable, matching the historical ``mongomock`` behavior (its + ``close()`` did nothing) that callers such as the builders rely on when + they query a Store after it has been closed. Resources are released and + any ephemeral MongoDB database backing the Store is dropped when the + Store is garbage collected or at interpreter exit (see + ``_drop_ephemeral_database``). Use ``force_reset=True`` on ``connect()`` + to explicitly discard the contents and start fresh. + """ @property def name(self): @@ -682,7 +796,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._connect_collection(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 +991,11 @@ 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.""" + if self._coll is not None: + self._coll.database.client.close() + @property def name(self) -> str: """Return a string representing this data source.""" diff --git a/tests/stores/test_mongolike.py b/tests/stores/test_mongolike.py index fb264953b..e3ae08a3f 100644 --- a/tests/stores/test_mongolike.py +++ b/tests/stores/test_mongolike.py @@ -248,13 +248,53 @@ def test_mongostore_newer_in(mongostore): # Memory store tests -def test_memory_store_connect(): - memorystore = MemoryStore() +def test_memory_store_connect_mongomock_fallback(): + # server_selection_timeout_ms=0 skips the probe and forces the mongomock backend + memorystore = MemoryStore(server_selection_timeout_ms=0) 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_connect_unreachable_falls_back(): + # an unreachable server should fall back to mongomock rather than raise + memorystore = MemoryStore(port=1, server_selection_timeout_ms=50) + 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_available(): + # a real MongoDB server is available in CI; when present it should be used + try: + pymongo.MongoClient(serverSelectionTimeoutMS=500).admin.command("ping") + except Exception: + pytest.skip("no MongoDB server reachable on localhost:27017") + + memorystore = MemoryStore() + 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,8 +562,11 @@ def test_jsonstore_orjson_options(test_dir): class SubFloat(float): pass + # Force the mongomock backend (server_selection_timeout_ms=0): a real MongoDB + # backend coerces the SubFloat subclass to a plain float on write/read, so the + # serialization_default option this test exercises would never be triggered. with ScratchDir("."): - jsonstore = JSONStore("d.json", read_only=False) + jsonstore = JSONStore("d.json", read_only=False, server_selection_timeout_ms=0) jsonstore.connect() with pytest.raises(orjson.JSONEncodeError): jsonstore.update({"wrong_field": SubFloat(1.1), "task_id": 3}) @@ -534,6 +577,7 @@ class SubFloat(float): read_only=False, serialization_option=None, serialization_default=lambda x: "test", + server_selection_timeout_ms=0, ) jsonstore.connect() jsonstore.update({"wrong_field": SubFloat(1.1), "task_id": 3}) From ee02544900335acea4b47d2a5987847ebd6714c1 Mon Sep 17 00:00:00 2001 From: Ryan Kingsbury Date: Wed, 8 Jul 2026 11:28:42 -0400 Subject: [PATCH 2/5] Add explanatory comments to best-effort except clauses Addresses GitHub AI security warnings about except clauses that only pass without an explanatory comment (the two best-effort ephemeral-database drops). Co-Authored-By: Claude Opus 4.8 --- src/maggma/stores/mongolike.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/maggma/stores/mongolike.py b/src/maggma/stores/mongolike.py index 0e5486c28..112d757a7 100644 --- a/src/maggma/stores/mongolike.py +++ b/src/maggma/stores/mongolike.py @@ -610,6 +610,9 @@ def _drop_ephemeral_database(host: str, port: int, database: str, mongoclient_kw 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 _connect_collection(self, force_reset: bool = False): @@ -622,6 +625,9 @@ def _connect_collection(self, force_reset: bool = False): try: old_client.drop_database(self._database) except Exception: + # Best-effort discard of the previous contents; if the drop + # fails the database will still be cleaned up by the finalizer + # on garbage collection or interpreter exit. pass old_client.close() client = self._get_memory_client() From 07ca34e263833db99787aac2ef967b5910533955 Mon Sep 17 00:00:00 2001 From: Ryan Kingsbury Date: Wed, 8 Jul 2026 11:43:41 -0400 Subject: [PATCH 3/5] Make MemoryStore.close() a genuine close; enforce query-before-close Previously MemoryStore.close() was a no-op so that callers could keep querying a "closed" store. This makes close() actually disconnect: querying a closed store now raises StoreError, and the store must be reconnected first. Data is preserved across a close()/connect() cycle -- the real-MongoDB backend keeps it in the ephemeral server-side database (dropped by the finalizer on GC/exit), and the mongomock backend keeps its in-process client alive -- so reconnect restores the contents while connect(force_reset=True) starts fresh. The Builder framework already only queries stores before closing them (run/serial/multi all do connect -> get_items -> update_targets -> finalize, and MapBuilder.finalize queries before super().finalize()). The only query-after-close was in the builder test suite, which relied on the old no-op close(); those tests now reconnect the target before asserting, matching the existing pattern in test_copy_builder.test_run. Co-Authored-By: Claude Opus 4.8 --- src/maggma/stores/mongolike.py | 87 ++++++++++++++--------- tests/builders/test_copy_builder.py | 2 + tests/builders/test_projection_builder.py | 2 + 3 files changed, 59 insertions(+), 32 deletions(-) diff --git a/src/maggma/stores/mongolike.py b/src/maggma/stores/mongolike.py index 112d757a7..07f81339f 100644 --- a/src/maggma/stores/mongolike.py +++ b/src/maggma/stores/mongolike.py @@ -541,7 +541,8 @@ def __init__( collection_name: name for the collection in memory. host: hostname to probe for a running MongoDB server to back the Store with. The data is written to an ephemeral database that - is dropped when the Store is closed. + is dropped when the Store is garbage collected or the + interpreter exits. port: TCP port to probe for a running MongoDB server. mongoclient_kwargs: Dict of extra kwargs to pass to MongoClient when a real MongoDB backend is used. @@ -559,6 +560,7 @@ def __init__( # 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) @@ -602,8 +604,10 @@ def _get_memory_client(self) -> MongoClient: @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.""" + """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) @@ -615,23 +619,33 @@ def _drop_ephemeral_database(host: str, port: int, database: str, mongoclient_kw # a finalizer/atexit hook never raises. pass - def _connect_collection(self, force_reset: bool = False): - """Establish (or re-establish) the underlying in-memory collection.""" - if force_reset and self._coll is not None: - old_client = self._coll.database.client - # on a forced reset, discard the previous contents (matching the - # historical mongomock behavior of starting from a fresh client) - if getattr(self, "_using_real_mongo", False): + 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: - old_client.drop_database(self._database) + self._client.drop_database(self._database) except Exception: - # Best-effort discard of the previous contents; if the drop - # fails the database will still be cleaned up by the finalizer - # on garbage collection or interpreter exit. + # Best-effort discard; the finalizer will still drop the + # database on garbage collection or interpreter exit. pass - old_client.close() - client = self._get_memory_client() - self._coll = client[self._database][self.collection_name] # type: ignore + 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): """ @@ -641,21 +655,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._connect_collection(force_reset=force_reset) + self._ensure_connected(force_reset=force_reset) def close(self): - """Close up all collections. - - For an in-memory Store this is intentionally a no-op that leaves the - Store usable, matching the historical ``mongomock`` behavior (its - ``close()`` did nothing) that callers such as the builders rely on when - they query a Store after it has been closed. Resources are released and - any ephemeral MongoDB database backing the Store is dropped when the - Store is garbage collected or at interpreter exit (see - ``_drop_ephemeral_database``). Use ``force_reset=True`` on ``connect()`` - to explicitly discard the contents and start fresh. - """ + """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): @@ -802,7 +824,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._connect_collection(force_reset=force_reset) + 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(): @@ -998,9 +1020,10 @@ def connect(self, force_reset: bool = False): self._coll = client[self.database_name][self.collection_name] def close(self): - """Close up the MontyDB client.""" + """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: 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 From 96cb7566d225ee005aa2135670ae2755b2a16b74 Mon Sep 17 00:00:00 2001 From: Ryan Kingsbury Date: Wed, 8 Jul 2026 12:16:28 -0400 Subject: [PATCH 4/5] Update S3Store test_close for genuine MemoryStore.close() test_close expected an AttributeError from querying a closed S3Store (because S3Store.close() sets s3_bucket=None). With MemoryStore.close() now being a genuine close, the S3 index is genuinely closed and raises StoreError when queried, before the s3_bucket access is reached. Expect StoreError, which is the more principled signal that the store was closed. Co-Authored-By: Claude Opus 4.8 --- tests/stores/test_aws.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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()) From 7e3fe4b58e670bc4659796aad8b7ed58c60de154 Mon Sep 17 00:00:00 2001 From: Ryan Kingsbury Date: Sun, 12 Jul 2026 22:50:12 -0400 Subject: [PATCH 5/5] Make the real-MongoDB backend opt-in via port MemoryStore previously probed localhost:27017 on every connect and used a real MongoDB server if one answered, falling back to mongomock otherwise. Make the real backend strictly opt-in instead: - The `port` kwarg now defaults to None, in which case the Store always uses the in-process mongomock backend -- no probing and no dependency on a running server, even if one happens to be listening on localhost. - Supplying a port (e.g. port=27017 for a local mongod) opts in to a real MongoDB server at host:port, writing to the ephemeral, auto-dropped database. Removes the auto-probe/ping/fallback path; server_selection_timeout_ms now only controls the real client's timeout. Tests and docstrings updated accordingly. Co-Authored-By: Claude Opus 4.8 --- src/maggma/stores/mongolike.py | 90 ++++++++++++++++------------------ tests/stores/test_mongolike.py | 33 ++++++------- 2 files changed, 57 insertions(+), 66 deletions(-) diff --git a/src/maggma/stores/mongolike.py b/src/maggma/stores/mongolike.py index 07f81339f..690421fa0 100644 --- a/src/maggma/stores/mongolike.py +++ b/src/maggma/stores/mongolike.py @@ -509,14 +509,14 @@ class MemoryStore(MongoStore): """ An in-memory Store that functions similarly to a MongoStore. - If a MongoDB server is reachable (by default on ``localhost:27017``), the - data is stored in a real, ephemeral MongoDB database for full MongoDB - compatibility and performance. That database is namespaced uniquely per - Store instance and is 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. If no server is reachable, the Store - transparently falls back to an in-process ``mongomock`` database, so it - works even with no MongoDB installed. + 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. """ #: Set to True by connect() when a real MongoDB backend is in use. Defined @@ -529,7 +529,7 @@ def __init__( self, collection_name: str = "memory_db", host: str = "localhost", - port: int = 27017, + port: int | None = None, mongoclient_kwargs: dict | None = None, server_selection_timeout_ms: int = 500, **kwargs, @@ -539,17 +539,17 @@ def __init__( Args: collection_name: name for the collection in memory. - host: hostname to probe for a running MongoDB server to back the - Store with. The data is written to an ephemeral database that - is dropped when the Store is garbage collected or the - interpreter exits. - port: TCP port to probe for a running MongoDB server. + 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: how long, in milliseconds, to wait when - probing ``host:port`` for a MongoDB server before falling back - to an in-process ``mongomock`` database. Set to 0 (or less) to - skip the probe entirely and always use ``mongomock``. + 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 @@ -569,38 +569,34 @@ def _get_memory_client(self) -> MongoClient: """ Return a client backing the in-memory Store. - Attempts to connect to a real MongoDB server at ``host:port`` for full - MongoDB compatibility and performance. If none is reachable within - ``server_selection_timeout_ms``, falls back to an in-process - ``mongomock`` client so the Store works without a MongoDB server. + 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 = dict(self.mongoclient_kwargs) mongoclient_kwargs.setdefault("serverSelectionTimeoutMS", self.server_selection_timeout_ms) - try: - client = MongoClient(host=self.host, port=self.port, **mongoclient_kwargs) - # force server selection to confirm a server is actually reachable - client.admin.command("ping") - 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 - except Exception: - self.logger.debug(f"{self.name}: no MongoDB server reachable, falling back to mongomock") - - self._using_real_mongo = False - return mongomock.MongoClient() # type: ignore + 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): diff --git a/tests/stores/test_mongolike.py b/tests/stores/test_mongolike.py index e3ae08a3f..93f5727d2 100644 --- a/tests/stores/test_mongolike.py +++ b/tests/stores/test_mongolike.py @@ -248,31 +248,26 @@ def test_mongostore_newer_in(mongostore): # Memory store tests -def test_memory_store_connect_mongomock_fallback(): - # server_selection_timeout_ms=0 skips the probe and forces the mongomock backend - memorystore = MemoryStore(server_selection_timeout_ms=0) +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_connect_unreachable_falls_back(): - # an unreachable server should fall back to mongomock rather than raise - memorystore = MemoryStore(port=1, server_selection_timeout_ms=50) - 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_available(): - # a real MongoDB server is available in CI; when present it should be used +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() + memorystore = MemoryStore(port=27017) memorystore.connect() assert memorystore._using_real_mongo is True assert isinstance(memorystore._collection, pymongo.collection.Collection) @@ -562,11 +557,12 @@ def test_jsonstore_orjson_options(test_dir): class SubFloat(float): pass - # Force the mongomock backend (server_selection_timeout_ms=0): a real MongoDB - # backend coerces the SubFloat subclass to a plain float on write/read, so the - # serialization_default option this test exercises would never be triggered. + # 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, server_selection_timeout_ms=0) + jsonstore = JSONStore("d.json", read_only=False) jsonstore.connect() with pytest.raises(orjson.JSONEncodeError): jsonstore.update({"wrong_field": SubFloat(1.1), "task_id": 3}) @@ -577,7 +573,6 @@ class SubFloat(float): read_only=False, serialization_option=None, serialization_default=lambda x: "test", - server_selection_timeout_ms=0, ) jsonstore.connect() jsonstore.update({"wrong_field": SubFloat(1.1), "task_id": 3})