From d962c62b23b8d57e531d5de6421d8b1edc4dd911 Mon Sep 17 00:00:00 2001 From: Ryan Kingsbury Date: Thu, 20 Jul 2023 17:45:08 -0400 Subject: [PATCH 1/3] MontyStore: fix db name not working --- src/maggma/stores/mongolike.py | 36 ++++++++++++++++++---------------- tests/stores/test_mongolike.py | 23 ++++++++++++++-------- 2 files changed, 34 insertions(+), 25 deletions(-) diff --git a/src/maggma/stores/mongolike.py b/src/maggma/stores/mongolike.py index 0d3a01c1c..1d4960804 100644 --- a/src/maggma/stores/mongolike.py +++ b/src/maggma/stores/mongolike.py @@ -10,7 +10,7 @@ from itertools import chain, groupby from socket import socket import warnings -from typing import Dict, Iterator, List, Optional, Tuple, Union, Any, Callable +from typing import Dict, Iterator, List, Optional, Tuple, Union, Any, Callable, Literal import mongomock import orjson @@ -27,9 +27,9 @@ from maggma.utils import confirm_field_index, to_dt try: - import montydb # type: ignore + from montydb import MontyClient, set_storage # type: ignore except ImportError: - montydb = None + MontyClient = None class SSHTunnel(MSONable): @@ -872,7 +872,7 @@ def __eq__(self, other: object) -> bool: @requires( - montydb, + MontyClient is not None, "MontyStore requires MontyDB to be installed. See the MontyDB repository for more " "information: https://github.com/davidlatwe/montydb", ) @@ -903,7 +903,7 @@ def __init__( collection_name, database_path: str = None, database_name: str = "db", - storage: str = "sqlite", + storage: Literal["sqlite", "flatfile", "lightning"] = "sqlite", storage_kwargs: Optional[dict] = None, client_kwargs: Optional[dict] = None, **kwargs, @@ -913,10 +913,12 @@ def __init__( Args: collection_name: Name for the collection. - database_path: Path to the directory containing the on-disk database files. - If None, the current working directory will be used. + database_path: Path to on-disk database files. If None, the current working + directory will be used. database_name: The database name. - storage: The storage type. Options include "sqlite", "lightning", "flatfile". + storage: The storage type. Options include "sqlite", "lightning", "flatfile". Note that + although MontyDB supports in memory storage, this capability is disabled in maggma to avoid unintended behavior, since multiple + in-memory MontyStore would actually point to the same data. storage_kwargs: Keyword arguments passed to ``montydb.set_storage``. client_kwargs: Keyword arguments passed to the ``montydb.MontyClient`` constructor. @@ -925,7 +927,6 @@ def __init__( if database_path is None: database_path = str(Path.cwd()) - self.database = "MontyDB" self.database_path = database_path self.database_name = database_name self.collection_name = collection_name @@ -935,8 +936,8 @@ def __init__( self.kwargs = kwargs self.storage = storage self.storage_kwargs = storage_kwargs or { - "use_bson": True, - "monty_version": "4.0", + "use_bson": True, # import pymongo's BSON; do not use montydb's + "mongo_version": "4.0", } self.client_kwargs = client_kwargs or {} super(MongoStore, self).__init__(**kwargs) # noqa @@ -948,17 +949,19 @@ def connect(self, force_reset: bool = False): Args: force_reset: Force connection reset. """ - from montydb import set_storage, MontyClient # type: ignore - - set_storage(self.database_path, storage=self.storage, **self.storage_kwargs) + # TODO - workaround, may be obviated by a future montydb update + if self.database_path != ":memory:": + set_storage(self.database_path, storage=self.storage, **self.storage_kwargs) client = MontyClient(self.database_path, **self.client_kwargs) if not self._coll or force_reset: - self._coll = client["db"][self.collection_name] + self._coll = client[self.database_name][self.collection_name] @property def name(self) -> str: """Return a string representing this data source.""" - return f"monty://{self.database_path}/{self.database}/{self.collection_name}" + return ( + f"monty://{self.database_path}/{self.database_name}/{self.collection_name}" + ) def count( self, @@ -973,7 +976,6 @@ def count( hint: Dictionary of indexes to use as hints for query optimizer. Keys are field names and values are 1 for ascending or -1 for descending. """ - criteria = criteria if criteria else {} hint_list = ( diff --git a/tests/stores/test_mongolike.py b/tests/stores/test_mongolike.py index 92fee44b5..5715b76b3 100644 --- a/tests/stores/test_mongolike.py +++ b/tests/stores/test_mongolike.py @@ -12,8 +12,7 @@ from pymongo.errors import ConfigurationError, DocumentTooLarge, OperationFailure from maggma.core import StoreError -from maggma.stores import JSONStore, MemoryStore, MongoStore, MongoURIStore -from maggma.stores.mongolike import MontyStore +from maggma.stores import JSONStore, MemoryStore, MongoStore, MongoURIStore, MontyStore from maggma.validators import JSONSchemaValidator @@ -304,6 +303,14 @@ def test_monty_store_connect(tmp_dir): assert montystore._collection is not None assert montystore.name is not None + # check that the kwargs work + with ScratchDir("."): + store = MontyStore("my_results", database_name="NotNamedDB") + store.connect() + store.update({"test": {"cow": "moo"}}, key="test") + store.close() + assert Path("NotNamedDB/my_results.collection").exists() + def test_monty_store_groupby(montystore): montystore.update( @@ -338,7 +345,7 @@ def test_monty_store_groupby(montystore): assert len(data) == 2 -def test_montystore_query(montystore): +def test_monty_store_query(montystore): montystore._collection.insert_one({"a": 1, "b": 2, "c": 3}) assert montystore.query_one(properties=["a"])["a"] == 1 assert montystore.query_one(properties=["a"])["a"] == 1 @@ -346,7 +353,7 @@ def test_montystore_query(montystore): assert montystore.query_one(properties=["c"])["c"] == 3 -def test_montystore_count(montystore): +def test_monty_store_count(montystore): montystore._collection.insert_one({"a": 1, "b": 2, "c": 3}) assert montystore.count() == 1 montystore._collection.insert_one({"aa": 1, "b": 2, "c": 3}) @@ -354,7 +361,7 @@ def test_montystore_count(montystore): assert montystore.count({"a": 1}) == 1 -def test_montystore_distinct(montystore): +def test_monty_store_distinct(montystore): montystore._collection.insert_one({"a": 1, "b": 2, "c": 3}) montystore._collection.insert_one({"a": 4, "d": 5, "e": 6, "g": {"h": 1}}) assert set(montystore.distinct("a")) == {1, 4} @@ -375,7 +382,7 @@ def test_montystore_distinct(montystore): assert montystore.distinct("i") == [None] -def test_montystore_update(montystore): +def test_monty_store_update(montystore): montystore.update({"e": 6, "d": 4}, key="e") assert ( montystore.query_one(criteria={"d": {"$exists": 1}}, properties=["d"])["d"] == 4 @@ -399,7 +406,7 @@ def test_montystore_update(montystore): montystore.update({"e": "abc", "d": 3}, key="e") -def test_montystore_remove_docs(montystore): +def test_monty_store_remove_docs(montystore): montystore._collection.insert_one({"a": 1, "b": 2, "c": 3}) montystore._collection.insert_one({"a": 4, "d": 5, "e": 6, "g": {"h": 1}}) montystore.remove_docs({"a": 1}) @@ -407,7 +414,7 @@ def test_montystore_remove_docs(montystore): assert len(list(montystore.query({"a": 1}))) == 0 -def test_montystore_last_updated(montystore): +def test_monty_store_last_updated(montystore): assert montystore.last_updated == datetime.min start_time = datetime.utcnow() montystore._collection.insert_one({montystore.key: 1, "a": 1}) From cdb7bd6faaf66eaa721d8b2a4979bfcda297b915 Mon Sep 17 00:00:00 2001 From: Ryan Kingsbury Date: Thu, 20 Jul 2023 18:06:22 -0400 Subject: [PATCH 2/3] python 3.7 compatibility --- src/maggma/stores/mongolike.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/src/maggma/stores/mongolike.py b/src/maggma/stores/mongolike.py index 1d4960804..a55d6b2ab 100644 --- a/src/maggma/stores/mongolike.py +++ b/src/maggma/stores/mongolike.py @@ -10,7 +10,22 @@ from itertools import chain, groupby from socket import socket import warnings -from typing import Dict, Iterator, List, Optional, Tuple, Union, Any, Callable, Literal + +try: + from typing import ( + Dict, + Iterator, + List, + Optional, + Tuple, + Union, + Any, + Callable, + Literal, + ) +except ImportError: + from typing import Dict, Iterator, List, Optional, Tuple, Union, Any, Callable + from typing_extensions import Literal import mongomock import orjson From 41b5a08051b93106e6a66585108098558dedea9d Mon Sep 17 00:00:00 2001 From: Ryan Kingsbury Date: Thu, 20 Jul 2023 18:08:32 -0400 Subject: [PATCH 3/3] python 3.7 compatibility --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index dd9995b0f..04bad4282 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,6 +9,7 @@ aioitertools==0.10.0 pydantic==1.10.2 fastapi==0.79.0 numpy==1.23.0;python_version>"3.8" +typing_extensions;python_version<"3.8" pyzmq==24.0.1 dnspython==2.2.1 uvicorn==0.18.3