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

nmdc_runtime/site/resources.py

@@ -14,12 +14,10 @@ from dagster import (
     StringSource,
     InitResourceContext,
 )
-from fastjsonschema import JsonSchemaValueException
 from frozendict import frozendict
 from linkml_runtime.dumpers import json_dumper
 from pydantic import BaseModel, AnyUrl
 from pymongo import MongoClient, ReplaceOne, InsertOne
-from terminusdb_client import WOQLClient
 from toolz import get_in
 from toolz import merge
 
@@ -28,7 +26,7 @@ from nmdc_runtime.api.models.object import DrsObject, AccessURL, DrsObjectIn
 from nmdc_runtime.api.models.operation import ListOperationsResponse
 from nmdc_runtime.api.models.util import ListRequest
 from nmdc_runtime.site.normalization.gold import normalize_gold_id
-from nmdc_runtime.util import unfreeze, nmdc_jsonschema_validator_noidpatterns
+from nmdc_runtime.util import unfreeze, get_nmdc_schema_validator
 from nmdc_schema import nmdc
 
 
@@ -110,7 +108,7 @@ class RuntimeApiUserClient(RuntimeApiClient):
             },
         )
         response.raise_for_status()
-        return response.json()["cursor"]["firstBatch"]
+        return response.json()["cursor"]["batch"]
 
     def get_omics_processing_records_by_gold_project_id(self, gold_project_id: str):
         gold_project_id = normalize_gold_id(gold_project_id)
@@ -127,19 +125,39 @@ class RuntimeApiUserClient(RuntimeApiClient):
             },
         )
         response.raise_for_status()
-        return response.json()["cursor"]["firstBatch"]
+        return response.json()["cursor"]["batch"]
 
     def get_biosamples_for_study(self, study_id: str):
+        # TODO: 10000 is an arbitrarily large number that has been chosen for the max_page_size param.
+        # The /nmdcschema/{collection-name} endpoint implements pagination via the page_token mechanism,
+        # but the tradeoff there is that we would need to make multiple requests to step through the
+        # each of the pages. By picking a large number for max_page_size, we can get all the results
+        # in a single request.
+        # This method previously used the /queries:run endpoint but the problem with that was that
+        # it used to truncate the number of results returned to 100.
         response = self.request(
-            "POST",
-            f"/queries:run",
+            "GET",
+            f"/nmdcschema/biosample_set",
             {
-                "find": "biosample_set",
-                "filter": {"part_of": {"$elemMatch": {"$eq": study_id}}},
+                "filter": json.dumps({"associated_studies": study_id}),
+                "max_page_size": 10000,
             },
         )
         response.raise_for_status()
-        return response.json()["cursor"]["firstBatch"]
+        return response.json()["resources"]
+
+    def get_data_generation_records_for_study(self, study_id: str):
+        # TODO: same as above, we are using a large max_page_size to avoid pagination.
+        response = self.request(
+            "GET",
+            f"/nmdcschema/data_generation_set",
+            {
+                "filter": json.dumps({"associated_studies": study_id}),
+                "max_page_size": 10000,
+            },
+        )
+        response.raise_for_status()
+        return response.json()["resources"]
 
     def get_omics_processing_by_name(self, name: str):
         response = self.request(
@@ -151,7 +169,19 @@ class RuntimeApiUserClient(RuntimeApiClient):
             },
         )
         response.raise_for_status()
-        return response.json()["cursor"]["firstBatch"]
+        return response.json()["cursor"]["batch"]
+
+    def get_study(self, study_id: str):
+        response = self.request(
+            "POST",
+            f"/queries:run",
+            {
+                "find": "study_set",
+                "filter": {"id": study_id},
+            },
+        )
+        response.raise_for_status()
+        return response.json()["cursor"]["batch"]
 
 
 class RuntimeApiSiteClient(RuntimeApiClient):
@@ -332,9 +362,26 @@ class GoldApiClient(BasicAuthClient):
         """
         return id.replace("gold:", "")
 
-    def fetch_biosamples_by_study(self, study_id: str) -> List[Dict[str, Any]]:
+    def fetch_biosamples_by_study(
+        self, study_id: str, include_project=True
+    ) -> List[Dict[str, Any]]:
         id = self._normalize_id(study_id)
         results = self.request("/biosamples", params={"studyGoldId": id})
+        if include_project:
+            projects = self.fetch_projects_by_study(id)
+            biosamples_by_id = {
+                biosample["biosampleGoldId"]: biosample for biosample in results
+            }
+            for project in projects:
+                sample_id = project.get("biosampleGoldId")
+                if not sample_id:
+                    continue
+                if sample_id not in biosamples_by_id:
+                    continue
+                biosample = biosamples_by_id[sample_id]
+                if "projects" not in biosample:
+                    biosample["projects"] = []
+                biosample["projects"].append(project)
         return results
 
     def fetch_projects_by_study(self, study_id: str) -> List[Dict[str, Any]]:
@@ -354,6 +401,18 @@ class GoldApiClient(BasicAuthClient):
             return None
         return results[0]
 
+    def fetch_projects_by_biosample(self, biosample_id: str) -> List[Dict[str, Any]]:
+        id = self._normalize_id(biosample_id)
+        results = self.request("/projects", params={"biosampleGoldId": id})
+        return results
+
+    def fetch_biosample_by_biosample_id(
+        self, biosample_id: str
+    ) -> List[Dict[str, Any]]:
+        id = self._normalize_id(biosample_id)
+        results = self.request("/biosamples", params={"biosampleGoldId": id})
+        return results
+
 
 @resource(
     config_schema={
@@ -372,15 +431,47 @@ def gold_api_client_resource(context: InitResourceContext):
 
 @dataclass
 class NmdcPortalApiClient:
+
     base_url: str
-    # Using a cookie for authentication is not ideal and should be replaced
-    # when this API has an another authentication method
-    session_cookie: str
+    refresh_token: str
+    access_token: Optional[str] = None
+    access_token_expires_at: Optional[datetime] = None
+
+    def _request(self, method: str, endpoint: str, **kwargs):
+        r"""
+        Submits a request to the specified API endpoint;
+        after refreshing the access token, if necessary.
+        """
+        if self.access_token is None or datetime.now() > self.access_token_expires_at:
+            refresh_response = requests.post(
+                f"{self.base_url}/auth/refresh",
+                json={"refresh_token": self.refresh_token},
+            )
+            refresh_response.raise_for_status()
+            refresh_body = refresh_response.json()
+            self.access_token_expires_at = datetime.now() + timedelta(
+                seconds=refresh_body["expires_in"]
+            )
+            self.access_token = refresh_body["access_token"]
+
+        headers = kwargs.get("headers", {})
+        headers["Authorization"] = f"Bearer {self.access_token}"
+        return requests.request(
+            method, f"{self.base_url}{endpoint}", **kwargs, headers=headers
+        )
 
     def fetch_metadata_submission(self, id: str) -> Dict[str, Any]:
-        response = requests.get(
-            f"{self.base_url}/api/metadata_submission/{id}",
-            cookies={"session": self.session_cookie},
+        response = self._request("GET", f"/api/metadata_submission/{id}")
+        response.raise_for_status()
+        return response.json()
+
+    def make_submission_images_public(
+        self, submission_id: str, *, study_id: str
+    ) -> Dict[str, Any]:
+        response = self._request(
+            "POST",
+            f"/api/metadata_submission/{submission_id}/image/make_public",
+            json={"study_id": study_id},
         )
         response.raise_for_status()
         return response.json()
@@ -389,13 +480,13 @@ class NmdcPortalApiClient:
 @resource(
     config_schema={
         "base_url": StringSource,
-        "session_cookie": StringSource,
+        "refresh_token": StringSource,
     }
 )
 def nmdc_portal_api_client_resource(context: InitResourceContext):
     return NmdcPortalApiClient(
         base_url=context.resource_config["base_url"],
-        session_cookie=context.resource_config["session_cookie"],
+        refresh_token=context.resource_config["refresh_token"],
     )
 
 
@@ -439,36 +530,49 @@ class MongoDB:
         self.db = self.client[dbname]
 
     def add_docs(self, docs, validate=True, replace=True):
-        try:
-            if validate:
-                nmdc_jsonschema_validator_noidpatterns(docs)
-            rv = {}
-            for collection_name, docs in docs.items():
-                rv[collection_name] = self.db[collection_name].bulk_write(
-                    [
-                        (
-                            ReplaceOne({"id": d["id"]}, d, upsert=True)
-                            if replace
-                            else InsertOne(d)
-                        )
-                        for d in docs
-                    ]
-                )
-                now = datetime.now(timezone.utc)
-                self.db.txn_log.insert_many(
-                    [
-                        {
-                            "tgt": {"id": d.get("id"), "c": collection_name},
-                            "type": "upsert",
-                            "ts": now,
-                            # "dtl": {},
-                        }
-                        for d in docs
-                    ]
-                )
-            return rv
-        except JsonSchemaValueException as e:
-            raise ValueError(e.message)
+        """
+        TODO: Document this function.
+        """
+        if validate:
+            validator = get_nmdc_schema_validator()
+            # Fail fast on first validation error.
+            for result in validator.iter_results(docs, target_class="Database"):
+                raise ValueError(result.message)
+        rv = {}
+        for collection_name, collection_docs in docs.items():
+            # If `collection_docs` is empty, abort this iteration.
+            #
+            # Note: We do this because the `bulk_write` method called below will raise
+            #       an `InvalidOperation` exception if it is passed 0 operations.
+            #
+            # Reference: https://pymongo.readthedocs.io/en/stable/api/pymongo/collection.html#pymongo.collection.Collection.bulk_write
+            #
+            if len(collection_docs) == 0:
+                continue
+
+            rv[collection_name] = self.db[collection_name].bulk_write(
+                [
+                    (
+                        ReplaceOne({"id": d["id"]}, d, upsert=True)
+                        if replace
+                        else InsertOne(d)
+                    )
+                    for d in collection_docs
+                ]
+            )
+            now = datetime.now(timezone.utc)
+            self.db.txn_log.insert_many(
+                [
+                    {
+                        "tgt": {"id": d.get("id"), "c": collection_name},
+                        "type": "upsert",
+                        "ts": now,
+                        # "dtl": {},
+                    }
+                    for d in collection_docs
+                ]
+            )
+        return rv
 
 
 @resource(
@@ -512,33 +616,3 @@ def get_mongo(run_config: frozendict):
         )
     )
     return mongo_resource(resource_context)
-
-
-class TerminusDB:
-    def __init__(self, server_url, user, key, account, dbid):
-        self.client = WOQLClient(server_url=server_url)
-        self.client.connect(user=user, key=key, account=account)
-        db_info = self.client.get_database(dbid=dbid, account=account)
-        if db_info is None:
-            self.client.create_database(dbid=dbid, accountid=account, label=dbid)
-            self.client.create_graph(graph_type="inference", graph_id="main")
-        self.client.connect(user=user, key=key, account=account, db=dbid)
-
-
-@resource(
-    config_schema={
-        "server_url": StringSource,
-        "user": StringSource,
-        "key": StringSource,
-        "account": StringSource,
-        "dbid": StringSource,
-    }
-)
-def terminus_resource(context):
-    return TerminusDB(
-        server_url=context.resource_config["server_url"],
-        user=context.resource_config["user"],
-        key=context.resource_config["key"],
-        account=context.resource_config["account"],
-        dbid=context.resource_config["dbid"],
-    )