c2cwsgiutils 5.1.7.dev20230901073305__py3-none-any.whl → 5.2.1.dev197__py3-none-any.whl

c2cwsgiutils/health_check.py

@@ -12,24 +12,50 @@ 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
 import sqlalchemy.engine
 import sqlalchemy.orm
+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."""
@@ -58,12 +84,12 @@ class _Binding:
     def name(self) -> str:
         raise NotImplementedError()
 
-    def __enter__(self) -> 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]:
@@ -78,27 +104,27 @@ class _NewBinding(_Binding):
     def name(self) -> str:
         return self.session.engine_name(self.readwrite)
 
-    def __enter__(self) -> sqlalchemy.orm.Session:
+    def __enter__(self) -> scoped_session:
         return self.session(None, self.readwrite)
 
 
 class _OldBinding(_Binding):
-    def __init__(self, session: sqlalchemy.orm.scoping.scoped_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
 
     def name(self) -> str:
-        return cast(str, self.engine.c2c_name)
+        return cast(str, self.engine.c2c_name)  # type: ignore
 
-    def __enter__(self) -> sqlalchemy.orm.Session:
-        self.prev_bind = self.session.bind
+    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]:
@@ -107,7 +133,7 @@ class _OldBinding(_Binding):
 
 
 def _get_binding_class(
-    session: Union[sqlalchemy.orm.scoping.scoped_session, c2cwsgiutils.db.SessionFactory],
+    session: Union[scoped_session, c2cwsgiutils.db.SessionFactory],
     ro_engin: sqlalchemy.engine.Engine,
     rw_engin: sqlalchemy.engine.Engine,
     readwrite: bool,
@@ -119,15 +145,15 @@ def _get_binding_class(
 
 
 def _get_bindings(
-    session: Union[sqlalchemy.orm.scoping.scoped_session, c2cwsgiutils.db.SessionFactory],
+    session: Union[scoped_session, c2cwsgiutils.db.SessionFactory],
     engine_type: EngineType,
-) -> List[sqlalchemy.engine.Engine]:
+) -> list[_Binding]:
     if isinstance(session, c2cwsgiutils.db.SessionFactory):
         ro_engin = session.ro_engine
         rw_engin = session.rw_engine
     else:
-        ro_engin = session.c2c_ro_bind
-        rw_engin = session.c2c_rw_bind
+        ro_engin = session.c2c_ro_bind  # type: ignore
+        rw_engin = session.c2c_rw_bind  # type: ignore
 
     if rw_engin == ro_engin:
         engine_type = EngineType.WRITE_ONLY
@@ -159,7 +185,9 @@ def _get_alembic_version(alembic_ini_path: str, name: str) -> str:
     ).decode("utf-8")
     out_match = ALEMBIC_HEAD_RE.match(out)
     if not out_match:
-        raise Exception("Cannot get the alembic HEAD version from: " + out)
+        raise Exception(  # pylint: disable=broad-exception-raised
+            "Cannot get the alembic HEAD version from: " + out
+        )
     return out_match.group(1)
 
 
@@ -175,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,
@@ -190,8 +218,8 @@ class HealthCheck:
 
     def add_db_session_check(
         self,
-        session: Union[sqlalchemy.orm.scoping.scoped_session, c2cwsgiutils.db.SessionFactory],
-        query_cb: Optional[Callable[[sqlalchemy.orm.scoping.scoped_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,
@@ -218,7 +246,7 @@ class HealthCheck:
 
     def add_alembic_check(
         self,
-        session: Union[sqlalchemy.orm.scoping.scoped_session, c2cwsgiutils.db.SessionFactory],
+        session: scoped_session,
         alembic_ini_path: str,
         level: int = 2,
         name: str = "alembic",
@@ -247,44 +275,38 @@ class HealthCheck:
 
         if version_schema is None:
             version_schema = config.get(name, "version_table_schema", fallback="public")
+        assert version_schema
 
         if version_table is None:
             version_table = config.get(name, "version_table", fallback="alembic_version")
+        assert version_table
 
         class _Check:
-            def __init__(self, session: sqlalchemy.orm.scoping.scoped_session) -> None:
+            def __init__(self, session: scoped_session) -> None:
                 self.session = session
 
             def __call__(self, request: pyramid.request.Request) -> str:
+                assert version_schema
+                assert version_table
                 for binding in _get_bindings(self.session, EngineType.READ_AND_WRITE):
-                    with binding as session:
-                        if stats.USE_TAGS:
-                            key = ["sql", "manual", "health_check", "alembic"]
-                            tags: Optional[Dict[str, str]] = dict(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):
-                            quote = session.bind.dialect.identifier_preparer.quote
-                            (actual_version,) = session.execute(
-                                "SELECT version_num FROM "  # nosec
-                                f"{quote(version_schema)}.{quote(version_table)}"
-                            ).fetchone()
-                            if stats.USE_TAGS:
-                                stats.increment_counter(
-                                    ["alembic_version"], 1, tags=dict(version=actual_version, name=name)
+                    with binding as binded_session:
+                        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
+                                    f"{sqlalchemy.sql.quoted_name(version_schema, True)}."
+                                    f"{sqlalchemy.sql.quoted_name(version_table, True)}"
                                 )
-                            else:
-                                stats.increment_counter(["alembic_version", name, actual_version], 1)
+                            ).fetchone()
+                            assert result is not None
+                            (actual_version,) = result
+                            _PROMETHEUS_ALEMBIC_VERSION.labels(
+                                version=actual_version, name=name, configuration=alembic_ini_path
+                            ).set(1)
                             if actual_version != version_:
-                                raise Exception(
+                                raise Exception(  # pylint: disable=broad-exception-raised
                                     f"Invalid alembic version (db: {actual_version}, code: {version_})"
                                 )
                 return version_
@@ -392,28 +414,20 @@ class HealthCheck:
         """
         Check that the version matches across all instances.
 
-        Arguments:
+        Keyword Arguments:
 
             name: the name of the check (defaults to "version")
             level: the level of the health check
-        :return:
         """
 
-        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=dict(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 dict(version=ref, count=len(versions))
+            return {"version": ref, "count": len(versions)}
 
         assert name
         self._checks.append((name, check, level))
@@ -439,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": {},
         }
@@ -463,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=dict(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=dict(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":
@@ -491,24 +499,20 @@ class HealthCheck:
     @staticmethod
     def _create_db_engine_check(
         binding: _Binding,
-        query_cb: Callable[[sqlalchemy.orm.scoping.scoped_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]] = dict(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.scoping.scoped_session], Any]:
-        def query(session: sqlalchemy.orm.scoping.scoped_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")
@@ -520,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()