From 998ee16a61b45d3414f6e2197ef8e13a47cda5db Mon Sep 17 00:00:00 2001
From: Riley Grant <grantrileyh@gmail.com>
Date: Thu, 21 Nov 2024 12:54:49 -0600
Subject: [PATCH] feat(deployctl): note when deployment is demo

---
 deploy/deployctl/shell.py                     | 10 ++--
 .../subcommands/browser_deployments.py        | 46 +++++++++++++++++--
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/deploy/deployctl/shell.py b/deploy/deployctl/shell.py
index 16c84fa3f..6b7d2c076 100644
--- a/deploy/deployctl/shell.py
+++ b/deploy/deployctl/shell.py
@@ -35,24 +35,26 @@ def kubectl(args: typing.List[str], **kwargs) -> str:
     )
 
 
-def get_k8s_deployments(selector: str) -> typing.List[str]:
+def get_k8s_deployments(selector: str) -> typing.List[typing.Tuple[str, str]]:
     result = kubectl(
         [
             "get",
             "deployments",
             f"--selector={selector}",
             "--sort-by={.metadata.creationTimestamp}",
-            "--output=jsonpath={range .items[*]}{.metadata.name}{'\\n'}",
+            "--output=jsonpath={range .items[*]}{.metadata.name} {.spec.template.spec.nodeSelector.cloud\\.google\\.com/gke-nodepool}{'\\n'}",
         ]
     )
-    return [line for line in result.splitlines() if line]
+    return [
+        (parts[0], parts[1]) for line in result.splitlines() for parts in [line.split(maxsplit=1)] if len(parts) == 2
+    ]
 
 
 def get_most_recent_k8s_deployment(selector: str) -> str:
     deployments = get_k8s_deployments(selector)
     if not deployments:
         raise RuntimeError(f"No deployment matching '{selector}' found")
-    return deployments[len(deployments) - 1]
+    return deployments[len(deployments) - 1][0]
 
 
 def k8s_deployment_exists(k8s_deployment_name: str) -> bool:
diff --git a/deploy/deployctl/subcommands/browser_deployments.py b/deploy/deployctl/subcommands/browser_deployments.py
index b1c9eedf2..d00cedeff 100644
--- a/deploy/deployctl/subcommands/browser_deployments.py
+++ b/deploy/deployctl/subcommands/browser_deployments.py
@@ -5,6 +5,7 @@
 import string
 import sys
 import typing
+import yaml
 
 from deployctl.config import config
 from deployctl.shell import kubectl, get_most_recent_tag, image_exists, get_k8s_deployments
@@ -115,19 +116,58 @@ def deployments_directory() -> str:
     return path
 
 
+def print_pool_name(pool: str) -> str:
+    if pool == "demo-pool":
+        return "(demo)"
+    if pool == "main-pool":
+        return ""
+
+    return pool
+
+
+def determine_deployment_pool(path: str) -> str:
+    with open(path) as f:
+        content = yaml.safe_load(f)
+
+    patches = content.get("patches", [])
+    for patch in patches:
+        if isinstance(patch, dict):
+            patch_content = yaml.safe_load(patch.get("patch", ""))
+        else:
+            patch_content = yaml.safe_load(patch)
+
+        if (
+            patch_content
+            and patch_content.get("kind") == "Deployment"
+            and patch_content.get("spec", {})
+            .get("template", {})
+            .get("spec", {})
+            .get("nodeSelector", {})
+            .get("cloud.google.com/gke-nodepool")
+            == "demo-pool"
+        ):
+            return "demo-pool"
+
+    return "main-pool"
+
+
 def list_deployments() -> None:
     print("Local configurations")
     print("====================")
     paths = reversed(sorted(glob.iglob(f"{deployments_directory()}/*/kustomization.yaml"), key=os.path.getmtime))
     for path in paths:
-        print(os.path.basename(os.path.dirname(path)))
+        name = os.path.basename(os.path.dirname(path))
+        pool = determine_deployment_pool(path)
+        print(f"{name} {print_pool_name(pool)}")
 
     print()
 
     print("Cluster deployments")
     print("===================")
-    for deployment in get_k8s_deployments("component=gnomad-browser"):
-        print(deployment[len("gnomad-browser-") :])
+    for deployment, pool in get_k8s_deployments("component=gnomad-browser"):
+        print(f"{deployment[len('gnomad-browser-'):]} {print_pool_name(pool)}")
+
+    print()
 
 
 def create_deployment(name: str, browser_tag: str = None, api_tag: str = None, demo: bool = False) -> None: