Merge "Check that the limit is never over 5000 in Loki queries"

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