From 33e29d33f6ca83a3eb9e4763e6dbea02b41ddb1c Mon Sep 17 00:00:00 2001
From: Juan Larriba <jlarriba@redhat.com>
Date: Wed, 3 Dec 2025 15:10:49 +0100
Subject: [PATCH] Check that the limit is never over 5000 in Loki queries

Loki has hard limit of 5000 rows it can return at once, and if the user
tries to bring more, it will just deny the request. Instead of
delegating the request denial to Loki, we add the proper checks before
the queries are done, so we can properly inform the user that more than
5000 rows is not acceptable.

Change-Id: I9b98bb7d0c51724b1a63e2ecbf01ffe4a4f1c8e6
Signed-off-by: Juan Larriba <jlarriba@redhat.com>
---
 cloudkitty/storage/v2/loki/__init__.py        |  9 ++++
 .../tests/storage/v2/test_storage_unit.py     | 41 +++++++++++++++++++
 requirements.txt                              |  1 +
 3 files changed, 51 insertions(+)

diff --git a/cloudkitty/storage/v2/loki/__init__.py b/cloudkitty/storage/v2/loki/__init__.py
index b704cb9b..17cb9412 100644
--- a/cloudkitty/storage/v2/loki/__init__.py
+++ b/cloudkitty/storage/v2/loki/__init__.py
@@ -17,6 +17,7 @@ import json
 
 from oslo_config import cfg
 from oslo_log import log as oslo_logging
+from werkzeug import exceptions as http_exceptions
 
 from cloudkitty import dataframe
 from cloudkitty.storage import v2 as v2_storage
@@ -161,6 +162,10 @@ class LokiStorage(v2_storage.BaseStorage):
                  filters=None,
                  metric_types=None,
                  offset=0, limit=1000, paginate=True):
+        if limit > 5000:
+            raise http_exceptions.BadRequest(
+                f"Limit {limit} exceeds maximum allowed limit of 5000 for "
+                f"Loki storage. Please reduce the limit parameter.")
         begin, end = self._local_to_utc(begin or tzutils.get_month_start(),
                                         end or tzutils.get_next_month())
         total, logs = self._conn.retrieve(
@@ -195,6 +200,10 @@ class LokiStorage(v2_storage.BaseStorage):
     def total(self, groupby=None, begin=None, end=None, metric_types=None,
               filters=None, custom_fields=None, offset=0, limit=1000,
               paginate=False):
+        if limit > 5000:
+            raise http_exceptions.BadRequest(
+                f"Limit {limit} exceeds maximum allowed limit of 5000 for "
+                f"Loki storage. Please reduce the limit parameter.")
         begin, end = self._local_to_utc(begin or tzutils.get_month_start(),
                                         end or tzutils.get_next_month())
 
diff --git a/cloudkitty/tests/storage/v2/test_storage_unit.py b/cloudkitty/tests/storage/v2/test_storage_unit.py
index 87bbb2e3..e10c0e5d 100644
--- a/cloudkitty/tests/storage/v2/test_storage_unit.py
+++ b/cloudkitty/tests/storage/v2/test_storage_unit.py
@@ -422,3 +422,44 @@ class StorageUnitTest(TestCase):
 
 
 StorageUnitTest.generate_scenarios()
+
+
+class LokiStorageLimitTest(TestCase):
+    """Test Loki-specific limit validation"""
+
+    @mock.patch(_LOKI_CLIENT_PATH, new=loki_utils.FakeLokiClient)
+    @mock.patch('cloudkitty.utils.load_conf', new=test_utils.load_conf)
+    def setUp(self):
+        super(LokiStorageLimitTest, self).setUp()
+        self.conf.set_override('backend', 'loki', 'storage')
+        self.conf.set_override('version', '2', 'storage')
+        self.storage = storage.get_storage(conf=test_utils.load_conf())
+        self.storage.init()
+
+    def test_retrieve_with_limit_over_5000_raises_exception(self):
+        from werkzeug import exceptions as http_exceptions
+        self.assertRaisesRegex(
+            http_exceptions.BadRequest,
+            r'Limit 5001 exceeds maximum allowed limit of 5000',
+            self.storage.retrieve,
+            limit=5001
+        )
+
+    def test_retrieve_with_limit_5000_succeeds(self):
+        # This should not raise an exception
+        result = self.storage.retrieve(limit=5000)
+        self.assertIsNotNone(result)
+
+    def test_total_with_limit_over_5000_raises_exception(self):
+        from werkzeug import exceptions as http_exceptions
+        self.assertRaisesRegex(
+            http_exceptions.BadRequest,
+            r'Limit 5001 exceeds maximum allowed limit of 5000',
+            self.storage.total,
+            limit=5001
+        )
+
+    def test_total_with_limit_5000_succeeds(self):
+        # This should not raise an exception
+        result = self.storage.total(limit=5000)
+        self.assertIsNotNone(result)
diff --git a/requirements.txt b/requirements.txt
index cf2bab68..1d49c431 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -32,6 +32,7 @@ influxdb>=5.3.1 # MIT
 influxdb-client>=1.36.0 # MIT
 Flask>=2.0.0 # BSD
 Flask-RESTful>=0.3.9 # BSD
+Werkzeug>=2.0.0 # BSD
 cotyledon>=1.7.3 # Apache-2.0
 futurist>=2.3.0 # Apache-2.0
 datetimerange>=0.6.1 # MIT