base telemetry idea

app/api/init_api.py

@@ -25,6 +25,7 @@
 import app.bg_loops
 import app.settings
 import app.state
+import app.telemetry
 import app.utils
 from app.api import api_router  # type: ignore[attr-defined]
 from app.api import domains
@@ -83,6 +84,8 @@ async def lifespan(asgi_app: BanchoAPI) -> AsyncIterator[Never]:
             Ansi.LRED,
         )
 
+    app.telemetry.hook_database_calls()
+
     await app.state.services.database.connect()
     await app.state.services.redis.initialize()
 

app/telemetry.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import dataclasses
+import hashlib
+import json
+import platform
+import time
+from dataclasses import dataclass
+from typing import Any
+
+import psutil
+from httpx import AsyncClient
+from tenacity import retry
+from tenacity.stop import stop_after_attempt
+
+from app import settings
+
+MIN_SLOW_QUERY_SECONDS = 5.0
+
+HTTP_CLIENT = AsyncClient()
+
+
+@dataclass
+class SoftwareInfo:
+    version: str
+    domain: str
+
+
+@dataclass
+class SystemInfo:
+    system: str
+    node: str
+    release: str
+    version: str
+    machine: str
+    processor: str
+
+
+@dataclass
+class LanguageInfo:
+    python_version: str
+    build_no: str
+    build_date: str
+    python_compiler: str
+    python_implementation: str
+
+
+@dataclass
+class SystemLoadInfo:
+    cpu_1min_average: float
+    cpu_5min_average: float
+    cpu_15min_average: float
+
+
+@dataclass
+class TelemetryEventReport:
+    software_info: SoftwareInfo
+    system_info: SystemInfo
+    language_info: LanguageInfo
+    system_load_info: SystemLoadInfo
+    event_data: dict[str, Any]
+
+
+@retry(reraise=True, stop=stop_after_attempt(3))
+async def report_event(event_data: dict[str, Any]) -> None:
+    cpu_1min_average, cpu_5min_average, cpu_15min_average = psutil.getloadavg()
+    event = TelemetryEventReport(
+        software_info=SoftwareInfo(
+            version=settings.VERSION,
+            domain=settings.DOMAIN,
+        ),
+        system_info=SystemInfo(
+            system=platform.system(),
+            node=platform.node(),
+            release=platform.release(),
+            version=platform.version(),
+            machine=platform.machine(),
+            processor=platform.processor(),
+        ),
+        language_info=LanguageInfo(
+            python_version=platform.python_version(),
+            build_no=platform.python_build()[0],
+            build_date=platform.python_build()[1],
+            python_compiler=platform.python_compiler(),
+            python_implementation=platform.python_implementation(),
+        ),
+        system_load_info=SystemLoadInfo(
+            cpu_1min_average=cpu_1min_average,
+            cpu_5min_average=cpu_5min_average,
+            cpu_15min_average=cpu_15min_average,
+        ),
+        event_data=event_data,
+    )
+    request_data = dataclasses.asdict(event)
+    idempotency_key = hashlib.sha256(
+        json.dumps(request_data, sort_keys=True).encode("utf-8"),
+    ).hexdigest()
+    response = await HTTP_CLIENT.post(
+        url="https://telemetry.cmyui.xyz/report",
+        headers={"Idempotency-Key": idempotency_key},
+        json=request_data,
+    )
+    response.raise_for_status()
+    return None
+
+
+def hook_database_calls() -> None:
+    def _wrap_database_call(func):
+        async def wrapper(*args, **kwargs):
+            start_time = time.perf_counter()
+            response = await func(*args, **kwargs)
+            end_time = time.perf_counter()
+
+            seconds_elapsed = end_time - start_time
+
+            if seconds_elapsed >= MIN_SLOW_QUERY_SECONDS:
+                event_data = {
+                    "query": kwargs.get("query", args and args[0]),
+                    "seconds_elapsed": seconds_elapsed,
+                }
+                await report_event(event_data)
+
+            return response
+
+        return wrapper
+
+    import app.state.services
+
+    for attr in ("execute", "execute_many", "fetch_one", "fetch_all"):
+        unwrapped_func = getattr(app.state.services.database, attr)
+        wrapped_func = _wrap_database_call(unwrapped_func)
+        setattr(app.state.services.database, attr, wrapped_func)
+
+    return None