nmdc-runtime 2.10.0__py3-none-any.whl → 2.11.1__py3-none-any.whl

nmdc_runtime/api/endpoints/find.py

@@ -1,21 +1,17 @@
-from operator import itemgetter
-from typing import List, Annotated
+import logging
+from typing import Annotated
 
 from fastapi import APIRouter, Depends, Path, Query
-from jinja2 import Environment, PackageLoader, select_autoescape
-from nmdc_runtime.util import get_nmdc_jsonschema_dict
 from pymongo.database import Database as MongoDatabase
-from starlette.responses import HTMLResponse
-from toolz import merge, assoc_in
 
 from nmdc_schema.get_nmdc_view import ViewGetter
 from nmdc_runtime.api.core.util import raise404_if_none
 from nmdc_runtime.api.db.mongo import (
     get_mongo_db,
-    activity_collection_names,
     get_planned_process_collection_names,
     get_nonempty_nmdc_schema_collection_names,
 )
+from nmdc_runtime.api.endpoints.nmdcschema import get_linked_instances
 from nmdc_runtime.api.endpoints.util import (
     find_resources,
     strip_oid,
@@ -25,9 +21,8 @@ from nmdc_runtime.api.models.metadata import Doc
 from nmdc_runtime.api.models.util import (
     FindResponse,
     FindRequest,
-    entity_attributes_to_index,
 )
-from nmdc_runtime.util import get_class_names_from_collection_spec
+
 
 router = APIRouter()
 
@@ -178,133 +173,71 @@ def find_data_objects_for_study(
         is a list of the `DataObject`s associated with that `Biosample`.
     """
     biosample_data_objects = []
-    study = raise404_if_none(
-        mdb.study_set.find_one({"id": study_id}, ["id"]), detail="Study not found"
-    )
-
-    # Note: With nmdc-schema v10 (legacy schema), we used the field named `part_of` here.
-    #       With nmdc-schema v11 (Berkeley schema), we use the field named `associated_studies` here.
-    biosamples = mdb.biosample_set.find({"associated_studies": study["id"]}, ["id"])
-    biosample_ids = [biosample["id"] for biosample in biosamples]
-
-    # SchemaView interface to NMDC Schema
-    nmdc_view = ViewGetter()
-    nmdc_sv = nmdc_view.get_view()
-    dg_descendants = [
-        (f"nmdc:{t}" if ":" not in t else t)
-        for t in nmdc_sv.class_descendants("DataGeneration")
-    ]
-
-    def collect_data_objects(doc_ids, collected_objects, unique_ids):
-        """Helper function to collect data objects from `has_input` and `has_output` references."""
-        for doc_id in doc_ids:
-            # Check if this is a DataObject by looking at the document's type directly
-            doc = mdb.alldocs.find_one({"id": doc_id}, {"type": 1})
-            if (
-                doc
-                and doc.get("type") == "nmdc:DataObject"
-                and doc_id not in unique_ids
-            ):
-                data_obj = mdb.data_object_set.find_one({"id": doc_id})
-                if data_obj:
-                    collected_objects.append(strip_oid(data_obj))
-                    unique_ids.add(doc_id)
-
-    # Another way in which DataObjects can be related to Biosamples is through the
-    # `was_informed_by` key/slot. We need to link records from the `workflow_execution_set`
-    # collection that are "informed" by the same DataGeneration records that created
-    # the outputs above. Then we need to get additional DataObject records that are
-    # created by this linkage.
-    def process_informed_by_docs(doc, collected_objects, unique_ids):
-        """Process documents linked by `was_informed_by` and collect relevant data objects."""
-        # Note: As of nmdc-schema 11.9.0, the `was_informed_by` field, if defined,
-        #       will contain a list of strings. In MongoDB, the `{k: v}` filter
-        #       can be used to check whether either (a) the value of field `f` is
-        #       an array containing `v` as one of its elements, or (b) the value
-        #       of field `f` is exactly equal to `v`. We rely on behavior (a) here.
-        informed_by_docs = mdb.workflow_execution_set.find(
-            {"was_informed_by": doc["id"]}
-        )
-        for informed_doc in informed_by_docs:
-            collect_data_objects(
-                informed_doc.get("has_input", []), collected_objects, unique_ids
-            )
-            collect_data_objects(
-                informed_doc.get("has_output", []), collected_objects, unique_ids
-            )
 
-    biosample_data_objects = []
+    # Respond with an error if the specified `Study` does not exist.
+    # Note: We project only the `_id` field, to minimize data transfer.
+    raise404_if_none(
+        mdb["study_set"].find_one({"id": study_id}, projection={"_id": 1}),
+        detail="Study not found",
+    )
 
-    for biosample_id in biosample_ids:
-        current_ids = [biosample_id]
-        collected_data_objects = []
-        unique_ids = set()
-
-        # Iterate over records in the `alldocs` collection. Look for
-        # records that have the given biosample_id as value on the
-        # `has_input` key/slot. The retrieved documents might also have a
-        # `has_output` key/slot associated with them. Get the value of the
-        # `has_output` key and check if it's type is `nmdc:DataObject`. If
-        # it's not, repeat the process till it is.
-        while current_ids:
-            new_current_ids = []
-            for current_id in current_ids:
-                # Query to find all documents with current_id as the value on
-                # `has_input` slot
-                for doc in mdb.alldocs.find({"has_input": current_id}):
-                    has_output = doc.get("has_output", [])
-
-                    # Process `DataGeneration` type documents linked by `was_informed_by`
-                    if not has_output and any(
-                        t in dg_descendants for t in doc.get("_type_and_ancestors", [])
-                    ):
-                        process_informed_by_docs(
-                            doc, collected_data_objects, unique_ids
-                        )
-                        continue
-
-                    collect_data_objects(has_output, collected_data_objects, unique_ids)
-                    # Add non-DataObject outputs to continue the chain
-                    for op in has_output:
-                        doc_check = mdb.alldocs.find_one({"id": op}, {"type": 1})
-                        if doc_check and doc_check.get("type") != "nmdc:DataObject":
-                            new_current_ids.append(op)
-
-                    if any(
-                        t in dg_descendants for t in doc.get("_type_and_ancestors", [])
-                    ):
-                        process_informed_by_docs(
-                            doc, collected_data_objects, unique_ids
-                        )
-
-                # Also check if current_id is a DataObject that serves as input to other processes
-                current_doc_type = mdb.alldocs.find_one({"id": current_id}, {"type": 1})
-                if (
-                    current_doc_type
-                    and current_doc_type.get("type") == "nmdc:DataObject"
-                ):
-                    # Find all documents in alldocs that have this DataObject as input
-                    for doc in mdb.alldocs.find({"has_input": current_id}):
-                        has_output = doc.get("has_output", [])
-                        # Process outputs from these documents
-                        collect_data_objects(
-                            has_output, collected_data_objects, unique_ids
-                        )
-                        # Add non-DataObject outputs to continue the chain
-                        for op in has_output:
-                            doc_check = mdb.alldocs.find_one({"id": op}, {"type": 1})
-                            if doc_check and doc_check.get("type") != "nmdc:DataObject":
-                                new_current_ids.append(op)
-
-            current_ids = new_current_ids
-
-        if collected_data_objects:
-            result = {
+    # Use the `get_linked_instances` function—which is the function that
+    # underlies the `/nmdcschema/linked_instances` API endpoint—to get all
+    # the `Biosample`s that are downstream of the specified `Study`.
+    #
+    # Note: The `get_linked_instances` function requires that a `max_page_size`
+    #       integer argument be passed in. In our case, we want to get _all_ of
+    #       the instances. Python has no "infinity" integer; and, even if it did,
+    #       if we were to specify too large of an integer, we'd get this error:
+    #       > "OverflowError: MongoDB can only handle up to 8-byte ints"
+    #       So, as a workaround, we pass in a number that is large enough that we
+    #       think it will account for all cases in practice (e.g., a study having
+    #       a trillion biosamples or a trillion data objects).
+    #
+    #       TODO: Update the `get_linked_instances` function to optionally impose _no_ limit.
+    #
+    large_max_page_size: int = 1_000_000_000_000
+    linked_biosamples_result: dict = get_linked_instances(
+        ids=[study_id],
+        types=["nmdc:Biosample"],
+        hydrate=False,  # we'll only use their `id` values
+        page_token=None,
+        max_page_size=large_max_page_size,
+        mdb=mdb,
+    )
+    biosample_ids = [d["id"] for d in linked_biosamples_result.get("resources", [])]
+    logging.debug(f"Found {len(biosample_ids)} Biosamples for Study {study_id}")
+
+    # Get all the `DataObject`s that are downstream from any of those `Biosample`s.
+    data_objects_by_biosample_id = {}
+    linked_data_objects_result: dict = get_linked_instances(
+        ids=biosample_ids,
+        types=["nmdc:DataObject"],
+        hydrate=True,  # we want the full `DataObject` documents
+        page_token=None,
+        max_page_size=large_max_page_size,
+        mdb=mdb,
+    )
+    for data_object in linked_data_objects_result.get("resources", []):
+        upstream_biosample_id = data_object["_downstream_of"][0]
+        if upstream_biosample_id not in data_objects_by_biosample_id.keys():
+            data_objects_by_biosample_id[upstream_biosample_id] = []
+
+        # Strip away the metadata fields injected by `get_linked_instances()`.
+        data_object.pop("_upstream_of", None)
+        data_object.pop("_downstream_of", None)
+        data_objects_by_biosample_id[upstream_biosample_id].append(data_object)
+
+    # Convert the `data_objects_by_biosample_id` dictionary into a list of dicts;
+    # i.e., into the format returned by the initial version of this API endpoint,
+    # which did not use the `get_linked_instances` function under the hood.
+    for biosample_id, data_objects in data_objects_by_biosample_id.items():
+        biosample_data_objects.append(
+            {
                 "biosample_id": biosample_id,
-                "data_objects": collected_data_objects,
+                "data_objects": data_objects,
             }
-            biosample_data_objects.append(result)
-
+        )
     return biosample_data_objects
 
 
@@ -699,96 +632,3 @@ def find_related_objects_for_workflow_execution(
     }
 
     return response
-
-
-jinja_env = Environment(
-    loader=PackageLoader("nmdc_runtime"), autoescape=select_autoescape()
-)
-
-
-def attr_index_sort_key(attr):
-    return "_" if attr == "id" else attr
-
-
-def documentation_links(jsonschema_dict, collection_names) -> dict:
-    """This function constructs a hierarchical catalog of (links to) schema classes and their slots.
-
-    The returned dictionary `doc_links` is used as input to the Jinja template `nmdc_runtime/templates/search.html`
-    in order to support user experience for `GET /search`.
-    """
-
-    # Note: All documentation URLs generated within this function will begin with this.
-    base_url = r"https://w3id.org/nmdc"
-
-    # Initialize dictionary in which to associate key/value pairs via the following for loop.
-    doc_links = {}
-
-    for collection_name in collection_names:
-        # Since a given collection can be associated with multiple classes, the `doc_links` dictionary
-        # will have a _list_ of values for each collection.
-        class_descriptors = []
-
-        # If the collection name is one that the `search.html` page has a dedicated section for,
-        # give it a top-level key; otherwise, nest it under `activity_set`.
-        key_hierarchy: List[str] = ["activity_set", collection_name]
-        if collection_name in ("biosample_set", "study_set", "data_object_set"):
-            key_hierarchy = [collection_name]
-
-        # Process the name of each class that the schema associates with this collection.
-        collection_spec = jsonschema_dict["$defs"]["Database"]["properties"][
-            collection_name
-        ]
-        class_names = get_class_names_from_collection_spec(collection_spec)
-        for idx, class_name in enumerate(class_names):
-            # Make a list of dictionaries, each of which describes one attribute of this class.
-            entity_attrs = list(jsonschema_dict["$defs"][class_name]["properties"])
-            entity_attr_descriptors = [
-                {"url": f"{base_url}/{attr_name}", "attr_name": attr_name}
-                for attr_name in entity_attrs
-            ]
-
-            # Make a dictionary describing this class.
-            class_descriptor = {
-                "collection_name": collection_name,
-                "entity_url": f"{base_url}/{class_name}",
-                "entity_name": class_name,
-                "entity_attrs": sorted(
-                    entity_attr_descriptors, key=itemgetter("attr_name")
-                ),
-            }
-
-            # Add that descriptor to this collection's list of class descriptors.
-            class_descriptors.append(class_descriptor)
-
-        # Add a key/value pair describing this collection to the `doc_links` dictionary.
-        # Reference: https://toolz.readthedocs.io/en/latest/api.html#toolz.dicttoolz.assoc_in
-        doc_links = assoc_in(doc_links, keys=key_hierarchy, value=class_descriptors)
-
-    return doc_links
-
-
-@router.get("/search", response_class=HTMLResponse, include_in_schema=False)
-def search_page(
-    mdb: MongoDatabase = Depends(get_mongo_db),
-):
-    template = jinja_env.get_template("search.html")
-    indexed_entity_attributes = merge(
-        {n: {"id"} for n in activity_collection_names(mdb)},
-        {
-            coll: sorted(attrs | {"id"}, key=attr_index_sort_key)
-            for coll, attrs in entity_attributes_to_index.items()
-        },
-    )
-    doc_links = documentation_links(
-        get_nmdc_jsonschema_dict(),
-        (
-            list(activity_collection_names(mdb))
-            + ["biosample_set", "study_set", "data_object_set"]
-        ),
-    )
-    html_content = template.render(
-        activity_collection_names=sorted(activity_collection_names(mdb)),
-        indexed_entity_attributes=indexed_entity_attributes,
-        doc_links=doc_links,
-    )
-    return HTMLResponse(content=html_content, status_code=200)

nmdc_runtime/api/endpoints/lib/linked_instances.py

@@ -0,0 +1,180 @@
+"""
+
+This module houses logic for the `GET /nmdcschema/linked_instances` endpoint, defined as
+`nmdc_runtime.api.endpoints.nmdcschema.linked_instances`, to avoid (further) bloating the
+`nmdc_runtime.api.endpoints.nmdcschema` module.
+
+"""
+
+from typing import Literal, Any
+
+from bson import ObjectId
+from pymongo.collection import Collection as MongoCollection
+from pymongo.database import Database as MongoDatabase
+from toolz import merge
+
+from nmdc_runtime.api.core.util import hash_from_str
+from nmdc_runtime.util import get_class_name_to_collection_names_map, nmdc_schema_view
+
+
+def hash_from_ids_and_types(ids: list[str], types: list[str]) -> str:
+    """A quick hash as a function of `ids` and `types`.
+
+    This will serve as part of a temporary mongo collection name.
+    Because it will only be "part of" the name, avoiding hash collisions isn't a priority.
+
+    Returns a hex digest truncated to 8 characters, so 16**8 ≈ 4M possible values.
+    """
+    return hash_from_str(
+        ",".join(sorted(ids)) + "." + ",".join(sorted(types)), algo="md5"
+    )[:8]
+
+
+def temp_linked_instances_collection_name(ids: list[str], types: list[str]) -> str:
+    """A name for a temporary mongo collection to store linked instances in service of an API request."""
+    return f"_runtime.tmp.linked_instances.{hash_from_ids_and_types(ids=ids,types=types)}.{ObjectId()}"
+
+
+def gather_linked_instances(
+    alldocs_collection: MongoCollection,
+    ids: list[str],
+    types: list[str],
+) -> str:
+    """Collect linked instances and stores them in a new temporary collection.
+
+    Run an aggregation pipeline over `alldocs_collection` that collects ∈`types` instances linked to `ids`.
+    The pipeline is run twice, once for each of {"downstream", "upstream"} directions.
+    """
+    merge_into_collection_name = temp_linked_instances_collection_name(
+        ids=ids, types=types
+    )
+    for direction in ["downstream", "upstream"]:
+        _ = list(
+            alldocs_collection.aggregate(
+                pipeline_for_direction(
+                    ids=ids,
+                    types=types,
+                    direction=direction,
+                    merge_into_collection_name=merge_into_collection_name,
+                ),
+                allowDiskUse=True,
+            )
+        )
+    return merge_into_collection_name
+
+
+def pipeline_for_direction(
+    ids: list[str],
+    types: list[str],
+    direction: Literal["downstream", "upstream"],
+    merge_into_collection_name: str,
+    alldocs_collection_name: str = "alldocs",
+) -> list:
+    """A pure function that returns the aggregation pipeline for `direction`.
+
+    The pipeline
+    - collects ∈`types` instances linked to `ids` along `direction`,
+    - retains only those document fields essential to the caller, and
+    - ensures the collected instances are present, and properly updated if applicable, in a merge-target collection.
+    """
+    return pipeline_for_instances_linked_to_ids_by_direction(
+        ids=ids,
+        types=types,
+        direction=direction,
+        alldocs_collection_name=alldocs_collection_name,
+    ) + [
+        {"$project": {"id": 1, "type": 1, f"_{direction}_of": 1}},
+        pipeline_stage_for_merging_instances_and_grouping_link_provenance_by_direction(
+            merge_into_collection_name=merge_into_collection_name, direction=direction
+        ),
+    ]
+
+
+def pipeline_for_instances_linked_to_ids_by_direction(
+    ids: list[str],
+    types: list[str],
+    direction: Literal["downstream", "upstream"],
+    alldocs_collection_name: str = "alldocs",
+    slim: bool = True,
+) -> list[dict[str, Any]]:
+    """
+    Returns an aggregation pipeline that:
+    - traverses the graph of documents in the alldocs collection, following `direction`-specific relationships
+      to discover documents linked to the documents given by `ids`.
+    - `$unwind`s the collected (via `$graphLookup`) docs,
+    - filters them by given `types` of interest,
+     - adds bookkeeping information about `direction`ality, and
+     - (optionally) projects only essential fields to reduce response latency and size.
+    """
+    return [
+        {"$match": {"id": {"$in": ids}}},
+        {
+            "$graphLookup": {
+                "from": alldocs_collection_name,
+                "startWith": f"$_{direction}.id",
+                "connectFromField": f"_{direction}.id",
+                "connectToField": "id",
+                "as": f"{direction}_docs",
+            }
+        },
+        {"$unwind": {"path": f"${direction}_docs"}},
+        {"$match": {f"{direction}_docs._type_and_ancestors": {"$in": types}}},
+        {"$addFields": {f"{direction}_docs._{direction}_of": ["$id"]}},
+        {"$replaceRoot": {"newRoot": f"${direction}_docs"}},
+    ] + ([{"$project": {"id": 1, "type": 1, f"_{direction}_of": 1}}] if slim else [])
+
+
+def pipeline_stage_for_merging_instances_and_grouping_link_provenance_by_direction(
+    merge_into_collection_name: str,
+    direction: Literal["downstream", "upstream"],
+) -> dict[str, Any]:
+    """
+    Returns an aggregation-pipeline step that merges its input document stream to a collection dedicated to serving
+    the caller in a manner amenable to pagination across multiple HTTP requests.
+    """
+    return {
+        "$merge": {
+            "into": merge_into_collection_name,
+            "on": "_id",
+            "whenMatched": [
+                {
+                    "$set": {
+                        f"_{direction}_of": {
+                            "$setUnion": [
+                                f"$_{direction}_of",
+                                f"$$new._{direction}_of",
+                            ]
+                        }
+                    }
+                }
+            ],
+            "whenNotMatched": "insert",
+        }
+    }
+
+
+def hydrated(resources: list[dict], mdb: MongoDatabase) -> list[dict]:
+    """Replace each `dict` in `resources` with a hydrated version.
+
+    Instead of returning the retrieved "full" documents as is, we merge each one with (a copy of) the corresponding
+    original document in *resources*, which includes additional fields, e.g. `_upstream_of` and `_downstream_of`.
+    """
+    class_name_to_collection_names_map = get_class_name_to_collection_names_map(
+        nmdc_schema_view()
+    )
+    types_of_resources = {r["type"] for r in resources}
+    full_docs_by_id = {}
+
+    for type in types_of_resources:
+        resource_ids_of_type = [d["id"] for d in resources if d["type"] == type]
+        schema_collection = mdb.get_collection(
+            # Note: We are assuming that documents of a given type are only allowed (by the schema) to reside in one
+            # collection. Based on that assumption, we will query only the _first_ collection whose name we get from
+            # the map. This assumption is continuously verified prior to code deployment via
+            # `test_get_class_name_to_collection_names_map_has_one_and_only_one_collection_name_per_class_name`.
+            class_name_to_collection_names_map[type.removeprefix("nmdc:")][0]
+        )
+        for doc in schema_collection.find({"id": {"$in": resource_ids_of_type}}):
+            full_docs_by_id[doc["id"]] = doc
+
+    return [merge(r, full_docs_by_id[r["id"]]) for r in resources]