From fbb70b95aa57cf3e69f01aaa70e281f0bc7e42d9 Mon Sep 17 00:00:00 2001
From: skelmis <ethanmkh2@gmail.com>
Date: Thu, 11 Jan 2024 23:47:49 +1300
Subject: [PATCH] feat: update garven to support ws stuff

---
 garven/routers/cluster.py       | 23 ++++++++++++++++++++++-
 garven/schema/cluster/health.py | 17 +++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/garven/routers/cluster.py b/garven/routers/cluster.py
index 38ae82e..563fa1a 100644
--- a/garven/routers/cluster.py
+++ b/garven/routers/cluster.py
@@ -14,7 +14,7 @@
 from garven.calculate_dev_cluster import get_dev_cluster
 from garven.dependencies import get_auth_header
 from garven.schema import Message
-from garven.schema.cluster import ClusterHealth, DevShare
+from garven.schema.cluster import ClusterHealth, DevShare, ClusterWSInfo
 
 if TYPE_CHECKING:
     from zonis.server import Server
@@ -46,6 +46,27 @@ async def cluster_status(request: Request):
     return ClusterHealth(clusters=d, partial_response=partial_response)
 
 
+@cluster_router.get("/latency/ws", response_model=ClusterWSInfo)
+async def cluster_status(request: Request):
+    partial_response = False
+    z: Server = request.app.zonis
+    d: dict[str, dict[str, str]] = await z.request_all("cluster_ws_status")
+    shard_data = {}
+    for _, item in deepcopy(d).items():
+        if isinstance(item, RequestFailed):
+            partial_response = True
+            log.error("/cluster/status/ws WS threw '%s'", item.response_data)
+            continue
+
+        for shard_id, value in item.items():
+            shard_data[shard_id] = value
+
+    if len(shard_data.values()) != int(os.environ.get("TOTAL_SHARDS")):
+        partial_response = True
+
+    return ClusterWSInfo(shards=shard_data, partial_response=partial_response)  # noqa
+
+
 @cluster_router.post(
     "/notify_devs", status_code=204, responses={503: {"model": Message}}
 )
diff --git a/garven/schema/cluster/health.py b/garven/schema/cluster/health.py
index 33b7392..85237b5 100644
--- a/garven/schema/cluster/health.py
+++ b/garven/schema/cluster/health.py
@@ -10,6 +10,23 @@ class ShardInfo(BaseModel):
     )
 
 
+class ClusterWSInfo(BaseModel):
+    partial_response: bool = Field(
+        False,
+        description="This will be true when the returned statistic "
+        "does not accurately represent the entire expected dataset.",
+    )
+    shards: dict[str, dict[str, str]] = Field(
+        description="A mapping of shard id's to self reported latencies"
+    )
+
+    class Config:
+        schema_extra = {
+            "partial_response": True,
+            "shards": {"1": {"ws": "1.1", "keepalive": "15.5"}},
+        }
+
+
 class ClusterInfo(BaseModel):
     cluster_is_up: bool = Field(
         description="Ignoring individual shards, is the cluster as a whole functional"