Back MemoryStore with a real MongoDB when available, fall back to mongomock-ng

[rkingsbury]

Summary

Building on #1115 (mongomock → mongomock-ng), this lets MemoryStore opt in to a real, ephemeral MongoDB backend when one is available, while keeping the in-process mongomock backend as the default — with no manual database setup either way.

Motivation (see #830, #788): mongomock is a reimplementation and doesn't cover every MongoDB feature. When a real mongod is available, opting in gives exact MongoDB semantics and better performance, while the default keeps working out-of-the-box with no server installed.

Behavior

MemoryStore takes an optional port:

• port=None (default) → data is held in an in-process mongomock database. No server is required or contacted, even if one happens to be listening on localhost. This is the historical behavior.
• port=<int> (e.g. port=27017) → the Store uses a real MongoDB server at host:port, writing to a database named maggma_memory_<uuid>, unique per store instance so instances never collide. That database is dropped automatically (via a weakref.finalize) when the store is garbage collected or the interpreter exits, so nothing is left behind and the user never creates or cleans up a database manually.

Other new, backward-compatible kwargs: host (used only when a port is supplied), mongoclient_kwargs, server_selection_timeout_ms (the real client's timeout). Existing MemoryStore("name") usage is unchanged and still uses mongomock.

close() semantics

MemoryStore.close() is a genuine close: after it, the store must be reconnected with connect() before it can be queried again (querying a closed store raises StoreError). Data is preserved across a close()/connect() cycle — the real-MongoDB backend keeps it in the ephemeral server-side database, and the mongomock backend keeps its in-process client alive — while connect(force_reset=True) discards the contents and starts fresh.

The build pipeline already only queries stores before closing them: Builder.run() and the serial/multiprocessing runners all do connect → get_items → update_targets → finalize(), and MapBuilder.finalize() queries before super().finalize(). A couple of builder tests previously queried the target after run() had closed it; they now reconnect first (matching the existing pattern in test_copy_builder.test_run).

Other notes

• JSONStore and FileStore inherit the behavior through the shared connection helper (and use the default mongomock backend); MontyStore keeps its own close().
• The test_jsonstore_orjson_options test relies on the mongomock backend, since a real MongoDB coerces the custom float subclass before write and the serialization_default option under test would never trigger (this is the backend-dependent case discussed in MemoryStore: replace mongomock with pymongo-inmemory #846).

Testing

• tests/stores/test_mongolike.py, tests/stores/test_file_store.py, tests/stores/test_aws.py, tests/builders with a real mongod running: 74 passed, 9 skipped, 1 xfailed, with no leftover maggma_memory_* databases. (The default-mongomock test runs with a server present and confirms the real backend is strictly opt-in.)
• Verified the port-supplied backend: querying a closed store raises, reconnect restores the data, and the ephemeral database is dropped by the finalizer; and that as_dict/from_dict, MontyDecoder, and deepcopy roundtrips still work.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 87.09677% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.78%. Comparing base (cf97834) to head (d3c0939).

Files with missing lines	Patch %	Lines
src/maggma/stores/mongolike.py	87.09%	8 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1121      +/-   ##
==========================================
+ Coverage   69.53%   69.78%   +0.24%     
==========================================
  Files          47       47              
  Lines        4103     4160      +57     
==========================================
+ Hits         2853     2903      +50     
- Misses       1250     1257       +7     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[rkingsbury]

@Andrew-S-Rosen any thoughts on this? I feel this is the best solution to allowing MemoryStore to more closely emulate true mongo behavior, rather than relying on mocking. Basically, if mongod is running, you'll get a real Mongo collection in memory. Otherwise, mongomock-ng will be used.

[rkingsbury]

Supersedes #846

[rkingsbury]

@utf I'd also value your input, given maggma's use in atomate2

[gpetretto]

Hi @rkingsbury,

I had an initial look at the PR and made a few tests. Since we use the JSONStore in jobflow-remote this might have an impact on us as well.

Here are some comments:

• from what I tested, if the local DB exists but is password protected and the credentials are not set, the MemoryStore just sets _using_real_mongo to True, but then it fails when it tries to insert/access some documents. Maybe it would be good to also check if the DB is accessible with read/write permission before setting _using_real_mongo?

• when the failure happens in inserting a new document the error is difficult to interpret because the methods in the base class MongoStore tries to handle it, but uses attributes that are not defined in the subclass (i.e. safe_update)

  OperationFailure: Command update requires authentication, full error: {'ok': 0.0, 'errmsg': 'Command update requires authentication', 'code': 13, 'codeName': 'Unauthorized'}

  During handling of the above exception, another exception occurred:

  AttributeError Traceback (most recent call last)
  Cell In[11], line 1
  ----> 1 ms.update({"test":1, "task_id": 1})

  File ~/tmp/maggma/src/maggma/stores/mongolike.py:399, in MongoStore.update(self, docs, key)
  397 self._collection.bulk_write(requests, ordered=False)
  398 except (OperationFailure, DocumentTooLarge) as e:
  --> 399 if self.safe_update:
  400 for req in requests:
  401 try:

  AttributeError: 'MemoryStore' object has no attribute 'safe_update'

  In this context it seems problematic that the __init__ of the MongoStore is not called. Not sure if it would be possible to actually call it, to ensure that all the attributes are properly set.

• Unless for an existing passwordless local mongodb, the basic configuration would not allow to connect. I wonder if it is a common occurrence for a user to have such a configuration. Wouldn't it make more sense that the local mongodb is an opt-in rather than an opt-out option?

I also have a question related to #1115. Having the limitation on the pymongo version was definitely an issue and it is good that it has been resolved. However, I see that the mongomock-ng package is a fork with 2 stars on github and is from a user with relatively little github contributions. Of course I mean no disrespect to the developer, that could be a reliable experienced developer, but I was wondering if you trust the repository and/or if you know the developer.
I am asking because, as mentioned above, the JSONStore is heavily used in jobflow-remote and it cannot rely on a real mongodb.

[rkingsbury]

Great feedback @gpetretto, thank you for testing!

Unless for an existing passwordless local mongodb, the basic configuration would not allow to connect. I wonder if it is a common occurrence for a user to have such a configuration. Wouldn't it make more sense that the local mongodb is an opt-in rather than an opt-out option?

You're absolutely right. I've modified the PR so that this is now an opt-in featured. If the new kwarg port=None (default), you get the existing / legacy behavior. Only if port is set will maggma try to use a real DB for the Store.

Maybe it would be good to also check if the DB is accessible with read/write permission before setting _using_real_mongo?

when the failure happens in inserting a new document the error is difficult to interpret because the methods in the base class MongoStore tries to handle it, but uses attributes that are not defined in the subclass (i.e. safe_update)

Again thank you for the detailed explanation of the errors. Maybe we (read: Claude :) can add some warnings or guard statements to verify the DB is read/writable and make this failure clearer.

I also have a question related to #1115. Having the limitation on the pymongo version was definitely an issue and it is good that it has been resolved. However, I see that the mongomock-ng package is a fork with 2 stars on github and is from a user with relatively little github contributions. Of course I mean no disrespect to the developer, that could be a reliable experienced developer, but I was wondering if you trust the repository and/or if you know the developer.
I am asking because, as mentioned above, the JSONStore is heavily used in jobflow-remote and it cannot rely on a real mongodb.

Yes I appreciate your thinking about this. I do not know the developer; however based on this issue and a few others in the base mongomock repo, my understanding is they have now been invited to take over maintenance of mongomock itself; I think the fork is sort of temporary until longer-term plans for mongomock solidify. While you rightly raise questions about the long-term reliability / robustness, in my opinion it is no worse than the previous state of affairs in which mongomock was more or less unmaintained. And since mongomock-ng is the fork of the future maintainer, it seems that it will be a good predictor of how mongomock itself will evolve.

FWIW, I did ask Claude to research some alternatives to mongomock, but each of those identified had some limitations that made them less desirable than sticking with mongomock.