nmdc-runtime 2.9.0__py3-none-any.whl → 2.11.0__py3-none-any.whl

nmdc_runtime/api/endpoints/jobs.py

@@ -0,0 +1,143 @@
+import json
+from typing import Optional, Annotated
+
+from pymongo.database import Database
+from fastapi import APIRouter, Depends, Query, HTTPException, Path
+from pymongo.errors import ConnectionFailure, OperationFailure
+from starlette import status
+
+from nmdc_runtime.api.core.util import (
+    raise404_if_none,
+)
+from nmdc_runtime.api.db.mongo import get_mongo_db
+from nmdc_runtime.api.endpoints.util import list_resources, _claim_job
+from nmdc_runtime.api.models.job import Job, JobClaim
+from nmdc_runtime.api.models.operation import Operation, MetadataT
+from nmdc_runtime.api.models.site import (
+    Site,
+    maybe_get_current_client_site,
+    get_current_client_site,
+)
+from nmdc_runtime.api.models.util import ListRequest, ListResponse, ResultT
+
+router = APIRouter()
+
+
+@router.get(
+    "/jobs", response_model=ListResponse[Job], response_model_exclude_unset=True
+)
+def list_jobs(
+    req: Annotated[ListRequest, Query()],
+    mdb: Database = Depends(get_mongo_db),
+    maybe_site: Optional[Site] = Depends(maybe_get_current_client_site),
+):
+    """List pre-configured workflow jobs.
+
+    If authenticated as a site client, `req.filter` defaults to fetch unclaimed jobs
+    that are claimable by the site client. This default can be overridden to view all jobs
+    by explicitly passing a `req.filter` of `{}`.
+    """
+    if isinstance(maybe_site, Site) and req.filter is None:
+        req.filter = json.dumps({"claims.site_id": {"$ne": maybe_site.id}})
+    return list_resources(req, mdb, "jobs")
+
+
+@router.get("/jobs/{job_id}", response_model=Job, response_model_exclude_unset=True)
+def get_job_info(
+    job_id: str,
+    mdb: Database = Depends(get_mongo_db),
+):
+    return raise404_if_none(mdb.jobs.find_one({"id": job_id}))
+
+
+@router.post("/jobs/{job_id}:claim", response_model=Operation[ResultT, MetadataT])
+def claim_job(
+    job_id: str,
+    mdb: Database = Depends(get_mongo_db),
+    site: Site = Depends(get_current_client_site),
+):
+    return _claim_job(job_id, mdb, site)
+
+
+@router.post("/jobs/{job_id}:release")
+def release_job(
+    job_id: Annotated[
+        str,
+        Path(
+            title="Job ID",
+            description="The `id` of the job.\n\n_Example_: `nmdc:f81d4fae-7dec-11d0-a765-00a0c91e6bf6`",
+            examples=["nmdc:f81d4fae-7dec-11d0-a765-00a0c91e6bf6"],
+        ),
+    ],
+    mdb: Database = Depends(get_mongo_db),
+    site: Site = Depends(get_current_client_site),
+) -> Optional[Job]:
+    r"""
+    Release the specified job.
+
+    Releasing a job cancels all the unfinished operations (of that job)
+    claimed by the `site` associated with the logged-in site client.
+
+    Return the updated job, reflecting that the aforementioned operations have been cancelled.
+    """
+    job = Job(**raise404_if_none(mdb.jobs.find_one({"id": job_id})))
+    active_job_claims_by_this_site = list(
+        mdb.operations.find(
+            {
+                "metadata.job.id": job_id,
+                "metadata.site_id": site.id,
+                "done": False,
+            },
+            ["id"],
+        )
+    )
+    job_claims_by_this_site_post_release = [
+        JobClaim(op_id=claim["id"], site_id=site.id, done=True, cancelled=True)
+        for claim in active_job_claims_by_this_site
+    ]
+    job_claims_not_by_this_site = [
+        claim for claim in job.claims if (claim.site_id != site.id)
+    ]
+
+    # Execute MongoDB transaction to ensure atomic change of job document plus relevant set of operations documents.
+    def transactional_update(session):
+        mdb.operations.update_many(
+            {"id": {"$in": [claim["id"] for claim in active_job_claims_by_this_site]}},
+            {"$set": {"metadata.cancelled": True, "metadata.done": True}},
+            session=session,
+        )
+        job_claim_subdocuments_post_release = [
+            claim.model_dump(exclude_unset=True)
+            for claim in (
+                job_claims_not_by_this_site + job_claims_by_this_site_post_release
+            )
+        ]
+        mdb.jobs.update_one(
+            {"id": job_id},
+            {"$set": {"claims": job_claim_subdocuments_post_release}},
+            session=session,
+        )
+
+    try:
+        with mdb.client.start_session() as session:
+            with session.start_transaction():
+                transactional_update(session)
+    except (ConnectionFailure, OperationFailure) as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Transaction failed: {e}",
+        )
+
+    # Return the updated `jobs` document.
+    #
+    # TODO: Consider retrieving the document within the transaction
+    #       to ensure it still exists.
+    #
+    updated_job = mdb.jobs.find_one({"id": job_id})
+    if updated_job is None:
+        # Note: We return `None` in this case because that's what the
+        #       endpoint originally did in this case, and we don't want
+        #       to introduce a breaking change as part of this refactor.
+        return None
+    else:
+        return Job(**updated_job)

nmdc_runtime/api/endpoints/lib/helpers.py

@@ -0,0 +1,274 @@
+import json
+import bson.json_util
+from typing import List
+
+from pymongo.database import Database
+from refscan.lib.Finder import Finder
+from refscan.lib.helpers import derive_schema_class_name_from_document
+from refscan.scanner import identify_referring_documents, scan_outgoing_references
+
+from nmdc_runtime.api.models.lib.helpers import derive_update_specs
+from nmdc_runtime.api.models.query import UpdateCommand, UpdateSpecs
+from nmdc_runtime.util import get_allowed_references, nmdc_schema_view
+
+
+def make_violation_message(
+    collection_name: str,
+    source_document_id: str,
+    source_field_name: str,
+    target_document_id: str,
+) -> str:
+    r"""
+    Constructs a violation message that indicates that a document would contain a broken reference.
+
+    :param collection_name: The name of the collection containing the document containing the broken reference
+    :param source_document_id: The `id` of the document containing the broken reference
+    :param source_field_name: The name of the field containing the broken reference
+    :param target_document_id: The `id` of the document that is being referenced
+
+    :return: A formatted string describing the violation
+    """
+    return (
+        f"The document having id='{source_document_id}' in "
+        f"the collection '{collection_name}' contains a "
+        f"reference (in its '{source_field_name}' field, "
+        f"referring to the document having id='{target_document_id}') "
+        f"which would be broken."
+    )
+
+
+def simulate_updates_and_check_references(
+    db: Database, update_cmd: UpdateCommand
+) -> List[str]:
+    r"""
+    Checks whether, if the specified updates were performed on the specified database,
+    both of the following things would be true afterward:
+    1. (Regarding outgoing references): The updated documents do not contain any
+       broken references.
+    2. (Regarding incoming references): The documents that originally _referenced_
+       any of the updated documents do not contain any broken references.
+       This check is necessary because update operations can currently change `id`
+       and `type` values, which can affect what can legally reference those documents.
+
+    This function checks those things by performing the updates within a MongoDB
+    transaction, leaving the transaction in the _pending_ (i.e. not committed) state,
+    and then performing various checks on the database in that _pending_ state.
+
+    :param db: The database on which to simulate performing the updates
+    :param update_cmd: The command that specifies the updates
+
+    :return: List of violation messages. If the list is empty, it means that—if
+             the updates had been performed (instead of only simulated) here—they
+             would not have left behind any broken references.
+    """
+
+    # Initialize the list of violation messages that we will return.
+    violation_messages: List[str] = []
+
+    # Instantiate a `Finder` bound to the Mongo database. This will be
+    # used later, to identify and check inter-document references.
+    finder = Finder(database=db)
+
+    # Extract the collection name from the command.
+    collection_name = update_cmd.update
+
+    # Derive the update specifications from the command.
+    update_specs: UpdateSpecs = derive_update_specs(update_cmd)
+
+    # Get a reference to a `SchemaView` bound to the NMDC schema, so we can
+    # use it to, for example, map `type` field values to schema class names.
+    schema_view = nmdc_schema_view()
+
+    # Get some data structures that indicate which fields of which documents
+    # can legally contain references, according to the NMDC schema.
+    legal_references = get_allowed_references()
+    reference_field_names_by_source_class_name = (
+        legal_references.get_reference_field_names_by_source_class_name()
+    )
+
+    # Start a "throwaway" MongoDB transaction so we can simulate the updates.
+    with db.client.start_session() as session:
+        with session.start_transaction():
+
+            # Make a list of the `_id`, `id`, and `type` values of the documents that
+            # the user wants to update.
+            projection = {"_id": 1, "id": 1, "type": 1}
+            subject_document_summaries_pre_update = list(
+                db[collection_name].find(
+                    filter={"$or": [spec["filter"] for spec in update_specs]},
+                    projection=projection,
+                    session=session,
+                )
+            )
+
+            # Make a set of the `_id` values of the subject documents so that (later) we can
+            # check whether a given _referring_ document is also one of the _subject_
+            # documents (i.e. is among the documents the user wants to update).
+            subject_document_object_ids = set(
+                tdd["_id"] for tdd in subject_document_summaries_pre_update
+            )
+
+            # Identify _all_ documents that reference any of the subject documents.
+            all_referring_document_descriptors_pre_update = []
+            for subject_document_summary in subject_document_summaries_pre_update:
+                # If the document summary lacks the "id" field, we already know that no
+                # documents reference it (since they would have to _use_ that "id" value to
+                # do so); so, we abort this iteration and move on to the next subject document.
+                if "id" not in subject_document_summary:
+                    continue
+
+                referring_document_descriptors = identify_referring_documents(
+                    document=subject_document_summary,  # expects at least "id" and "type"
+                    schema_view=schema_view,
+                    references=legal_references,
+                    finder=finder,
+                    client_session=session,
+                )
+                all_referring_document_descriptors_pre_update.extend(
+                    referring_document_descriptors
+                )
+
+            # Simulate the updates (i.e. apply them within the context of the transaction).
+            db.command(
+                # Note: This expression was copied from the `_run_mdb_cmd` function in `queries.py`.
+                # TODO: Document this expression (i.e. the Pydantic->JSON->BSON chain).
+                bson.json_util.loads(
+                    json.dumps(update_cmd.model_dump(exclude_unset=True))
+                ),
+                session=session,
+            )
+            # For each referring document, check whether any of its outgoing references
+            # is broken (in the context of the transaction).
+            for descriptor in all_referring_document_descriptors_pre_update:
+                referring_document_oid = descriptor["source_document_object_id"]
+                referring_document_id = descriptor["source_document_id"]
+                referring_collection_name = descriptor["source_collection_name"]
+                # If the referring document is among the documents that the user wanted to
+                # update, we skip it for now. We will check its outgoing references later
+                # (i.e. when we check the outgoing references of _all_ updated documents).
+                if referring_document_oid in subject_document_object_ids:
+                    continue
+                # Get the referring document, so we can check its outgoing references.
+                # Note: We project only the fields that can legally contain references,
+                #       plus other fields involved in referential integrity checking.
+                referring_document_reference_field_names = (
+                    reference_field_names_by_source_class_name[
+                        descriptor["source_class_name"]
+                    ]
+                )
+                projection = {
+                    field_name: 1
+                    for field_name in referring_document_reference_field_names
+                } | {
+                    "_id": 1,
+                    "id": 1,
+                    "type": 1,
+                }  # note: `|` unions the dicts
+                referring_document = db[referring_collection_name].find_one(
+                    {"_id": referring_document_oid},
+                    projection=projection,
+                    session=session,
+                )
+                # Note: We assert that the referring document exists (to satisfy the type checker).
+                assert (
+                    referring_document is not None
+                ), "A referring document has vanished."
+                violations = scan_outgoing_references(
+                    document=referring_document,
+                    source_collection_name=referring_collection_name,
+                    schema_view=schema_view,
+                    references=legal_references,
+                    finder=finder,
+                    client_session=session,  # so it uses the pending transaction's session
+                )
+                # For each violation (i.e. broken reference) that exists, add a violation message
+                # to the list of violation messages.
+                #
+                # TODO: The violation might not involve a reference to one of the
+                #       subject documents. The `scan_outgoing_references` function
+                #       scans _all_ references emanating from the document.
+                #
+                for violation in violations:
+                    source_field_name = violation.source_field_name
+                    target_id = violation.target_id
+                    violation_messages.append(
+                        make_violation_message(
+                            collection_name=referring_collection_name,
+                            source_document_id=referring_document_id,
+                            source_field_name=source_field_name,
+                            target_document_id=target_id,
+                        )
+                    )
+
+            # For each updated document, check whether any of its outgoing references
+            # is broken (in the context of the transaction).
+            for subject_document_summary in subject_document_summaries_pre_update:
+                subject_document_oid = subject_document_summary["_id"]
+                subject_document_id = subject_document_summary["id"]
+                subject_document_class_name = derive_schema_class_name_from_document(
+                    document=subject_document_summary,
+                    schema_view=schema_view,
+                )
+                assert (
+                    subject_document_class_name is not None
+                ), "The updated document does not represent a valid schema class instance."
+                subject_collection_name = (
+                    collection_name  # makes a disambiguating alias
+                )
+                # Get the updated document, so we can check its outgoing references.
+                # Note: We project only the fields that can legally contain references,
+                #       plus other fields involved in referential integrity checking.
+                updated_document_reference_field_names = (
+                    reference_field_names_by_source_class_name[
+                        subject_document_class_name
+                    ]
+                )
+                projection = {
+                    field_name: 1
+                    for field_name in updated_document_reference_field_names
+                } | {
+                    "_id": 1,
+                    "id": 1,
+                    "type": 1,
+                }  # note: `|` unions the dicts
+                updated_document = db[subject_collection_name].find_one(
+                    {"_id": subject_document_oid},
+                    projection=projection,
+                    session=session,
+                )
+                # Note: We assert that the updated document exists (to satisfy the type checker).
+                assert updated_document is not None, "An updated document has vanished."
+                violations = scan_outgoing_references(
+                    document=updated_document,
+                    source_collection_name=subject_collection_name,
+                    schema_view=schema_view,
+                    references=legal_references,
+                    finder=finder,
+                    client_session=session,  # so it uses the pending transaction's session
+                )
+                # For each violation (i.e. broken reference) that exists, add a violation message
+                # to the list of violation messages.
+                for violation in violations:
+                    source_field_name = violation.source_field_name
+                    target_id = violation.target_id
+                    violation_messages.append(
+                        make_violation_message(
+                            collection_name=subject_collection_name,
+                            source_document_id=subject_document_id,
+                            source_field_name=source_field_name,
+                            target_document_id=target_id,
+                        )
+                    )
+
+            # Whatever happens (i.e. whether there are violations or not), abort the transaction.
+            #
+            # Note: If an exception was raised within this `with` block, the transaction
+            #       will already have been aborted automatically (and execution will not
+            #       have reached this statement). On the other hand, if no exception
+            #       was raised, we explicitly abort the transaction so that the updates
+            #       that we "simulated" in this block do not get applied to the real database.
+            #       Reference: https://pymongo.readthedocs.io/en/stable/api/pymongo/client_session.html
+            #
+            session.abort_transaction()
+
+    return violation_messages

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]