Merge pull request #258 from jmmshn/master

added context manager for stores

requirements-docs.txt

@@ -1,5 +1,5 @@
 mkdocs==1.1.2
-mkdocs-material==5.5.7
+mkdocs-material==5.5.8
 mkdocs-minify-plugin==0.3.0
 mkdocstrings==0.13.0
 pymdown-extensions==8.0

src/maggma/core/store.py

@@ -346,6 +346,13 @@ def __setstate__(self, d):
         d = MontyDecoder().process_decoded(d)
         self.__init__(**d)
 
+    def __enter__(self):
+        self.connect()
+        return self
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        self.close()
+
 
 class StoreError(Exception):
     """ General Store-related error """

src/maggma/stores/aws.py

@@ -7,6 +7,7 @@
 import zlib
 from concurrent.futures import wait
 from concurrent.futures.thread import ThreadPoolExecutor
+import warnings
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import msgpack  # type: ignore
@@ -39,6 +40,7 @@ def __init__(
         endpoint_url: str = None,
         sub_dir: str = None,
         s3_workers: int = 1,
+        key: str = "task_id",
         **kwargs,
     ):
         """
@@ -66,10 +68,18 @@ def __init__(
         self.s3_bucket = None  # type: Any
         self.s3_workers = s3_workers
         # Force the key to be the same as the index
+        assert isinstance(
+            index.key, str
+        ), "Since we are using the key as a file name in S3, they key must be a string"
+        if key != index.key:
+            warnings.warn(
+                f'The desired S3Store key "{key}" does not match the index key "{index.key},"'
+                "the index key will be used",
+                UserWarning,
+            )
         kwargs["key"] = str(index.key)
 
         self._thread_local = threading.local()
-
         super(S3Store, self).__init__(**kwargs)
 
     def name(self) -> str:
@@ -279,7 +289,7 @@ def update(self, docs: Union[List[Dict], Dict], key: Union[List, str, None] = No
             for sdoc in fs:
                 search_docs.append(sdoc.result())
         # Use store's update to remove key clashes
-        self.index.update(search_docs)
+        self.index.update(search_docs, key=self.key)
 
     def get_bucket(self):
         if threading.current_thread().name == "MainThread":

tests/stores/test_aws.py

@@ -63,13 +63,30 @@ def s3store_multi():
         conn = boto3.client("s3")
         conn.create_bucket(Bucket="bucket1")
 
-        index = MemoryStore("index'")
+        index = MemoryStore("index")
         store = S3Store(index, "bucket1", s3_workers=4)
         store.connect()
 
         yield store
 
 
+def test_keys():
+    with mock_s3():
+        conn = boto3.client("s3")
+        conn.create_bucket(Bucket="bucket1")
+        index = MemoryStore("index", key=1)
+        with pytest.raises(AssertionError, match=r"Since we are.*"):
+            store = S3Store(index, "bucket1", s3_workers=4, key=1)
+        index = MemoryStore("index", key="key1")
+        with pytest.warns(UserWarning, match=r"The desired S3Store.*$"):
+            store = S3Store(index, "bucket1", s3_workers=4, key="key2")
+        store.connect()
+        store.update({"key1": "mp-1", "data": "1234"})
+        with pytest.raises(KeyError):
+            store.update({"key2": "mp-2", "data": "1234"})
+        assert store.key == store.index.key == "key1"
+
+
 def test_multi_update(s3store, s3store_multi):
     data = [{"task_id": str(j), "data": "DATA"} for j in range(32)]
 
@@ -137,8 +154,7 @@ def tests_msonable_read_write(s3store):
 def test_remove(s3store):
     s3store.update([{"task_id": "mp-2", "data": "asd"}])
     s3store.update([{"task_id": "mp-4", "data": "asd"}])
-    s3store.update([{"task_id": "mp-5", "data": "aaa"}])
-
+    s3store.update({"task_id": "mp-5", "data": "aaa"})
     assert s3store.query_one({"task_id": "mp-2"}) is not None
     assert s3store.query_one({"task_id": "mp-4"}) is not None