nmdc-runtime 2.6.0__py3-none-any.whl → 2.12.0__py3-none-any.whl

nmdc_runtime/api/models/query.py

@@ -0,0 +1,246 @@
+import json
+import logging
+from typing import Optional, Any, Dict, List, Union, TypedDict
+
+import bson
+import bson.json_util
+from pydantic import (
+    model_validator,
+    Field,
+    BaseModel,
+    PositiveInt,
+    NonNegativeInt,
+    field_validator,
+    WrapSerializer,
+)
+from toolz import assoc, assoc_in
+from typing_extensions import Annotated
+
+from nmdc_runtime.api.core.util import pick
+
+
+def bson_to_json(doc: Any, handler) -> dict:
+    """Ensure a dict with e.g. mongo ObjectIds will serialize as JSON."""
+    return json.loads(bson.json_util.dumps(doc))
+
+
+Document = Annotated[Dict[str, Any], WrapSerializer(bson_to_json)]
+
+OneOrZero = Annotated[int, Field(ge=0, le=1)]
+One = Annotated[int, Field(ge=1, le=1)]
+MinusOne = Annotated[int, Field(ge=-1, le=-1)]
+OneOrMinusOne = Union[One, MinusOne]
+
+
+class CommandBase(BaseModel):
+    comment: Optional[Any] = None
+
+
+class CollStatsCommand(CommandBase):
+    collStats: str
+    scale: Optional[int] = 1
+
+
+class CountCommand(CommandBase):
+    count: str
+    query: Optional[Document] = None
+
+
+class FindCommand(CommandBase):
+    find: str
+    filter: Optional[Document] = None
+    projection: Optional[Dict[str, OneOrZero]] = None
+    allowPartialResults: Optional[bool] = True
+    batchSize: Optional[PositiveInt] = 101
+    sort: Optional[Dict[str, OneOrMinusOne]] = None
+    limit: Optional[NonNegativeInt] = None
+
+
+class AggregateCommand(CommandBase):
+    aggregate: str
+    pipeline: List[Document]
+    allowDiskUse: Optional[bool] = False
+    cursor: Optional[Document] = None
+
+    @field_validator("pipeline")
+    @classmethod
+    def disallow_invalid_pipeline_stages(
+        cls, pipeline: List[Document]
+    ) -> List[Document]:
+        deny_list = ["$out", "$merge"]
+
+        if any(
+            key in deny_list for pipeline_stage in pipeline for key in pipeline_stage
+        ):
+            raise ValueError("$Out and $merge pipeline stages are not allowed.")
+
+        return pipeline
+
+    @model_validator(mode="before")
+    @classmethod
+    def ensure_default_value_for_cursor(cls, data: Any) -> Document:
+        if isinstance(data, dict) and "cursor" not in data:
+            return assoc(data, "cursor", {"batchSize": 25})
+        return data
+
+
+class GetMoreCommand(CommandBase):
+    # Note: No `collection` field. See `QueryContinuation` for inter-API-request "sessions" are modeled.
+    getMore: str  # Note: runtime uses a `str` id, not an `int` like mongo's native session cursors.
+    batchSize: Optional[PositiveInt] = None
+
+
+class CommandResponse(BaseModel):
+    ok: OneOrZero
+
+
+class CollStatsCommandResponse(CommandResponse):
+    ns: str
+    size: float
+    count: float
+    avgObjSize: Optional[float] = None
+    storageSize: float
+    totalIndexSize: float
+    totalSize: float
+    scaleFactor: float
+
+
+class CountCommandResponse(CommandResponse):
+    n: NonNegativeInt
+
+
+class CommandResponseCursor(BaseModel):
+    # Note: No `ns` field, `id` is a `str`, and `partialResultsReturned` aliased to `queriedShardsUnavailable` to be
+    # less confusing to Runtime API clients. See `QueryContinuation` for inter-API-request "sessions" are modeled.
+    batch: List[Document]
+    partialResultsReturned: Optional[bool] = Field(
+        None, alias="queriedShardsUnavailable"
+    )
+    id: Optional[str] = None
+
+    @field_validator("id", mode="before")
+    @classmethod
+    def coerce_int_to_str(cls, value: Any) -> Any:
+        if isinstance(value, int):
+            return str(value)
+        else:
+            return value
+
+
+class CursorYieldingCommandResponse(CommandResponse):
+    cursor: CommandResponseCursor
+
+    @classmethod
+    def slimmed(cls, cmd_response) -> Optional["CursorYieldingCommandResponse"]:
+        """Create a new response object that retains only the `_id` for each cursor batch document."""
+        dump: dict = cmd_response.model_dump(exclude_unset=True)
+
+        # If any dictionary in this batch lacks an `_id` key, log a warning and return `None`.`
+        id_list = [pick(["_id"], batch_doc) for batch_doc in dump["cursor"]["batch"]]
+        if any("_id" not in doc for doc in id_list):
+            logging.warning("Some documents in the batch lack an `_id` field.")
+            return None
+
+        dump = assoc_in(
+            dump,
+            ["cursor", "batch"],
+            id_list,
+        )
+        return cls(**dump)
+
+
+class DeleteStatement(BaseModel):
+    q: Document
+    # `limit` is required: https://www.mongodb.com/docs/manual/reference/command/delete/#std-label-deletes-array-limit
+    limit: OneOrZero
+    hint: Optional[Dict[str, OneOrMinusOne]] = None
+
+
+class DeleteCommand(CommandBase):
+    delete: str
+    deletes: List[DeleteStatement]
+
+
+class DeleteCommandResponse(CommandResponse):
+    ok: OneOrZero
+    n: NonNegativeInt
+    writeErrors: Optional[List[Document]] = None
+
+
+# Custom types for the `delete_specs` derived from `DeleteStatement`s.
+DeleteSpec = TypedDict("DeleteSpec", {"filter": Document, "limit": OneOrZero})
+DeleteSpecs = List[DeleteSpec]
+
+
+# If `multi==True` all documents that meet the query criteria will be updated.
+# Else only a single document that meets the query criteria will be updated.
+class UpdateStatement(BaseModel):
+    q: Document
+    u: Document
+    upsert: bool = False
+    multi: bool = False
+    hint: Optional[Dict[str, OneOrMinusOne]] = None
+
+
+# Custom types for the `update_specs` derived from `UpdateStatement`s.
+UpdateSpec = TypedDict("UpdateSpec", {"filter": Document, "limit": OneOrZero})
+UpdateSpecs = List[UpdateSpec]
+
+
+class UpdateCommand(CommandBase):
+    update: str
+    updates: List[UpdateStatement]
+
+
+class DocumentUpserted(BaseModel):
+    index: NonNegativeInt
+    _id: bson.ObjectId
+
+
+class UpdateCommandResponse(CommandResponse):
+    ok: OneOrZero
+    n: NonNegativeInt
+    nModified: NonNegativeInt
+    upserted: Optional[List[DocumentUpserted]] = None
+    writeErrors: Optional[List[Document]] = None
+
+
+QueryCmd = Union[FindCommand, AggregateCommand]
+
+CursorYieldingCommand = Union[
+    QueryCmd,
+    GetMoreCommand,
+]
+
+
+Cmd = Union[
+    CursorYieldingCommand,
+    CollStatsCommand,
+    CountCommand,
+    DeleteCommand,
+    UpdateCommand,
+]
+
+CommandResponseOptions = Union[
+    CursorYieldingCommandResponse,
+    CollStatsCommandResponse,
+    CountCommandResponse,
+    DeleteCommandResponse,
+    UpdateCommandResponse,
+]
+
+
+def command_response_for(type_):
+    r"""
+    TODO: Add a docstring and type hints to this function.
+    """
+    if issubclass(type_, CursorYieldingCommand):
+        return CursorYieldingCommandResponse
+
+    d = {
+        CollStatsCommand: CollStatsCommandResponse,
+        CountCommand: CountCommandResponse,
+        DeleteCommand: DeleteCommandResponse,
+        UpdateCommand: UpdateCommandResponse,
+    }
+    return d.get(type_)

nmdc_runtime/api/models/query_continuation.py

@@ -0,0 +1,111 @@
+"""
+A *query continuation* is a means to effectively resume a query, i.e. a `find` or `aggregate` MongoDB database command.
+
+A *query continuation* document represents a *continuation* (cf. <https://en.wikipedia.org/wiki/Continuation>) for a
+query and uses a stored value ("cursor") for MongoDB's guaranteed unique-valued document field, `_id`,
+such that the documents returned by the command are guaranteed to be sorted in ascending order by `_id`.
+
+In this way, an API client may retrieve all documents defined by a `find` or `aggregate` command over multiple HTTP
+requests. One can think of this process as akin to pagination; however, with "cursor-based" pagination, there are no
+guarantees wrt a fixed "page size".
+
+"""
+
+import datetime
+import logging
+import json
+
+from pydantic import BaseModel, Field
+from pymongo.database import Database as MongoDatabase
+
+from nmdc_runtime.api.core.idgen import generate_one_id
+from nmdc_runtime.api.core.util import now
+from nmdc_runtime.api.db.mongo import get_mongo_db
+from nmdc_runtime.api.models.query import (
+    CommandResponse,
+    QueryCmd,
+)
+
+COLLECTION_NAME_FOR_QUERY_CONTINUATIONS = "_runtime.query_continuations"
+
+_mdb: MongoDatabase = get_mongo_db()
+_qc_collection = _mdb[COLLECTION_NAME_FOR_QUERY_CONTINUATIONS]
+
+# Ensure one-hour TTL on `_runtime.query_continuations` documents via TTL Index.
+# Reference: https://www.mongodb.com/docs/manual/core/index-ttl/
+_qc_collection.create_index({"last_modified": 1}, expireAfterSeconds=3600)
+
+
+def not_empty(lst: list) -> bool:
+    return len(lst) > 0
+
+
+class QueryContinuation(BaseModel):
+    """A query that has not completed, and that may be resumed, using `cursor` to modify `query_cmd`.
+
+    This model is intended to represent the state of a logical "session" to "page" through a query's results
+    over several HTTP requests, and may be discarded after fetching all "batches" of documents.
+
+    Thus, a mongo collection tracking query continuations may be reasonably given e.g. a so-called "TTL Index"
+    for the `last_modified` field, assuming that `last_modified` is updated each time `query` is updated.
+    """
+
+    id: str = Field(..., alias="_id")
+    query_cmd: QueryCmd
+    cursor: str
+    last_modified: datetime.datetime
+
+
+class QueryContinuationError(Exception):
+    def __init__(self, detail: str):
+        self.detail = detail
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}: {self.detail})"
+
+
+def dump_qc(m: BaseModel):
+    return m.model_dump(by_alias=True, exclude_unset=True)
+
+
+def create_qc(query_cmd: QueryCmd, cmd_response: CommandResponse) -> QueryContinuation:
+    """Creates query continuation from command and response, and persists continuation to database."""
+
+    logging.info(f"cmd_response: {cmd_response}")
+    last_id = json.dumps(cmd_response.cursor.batch[-1]["_id"])
+    logging.info(f"Last document ID for query continuation: {last_id}")
+    cc = QueryContinuation(
+        _id=generate_one_id(_mdb, "query_continuation"),
+        query_cmd=query_cmd,
+        cursor=last_id,
+        last_modified=now(),
+    )
+    _qc_collection.insert_one(dump_qc(cc))
+    return cc
+
+
+def get_qc_by__id(_id: str) -> QueryContinuation | None:
+    r"""
+    Returns the `QueryContinuation` having the specified `_id` value, raising an exception
+    if the corresponding document does not exist in the database.
+    """
+    doc = _qc_collection.find_one({"_id": _id})
+    if doc is None:
+        raise QueryContinuationError(f"cannot find cc with id {_id}")
+    return QueryContinuation(**doc)
+
+
+def get_last_doc__id_for_qc(query_continuation: QueryContinuation) -> str:
+    """
+    Retrieve the last document `_id` for the given `QueryContinuation`.
+    """
+    # Assuming `query_continuation` has an attribute `cursor` that stores the last document _id
+    logging.info(f"Cursor for last doc query continuation: {query_continuation.cursor}")
+    return json.loads(query_continuation.cursor)
+
+
+def get_initial_query_for_qc(query_continuation: QueryContinuation) -> QueryCmd:
+    """
+    Retrieve the initial query command for the given `QueryContinuation`.
+    """
+    return query_continuation.query_cmd

nmdc_runtime/api/models/run.py

@@ -0,0 +1,161 @@
+from enum import Enum
+import os
+from functools import lru_cache
+from typing import List, Optional
+
+from dagster_graphql import DagsterGraphQLClient
+from pydantic import BaseModel
+from pymongo.database import Database as MongoDatabase
+from toolz import merge
+
+from nmdc_runtime.api.core.idgen import generate_one_id
+from nmdc_runtime.api.core.util import now, now_str, raise404_if_none, pick
+from nmdc_runtime.api.models.user import User
+
+PRODUCER_URL_BASE_DEFAULT = (
+    "https://github.com/microbiomedata/nmdc-runtime/tree/main/nmdc_runtime/"
+)
+SCHEMA_URL_BASE_DEFAULT = (
+    "https://github.com/microbiomedata/nmdc-runtime/tree/main/nmdc_runtime/"
+)
+
+PRODUCER_URL = PRODUCER_URL_BASE_DEFAULT.replace("/main/", "/v0-0-1/") + "producer"
+SCHEMA_URL = SCHEMA_URL_BASE_DEFAULT.replace("/main/", "/v0-0-1/") + "schema.json"
+
+
+class OpenLineageBase(BaseModel):
+    producer: str
+    schemaURL: str
+
+
+class RunUserSpec(BaseModel):
+    job_id: str
+    run_config: dict = {}
+    inputs: List[str] = []
+
+
+class JobSummary(OpenLineageBase):
+    id: str
+    description: str
+
+
+class Run(BaseModel):
+    id: str
+    facets: Optional[dict] = None
+
+
+class RunEventType(str, Enum):
+    REQUESTED = "REQUESTED"
+    STARTED = "STARTED"
+    FAIL = "FAIL"
+    COMPLETE = "COMPLETE"
+
+
+class RunSummary(OpenLineageBase):
+    id: str
+    status: RunEventType
+    started_at_time: str
+    was_started_by: str
+    inputs: List[str]
+    outputs: List[str]
+    job: JobSummary
+
+
+class RunEvent(OpenLineageBase):
+    run: Run
+    job: JobSummary
+    type: RunEventType
+    time: str
+    inputs: Optional[List[str]] = []
+    outputs: Optional[List[str]] = []
+
+
+@lru_cache
+def get_dagster_graphql_client() -> DagsterGraphQLClient:
+    hostname, port_str = os.getenv("DAGIT_HOST").split("://", 1)[-1].split(":", 1)
+    port_number = int(port_str)
+    return DagsterGraphQLClient(hostname=hostname, port_number=port_number)
+
+
+def _add_run_requested_event(run_spec: RunUserSpec, mdb: MongoDatabase, user: User):
+    # XXX what we consider a "job" here, is currently a "workflow" elsewhere...
+    job = raise404_if_none(mdb.workflows.find_one({"id": run_spec.job_id}))
+    run_id = generate_one_id(mdb, "runs")
+    event = RunEvent(
+        producer=user.username,
+        schemaURL=SCHEMA_URL,
+        run=Run(id=run_id, facets={"nmdcRuntime_runConfig": run_spec.run_config}),
+        job=merge(
+            pick(["id", "description"], job),
+            {"producer": PRODUCER_URL, "schemaURL": SCHEMA_URL},
+        ),
+        type=RunEventType.REQUESTED,
+        time=now_str(),
+        inputs=run_spec.inputs,
+    )
+    mdb.run_events.insert_one(event.model_dump())
+    return run_id
+
+
+def _add_run_started_event(run_id: str, mdb: MongoDatabase):
+    requested: RunEvent = RunEvent(
+        **raise404_if_none(
+            mdb.run_events.find_one(
+                {"run.id": run_id, "type": "REQUESTED"}, sort=[("time", -1)]
+            )
+        )
+    )
+    mdb.run_events.insert_one(
+        RunEvent(
+            producer=PRODUCER_URL,
+            schemaURL=SCHEMA_URL,
+            run=requested.run,
+            job=requested.job,
+            type=RunEventType.STARTED,
+            time=now_str(),
+        ).model_dump()
+    )
+    return run_id
+
+
+def _add_run_fail_event(run_id: str, mdb: MongoDatabase):
+    requested: RunEvent = RunEvent(
+        **raise404_if_none(
+            mdb.run_events.find_one(
+                {"run.id": run_id, "type": "REQUESTED"}, sort=[("time", -1)]
+            )
+        )
+    )
+    mdb.run_events.insert_one(
+        RunEvent(
+            producer=PRODUCER_URL,
+            schemaURL=SCHEMA_URL,
+            run=requested.run,
+            job=requested.job,
+            type=RunEventType.FAIL,
+            time=now_str(),
+        ).model_dump()
+    )
+    return run_id
+
+
+def _add_run_complete_event(run_id: str, mdb: MongoDatabase, outputs: List[str]):
+    started: RunEvent = RunEvent(
+        **raise404_if_none(
+            mdb.run_events.find_one(
+                {"run.id": run_id, "type": "STARTED"}, sort=[("time", -1)]
+            )
+        )
+    )
+    mdb.run_events.insert_one(
+        RunEvent(
+            producer=PRODUCER_URL,
+            schemaURL=SCHEMA_URL,
+            run=started.run,
+            job=started.job,
+            type=RunEventType.COMPLETE,
+            time=now_str(),
+            outputs=outputs,
+        ).model_dump()
+    )
+    return run_id

nmdc_runtime/api/models/site.py

@@ -0,0 +1,87 @@
+from typing import List, Optional
+
+import pymongo.database
+from fastapi import Depends
+from jose import JWTError, jwt
+from pydantic import BaseModel
+
+from nmdc_runtime.api.core.auth import (
+    verify_password,
+    TokenData,
+    optional_oauth2_scheme,
+)
+from nmdc_runtime.api.db.mongo import get_mongo_db
+from nmdc_runtime.api.models.user import (
+    oauth2_scheme,
+    credentials_exception,
+    SECRET_KEY,
+    ALGORITHM,
+)
+
+
+class Site(BaseModel):
+    id: str
+    capability_ids: List[str] = []
+
+
+class SiteClientInDB(BaseModel):
+    id: str
+    hashed_secret: str
+
+
+class SiteInDB(Site):
+    clients: List[SiteClientInDB] = []
+
+
+def get_site(mdb, client_id: str) -> Optional[SiteInDB]:
+    r"""
+    Returns the site, if any, for which the specified `client_id` was generated.
+    """
+
+    site = mdb.sites.find_one({"clients.id": client_id})
+    if site is not None:
+        return SiteInDB(**site)
+
+
+def authenticate_site_client(mdb, client_id: str, client_secret: str):
+    site = get_site(mdb, client_id)
+    if not site:
+        return False
+    hashed_secret = next(
+        client.hashed_secret for client in site.clients if client.id == client_id
+    )
+    if not verify_password(client_secret, hashed_secret):
+        return False
+    return site
+
+
+async def get_current_client_site(
+    token: str = Depends(oauth2_scheme),
+    mdb: pymongo.database.Database = Depends(get_mongo_db),
+):
+    if mdb.invalidated_tokens.find_one({"_id": token}):
+        raise credentials_exception
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        subject: str = payload.get("sub")
+        if subject is None:
+            raise credentials_exception
+        if not subject.startswith("client:"):
+            raise credentials_exception
+        client_id = subject.split("client:", 1)[1]
+        token_data = TokenData(subject=client_id)
+    except JWTError:
+        raise credentials_exception
+    site = get_site(mdb, client_id=token_data.subject)
+    if site is None:
+        raise credentials_exception
+    return site
+
+
+async def maybe_get_current_client_site(
+    token: str = Depends(optional_oauth2_scheme),
+    mdb: pymongo.database.Database = Depends(get_mongo_db),
+):
+    if token is None:
+        return None
+    return await get_current_client_site(token, mdb)

nmdc_runtime/api/models/trigger.py

@@ -0,0 +1,13 @@
+import datetime
+
+from pydantic import BaseModel
+
+
+class TriggerBase(BaseModel):
+    object_type_id: str
+    workflow_id: str
+
+
+class Trigger(TriggerBase):
+    id: str
+    created_at: datetime.datetime