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

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, 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(as_str=True),
+        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(as_str=True),
+        ).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(as_str=True),
+        ).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(as_str=True),
+            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

nmdc_runtime/api/models/user.py

@@ -0,0 +1,140 @@
+from typing import List, Optional, Union
+
+import pymongo.database
+from fastapi import Depends, HTTPException
+from jose import JWTError, jwt
+from pydantic import BaseModel
+
+from nmdc_runtime.api.core.auth import (
+    verify_password,
+    SECRET_KEY,
+    ALGORITHM,
+    oauth2_scheme,
+    credentials_exception,
+    TokenData,
+    bearer_scheme,
+)
+
+from nmdc_runtime.api.models.site import get_site
+
+from nmdc_runtime.api.db.mongo import get_mongo_db
+
+
+class User(BaseModel):
+    username: str
+    email: Optional[str] = None
+    full_name: Optional[str] = None
+    site_admin: Optional[List[str]] = []
+    disabled: Optional[bool] = False
+
+
+class UserIn(User):
+    password: str
+
+
+class UserInDB(User):
+    hashed_password: str
+
+
+def get_user(mdb, username: str) -> Optional[UserInDB]:
+    r"""
+    Returns the user having the specified username.
+    """
+
+    user = mdb.users.find_one({"username": username})
+    if user is not None:
+        return UserInDB(**user)
+
+
+def authenticate_user(mdb, username: str, password: str) -> Union[UserInDB, bool]:
+    r"""
+    Returns the user, if any, having the specified username/password combination.
+    """
+
+    user = get_user(mdb, username)
+    if not user:
+        return False
+    if not verify_password(password, user.hashed_password):
+        return False
+    return user
+
+
+async def get_current_user(
+    token: str = Depends(oauth2_scheme),
+    bearer_credentials: str = Depends(bearer_scheme),
+    mdb: pymongo.database.Database = Depends(get_mongo_db),
+) -> UserInDB:
+    r"""
+    Returns a user based upon the provided token.
+
+    If the token belongs to a site client, the returned user is an ephemeral "user"
+    whose username is the site client's `client_id`.
+
+    Raises an exception if the token is invalid.
+    """
+
+    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("user:") and not subject.startswith("client:"):
+            raise credentials_exception
+
+        # subject is in the form "user:foo" or "client:bar"
+        username = subject.split(":", 1)[1]
+        token_data = TokenData(subject=username)
+    except (JWTError, AttributeError) as e:
+        print(f"jwt error: {e}")
+        raise credentials_exception
+
+    # Coerce a "client" into a "user"
+    # TODO: consolidate the client/user distinction.
+    if subject.startswith("user:"):
+        user = get_user(mdb, username=token_data.subject)
+    elif subject.startswith("client:"):
+        # construct a user from the client_id
+        user = get_client_user(mdb, client_id=token_data.subject)
+    else:
+        raise credentials_exception
+    if user is None:
+        raise credentials_exception
+    return user
+
+
+def get_client_user(mdb, client_id: str) -> UserInDB:
+    r"""
+    Returns an ephemeral "user" whose username is the specified `client_id`
+    and whose password is the hashed secret of the client; provided that the
+    specified `client_id` is associated with a site in the database.
+
+    TODO: Clarify the above summary of the function.
+    """
+
+    # Get the site associated with the identified client.
+    site = get_site(mdb, client_id)
+    if site is None:
+        raise credentials_exception
+
+    # Get the client, itself, via the site.
+    client = next(client for client in site.clients if client.id == client_id)
+    if client is None:
+        raise credentials_exception
+
+    # Make an ephemeral "user" whose username matches the client's `id`.
+    user = UserInDB(username=client.id, hashed_password=client.hashed_secret)
+    return user
+
+
+async def get_current_active_user(
+    current_user: UserInDB = Depends(get_current_user),
+) -> UserInDB:
+    r"""
+    Returns the current user, provided their user account is not disabled.
+    """
+
+    if current_user.disabled:
+        raise HTTPException(status_code=400, detail="Inactive user")
+    return current_user