diff --git a/.gitignore b/.gitignore index 521ba53..1f9fe3d 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ MANIFEST # Unit test / coverage reports .pytest_cache/ +__pycache__/ # Environments .env diff --git a/lambda_powertools/prometheus.py b/lambda_powertools/prometheus.py new file mode 100644 index 0000000..906bdde --- /dev/null +++ b/lambda_powertools/prometheus.py @@ -0,0 +1,72 @@ +import os +import json +import prometheus_client +from prometheus_client import CollectorRegistry + +prometheus_client.disable_created_metrics() +prometheus_client.REGISTRY.unregister(prometheus_client.GC_COLLECTOR) +prometheus_client.REGISTRY.unregister(prometheus_client.PLATFORM_COLLECTOR) +prometheus_client.REGISTRY.unregister(prometheus_client.PROCESS_COLLECTOR) + + +def reset(): + for collector in prometheus_client.REGISTRY._collector_to_names: + if "_value" in dir(collector): + # case Counter without labels + collector._value.set(0) + elif "_metrics" in dir(collector): + for metric in collector._metrics.values(): + if "_value" in dir(metric): + # case Counter with labels + metric._value.set(0) + elif "_buckets" in dir(metric) and "_sum" in dir(metric): + # case Histogram with labels + metric._sum.set(0) + for bucket in metric._buckets: + bucket.set(0) + elif "_buckets" in dir(collector) and "_sum" in dir(collector): + # case Histogram without labels + collector._sum.set(0) + for bucket in collector._buckets: + bucket.set(0) + + +def filter_metrics(m): + # Other metric types are not supported so far + if m._type not in ("counter", "histogram"): + return False + + # Empty metrics with labels are exported without values + dir_m = dir(m) + if "_value" in dir_m and m._value._value == 0: + return False + if "_buckets" in dir_m and m._sum._value == 0: + return False + if "_metrics" in dir_m: + for metric in m._metrics.values(): + if ("_value" in dir(metric) and metric._value._value > 0) or ("_sum" in dir(metric) and metric._sum._value > 0): + return True + return False + + return True + + +def get_metrics(): + new_registry = CollectorRegistry(auto_describe=True) + + for collector in filter(filter_metrics, prometheus_client.REGISTRY._collector_to_names): + new_registry.register(collector) + + return prometheus_client.generate_latest(new_registry).decode("utf-8") + + +def flush_metrics(): + metrics = get_metrics() + + if not metrics or len(metrics) == 0: + return + + if os.environ.get("PYTEST_CURRENT_TEST"): + return + + print("PROMLOG [" + json.dumps(metrics) + "]") diff --git a/lambda_powertools/runtime.py b/lambda_powertools/runtime.py index 4bbcfc7..7ce4d5b 100644 --- a/lambda_powertools/runtime.py +++ b/lambda_powertools/runtime.py @@ -1,4 +1,5 @@ import os +import lambda_powertools.prometheus as prometheus from lambda_powertools.logger import logger @@ -8,6 +9,8 @@ def wrapper(event, context): logger.reset() logger.capture(env=os.environ, event=event, context=context) + prometheus.reset() + response = None error = None try: @@ -15,6 +18,8 @@ def wrapper(event, context): except Exception as e: error = e + prometheus.flush_metrics() + if error is not None: raise error diff --git a/setup.py b/setup.py index 076e50f..6cc687d 100644 --- a/setup.py +++ b/setup.py @@ -2,11 +2,11 @@ setup( name='lambda_powertools', packages=find_packages(include=['lambda_powertools']), - version='0.1.0', + version='0.2.0', description='Lambda Powertools is a package encapsulating utilities and best practices used to write Python Lambda functions at Spreaker.', author='Spreaker', license='MIT', - install_requires=[], + install_requires=['prometheus-client==0.18.0'], setup_requires=['pytest-runner'], tests_require=['pytest==7.2.2', 'pytest-mock==3.10.0'], test_suite='tests' diff --git a/tests/test_prometheus.py b/tests/test_prometheus.py new file mode 100644 index 0000000..4a763bd --- /dev/null +++ b/tests/test_prometheus.py @@ -0,0 +1,108 @@ +import os +import pytest +from unittest.mock import patch + +from lambda_powertools.prometheus import get_metrics, reset, flush_metrics +from prometheus_client import Counter, Histogram, Gauge + +counter_no_labels = Counter( + name="prometheus_spec_counter_no_labels", + documentation="Prometheus example counter without labels" +) + +counter_no_labels_no_value = Counter( + name="prometheus_spec_counter_no_labels_no_value", + documentation="Prometheus example counter without labels" +) + +counter_with_labels = Counter( + name="prometheus_spec_counter_with_labels", + documentation="Prometheus example counter with labels", + labelnames=["foo"] +) + +histogram_no_labels = Histogram( + name="prometheus_spec_histogram_no_labels", + documentation="Prometheus example histogram without labels", + buckets=[1, 2, 5] +) + +histogram_with_labels = Histogram( + name="prometheus_spec_histogram_with_labels", + documentation="Prometheus example histogram with labels", + buckets=[1, 2, 5], + labelnames=["foo"] +) + +gauge = Gauge( + name="prometheus_spec_gauge_metric_name", + documentation="gauge_metric_documentation" +) + + +@pytest.fixture(autouse=True) +def before_each(monkeypatch): + reset() + + +def test_prometheus_get_metrics_does_not_return_empty_metrics(): + gauge.set(1) # Gauge is not supported yet + counter_no_labels.inc(1) + counter_with_labels.labels("bar").inc(2) + histogram_no_labels.observe(2) + histogram_no_labels.observe(5) + histogram_with_labels.labels("bar").observe(2) + + reset() + + metrics = get_metrics() + + assert metrics == "" + + +def test_prometheus_get_metrics_returns_non_empty_metrics(): + gauge.set(1) # Gauge is not supported yet + counter_no_labels.inc(1) + counter_with_labels.labels("bar").inc(2) + histogram_no_labels.observe(2) + histogram_no_labels.observe(5) + histogram_with_labels.labels("bar").observe(2) + + metrics = get_metrics() + + expected = """# HELP prometheus_spec_counter_no_labels_total Prometheus example counter without labels +# TYPE prometheus_spec_counter_no_labels_total counter +prometheus_spec_counter_no_labels_total 1.0 +# HELP prometheus_spec_counter_with_labels_total Prometheus example counter with labels +# TYPE prometheus_spec_counter_with_labels_total counter +prometheus_spec_counter_with_labels_total{foo="bar"} 2.0 +# HELP prometheus_spec_histogram_no_labels Prometheus example histogram without labels +# TYPE prometheus_spec_histogram_no_labels histogram +prometheus_spec_histogram_no_labels_bucket{le="1.0"} 0.0 +prometheus_spec_histogram_no_labels_bucket{le="2.0"} 1.0 +prometheus_spec_histogram_no_labels_bucket{le="5.0"} 2.0 +prometheus_spec_histogram_no_labels_bucket{le="+Inf"} 2.0 +prometheus_spec_histogram_no_labels_count 2.0 +prometheus_spec_histogram_no_labels_sum 7.0 +# HELP prometheus_spec_histogram_with_labels Prometheus example histogram with labels +# TYPE prometheus_spec_histogram_with_labels histogram +prometheus_spec_histogram_with_labels_bucket{foo="bar",le="1.0"} 0.0 +prometheus_spec_histogram_with_labels_bucket{foo="bar",le="2.0"} 1.0 +prometheus_spec_histogram_with_labels_bucket{foo="bar",le="5.0"} 1.0 +prometheus_spec_histogram_with_labels_bucket{foo="bar",le="+Inf"} 1.0 +prometheus_spec_histogram_with_labels_count{foo="bar"} 1.0 +prometheus_spec_histogram_with_labels_sum{foo="bar"} 2.0 +""" + + assert metrics == expected + + +@patch('builtins.print') +@patch.dict(os.environ, {"PYTEST_CURRENT_TEST": ""}) +def test_prometheus_flush_metrics(mock_print): + counter_no_labels.inc(1) + counter_with_labels.labels("bar").inc(2) + + flush_metrics() + + mock_print.assert_called_with('PROMLOG ["# HELP prometheus_spec_counter_no_labels_total Prometheus example counter without labels\\n# TYPE prometheus_spec_counter_no_labels_total counter\\nprometheus_spec_counter_no_labels_total 1.0\\n# HELP prometheus_spec_counter_with_labels_total Prometheus example counter with labels\\n# TYPE prometheus_spec_counter_with_labels_total counter\\nprometheus_spec_counter_with_labels_total{foo=\\"bar\\"} 2.0\\n"]')