From 7bf0565bf3d80c721d5261d31386114f5590a40c Mon Sep 17 00:00:00 2001
From: Heather Ward <heather@dnastack.com>
Date: Thu, 18 May 2023 16:51:46 -0400
Subject: [PATCH] Only delete custom workflows registered by wdl-ci

Do not delete other custom workflows in the user's namespace.
---
 src/wdlci/cli/cleanup.py                      |  7 +++---
 src/wdlci/cli/submit.py                       |  9 +++++--
 src/wdlci/model/submission_state.py           | 15 ++++++++---
 .../workbench/workflow_service_client.py      | 25 ++++++++++++++++++-
 4 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/src/wdlci/cli/cleanup.py b/src/wdlci/cli/cleanup.py
index 3e9ebea..1190eb6 100644
--- a/src/wdlci/cli/cleanup.py
+++ b/src/wdlci/cli/cleanup.py
@@ -1,3 +1,4 @@
+import jsonpickle
 import os
 import sys
 from wdlci.auth.refresh_token_auth import RefreshTokenAuth
@@ -40,9 +41,9 @@ def cleanup_handler(kwargs):
         )
         workflow_service_client = WorkflowServiceClient(workflow_service_auth)
 
-        # purge all custom workflows in the namespace
-        # TODO maybe add option here to purge all custom workflows, or only those registered in state file
-        workflow_service_client.purge_custom_workflows()
+        submission_state = jsonpickle.decode(open(SUBMISSION_JSON, "r").read())
+        for workflow in submission_state.workflows.values():
+            workflow_service_client.delete_custom_workflow(workflow)
 
         print("Cleanup complete. Custom workflows purged from namespace")
 
diff --git a/src/wdlci/cli/submit.py b/src/wdlci/cli/submit.py
index 794f97e..cec496f 100644
--- a/src/wdlci/cli/submit.py
+++ b/src/wdlci/cli/submit.py
@@ -116,12 +116,17 @@ def submit_handler(kwargs):
                             )
                             sys.exit(e.exit_code)
 
-                        workflow_id = workflow_service_client.register_workflow(
+                        (
+                            workflow_id,
+                            workflow_etag,
+                        ) = workflow_service_client.register_workflow(
                             test_key,
                             workflow_config,
                             transient=True,
                         )
-                        submission_state.add_workflow(test_key, workflow_id)
+                        submission_state.add_workflow(
+                            test_key, workflow_id, workflow_etag
+                        )
 
                         tasks_to_test[test_key] = {
                             "task": task,
diff --git a/src/wdlci/model/submission_state.py b/src/wdlci/model/submission_state.py
index 74d97b1..e6a586b 100644
--- a/src/wdlci/model/submission_state.py
+++ b/src/wdlci/model/submission_state.py
@@ -7,9 +7,9 @@ def __init__(self):
         self.engines = {}
         self.workflow_runs = []
 
-    def add_workflow(self, workflow_key, workflow_id):
+    def add_workflow(self, workflow_key, workflow_id, workflow_etag):
         self.workflows[workflow_key] = SubmissionStateWorkflow(
-            workflow_key, workflow_id
+            workflow_key, workflow_id, workflow_etag
         )
 
     def add_engine(self, engine_id, engine_json):
@@ -31,9 +31,10 @@ def submit_all(self):
 
 
 class SubmissionStateWorkflow(object):
-    def __init__(self, workflow_key, workflow_id):
+    def __init__(self, workflow_key, workflow_id, workflow_etag):
         self._workflow_key = workflow_key
         self._workflow_id = workflow_id
+        self._workflow_etag = workflow_etag
 
     @property
     def workflow_key(self):
@@ -51,6 +52,14 @@ def workflow_id(self):
     def workflow_id(self, workflow_id):
         self._workflow_id = workflow_id
 
+    @property
+    def workflow_etag(self):
+        return self._workflow_etag
+
+    @workflow_etag.setter
+    def workflow_etag(self, workflow_etag):
+        return self._workflow_etag
+
 
 class SubmissionStateWorkflowRun(object):
     STATUS_UNSUBMITTED = "UNSUBMITTED"
diff --git a/src/wdlci/workbench/workflow_service_client.py b/src/wdlci/workbench/workflow_service_client.py
index ab3d80d..91b569f 100644
--- a/src/wdlci/workbench/workflow_service_client.py
+++ b/src/wdlci/workbench/workflow_service_client.py
@@ -43,8 +43,31 @@ def register_workflow(self, workflow_key, workflow_config, transient=False):
                 f"Could not register workflow [{workflow_key}] on Workbench", 1
             )
 
-        return response.json()["id"]
+        workflow_id = response.json()["id"]
+        workflow_etag = response.headers.get("Etag").strip('"')
+        return workflow_id, workflow_etag
 
+    # Delete individual custom workflows
+    def delete_custom_workflow(self, workflow):
+        env = Config.instance().env
+        base_url, namespace = (
+            env.workbench_workflow_service_url,
+            env.workbench_namespace,
+        )
+        url = f"{base_url}/{namespace}/workflows/{workflow.workflow_id}"
+        headers = {
+            "Authorization": "Bearer " + self.workflow_service_auth.access_token,
+            "If-Match": workflow.workflow_etag,
+        }
+
+        response = requests.delete(url, headers=headers)
+        if response.status_code != 204:
+            raise WdlTestCliExitException(
+                f"Error when deleting custom workflow {workflow.workflow_key} ({workflow.workflow_id})",
+                1,
+            )
+
+    # Purge all custom workflows
     def purge_custom_workflows(self):
         env = Config.instance().env
         base_url, namespace = (