decnet/telemetry.py

@@ -209,163 +209,38 @@ def _wrap(fn: F, span_name: str | None) -> F:
 # ---------------------------------------------------------------------------
 
 def wrap_repository(repo: Any) -> Any:
-    """Wrap *repo* in a tracing proxy.  Returns *repo* unchanged when disabled."""
+    """Wrap *repo* in a dynamic tracing proxy.  Returns *repo* unchanged when disabled.
+
+    Instead of mirroring every method signature (which drifts when concrete
+    repos add extra kwargs beyond the ABC), this proxy introspects the inner
+    repo at construction time and wraps every public async method in a span
+    via ``__getattr__``.  Sync attributes are forwarded directly.
+    """
     if not _ENABLED:
         return repo
 
-    from decnet.web.db.repository import BaseRepository
+    tracer = get_tracer("db")
 
-    class TracedRepository(BaseRepository):
+    class TracedRepository:
-        """Proxy that creates a DB span around every BaseRepository call."""
+        """Dynamic proxy — wraps every async method call in a DB span."""
 
-        def __init__(self, inner: BaseRepository) -> None:
+        def __init__(self, inner: Any) -> None:
             self._inner = inner
-            self._tracer = get_tracer("db")
-
-        # --- Forward every ABC method through a span ---
-
-        async def initialize(self) -> None:
-            with self._tracer.start_as_current_span("db.initialize"):
-                return await self._inner.initialize()
-
-        async def add_log(self, log_data):
-            with self._tracer.start_as_current_span("db.add_log"):
-                return await self._inner.add_log(log_data)
-
-        async def get_logs(self, limit=50, offset=0, search=None):
-            with self._tracer.start_as_current_span("db.get_logs") as span:
-                span.set_attribute("db.limit", limit)
-                span.set_attribute("db.offset", offset)
-                return await self._inner.get_logs(limit=limit, offset=offset, search=search)
-
-        async def get_total_logs(self, search=None):
-            with self._tracer.start_as_current_span("db.get_total_logs"):
-                return await self._inner.get_total_logs(search=search)
-
-        async def get_stats_summary(self):
-            with self._tracer.start_as_current_span("db.get_stats_summary"):
-                return await self._inner.get_stats_summary()
-
-        async def get_deckies(self):
-            with self._tracer.start_as_current_span("db.get_deckies"):
-                return await self._inner.get_deckies()
-
-        async def get_user_by_username(self, username):
-            with self._tracer.start_as_current_span("db.get_user_by_username"):
-                return await self._inner.get_user_by_username(username)
-
-        async def get_user_by_uuid(self, uuid):
-            with self._tracer.start_as_current_span("db.get_user_by_uuid"):
-                return await self._inner.get_user_by_uuid(uuid)
-
-        async def create_user(self, user_data):
-            with self._tracer.start_as_current_span("db.create_user"):
-                return await self._inner.create_user(user_data)
-
-        async def update_user_password(self, uuid, password_hash, must_change_password=False):
-            with self._tracer.start_as_current_span("db.update_user_password"):
-                return await self._inner.update_user_password(uuid, password_hash, must_change_password)
-
-        async def list_users(self):
-            with self._tracer.start_as_current_span("db.list_users"):
-                return await self._inner.list_users()
-
-        async def delete_user(self, uuid):
-            with self._tracer.start_as_current_span("db.delete_user"):
-                return await self._inner.delete_user(uuid)
-
-        async def update_user_role(self, uuid, role):
-            with self._tracer.start_as_current_span("db.update_user_role"):
-                return await self._inner.update_user_role(uuid, role)
-
-        async def purge_logs_and_bounties(self):
-            with self._tracer.start_as_current_span("db.purge_logs_and_bounties"):
-                return await self._inner.purge_logs_and_bounties()
-
-        async def add_bounty(self, bounty_data):
-            with self._tracer.start_as_current_span("db.add_bounty"):
-                return await self._inner.add_bounty(bounty_data)
-
-        async def get_bounties(self, limit=50, offset=0, bounty_type=None, search=None):
-            with self._tracer.start_as_current_span("db.get_bounties") as span:
-                span.set_attribute("db.limit", limit)
-                span.set_attribute("db.offset", offset)
-                return await self._inner.get_bounties(limit=limit, offset=offset, bounty_type=bounty_type, search=search)
-
-        async def get_total_bounties(self, bounty_type=None, search=None):
-            with self._tracer.start_as_current_span("db.get_total_bounties"):
-                return await self._inner.get_total_bounties(bounty_type=bounty_type, search=search)
-
-        async def get_state(self, key):
-            with self._tracer.start_as_current_span("db.get_state") as span:
-                span.set_attribute("db.state_key", key)
-                return await self._inner.get_state(key)
-
-        async def set_state(self, key, value):
-            with self._tracer.start_as_current_span("db.set_state") as span:
-                span.set_attribute("db.state_key", key)
-                return await self._inner.set_state(key, value)
-
-        async def get_max_log_id(self):
-            with self._tracer.start_as_current_span("db.get_max_log_id"):
-                return await self._inner.get_max_log_id()
-
-        async def get_logs_after_id(self, last_id, limit=500):
-            with self._tracer.start_as_current_span("db.get_logs_after_id") as span:
-                span.set_attribute("db.last_id", last_id)
-                span.set_attribute("db.limit", limit)
-                return await self._inner.get_logs_after_id(last_id, limit=limit)
-
-        async def get_all_bounties_by_ip(self):
-            with self._tracer.start_as_current_span("db.get_all_bounties_by_ip"):
-                return await self._inner.get_all_bounties_by_ip()
-
-        async def get_bounties_for_ips(self, ips):
-            with self._tracer.start_as_current_span("db.get_bounties_for_ips") as span:
-                span.set_attribute("db.ip_count", len(ips))
-                return await self._inner.get_bounties_for_ips(ips)
-
-        async def upsert_attacker(self, data):
-            with self._tracer.start_as_current_span("db.upsert_attacker"):
-                return await self._inner.upsert_attacker(data)
-
-        async def upsert_attacker_behavior(self, attacker_uuid, data):
-            with self._tracer.start_as_current_span("db.upsert_attacker_behavior"):
-                return await self._inner.upsert_attacker_behavior(attacker_uuid, data)
-
-        async def get_attacker_behavior(self, attacker_uuid):
-            with self._tracer.start_as_current_span("db.get_attacker_behavior"):
-                return await self._inner.get_attacker_behavior(attacker_uuid)
-
-        async def get_behaviors_for_ips(self, ips):
-            with self._tracer.start_as_current_span("db.get_behaviors_for_ips") as span:
-                span.set_attribute("db.ip_count", len(ips))
-                return await self._inner.get_behaviors_for_ips(ips)
-
-        async def get_attacker_by_uuid(self, uuid):
-            with self._tracer.start_as_current_span("db.get_attacker_by_uuid"):
-                return await self._inner.get_attacker_by_uuid(uuid)
-
-        async def get_attackers(self, limit=50, offset=0, search=None, sort_by="recent", service=None):
-            with self._tracer.start_as_current_span("db.get_attackers") as span:
-                span.set_attribute("db.limit", limit)
-                span.set_attribute("db.offset", offset)
-                return await self._inner.get_attackers(limit=limit, offset=offset, search=search, sort_by=sort_by, service=service)
-
-        async def get_total_attackers(self, search=None, service=None):
-            with self._tracer.start_as_current_span("db.get_total_attackers"):
-                return await self._inner.get_total_attackers(search=search, service=service)
-
-        async def get_attacker_commands(self, uuid, limit=50, offset=0, service=None):
-            with self._tracer.start_as_current_span("db.get_attacker_commands") as span:
-                span.set_attribute("db.limit", limit)
-                span.set_attribute("db.offset", offset)
-                return await self._inner.get_attacker_commands(uuid, limit=limit, offset=offset, service=service)
-
-        # --- Catch-all for methods defined on concrete subclasses but not
-        #     in the ABC (e.g. get_log_histogram). ---
 
         def __getattr__(self, name: str) -> Any:
-            return getattr(self._inner, name)
+            attr = getattr(self._inner, name)
+
+            if asyncio.iscoroutinefunction(attr):
+                @functools.wraps(attr)
+                async def _traced_method(*args: Any, **kwargs: Any) -> Any:
+                    with tracer.start_as_current_span(f"db.{name}") as span:
+                        try:
+                            return await attr(*args, **kwargs)
+                        except Exception as exc:
+                            span.record_exception(exc)
+                            raise
+                return _traced_method
+
+            return attr
 
     return TracedRepository(repo)

tests/conftest.py

@@ -9,6 +9,7 @@ import os
 os.environ["DECNET_JWT_SECRET"] = "stable-test-secret-key-at-least-32-chars-long"
 os.environ["DECNET_ADMIN_PASSWORD"] = "test-password-123"
 os.environ["DECNET_DEVELOPER"] = "true"
+os.environ["DECNET_DEVELOPER_TRACING"] = "false"
 os.environ["DECNET_DB_TYPE"] = "sqlite"
 
 import pytest