From 98539210c886391162e703df2cd977e3bb5cbd40 Mon Sep 17 00:00:00 2001
From: moe-berri <moe@berri.ai>
Date: Fri, 6 Nov 2026 08:35:15 +0800
Subject: [PATCH] fix: preserve tracing for bounded activity queries

---
 .../repositories/bounded_daily_activity.py    | 10 +++--
 tests/unit/proxy/db/test_prisma_query_span.py | 45 ++++++++++++++++++-
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/litellm/repositories/bounded_daily_activity.py b/litellm/repositories/bounded_daily_activity.py
index c0832c8e2df..c3972bd2b45 100644
--- a/litellm/repositories/bounded_daily_activity.py
+++ b/litellm/repositories/bounded_daily_activity.py
@@ -67,12 +67,16 @@ class BoundedDailyActivityRepository:
     async def aggregate(
         self, scope: DailyActivityScope, *, model_limit: int | None = None
     ) -> BoundedDailyActivityResponse:
+        from litellm.proxy.db.db_span import db_span
+        from litellm.proxy.db.prisma_query_span import sql_relation
+
         if model_limit is not None and not 1 <= model_limit <= 100:
             raise ValueError("model_limit must be between 1 and 100")
         query: Final = build_bounded_daily_activity_sql(scope, model_limit=model_limit)
-        rows: Final = TypeAdapter(tuple[_ActivityRow, ...]).validate_python(
-            await self.database.query_raw(query.sql, *query.params)
-        )
+        async with db_span("daily_activity_query", sql_relation(query.sql)):
+            rows: Final = TypeAdapter(tuple[_ActivityRow, ...]).validate_python(
+                await self.database.query_raw(query.sql, *query.params)
+            )
         total: Final = next(row for row in rows if row.kind == "total")
         other: Final = next((row for row in rows if row.kind == "other"), None)
         days: Final = tuple(
diff --git a/tests/unit/proxy/db/test_prisma_query_span.py b/tests/unit/proxy/db/test_prisma_query_span.py
index 55ba03d9f2c..9520a3f7c8b 100644
--- a/tests/unit/proxy/db/test_prisma_query_span.py
+++ b/tests/unit/proxy/db/test_prisma_query_span.py
@@ -284,7 +284,10 @@ def _parameters(parents: tuple[ast.AST, ...]) -> frozenset[str]:
 
 
 def _argument_for(call: ast.Call, function: ast.AsyncFunctionDef | ast.FunctionDef, parameter: str) -> ast.expr | None:
-    positional: Final = tuple(arg.arg for arg in function.args.args)
+    bound: Final = (
+        isinstance(call.func, ast.Attribute) and bool(function.args.args) and function.args.args[0].arg == "self"
+    )
+    positional: Final = tuple(arg.arg for arg in function.args.args[1 if bound else 0 :])
     by_keyword: Final = next((k.value for k in call.keywords if k.arg == parameter), None)
     if by_keyword is not None or parameter not in positional:
         return by_keyword
@@ -324,7 +327,16 @@ def _caller_site(
     wrapped: Final = _wrapper_site(call, parents, method, location, module)
     if wrapped is not None:
         return wrapped
-    text: Final = _sql_text(_argument_for(call, function, parameter), _scope(module, parents))
+    argument: Final = _argument_for(call, function, parameter)
+    caller: Final = next((p for p in parents if isinstance(p, (ast.AsyncFunctionDef, ast.FunctionDef))), None)
+    if (
+        caller is not None
+        and caller.name in _RAW_METHODS
+        and isinstance(argument, ast.Name)
+        and argument.id in _parameters(parents)
+    ):
+        return _PrismaCallSite(location, method, "forwarder", f"(callers of {caller.name})")
+    text: Final = _sql_text(argument, _scope(module, parents))
     return _PrismaCallSite(location, method, "engine", _render_statements(method, text) if text is not None else None)
 
 
@@ -449,6 +461,35 @@ def prisma_call_sites() -> tuple[_PrismaCallSite, ...]:
     return tuple(site for module in _modules() for site in _call_sites(module))  # comprehension-ok: flatten
 
 
+@pytest.mark.parametrize(
+    ("caller", "expression", "expected"),
+    [
+        ("query_raw", "query", "(callers of query_raw)"),
+        ("query_raw", "query=query", "(callers of query_raw)"),
+        ("run", "'SELECT * FROM LiteLLM_UserTable'", "postgres.select LiteLLM_UserTable"),
+        ("run", "'not a SQL statement'", None),
+        ("run", "query", None),
+    ],
+)
+def test_query_wrapper_keeps_literal_and_forwarded_statements_checked(
+    caller: str, expression: str, expected: str | None
+) -> None:
+    tree: Final = ast.parse(
+        f"""
+class Database:
+    async def {caller}(self, query):
+        return await self._execute({expression})
+
+    async def _execute(self, query):
+        return await engine.query_raw(query)
+"""
+    )
+    module: Final = _Module(_REPO / "litellm" / "query_wrapper.py", tree, _assignments(tree.body))
+    sites: Final = tuple(_call_sites(module))
+    assert len(sites) == 1
+    assert sites[0].rendered == expected
+
+
 def test_every_prisma_call_site_in_the_proxy_renders_a_bounded_postgres_span_name() -> None:
     """A raw ``query_raw``/``execute_raw``/``query_first`` or a direct model call that no producer
     wraps is named by the engine from its payload; this scan replays that naming (and the wrappers')