Skip to content

[#12953] feat(mcp-server): select the metalake per tool call - #12960

Merged
jerryshao merged 6 commits into
apache:mainfrom
jerryshao:mcp-metalake-per-request
Sep 9, 2026
Merged

[#12953] feat(mcp-server): select the metalake per tool call#12960
jerryshao merged 6 commits into
apache:mainfrom
jerryshao:mcp-metalake-per-request

Conversation

@jerryshao

@jerryshao jerryshao commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

The MCP server's metalake was fixed at startup (--metalake) and every tool
call resolved against that one metalake for the life of the process. This PR
makes it a per-call choice: every tool accepts an optional metalake
argument that takes priority over the startup default.

No tool declares the argument. MetalakeArgumentMiddleware advertises it in
each 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 GravitinoContext resolves against: the call's argument, then the
--metalake default, then an error. Nothing is remembered between calls, so
concurrent callers stay isolated and the server keeps no state that would have
to be shared between replicas.

A list_metalakes tool is added 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 with no --metalake at all - the deployment that needs it most.

Two smaller changes come along:

  • The statistic tools took their own metalake_name argument, shadowing the
    value 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.
  • Audit entries record the resolved metalake, since "which tenant did this
    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-Metalake request 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?

  • --metalake is now optional on every transport. It is the default used by
    any call that does not name one.
  • Every metalake-scoped tool accepts an optional metalake argument.
  • New list_metalakes tool, tag metalake.
  • Breaking: list_statistics_for_metadata and list_statistics_for_partition
    no longer accept metalake_name; callers pass metalake instead, like every
    other tool.
  • A call that names no metalake on a server with no default now fails with a
    message telling the agent to call list_metalakes and retry - or, when a tag
    filter hides that tool, to ask the user which metalake to use.
  • Audit records gain a metalake field.

How was this patch tested?

257 unit tests pass; isort, black and pylint (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:

  • Concurrency, through the real MCP protocol: two overlapping calls naming
    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.
  • No leakage between sequential calls, which is what would silently turn
    per-call plumbing into session state.
  • Cold start with no --metalake: list_metalakes works, other tools return a
    recoverable error, and naming a metalake per call works.
  • The new client operation is exercised against a mocked httpx client, pinning
    the /api/metalakes endpoint and the response key.

…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>
Copilot AI lite review requested due to automatic review settings September 7, 2026 09:17
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>
Comment thread docs/gravitino-mcp-server.md Outdated
Comment thread docs/gravitino-mcp-server.md Outdated

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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-Metalake request header support and per-(identity, metalake) REST client caching with bounded LRU eviction.
  • Make --metalake optional for HTTP transport (still required for stdio), and validate metalake/OAuth settings at GravitinoContext construction.
  • 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.

Comment thread docs/gravitino-mcp-server.md Outdated
Comment thread mcp-server/mcp_server/core/context.py Outdated
Comment thread mcp-server/mcp_server/core/context.py Outdated
Comment thread mcp-server/tests/unit/test_per_request_metalake.py Outdated
Comment thread mcp-server/mcp_server/core/context.py Outdated
jerryshao and others added 2 commits September 7, 2026 17:56
…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>
@jerryshao

Copy link
Copy Markdown
Contributor Author

This solution has some limitations: the metalake cannot be changed during runtime. I need to investigate more.

@github-actions

github-actions Bot commented Sep 7, 2026

Copy link
Copy Markdown

Code Coverage Report

Overall Project 69.7% +0.55% 🟢
Files changed 78.44% 🟢

Module Coverage
aliyun 19.74% 🔴
api 51.61% 🟢
authorization-common 85.96% 🟢
authorization-ranger 4.38% 🔴
aws 53.54% 🟢
azure 32.1% 🔴
catalog-common 27.9% -0.04% 🔴
catalog-fileset 82.17% -1.05% 🟢
catalog-glue 70.84% +2.38% 🟢
catalog-hive 82.96% -2.87% 🟢
catalog-jdbc-common 45.09% +0.7% 🟢
catalog-jdbc-doris 82.69% +2.28% 🟢
catalog-jdbc-mysql 79.33% +13.02% 🟢
catalog-jdbc-postgresql 83.83% +5.05% 🟢
catalog-jdbc-starrocks 79.16% 🟢
catalog-kafka 76.99% +18.92% 🟢
catalog-lakehouse-generic 60.88% -29.12% 🟢
catalog-lakehouse-hudi 79.1% 🟢
catalog-lakehouse-iceberg 85.97% 🟢
catalog-lakehouse-paimon 84.29% -2.98% 🟢
catalog-model 77.99% +44.06% 🟢
cli 44.51% 🟢
client-java 77.5% 🟢
common 57.86% 🟢
core 84.36% 🟢
filesystem-hadoop3 76.48% 🟢
flink 0.0% 🔴
flink-common 53.22% 🟢
flink-runtime 0.0% 🔴
gcp 32.2% 🔴
hadoop-auth 68.0% 🟢
hadoop-common 17.84% -0.34% 🔴
hive-metastore-common 53.52% +4.58% 🟢
iceberg-aliyun-bundle 0.0% 🔴
iceberg-common 66.89% -0.06% 🟢
iceberg-rest-server 76.72% +2.02% 🟢
idp-basic 86.75% 🟢
integration-test-common 0.0% 🔴
jobs 62.92% 🟢
lance-common 34.02% 🔴
lance-rest-server 68.12% 🟢
lineage 59.39% 🟢
optimizer 83.17% 🟢
optimizer-api 21.95% 🔴
server 89.87% 🟢
server-common 81.35% 🟢
spark 56.27% 🟢
tencent 81.78% 🟢
trino-connector 58.36% +0.22% 🟢
Files
Module File Coverage
catalog-common IcebergConstants.java 0.0% 🔴
catalog-fileset FilesetCatalogOperations.java 82.0% 🟢
catalog-glue GlueExceptionConverter.java 100.0% 🟢
GlueTableOperations.java 90.52% 🟢
catalog-hive TrinoNativeViewCodec.java 92.96% 🟢
HiveCatalogOperations.java 81.64% 🟢
HiveViewCatalogOperations.java 78.9% 🟢
HiveTableOperations.java 77.33% 🟢
catalog-jdbc-common JdbcTable.java 70.37% 🟢
catalog-jdbc-doris DorisTableOperations.java 84.91% 🟢
catalog-jdbc-mysql MysqlTableOperations.java 85.68% 🟢
catalog-jdbc-postgresql PostgreSqlTableOperations.java 86.99% 🟢
catalog-kafka KafkaCatalogOperations.java 81.03% 🟢
catalog-lakehouse-generic LanceTableOperations.java 62.92% 🟢
GenericCatalogOperations.java 28.69% 🔴
catalog-lakehouse-paimon CatalogUtils.java 43.24% 🔴
catalog-model ModelCatalogOperations.java 85.9% 🟢
hadoop-common FileSystemUtils.java 33.33% 🔴
HDFSFileSystemProxy.java 0.0% 🔴
hive-metastore-common HiveClientFactory.java 85.22% 🟢
ProxyHiveClientImpl.java 80.95% 🟢
HiveExceptionConverter.java 80.72% 🟢
HmsKerberosClient.java 80.7% 🟢
Util.java 71.79% 🟢
HiveClientImpl.java 0.0% 🔴
iceberg-common IcebergCatalogUtil.java 66.41% 🟢
iceberg-rest-server IcebergCatalogWrapperManager.java 90.48% 🟢
FederatedCatalogWrapper.java 90.28% 🟢
IcebergRESTUtils.java 88.73% 🟢
IcebergTableOperations.java 81.28% 🟢
trino-connector MySQLExternalDataType.java 93.1% 🟢

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>
@jerryshao jerryshao changed the title [#12953] feat(mcp-server): resolve metalake per HTTP request instead of pinning it at startup [#12953] feat(mcp-server): select the metalake per tool call Sep 8, 2026
@jerryshao
jerryshao marked this pull request as ready for review September 8, 2026 07:24
@jerryshao
jerryshao requested a lite review from Copilot September 8, 2026 07:54

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟢 Approval recommended

The functional changes are well-covered by targeted tests; remaining findings are minor docstring/spelling fixes.

Review details
  • Files reviewed: 22/22 changed files
  • Comments generated: 2
  • Review effort level: Lite

Comment thread mcp-server/mcp_server/core/audit.py Outdated
Comment thread mcp-server/mcp_server/tools/statistic.py Outdated

@yuqi1129 yuqi1129 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread mcp-server/mcp_server/core/context.py Outdated
Comment thread mcp-server/mcp_server/core/middleware.py
Comment thread mcp-server/mcp_server/core/context.py Outdated
Comment thread mcp-server/mcp_server/tools/statistic.py
…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>
@jerryshao
jerryshao requested a review from yuqi1129 September 9, 2026 08:46
@jerryshao jerryshao self-assigned this Sep 9, 2026
@jerryshao
jerryshao merged commit 7694553 into apache:main Sep 9, 2026
35 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Improvement] Allow the MCP server to serve more than one metalake per instance

3 participants