diff --git a/kf_utils/dataservice/descendants.py b/kf_utils/dataservice/descendants.py index dd0cce8..4ce02b4 100644 --- a/kf_utils/dataservice/descendants.py +++ b/kf_utils/dataservice/descendants.py @@ -286,6 +286,13 @@ def find_descendants_by_kfids( table_to_endpoint = {v: k for k, v in endpoint_to_table.items()} parent_type = endpoint_to_table[parent_endpoint] + if use_api: + descendancy = _api_descendancy + else: + descendancy = _db_descendancy + db_conn = psycopg2.connect(api_or_db_url) + db_cur = db_conn.cursor(cursor_factory=psycopg2.extras.DictCursor) + if isinstance(parents, str): parents = [parents] @@ -294,14 +301,19 @@ def find_descendants_by_kfids( descendants = {parent_type: {p["kf_id"]: p for p in parents}} else: parent_kfids = set(parents) - descendants = {parent_type: {k: k for k in parent_kfids}} - - if use_api: - descendancy = _api_descendancy - else: - descendancy = _db_descendancy - db_conn = psycopg2.connect(api_or_db_url) - db_cur = db_conn.cursor(cursor_factory=psycopg2.extras.DictCursor) + if use_api: + descendants = { + parent_type: { + e["kf_id"]: e + for e in yield_entities(api_or_db_url, None, parent_kfids) + } + } + else: + query = f"select distinct * from {parent_type} where kf_id in %s" + db_cur.execute(query, (tuple(parent_kfids | {None}),)) + descendants = { + parent_type: {p["kf_id"]: dict(p) for p in db_cur.fetchall()} + } done = set() for t in descendancy.keys(): @@ -350,7 +362,7 @@ def _inner(parent_type, parent_kfids, descendants): f" on {child_type}.{link_on_child} = {parent_type}.{link_on_parent}" f" where {parent_type}.kf_id in %s" ) - db_cur.execute(query, (tuple(parent_kfids),)) + db_cur.execute(query, (tuple(parent_kfids | {None}),)) children = {c["kf_id"]: dict(c) for c in db_cur.fetchall()} if children: @@ -407,6 +419,7 @@ def find_descendants_by_filter( things = list(yield_entities(api_url, endpoint, filter, show_progress=True)) if kfids_only: things = [t["kf_id"] for t in things] + descendants = find_descendants_by_kfids( db_url or api_url, endpoint, diff --git a/tests/integration/test_dataservice_descendants.py b/tests/integration/test_dataservice_descendants.py new file mode 100644 index 0000000..bbeaacc --- /dev/null +++ b/tests/integration/test_dataservice_descendants.py @@ -0,0 +1,49 @@ +from kf_utils.dataservice.descendants import find_descendants_by_kfids, find_descendants_by_filter +from tests.conftest import DATASERVICE_URL, populate_data + + +def test_by_kfids_simple(dataservice_setup): + n = 4 + populate_data(n) + + desc = find_descendants_by_kfids( + DATASERVICE_URL, "participants", ["PT_11111111"], + ignore_gfs_with_hidden_external_contribs=False, + kfids_only=False + ) + assert len(desc["participants"]) == 1 + assert len(desc["biospecimens"]) == n + # parent level should be populated + assert isinstance(desc["participants"]["PT_11111111"], dict) + + desc = find_descendants_by_kfids( + DATASERVICE_URL, "participants", [ + "PT_10111111", "PT_11111111", "PT_12111111", "PT_13111111" + ], + ignore_gfs_with_hidden_external_contribs=False, + kfids_only=False + ) + assert len(desc["participants"]) == n + assert len(desc["biospecimens"]) == n*n + + desc = find_descendants_by_kfids( + DATASERVICE_URL, "studies", ["SD_11111111"], + ignore_gfs_with_hidden_external_contribs=False, + kfids_only=False + ) + assert len(desc["participants"]) == n + assert len(desc["biospecimens"]) == n*n + + +def test_by_filter_simple(dataservice_setup): + n = 4 + populate_data(n) + desc = find_descendants_by_filter( + DATASERVICE_URL, "participants", {"study_id": "SD_11111111"}, + ignore_gfs_with_hidden_external_contribs=False, + kfids_only=False + ) + assert len(desc["participants"]) == n + assert len(desc["biospecimens"]) == n*n + # parent level should be populated + assert isinstance(desc["participants"]["PT_11111111"], dict)