Add prometheus to services methods calls (fixes #62)

Author: leplatrem

src/jbi/services.py

@@ -1,9 +1,11 @@
 """Services and functions that can be used to create custom actions"""
+import contextlib
 import logging
 from typing import Dict, List
 
 import bugzilla as rh_bugzilla
 from atlassian import Jira
+from prometheus_client import Counter, Summary
 
 from src.app import environment
 from src.jbi.models import Actions
@@ -16,14 +18,61 @@
 ServiceHealth = Dict[str, bool]
 
 
+class InstrumentedClient:
+    """This class wraps an object and increments a counter everytime
+    the specieds methods are called, and times their execution.
+    """
+
+    counters: Dict[str, Counter] = {}
+    timers: Dict[str, Counter] = {}
+
+    def __init__(self, wrapped, prefix, methods):
+        self.wrapped = wrapped
+        self.methods = methods
+
+        # We have a single instrument per prefix. Methods are reported as labels.
+        counter_name = prefix + "_methods_total"
+        if counter_name not in self.counters:
+            self.counters[counter_name] = Counter(
+                counter_name, f"{prefix} method calls", labelnames=["method"]
+            )
+        self.counter = self.counters[counter_name]
+
+        timer_name = prefix + "_methods_milliseconds"
+        if timer_name not in self.timers:
+            self.timers[timer_name] = Summary(
+                timer_name, f"{prefix} method timing", labelnames=["method"]
+            )
+        self.timer = self.timers[timer_name]
+
+    def __getattr__(self, attr):
+        if attr in self.methods:
+            self.counter.labels(method=attr).inc()
+            timer_cm = self.timer.labels(method=attr).time()
+        else:
+            timer_cm = contextlib.nullcontext()
+
+        with timer_cm:
+            return getattr(self.wrapped, attr)
+
+
 def get_jira():
     """Get atlassian Jira Service"""
-    return Jira(
+    jira_client = Jira(
         url=settings.jira_base_url,
         username=settings.jira_username,
         password=settings.jira_api_key,  # package calls this param 'password' but actually expects an api key
         cloud=True,  # we run against an instance of Jira cloud
     )
+    instrumented_methods = (
+        "update_issue_field",
+        "set_issue_status",
+        "issue_add_comment",
+        "create_issue",
+    )
+    return InstrumentedClient(
+        wrapped=jira_client, prefix="jbi_jira", methods=instrumented_methods
+    )
 
 
 def jira_visible_projects(jira=None) -> List[Dict]:
@@ -35,9 +84,17 @@ def jira_visible_projects(jira=None) -> List[Dict]:
 
 def get_bugzilla():
     """Get bugzilla service"""
-    return rh_bugzilla.Bugzilla(
+    bugzilla_client = rh_bugzilla.Bugzilla(
         settings.bugzilla_base_url, api_key=str(settings.bugzilla_api_key)
     )
+    instrumented_methods = (
+        "get_bug",
+        "get_comments",
+        "update_bugs",
+    )
+    return InstrumentedClient(
+        wrapped=bugzilla_client, prefix="jbi_bugzilla", methods=instrumented_methods
+    )
 
 
 def _bugzilla_check_health() -> ServiceHealth:

tests/unit/jbi/test_services.py

@@ -0,0 +1,47 @@
+"""
+Module for testing src/jbi/services.py
+"""
+from unittest import mock
+
+import pytest
+
+from src.jbi.services import get_bugzilla, get_jira
+
+
+def test_counter_is_incremented_on_jira_create_issue():
+    with mock.patch("src.jbi.services.Jira.create_issue"):
+        jira_client = get_jira()
+
+        labelled_counter = jira_client.counters["jbi_jira_methods_total"].labels(
+            method="create_issue"
+        )
+        with mock.patch.object(labelled_counter, "inc") as mocked:
+            jira_client.create_issue({})
+
+    assert mocked.called, "Counter was incremented on create_issue()"
+
+
+def test_timer_is_used_on_jira_create_issue():
+    with mock.patch("src.jbi.services.Jira.create_issue"):
+        jira_client = get_jira()
+
+        labelled_timer = jira_client.timers["jbi_jira_methods_milliseconds"].labels(
+            method="create_issue"
+        )
+        with mock.patch.object(labelled_timer, "time") as mocked:
+            jira_client.create_issue({})
+
+    assert mocked.called, "Timer was used on create_issue()"
+
+
+def test_timer_is_used_on_bugzilla_getcomments():
+    with mock.patch("src.jbi.services.rh_bugzilla.Bugzilla"):
+        bugzilla_client = get_bugzilla()
+
+        labelled_timer = bugzilla_client.timers[
+            "jbi_bugzilla_methods_milliseconds"
+        ].labels(method="get_comments")
+        with mock.patch.object(labelled_timer, "time") as mocked:
+            bugzilla_client.get_comments([])
+
+    assert mocked.called, "Timer was used on get_comments()"