✨ Add yielding entity contents from kfids

README.md

@@ -28,17 +28,38 @@ from kf_utils.dataservice.delete import delete_entities
 ```
 
 ```Python
-# Yield all entities from the given endpoint matching the given query
+# Yield all entities from the given endpoint matching the given filter query
+for e in yield_entities_from_filter(
+    kf_api_url, "participants", {"study_id": "SD_12345678"}, show_progress=True
+):
+  ...
+```
+
+```Python
+# Yield all entities matching the given kfids
+for e in yield_entities_from_kfids(
+    kf_api_url, ["PT_12345678", "BS_99999999"], show_progress=True
+):
+  ...
+```
+
+```Python
+# Compatibility wrapper for yielding entities from kfid or a filter query
+for e in yield_entities(
+    kf_api_url, "participants", {"study_id": "SD_12345678"}, show_progress=True
+):
+  ...
+
 for e in yield_entities(
-    kf_api_url, "participants", {"study_id": "SD_12345678"}
+    kf_api_url, None, ["PT_12345678", "BS_99999999"], show_progress=True
 ):
   ...
 ```
 
 ```Python
-# Like yield_entities but just yields the kfids
+# Like yield_entities_from_filter but just yields the kfids
 for kfid in yield_kfids(
-    kf_api_url, "participants", {"study_id": "SD_12345678"}
+    kf_api_url, "participants", {"study_id": "SD_12345678"}, show_progress=True
 ):
   ...
 ```

kf_utils/dataservice/meta.py

@@ -26,3 +26,7 @@
 
 def prefix(kfid):
     return kfid.split("_")[0]
+
+
+def get_endpoint(kfid):
+    return prefix_endpoints[prefix(kfid)]

kf_utils/dataservice/scrape.py

@@ -1,7 +1,10 @@
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
 from d3b_utils.requests_retry import Session
+from kf_utils.dataservice.meta import get_endpoint
 
 
-def yield_entities(host, endpoint, filters, show_progress=False):
+def yield_entities_from_filter(host, endpoint, filters, show_progress=False):
     """
     Scrape the dataservice for paginated entities matching the filter params.
 
@@ -52,7 +55,69 @@ def yield_entities(host, endpoint, filters, show_progress=False):
     assert expected == num, f"FOUND {num} ENTITIES BUT EXPECTED {expected}"
 
 
+def yield_entities_from_kfids(host, kfids, show_progress=False):
+    """Fetch the given entities from the dataservice quickly.
+
+    :param host: dataservice base url string (e.g. "http://localhost:5000")
+    :param kfids: kfids to request entities for
+    :raises Exception: if the dataservice doesn't return status 200
+    :yields: entities for the given kfids
+    """
+    host = host.strip("/")
+
+    def do_get(url):
+        response = Session().get(url)
+        if response.status_code != 200:
+            raise Exception(response.text)
+        body = response.json()
+        res = body["results"]
+        res["_links"] = body["_links"]
+        return res
+
+    with ThreadPoolExecutor() as tpex:
+        futures = [
+            tpex.submit(do_get, f"{host}/{get_endpoint(k)}/{k}") for k in kfids
+        ]
+        for f in as_completed(futures):
+            if show_progress:
+                print(".", end="", flush=True)
+            yield f.result()
+
+
+def yield_entities(
+    host, endpoint_if_filter, filters_or_kfids, show_progress=False
+):
+    """Combined call for yield_entities_from_filter and
+    yield_entities_from_kfids to preserve backward compatibility because
+    yield_entities_from_filter was previously called yield_entities.
+
+    :param host: dataservice base url string (e.g. "http://localhost:5000")
+    :param endpoint_if_filter: None or dataservice endpoint string
+        (e.g. "genomic-files")
+    :param filters_or_kfids: dict of filters to winnow results from the
+        dataservice (e.g. {"study_id": "SD_DYPMEHHF", "external_id": "foo"})
+        or a list of kfids
+    :raises Exception: if the dataservice doesn't return status 200
+    :yields: matching entities
+    """
+    if isinstance(filters_or_kfids, str):
+        filters_or_kfids = [filters_or_kfids]
+
+    if isinstance(filters_or_kfids, dict):
+        assert endpoint_if_filter, "If filtering, must specify an endpoint"
+        return yield_entities_from_filter(
+            host,
+            endpoint_if_filter,
+            filters_or_kfids,
+            show_progress=show_progress,
+        )
+    else:
+        return yield_entities_from_kfids(
+            host, filters_or_kfids, show_progress=show_progress
+        )
+
+
 def yield_kfids(host, endpoint, filters, show_progress=False):
-    """Simple wrapper around yield_entities that yields just the KFIDs"""
-    for e in yield_entities(host, endpoint, filters, show_progress):
+    """Wrapper around yield_entities_from_filter that yields just KFIDs"""
+    for e in yield_entities_from_filter(host, endpoint, filters, show_progress):
         yield e["kf_id"]

tests/conftest.py

@@ -1,7 +1,8 @@
 import pytest
 from kf_utils.dataservice.delete import delete_entities
 
-DATASERVICE_URL="http://localhost:5000"
+DATASERVICE_URL = "http://localhost:5000"
+
 
 @pytest.fixture
 def dataservice_setup():
@@ -11,4 +12,3 @@ def dataservice_setup():
     delete_entities(DATASERVICE_URL)
     yield
     delete_entities(DATASERVICE_URL)
-

tests/integration/test_dataservice_utils.py → tests/integration/test_dataservice_delete.py

@@ -17,17 +17,18 @@ def test_delete_entities(dataservice_setup):
             for i in range(n_studies)
         ],
         "participants": [
-            {"kf_id": f"PT_{i}{j}111111", "gender": "Female",
-             "study_id": f"SD_{i}1111111"}
+            {
+                "kf_id": f"PT_{i}{j}111111",
+                "gender": "Female",
+                "study_id": f"SD_{i}1111111",
+            }
             for i in range(n_studies)
             for j in range(n_participants)
-        ]
+        ],
     }
     for endpoint, payloads in data.items():
         for p in payloads:
-            resp = requests.post(
-                f"{DATASERVICE_URL}/{endpoint}", json=p
-            )
+            resp = requests.post(f"{DATASERVICE_URL}/{endpoint}", json=p)
 
     # Delete first two studies
     sids = [s["kf_id"] for s in data["studies"][0:2]]

tests/integration/test_dataservice_scrape.py

@@ -0,0 +1,95 @@
+import requests
+from kf_utils.dataservice.scrape import (
+    yield_entities,
+    yield_entities_from_filter,
+    yield_entities_from_kfids,
+)
+from tests.conftest import DATASERVICE_URL
+
+
+def create_sequencing_center():
+    requests.post(
+        f"{DATASERVICE_URL}/sequencing-centers",
+        json={"kf_id": "SC_11111111", "external_id": "x", "name": "x"},
+    )
+
+
+def create_study(si):
+    requests.post(
+        f"{DATASERVICE_URL}/studies",
+        json={"kf_id": f"SD_{si}1111111", "external_id": f"{si}"},
+    )
+
+
+def create_participant(si, pi):
+    requests.post(
+        f"{DATASERVICE_URL}/participants",
+        json={"kf_id": f"PT_{si}{pi}111111", "study_id": f"SD_{si}1111111", "external_id": f"{pi}"},
+    )
+
+
+def create_biospecimen(si, pi, bi):
+    requests.post(
+        f"{DATASERVICE_URL}/biospecimens",
+        json={
+            "kf_id": f"BS_{si}{pi}{bi}11111",
+            "participant_id": f"PT_{si}{pi}111111",
+            "external_sample_id": f"{pi}{bi}",
+            "external_aliquot_id": f"{pi}{bi}",
+            "sequencing_center_id": "SC_11111111",
+            "analyte_type": "DNA",
+        },
+    )
+
+
+def populate_data(n):
+    # Create some data in dataservice
+    create_sequencing_center()
+    for si in range(n):
+        create_study(si)
+        for pi in range(n):
+            create_participant(si, pi)
+            for bi in range(n):
+                create_biospecimen(si, pi, bi)
+
+
+def test_yield_entities_from_filter(dataservice_setup):
+    n = 4
+    populate_data(n)
+
+    si = 1
+    filter = {"study_id": f"SD_{si}1111111"}
+
+    # Get all participants from one study
+    endpoint = "participants"
+    for ps in [
+        list(yield_entities_from_filter(DATASERVICE_URL, endpoint, filter)),
+        list(yield_entities(DATASERVICE_URL, endpoint, filter)),
+    ]:
+        assert len(ps) == n
+        for p in ps:
+            assert p["kf_id"].startswith(f"PT_{si}")
+
+    # Get all biospecimens from one study
+    endpoint = "biospecimens"
+    for bs in [
+        list(yield_entities_from_filter(DATASERVICE_URL, endpoint, filter)),
+        list(yield_entities(DATASERVICE_URL, endpoint, filter)),
+    ]:
+        assert len(bs) == (n * n)
+        for b in bs:
+            assert b["kf_id"].startswith(f"BS_{si}")
+
+
+def test_yield_entities_from_kfids(dataservice_setup):
+    n = 4
+    populate_data(n)
+
+    kfid_set = {"SD_11111111", "PT_11111111", "BS_11111111"}
+    for es in [
+        list(yield_entities_from_kfids(DATASERVICE_URL, kfid_set)),
+        list(yield_entities(DATASERVICE_URL, None, kfid_set)),
+    ]:
+        assert len(es) == len(kfid_set)
+        found_kfids = {e["kf_id"] for e in es}
+        assert kfid_set == found_kfids

tests/test_dataservice_delete.py

@@ -25,27 +25,21 @@
         ("http://prd.dataservice.org", False, None, {"foo": "bar"}),
         ("http://10.10.1.191", True, MagicMock(), {"foo": "bar"}),
         ("http://10.10.1.191", False, None, {"foo": "bar"}),
-    ]
+    ],
 )
 def test_safe_delete(mocker, url, safety_check, mysession, kwargs):
     """
     Test kf_utils.dataservice.delete.safe_delete
     """
     # Setup mocks
-    mock_session = mocker.patch(
-        "kf_utils.dataservice.delete.Session"
-    )()
+    mock_session = mocker.patch("kf_utils.dataservice.delete.Session")()
 
     if safety_check and ("localhost" not in url):
         with pytest.raises(Exception) as e:
-            safe_delete(
-                url, safety_check=safety_check, session=session
-            )
+            safe_delete(url, safety_check=safety_check, session=session)
             assert "Safe delete is ENABLED" in str(e)
     else:
-        safe_delete(
-            url, safety_check=safety_check, session=mysession, **kwargs
-        )
+        safe_delete(url, safety_check=safety_check, session=mysession, **kwargs)
         if mysession:
             mock_session = mysession
         mock_session.delete.assert_called_with(url, **kwargs)
@@ -56,9 +50,7 @@ def test_delete_kfids(mocker):
     Test kf_utils.dataservice.delete.delete_kfids
     """
     # Setup mocks
-    mock_session = mocker.patch(
-        "kf_utils.dataservice.delete.Session"
-    )()
+    mock_session = mocker.patch("kf_utils.dataservice.delete.Session")()
     mock_resp = MagicMock()
     mock_session.delete.return_value = mock_resp
     kfids = [f"kfid{i}" for i in range(2)]
@@ -84,12 +76,8 @@ def test_delete_entities(mocker):
     """
     Test kf_utils.dataservice.delete.delete_entities
     """
-    mock_yield_kfids = mocker.patch(
-        "kf_utils.dataservice.delete.yield_kfids"
-    )
-    mock_delete_kfids = mocker.patch(
-        "kf_utils.dataservice.delete.delete_kfids"
-    )
+    mock_yield_kfids = mocker.patch("kf_utils.dataservice.delete.yield_kfids")
+    mock_delete_kfids = mocker.patch("kf_utils.dataservice.delete.delete_kfids")
     kfids = [f"kfid{i}" for i in range(2)]
     study_ids = [f"study{i}" for i in range(2)]
     mock_yield_kfids.return_value = kfids
@@ -99,14 +87,16 @@ def test_delete_entities(mocker):
     errors = delete_entities(HOST, study_ids=study_ids)
     assert not errors
     mock_delete_kfids.assert_has_calls(
-        [call(HOST, STUDIES, [sid], safety_check=True)
-         for sid in study_ids],
-        any_order=True
+        [call(HOST, STUDIES, [sid], safety_check=True) for sid in study_ids],
+        any_order=True,
     )
     mock_yield_kfids.assert_has_calls(
-        [call(HOST, e, {"study_id": sid})
-         for e in ENDPOINTS for sid in study_ids],
-        any_order=True
+        [
+            call(HOST, e, {"study_id": sid})
+            for e in ENDPOINTS
+            for sid in study_ids
+        ],
+        any_order=True,
     )
     mock_yield_kfids.reset_mock()
     mock_delete_kfids.reset_mock()
@@ -116,10 +106,13 @@ def test_delete_entities(mocker):
     errors = delete_entities(HOST, study_ids=None)
     assert not errors
     mock_yield_kfids.assert_has_calls(
-        [call(HOST, STUDIES, {})] +
-        [call(HOST, e, {"study_id": sid})
-         for e in ENDPOINTS for sid in study_ids],
-        any_order=True
+        [call(HOST, STUDIES, {})]
+        + [
+            call(HOST, e, {"study_id": sid})
+            for e in ENDPOINTS
+            for sid in study_ids
+        ],
+        any_order=True,
     )
     mock_yield_kfids.reset_mock()
     mock_delete_kfids.reset_mock()