From 0d6f10d4f0047533b1e26b4b167eae48a5807247 Mon Sep 17 00:00:00 2001
From: Nareh Sahakian <sahakian@broadinstitute.org>
Date: Tue, 21 Jan 2025 13:00:14 -0500
Subject: [PATCH] update logic for populating workspace metrics

---
 .../set_up_staging_workspace_and_dataset.py   | 53 +++++++++++++------
 1 file changed, 37 insertions(+), 16 deletions(-)

diff --git a/python/set_up_staging_workspace_and_dataset.py b/python/set_up_staging_workspace_and_dataset.py
index 161379a..b4660d0 100644
--- a/python/set_up_staging_workspace_and_dataset.py
+++ b/python/set_up_staging_workspace_and_dataset.py
@@ -238,23 +238,44 @@ def _create_dataset_properties(self) -> dict:
         return additional_properties
 
     def _add_row_to_table(self, dataset_id: str) -> None:
-        ingest_records = [
-            {
-                "key": "Staging Workspace",
-                "value": f"{self.terra_billing_project}/{self.workspace_name}"
-            },
-            {
-                "key": "Authorization Group",
-                "value": self.auth_group
-            }
-        ]
-        if workspace_version:
-            ingest_records.append(
-                {
-                    "key": f"Staging Workspace Version {self.workspace_version}",
-                    "value": f"{self.terra_billing_project}/{self.workspace_name}"
-                }
+        dataset_metrics = tdr.get_dataset_table_metrics(dataset_id=dataset_id, target_table_name=self.REFERENCE_TABLE)
+        workspace_billing_combo = f"{self.terra_billing_project}/{self.workspace_name}"
+
+        ingest_records = []
+        if not dataset_metrics:
+            ingest_records.extend(
+                [
+                    {
+                        "key": "Staging Workspace",
+                        "value": workspace_billing_combo,
+                    },
+                    {
+                        "key": "Authorization Group",
+                        "value": self.auth_group
+                    }
+                ]
             )
+        else:
+            linked_workspaces = [w["value"] for w in dataset_metrics if dataset_metrics]
+            if workspace_billing_combo not in linked_workspaces:
+                if not self.workspace_version:
+                    ingest_records.extend(
+                        [
+                            {
+                                "key": "Staging Workspace",
+                                "value": workspace_billing_combo
+                            }
+                        ]
+                    )
+                else:
+                    ingest_records.extend(
+                        [
+                            {
+                                "key": f"Staging Workspace Version {self.workspace_version}",
+                                "value": workspace_billing_combo
+                            }
+                        ]
+                    )
 
         StartAndMonitorIngest(
             ingest_records=ingest_records,