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