[Data - 3127] Add tag to cohort definitions

charts/d2e-services/templates/atlas-db-init-cm.yaml

@@ -304,3 +304,66 @@ data:
     -- Do not use psql meta-commands (\gset, \if) in this file: trex applies it over
     -- the wire protocol, and the simple query protocol parses the whole file before
     -- executing any of it, so one meta-command stops every statement from running.
+  230_imported_cohort_metadata_tag_group.sql: |
+    -- Tag groups for cohorts imported from external libraries.
+    --
+    -- WebAPI cannot create a root tag through its API (POST /tag requires a parent
+    -- group), so the groups have to be seeded here. Two of them, because their tags
+    -- have opposite cardinality: a cohort's source coexists with other tags, while
+    -- its review status is one-at-a-time. AbstractDaoService.assignTag clears the
+    -- other tags in a single-selection group, which is what retires the previous
+    -- status on re-import. Both groups need allow_custom to accept children, and
+    -- neither is ever assigned to a cohort.
+    DO $$
+    DECLARE
+        source_group_id integer;
+        status_group_id integer;
+    BEGIN
+        -- Group holding the source/provenance tag.
+        SELECT id INTO source_group_id FROM webapi.tag
+        WHERE lower(name) = lower('Imported Cohort Metadata');
+
+        IF source_group_id IS NULL THEN
+            INSERT INTO webapi.tag (name, type, count, show_group, multi_selection,
+                                    permission_protected, mandatory, allow_custom, description)
+            VALUES ('Imported Cohort Metadata', 0, 0, false, true,
+                    false, false, true,
+                    'Container for tags recording where an imported cohort came from')
+            RETURNING id INTO source_group_id;
+            RAISE NOTICE 'Created "Imported Cohort Metadata" tag group (id: %)', source_group_id;
+        ELSE
+            UPDATE webapi.tag
+            SET allow_custom = true, show_group = false, multi_selection = true
+            WHERE id = source_group_id
+              AND (allow_custom IS NOT TRUE OR show_group IS NOT FALSE
+                   OR multi_selection IS NOT TRUE);
+        END IF;
+
+        -- Group holding the mutually exclusive review-status tags.
+        SELECT id INTO status_group_id FROM webapi.tag
+        WHERE lower(name) = lower('Cohort Review Status');
+
+        IF status_group_id IS NULL THEN
+            INSERT INTO webapi.tag (name, type, count, show_group, multi_selection,
+                                    permission_protected, mandatory, allow_custom, description)
+            VALUES ('Cohort Review Status', 0, 0, false, false,
+                    false, false, true,
+                    'Container for the review status of an imported cohort; one applies at a time')
+            RETURNING id INTO status_group_id;
+            RAISE NOTICE 'Created "Cohort Review Status" tag group (id: %)', status_group_id;
+        ELSE
+            UPDATE webapi.tag
+            SET allow_custom = true, show_group = false, multi_selection = false
+            WHERE id = status_group_id
+              AND (allow_custom IS NOT TRUE OR show_group IS NOT FALSE
+                   OR multi_selection IS NOT FALSE);
+        END IF;
+
+    END $$;
+
+    SELECT g.name AS tag_group, g.multi_selection, count(tg.tag_id) AS members
+    FROM webapi.tag g
+    LEFT JOIN webapi.tag_group tg ON tg.group_id = g.id
+    WHERE g.name IN ('Imported Cohort Metadata', 'Cohort Review Status')
+    GROUP BY g.name, g.multi_selection ORDER BY g.name;
+

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/230_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/230_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}"
+            )