From 814ccf6b04e97659c30ecc18dd63607a3ba485e6 Mon Sep 17 00:00:00 2001
From: moson-mo <mo-son@mailbox.org>
Date: Tue, 4 Jul 2023 09:40:39 +0200
Subject: [PATCH] feat: add Prometheus metrics for Redis cache

Adding a Prometheus counter to be able to monitor cache hits/misses
for search queries

Signed-off-by: moson-mo <mo-son@mailbox.org>
---
 aurweb/cache.py      | 13 +++++++++++--
 test/test_metrics.py | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 2 deletions(-)
 create mode 100644 test/test_metrics.py

diff --git a/aurweb/cache.py b/aurweb/cache.py
index 56bb45b7..fe1e5f1d 100644
--- a/aurweb/cache.py
+++ b/aurweb/cache.py
@@ -1,5 +1,6 @@
 import pickle
 
+from prometheus_client import Counter
 from sqlalchemy import orm
 
 from aurweb import config
@@ -7,6 +8,11 @@ from aurweb.aur_redis import redis_connection
 
 _redis = redis_connection()
 
+# Prometheus metrics
+SEARCH_REQUESTS = Counter(
+    "search_requests", "Number of search requests by cache hit/miss", ["cache"]
+)
+
 
 async def db_count_cache(key: str, query: orm.Query, expire: int = None) -> int:
     """Store and retrieve a query.count() via redis cache.
@@ -24,7 +30,7 @@ async def db_count_cache(key: str, query: orm.Query, expire: int = None) -> int:
     return int(result)
 
 
-async def db_query_cache(key: str, query: orm.Query, expire: int = None):
+async def db_query_cache(key: str, query: orm.Query, expire: int = None) -> list:
     """Store and retrieve query results via redis cache.
 
     :param key: Redis key
@@ -34,10 +40,13 @@ async def db_query_cache(key: str, query: orm.Query, expire: int = None):
     """
     result = _redis.get(key)
     if result is None:
+        SEARCH_REQUESTS.labels(cache="miss").inc()
         if _redis.dbsize() > config.getint("cache", "max_search_entries", 50000):
             return query.all()
-        _redis.set(key, (result := pickle.dumps(query.all())), ex=expire)
+        _redis.set(key, (result := pickle.dumps(query.all())))
         if expire:
             _redis.expire(key, expire)
+    else:
+        SEARCH_REQUESTS.labels(cache="hit").inc()
 
     return pickle.loads(result)
diff --git a/test/test_metrics.py b/test/test_metrics.py
new file mode 100644
index 00000000..1859d8cb
--- /dev/null
+++ b/test/test_metrics.py
@@ -0,0 +1,40 @@
+import pytest
+from prometheus_client import REGISTRY, generate_latest
+
+from aurweb import db
+from aurweb.cache import db_query_cache
+from aurweb.models.account_type import USER_ID
+from aurweb.models.user import User
+
+
+@pytest.fixture(autouse=True)
+def setup(db_test):
+    return
+
+
+@pytest.fixture
+def user() -> User:
+    with db.begin():
+        user = db.create(
+            User,
+            Username="test",
+            Email="test@example.org",
+            RealName="Test User",
+            Passwd="testPassword",
+            AccountTypeID=USER_ID,
+        )
+    yield user
+
+
+@pytest.mark.asyncio
+async def test_search_cache_metrics(user: User):
+    # Fire off 3 identical queries for caching
+    for _ in range(3):
+        await db_query_cache("key", db.query(User))
+
+    # Get metrics
+    metrics = str(generate_latest(REGISTRY))
+
+    # We should have 1 miss and 2 hits
+    assert 'search_requests_total{cache="miss"} 1.0' in metrics
+    assert 'search_requests_total{cache="hit"} 2.0' in metrics