import gzip import pickle from datetime import datetime, timedelta from io import BytesIO, StringIO import boto3 import jsonlines import pandas as pd import pytest from bson import json_util from moto import mock_aws from maggma.stores.open_data import OpenDataStore, PandasMemoryStore, S3IndexStore, TasksOpenDataStore pd.set_option("future.no_silent_downcasting", True) # PandasMemoryStore tests @pytest.fixture() def memstore(): store = PandasMemoryStore(key="task_id") store.update( pd.DataFrame( [ { store.key: "mp-1", store.last_updated_field: datetime.utcnow(), "data": "asd", "int_val": 1, }, { store.key: "mp-3", store.last_updated_field: datetime.utcnow(), "data": "sdf", "int_val": 3, }, ] ) ) return store def test_pdmems_pickle(memstore): sobj = pickle.dumps(memstore) dobj = pickle.loads(sobj) assert hash(dobj) == hash(memstore) assert dobj == memstore def test_pdmems_query(memstore): # bad criteria with pytest.raises(AttributeError, match=r".*only support query or is_in"): memstore.query(criteria={"boo": "hoo"}) with pytest.raises(AttributeError, match=r".*please just use one or the other"): memstore.query( criteria={ "query": f"{memstore.key} == 'mp-1'", "is_in": ("data", ["sdf", "fdr"]), "boo": "hoo", } ) # all pd = memstore.query() assert len(pd) == 2 # query criteria pd = memstore.query(criteria={"query": f"{memstore.key} == 'mp-1'"}) assert len(pd) == 1 assert pd[memstore.key].iloc[0] == "mp-1" assert pd["data"].iloc[0] == "asd" # is_in criteria pd = memstore.query(criteria={"is_in": ("data", ["sdf", "fdr"]), "boo": "hoo"}) assert len(pd) == 1 assert pd[memstore.key].iloc[0] == "mp-3" assert pd["data"].iloc[0] == "sdf" # properties pd = memstore.query(properties=[memstore.key, memstore.last_updated_field]) assert len(pd) == 2 assert len(pd[memstore.key]) == 2 with pytest.raises(KeyError): assert pd["data"] with pytest.raises(KeyError): memstore.query(properties=["fake"]) # sort pd = memstore.query(sort={memstore.key: -1}) assert len(pd) == 2 assert pd[memstore.key].iloc[0] == "mp-3" # skip pd = memstore.query(sort={memstore.key: -1}, skip=1) assert len(pd) == 1 assert pd[memstore.key].iloc[0] == "mp-1" # limit pd = memstore.query(sort={memstore.key: -1}, limit=1) assert len(pd) == 1 assert pd[memstore.key].iloc[0] == "mp-3" # all pd = memstore.query( criteria={ "is_in": ("data", ["sdf", "fdr", "asd"]), "boo": "hoo", }, properties=[memstore.key], sort={"data": -1}, skip=1, limit=1, ) assert len(pd) == 1 assert pd[memstore.key].iloc[0] == "mp-1" with pytest.raises(KeyError): assert pd["data"] def test_pdmems_count(memstore): assert memstore.count() == 2 assert memstore.count(criteria={"query": f"{memstore.key} == 'mp-1'"}) == 1 assert memstore.count(criteria={"is_in": ("data", ["sdf", "fdr"]), "boo": "hoo"}) == 1 assert PandasMemoryStore(key="task_id").count() == 0 @pytest.fixture() def memstore2(): store = PandasMemoryStore(key="task_id") store.update( pd.DataFrame( [ { store.key: "mp-1", store.last_updated_field: datetime.utcnow() - timedelta(hours=1), "data": "asd", }, { store.key: "mp-2", store.last_updated_field: datetime.utcnow() + timedelta(hours=1), "data": "asd", }, { store.key: "mp-3", store.last_updated_field: datetime.utcnow() + timedelta(hours=1), "data": "sdf", }, ] ) ) return store def test_pdmems_distinct(memstore2): assert len(memstore2.distinct(field=memstore2.key)) == len(memstore2._data) assert len(memstore2.distinct(field="data")) == 2 assert len(memstore2.distinct(field=memstore2.key, criteria={"query": "data == 'asd'"})) == 2 def test_pdmems_last_updated(memstore2): assert memstore2._data[memstore2.last_updated_field].iloc[2] == memstore2.last_updated store = PandasMemoryStore(key="task_id") assert store.last_updated == datetime.min def test_pdmems_newer_in(memstore, memstore2): s = memstore.newer_in(memstore2) assert len(s) == 2 assert "mp-2" in s.unique() assert "mp-3" in s.unique() s = memstore.newer_in(target=memstore2, criteria={"query": "data == 'asd'"}, exhaustive=True) assert len(s) == 1 assert "mp-2" in s.unique() with pytest.raises(AttributeError): memstore.newer_in(target=memstore2, criteria={"query": "data == 'asd'"}, exhaustive=False) def test_pdmems_update(memstore): df = pd.DataFrame( [ { memstore.key: "mp-1", memstore.last_updated_field: datetime.utcnow(), "data": "boo", "int_val": 1, } ] ) df2 = memstore.update(df) assert len(memstore._data) == 2 assert memstore.query(criteria={"query": f"{memstore.key} == 'mp-1'"})["data"].iloc[0] == "boo" assert memstore.query(criteria={"query": f"{memstore.key} == 'mp-1'"})["int_val"].iloc[0] == 1 assert df2.equals(df) df = pd.DataFrame( [ { memstore.key: "mp-2", memstore.last_updated_field: datetime.utcnow(), "data": "boo", "int_val": 2, } ] ) df2 = memstore.update(df) assert len(memstore._data) == 3 assert memstore.query(criteria={"query": f"{memstore.key} == 'mp-2'"})["data"].iloc[0] == "boo" assert memstore.query(criteria={"query": f"{memstore.key} == 'mp-2'"})["int_val"].iloc[0] == 2 assert df2.equals(df) @pytest.fixture() def s3indexstore(): data = [{"task_id": "mp-1", "last_updated": datetime.utcnow()}] with mock_aws(): conn = boto3.resource("s3", region_name="us-east-1") conn.create_bucket(Bucket="bucket1") client = boto3.client("s3", region_name="us-east-1") string_io = StringIO() with jsonlines.Writer(string_io, dumps=json_util.dumps) as writer: for _, row in pd.DataFrame(data).iterrows(): writer.write(row.to_dict()) client.put_object( Bucket="bucket1", Body=BytesIO(string_io.getvalue().encode("utf-8")), Key="manifest.jsonl", ) store = S3IndexStore(collection_name="index", bucket="bucket1", key="task_id") store.connect() yield store def test_s3is_pickle(s3indexstore): sobj = pickle.dumps(s3indexstore) dobj = pickle.loads(sobj) assert hash(dobj) == hash(s3indexstore) assert dobj == s3indexstore def test_s3is_connect_retrieve_manifest(s3indexstore): assert s3indexstore.retrieve_manifest().equals(s3indexstore._data) with mock_aws(): with pytest.raises(s3indexstore.s3_client.exceptions.NoSuchBucket): S3IndexStore(collection_name="foo", bucket="bucket2", key="task_id").retrieve_manifest() with pytest.raises(RuntimeError): S3IndexStore(collection_name="foo", bucket="bucket2", key="task_id").connect() conn = boto3.resource("s3", region_name="us-east-1") conn.create_bucket(Bucket="bucket2") s3is = S3IndexStore(collection_name="foo", bucket="bucket2", key="task_id") s3is.connect() assert s3is._data is None assert s3is.retrieve_manifest() is None assert s3is.count() == 0 def test_s3is_store_manifest(): with mock_aws(): conn = boto3.resource("s3", region_name="us-east-1") conn.create_bucket(Bucket="bucket2") s3is = S3IndexStore(collection_name="foo", bucket="bucket2", key="task_id") s3is.connect() s3is.update(pd.DataFrame([{"task_id": "mp-2", "last_updated": "now"}])) s3is.store_manifest() df = s3is.retrieve_manifest() assert len(df) == 1 assert df.equals(s3is._data) s3is.update(pd.DataFrame([{"task_id": "mp-3", "last_updated": "later"}])) df = s3is.retrieve_manifest() assert not df.equals(s3is._data) def test_s3is_close(s3indexstore): s3indexstore.close() assert len(s3indexstore.query()) == 1 # actions auto-reconnect s3indexstore.update(pd.DataFrame([{"task_id": "mp-2", "last_updated": "now"}])) assert len(s3indexstore.query()) == 2 s3indexstore.close() assert len(s3indexstore.query()) == 2 # actions auto-reconnect s3indexstore.close() s3indexstore.connect() assert len(s3indexstore.query()) == 1 # explicit connect reloads manifest @pytest.fixture() def s3store(): with mock_aws(): conn = boto3.resource("s3", region_name="us-east-1") conn.create_bucket(Bucket="bucket1") store = OpenDataStore( collection_name="index", bucket="bucket1", key="task_id", object_grouping=["group_level_two", "task_id"] ) store.connect() store.update( pd.DataFrame( [ { "task_id": "mp-1", "data": "asd", store.last_updated_field: datetime.utcnow(), "group": {"level_two": 4}, }, { "task_id": "mp-3", "data": "sdf", store.last_updated_field: datetime.utcnow(), "group": {"level_two": 4}, }, ] ) ) yield store @pytest.fixture() def s3store_w_subdir(): with mock_aws(): conn = boto3.resource("s3", region_name="us-east-1") conn.create_bucket(Bucket="bucket1") conn.create_bucket(Bucket="bucket2") index = S3IndexStore(collection_name="index", bucket="bucket1", key="task_id") store = OpenDataStore( index=index, collection_name=index.collection_name, bucket="bucket2", key=index.key, prefix="subdir1", object_grouping=["data_foo"], ) store.connect() store.update( pd.DataFrame( [ { "task_id": "mp-1", "data": {"foo": "asd"}, store.last_updated_field: datetime.utcnow(), }, { "task_id": "mp-3", "data": {"foo": "sdf"}, store.last_updated_field: datetime.utcnow(), }, ] ) ) yield store def test_pickle(s3store): sobj = pickle.dumps(s3store) dobj = pickle.loads(sobj) assert hash(dobj) == hash(s3store) assert dobj == s3store assert len(dobj.query(criteria={"query": "task_id == 'mp-2'"})) == 0 assert dobj.query(criteria={"query": "task_id == 'mp-1'"})["data"].iloc[0] == "asd" def test_index_property(s3store, s3store_w_subdir): assert s3store.index != s3store assert s3store_w_subdir.index != s3store_w_subdir def test_read_doc_from_s3(): data = [ {"task_id": "mp-1", "last_updated": datetime.utcnow(), "group": "foo", "also_group": "boo"}, {"task_id": "mp-2", "last_updated": datetime.utcnow(), "group": "foo", "also_group": "boo"}, ] with mock_aws(): conn = boto3.resource("s3", region_name="us-east-1") conn.create_bucket(Bucket="bucket1") client = boto3.client("s3", region_name="us-east-1") string_io = StringIO() with jsonlines.Writer(string_io, dumps=json_util.dumps) as writer: for _, row in pd.DataFrame(data).iterrows(): writer.write(row.to_dict()) client.put_object( Bucket="bucket1", Body=BytesIO(gzip.compress(string_io.getvalue().encode("utf-8"))), Key="group=foo/also_group=boo.jsonl.gz", ) store = OpenDataStore( collection_name="index", bucket="bucket1", key="task_id", object_grouping=["group", "also_group"] ) store.connect() df = store._read_doc_from_s3(file_id="group=foo/also_group=boo.jsonl.gz") assert len(df) == 2 assert store._get_full_key_path(index=df) == "group=foo/also_group=boo.jsonl.gz" assert (df["task_id"] == "mp-1").any() assert (df["task_id"] == "mp-2").any() @pytest.mark.xfail( reason="Known issue, the store is in a deprecated state, and in particular may be incompatible with numpy 2.0+" ) def test_update(s3store): assert len(s3store.index_data) == 2 s3store.update( pd.DataFrame( [ { "task_id": "mp-199999", "data": "asd", "group": {"level_two": 4}, s3store.last_updated_field: datetime.utcnow(), } ] ) ) assert len(s3store.index_data) == 3 with pytest.raises(KeyError): assert s3store.index_data.query("task_id == 'mp-199999'")["data"].iloc[0] == "asd" s3store.update( pd.DataFrame( [ { "task_id": "mp-199999", "data": "foo", "group": {"level_two": 4}, s3store.last_updated_field: datetime.utcnow(), } ] ) ) assert len(s3store.index_data) == 3 assert len(s3store.index_data.query("task_id == 'mp-199999'")) == 1 mp4 = [{"task_id": "mp-4", "data": "asd", "group": {"level_two": 4}, s3store.last_updated_field: datetime.utcnow()}] s3store.update(pd.DataFrame(mp4)) assert len(s3store.index_data) == 4 mp4_index = [{"task_id": "mp-4", "group_level_two": 4, s3store.last_updated_field: datetime.utcnow()}] assert s3store._get_full_key_path(pd.DataFrame(mp4_index)) == "group_level_two=4/task_id=mp-4.jsonl.gz" s3store.s3_client.head_object(Bucket=s3store.bucket, Key=s3store._get_full_key_path(pd.DataFrame(mp4_index))) def test_query(s3store): assert len(s3store.query(criteria={"query": "task_id == 'mp-2'"})) == 0 assert s3store.query(criteria={"query": "task_id == 'mp-1'"})["data"].iloc[0] == "asd" assert s3store.query(criteria={"query": "task_id == 'mp-3'"})["data"].iloc[0] == "sdf" assert s3store.query(criteria={"query": "task_id == 'mp-1'"}, properties=["task_id"])["task_id"].iloc[0] == "mp-1" assert ( s3store.query(criteria={"query": "task_id == 'mp-1'"}, properties=["task_id", "data"])["data"].iloc[0] == "asd" ) assert len(s3store.query()) == 2 # will use optimized search df = s3store.query(criteria={"query": "task_id == 'mp-1'"}, properties=["task_id"], criteria_fields=["task_id"]) assert len(df) == 1 assert df["task_id"].iloc[0] == "mp-1" with pytest.raises(KeyError): assert df["data"].iloc[0] == "asd" # will use optimized search df = s3store.query(properties=["task_id"], criteria_fields=[]) assert len(df) == 2 # will not use optimized search even with hints since data is not in the searchable_fields df = s3store.query(criteria={"query": "data == 'asd'"}, properties=["task_id"], criteria_fields=["data"]) assert len(df) == 1 assert df["task_id"].iloc[0] == "mp-1" with pytest.raises(KeyError): assert df["data"].iloc[0] == "asd" def test_rebuild_index_from_s3_data(s3store): data = [ {"task_id": "mp-2", "data": "asd", s3store.last_updated_field: datetime.utcnow(), "group": {"level_two": 4}} ] client = boto3.client("s3", region_name="us-east-1") string_io = StringIO() with jsonlines.Writer(string_io, dumps=json_util.dumps) as writer: for _, row in pd.DataFrame(data).iterrows(): writer.write(row.to_dict()) data = [ {"task_id": "mp-99", "data": "asd", s3store.last_updated_field: datetime.utcnow(), "group": {"level_two": 4}} ] string_io2 = StringIO() with jsonlines.Writer(string_io2, dumps=json_util.dumps) as writer: for _, row in pd.DataFrame(data).iterrows(): writer.write(row.to_dict()) client.put_object( Bucket="bucket1", Body=BytesIO(gzip.compress(string_io.getvalue().encode("utf-8"))), Key="group_level_two=4/task_id=mp-2.jsonl.gz", ) # creating file that should not be indexed to test that it gets skipped client.put_object( Bucket="bucket1", Body=BytesIO(gzip.compress(string_io2.getvalue().encode("utf-8"))), Key="task_id=mp-99.gz", ) assert len(s3store.index.index_data) == 2 index_docs = s3store.rebuild_index_from_s3_data() assert len(index_docs) == 3 assert len(s3store.index.index_data) == 3 for key in index_docs.columns: assert key == "task_id" or key == "last_updated" or key == "group_level_two" def test_rebuild_index_from_data(s3store): data = [ {"task_id": "mp-2", "data": "asd", s3store.last_updated_field: datetime.utcnow(), "group": {"level_two": 4}} ] index_docs = s3store.rebuild_index_from_data(pd.DataFrame(data)) assert len(index_docs) == 1 assert len(s3store.index.index_data) == 1 for key in index_docs.columns: assert key == "task_id" or key == "last_updated" or key == "group_level_two" def test_count_subdir(s3store_w_subdir): s3store_w_subdir.update( pd.DataFrame( [{"task_id": "mp-1", "data": {"foo": "asd"}, s3store_w_subdir.last_updated_field: datetime.utcnow()}] ) ) s3store_w_subdir.update( pd.DataFrame( [{"task_id": "mp-2", "data": {"foo": "asd"}, s3store_w_subdir.last_updated_field: datetime.utcnow()}] ) ) assert len(s3store_w_subdir.query()) == 3 assert len(s3store_w_subdir.query(criteria=None)) == 3 assert s3store_w_subdir.count() == 3 assert s3store_w_subdir.count({"query": "task_id == 'mp-2'"}) == 1 def test_subdir_storage(s3store_w_subdir): def objects_in_bucket(key): objs = s3store_w_subdir.s3_client.list_objects_v2(Bucket=s3store_w_subdir.bucket, Prefix=key) return key in [o["Key"] for o in objs["Contents"]] s3store_w_subdir.update( pd.DataFrame( [{"task_id": "mp-1", "data": {"foo": "asd"}, s3store_w_subdir.last_updated_field: datetime.utcnow()}] ) ) s3store_w_subdir.update( pd.DataFrame( [{"task_id": "mp-2", "data": {"foo": "asd"}, s3store_w_subdir.last_updated_field: datetime.utcnow()}] ) ) assert objects_in_bucket("subdir1/data_foo=asd.jsonl.gz") assert objects_in_bucket("subdir1/data_foo=sdf.jsonl.gz") def test_additional_metadata(s3store): tic = datetime(2018, 4, 12, 16) data = [{"task_id": f"mp-{i}", "a": i, s3store.last_updated_field: tic} for i in range(4)] with pytest.raises(TypeError): s3store.update(data, key="a", additional_metadata="task_id") def test_rebuild_index_from_s3_for_tasks_store(): data = [{"task_id": "mp-2", "data": "asd", "last_updated": datetime.utcnow(), "group": {"level_two": 4}}] with mock_aws(): conn = boto3.resource("s3", region_name="us-east-1") conn.create_bucket(Bucket="bucket1") string_io = StringIO() with jsonlines.Writer(string_io, dumps=json_util.dumps) as writer: for _, row in pd.DataFrame(data).iterrows(): writer.write(row.to_dict()) client = boto3.client("s3", region_name="us-east-1") client.put_object( Bucket="bucket1", Body=BytesIO(gzip.compress(string_io.getvalue().encode("utf-8"))), Key="group_level_two=4/dt=some_random_data.jsonl.gz", ) store = TasksOpenDataStore( collection_name="index", bucket="bucket1", key="task_id", object_grouping=["group_level_two", "dt"] ) store.connect() index_docs = store.rebuild_index_from_s3_data() assert len(index_docs) == 1 assert len(store.index.index_data) == 1 for key in index_docs.columns: assert key == "task_id" or key == "last_updated" or key == "group_level_two" or key == "dt" assert index_docs["dt"].iloc[0] == "some_random_data" def test_no_update_for_tasks_store(): with mock_aws(): conn = boto3.resource("s3", region_name="us-east-1") conn.create_bucket(Bucket="bucket1") store = TasksOpenDataStore( collection_name="index", bucket="bucket1", key="task_id", object_grouping=["group_level_two", "dt"] ) store.connect() with pytest.raises(NotImplementedError): store.update( pd.DataFrame( [ { "task_id": "mp-199999", "data": "foo", "group": {"level_two": 4}, store.last_updated_field: datetime.utcnow(), } ] ) )