Skip to content

Commit

Permalink
Merge pull request #497 from allegro/am_prometheus_fix_path
Browse files Browse the repository at this point in the history
Add victoriametrics support + possibility to define global labels for Prometheus metrics
  • Loading branch information
ziollek committed Sep 15, 2023
2 parents 3126908 + 96693e7 commit b85114f
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 4 deletions.
9 changes: 7 additions & 2 deletions vaas/vaas/metrics/prometheus.py
Expand Up @@ -13,15 +13,20 @@

class PrometheusClient:
def __init__(self) -> None:
self.host: str = f'{settings.PROMETHEUS_GATEWAY_HOST}:{settings.PROMETHEUS_GATEWAY_PORT}'
self.host: str = (f'{settings.PROMETHEUS_GATEWAY_HOST}:{settings.PROMETHEUS_GATEWAY_PORT}'
f'{settings.VICTORIAMETRICS_PATH if settings.VICTORIAMETRICS_SUPPORT else ""}')
self.job: str = settings.PROMETHEUS_GATEWAY_JOB
self.labels: Dict[str, str] = settings.PROMETHEUS_GATEWAY_LABELS
self.metrics_bucket: Dict[str, Kind] = {}
self.registry: CollectorRegistry = CollectorRegistry()

def get_or_create(self, name: str, kind: Type[Kind]) -> Kind:
metric: Optional[Kind] = self.metrics_bucket.get(name)
if not metric:
new_metrics: Kind = kind(name, name, registry=self.registry)
if self.labels:
new_metrics: Kind = kind(name, name, self.labels.keys(), registry=self.registry).labels(**self.labels)
else:
new_metrics: Kind = kind(name, name, registry=self.registry)
self.metrics_bucket[name] = new_metrics
return new_metrics
return metric
Expand Down
24 changes: 23 additions & 1 deletion vaas/vaas/metrics/tests/test_prometheus.py
Expand Up @@ -6,7 +6,7 @@


class TestPrometheusClient(SimpleTestCase):
def test_create_new_metric_if_not_exists_in_registry(self):
def test_create_new_lableless_metric_if_not_exists_in_registry(self):
client = PrometheusClient()
self.assertEqual(len(client.metrics_bucket), 0)

Expand All @@ -15,6 +15,7 @@ def test_create_new_metric_if_not_exists_in_registry(self):
self.assertTrue(isinstance(metric, Summary))
self.assertEqual(metric._name, 'test_sum')
self.assertEqual(len(client.metrics_bucket), 1)
self.assertEqual(len(metric._labelnames), 0)

def test_return_existsing_metric_from_registry(self):
client = PrometheusClient()
Expand All @@ -25,6 +26,16 @@ def test_return_existsing_metric_from_registry(self):
self.assertEqual(len(client.metrics_bucket), 1)
self.assertEqual(first_metric, second_metric)

@override_settings(PROMETHEUS_GATEWAY_LABELS={'key': 'value'})
def test_return_metric_with_defined_labels_from_registry(self):
client = PrometheusClient()

metrics = client.get_or_create('test_gauge', Gauge)

self.assertEqual(len(client.metrics_bucket), 1)
self.assertEqual(metrics._labelnames, ('key', ))
self.assertEqual(metrics._labelvalues, ('value', ))

@override_settings(PROMETHEUS_GATEWAY_HOST='1.2.3.4', PROMETHEUS_GATEWAY_PORT='123', PROMETHEUS_GATEWAY_JOB='')
@patch('vaas.metrics.prometheus.CollectorRegistry')
@patch('vaas.metrics.prometheus.push_to_gateway')
Expand All @@ -34,6 +45,17 @@ def test_should_push_metrics_to_prometheus_gateway(self, mock_push_to_gateway, r
client.push()
mock_push_to_gateway.assert_called_with(gateway='1.2.3.4:123', job='', registry=registry_mock)

@override_settings(PROMETHEUS_GATEWAY_HOST='1.2.3.4', PROMETHEUS_GATEWAY_PORT='123', PROMETHEUS_GATEWAY_JOB='',
VICTORIAMETRICS_SUPPORT=True, VICTORIAMETRICS_PATH='/v1/import/prometheus')
@patch('vaas.metrics.prometheus.CollectorRegistry')
@patch('vaas.metrics.prometheus.push_to_gateway')
def test_should_push_metrics_to_victoriametrics_gateway(self, mock_push_to_gateway, registry_mock):
client = PrometheusClient()
client.registry = registry_mock
client.push()
mock_push_to_gateway.assert_called_with(gateway='1.2.3.4:123/v1/import/prometheus',
job='', registry=registry_mock)

@patch('vaas.metrics.prometheus.push_to_gateway', Mock(side_effect=Exception))
def test_should_continue_after_push_to_gateway_raise_an_exception(self):
with patch('vaas.metrics.prometheus.logger') as mock_logger:
Expand Down
10 changes: 9 additions & 1 deletion vaas/vaas/settings/base.py
Expand Up @@ -245,7 +245,15 @@ def serialize(value: any) -> str:
PROMETHEUS_ENABLE = env.bool('PROMETHEUS_ENABLE', default=False)
PROMETHEUS_GATEWAY_HOST = env.str('PROMETHEUS_GATEWAY_HOST', default='localhost')
PROMETHEUS_GATEWAY_PORT = env.int('PROMETHEUS_GATEWAY_PORT', default=9091)
PROMETHEUS_GATEWAY_JOB = env.str('PROMETHEUS_GATEWAY_JOB', default='')
PROMETHEUS_GATEWAY_JOB = env.str('PROMETHEUS_GATEWAY_JOB', default='vaas')
# Please in env definition stick to the proscription: name=value,second=value
# It will transfer to: {name: value, second: value}.
PROMETHEUS_GATEWAY_LABELS = env.dict('PROMETHEUS_GATEWAY_LABELS', default={})

# We also allow push metric via victoriametrics agent.
# https://docs.victoriametrics.com/#how-to-import-data-in-prometheus-exposition-format
VICTORIAMETRICS_SUPPORT = env.bool('VICTORIAMETRICS_SUPPORT', default=False)
VICTORIAMETRICS_PATH = env.str('VICTORIAMETRICS_PATH', default="/api/v1/import/prometheus")

# HEADER FOR PERMIT ACCESS TO /vaas/ ENDPOINT
ALLOW_METRICS_HEADER = env.bool('ALLOW_METRICS_HEADER', default='x-allow-metric-header')
Expand Down

0 comments on commit b85114f

Please sign in to comment.