PowerHandler with Prometheus

.gitignore

@@ -7,6 +7,7 @@ MANIFEST
 
 # Unit test / coverage reports
 .pytest_cache/
+__pycache__/
 
 # Environments
 .env

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) + "]")

lambda_powertools/runtime.py

@@ -1,4 +1,5 @@
 import os
+import lambda_powertools.prometheus as prometheus
 from lambda_powertools.logger import logger
 
 
@@ -8,13 +9,17 @@ def wrapper(event, context):
         logger.reset()
         logger.capture(env=os.environ, event=event, context=context)
 
+        prometheus.reset()
+
         response = None
         error = None
         try:
             response = wrapped_handler(event, context)
         except Exception as e:
             error = e
 
+        prometheus.flush_metrics()
+
         if error is not None:
             raise error
 

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'

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"]')