Updates

ETB-API/venv/lib/python3.12/site-packages/django_prometheus/db/__init__.py

@@ -0,0 +1,20 @@
+# Import all metrics
+from django_prometheus.db.metrics import (
+    Counter,
+    connection_errors_total,
+    connections_total,
+    errors_total,
+    execute_many_total,
+    execute_total,
+    query_duration_seconds,
+)
+
+__all__ = [
+    "Counter",
+    "connection_errors_total",
+    "connections_total",
+    "errors_total",
+    "execute_many_total",
+    "execute_total",
+    "query_duration_seconds",
+]

ETB-API/venv/lib/python3.12/site-packages/django_prometheus/db/backends/mysql/base.py

@@ -0,0 +1,16 @@
+from django.db.backends.mysql import base
+
+from django_prometheus.db.common import DatabaseWrapperMixin, ExportingCursorWrapper
+
+
+class DatabaseFeatures(base.DatabaseFeatures):
+    """Our database has the exact same features as the base one."""
+
+
+class DatabaseWrapper(DatabaseWrapperMixin, base.DatabaseWrapper):
+    CURSOR_CLASS = base.CursorWrapper
+
+    def create_cursor(self, name=None):
+        cursor = self.connection.cursor()
+        CursorWrapper = ExportingCursorWrapper(self.CURSOR_CLASS, self.alias, self.vendor)
+        return CursorWrapper(cursor)

ETB-API/venv/lib/python3.12/site-packages/django_prometheus/db/backends/postgis/base.py

@@ -0,0 +1,21 @@
+from django.contrib.gis.db.backends.postgis import base
+from django.db.backends.postgresql.base import Cursor
+
+from django_prometheus.db.common import DatabaseWrapperMixin, ExportingCursorWrapper
+
+
+class DatabaseWrapper(DatabaseWrapperMixin, base.DatabaseWrapper):
+    def get_new_connection(self, *args, **kwargs):
+        conn = super().get_new_connection(*args, **kwargs)
+        conn.cursor_factory = ExportingCursorWrapper(
+            conn.cursor_factory or Cursor(),
+            "postgis",
+            self.vendor,
+        )
+
+        return conn
+
+    def create_cursor(self, name=None):
+        # cursor_factory is a kwarg to connect() so restore create_cursor()'s
+        # default behavior
+        return base.DatabaseWrapper.create_cursor(self, name=name)

ETB-API/venv/lib/python3.12/site-packages/django_prometheus/db/backends/postgresql/base.py

@@ -0,0 +1,21 @@
+from django.db.backends.postgresql import base
+from django.db.backends.postgresql.base import Cursor
+
+from django_prometheus.db.common import DatabaseWrapperMixin, ExportingCursorWrapper
+
+
+class DatabaseWrapper(DatabaseWrapperMixin, base.DatabaseWrapper):
+    def get_new_connection(self, *args, **kwargs):
+        conn = super().get_new_connection(*args, **kwargs)
+        conn.cursor_factory = ExportingCursorWrapper(
+            conn.cursor_factory or Cursor(),
+            self.alias,
+            self.vendor,
+        )
+
+        return conn
+
+    def create_cursor(self, name=None):
+        # cursor_factory is a kwarg to connect() so restore create_cursor()'s
+        # default behavior
+        return base.DatabaseWrapper.create_cursor(self, name=name)

ETB-API/venv/lib/python3.12/site-packages/django_prometheus/db/backends/spatialite/base.py

@@ -0,0 +1,12 @@
+from django.contrib.gis.db.backends.spatialite import base, features
+from django.db.backends.sqlite3 import base as sqlite_base
+
+from django_prometheus.db.common import DatabaseWrapperMixin
+
+
+class DatabaseFeatures(features.DatabaseFeatures):
+    """Our database has the exact same features as the base one."""
+
+
+class DatabaseWrapper(DatabaseWrapperMixin, base.DatabaseWrapper):
+    CURSOR_CLASS = sqlite_base.SQLiteCursorWrapper

ETB-API/venv/lib/python3.12/site-packages/django_prometheus/db/backends/sqlite3/base.py

@@ -0,0 +1,11 @@
+from django.db.backends.sqlite3 import base
+
+from django_prometheus.db.common import DatabaseWrapperMixin
+
+
+class DatabaseFeatures(base.DatabaseFeatures):
+    """Our database has the exact same features as the base one."""
+
+
+class DatabaseWrapper(DatabaseWrapperMixin, base.DatabaseWrapper):
+    CURSOR_CLASS = base.SQLiteCursorWrapper

ETB-API/venv/lib/python3.12/site-packages/django_prometheus/db/common.py

@@ -0,0 +1,80 @@
+from django_prometheus.db import (
+    connection_errors_total,
+    connections_total,
+    errors_total,
+    execute_many_total,
+    execute_total,
+    query_duration_seconds,
+)
+
+
+class ExceptionCounterByType:
+    """A context manager that counts exceptions by type.
+
+    Exceptions increment the provided counter, whose last label's name
+    must match the `type_label` argument.
+
+    In other words:
+
+    c = Counter('http_request_exceptions_total', 'Counter of exceptions',
+                ['method', 'type'])
+    with ExceptionCounterByType(c, extra_labels={'method': 'GET'}):
+        handle_get_request()
+    """
+
+    def __init__(self, counter, type_label="type", extra_labels=None):
+        self._counter = counter
+        self._type_label = type_label
+        self._labels = dict(extra_labels)  # Copy labels since we modify them.
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, typ, value, traceback):
+        if typ is not None:
+            self._labels.update({self._type_label: typ.__name__})
+            self._counter.labels(**self._labels).inc()
+
+
+class DatabaseWrapperMixin:
+    """Extends the DatabaseWrapper to count connections and cursors."""
+
+    def get_new_connection(self, *args, **kwargs):
+        connections_total.labels(self.alias, self.vendor).inc()
+        try:
+            return super().get_new_connection(*args, **kwargs)
+        except Exception:
+            connection_errors_total.labels(self.alias, self.vendor).inc()
+            raise
+
+    def create_cursor(self, name=None):
+        return self.connection.cursor(factory=ExportingCursorWrapper(self.CURSOR_CLASS, self.alias, self.vendor))
+
+
+def ExportingCursorWrapper(cursor_class, alias, vendor):
+    """Returns a CursorWrapper class that knows its database's alias and
+    vendor name.
+    """
+    labels = {"alias": alias, "vendor": vendor}
+
+    class CursorWrapper(cursor_class):
+        """Extends the base CursorWrapper to count events."""
+
+        def execute(self, *args, **kwargs):
+            execute_total.labels(alias, vendor).inc()
+            with (
+                query_duration_seconds.labels(**labels).time(),
+                ExceptionCounterByType(errors_total, extra_labels=labels),
+            ):
+                return super().execute(*args, **kwargs)
+
+        def executemany(self, query, param_list, *args, **kwargs):
+            execute_total.labels(alias, vendor).inc(len(param_list))
+            execute_many_total.labels(alias, vendor).inc(len(param_list))
+            with (
+                query_duration_seconds.labels(**labels).time(),
+                ExceptionCounterByType(errors_total, extra_labels=labels),
+            ):
+                return super().executemany(query, param_list, *args, **kwargs)
+
+    return CursorWrapper

ETB-API/venv/lib/python3.12/site-packages/django_prometheus/db/metrics.py

@@ -0,0 +1,48 @@
+from prometheus_client import Counter, Histogram
+
+from django_prometheus.conf import NAMESPACE, PROMETHEUS_LATENCY_BUCKETS
+
+connections_total = Counter(
+    "django_db_new_connections_total",
+    "Counter of created connections by database and by vendor.",
+    ["alias", "vendor"],
+    namespace=NAMESPACE,
+)
+
+connection_errors_total = Counter(
+    "django_db_new_connection_errors_total",
+    "Counter of connection failures by database and by vendor.",
+    ["alias", "vendor"],
+    namespace=NAMESPACE,
+)
+
+execute_total = Counter(
+    "django_db_execute_total",
+    ("Counter of executed statements by database and by vendor, including bulk executions."),
+    ["alias", "vendor"],
+    namespace=NAMESPACE,
+)
+
+
+execute_many_total = Counter(
+    "django_db_execute_many_total",
+    ("Counter of executed statements in bulk operations by database and by vendor."),
+    ["alias", "vendor"],
+    namespace=NAMESPACE,
+)
+
+
+errors_total = Counter(
+    "django_db_errors_total",
+    ("Counter of execution errors by database, vendor and exception type."),
+    ["alias", "vendor", "type"],
+    namespace=NAMESPACE,
+)
+
+query_duration_seconds = Histogram(
+    "django_db_query_duration_seconds",
+    ("Histogram of query duration by database and vendor."),
+    ["alias", "vendor"],
+    buckets=PROMETHEUS_LATENCY_BUCKETS,
+    namespace=NAMESPACE,
+)