Instrument the metadata service to trace database queries per request

[valayDave]

Each HTTP request to the metadata service triggers one or more database queries. Currently there is no per-request visibility into how many queries are executed, how many rows are returned, or how long queries take. For agent-facing workloads — where a single logical client operation fans out into many HTTP requests, each triggering its own DB activity — this hidden multiplier is the key unknown.

All database access in the metadata service flows through the execute_sql() method in postgres_async_db.py. Adding opt-in query-level logging here would make the full DB cost of any API operation measurable, forming the server-side counterpart to the client-side instrumentation in issue #4.

Goals

• Add opt-in per-request query logging to execute_sql() that records: query count, rows returned, and execution time
• Verify this captures DB activity for at least two different API endpoints (e.g. GET /flows/{id}/runs and GET /flows/{id}/runs/{id}/steps/{name}/tasks)
• Confirm the instrumentation is a no-op when disabled

Instructions

This issue is a venue for questions and design discussion. Keep the instrumentation orthogonal to any changes in API behaviour or schema — this is observability only.

Use the dev stack (see Resources) to run the service locally. The Swagger UI at http://localhost:8080/api/doc is useful for manually triggering individual endpoints while observing query logs.

When you have met the goals, show your work via a fork with an issue / PRs explaining your work. You can add comments with what you achieve on the issue thread. You can reference the Fork/PR/issue in this issue.

Resources

Dev stack:

• Setting Up the Dev Stack — run metaflow-dev up to get a full local environment; use metaflow-dev shell for a pre-configured shell pointing at it
• Swagger UI: http://localhost:8080/api/doc (once the stack is running) — fire individual endpoints and observe query behaviour

Key code:

• services/data/postgres_async_db.py#L248 — execute_sql(), the single path for all DB queries
• services/data/postgres_async_db.py#L149 — AsyncPostgresDB singleton, manages the connection pool
• services/data/db_utils.py#L10 — DBResponse and DBPagination named tuples returned by all queries
• services/metadata_service/api/task.py — a good endpoint to trace end-to-end as a starting point

Docs:

• Setting Up the Dev Stack
• Service architecture overview

[mandeepsingh2007]

hii @valayDave

I'd like to propose an architectural approach for this. Similar to the concerns raised in Issue #4 regarding review cost and core modifications, we should ensure this DB instrumentation has zero performance penalty for normal workloads.

Instead of hardcoding logging logic directly inside execute_sql(), we could use an Opt-In Async Context Manager injected at the API routing layer (e.g., via a lightweight FastAPI/Tornado middleware depending on the stack).

Proposed Flow:

A request comes in with a specific header (e.g., X-Metaflow-Trace-DB: true).

The middleware detects this and initializes a contextvars.ContextVar to hold the query count, rows returned, and execution time.

Inside postgres_async_db.py, execute_sql() simply checks if this context variable is active. If yes, it appends the metrics; if not, it bypasses completely (zero overhead).

The middleware intercepts the response and injects the DB metrics into the HTTP response headers or a specific log stream.

This keeps the instrumentation entirely orthogonal to the schema and business logic. Let me know if this aligns with your expectations

[shahidsi7]

Hi @valayDave, I've been reading through this issue and the approach suggested above by @mandeepsingh2007 . It makes sense to me by using a ContextVar to pass metrics through the request and surfacing them in response headers is clean and keeps things orthogonal.
While exploring execute_sql() I was also thinking about a simpler alternative: just gating a log statement inside execute_sql() behind an environment variable like METAFLOW_DB_QUERY_LOG=true. When disabled it's a no-op, when enabled it logs query count, rows returned, and execution time to a named logger. No middleware, no context passing — stays entirely within the one method the issue points to.
I'm not sure which direction you'd prefer, or whether both could coexist for different use cases. Before I go further I wanted to check does this direction seem worth exploring alongside the ContextVar approach, or would you rather keep it to one design?

Happy to open a fork and demonstrate it if that's useful.

[mandeepsingh2007]

hii @shahidsi7
Thanks for jumping in! The environment variable approach is definitely simpler to implement and great for local, single-user debugging.

However, my concern with a global toggle like METAFLOW_DB_QUERY_LOG=true is how it behaves in a concurrent environment. The issue specifically highlights the need for "per-request visibility".

If we use a process-wide environment variable, enabling it would log all queries across all concurrent HTTP requests hitting the metadata service. In a busy environment with multiple agents/users, the log stream would be highly interleaved, making it nearly impossible to attribute a specific DB query to its parent HTTP request without complex log-parsing and request-ID propagation.

This is exactly where ContextVar shines. Since it provides state isolation per asynchronous task (or thread), it naturally scopes the tracing to the individual HTTP request lifecycle. Furthermore, the middleware approach allows us to aggregate the metrics (e.g., total_queries_count, total_execution_time) and emit just one clean log entry per HTTP request, rather than spamming the logger with individual entries for every single execute_sql() invocation.

Would love to hear @valayDave 's thoughts on whether the goal here is a simple global debug toggle, or true concurrent request-level observability

[Youssef-SH]

Hi @valayDave,

Implemented opt-in per-request DB tracing at execute_sql(), recording query count, rows returned, and execution time. Scoped per request using ContextVar, and a no-op when disabled.

Validation (dev stack):

GET /flows/TestFlow/runs
query_count=1, total_rows=3

GET /flows/TestFlow/runs/3/steps/start/tasks
query_count=2, total_rows=2

PR:
Netflix/metaflow-service#471

[ahmedelgyar1]

Hi @mandeepsingh2007 , great point about the interleaved logs! I completely agree that a simple global toggle wouldn't work well for concurrent workloads.

I actually went ahead and implemented this exact ContextVar + middleware approach. It successfully isolates the queries per request and emits one clean summary log at the end, with zero overhead when the QUERY_TRACING_ENABLED env var is off.

@valayDave I've opened a PR with this implementation here: Netflix/metaflow-service#476 Would love your thoughts on it!