From 9709a9602769b7e1b0da810e58d3edcea0728c7a Mon Sep 17 00:00:00 2001
From: skelmis <ethanmkh2@gmail.com>
Date: Thu, 4 Jan 2024 16:04:49 +1300
Subject: [PATCH] feat: expose cached item counters

---
 garven/routers/aggregate.py | 33 ++++++++++++++++++++++++++++++++-
 garven/schema/__init__.py   |  2 +-
 garven/schema/statistic.py  | 10 ++++++++++
 3 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/garven/routers/aggregate.py b/garven/routers/aggregate.py
index a9b07c2..156caa4 100644
--- a/garven/routers/aggregate.py
+++ b/garven/routers/aggregate.py
@@ -9,7 +9,7 @@
 from zonis import RequestFailed
 
 from garven.dependencies import get_auth_header
-from garven.schema import Statistic
+from garven.schema import Statistic, CachedItemsStatistic
 
 if TYPE_CHECKING:
     from zonis.server import Server
@@ -45,3 +45,34 @@ async def guild_count(request: Request):
         )
 
     return statistic
+
+
+@aggregate_router.get(
+    "/cached/count", description="Fetch the counts of cached items in each cluster."
+)
+async def cached_item_counter(request: Request):
+    z: Server = request.app.zonis
+    partial_response = False
+    data: dict[str, dict[str, int]] = await z.request_all("cached_item_count")
+    totals: dict[str, int] = {k: 0 for k in data.keys()}
+
+    for key, value in data.items():
+        if isinstance(value, RequestFailed):
+            partial_response = True
+            log.error("/cached/count WS threw '%s'", value.response_data)
+            continue
+
+        totals[key] += value
+
+    statistic = CachedItemsStatistic(
+        per_cluster=data, partial_response=partial_response, total_counts=totals
+    )
+
+    cluster_count = int(os.environ["CLUSTER_COUNT"])
+    if len(data.keys()) != cluster_count:
+        statistic.partial_response = True
+        log.error(
+            "/cached/count did not get a response from all %s clusters", cluster_count
+        )
+
+    return statistic
diff --git a/garven/schema/__init__.py b/garven/schema/__init__.py
index c6f68fe..2ed62af 100644
--- a/garven/schema/__init__.py
+++ b/garven/schema/__init__.py
@@ -1,3 +1,3 @@
 from .message import Message
-from .statistic import Statistic
+from .statistic import Statistic, CachedItemsStatistic
 from .ratelimited import RateLimited
diff --git a/garven/schema/statistic.py b/garven/schema/statistic.py
index 89b72e7..66ab25d 100644
--- a/garven/schema/statistic.py
+++ b/garven/schema/statistic.py
@@ -11,3 +11,13 @@ class Statistic(BaseModel):
 
     class Config:
         schema_extra = {"example": {"statistic": 500, "partial_response": False}}
+
+
+class CachedItemsStatistic(BaseModel):
+    total_counts: dict[str, int]
+    per_cluster: dict[str, dict[str, int]]
+    partial_response: bool = Field(
+        False,
+        description="This will be true when the returned statistic "
+        "does not accurately represent the entire expected dataset.",
+    )