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

@@ -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())
 

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)

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