[#12953] feat(mcp-server): select the metalake per tool call - #12960
Conversation
…stead of pinning it at startup The MCP server takes a single metalake at startup and serves it for the life of the process, so a multi-metalake install needs one MCP server per metalake. Add an X-Gravitino-Metalake request header (HTTP transport only) that takes priority over the --metalake startup default, resolved statelessly per request so the server stays correct regardless of how many replicas it runs as. stdio is unchanged - it has no per-request header channel, so --metalake remains required there. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The cell was far longer than every other row and packed in two code spans plus a semicolon clause, breaking the table's rendering. Point to the existing "Per-request metalake (HTTP)" section for the detail instead of duplicating it in the cell. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Pull request overview
This PR enables per-request metalake resolution for the MCP HTTP transport via an X-Gravitino-Metalake header (overriding the startup default), while keeping stdio pinned to a required startup --metalake.
Changes:
- Add
X-Gravitino-Metalakerequest header support and per-(identity, metalake) REST client caching with bounded LRU eviction. - Make
--metalakeoptional for HTTP transport (still required for stdio), and validate metalake/OAuth settings atGravitinoContextconstruction. - Add unit tests and documentation describing per-request metalake behavior and precedence.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| mcp-server/mcp_server/core/context.py | Resolve metalake per request and cache REST clients by (Authorization, metalake) plus service-identity-per-metalake. |
| mcp-server/mcp_server/core/setting.py | Make metalake optional by default, strip whitespace, and add validate_metalake() enforcing stdio requirements. |
| mcp-server/mcp_server/main.py | Accept missing --metalake by default, validate metalake at startup, and log the metalake policy. |
| mcp-server/mcp_server/server.py | Add startup logging describing how metalake selection works (default vs per-request header). |
| mcp-server/tests/unit/test_per_request_metalake.py | Add unit coverage for header resolution, precedence, caching/eviction, and construction-time validation. |
| docs/gravitino-mcp-server.md | Document --metalake semantics change and add a “Per-request metalake (HTTP)” section. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
…t metalake
- Merge the service-identity client cache into _clients_by_auth, keyed by
("", metalake), so _MAX_CACHED_CLIENTS bounds total open connection pools
instead of applying separately to two caches.
- Tighten the cache type annotation to OrderedDict[tuple[str, str], object].
- Make the concurrency test genuinely concurrent: two overlapping asyncio
tasks each with their own request contextvar, verified by mutation test to
fail when resolution leaks across tasks.
- Reformat the configuration table so every row shares the column widths.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Answers the review question directly: required for stdio, optional for HTTP, and spell out the resolution order rather than leaving it implied across three paragraphs. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
|
This solution has some limitations: the metalake cannot be changed during runtime. I need to investigate more. |
Code Coverage Report
Files |
Replace the X-Gravitino-Metalake request header with an optional `metalake` argument on every tool. The header only worked over HTTP and was fixed for the life of a client connection, so an agent could not switch metalake while running; a tool argument works on stdio too and can differ per call. No tool declares the argument. MetalakeArgumentMiddleware advertises it in each tool's input schema, strips it from the incoming arguments, and publishes it in a context variable scoped to one call, which GravitinoContext resolves against: call argument, then the --metalake default, then an error naming the recovery path. Nothing is remembered between calls, so concurrent callers stay isolated and the server needs no shared state across replicas. Add a `list_metalakes` tool so an agent can discover what that argument accepts. It is the one tool that resolves no metalake, so it works on a server started without --metalake - the case that needs it most. Also unify the two statistic tools, which took their own `metalake_name` argument and shadowed the value their client was already built with, and record the resolved metalake in audit entries now that one server can serve several tenants. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
yuqi1129
left a comment
There was a problem hiding this comment.
Reviewed commit 67773f0. All 257 unit tests passed. I also checked concurrent calls over one HTTP MCP session with different identities and metalakes; those remained isolated. Additional probes reproduced the issues below. Please address the connection-lifetime, input-validation, and OAuth refresh issues before merging, and retain compatibility for the statistic tools unless the breaking change has an agreed migration plan. The probes used a local HTTP server and controlled mocks, not a live Gravitino/IdP deployment.
…election - Do not close a cached client while a call is still using it. Eviction scheduled close() immediately, so a slow call lost its connection once other calls filled the cache - and for a write the backend may already have committed. Borrows are now tracked per call and an evicted client is closed by its last user. - Reject an explicitly supplied non-string metalake. Popping the argument removes it from FastMCP's schema validation, so `false`, `0` and `[]` silently routed calls to the default metalake and `42` raised AttributeError outside the audit and error middleware. The raw value is now published and validated during resolution, inside both, before any REST call. - Build the service OAuth auth object once per context. One instance per metalake shared the token cache key but not the refresh lock, so a cold cache hit the IdP once per metalake. - Keep accepting `metalake_name` on the two statistic tools as a deprecated alias, since it has shipped since v1.0.0. It is folded into the shared metalake and no longer advertised; supplying both with different values is rejected. - Correct the audit docstring, which still described the old behaviour, and a docstring typo in the statistic tools. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
What changes were proposed in this pull request?
The MCP server's metalake was fixed at startup (
--metalake) and every toolcall resolved against that one metalake for the life of the process. This PR
makes it a per-call choice: every tool accepts an optional
metalakeargument that takes priority over the startup default.
No tool declares the argument.
MetalakeArgumentMiddlewareadvertises it ineach tool's input schema, strips it from the incoming arguments before the
tool function runs, and publishes it in a context variable scoped to a single
call, which
GravitinoContextresolves against: the call's argument, then the--metalakedefault, then an error. Nothing is remembered between calls, soconcurrent callers stay isolated and the server keeps no state that would have
to be shared between replicas.
A
list_metalakestool is added so an agent can discover what that argumentaccepts. It is the one tool that resolves no metalake, so it works on a server
started with no
--metalakeat all - the deployment that needs it most.Two smaller changes come along:
metalake_nameargument, shadowing thevalue their client was already constructed with. They now use the shared
resolution like every other tool, so there is one way to name a metalake.
touch" is no longer answerable from the server configuration alone. Tools
that are not metalake-scoped record none rather than the default.
An earlier revision of this PR used an
X-Gravitino-Metalakerequest header.That was dropped: the header only worked over HTTP, was fixed for the life of
a client connection (so an agent could not switch metalake mid-conversation),
and depended on MCP clients forwarding custom headers on every request, which
they do inconsistently. A tool argument travels in the JSON-RPC body, works
identically on stdio and HTTP, and can differ per call.
Why are the changes needed?
A metalake is Gravitino's tenant boundary, and every other interface takes it
per request. The MCP server was the only one that pinned it at startup, so an
installation with more than one metalake needed one server process per
metalake, or had to accept that agents could reach only one of them.
Fix: #12953
Does this PR introduce any user-facing change?
--metalakeis now optional on every transport. It is the default used byany call that does not name one.
metalakeargument.list_metalakestool, tagmetalake.list_statistics_for_metadataandlist_statistics_for_partitionno longer accept
metalake_name; callers passmetalakeinstead, like everyother tool.
message telling the agent to call
list_metalakesand retry - or, when a tagfilter hides that tool, to ask the user which metalake to use.
metalakefield.How was this patch tested?
257 unit tests pass;
isort,blackandpylint(10.00/10) are clean.Coverage of the changed modules is 96%, with the middleware, settings and new
tool modules at 100%.
The tests that matter most:
different metalakes each get their own, held simultaneously in flight by a
barrier so the overlap is real. Mutation-tested - replacing the context
variable with process-global state fails these tests.
per-call plumbing into session state.
--metalake:list_metalakesworks, other tools return arecoverable error, and naming a metalake per call works.
the
/api/metalakesendpoint and the response key.