From 5d3fd520b72e72ab8dc39c371d924e0e9d5ce5cc Mon Sep 17 00:00:00 2001
From: Oleg Komarov <oleg@blender.org>
Date: Mon, 2 Sep 2024 19:26:34 +0200
Subject: [PATCH] Metrics: collect application metrics via vector

As a first step, start collecting db query counts and timings.

This code can be extended with custom metrics support.
Later, once the interface stabilizes, it should be extracted into a library
reused by other django projects.
---
 blender_extensions/middleware.py | 89 ++++++++++++++++++++++++++++++++
 blender_extensions/settings.py   |  4 ++
 2 files changed, 93 insertions(+)
 create mode 100644 blender_extensions/middleware.py

diff --git a/blender_extensions/middleware.py b/blender_extensions/middleware.py
new file mode 100644
index 00000000..8cccb17e
--- /dev/null
+++ b/blender_extensions/middleware.py
@@ -0,0 +1,89 @@
+from contextlib import ExitStack
+import json
+import socket
+import time
+
+from django.conf import settings
+from django.db import connections
+
+
+class VectorClient:
+    def __init__(self):
+        self._socket = None
+        self.hostname = socket.gethostname()
+
+    @property
+    def socket(self):
+        if not self._socket:
+            self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        return self._socket
+
+    def send(self, data):
+        self.socket.sendto(
+            bytes(json.dumps(data), "utf-8"),
+            ("127.0.0.1", settings.VECTOR_UDP_PORT),
+        )
+
+
+vector_client = VectorClient()
+
+
+class QueryLogger:
+    def __init__(self, alias):
+        self.alias = alias
+        self.durations = []
+
+    def __call__(self, execute, sql, params, many, context):
+        start = time.perf_counter()
+        try:
+            return execute(sql, params, many, context)
+        finally:
+            duration = time.perf_counter() - start
+            self.durations.append(duration)
+
+
+def metrics_middleware(get_response):
+    """All timings reported as integer milliseconds."""
+
+    def middleware(request):
+        # pre-request setup
+        start = time.perf_counter()
+        context_managers = []
+        query_loggers = []
+        for connection in connections.all():
+            query_logger = QueryLogger(connection.alias)
+            query_loggers.append(query_logger)
+            context_managers.append(connection.execute_wrapper(query_logger))
+
+        # request processing
+        with ExitStack() as stack:
+            for context_manager in context_managers:
+                stack.enter_context(context_manager)
+            response = get_response(request)
+
+        # post-request reporting
+        request_time = int(1000 * (time.perf_counter() - start))
+        db_metrics = {}
+        for query_logger in query_loggers:
+            db_metrics[query_logger.alias] = {
+                "query_count": len(query_logger.durations),
+                "query_time_sum": int(1000 * sum(query_logger.durations)),
+            }
+
+        data = {
+            "db": db_metrics,
+            "hostname": vector_client.hostname,
+            "http": {
+                "path": request.path,
+                "request_time": request_time,
+                "remote_addr": request.META.get("REMOTE_ADDR"),
+                "status_code": response.status_code,
+                "user_agent": request.headers.get("user-agent"),
+            },
+            "service_name": settings.SERVICE_NAME,
+        }
+        vector_client.send(data)
+
+        return response
+
+    return middleware
diff --git a/blender_extensions/settings.py b/blender_extensions/settings.py
index 137e868d..689f5c7f 100644
--- a/blender_extensions/settings.py
+++ b/blender_extensions/settings.py
@@ -79,6 +79,7 @@ INSTALLED_APPS = [
 ]
 
 MIDDLEWARE = [
+    'blender_extensions.middleware.metrics_middleware',
     'django.middleware.security.SecurityMiddleware',
     'django.contrib.sessions.middleware.SessionMiddleware',
     'django.middleware.common.CommonMiddleware',
@@ -354,3 +355,6 @@ if os.environ.get('ADMINS') is not None:
     ADMINS = [[_.strip() for _ in adm.split(':')] for adm in os.environ.get('ADMINS').split(',')]
     EMAIL_SUBJECT_PREFIX = f'[{ALLOWED_HOSTS[0]}]'
     SERVER_EMAIL = f'django@{ALLOWED_HOSTS[0]}'
+
+SERVICE_NAME = os.getenv('SERVICE_NAME', 'dummy_service_name')
+VECTOR_UDP_PORT = 18125
-- 
2.30.2