1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
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
|
