diff --git a/superset/utils/decorators.py b/superset/utils/decorators.py index f14335f2c..073139574 100644 --- a/superset/utils/decorators.py +++ b/superset/utils/decorators.py @@ -45,7 +45,17 @@ def statsd_gauge(metric_prefix: Optional[str] = None) -> Callable[..., Any]: current_app.config["STATS_LOGGER"].gauge(f"{metric_prefix_}.ok", 1) return result except Exception as ex: - current_app.config["STATS_LOGGER"].gauge(f"{metric_prefix_}.error", 1) + if ( + hasattr(ex, "status") + and ex.status < 500 # type: ignore # pylint: disable=no-member + ): + current_app.config["STATS_LOGGER"].gauge( + f"{metric_prefix_}.warning", 1 + ) + else: + current_app.config["STATS_LOGGER"].gauge( + f"{metric_prefix_}.error", 1 + ) raise ex return wrapped diff --git a/superset/views/base_api.py b/superset/views/base_api.py index 29bac574a..57d7e1736 100644 --- a/superset/views/base_api.py +++ b/superset/views/base_api.py @@ -112,7 +112,13 @@ def statsd_metrics(f: Callable[..., Any]) -> Callable[..., Any]: try: duration, response = time_function(f, self, *args, **kwargs) except Exception as ex: - self.incr_stats("error", func_name) + if ( + hasattr(ex, "status") + and ex.status < 500 # type: ignore # pylint: disable=no-member + ): + self.incr_stats("warning", func_name) + else: + self.incr_stats("error", func_name) raise ex self.send_stats_metrics(response, func_name, duration) @@ -205,6 +211,8 @@ class BaseSupersetApiMixin: """ if 200 <= response.status_code < 400: self.incr_stats("success", key) + elif 400 <= response.status_code < 500: + self.incr_stats("warning", key) else: self.incr_stats("error", key) if time_delta: diff --git a/tests/integration_tests/base_tests.py b/tests/integration_tests/base_tests.py index 999f22dd2..f70f0f63b 100644 --- a/tests/integration_tests/base_tests.py +++ b/tests/integration_tests/base_tests.py @@ -90,6 +90,8 @@ def post_assert_metric( rv = client.post(uri, json=data) if 200 <= rv.status_code < 400: mock_method.assert_called_once_with("success", func_name) + elif 400 <= rv.status_code < 500: + mock_method.assert_called_once_with("warning", func_name) else: mock_method.assert_called_once_with("error", func_name) return rv @@ -455,6 +457,8 @@ class SupersetTestCase(TestCase): rv = self.client.get(uri) if 200 <= rv.status_code < 400: mock_method.assert_called_once_with("success", func_name) + elif 400 <= rv.status_code < 500: + mock_method.assert_called_once_with("warning", func_name) else: mock_method.assert_called_once_with("error", func_name) return rv @@ -474,6 +478,8 @@ class SupersetTestCase(TestCase): rv = self.client.delete(uri) if 200 <= rv.status_code < 400: mock_method.assert_called_once_with("success", func_name) + elif 400 <= rv.status_code < 500: + mock_method.assert_called_once_with("warning", func_name) else: mock_method.assert_called_once_with("error", func_name) return rv @@ -501,6 +507,8 @@ class SupersetTestCase(TestCase): rv = self.client.put(uri, json=data) if 200 <= rv.status_code < 400: mock_method.assert_called_once_with("success", func_name) + elif 400 <= rv.status_code < 500: + mock_method.assert_called_once_with("warning", func_name) else: mock_method.assert_called_once_with("error", func_name) return rv diff --git a/tests/integration_tests/utils/decorators_tests.py b/tests/integration_tests/utils/decorators_tests.py deleted file mode 100644 index d0ab6f984..000000000 --- a/tests/integration_tests/utils/decorators_tests.py +++ /dev/null @@ -1,61 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -from unittest.mock import call, Mock, patch - -import pytest -from flask import current_app - -from superset.utils import decorators -from tests.integration_tests.base_tests import SupersetTestCase - - -class UtilsDecoratorsTests(SupersetTestCase): - def test_debounce(self): - mock = Mock() - - @decorators.debounce() - def myfunc(arg1: int, arg2: int, kwarg1: str = "abc", kwarg2: int = 2): - mock(arg1, kwarg1) - return arg1 + arg2 + kwarg2 - - # should be called only once when arguments don't change - myfunc(1, 1) - myfunc(1, 1) - result = myfunc(1, 1) - mock.assert_called_once_with(1, "abc") - self.assertEqual(result, 4) - - # kwarg order shouldn't matter - myfunc(1, 0, kwarg2=2, kwarg1="haha") - result = myfunc(1, 0, kwarg1="haha", kwarg2=2) - mock.assert_has_calls([call(1, "abc"), call(1, "haha")]) - self.assertEqual(result, 3) - - def test_statsd_gauge(self): - @decorators.statsd_gauge("custom.prefix") - def my_func(fail: bool, *args, **kwargs): - if fail: - raise ValueError("Error") - return "OK" - - with patch.object(current_app.config["STATS_LOGGER"], "gauge") as mock: - my_func(False, 1, 2) - mock.assert_called_once_with("custom.prefix.ok", 1) - - with pytest.raises(ValueError): - my_func(True, 1, 2) - mock.assert_called_once_with("custom.prefix.error", 1) diff --git a/tests/unit_tests/utils/test_decorators.py b/tests/unit_tests/utils/test_decorators.py new file mode 100644 index 000000000..3aafc7a91 --- /dev/null +++ b/tests/unit_tests/utils/test_decorators.py @@ -0,0 +1,87 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +from contextlib import nullcontext +from enum import Enum +from inspect import isclass +from typing import Any, Optional +from unittest.mock import call, Mock, patch + +import pytest + +from superset import app +from superset.utils import decorators + + +class ResponseValues(str, Enum): + FAIL = "fail" + WARN = "warn" + OK = "ok" + + +def test_debounce() -> None: + mock = Mock() + + @decorators.debounce() + def myfunc(arg1: int, arg2: int, kwarg1: str = "abc", kwarg2: int = 2) -> int: + mock(arg1, kwarg1) + return arg1 + arg2 + kwarg2 + + # should be called only once when arguments don't change + myfunc(1, 1) + myfunc(1, 1) + result = myfunc(1, 1) + mock.assert_called_once_with(1, "abc") + assert result == 4 + + # kwarg order shouldn't matter + myfunc(1, 0, kwarg2=2, kwarg1="haha") + result = myfunc(1, 0, kwarg1="haha", kwarg2=2) + mock.assert_has_calls([call(1, "abc"), call(1, "haha")]) + assert result == 3 + + +@pytest.mark.parametrize( + "response_value, expected_exception, expected_result", + [ + (ResponseValues.OK, None, "custom.prefix.ok"), + (ResponseValues.FAIL, ValueError, "custom.prefix.error"), + (ResponseValues.WARN, FileNotFoundError, "custom.prefix.warn"), + ], +) +def test_statsd_gauge( + response_value: str, expected_exception: Optional[Exception], expected_result: str +) -> None: + @decorators.statsd_gauge("custom.prefix") + def my_func(response: ResponseValues, *args: Any, **kwargs: Any) -> str: + if response == ResponseValues.FAIL: + raise ValueError("Error") + if response == ResponseValues.WARN: + raise FileNotFoundError("Not found") + return "OK" + + with patch.object(app.config["STATS_LOGGER"], "gauge") as mock: + cm = ( + pytest.raises(expected_exception) + if isclass(expected_exception) and issubclass(expected_exception, Exception) + else nullcontext() + ) + + with cm: + my_func(response_value, 1, 2) + mock.assert_called_once_with(expected_result, 1)