[Data - 3127] Add tag to cohort definitions

plugins/flows/_shared_flow_utils/api/PhenotypeAPI.py

@@ -10,42 +10,84 @@ def __init__(self):
         self.url = self.get_service_route("d2e-webapi")  
         self.cohort_definition_url = self.url + 'cohortdefinition'
         self.headers = self.get_options()
+
+    def get_cohort_name_index(self, dataset_id: str) -> dict:
+        headers = self.headers.copy()
+        headers["datasetId"] = dataset_id
+
+        response = requests.get(
+            self.cohort_definition_url,
+            headers=headers,
+            verify=self.get_verify_value()
+        )
+
+        if response.status_code != 200:
+            raise Exception(
+                f"Failed to list cohort definitions: "
+                f"{response.status_code} - {response.text}"
+            )
+
+        return {
+            cohort["name"]: cohort["id"]
+            for cohort in response.json()
+            if cohort.get("name") is not None
+        }
 
-    def create_single_cohort_definition(self, cohort_def: json, dataset_id: str, user_name: str):
+    def create_single_cohort_definition(self, cohort_def: dict, dataset_id: str, user_name: str,
+                                        name_index: dict):
         logger = get_run_logger()
         current_time = int(time.time() * 1000)
         headers = self.headers.copy()
         headers["datasetId"] = dataset_id
 
         # Parse the JSON expression
         expression = json.loads(cohort_def['json'])
+        name = f"{cohort_def['cohortId']}_{cohort_def['cohortName']}"
+        existing_id = name_index.get(name)
         payload = {
-            "id": cohort_def['cohortId'],
-            "name": f"{cohort_def['cohortId']}_{cohort_def['cohortName']}",
+            "id": 0,
+            "name": name,
             "description": f"Phenotype Library cohort: {cohort_def['cohortName']}",
             "expressionType": "SIMPLE_EXPRESSION",
             "expression": expression,
             "createdBy": user_name,
             "createdDate": current_time,
             "modifiedBy": user_name,
             "modifiedDate": current_time,
-            "datasetId": dataset_id,
-            "tags": ["phenotype_library"],
+            "tags": [],
         }
-
-        logger.info(f"Creating cohort: {cohort_def['cohortName']} (ID: {cohort_def['cohortId']})")
-        response = requests.post(
-            self.cohort_definition_url, 
-            headers=headers, 
-            json=payload, 
-            verify=self.get_verify_value()
-        )
-
+        # datasetId travels in the header, not the body.
+
+        verb = "Updating" if existing_id else "Creating"
+        logger.info(f"{verb} cohort: {cohort_def['cohortName']} (ID: {cohort_def['cohortId']})")
+
+        if existing_id:
+            response = requests.put(
+                f"{self.cohort_definition_url}/{existing_id}",
+                headers=headers,
+                json=payload,
+                verify=self.get_verify_value()
+            )
+        else:
+            response = requests.post(
+                self.cohort_definition_url,
+                headers=headers,
+                json=payload,
+                verify=self.get_verify_value()
+            )
+
         if response.status_code in [200, 201]:
             result = response.json()
-            logger.info(f"Successfully created cohort {cohort_def['cohortId']}")
+            name_index[name] = result["id"]
+            logger.info(
+                f"{'Updated' if existing_id else 'Created'} WebAPI cohort definition "
+                f"{result["id"]} for phenotype cohort {cohort_def['cohortId']}"
+            )
             return result
-        else:
-            error_msg = f"Failed to create cohort {cohort_def['cohortId']}: {response.status_code} - {response.text}"
-            logger.error(error_msg)
-            raise Exception(error_msg)
+
+        error_msg = (
+            f"Failed to {'update' if existing_id else 'create'} cohort "
+            f"{cohort_def['cohortId']}: {response.status_code} - {response.text}"
+        )
+        logger.error(error_msg)
+        raise Exception(error_msg)

plugins/flows/_shared_flow_utils/api/PhenotypeTagAPI.py

@@ -0,0 +1,149 @@
+import requests
+from prefect.logging import get_run_logger
+
+from _shared_flow_utils.api.BaseAPI import BaseAPI
+
+
+class PhenotypeTagAPI(BaseAPI):
+    """Resolve the WebAPI tags applied to Phenotype Library cohort imports.
+
+    Each cohort carries two: its source, and its review status. The groups holding
+    them are seeded in SQL and never assigned to a cohort themselves.
+    """
+
+    # Seeded by services/atlas-db-init/220_imported_cohort_metadata_tag_group.sql.
+    # Two groups because multi_selection is a property of the group: statuses are mutually exclusive, the source tag is not.
+    REQUEST_TIMEOUT = (10, 30)
+
+    SOURCE_GROUP = "Imported Cohort Metadata"
+    STATUS_GROUP = "Cohort Review Status"
+    PHENOTYPE_LIBRARY_TAG = "Phenotype Library"
+
+    def __init__(self):
+        super().__init__()
+        # The d2e-webapi plugin exposes no /tag routes, so tag work goes straight
+        # to WebAPI through the d2e-compat shim, which performs the same token
+        # exchange the plugin's own calls trigger.
+        self.tag_url = self.get_service_route("webapi") + "tag"
+        self.headers = self.get_options()
+
+    def _get_tags(self, dataset_id: str) -> list:
+        headers = self.headers.copy()
+        headers["datasetId"] = dataset_id
+        response = requests.get(
+            self.tag_url, headers=headers, verify=self.get_verify_value(),
+            timeout=self.REQUEST_TIMEOUT
+        )
+        if response.status_code != 200:
+            raise Exception(
+                f"Failed to list tags: {response.status_code} - {response.text}"
+            )
+        return response.json()
+
+    def _create_tag(self, dataset_id: str, name: str, group_id: int) -> dict:
+        """Create a tag inside one of the seeded groups.
+
+        groups must be non-empty: WebAPI 400s on an empty list and 500s if the
+        field is absent, which is why groups cannot be created through the API.
+        """
+        headers = self.headers.copy()
+        headers["datasetId"] = dataset_id
+        payload = {
+            "name": name,
+            # The nested "groups": [] is required, not noise. TagDTOToTagConverter
+            # converts each group reference recursively through itself, and that
+            # second pass dereferences source.getGroups() -- a group reference
+            # without the field NPEs into a 500 ConversionFailedException.
+            "groups": [{"id": group_id, "groups": []}],
+            "allowCustom": False,
+            "showGroup": False,
+            "multiSelection": False,
+            "permissionProtected": False,
+            "mandatory": False,
+        }
+        response = requests.post(
+            self.tag_url,
+            headers=headers,
+            json=payload,
+            verify=self.get_verify_value(),
+            timeout=self.REQUEST_TIMEOUT,
+        )
+        if response.status_code not in [200, 201]:
+            raise Exception(
+                f"Failed to create tag '{name}': "
+                f"{response.status_code} - {response.text}"
+            )
+        return response.json()
+
+    def _require_group(self, by_name: dict, group_name: str) -> dict:
+        """Look up a seeded tag group, failing with something actionable."""
+        group = by_name.get(group_name.lower())
+        if group is None:
+            # Not created here: WebAPI refuses to create a tag with no parent
+            # group, so a root group cannot come from the API at all.
+            raise Exception(
+                f"Tag group '{group_name}' does not exist in WebAPI. It cannot be "
+                f"created through the API; apply "
+                f"services/atlas-db-init/220_imported_cohort_metadata_tag_group.sql "
+                f"(restarting the webapi-init container does this) and retry."
+            )
+        if not group.get("allowCustom", False):
+            raise Exception(
+                f"Tag group '{group_name}' (id {group['id']}) has allowCustom "
+                f"disabled, so tags cannot be attached to it."
+            )
+        return group
+
+    def resolve_import_tags(self, dataset_id: str, statuses) -> tuple[dict, dict]:
+        """Return the provenance tag and a {status: tag} map, creating what is missing.
+
+        Tag names are unique case-insensitively (tags_name_idx on lower(name)), so
+        an existing tag is reused rather than recreated.
+        """
+        logger = get_run_logger()
+        by_name = {tag["name"].lower(): tag for tag in self._get_tags(dataset_id)}
+
+        source_group = self._require_group(by_name, self.SOURCE_GROUP)
+        status_group = self._require_group(by_name, self.STATUS_GROUP)
+
+        wanted = [(self.PHENOTYPE_LIBRARY_TAG, source_group)]
+        wanted += [(status, status_group) for status in sorted(set(statuses))]
+
+        for name, group in wanted:
+            if name.lower() not in by_name:
+                by_name[name.lower()] = self._create_tag(dataset_id, name, group["id"])
+                logger.info(f"Created cohort tag '{name}' under '{group['name']}'")
+
+        status_tags = {
+            status: by_name[status.lower()] for status in sorted(set(statuses))
+        }
+        return by_name[self.PHENOTYPE_LIBRARY_TAG.lower()], status_tags
+
+    def assign_tags_to_cohorts(self, dataset_id: str, tag_ids: list, cohort_ids: list) -> None:
+        """Attach a set of tags to a set of cohort definitions in one request.
+
+        Not the per-cohort POST /cohortdefinition/{id}/tag: that takes one tag id,
+        so a full import costs ~2200 requests and exceeds Trex's rate limit of
+        5000 per 15 minutes. WebAPI routes each pair through the same assignTag,
+        so re-assigning stays a no-op and the single-selection swap still applies.
+        """
+        if not tag_ids or not cohort_ids:
+            return
+        headers = self.headers.copy()
+        headers["datasetId"] = dataset_id
+        payload = {
+            "tags": list(tag_ids),
+            "assets": {"cohorts": list(cohort_ids)},
+        }
+        response = requests.post(
+            f"{self.tag_url}/multiAssign",
+            headers=headers,
+            json=payload,
+            verify=self.get_verify_value(),
+            timeout=self.REQUEST_TIMEOUT,
+        )
+        if response.status_code not in [200, 201, 204]:
+            raise Exception(
+                f"Failed to assign tags {list(tag_ids)} to {len(cohort_ids)} cohort "
+                f"definitions: {response.status_code} - {response.text}"
+            )

plugins/flows/base/phenotype_plugin/flow.py

@@ -1,6 +1,5 @@
 import os, logging
 from rpy2.rinterface_lib.callbacks import logger as rpy2_logger
-from rpy2.robjects import pandas2ri, numpy2ri
 from rpy2 import robjects
 
 from prefect import flow, task
@@ -10,6 +9,7 @@
 from _shared_flow_utils.types import UserType
 from _shared_flow_utils.dao.DBDao import DBDao
 from _shared_flow_utils.api.PhenotypeAPI import PhenotypeAPI
+from _shared_flow_utils.api.PhenotypeTagAPI import PhenotypeTagAPI
 
 os.environ["plugin_name"] = "phenotype_plugin"
 
@@ -60,11 +60,9 @@ def get_cohort_definitions(cohorts_id: str, vocabschema_name: str, materialize:
     Returns:
         list: List of cohort definitions with cohortId, cohortName, json, and sql.
     """
-    pandas2ri.activate()
-    numpy2ri.activate()
     r_script_path = os.path.join(os.path.dirname(__file__), 'get_cohort_definitions.R')
 
-    with robjects.conversion.localconverter(robjects.default_converter):
+    with robjects.default_converter.context():
 
         # Source the R script to load the function
         robjects.r(f'source("{r_script_path}")')
@@ -86,6 +84,7 @@ def get_cohort_definitions(cohorts_id: str, vocabschema_name: str, materialize:
                     "cohortName": str(result[i].rx2("cohortName")[0]),
                     "json": str(result[i].rx2("json")[0]),
                     "sql": str(result[i].rx2("sql")[0]),
+                    "status": str(result[i].rx2("status")[0]),
                 }
                 cohort_definitions.append(cohort_def)
             return cohort_definitions
@@ -109,23 +108,71 @@ def atlas_cohort_definitions(
     """
     logger = get_run_logger()
     phenotype_api = PhenotypeAPI()
+    phenotype_tag_api = PhenotypeTagAPI()
     created_cohorts = []
+    name_index = phenotype_api.get_cohort_name_index(dataset_id)
+    logger.info(f"Indexed {len(name_index)} existing WebAPI cohort definitions")
+
+    # get_cohort_definitions.R already maps blank/NA to "Unspecified".
+    statuses = {cohort_def["status"] for cohort_def in cohort_definitions}
+    phenotype_library_tag, status_tags = phenotype_tag_api.resolve_import_tags(
+        dataset_id, statuses
+    )
+
+    # Collected while writing, applied in bulk afterwards -- see tag_cohort_definitions.
+    cohort_ids_by_status = {}
 
     for cohort_def in cohort_definitions:
         try:
+            status = cohort_def["status"]
             result = phenotype_api.create_single_cohort_definition(
-                cohort_def, dataset_id, user_name
+                cohort_def,
+                dataset_id,
+                user_name,
+                name_index,
             )
             created_cohorts.append(result)
+            cohort_ids_by_status.setdefault(status, []).append(result["id"])
         except Exception as e:
             error_message = (
-                f"Failed to create cohort {cohort_def['cohortId']}: {str(e)}"
+                f"Failed to save cohort {cohort_def['cohortId']}: {str(e)}. "
+                f"{len(created_cohorts)} cohorts were written before this and are "
+                f"still untagged; re-run the flow to finish them."
             )
             logger.error(error_message)
             raise Exception(error_message) from e
+
+    tag_cohort_definitions(
+        phenotype_tag_api, dataset_id, phenotype_library_tag, status_tags,
+        cohort_ids_by_status,
+    )
     return created_cohorts
 
 
+def tag_cohort_definitions(phenotype_tag_api, dataset_id: str, phenotype_library_tag: dict,
+                           status_tags: dict, cohort_ids_by_status: dict) -> None:
+    """Record where the cohorts came from, and what their review status is.
+    """
+    logger = get_run_logger()
+    all_cohort_ids = [i for ids in cohort_ids_by_status.values() for i in ids]
+    if not all_cohort_ids:
+        return
+
+    phenotype_tag_api.assign_tags_to_cohorts(
+        dataset_id, [phenotype_library_tag["id"]], all_cohort_ids
+    )
+    logger.info(
+        f"Tagged {len(all_cohort_ids)} cohort definitions as "
+        f"'{phenotype_library_tag['name']}'"
+    )
+
+    for status, cohort_ids in sorted(cohort_ids_by_status.items()):
+        phenotype_tag_api.assign_tags_to_cohorts(
+            dataset_id, [status_tags[status]["id"]], cohort_ids
+        )
+        logger.info(f"Tagged {len(cohort_ids)} cohort definitions as '{status}'")
+
+
 @task(log_prints=True)
 def materialize_cohort_definitions(
     dbdao: DBDao,