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

docs/gravitino-mcp-server.md

@@ -144,7 +144,7 @@ You could config Gravitino MCP server by arguments, `uv run mcp_server -h` shows
 
 | Argument                         | Description                                                                                                                     | Default value               | Required |
 |----------------------------------|---------------------------------------------------------------------------------------------------------------------------------|-----------------------------|----------|
-| `--metalake`                     | The Gravitino metalake name.                                                                                                    | none                        | Yes      |
+| `--metalake`                     | Default Gravitino metalake, used when a request names none. Required for `stdio`; optional for HTTP, where each request may instead name a metalake via the `X-Gravitino-Metalake` header. | none                        | stdio only |
 | `--gravitino-uri`                | The URI of Gravitino server.                                                                                                    | `http://127.0.0.1:8090`     | No       |
 | `--transport`                    | Transport protocol: stdio (local), http / streamable-http (Streamable HTTP).                                                    | `stdio`                     | No       |
 | `--mcp-url`                      | The URL of MCP server if using HTTP transport.                                                                                  | `http://127.0.0.1:8000/mcp` | No       |
@@ -225,6 +225,14 @@ For exposed or multi-caller HTTP deployments, set `--no-service-identity-fallbac
 
 Authorization itself is always enforced by Gravitino: the MCP server forwards the identity but does not make access-control decisions of its own.
 
+### Per-request metalake (HTTP)
+
+When the server runs with HTTP transport, a request may name the metalake to operate on with the `X-Gravitino-Metalake` header, taking priority over the `--metalake` default configured at startup. This lets one server instance serve more than one metalake: each request independently resolves its own metalake from its own header, so the server holds no per-connection or per-session metalake state and stays correct regardless of how many replicas it runs as.
+
+Falls back to `--metalake` when the header is absent. If neither is set, the call fails with an error naming the missing argument. Authorization is unchanged — the caller's identity (see above) determines what it may see in the requested metalake exactly as it would through the REST API.
+
+stdio transport has no per-request header, so `--metalake` remains the only source there; switching metalake means starting another stdio process with a different `--metalake`.
+
 ### Serving over HTTPS (TLS)
 
 To serve the MCP HTTP endpoint (the `--mcp-url`, not the `--gravitino-uri`) over TLS, provide both `--tls-cert` and `--tls-key` and use an `https://` `--mcp-url`. The certificate and key must be provided together, and the URL scheme must match the TLS setting (an `https://` URL without a cert/key, or a cert/key behind an `http://` URL, is rejected at startup).

mcp-server/mcp_server/core/context.py

@@ -50,6 +50,10 @@
     "negotiate": "Negotiate",
 }
 
+# HTTP header a request uses to name the metalake it wants to operate on.
+# Headers are matched case-insensitively by the underlying request object.
+METALAKE_HEADER = "X-Gravitino-Metalake"
+
 
 class ServiceIdentityFallbackDisabled(RuntimeError):
     """HTTP omitted Authorization while service-identity fallback is disabled."""
@@ -87,6 +91,24 @@ def _get_request_authorization() -> str:
         return ""
 
 
+def _get_request_metalake() -> str:
+    """Return the ``X-Gravitino-Metalake`` header of the current HTTP request.
+
+    Returns an empty string in stdio mode or when the header is absent, in
+    which case the caller falls back to the configured startup default.
+    """
+    try:
+        # Imported lazily: only available within an HTTP request context.
+        # pylint: disable=import-outside-toplevel
+        from fastmcp.server.dependencies import get_http_request
+
+        # Stripped so a whitespace-only header (e.g. an empty templated
+        # value) is treated as absent, not as a literal metalake name.
+        return get_http_request().headers.get(METALAKE_HEADER, "").strip()
+    except (LookupError, RuntimeError):
+        return ""
+
+
 def startup_authorization(setting: Setting) -> str:
     """The static --token rendered as an ``Authorization`` header value.
 
@@ -147,38 +169,64 @@ def _service_auth(setting: Setting):
 
 class GravitinoContext:
     def __init__(self, setting: Setting):
+        # Enforced here (not only in do_main()) so any path that constructs a
+        # GravitinoContext directly - not just the CLI entrypoint - fails
+        # fast on an invalid Setting, matching the pre-per-request-metalake
+        # behavior where a bad Setting couldn't be constructed at all.
+        setting.validate_metalake()
+        setting.validate_oauth()
         self._setting = setting
-        self._default_client = RESTClientFactory.create_rest_client(
-            setting.metalake,
-            setting.gravitino_uri,
-            startup_authorization(setting),
-            auth=_service_auth(setting),
+        # Eagerly built only when a startup default is configured, so the
+        # common single-metalake deployment pays no extra cost. Left unset
+        # (None) when metalake resolution must come from a per-request header
+        # on every call (HTTP transport with no --metalake default).
+        self._default_client = (
+            RESTClientFactory.create_rest_client(
+                setting.metalake,
+                setting.gravitino_uri,
+                startup_authorization(setting),
+                auth=_service_auth(setting),
+            )
+            if setting.metalake
+            else None
         )
-        # LRU cache of per-principal clients keyed by the raw Authorization header.
-        # Safe without locking: rest_client() runs on the single asyncio event
-        # loop and never awaits between lookup and insert.
-        self._clients_by_auth: "OrderedDict[str, object]" = OrderedDict()
+        # LRU cache of per-principal clients keyed by (Authorization header,
+        # metalake). Safe without locking: rest_client() runs on the single
+        # asyncio event loop and never awaits between lookup and insert.
+        self._clients_by_auth: "OrderedDict[tuple, object]" = OrderedDict()
+        # LRU cache of service-identity clients (static token / OAuth) keyed by
+        # metalake, for requests that name a non-default metalake but carry no
+        # per-request Authorization header. The startup default metalake is
+        # served by _default_client instead and never enters this cache.
+        self._service_clients: "OrderedDict[str, object]" = OrderedDict()
         # Strong references to in-flight background close tasks; the event loop
         # only keeps weak references, so without this they could be GC'd before
         # running. Entries are discarded when each task completes.
         self._pending_closes: "set[asyncio.Task]" = set()
 
     def rest_client(self):
-        """Return a REST client carrying the correct identity for this request.
+        """Return a REST client carrying the correct identity and metalake.
+
+        The metalake is resolved per request: an HTTP request's
+        ``X-Gravitino-Metalake`` header takes priority, falling back to the
+        configured startup default (``--metalake``). Raises ``ValueError``
+        when neither is available.
 
-        In HTTP transport mode the incoming request's ``Authorization`` header is
-        forwarded verbatim to Gravitino, taking priority over the static startup
-        token. This keeps concurrent sessions with different principals fully
-        isolated — one principal's identity never leaks into another's calls.
+        Identity resolution is unchanged: in HTTP transport mode the incoming
+        request's ``Authorization`` header is forwarded verbatim to Gravitino,
+        taking priority over the static startup token. This keeps concurrent
+        sessions with different principals and/or metalakes fully isolated —
+        one caller's identity or metalake never leaks into another's calls.
 
-        Falls back to the shared default client (static token or OAuth
+        Falls back to a service-identity client (static token or OAuth
         client-credentials) when:
         - running in stdio mode (no HTTP request context), or
         - the incoming request carries no Authorization header.
 
-        Per-principal clients are cached (and their connection pools reused) so a
-        new pool is not opened on every tool call.
+        Clients are cached per (identity, metalake) combination (and their
+        connection pools reused) so a new pool is not opened on every call.
         """
+        metalake = self._resolve_metalake()
         authorization = _get_request_authorization()
         if not authorization:
             if (
@@ -190,24 +238,68 @@ def rest_client(self):
                     "HTTP request omitted Authorization and "
                     "--no-service-identity-fallback is set"
                 )
-            return self._default_client
+            return self._service_client(metalake)
 
-        cached = self._clients_by_auth.get(authorization)
+        key = (authorization, metalake)
+        cached = self._clients_by_auth.get(key)
         if cached is not None:
-            self._clients_by_auth.move_to_end(authorization)
+            self._clients_by_auth.move_to_end(key)
             return cached
 
         client = RESTClientFactory.create_rest_client(
-            self._setting.metalake,
+            metalake,
             self._setting.gravitino_uri,
             authorization,
         )
-        self._clients_by_auth[authorization] = client
-        if len(self._clients_by_auth) > _MAX_CACHED_CLIENTS:
-            _, evicted = self._clients_by_auth.popitem(last=False)
-            self._schedule_close(evicted)
+        self._cache_put(self._clients_by_auth, key, client)
         return client
 
+    def _resolve_metalake(self) -> str:
+        """Resolve the metalake for the current call, header first.
+
+        Raises ``ValueError`` (an invalid/missing request parameter, mapped
+        by FastMCP's error middleware to a client-facing "Invalid params"
+        error rather than an internal-error code) when the request names
+        none and no startup default (``--metalake``) is configured.
+        """
+        metalake = _get_request_metalake() or self._setting.metalake
+        if not metalake:
+            raise ValueError(
+                f"No metalake specified: the request omitted the "
+                f"{METALAKE_HEADER!r} header and no --metalake default is "
+                "configured."
+            )
+        return metalake
+
+    def _service_client(self, metalake: str):
+        """Return the service-identity client (static token / OAuth) for ``metalake``."""
+        if (
+            metalake == self._setting.metalake
+            and self._default_client is not None
+        ):
+            return self._default_client
+
+        cached = self._service_clients.get(metalake)
+        if cached is not None:
+            self._service_clients.move_to_end(metalake)
+            return cached
+
+        client = RESTClientFactory.create_rest_client(
+            metalake,
+            self._setting.gravitino_uri,
+            startup_authorization(self._setting),
+            auth=_service_auth(self._setting),
+        )
+        self._cache_put(self._service_clients, metalake, client)
+        return client
+
+    def _cache_put(self, cache: "OrderedDict", key, client) -> None:
+        """Insert into an LRU cache, evicting (and closing) the oldest past the cap."""
+        cache[key] = client
+        if len(cache) > _MAX_CACHED_CLIENTS:
+            _, evicted = cache.popitem(last=False)
+            self._schedule_close(evicted)
+
     def _schedule_close(self, client) -> None:
         """Best-effort close of an evicted client's connection pool.
 

mcp-server/mcp_server/core/setting.py

@@ -28,7 +28,11 @@ class DefaultSetting:
 
 @dataclass
 class Setting:  # pylint: disable=too-many-instance-attributes
-    metalake: str
+    # Default metalake used when a request names none. Optional for HTTP
+    # transport, where each request can instead name a metalake via the
+    # X-Gravitino-Metalake header; required for stdio, which has no
+    # per-request channel to name one.
+    metalake: str = ""
     gravitino_uri: str = DefaultSetting.default_gravitino_uri
     tags: Set[str] = field(default_factory=set)
     transport: str = DefaultSetting.default_transport
@@ -56,6 +60,12 @@ class Setting:  # pylint: disable=too-many-instance-attributes
     # --token is configured instead of falling back to the service identity.
     no_service_identity_fallback: bool = False
 
+    def __post_init__(self) -> None:
+        # A whitespace-only --metalake (e.g. a shell-quoting mistake) must be
+        # treated as "no default configured", the same as an empty string,
+        # rather than silently used as a nonsensical metalake name.
+        self.metalake = self.metalake.strip()
+
     def has_oauth_client(self) -> bool:
         """Return True when client-credentials is fully configured."""
         return bool(
@@ -68,6 +78,19 @@ def has_service_identity(self) -> bool:
         """Return True when a static token or OAuth client-credentials is set."""
         return bool(self.token.strip()) or self.has_oauth_client()
 
+    def validate_metalake(self) -> None:
+        """Reject stdio transport with no default metalake configured.
+
+        stdio has no per-request channel to name a metalake, so --metalake is
+        the only source there. HTTP transport can rely on the per-request
+        X-Gravitino-Metalake header instead, so an empty default is legitimate.
+        """
+        if self.transport == "stdio" and not self.metalake:
+            raise ValueError(
+                "--metalake is required for stdio transport (stdio has no "
+                "per-request way to select a metalake)."
+            )
+
     def validate_oauth(self) -> None:
         """Reject a partial OAuth client-credentials configuration."""
         filled = [

mcp-server/mcp_server/main.py

@@ -22,6 +22,7 @@
 from mcp_server.core.setting import DefaultSetting, Setting
 from mcp_server.server import (
     GravitinoMCPServer,
+    log_metalake_policy,
     log_service_identity_fallback_policy,
 )
 from mcp_server.tools import SUPPORTED_TOOL_TAGS
@@ -50,10 +51,12 @@ def do_main():
     )
     _init_logging(setting)
     try:
+        setting.validate_metalake()
         setting.validate_oauth()
     except ValueError as exc:
         logging.error("%s", exc)
         raise SystemExit(1) from None
+    log_metalake_policy(setting)
     log_service_identity_fallback_policy(setting)
     logging.info("Gravitino MCP server setting: %s", setting)
     server = GravitinoMCPServer(setting)
@@ -90,8 +93,11 @@ def _parse_args():
     parser.add_argument(
         "--metalake",
         type=str,
-        required=True,
-        help="Gravitino metalake name.",
+        default="",
+        help="Default Gravitino metalake name, used when a request names "
+        "none. Required for stdio transport. Optional for HTTP transport, "
+        "where each request can instead name a metalake via the "
+        "X-Gravitino-Metalake header.",
     )
     parser.add_argument(
         "--gravitino-uri",

mcp-server/mcp_server/server.py

@@ -37,6 +37,7 @@
 
 from mcp_server.core import audit
 from mcp_server.core.context import (
+    METALAKE_HEADER,
     GravitinoContext,
     _get_request_authorization,
     service_fallback_authorization,
@@ -152,6 +153,26 @@ def log_service_identity_fallback_policy(setting: Setting) -> None:
         )
 
 
+def log_metalake_policy(setting: Setting) -> None:
+    """Log how the metalake is resolved for requests at startup."""
+    if setting.metalake:
+        if setting.transport != "stdio":
+            logging.info(
+                "Default metalake '%s' configured; HTTP requests may "
+                "override it with the %s header.",
+                setting.metalake,
+                METALAKE_HEADER,
+            )
+        return
+    # No default: only reachable for HTTP transport (validate_metalake
+    # rejects an empty metalake for stdio before this runs).
+    logging.info(
+        "No default --metalake configured; every request must name one "
+        "via the %s header.",
+        METALAKE_HEADER,
+    )
+
+
 def _parse_mcp_url(url: str) -> tuple[str, int, str]:
     try:
         parsed = urlparse(url)