c2cwsgiutils 5.2.1__py3-none-any.whl → 5.2.1.dev197__py3-none-any.whl

c2cwsgiutils/health_check.py

@@ -12,11 +12,12 @@ import re
 import subprocess  # nosec
 import time
 import traceback
-from collections import Counter
+from collections.abc import Mapping
 from enum import Enum
 from types import TracebackType
-from typing import Any, Callable, Dict, List, Literal, Mapping, Optional, Tuple, Type, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union, cast
 
+import prometheus_client
 import pyramid.config
 import pyramid.request
 import requests
@@ -26,11 +27,35 @@ import sqlalchemy.sql
 from pyramid.httpexceptions import HTTPNotFound
 
 import c2cwsgiutils.db
-from c2cwsgiutils import auth, broadcast, config_utils, redis_utils, stats, version
+from c2cwsgiutils import auth, broadcast, config_utils, prometheus, redis_utils, version
+
+if TYPE_CHECKING:
+    scoped_session = sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session]
+else:
+    scoped_session = sqlalchemy.orm.scoped_session
 
 LOG = logging.getLogger(__name__)
 ALEMBIC_HEAD_RE = re.compile(r"^([a-f0-9]+) \(head\)\n$")
 
+_PROMETHEUS_DB_SUMMARY = prometheus_client.Summary(
+    prometheus.build_metric_name("health_check_db"),
+    "The to do a database query",
+    ["configuration", "connection", "check"],
+    unit="seconds",
+)
+_PROMETHEUS_ALEMBIC_VERSION = prometheus_client.Gauge(
+    prometheus.build_metric_name("alembic_version"),
+    "The alembic version of the database",
+    ["version", "name", "configuration"],
+    multiprocess_mode="liveall",
+)
+_PROMETHEUS_HEALTH_CHECKS_FAILURE = prometheus_client.Gauge(
+    prometheus.build_metric_name("health_check_failure"),
+    "The health check",
+    ["name"],
+    multiprocess_mode="livemax",
+)
+
 
 class EngineType(Enum):
     """The type of engine."""
@@ -59,12 +84,12 @@ class _Binding:
     def name(self) -> str:
         raise NotImplementedError()
 
-    def __enter__(self) -> sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session]:
+    def __enter__(self) -> scoped_session:
         raise NotImplementedError()
 
     def __exit__(
         self,
-        exc_type: Optional[Type[BaseException]],
+        exc_type: Optional[type[BaseException]],
         exc_value: Optional[BaseException],
         exc_traceback: Optional[TracebackType],
     ) -> Literal[False]:
@@ -79,14 +104,12 @@ class _NewBinding(_Binding):
     def name(self) -> str:
         return self.session.engine_name(self.readwrite)
 
-    def __enter__(self) -> sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session]:
+    def __enter__(self) -> scoped_session:
         return self.session(None, self.readwrite)
 
 
 class _OldBinding(_Binding):
-    def __init__(
-        self, session: sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session], engine: sqlalchemy.engine.Engine
-    ):
+    def __init__(self, session: scoped_session, engine: sqlalchemy.engine.Engine):
         self.session = session
         self.engine = engine
         self.prev_bind = None
@@ -94,14 +117,14 @@ class _OldBinding(_Binding):
     def name(self) -> str:
         return cast(str, self.engine.c2c_name)  # type: ignore
 
-    def __enter__(self) -> sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session]:
+    def __enter__(self) -> scoped_session:
         self.prev_bind = self.session.bind  # type: ignore
         self.session.bind = self.engine
         return self.session
 
     def __exit__(
         self,
-        exc_type: Optional[Type[BaseException]],
+        exc_type: Optional[type[BaseException]],
         exc_value: Optional[BaseException],
         exc_traceback: Optional[TracebackType],
     ) -> Literal[False]:
@@ -110,7 +133,7 @@ class _OldBinding(_Binding):
 
 
 def _get_binding_class(
-    session: Union[sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session], c2cwsgiutils.db.SessionFactory],
+    session: Union[scoped_session, c2cwsgiutils.db.SessionFactory],
     ro_engin: sqlalchemy.engine.Engine,
     rw_engin: sqlalchemy.engine.Engine,
     readwrite: bool,
@@ -122,9 +145,9 @@ def _get_binding_class(
 
 
 def _get_bindings(
-    session: Union[sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session], c2cwsgiutils.db.SessionFactory],
+    session: Union[scoped_session, c2cwsgiutils.db.SessionFactory],
     engine_type: EngineType,
-) -> List[_Binding]:
+) -> list[_Binding]:
     if isinstance(session, c2cwsgiutils.db.SessionFactory):
         ro_engin = session.ro_engine
         rw_engin = session.rw_engine
@@ -180,7 +203,7 @@ class HealthCheck:
             "c2c_health_check", config_utils.get_base_path(config) + r"/health_check", request_method="GET"
         )
         config.add_view(self._view, route_name="c2c_health_check", renderer="fast_json", http_cache=0)
-        self._checks: List[Tuple[str, Callable[[pyramid.request.Request], Any], int]] = []
+        self._checks: list[tuple[str, Callable[[pyramid.request.Request], Any], int]] = []
 
         self.name = config_utils.env_or_config(
             config,
@@ -195,8 +218,8 @@ class HealthCheck:
 
     def add_db_session_check(
         self,
-        session: Union[sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session], c2cwsgiutils.db.SessionFactory],
-        query_cb: Optional[Callable[[sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session]], Any]] = None,
+        session: Union[scoped_session, c2cwsgiutils.db.SessionFactory],
+        query_cb: Optional[Callable[[scoped_session], Any]] = None,
         at_least_one_model: Optional[object] = None,
         level: int = 1,
         engine_type: EngineType = EngineType.READ_AND_WRITE,
@@ -223,7 +246,7 @@ class HealthCheck:
 
     def add_alembic_check(
         self,
-        session: sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session],
+        session: scoped_session,
         alembic_ini_path: str,
         level: int = 2,
         name: str = "alembic",
@@ -259,7 +282,7 @@ class HealthCheck:
         assert version_table
 
         class _Check:
-            def __init__(self, session: sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session]) -> None:
+            def __init__(self, session: scoped_session) -> None:
                 self.session = session
 
             def __call__(self, request: pyramid.request.Request) -> str:
@@ -267,20 +290,9 @@ class HealthCheck:
                 assert version_table
                 for binding in _get_bindings(self.session, EngineType.READ_AND_WRITE):
                     with binding as binded_session:
-                        if stats.USE_TAGS:
-                            key = ["sql", "manual", "health_check", "alembic"]
-                            tags: Optional[Dict[str, str]] = {"conf": alembic_ini_path, "con": binding.name()}
-                        else:
-                            key = [
-                                "sql",
-                                "manual",
-                                "health_check",
-                                "alembic",
-                                alembic_ini_path,
-                                binding.name(),
-                            ]
-                            tags = None
-                        with stats.timer_context(key, tags):
+                        with _PROMETHEUS_DB_SUMMARY.labels(
+                            configuration=alembic_ini_path, connection=binding.name(), check="alembic"
+                        ).time():
                             result = binded_session.execute(
                                 sqlalchemy.text(
                                     "SELECT version_num FROM "  # nosec
@@ -290,12 +302,9 @@ class HealthCheck:
                             ).fetchone()
                             assert result is not None
                             (actual_version,) = result
-                            if stats.USE_TAGS:
-                                stats.increment_counter(
-                                    ["alembic_version"], 1, tags={"version": actual_version, "name": name}
-                                )
-                            else:
-                                stats.increment_counter(["alembic_version", name, actual_version], 1)
+                            _PROMETHEUS_ALEMBIC_VERSION.labels(
+                                version=actual_version, name=name, configuration=alembic_ini_path
+                            ).set(1)
                             if actual_version != version_:
                                 raise Exception(  # pylint: disable=broad-exception-raised
                                     f"Invalid alembic version (db: {actual_version}, code: {version_})"
@@ -411,19 +420,12 @@ class HealthCheck:
             level: the level of the health check
         """
 
-        def check(request: pyramid.request.Request) -> Dict[str, Any]:
+        def check(request: pyramid.request.Request) -> dict[str, Any]:
+            ref = version.get_version()
             all_versions = _get_all_versions()
             assert all_versions
             versions = [e for e in all_versions if e is not None]
-            # Output the versions we see on the monitoring
-            v: Optional[str]
-            for v, count in Counter(versions).items():
-                if stats.USE_TAGS:
-                    stats.increment_counter(["version"], count, tags={"version": v})
-                else:
-                    stats.increment_counter(["version", v], count)
-
-                ref = versions[0]
+
             assert all(v == ref for v in versions), "Non identical versions: " + ", ".join(versions)
             return {"version": ref, "count": len(versions)}
 
@@ -451,7 +453,7 @@ class HealthCheck:
     def _view(self, request: pyramid.request.Request) -> Mapping[str, Any]:
         max_level = int(request.params.get("max_level", "1"))
         is_auth = auth.is_auth(request)
-        results: Dict[str, Dict[str, Any]] = {
+        results: dict[str, dict[str, Any]] = {
             "failures": {},
             "successes": {},
         }
@@ -475,25 +477,19 @@ class HealthCheck:
         level: int,
         name: str,
         request: pyramid.request.Request,
-        results: Dict[str, Dict[str, Any]],
+        results: dict[str, dict[str, Any]],
     ) -> None:
-        start = time.monotonic()
+        start = time.perf_counter()
         try:
             result = check(request)
-            results["successes"][name] = {"timing": time.monotonic() - start, "level": level}
+            results["successes"][name] = {"timing": time.perf_counter() - start, "level": level}
             if result is not None:
                 results["successes"][name]["result"] = result
-            if stats.USE_TAGS:
-                stats.increment_counter(["health_check"], 1, tags={"name": name, "outcome": "success"})
-            else:
-                stats.increment_counter(["health_check", name, "success"], 1)
+            _set_success(check_name=name)
         except Exception as e:  # pylint: disable=broad-except
-            if stats.USE_TAGS:
-                stats.increment_counter(["health_check"], 1, tags={"name": name, "outcome": "failure"})
-            else:
-                stats.increment_counter(["health_check", name, "failure"], 1)
+            _PROMETHEUS_HEALTH_CHECKS_FAILURE.labels(name=name).set(1)
             LOG.warning("Health check %s failed", name, exc_info=True)
-            failure = {"message": str(e), "timing": time.monotonic() - start, "level": level}
+            failure = {"message": str(e), "timing": time.perf_counter() - start, "level": level}
             if isinstance(e, JsonCheckException) and e.json_data() is not None:
                 failure["result"] = e.json_data()
             if is_auth or os.environ.get("DEVELOPMENT", "0") != "0":
@@ -503,24 +499,20 @@ class HealthCheck:
     @staticmethod
     def _create_db_engine_check(
         binding: _Binding,
-        query_cb: Callable[[sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session]], None],
-    ) -> Tuple[str, Callable[[pyramid.request.Request], None]]:
+        query_cb: Callable[[scoped_session], None],
+    ) -> tuple[str, Callable[[pyramid.request.Request], None]]:
         def check(request: pyramid.request.Request) -> None:
             with binding as session:
-                if stats.USE_TAGS:
-                    key = ["sql", "manual", "health_check", "db"]
-                    tags: Optional[Dict[str, str]] = {"con": binding.name()}
-                else:
-                    key = ["sql", "manual", "health_check", "db", binding.name()]
-                    tags = None
-                with stats.timer_context(key, tags):
+                with _PROMETHEUS_DB_SUMMARY.labels(
+                    connection=binding.name(), check="database", configuration="<default>"
+                ).time():
                     return query_cb(session)
 
         return "db_engine_" + binding.name(), check
 
     @staticmethod
-    def _at_least_one(model: Any) -> Callable[[sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session]], Any]:
-        def query(session: sqlalchemy.orm.scoped_session[sqlalchemy.orm.Session]) -> None:
+    def _at_least_one(model: Any) -> Callable[[scoped_session], Any]:
+        def query(session: scoped_session) -> None:
             result = session.query(model).first()
             if result is None:
                 raise HTTPNotFound(model.__name__ + " record not found")
@@ -532,6 +524,13 @@ def _maybe_function(what: Any, request: pyramid.request.Request) -> Any:
     return what(request) if callable(what) else what
 
 
+@broadcast.decorator(expect_answers=False)
+def _set_success(check_name: str) -> None:
+    """Set check in success in all process."""
+
+    _PROMETHEUS_HEALTH_CHECKS_FAILURE.labels(name=check_name).set(0)
+
+
 @broadcast.decorator(expect_answers=True)
 def _get_all_versions() -> Optional[str]:
     return version.get_version()