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

nmdc_runtime/api/models/nmdc_schema.py

@@ -0,0 +1,146 @@
+import inspect
+from enum import Enum
+from typing import List, Any, Dict, Optional
+
+from pydantic import BaseModel, Field, create_model
+from refscan.lib.helpers import get_collection_names_from_schema
+
+from nmdc_runtime.util import nmdc_schema_view
+
+
+class FileTypeEnum(str, Enum):
+    ft_icr_ms_analysis_results = "FT ICR-MS Analysis Results"
+    gc_ms_metabolomics_results = "GC-MS Metabolomics Results"
+    metaproteomics_workflow_statistics = "Metaproteomics Workflow Statistics"
+    protein_report = "Protein Report"
+    peptide_report = "Peptide Report"
+    unfiltered_metaproteomics_results = "Unfiltered Metaproteomics Results"
+    read_count_and_rpkm = "Read Count and RPKM"
+    qc_non_rrna_r2 = "QC non-rRNA R2"
+    qc_non_rrna_r1 = "QC non-rRNA R1"
+    metagenome_bins = "Metagenome Bins"
+    checkm_statistics = "CheckM Statistics"
+    gottcha2_krona_plot = "GOTTCHA2 Krona Plot"
+    kraken2_krona_plot = "Kraken2 Krona Plot"
+    centrifuge_krona_plot = "Centrifuge Krona Plot"
+    kraken2_classification_report = "Kraken2 Classification Report"
+    kraken2_taxonomic_classification = "Kraken2 Taxonomic Classification"
+    centrifuge_classification_report = "Centrifuge Classification Report"
+    centrifuge_taxonomic_classification = "Centrifuge Taxonomic Classification"
+    structural_annotation_gff = "Structural Annotation GFF"
+    functional_annotation_gff = "Functional Annotation GFF"
+    annotation_amino_acid_fasta = "Annotation Amino Acid FASTA"
+    annotation_enzyme_commission = "Annotation Enzyme Commission"
+    annotation_kegg_orthology = "Annotation KEGG Orthology"
+    assembly_coverage_bam = "Assembly Coverage BAM"
+    assembly_agp = "Assembly AGP"
+    assembly_scaffolds = "Assembly Scaffolds"
+    assembly_contigs = "Assembly Contigs"
+    assembly_coverage_stats = "Assembly Coverage Stats"
+    filtered_sequencing_reads = "Filtered Sequencing Reads"
+    qc_statistics = "QC Statistics"
+    tigrfam_annotation_gff = "TIGRFam Annotation GFF"
+    clusters_of_orthologous_groups__cog__annotation_gff = (
+        "Clusters of Orthologous Groups (COG) Annotation GFF"
+    )
+    cath_funfams__functional_families__annotation_gff = (
+        "CATH FunFams (Functional Families) Annotation GFF"
+    )
+    superfam_annotation_gff = "SUPERFam Annotation GFF"
+    smart_annotation_gff = "SMART Annotation GFF"
+    pfam_annotation_gff = "Pfam Annotation GFF"
+    direct_infusion_ft_icr_ms_raw_data = "Direct Infusion FT ICR-MS Raw Data"
+
+
+class DataObject(BaseModel):
+    id: str = Field(None)
+    name: str = Field(None, description="A human readable label for an entity")
+    description: str = Field(
+        None, description="a human-readable description of a thing"
+    )
+    alternative_identifiers: List[str] = Field(
+        None, description="A list of alternative identifiers for the entity."
+    )
+    compression_type: str = Field(
+        None, description="If provided, specifies the compression type"
+    )
+    data_object_type: FileTypeEnum = Field(None)
+    file_size_bytes: int = Field(None, description="Size of the file in bytes")
+    md5_checksum: str = Field(None, description="MD5 checksum of file (pre-compressed)")
+    type: str = Field(
+        "nmdc:DataObject",
+        description="An optional string that specifies the type object.  This is used to allow for searches for different kinds of objects.",
+    )
+    url: str = Field(None)
+    was_generated_by: str = Field(None)
+
+
+list_request_ops_per_field_type = {
+    (object,): ["$eq", "$neq", "$in", "$nin"],
+    (str,): ["$regex"],
+    (int, float): ["$gt", "$gte", "$lt", "$lte"],
+}
+
+list_request_ops_with_many_args = {"$in", "$nin"}
+
+
+def create_list_request_model_for(cls):
+    field_filters = []
+    sig = inspect.signature(cls)
+    for p in sig.parameters.values():
+        field_name, field_type = p.name, p.annotation
+        if hasattr(field_type, "__origin__"):  # a GenericAlias object, e.g. List[str].
+            field_type = field_type.__args__[0]
+        field_default = (
+            None if p.default in (inspect.Parameter.empty, []) else p.default
+        )
+        field_ops = []
+        for types_ok, type_ops in list_request_ops_per_field_type.items():
+            if field_type in types_ok or isinstance(field_type, types_ok):
+                field_ops.extend(type_ops)
+        for op in field_ops:
+            field_filters.append(
+                {
+                    "name": field_name,
+                    "arg_type": field_type,
+                    "default": field_default,
+                    "op": op,
+                }
+            )
+    create_model_kwargs = {"max_page_size": (int, 20), "page_token": (str, None)}
+    for ff in field_filters:
+        model_field_name = f'{ff["name"]}_{ff["op"][1:]}'
+        model_field_type = ff["arg_type"]
+        create_model_kwargs[model_field_name] = (model_field_type, ff["default"])
+    return create_model(f"{cls.__name__}ListRequest", **create_model_kwargs)
+
+
+def list_request_filter_to_mongo_filter(req: dict):
+    filter_ = {}
+    for k, v in req.items():
+        if not v:
+            continue
+        field, op = k.rsplit("_", maxsplit=1)
+        op = f"${op}"
+        if field not in filter_:
+            filter_[field] = {}
+        if op not in filter_[field]:
+            if op in list_request_ops_with_many_args:
+                filter_[field][op] = v.split(",")
+            else:
+                filter_[field][op] = v
+    return filter_
+
+
+DataObjectListRequest = create_list_request_model_for(DataObject)
+
+SimplifiedDocument = Dict[str, Any]
+
+schema_view = nmdc_schema_view()
+SimplifiedNMDCDatabase = create_model(
+    "NMDCDatabase",
+    **{
+        coll_name: Optional[list[SimplifiedDocument]]
+        for coll_name in get_collection_names_from_schema(schema_view)
+    },
+)

nmdc_runtime/api/models/object.py

@@ -0,0 +1,180 @@
+import datetime
+import hashlib
+import http
+from enum import Enum
+from typing import Optional, List, Dict
+
+from pydantic import (
+    field_validator,
+    model_validator,
+    Field,
+    StringConstraints,
+    BaseModel,
+    AnyUrl,
+    HttpUrl,
+    field_serializer,
+)
+from typing_extensions import Annotated
+
+
+class AccessMethodType(str, Enum):
+    s3 = "s3"
+    gs = "gs"
+    ftp = "ftp"
+    gsiftp = "gsiftp"
+    globus = "globus"
+    htsget = "htsget"
+    https = "https"
+    file = "file"
+
+
+class AccessURL(BaseModel):
+    headers: Optional[Dict[str, str]] = None
+    url: AnyUrl
+
+    @field_serializer("url")
+    def serialize_url(self, url: AnyUrl, _info):
+        return str(url)
+
+
+class AccessMethod(BaseModel):
+    access_id: Optional[Annotated[str, StringConstraints(min_length=1)]] = None
+    access_url: Optional[AccessURL] = None
+    region: Optional[str] = None
+    type: AccessMethodType = AccessMethodType.https
+
+    @model_validator(mode="before")
+    def at_least_one_of_access_id_and_url(cls, values):
+        access_id, access_url = values.get("access_id"), values.get("access_url")
+        if access_id is None and access_url is None:
+            raise ValueError(
+                "At least one of access_url and access_id must be provided."
+            )
+        return values
+
+
+ChecksumType = Annotated[
+    str,
+    StringConstraints(
+        pattern=rf"(?P<checksumtype>({'|'.join(sorted(hashlib.algorithms_guaranteed))}))"
+    ),
+]
+
+
+class Checksum(BaseModel):
+    checksum: Annotated[str, StringConstraints(min_length=1)]
+    type: ChecksumType
+
+
+DrsId = Annotated[str, StringConstraints(pattern=r"^[A-Za-z0-9._~\-]+$")]
+PortableFilename = Annotated[str, StringConstraints(pattern=r"^[A-Za-z0-9._\-]+$")]
+
+
+class ContentsObject(BaseModel):
+    contents: Optional[List["ContentsObject"]] = None
+    drs_uri: Optional[List[AnyUrl]] = None
+    id: Optional[DrsId] = None
+    name: PortableFilename
+
+    @model_validator(mode="before")
+    def no_contents_means_single_blob(cls, values):
+        contents, id_ = values.get("contents"), values.get("id")
+        if contents is None and id_ is None:
+            raise ValueError("no contents means no further nesting, so id required")
+        return values
+
+    @field_serializer("drs_uri")
+    def serialize_url(self, drs_uri: Optional[List[AnyUrl]], _info):
+        if drs_uri is not None and len(drs_uri) > 0:
+            return [str(u) for u in drs_uri]
+        return drs_uri
+
+
+# Note: Between Pydantic v1 and v2, the `update_forward_refs` method was renamed to `model_rebuild`.
+#       Reference: https://docs.pydantic.dev/2.11/migration/#changes-to-pydanticbasemodel
+ContentsObject.model_rebuild()
+
+Mimetype = Annotated[str, StringConstraints(pattern=r"^\w+/[-+.\w]+$")]
+SizeInBytes = Annotated[int, Field(strict=True, ge=0)]
+
+
+class Error(BaseModel):
+    msg: Optional[str] = None
+    status_code: http.HTTPStatus
+
+
+class DrsObjectBase(BaseModel):
+    aliases: Optional[List[str]] = None
+    description: Optional[str] = None
+    mime_type: Optional[Mimetype] = None
+    name: Optional[PortableFilename] = None
+
+
+class DrsObjectIn(DrsObjectBase):
+    access_methods: Optional[List[AccessMethod]] = None
+    checksums: List[Checksum]
+    contents: Optional[List[ContentsObject]] = None
+    created_time: datetime.datetime
+    size: SizeInBytes
+    updated_time: Optional[datetime.datetime] = None
+    version: Optional[str] = None
+
+    @model_validator(mode="before")
+    def no_contents_means_single_blob(cls, values):
+        contents, access_methods = values.get("contents"), values.get("access_methods")
+        if contents is None and access_methods is None:
+            raise ValueError(
+                "no contents means single blob, which requires access_methods"
+            )
+        return values
+
+    @field_validator("checksums")
+    @classmethod
+    def at_least_one_checksum(cls, v):
+        if not len(v) >= 1:
+            raise ValueError("At least one checksum requried")
+        return v
+
+
+class DrsObject(DrsObjectIn):
+    id: DrsId
+    self_uri: AnyUrl
+
+    @field_serializer("self_uri")
+    def serialize_url(self, self_uri: AnyUrl, _info):
+        return str(self_uri)
+
+
+Seconds = Annotated[int, Field(strict=True, gt=0)]
+
+
+class ObjectPresignedUrl(BaseModel):
+    url: HttpUrl
+    expires_in: Seconds = 300
+
+    @field_serializer("url")
+    def serialize_url(self, url: HttpUrl, _info):
+        return str(url)
+
+
+class DrsObjectOutBase(DrsObjectBase):
+    checksums: List[Checksum]
+    created_time: datetime.datetime
+    id: DrsId
+    self_uri: AnyUrl
+    size: SizeInBytes
+    updated_time: Optional[datetime.datetime] = None
+    version: Optional[str] = None
+
+    @field_serializer("self_uri")
+    def serialize_url(self, self_uri: AnyUrl, _info):
+        return str(self_uri)
+
+
+class DrsObjectBlobOut(DrsObjectOutBase):
+    access_methods: List[AccessMethod]
+
+
+class DrsObjectBundleOut(DrsObjectOutBase):
+    access_methods: Optional[List[AccessMethod]] = None
+    contents: List[ContentsObject]

nmdc_runtime/api/models/object_type.py

@@ -0,0 +1,20 @@
+import datetime
+from typing import Optional, List
+
+from pydantic import BaseModel
+
+from nmdc_runtime.api.models.object import DrsObject
+
+
+class ObjectTypeBase(BaseModel):
+    name: Optional[str] = None
+    description: Optional[str] = None
+
+
+class ObjectType(ObjectTypeBase):
+    id: str
+    created_at: datetime.datetime
+
+
+class DrsObjectWithTypes(DrsObject):
+    types: Optional[List[str]] = None

nmdc_runtime/api/models/operation.py

@@ -0,0 +1,66 @@
+import datetime
+from typing import Generic, TypeVar, Optional, List, Any, Union
+
+from pydantic import StringConstraints, BaseModel, HttpUrl, field_serializer
+
+from nmdc_runtime.api.models.util import ResultT
+from typing_extensions import Annotated
+
+MetadataT = TypeVar("MetadataT")
+
+
+PythonImportPath = Annotated[str, StringConstraints(pattern=r"^[A-Za-z0-9_.]+$")]
+
+
+class OperationError(BaseModel):
+    code: str
+    message: str
+    details: Any = None
+
+
+class Operation(BaseModel, Generic[ResultT, MetadataT]):
+    id: str
+    done: bool = False
+    expire_time: datetime.datetime
+    result: Optional[Union[ResultT, OperationError]] = None
+    metadata: Optional[MetadataT] = None
+
+
+class UpdateOperationRequest(BaseModel, Generic[ResultT, MetadataT]):
+    done: bool = False
+    result: Optional[Union[ResultT, OperationError]] = None
+    metadata: Optional[MetadataT] = {}
+
+
+class ListOperationsResponse(BaseModel, Generic[ResultT, MetadataT]):
+    resources: List[Operation[ResultT, MetadataT]]
+    next_page_token: Optional[str] = None
+
+
+class Result(BaseModel):
+    model: Optional[PythonImportPath] = None
+
+
+class EmptyResult(Result):
+    pass
+
+
+class Metadata(BaseModel):
+    # XXX alternative: set model field using __class__ on __init__()?
+    model: Optional[PythonImportPath] = None
+    cancelled: Optional[bool] = None
+
+
+class PausedOrNot(Metadata):
+    paused: bool
+
+
+class ObjectPutMetadata(Metadata):
+    object_id: str
+    site_id: str
+    url: HttpUrl
+    expires_in_seconds: int
+
+    @field_serializer("url")
+    def serialize_url(self, url: HttpUrl, _info):
+        return str(url)

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_)