This is an automated email from the ASF dual-hosted git repository.
mchades pushed a commit to branch branch-1.3
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/branch-1.3 by this push:
new d853f1a0a5 [Cherry-pick to branch-1.3]
[#11565][#11566][#11568][#11572][#11575] feat(mcp-server): authentication,
per-request authorization and audit logging (#11622) (#11689)
d853f1a0a5 is described below
commit d853f1a0a58c7cc44d9f1e020ddb43315657e401
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Wed Jun 17 09:44:25 2026 +0800
[Cherry-pick to branch-1.3] [#11565][#11566][#11568][#11572][#11575]
feat(mcp-server): authentication, per-request authorization and audit logging
(#11622) (#11689)
**Cherry-pick Information:**
- Original commit: a98e74bc5fbccf4a6c56002b7668142d11f568e6
- Target branch: `branch-1.3`
- Status: ✅ Clean cherry-pick (no conflicts)
Co-authored-by: Qi Yu <[email protected]>
---
.github/workflows/mcp-integration-test.yml | 82 +++++++
mcp-server/.gitignore | 3 +
mcp-server/dev/INSPECTOR_DEMO.md | 208 +++++++++++++++++
mcp-server/dev/run_authz_integration_test.sh | 179 ++++++++++++++
mcp-server/dev/start_inspector_demo.sh | 260 +++++++++++++++++++++
mcp-server/dev/stop_inspector_demo.sh | 77 ++++++
mcp-server/mcp_server/client/factory.py | 7 +-
.../client/plain/plain_rest_client_operation.py | 23 +-
mcp-server/mcp_server/core/audit.py | 81 +++++++
mcp-server/mcp_server/core/context.py | 118 +++++++++-
mcp-server/mcp_server/core/setting.py | 21 +-
mcp-server/mcp_server/main.py | 46 +++-
mcp-server/mcp_server/server.py | 123 +++++++++-
mcp-server/pyproject.toml | 7 +
.../context.py => tests/integration/__init__.py} | 13 --
mcp-server/tests/integration/conftest.py | 73 ++++++
mcp-server/tests/integration/gravitino_setup.py | 145 ++++++++++++
mcp-server/tests/integration/test_authz_e2e.py | 175 ++++++++++++++
mcp-server/tests/unit/test_audit.py | 232 ++++++++++++++++++
mcp-server/tests/unit/test_auth_flow.py | 192 +++++++++++++++
mcp-server/tests/unit/test_per_request_token.py | 242 +++++++++++++++++++
mcp-server/tests/unit/test_transport_tls.py | 153 ++++++++++++
mcp-server/tests/unit/tools/mock_operation.py | 2 +-
23 files changed, 2427 insertions(+), 35 deletions(-)
diff --git a/.github/workflows/mcp-integration-test.yml
b/.github/workflows/mcp-integration-test.yml
new file mode 100644
index 0000000000..959d439e82
--- /dev/null
+++ b/.github/workflows/mcp-integration-test.yml
@@ -0,0 +1,82 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+name: mcp-integration-test
+
+on:
+ push:
+ branches: [ "main", "branch-*" ]
+ pull_request:
+ branches: [ "main", "branch-*" ]
+
+concurrency:
+ group: ${{ github.workflow }}-${{ github.event.pull_request.number ||
github.ref }}
+ cancel-in-progress: true
+
+jobs:
+ changes:
+ runs-on: ubuntu-latest
+ outputs:
+ mcp_or_authz_changes: ${{ steps.filter.outputs.mcp_or_authz_changes }}
+ steps:
+ - uses: actions/checkout@v4
+ - uses: dorny/paths-filter@de90cc6fb38fc0963ad72b210f1f284cd68cea36
+ id: filter
+ with:
+ filters: |
+ mcp_or_authz_changes:
+ - 'mcp-server/**'
+ - 'server/**'
+ - 'core/src/main/java/org/apache/gravitino/authorization/**'
+ -
'server-common/src/main/java/org/apache/gravitino/server/authentication/**'
+ - '.github/workflows/mcp-integration-test.yml'
+
+ mcp-authz-integration-test:
+ runs-on: ubuntu-latest
+ timeout-minutes: 60
+ needs: changes
+ if: needs.changes.outputs.mcp_or_authz_changes == 'true'
+ steps:
+ - uses: actions/checkout@v4
+
+ - uses: ./.github/actions/setup-java-toolchains
+ with:
+ java-version: 17
+
+ - name: Install uv
+ uses: astral-sh/setup-uv@v5
+
+ - name: Build Gravitino distribution
+ run: ./gradlew compileDistribution -x test -PskipWeb=true
+
+ - name: Run MCP authorization integration test
+ env:
+ GRAVITINO_HOME: ${{ github.workspace }}/distribution/package
+ run: |
+ ./mcp-server/dev/run_authz_integration_test.sh
+
+ - name: Upload server logs on failure
+ if: failure()
+ uses: actions/upload-artifact@v4
+ with:
+ name: mcp-authz-it-logs
+ path: |
+ distribution/package/logs/**
+ mcp-server/gravitino-mcp-audit.log
+ mcp-server/gravitino-mcp.log
+ if-no-files-found: ignore
diff --git a/mcp-server/.gitignore b/mcp-server/.gitignore
index f2a99ec686..6ec30088be 100644
--- a/mcp-server/.gitignore
+++ b/mcp-server/.gitignore
@@ -34,6 +34,9 @@ logs/
*.err
gravitino_mcp_server.egg-info
+# Local run artifacts from dev/start_inspector_demo.sh (.out logs and .pid
files)
+.inspector-demo-*
+
# Unit test / coverage reports
htmlcov/
.tox/
diff --git a/mcp-server/dev/INSPECTOR_DEMO.md b/mcp-server/dev/INSPECTOR_DEMO.md
new file mode 100644
index 0000000000..d3a86f9b46
--- /dev/null
+++ b/mcp-server/dev/INSPECTOR_DEMO.md
@@ -0,0 +1,208 @@
+<!--
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+-->
+
+# Hands-on MCP Authorization Demo (with MCP Inspector)
+
+Drive the Gravitino MCP server interactively through the
+[MCP Inspector](https://github.com/modelcontextprotocol/inspector) and watch
+per-user authorization and audit work end to end, as two principals
+(`admin` and `bob`).
+
+This exercises the three governance moments:
+
+1. **Scoped discovery** — `admin` and `bob` run the same call and get
different,
+ authorization-scoped results.
+2. **Write denied** — `bob` (read-only) attempts a write and is denied by
+ Gravitino authorization, surfaced as an explicit error through MCP.
+3. **Audit trail** — every call produces an audit record attributed to the
+ correct principal with an allow/deny outcome.
+
+---
+
+## 1. Prerequisites
+
+- A built Gravitino distribution. If you don't have one:
+ ```bash
+ ./gradlew compileDistribution -x test -PskipWeb=true
+ # produces distribution/package/
+ ```
+- `uv` available in the `mcp-server/` directory (the project already uses it).
+- Node.js (the start script launches the Inspector via `npx
@modelcontextprotocol/inspector`).
+- If you run behind an HTTP proxy, the scripts already export
+ `NO_PROXY=localhost,127.0.0.1`; make sure your shell doesn't force the proxy
+ for loopback in some other way.
+
+---
+
+## 2. Start the demo environment
+
+From the `mcp-server/` directory:
+
+```bash
+./dev/start_inspector_demo.sh
+```
+
+This script (idempotent — safe to re-run):
+
+1. Enables `simple` auth + authorization in
`distribution/package/conf/gravitino.conf`
+ (backing up the original).
+2. Starts Gravitino (or reuses a running one).
+3. Provisions demo data into metalake **`mcp_authz_it`**:
+ - `cat_allowed` and `cat_denied` (two model catalogs)
+ - user **`bob`**
+ - a reader role granting `bob` `USE_CATALOG` on `cat_allowed` **only**
+4. Starts the MCP server in HTTP mode at `http://127.0.0.1:8000/mcp`.
+5. Starts the **MCP Inspector** at `http://localhost:6274/` (with the
session-token
+ requirement disabled, so the plain URL works directly).
+
+On success it prints the connection details and the two principals' tokens. You
+should see the provisioning summary confirming the different slices:
+
+```
+[demo] admin sees catalogs: ['cat_allowed', 'cat_denied']
+[demo] bob sees catalogs: ['cat_allowed']
+```
+
+> If a distribution isn't found, set `GRAVITINO_HOME=/path/to/distribution`.
+
+---
+
+## 3. Open the Inspector
+
+The start script already launched it. Just open:
+
+```
+http://localhost:6274/
+```
+
+(No session token needed — the script starts it with
`DANGEROUSLY_OMIT_AUTH=true`
+for local convenience.)
+
+### Connect
+
+| Field | Value |
+|----------------|------------------------------------------|
+| Transport Type | `Streamable HTTP` |
+| URL | `http://127.0.0.1:8000/mcp` |
+| Header Name | `Authorization` |
+| Header Value | `Basic YWRtaW46ZHVtbXk=` (this is `admin`) |
+
+The two principal tokens (these are Gravitino simple-auth headers, i.e.
+`Basic base64("<user>:dummy")`):
+
+| Principal | Authorization header value |
+|-----------|--------------------------------|
+| `admin` | `Basic YWRtaW46ZHVtbXk=` |
+| `bob` | `Basic Ym9iOmR1bW15` |
+
+Click **Connect**, then **List Tools** — you should see the full read + write
+tool surface (catalogs, schemas, tables, filesets, topics, models, tags, …).
+
+> **Identity is the header.** The Inspector sends your `Authorization` header
on
+> every request; the MCP server forwards it verbatim to Gravitino, which
+> authorizes against that principal. To switch principals, change the header
+> value and reconnect.
+
+---
+
+## 4. The three scenarios
+
+Keep a terminal tailing the audit log while you click:
+
+```bash
+tail -f gravitino-mcp-audit.log
+```
+
+### Scenario 1 — Scoped discovery
+
+1. Connected as **admin**, run tool **`get_list_of_catalogs`**.
+ → Returns **both** `cat_allowed` and `cat_denied`.
+2. Reconnect with the **bob** header (`Basic Ym9iOmR1bW15`), run
+ **`get_list_of_catalogs`** again.
+ → Returns **only** `cat_allowed`.
+
+Same call, different results — sourced entirely from Gravitino's list
filtering,
+not from any logic in the MCP server.
+
+### Scenario 2 — Write denied by authorization
+
+Still connected as **bob**, run **`create_tag`** with arguments:
+
+```json
+{ "name": "test_tag", "comment": "x", "properties": {} }
+```
+
+→ You get an explicit error, e.g.
+`User 'bob' is not authorized to perform operation 'createTag' on metadata
'mcp_authz_it'`.
+
+It's a real authorization denial, not a hidden tool or a silent no-op.
Reconnect
+as **admin** and run the same `create_tag` — it succeeds (admin owns the
metalake).
+
+### Scenario 3 — Audit trail
+
+Look at the audit log you've been tailing. You should see discrete, correctly
+attributed records, for example:
+
+```json
+{"timestamp": "...", "principal": "admin", "tool": "get_list_of_catalogs",
"outcome": "allow"}
+{"timestamp": "...", "principal": "bob", "tool": "get_list_of_catalogs",
"outcome": "allow"}
+{"timestamp": "...", "principal": "bob", "tool": "create_tag", "outcome":
"deny", "error_type": "McpError"}
+```
+
+`admin`'s reads attributed to `admin`, `bob`'s reads to `bob`, and `bob`'s
denied
+write to `bob` with a `deny` outcome.
+
+---
+
+## 5. Tear down
+
+```bash
+./dev/stop_inspector_demo.sh
+```
+
+Stops the Inspector, the MCP server, and Gravitino, and restores the original
+`gravitino.conf`.
+
+---
+
+## 6. Troubleshooting
+
+**The audit log looks empty / stuck at 0 bytes.**
+The MCP server opens the audit file once at startup and holds the handle. If
you
+`rm` the file while the server is running, the process keeps writing to the now
+unlinked inode and a new empty file appears at the path — so you see nothing.
+Don't delete it mid-run; truncate instead (`: > gravitino-mcp-audit.log`), or
+restart the server. The start script truncates rather than deletes for exactly
+this reason.
+
+**`HTTP 000` / `502` when curling localhost.**
+An HTTP proxy is intercepting loopback traffic. Export
+`NO_PROXY=localhost,127.0.0.1` (the scripts already do) or pass `curl
--noproxy '*'`.
+
+**MCP server fails to bind (`can't assign requested address`).**
+`localhost` resolved to a non-loopback address. The scripts bind to the
+`127.0.0.1` literal to avoid this.
+
+**`ModuleNotFoundError: No module named 'mcp_server'`.**
+Launch with `uv run python -m mcp_server ...` (the scripts do); `uv run
mcp_server`
+needs an editable install.
+
+**Re-running the start script.**
+It's idempotent: it reuses a running Gravitino/MCP server and drops+recreates
the
+demo metalake, so you can re-run it freely.
diff --git a/mcp-server/dev/run_authz_integration_test.sh
b/mcp-server/dev/run_authz_integration_test.sh
new file mode 100755
index 0000000000..82673806b7
--- /dev/null
+++ b/mcp-server/dev/run_authz_integration_test.sh
@@ -0,0 +1,179 @@
+#!/bin/bash
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# End-to-end authorization integration test for the Gravitino MCP server.
+#
+# This script:
+# 1. Builds the Gravitino distribution (unless GRAVITINO_HOME is provided).
+# 2. Enables simple authentication + authorization (serviceAdmins=admin).
+# 3. Starts the Gravitino server.
+# 4. Starts the MCP server in HTTP transport mode.
+# 5. Runs the pytest integration suite (which provisions metadata as admin
and
+# verifies per-user authorization through MCP).
+# 6. Tears everything down and restores the original config.
+#
+# Usage:
+# ./dev/run_authz_integration_test.sh
+# GRAVITINO_HOME=/path/to/distribution ./dev/run_authz_integration_test.sh
+
+set -euo pipefail
+
+# All services run on localhost; never route them through an HTTP proxy.
+# httpx (MCP server -> Gravitino, test client) and curl both honour NO_PROXY.
+export NO_PROXY="localhost,127.0.0.1"
+export no_proxy="localhost,127.0.0.1"
+
+MCP_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+REPO_ROOT="$(cd "${MCP_DIR}/.." && pwd)"
+
+# Use the loopback literal (not "localhost") so the MCP server binds to a
+# guaranteed-assignable address; on some hosts "localhost" resolves to a LAN
IP.
+GRAVITINO_PORT="${GRAVITINO_PORT:-8090}"
+GRAVITINO_URI="http://127.0.0.1:${GRAVITINO_PORT}";
+MCP_PORT="${MCP_PORT:-8000}"
+MCP_URL="http://127.0.0.1:${MCP_PORT}/mcp";
+MCP_METALAKE="${MCP_METALAKE:-mcp_authz_it}"
+MCP_AUDIT_LOG="${MCP_DIR}/gravitino-mcp-audit.log"
+
+MCP_PID=""
+GRAVITINO_STARTED="false"
+CONF_BACKUP=""
+
+log() { echo "[authz-it] $*"; }
+
+cleanup() {
+ log "Tearing down..."
+ if [[ -n "${MCP_PID}" ]] && kill -0 "${MCP_PID}" 2>/dev/null; then
+ kill "${MCP_PID}" 2>/dev/null || true
+ wait "${MCP_PID}" 2>/dev/null || true
+ fi
+ if [[ "${GRAVITINO_STARTED}" == "true" ]]; then
+ "${GRAVITINO_HOME}/bin/gravitino.sh" stop || true
+ fi
+ if [[ -n "${CONF_BACKUP}" && -f "${CONF_BACKUP}" ]]; then
+ mv "${CONF_BACKUP}" "${GRAVITINO_HOME}/conf/gravitino.conf"
+ log "Restored original gravitino.conf"
+ fi
+}
+trap cleanup EXIT
+
+# ---------------------------------------------------------------------------
+# 1. Resolve / build the Gravitino distribution
+# ---------------------------------------------------------------------------
+if [[ -z "${GRAVITINO_HOME:-}" ]]; then
+ log "GRAVITINO_HOME not set; building distribution..."
+ (cd "${REPO_ROOT}" && ./gradlew compileDistribution -x test)
+ GRAVITINO_HOME="${REPO_ROOT}/distribution/package"
+fi
+log "Using GRAVITINO_HOME=${GRAVITINO_HOME}"
+
+if [[ ! -x "${GRAVITINO_HOME}/bin/gravitino.sh" ]]; then
+ log "ERROR: ${GRAVITINO_HOME}/bin/gravitino.sh not found"
+ exit 1
+fi
+
+# ---------------------------------------------------------------------------
+# 2. Enable simple auth + authorization
+# ---------------------------------------------------------------------------
+CONF="${GRAVITINO_HOME}/conf/gravitino.conf"
+CONF_BACKUP="${CONF}.authz-it.bak"
+cp "${CONF}" "${CONF_BACKUP}"
+
+# Remove any pre-existing values for the keys we manage, then append ours.
+sed -i.tmp \
+ -e '/^gravitino.authenticators/d' \
+ -e '/^gravitino.authorization.enable/d' \
+ -e '/^gravitino.authorization.serviceAdmins/d' \
+ "${CONF}"
+rm -f "${CONF}.tmp"
+cat >> "${CONF}" <<EOF
+
+# --- injected by run_authz_integration_test.sh ---
+gravitino.authenticators = simple
+gravitino.authorization.enable = true
+gravitino.authorization.serviceAdmins = admin
+EOF
+log "Configured simple auth + authorization (serviceAdmins=admin)"
+
+# ---------------------------------------------------------------------------
+# 3. Start Gravitino
+# ---------------------------------------------------------------------------
+log "Starting Gravitino server..."
+"${GRAVITINO_HOME}/bin/gravitino.sh" start
+GRAVITINO_STARTED="true"
+
+# With simple auth enabled every request needs an Authorization header.
+ADMIN_AUTH="Basic $(printf '%s' 'admin:dummy' | base64)"
+
+log "Waiting for Gravitino to become healthy..."
+for i in $(seq 1 60); do
+ if curl -sf --noproxy '*' -H "Authorization: ${ADMIN_AUTH}" \
+ "${GRAVITINO_URI}/api/version" >/dev/null 2>&1; then
+ log "Gravitino is up."
+ break
+ fi
+ if [[ "${i}" == "60" ]]; then
+ log "ERROR: Gravitino did not become healthy in time"
+ exit 1
+ fi
+ sleep 2
+done
+
+# ---------------------------------------------------------------------------
+# 4. Start the MCP server in HTTP transport mode
+# ---------------------------------------------------------------------------
+log "Starting MCP server (HTTP) on ${MCP_URL}..."
+rm -f "${MCP_AUDIT_LOG}"
+(
+ cd "${MCP_DIR}"
+ uv run python -m mcp_server \
+ --metalake "${MCP_METALAKE}" \
+ --gravitino-uri "${GRAVITINO_URI}" \
+ --transport http \
+ --mcp-url "${MCP_URL}"
+) &
+MCP_PID=$!
+
+log "Waiting for MCP server to become reachable..."
+for i in $(seq 1 30); do
+ if nc -z localhost "${MCP_PORT}" 2>/dev/null; then
+ log "MCP server is up."
+ break
+ fi
+ if [[ "${i}" == "30" ]]; then
+ log "ERROR: MCP server did not start in time"
+ exit 1
+ fi
+ sleep 1
+done
+
+# ---------------------------------------------------------------------------
+# 5. Run the pytest integration suite
+# ---------------------------------------------------------------------------
+log "Running integration tests..."
+(
+ cd "${MCP_DIR}"
+ GRAVITINO_URI="${GRAVITINO_URI}" \
+ MCP_URL="${MCP_URL}" \
+ MCP_METALAKE="${MCP_METALAKE}" \
+ MCP_AUDIT_LOG="${MCP_AUDIT_LOG}" \
+ uv run pytest tests/integration -v
+)
+
+log "Integration tests passed."
diff --git a/mcp-server/dev/start_inspector_demo.sh
b/mcp-server/dev/start_inspector_demo.sh
new file mode 100755
index 0000000000..65ac519f72
--- /dev/null
+++ b/mcp-server/dev/start_inspector_demo.sh
@@ -0,0 +1,260 @@
+#!/bin/bash
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# Bring up a persistent local demo environment for hands-on MCP authorization
+# testing with the MCP Inspector:
+#
+# 1. Gravitino server with simple auth + authorization enabled.
+# 2. Demo data: metalake + two catalogs, user "bob", a role granting bob
+# access to only one catalog (so admin and bob see different slices).
+# 3. MCP server in HTTP transport mode.
+#
+# Unlike run_authz_integration_test.sh, this script LEAVES everything running
so
+# you can drive it from the Inspector. Run stop_inspector_demo.sh to tear down.
+#
+# Usage:
+# ./dev/start_inspector_demo.sh
+# GRAVITINO_HOME=/path/to/distribution ./dev/start_inspector_demo.sh
+
+set -euo pipefail
+
+# Everything is local; never route through an HTTP proxy.
+export NO_PROXY="localhost,127.0.0.1"
+export no_proxy="localhost,127.0.0.1"
+
+MCP_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+REPO_ROOT="$(cd "${MCP_DIR}/.." && pwd)"
+
+GRAVITINO_PORT="${GRAVITINO_PORT:-8090}"
+GRAVITINO_URI="http://127.0.0.1:${GRAVITINO_PORT}";
+MCP_PORT="${MCP_PORT:-8000}"
+MCP_URL="http://127.0.0.1:${MCP_PORT}/mcp";
+MCP_METALAKE="${MCP_METALAKE:-mcp_authz_it}"
+MCP_PID_FILE="${MCP_DIR}/.inspector-demo-mcp.pid"
+ADMIN_AUTH="Basic $(printf '%s' 'admin:dummy' | base64)"
+
+# MCP Inspector ports (defaults baked into @modelcontextprotocol/inspector).
+INSPECTOR_UI_PORT="${INSPECTOR_UI_PORT:-6274}"
+INSPECTOR_PROXY_PORT="${INSPECTOR_PROXY_PORT:-6277}"
+INSPECTOR_PID_FILE="${MCP_DIR}/.inspector-demo-inspector.pid"
+
+log() { echo "[demo] $*"; }
+
+# ---------------------------------------------------------------------------
+# 1. Resolve the Gravitino distribution
+# ---------------------------------------------------------------------------
+if [[ -z "${GRAVITINO_HOME:-}" ]]; then
+ GRAVITINO_HOME="${REPO_ROOT}/distribution/package"
+fi
+if [[ ! -x "${GRAVITINO_HOME}/bin/gravitino.sh" ]]; then
+ log "Distribution not found at ${GRAVITINO_HOME}."
+ log "Build it first: ./gradlew compileDistribution -x test -PskipWeb=true"
+ log "Or set GRAVITINO_HOME to an existing distribution."
+ exit 1
+fi
+log "Using GRAVITINO_HOME=${GRAVITINO_HOME}"
+
+# ---------------------------------------------------------------------------
+# 2. Enable simple auth + authorization (idempotent)
+# ---------------------------------------------------------------------------
+CONF="${GRAVITINO_HOME}/conf/gravitino.conf"
+if ! grep -q "gravitino.authorization.enable = true" "${CONF}"; then
+ cp "${CONF}" "${CONF}.inspector-demo.bak"
+ sed -i.tmp \
+ -e '/^gravitino.authenticators/d' \
+ -e '/^gravitino.authorization.enable/d' \
+ -e '/^gravitino.authorization.serviceAdmins/d' \
+ "${CONF}"
+ rm -f "${CONF}.tmp"
+ cat >> "${CONF}" <<EOF
+
+# --- injected by start_inspector_demo.sh (restored on stop) ---
+gravitino.authenticators = simple
+gravitino.authorization.enable = true
+gravitino.authorization.serviceAdmins = admin
+EOF
+ log "Configured simple auth + authorization (serviceAdmins=admin)"
+else
+ log "Authorization already enabled in config"
+fi
+
+# ---------------------------------------------------------------------------
+# 3. Start Gravitino (if not already up)
+# ---------------------------------------------------------------------------
+if curl -sf --noproxy '*' -H "Authorization: ${ADMIN_AUTH}" \
+ "${GRAVITINO_URI}/api/version" >/dev/null 2>&1; then
+ log "Gravitino already running on ${GRAVITINO_PORT}"
+else
+ log "Starting Gravitino server..."
+ "${GRAVITINO_HOME}/bin/gravitino.sh" start
+ log "Waiting for Gravitino to become healthy..."
+ for i in $(seq 1 60); do
+ if curl -sf --noproxy '*' -H "Authorization: ${ADMIN_AUTH}" \
+ "${GRAVITINO_URI}/api/version" >/dev/null 2>&1; then
+ log "Gravitino is up."
+ break
+ fi
+ if [[ "${i}" == "60" ]]; then
+ log "ERROR: Gravitino did not become healthy in time"
+ exit 1
+ fi
+ sleep 2
+ done
+fi
+
+# ---------------------------------------------------------------------------
+# 4. Provision demo data (idempotent: drop + recreate the metalake)
+# ---------------------------------------------------------------------------
+log "Provisioning demo data into metalake '${MCP_METALAKE}'..."
+(
+ cd "${MCP_DIR}"
+ GRAVITINO_URI="${GRAVITINO_URI}" MCP_METALAKE="${MCP_METALAKE}" \
+ uv run python <<'PY'
+import base64
+import os
+
+import httpx
+
+from tests.integration.gravitino_setup import GravitinoFixture
+
+uri = os.environ["GRAVITINO_URI"]
+ml = os.environ["MCP_METALAKE"]
+
+
+def hdr(user):
+ return "Basic " + base64.b64encode(f"{user}:dummy".encode()).decode()
+
+
+# Drop an existing metalake so the script can be re-run cleanly.
+client = httpx.Client(
+ base_url=uri, headers={"Authorization": hdr("admin")}, timeout=30
+)
+if client.get(f"/api/metalakes/{ml}").status_code == 200:
+ client.put(
+ f"/api/metalakes/{ml}",
+ json={"updates": [{"@type": "setProperty", "property": "in-use",
"value": "false"}]},
+ )
+ client.delete(f"/api/metalakes/{ml}?force=true")
+ print(f"[demo] removed existing metalake '{ml}'")
+client.close()
+
+GravitinoFixture(uri, ml).provision()
+print(f"[demo] provisioned metalake '{ml}': cat_allowed, cat_denied, user bob,
reader role")
+
+# Show the resulting authorization slice.
+for user in ("admin", "bob"):
+ r = httpx.get(
+ f"{uri}/api/metalakes/{ml}/catalogs?details=true",
+ headers={"Authorization": hdr(user)},
+ )
+ names = [c["name"] for c in r.json().get("catalogs", [])]
+ print(f"[demo] {user} sees catalogs: {names}")
+PY
+)
+
+# ---------------------------------------------------------------------------
+# 5. Start the MCP server in HTTP mode (if not already up)
+# ---------------------------------------------------------------------------
+if nc -z 127.0.0.1 "${MCP_PORT}" 2>/dev/null; then
+ log "An MCP server is already listening on ${MCP_PORT}; leaving it as-is."
+else
+ log "Starting MCP server (HTTP) on ${MCP_URL}..."
+ (
+ cd "${MCP_DIR}"
+ # Truncate (not delete) the audit log so the running process keeps its
handle.
+ : > gravitino-mcp-audit.log
+ nohup uv run python -m mcp_server \
+ --metalake "${MCP_METALAKE}" \
+ --gravitino-uri "${GRAVITINO_URI}" \
+ --transport http \
+ --mcp-url "${MCP_URL}" > "${MCP_DIR}/.inspector-demo-mcp.out" 2>&1 &
+ echo $! > "${MCP_PID_FILE}"
+ )
+ for i in $(seq 1 30); do
+ nc -z 127.0.0.1 "${MCP_PORT}" 2>/dev/null && break
+ sleep 1
+ done
+ if nc -z 127.0.0.1 "${MCP_PORT}" 2>/dev/null; then
+ log "MCP server is up (pid $(cat "${MCP_PID_FILE}"))."
+ else
+ log "ERROR: MCP server did not start; see
${MCP_DIR}/.inspector-demo-mcp.out"
+ exit 1
+ fi
+fi
+
+# ---------------------------------------------------------------------------
+# 6. Start the MCP Inspector (UI on :6274, proxy on :6277)
+# ---------------------------------------------------------------------------
+# DANGEROUSLY_OMIT_AUTH disables the session-token requirement so the plain
+# http://localhost:6274/ URL works without a token (fine for a local demo).
+# MCP_AUTO_OPEN_ENABLED=false keeps it from popping a browser when
backgrounded.
+# The Inspector binds to "localhost" (may be IPv6 ::1), so probe via localhost.
+if nc -z localhost "${INSPECTOR_UI_PORT}" 2>/dev/null; then
+ log "An Inspector is already listening on ${INSPECTOR_UI_PORT}; leaving it
as-is."
+else
+ log "Starting MCP Inspector on http://localhost:${INSPECTOR_UI_PORT} ..."
+ (
+ cd "${MCP_DIR}"
+ DANGEROUSLY_OMIT_AUTH=true \
+ MCP_AUTO_OPEN_ENABLED=false \
+ nohup npx @modelcontextprotocol/inspector \
+ > "${MCP_DIR}/.inspector-demo-inspector.out" 2>&1 &
+ echo $! > "${INSPECTOR_PID_FILE}"
+ )
+ for i in $(seq 1 60); do
+ nc -z localhost "${INSPECTOR_UI_PORT}" 2>/dev/null && break
+ sleep 1
+ done
+ if nc -z localhost "${INSPECTOR_UI_PORT}" 2>/dev/null; then
+ log "Inspector is up (pid $(cat "${INSPECTOR_PID_FILE}"))."
+ else
+ log "WARN: Inspector did not come up; see
${MCP_DIR}/.inspector-demo-inspector.out"
+ log " (You can still start it manually: npx
@modelcontextprotocol/inspector)"
+ fi
+fi
+
+# ---------------------------------------------------------------------------
+# 7. Print connection details
+# ---------------------------------------------------------------------------
+cat <<EOF
+
+============================================================
+ Demo environment is ready.
+============================================================
+ Gravitino : ${GRAVITINO_URI} (simple auth + authorization)
+ MCP server: ${MCP_URL} (Streamable HTTP)
+ Metalake : ${MCP_METALAKE}
+ Audit log : ${MCP_DIR}/gravitino-mcp-audit.log
+
+ >>> Open the Inspector: http://localhost:${INSPECTOR_UI_PORT}/
+
+ In the Inspector, connect with:
+ Transport Type : Streamable HTTP
+ URL : ${MCP_URL}
+ Header Name : Authorization
+ Header Value : ${ADMIN_AUTH} <- admin
+ Basic $(printf '%s' 'bob:dummy' | base64) <- bob
+
+ Watch audit records live:
+ tail -f ${MCP_DIR}/gravitino-mcp-audit.log
+
+ Tear everything down (includes the Inspector):
+ ./dev/stop_inspector_demo.sh
+============================================================
+EOF
diff --git a/mcp-server/dev/stop_inspector_demo.sh
b/mcp-server/dev/stop_inspector_demo.sh
new file mode 100755
index 0000000000..0335dac0f2
--- /dev/null
+++ b/mcp-server/dev/stop_inspector_demo.sh
@@ -0,0 +1,77 @@
+#!/bin/bash
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# Tear down the demo environment started by start_inspector_demo.sh:
+# stops the MCP server, stops Gravitino, and restores the original config.
+
+set -uo pipefail
+
+MCP_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+REPO_ROOT="$(cd "${MCP_DIR}/.." && pwd)"
+
+MCP_PORT="${MCP_PORT:-8000}"
+MCP_PID_FILE="${MCP_DIR}/.inspector-demo-mcp.pid"
+INSPECTOR_UI_PORT="${INSPECTOR_UI_PORT:-6274}"
+INSPECTOR_PROXY_PORT="${INSPECTOR_PROXY_PORT:-6277}"
+INSPECTOR_PID_FILE="${MCP_DIR}/.inspector-demo-inspector.pid"
+
+log() { echo "[demo] $*"; }
+
+if [[ -z "${GRAVITINO_HOME:-}" ]]; then
+ GRAVITINO_HOME="${REPO_ROOT}/distribution/package"
+fi
+
+# 1. Stop the MCP Inspector (npx spawns child processes; free its ports too).
+if [[ -f "${INSPECTOR_PID_FILE}" ]]; then
+ INSPECTOR_PID="$(cat "${INSPECTOR_PID_FILE}")"
+ kill "${INSPECTOR_PID}" 2>/dev/null && \
+ log "Stopped Inspector (pid ${INSPECTOR_PID})" || true
+ rm -f "${INSPECTOR_PID_FILE}"
+fi
+for port in "${INSPECTOR_UI_PORT}" "${INSPECTOR_PROXY_PORT}"; do
+ lsof -ti :"${port}" 2>/dev/null | xargs kill -9 2>/dev/null && \
+ log "Freed Inspector port ${port}" || true
+done
+
+# 2. Stop the MCP server.
+if [[ -f "${MCP_PID_FILE}" ]]; then
+ MCP_PID="$(cat "${MCP_PID_FILE}")"
+ if kill "${MCP_PID}" 2>/dev/null; then
+ log "Stopped MCP server (pid ${MCP_PID})"
+ fi
+ rm -f "${MCP_PID_FILE}"
+fi
+# Belt and suspenders: free the port if anything is still bound.
+lsof -ti :"${MCP_PORT}" 2>/dev/null | xargs kill -9 2>/dev/null && \
+ log "Freed port ${MCP_PORT}" || true
+
+# 3. Stop Gravitino.
+if [[ -x "${GRAVITINO_HOME}/bin/gravitino.sh" ]]; then
+ "${GRAVITINO_HOME}/bin/gravitino.sh" stop || true
+ log "Stopped Gravitino"
+fi
+
+# 4. Restore the original config.
+CONF="${GRAVITINO_HOME}/conf/gravitino.conf"
+if [[ -f "${CONF}.inspector-demo.bak" ]]; then
+ mv "${CONF}.inspector-demo.bak" "${CONF}"
+ log "Restored original gravitino.conf"
+fi
+
+log "Teardown complete."
diff --git a/mcp-server/mcp_server/client/factory.py
b/mcp-server/mcp_server/client/factory.py
index d990284d28..201625540d 100644
--- a/mcp-server/mcp_server/client/factory.py
+++ b/mcp-server/mcp_server/client/factory.py
@@ -29,7 +29,7 @@ class RESTClientFactory:
@classmethod
def create_rest_client(
- cls, metalake_name: str, uri: str
+ cls, metalake_name: str, uri: str, authorization: str = ""
) -> "PlainRESTClientOperation":
"""
Create a new rest client instance with the specified parameters.
@@ -37,11 +37,14 @@ class RESTClientFactory:
Args:
metalake_name: Name of the metalake
uri: URI of the Gravitino server endpoint
+ authorization: Full Authorization header value forwarded verbatim
+ (e.g. "Bearer <token>" or "Basic <base64(user:dummy)>").
+ Empty string means anonymous.
Returns:
New instance of the configured rest client class
"""
- return cls._rest_client_class(metalake_name, uri)
+ return cls._rest_client_class(metalake_name, uri, authorization)
@classmethod
def set_rest_client(cls, rest_client_class: type) -> None:
diff --git a/mcp-server/mcp_server/client/plain/plain_rest_client_operation.py
b/mcp-server/mcp_server/client/plain/plain_rest_client_operation.py
index 205562b6d8..a6adc0e099 100644
--- a/mcp-server/mcp_server/client/plain/plain_rest_client_operation.py
+++ b/mcp-server/mcp_server/client/plain/plain_rest_client_operation.py
@@ -62,8 +62,23 @@ from mcp_server.client.topic_operation import TopicOperation
# pylint: disable=too-many-instance-attributes
class PlainRESTClientOperation(GravitinoOperation):
- def __init__(self, metalake_name: str, uri: str):
- _rest_client = httpx.AsyncClient(base_url=uri)
+ def __init__(self, metalake_name: str, uri: str, authorization: str = ""):
+ """Create a REST client for one identity.
+
+ Args:
+ metalake_name: Name of the metalake.
+ uri: Gravitino server URI.
+ authorization: Full ``Authorization`` header value forwarded
verbatim
+ on every request (for example ``"Bearer <token>"`` for OAuth2
or
+ ``"Basic <base64(user:secret)>"`` for simple or Basic auth).
+ Empty string means anonymous (no header sent).
+ """
+ headers = {}
+ if authorization:
+ headers["Authorization"] = authorization
+ _rest_client = httpx.AsyncClient(base_url=uri, headers=headers)
+ # Kept so the shared connection pool can be closed (see close()).
+ self._rest_client = _rest_client
self._catalog_operation = PlainRESTClientCatalogOperation(
metalake_name, _rest_client
)
@@ -95,6 +110,10 @@ class PlainRESTClientOperation(GravitinoOperation):
metalake_name, _rest_client
)
+ async def close(self) -> None:
+ """Close the shared httpx client and release its connection pool."""
+ await self._rest_client.aclose()
+
def as_catalog_operation(self) -> CatalogOperation:
return self._catalog_operation
diff --git a/mcp-server/mcp_server/core/audit.py
b/mcp-server/mcp_server/core/audit.py
new file mode 100644
index 0000000000..2de8a137dc
--- /dev/null
+++ b/mcp-server/mcp_server/core/audit.py
@@ -0,0 +1,81 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import base64
+import binascii
+import json
+import logging
+from datetime import datetime, timezone
+
+_audit_logger = logging.getLogger("gravitino.mcp.audit")
+
+
+def _extract_principal(authorization: str) -> str:
+ """Derive a display principal from a raw Authorization header value.
+
+ - "Basic <base64(user:secret)>" → "<user>" (Gravitino simple auth)
+ - "Bearer <token>" → "bearer:<first-8-chars-of-token>"
+ - empty / missing / unparsable → "anonymous"
+ """
+ if not authorization:
+ return "anonymous"
+ parts = authorization.split()
+ if len(parts) != 2:
+ return "anonymous"
+ scheme, credential = parts[0].lower(), parts[1]
+ if scheme == "basic":
+ try:
+ decoded = base64.b64decode(credential, validate=True).decode(
+ "utf-8"
+ )
+ except (binascii.Error, UnicodeDecodeError, ValueError):
+ return "anonymous"
+ user = decoded.split(":", 1)[0]
+ return user if user else "anonymous"
+ if scheme == "bearer":
+ return f"bearer:{credential[:8]}"
+ return "anonymous"
+
+
+def emit(
+ *,
+ principal: str,
+ tool: str,
+ outcome: str,
+ error_type: str = "",
+) -> None:
+ """Write one structured JSON audit record to the audit logger.
+
+ Args:
+ principal: Identity derived from the request (e.g. "bearer:abc12345"
or "anonymous").
+ tool: MCP tool name that was invoked.
+ outcome: "allow" for successful calls, "deny" for failed calls. Note
the
+ AuditMiddleware emits "deny" for any tool-call exception (an
+ authorization denial being the common case), not only
+ authorization failures; inspect error_type to disambiguate.
+ error_type: Exception class name when outcome is "deny", empty
otherwise.
+ """
+ record = {
+ "timestamp": datetime.now(timezone.utc).isoformat(),
+ "principal": principal,
+ "tool": tool,
+ "outcome": outcome,
+ }
+ if error_type:
+ record["error_type"] = error_type
+
+ _audit_logger.info(json.dumps(record))
diff --git a/mcp-server/mcp_server/core/context.py
b/mcp-server/mcp_server/core/context.py
index 6e540a0b51..fd6731b495 100644
--- a/mcp-server/mcp_server/core/context.py
+++ b/mcp-server/mcp_server/core/context.py
@@ -15,15 +15,127 @@
# specific language governing permissions and limitations
# under the License.
+import asyncio
+import logging
+from collections import OrderedDict
+
from mcp_server.client.factory import RESTClientFactory
from mcp_server.core.setting import Setting
+_LOG = logging.getLogger(__name__)
+
+# Upper bound on the number of per-principal REST clients kept alive at once.
+# Each client owns an httpx connection pool; caching by Authorization header
lets
+# repeated calls from the same principal reuse a pool instead of opening a new
one
+# per tool call, while the LRU bound keeps memory/sockets in check as
principals
+# (e.g. rotating tokens) come and go.
+_MAX_CACHED_CLIENTS = 128
+
+
+def _get_request_authorization() -> str:
+ """Return the raw ``Authorization`` header of the current HTTP request.
+
+ The header is forwarded to Gravitino verbatim so the auth scheme chosen by
+ the agent (``Basic`` for simple or Basic auth, ``Bearer`` for OAuth2,
+ ``Negotiate`` for Kerberos) is preserved. Returns an empty string in stdio
+ mode or when the header is absent.
+ """
+ try:
+ # Imported lazily: only available within an HTTP request context.
+ # pylint: disable=import-outside-toplevel
+ from fastmcp.server.dependencies import get_http_request
+
+ return get_http_request().headers.get("authorization", "")
+ except (LookupError, RuntimeError):
+ # No active HTTP request: stdio mode (get_http_request raises
+ # RuntimeError) or missing request context (LookupError).
+ return ""
+
+
+def startup_authorization(setting: Setting) -> str:
+ """The static --token rendered as an ``Authorization`` header value.
+
+ The CLI token is treated as an OAuth2 Bearer token. Empty string when no
+ token is configured (anonymous). This is the identity used in stdio mode
and
+ the fallback for HTTP requests that carry no ``Authorization`` header.
+ """
+ return f"Bearer {setting.token}" if setting.token else ""
+
class GravitinoContext:
def __init__(self, setting: Setting):
- self.gravitino_client = RESTClientFactory.create_rest_client(
- setting.metalake, setting.gravitino_uri
+ self._setting = setting
+ self._default_client = RESTClientFactory.create_rest_client(
+ setting.metalake,
+ setting.gravitino_uri,
+ startup_authorization(setting),
)
+ # 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()
+ # 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 self.gravitino_client
+ """Return a REST client carrying the correct identity for this request.
+
+ 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.
+
+ Falls back to the shared default client (static startup token) 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.
+ """
+ authorization = _get_request_authorization()
+ if not authorization:
+ return self._default_client
+
+ cached = self._clients_by_auth.get(authorization)
+ if cached is not None:
+ self._clients_by_auth.move_to_end(authorization)
+ return cached
+
+ client = RESTClientFactory.create_rest_client(
+ self._setting.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)
+ return client
+
+ def _schedule_close(self, client) -> None:
+ """Best-effort close of an evicted client's connection pool.
+
+ Closing is async; schedule it on the running event loop if there is one
+ (the normal HTTP-serving case). With no running loop (stdio mode/tests)
+ there is nothing to schedule and the client is left for GC.
+ """
+ close = getattr(client, "close", None)
+ if close is None:
+ return
+ try:
+ loop = asyncio.get_running_loop()
+ except RuntimeError:
+ return
+ task = loop.create_task(close())
+ # Hold a strong reference until the task finishes (see
_pending_closes).
+ self._pending_closes.add(task)
+ task.add_done_callback(self._on_close_done)
+
+ def _on_close_done(self, task: "asyncio.Task") -> None:
+ """Drop the finished close task and log any failure."""
+ self._pending_closes.discard(task)
+ exc = task.exception()
+ if exc is not None:
+ _LOG.warning("Failed to close evicted REST client: %s", exc)
diff --git a/mcp-server/mcp_server/core/setting.py
b/mcp-server/mcp_server/core/setting.py
index 474fdbc6b4..0659c8f602 100644
--- a/mcp-server/mcp_server/core/setting.py
+++ b/mcp-server/mcp_server/core/setting.py
@@ -27,9 +27,28 @@ class DefaultSetting:
@dataclass
-class Setting:
+class Setting: # pylint: disable=too-many-instance-attributes
metalake: str
gravitino_uri: str = DefaultSetting.default_gravitino_uri
tags: Set[str] = field(default_factory=set)
transport: str = DefaultSetting.default_transport
mcp_url: str = DefaultSetting.default_mcp_url
+ # Static OAuth2 Bearer token. Sent on every request in stdio mode; in HTTP
+ # mode it is only the fallback used when an incoming request carries no
+ # Authorization header (per-request identity takes priority).
+ # Empty string means anonymous (no Authorization header sent).
+ # repr=False keeps the raw value out of the dataclass-generated __repr__.
+ token: str = field(default="", repr=False)
+ # TLS certificate/key paths for serving the HTTP endpoint over HTTPS.
+ # Both must be set to enable TLS; empty means plain HTTP.
+ tls_cert: str = ""
+ tls_key: str = ""
+
+ def __str__(self) -> str:
+ token_display = "***" if self.token else ""
+ return (
+ f"Setting(metalake={self.metalake},
gravitino_uri={self.gravitino_uri}, "
+ f"tags={self.tags}, transport={self.transport},
mcp_url={self.mcp_url}, "
+ f"token={token_display}, tls_cert={self.tls_cert}, "
+ f"tls_key={self.tls_key})"
+ )
diff --git a/mcp-server/mcp_server/main.py b/mcp-server/mcp_server/main.py
index 73a5aa265f..903d9d30b0 100644
--- a/mcp-server/mcp_server/main.py
+++ b/mcp-server/mcp_server/main.py
@@ -17,6 +17,7 @@
import argparse
import logging
+import os
from mcp_server.core.setting import DefaultSetting, Setting
from mcp_server.server import GravitinoMCPServer
@@ -30,6 +31,9 @@ def do_main():
tags=args.include_tool_tags,
transport=args.transport,
mcp_url=args.mcp_url,
+ token=args.token,
+ tls_cert=args.tls_cert,
+ tls_key=args.tls_key,
)
_init_logging(setting)
logging.info("Gravitino MCP server setting: %s", setting)
@@ -44,6 +48,12 @@ def _init_logging(setting: Setting):
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
+ # Separate file handler for structured audit records (one JSON line per
entry).
+ audit_handler = logging.FileHandler("gravitino-mcp-audit.log")
+ audit_handler.setLevel(logging.INFO)
+ audit_handler.setFormatter(logging.Formatter("%(message)s"))
+ logging.getLogger("gravitino.mcp.audit").addHandler(audit_handler)
+ logging.getLogger("gravitino.mcp.audit").propagate = False
def _comma_separated_set(value) -> set:
@@ -83,9 +93,10 @@ def _parse_args():
parser.add_argument(
"--transport",
type=str,
- choices=["stdio", "http"],
+ choices=["stdio", "http", "streamable-http"],
default=DefaultSetting.default_transport,
- help=f"Transport protocol type: stdio (local), http (Streamable HTTP).
"
+ help="Transport protocol type: stdio (local), http / streamable-http "
+ "(networked Streamable HTTP; the two names are equivalent). "
f"(default: {DefaultSetting.default_transport})",
)
@@ -93,7 +104,36 @@ def _parse_args():
"--mcp-url",
type=str,
default=DefaultSetting.default_mcp_url,
- help=f"The url of MCP server if using http transport. (default:
{DefaultSetting.default_mcp_url})",
+ help="The url of MCP server if using http transport, http:// or
https://. "
+ f"(default: {DefaultSetting.default_mcp_url})",
+ )
+
+ parser.add_argument(
+ "--token",
+ type=str,
+ default=os.environ.get("GRAVITINO_TOKEN", ""),
+ help="Static OAuth2 Bearer token used to authenticate to Gravitino. "
+ "In stdio mode it is sent on every request; in HTTP mode it is only
the "
+ "fallback when an incoming request carries no Authorization header "
+ "(per-request identity takes priority). "
+ "Can also be set via the GRAVITINO_TOKEN environment variable. "
+ "When omitted, requests are sent without authentication.",
+ )
+
+ parser.add_argument(
+ "--tls-cert",
+ type=str,
+ default="",
+ help="Path to the TLS certificate (PEM) for serving the HTTP endpoint "
+ "over HTTPS. Requires --tls-key. When omitted, the endpoint serves
plain HTTP.",
+ )
+
+ parser.add_argument(
+ "--tls-key",
+ type=str,
+ default="",
+ help="Path to the TLS private key (PEM) for serving the HTTP endpoint "
+ "over HTTPS. Requires --tls-cert.",
)
args = parser.parse_args()
diff --git a/mcp-server/mcp_server/server.py b/mcp-server/mcp_server/server.py
index 9cce5f556b..b8909120da 100644
--- a/mcp-server/mcp_server/server.py
+++ b/mcp-server/mcp_server/server.py
@@ -21,18 +21,71 @@ from contextlib import asynccontextmanager
from typing import AsyncIterator
from urllib.parse import urlparse
+import mcp.types as mt
from fastmcp import FastMCP
from fastmcp.server.middleware.error_handling import ErrorHandlingMiddleware
from fastmcp.server.middleware.logging import (
LoggingMiddleware,
)
+from fastmcp.server.middleware.middleware import (
+ CallNext,
+ Middleware,
+ MiddlewareContext,
+)
from fastmcp.server.middleware.timing import TimingMiddleware
+from fastmcp.tools.base import ToolResult
-from mcp_server.core.context import GravitinoContext
+from mcp_server.core import audit
+from mcp_server.core.context import (
+ GravitinoContext,
+ _get_request_authorization,
+ startup_authorization,
+)
from mcp_server.core.setting import Setting
from mcp_server.tools import load_tools
+def _get_principal_from_request(fallback_authorization: str = "") -> str:
+ """Derive a display principal for audit attribution.
+
+ Uses the incoming HTTP request's Authorization header when present;
+ otherwise falls back to the static startup identity (``--token``), which is
+ what actually authenticates the call in stdio mode or in HTTP requests that
+ carry no Authorization header. Returns "anonymous" when neither is set.
+ """
+ authorization = _get_request_authorization() or fallback_authorization
+ # pylint: disable=protected-access
+ return audit._extract_principal(authorization)
+
+
+class AuditMiddleware(Middleware):
+ """Emit a structured audit record for every tool invocation."""
+
+ def __init__(self, fallback_authorization: str = ""):
+ super().__init__()
+ self._fallback_authorization = fallback_authorization
+
+ async def on_call_tool(
+ self,
+ context: MiddlewareContext[mt.CallToolRequestParams],
+ call_next: CallNext[mt.CallToolRequestParams, ToolResult],
+ ) -> ToolResult:
+ tool_name = context.message.name if context.message else "unknown"
+ principal = _get_principal_from_request(self._fallback_authorization)
+ try:
+ result = await call_next(context)
+ audit.emit(principal=principal, tool=tool_name, outcome="allow")
+ return result
+ except Exception as exc:
+ audit.emit(
+ principal=principal,
+ tool=tool_name,
+ outcome="deny",
+ error_type=type(exc).__name__,
+ )
+ raise
+
+
def _create_lifespan_manager(gravitino_context: GravitinoContext):
@asynccontextmanager
@@ -56,10 +109,10 @@ def _create_gravitino_mcp(setting: Setting) -> FastMCP:
lifespan=_create_lifespan_manager(GravitinoContext(setting)),
)
+ mcp.add_middleware(AuditMiddleware(startup_authorization(setting)))
mcp.add_middleware(
LoggingMiddleware(include_payloads=True, max_payload_length=1000)
)
-
mcp.add_middleware(TimingMiddleware())
mcp.add_middleware(
ErrorHandlingMiddleware(
@@ -70,17 +123,20 @@ def _create_gravitino_mcp(setting: Setting) -> FastMCP:
return mcp
-def _parse_mcp_url(url: str) -> ():
+def _parse_mcp_url(url: str) -> tuple[str, int, str]:
try:
parsed = urlparse(url)
- if parsed.scheme.lower() != "http":
- raise ValueError(f"Not support: {parsed.scheme}，only support HTTP")
+ scheme = parsed.scheme.lower()
+ if scheme not in ("http", "https"):
+ raise ValueError(
+ f"Not supported: {parsed.scheme}, only http/https are
supported"
+ )
host = parsed.hostname or "0.0.0.0"
port = parsed.port
if port is None:
- port = 80
+ port = 443 if scheme == "https" else 80
path = parsed.path
if not path.startswith("/"):
@@ -109,8 +165,55 @@ class GravitinoMCPServer:
def _run_http(self):
_host, _port, _path = _parse_mcp_url(self.setting.mcp_url)
- asyncio.run(
- self.mcp.run_async(
- transport="http", host=_host, port=_port, path=_path
- )
+ self._validate_tls_config()
+ # FastMCP accepts "http" and "streamable-http" as equivalent aliases.
+ transport = (
+ "streamable-http"
+ if self.setting.transport == "streamable-http"
+ else "http"
)
+
+ run_kwargs = {
+ "transport": transport,
+ "host": _host,
+ "port": _port,
+ "path": _path,
+ }
+
+ # Serve over TLS when both certificate and key are provided. FastMCP
+ # forwards uvicorn_config to the underlying uvicorn server.
+ if self.setting.tls_cert and self.setting.tls_key:
+ run_kwargs["uvicorn_config"] = {
+ "ssl_certfile": self.setting.tls_cert,
+ "ssl_keyfile": self.setting.tls_key,
+ }
+ logging.info(
+ "Serving MCP endpoint over TLS (cert=%s)",
self.setting.tls_cert
+ )
+
+ asyncio.run(self.mcp.run_async(**run_kwargs))
+
+ def _validate_tls_config(self):
+ """Reject inconsistent TLS configuration before starting the server.
+
+ Guards against serving plain HTTP on an ``https://`` URL (or TLS behind
+ an ``http://`` URL), and against providing only one of cert/key.
+ """
+ cert, key = self.setting.tls_cert, self.setting.tls_key
+ if bool(cert) != bool(key):
+ raise ValueError(
+ "Both --tls-cert and --tls-key must be provided together "
+ "(or neither)."
+ )
+ scheme = urlparse(self.setting.mcp_url).scheme.lower()
+ tls_enabled = bool(cert and key)
+ if scheme == "https" and not tls_enabled:
+ raise ValueError(
+ f"mcp_url '{self.setting.mcp_url}' uses https but TLS is not "
+ "configured; provide --tls-cert and --tls-key."
+ )
+ if scheme == "http" and tls_enabled:
+ raise ValueError(
+ "--tls-cert/--tls-key are set but mcp_url "
+ f"'{self.setting.mcp_url}' uses http; use an https URL."
+ )
diff --git a/mcp-server/pyproject.toml b/mcp-server/pyproject.toml
index 5966f88b54..958bada491 100644
--- a/mcp-server/pyproject.toml
+++ b/mcp-server/pyproject.toml
@@ -34,6 +34,13 @@ dependencies = [
"pylint>=2.20.0",
]
+# Restrict packaging to the importable library. Without this, setuptools'
+# flat-layout auto-discovery also picks up sibling dirs like `dev/` (it only
+# excludes `tests*` by default), failing the build with
+# "Multiple top-level packages discovered in a flat-layout".
+[tool.setuptools.packages.find]
+include = ["mcp_server", "mcp_server.*"]
+
[tool.isort]
profile = "black"
line_length = 80
diff --git a/mcp-server/mcp_server/core/context.py
b/mcp-server/tests/integration/__init__.py
similarity index 68%
copy from mcp-server/mcp_server/core/context.py
copy to mcp-server/tests/integration/__init__.py
index 6e540a0b51..13a83393a9 100644
--- a/mcp-server/mcp_server/core/context.py
+++ b/mcp-server/tests/integration/__init__.py
@@ -14,16 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-
-from mcp_server.client.factory import RESTClientFactory
-from mcp_server.core.setting import Setting
-
-
-class GravitinoContext:
- def __init__(self, setting: Setting):
- self.gravitino_client = RESTClientFactory.create_rest_client(
- setting.metalake, setting.gravitino_uri
- )
-
- def rest_client(self):
- return self.gravitino_client
diff --git a/mcp-server/tests/integration/conftest.py
b/mcp-server/tests/integration/conftest.py
new file mode 100644
index 0000000000..6dda99faea
--- /dev/null
+++ b/mcp-server/tests/integration/conftest.py
@@ -0,0 +1,73 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""pytest fixtures for the MCP authorization integration test.
+
+The orchestration script (``dev/run_authz_integration_test.sh``) starts a real
+Gravitino server and the MCP server in HTTP mode, then exports the connection
+details below. When these env vars are absent the whole integration suite is
+skipped so a plain ``pytest`` run stays green without external services.
+
+Required environment variables:
+ GRAVITINO_URI e.g. http://localhost:8090
+ MCP_URL e.g. http://localhost:8000/mcp
+ MCP_METALAKE metalake name the MCP server was launched against
+"""
+
+import os
+
+import pytest
+
+from tests.integration.gravitino_setup import GravitinoFixture
+
+# Fixtures receive other fixtures as same-named parameters; this is the
standard
+# pytest pattern, not an accidental shadowing.
+# pylint: disable=redefined-outer-name
+
+_REQUIRED_ENV = ("GRAVITINO_URI", "MCP_URL", "MCP_METALAKE")
+
+
+def _missing_env() -> list:
+ return [name for name in _REQUIRED_ENV if not os.environ.get(name)]
+
+
[email protected](scope="session")
+def integration_env() -> dict:
+ missing = _missing_env()
+ if missing:
+ pytest.skip(
+ "Integration test requires a running Gravitino + MCP server. "
+ f"Missing env: {', '.join(missing)}. "
+ "Run via dev/run_authz_integration_test.sh."
+ )
+ return {
+ "gravitino_uri": os.environ["GRAVITINO_URI"],
+ "mcp_url": os.environ["MCP_URL"],
+ "metalake": os.environ["MCP_METALAKE"],
+ }
+
+
[email protected](scope="session")
+def gravitino_fixture(integration_env: dict) -> GravitinoFixture:
+ """Provision metalake/catalogs/user/role/grant once for the whole suite."""
+ fixture = GravitinoFixture(
+ gravitino_uri=integration_env["gravitino_uri"],
+ metalake=integration_env["metalake"],
+ )
+ fixture.provision()
+ yield fixture
+ fixture.close()
diff --git a/mcp-server/tests/integration/gravitino_setup.py
b/mcp-server/tests/integration/gravitino_setup.py
new file mode 100644
index 0000000000..38510ac2c1
--- /dev/null
+++ b/mcp-server/tests/integration/gravitino_setup.py
@@ -0,0 +1,145 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Provision Gravitino metadata and authorization fixtures for the integration
test.
+
+All requests are issued as the service admin using Gravitino simple
+authentication (``Authorization: Basic base64(user:dummy)``). The fixture
+creates a metalake with two model catalogs, a non-admin user ``bob``, and a
role
+that grants ``bob`` access to only one of the two catalogs. This produces a
+visibly different authorization slice between the admin and ``bob`` principals.
+"""
+
+import base64
+
+import httpx
+
+
+def basic_auth_header(user: str) -> str:
+ """Build a Gravitino simple-auth header value for ``user``."""
+ credential = base64.b64encode(f"{user}:dummy".encode("utf-8")).decode(
+ "ascii"
+ )
+ return f"Basic {credential}"
+
+
+class GravitinoFixture: # pylint: disable=too-many-instance-attributes
+ """Sets up metalake/catalogs/user/role/grant via the Gravitino REST API."""
+
+ def __init__( # pylint:
disable=too-many-positional-arguments,too-many-arguments
+ self,
+ gravitino_uri: str,
+ metalake: str,
+ admin_user: str = "admin",
+ granted_user: str = "bob",
+ catalog_allowed: str = "cat_allowed",
+ catalog_denied: str = "cat_denied",
+ role_name: str = "reader_role",
+ ):
+ self.gravitino_uri = gravitino_uri.rstrip("/")
+ self.metalake = metalake
+ self.admin_user = admin_user
+ self.granted_user = granted_user
+ self.catalog_allowed = catalog_allowed
+ self.catalog_denied = catalog_denied
+ self.role_name = role_name
+ self._client = httpx.Client(
+ base_url=self.gravitino_uri,
+ headers={"Authorization": basic_auth_header(admin_user)},
+ timeout=30.0,
+ )
+
+ def _post(self, path: str, body: dict) -> httpx.Response:
+ response = self._client.post(path, json=body)
+ response.raise_for_status()
+ return response
+
+ def _put(self, path: str, body: dict) -> httpx.Response:
+ response = self._client.put(path, json=body)
+ response.raise_for_status()
+ return response
+
+ def provision(self) -> None:
+ """Create all metadata and authorization fixtures.
+
+ Not idempotent: every step raises on a non-2xx response, so re-running
+ against an already-provisioned metalake fails (e.g. HTTP 409). Expects
a
+ clean Gravitino instance.
+ """
+ self._create_metalake()
+ self._create_model_catalog(self.catalog_allowed)
+ self._create_model_catalog(self.catalog_denied)
+ self._add_user(self.granted_user)
+ self._create_reader_role()
+ self._grant_role_to_user()
+
+ def _create_metalake(self) -> None:
+ self._post(
+ "/api/metalakes",
+ {
+ "name": self.metalake,
+ "comment": "MCP authz integration test metalake",
+ "properties": {},
+ },
+ )
+
+ def _create_model_catalog(self, name: str) -> None:
+ self._post(
+ f"/api/metalakes/{self.metalake}/catalogs",
+ {
+ "name": name,
+ "type": "MODEL",
+ "provider": "model",
+ "comment": "model catalog for authz test",
+ "properties": {},
+ },
+ )
+
+ def _add_user(self, user: str) -> None:
+ self._post(
+ f"/api/metalakes/{self.metalake}/users",
+ {"name": user},
+ )
+
+ def _create_reader_role(self) -> None:
+ # Grant bob USE_CATALOG on the allowed catalog only.
+ self._post(
+ f"/api/metalakes/{self.metalake}/roles",
+ {
+ "name": self.role_name,
+ "properties": {},
+ "securableObjects": [
+ {
+ "fullName": self.catalog_allowed,
+ "type": "CATALOG",
+ "privileges": [
+ {"name": "USE_CATALOG", "condition": "ALLOW"}
+ ],
+ }
+ ],
+ },
+ )
+
+ def _grant_role_to_user(self) -> None:
+ self._put(
+ f"/api/metalakes/{self.metalake}/permissions"
+ f"/users/{self.granted_user}/grant/",
+ {"roleNames": [self.role_name]},
+ )
+
+ def close(self) -> None:
+ self._client.close()
diff --git a/mcp-server/tests/integration/test_authz_e2e.py
b/mcp-server/tests/integration/test_authz_e2e.py
new file mode 100644
index 0000000000..488e6245db
--- /dev/null
+++ b/mcp-server/tests/integration/test_authz_e2e.py
@@ -0,0 +1,175 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""End-to-end authorization integration test through the live MCP HTTP server.
+
+Validates the three acceptance scenarios against a real Gravitino with
+authorization enabled:
+
+ 1. Two principals run the same discovery call and get correctly different,
+ authorization-scoped results.
+ 2. A read-only principal attempts a write through MCP and is denied by
+ Gravitino authorization.
+ 3. Both the reads and the denied write appear as audit records attributed to
+ the correct principal.
+
+These tests only run when GRAVITINO_URI / MCP_URL / MCP_METALAKE are set (see
+conftest.py); otherwise the suite is skipped.
+"""
+
+import asyncio
+import json
+import os
+import time
+
+import pytest
+from fastmcp import Client
+from fastmcp.client.transports import StreamableHttpTransport
+
+from tests.integration.gravitino_setup import basic_auth_header
+
+ADMIN = "admin"
+BOB = "bob"
+CATALOG_ALLOWED = "cat_allowed"
+CATALOG_DENIED = "cat_denied"
+
+
+def _client_for(principal: str, mcp_url: str) -> Client:
+ """Build an MCP client that authenticates as ``principal`` (simple
auth)."""
+ transport = StreamableHttpTransport(
+ url=mcp_url,
+ headers={"Authorization": basic_auth_header(principal)},
+ )
+ return Client(transport)
+
+
+async def _list_catalog_names(principal: str, mcp_url: str) -> set:
+ async with _client_for(principal, mcp_url) as client:
+ result = await client.call_tool("get_list_of_catalogs")
+ payload = json.loads(result.content[0].text)
+ return {entry["name"] for entry in payload}
+
+
+def test_authorization_scoped_discovery(gravitino_fixture, integration_env):
+ """Admin and bob get different, correctly scoped catalog lists."""
+ mcp_url = integration_env["mcp_url"]
+
+ admin_catalogs = asyncio.run(_list_catalog_names(ADMIN, mcp_url))
+ bob_catalogs = asyncio.run(_list_catalog_names(BOB, mcp_url))
+
+ # Admin owns the metalake and sees both catalogs.
+ assert CATALOG_ALLOWED in admin_catalogs
+ assert CATALOG_DENIED in admin_catalogs
+
+ # Bob was granted USE_CATALOG on the allowed catalog only.
+ assert CATALOG_ALLOWED in bob_catalogs
+ assert CATALOG_DENIED not in bob_catalogs
+
+ # The two principals must receive different results.
+ assert admin_catalogs != bob_catalogs
+
+
+def test_write_denied_for_readonly_principal(
+ gravitino_fixture, integration_env
+):
+ """Bob (no write grant) is denied when creating a tag through MCP."""
+ mcp_url = integration_env["mcp_url"]
+
+ async def _attempt_write():
+ async with _client_for(BOB, mcp_url) as client:
+ await client.call_tool(
+ "create_tag",
+ {
+ "name": "denied_tag",
+ "comment": "should be denied",
+ "properties": {},
+ },
+ )
+
+ with pytest.raises(Exception) as exc_info: # noqa: B017
+ asyncio.run(_attempt_write())
+
+ # The failure must be an authorization denial, not an unrelated error
+ # (transport failure, missing tool, server crash, ...).
+ message = str(exc_info.value).lower()
+ assert any(
+ token in message
+ for token in (
+ "forbidden",
+ "unauthorized",
+ "not authorized",
+ "permission",
+ "denied",
+ "access",
+ "403",
+ )
+ ), f"expected an authorization denial, got: {exc_info.value!r}"
+
+
+def test_audit_trail_attribution(gravitino_fixture, integration_env):
+ """Audit log records reads/writes attributed to the right principal."""
+ audit_log = os.environ.get("MCP_AUDIT_LOG")
+ if not audit_log or not os.path.exists(audit_log):
+ pytest.skip("MCP_AUDIT_LOG not set or file missing")
+
+ mcp_url = integration_env["mcp_url"]
+
+ async def _attempt_write():
+ async with _client_for(BOB, mcp_url) as client:
+ await client.call_tool(
+ "create_tag",
+ {"name": "audit_tag", "comment": "", "properties": {}},
+ )
+
+ # Generate one allowed read (admin) and one denied write (bob).
+ asyncio.run(_list_catalog_names(ADMIN, mcp_url))
+ try:
+ asyncio.run(_attempt_write())
+ except Exception: # pylint: disable=broad-exception-caught
+ # Expected denial for an unauthorized principal.
+ pass
+
+ # Give the server a moment to flush the audit handler.
+ time.sleep(1.0)
+
+ records = []
+ with open(audit_log, "r", encoding="utf-8") as fh:
+ for line in fh:
+ line = line.strip()
+ if line:
+ try:
+ records.append(json.loads(line))
+ except json.JSONDecodeError:
+ continue
+
+ admin_allows = [
+ r
+ for r in records
+ if r.get("principal") == ADMIN
+ and r.get("tool") == "get_list_of_catalogs"
+ and r.get("outcome") == "allow"
+ ]
+ bob_denies = [
+ r
+ for r in records
+ if r.get("principal") == BOB
+ and r.get("tool") == "create_tag"
+ and r.get("outcome") == "deny"
+ ]
+
+ assert admin_allows, "expected an allow record attributed to admin"
+ assert bob_denies, "expected a deny record attributed to bob"
diff --git a/mcp-server/tests/unit/test_audit.py
b/mcp-server/tests/unit/test_audit.py
new file mode 100644
index 0000000000..b338bab890
--- /dev/null
+++ b/mcp-server/tests/unit/test_audit.py
@@ -0,0 +1,232 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import asyncio
+import json
+import logging
+import unittest
+
+from fastmcp import Client
+
+from mcp_server.client.factory import RESTClientFactory
+from mcp_server.client.plain.exception import GravitinoException
+from mcp_server.client.plain.plain_rest_client_operation import (
+ PlainRESTClientOperation,
+)
+from mcp_server.core import audit
+from mcp_server.core.setting import Setting
+from mcp_server.server import GravitinoMCPServer
+from tests.unit.tools import MockOperation
+
+# Tests intentionally exercise module-private helpers (e.g.
audit._extract_principal)
+# and client internals; protected access is expected here.
+# pylint: disable=protected-access
+
+
+class TestAuditEmit(unittest.TestCase):
+ """Unit tests for the audit.emit() function."""
+
+ def setUp(self):
+ # tests/unit/tools/__init__.py calls logging.disable(logging.INFO).
+ # Re-enable here so audit records flow through; restore on teardown.
+ logging.disable(logging.NOTSET)
+ self.log_records = []
+ self.handler = _CapturingHandler(self.log_records)
+ audit_logger = logging.getLogger("gravitino.mcp.audit")
+ audit_logger.addHandler(self.handler)
+ audit_logger.setLevel(logging.INFO)
+ audit_logger.propagate = False
+
+ def tearDown(self):
+ logging.disable(logging.INFO)
+ audit_logger = logging.getLogger("gravitino.mcp.audit")
+ audit_logger.removeHandler(self.handler)
+ audit_logger.propagate = True
+
+ def test_allow_record_structure(self):
+ """emit() writes a JSON record with all required fields on allow."""
+ audit.emit(
+ principal="bearer:abc12345", tool="list_catalogs", outcome="allow"
+ )
+
+ self.assertEqual(len(self.log_records), 1)
+ record = json.loads(self.log_records[0])
+ self.assertEqual(record["principal"], "bearer:abc12345")
+ self.assertEqual(record["tool"], "list_catalogs")
+ self.assertEqual(record["outcome"], "allow")
+ self.assertIn("timestamp", record)
+ self.assertNotIn("error_type", record)
+
+ def test_deny_record_includes_error_type(self):
+ """emit() includes error_type in the record when outcome is deny."""
+ audit.emit(
+ principal="bearer:xyz99999",
+ tool="create_tag",
+ outcome="deny",
+ error_type="GravitinoException",
+ )
+
+ record = json.loads(self.log_records[0])
+ self.assertEqual(record["outcome"], "deny")
+ self.assertEqual(record["error_type"], "GravitinoException")
+
+ def test_anonymous_principal(self):
+ """emit() works with anonymous principal."""
+ audit.emit(
+ principal="anonymous", tool="get_list_of_catalogs", outcome="allow"
+ )
+ record = json.loads(self.log_records[0])
+ self.assertEqual(record["principal"], "anonymous")
+
+
+class TestExtractPrincipal(unittest.TestCase):
+ """Unit tests for audit._extract_principal()."""
+
+ def test_bearer_token_truncated_to_8_chars(self):
+ self.assertEqual(
+ audit._extract_principal("Bearer abcdefghijklmnop"),
+ "bearer:abcdefgh",
+ )
+
+ def test_empty_header_returns_anonymous(self):
+ self.assertEqual(audit._extract_principal(""), "anonymous")
+
+ def test_none_like_empty_returns_anonymous(self):
+ self.assertEqual(audit._extract_principal(None), "anonymous")
+
+ def test_short_token_uses_full_token(self):
+ self.assertEqual(audit._extract_principal("Bearer abc"), "bearer:abc")
+
+ def test_basic_auth_decodes_user(self):
+ """Simple auth header 'Basic base64(alice:dummy)' -> principal
'alice'."""
+ # base64("alice:dummy") == "YWxpY2U6ZHVtbXk="
+ self.assertEqual(
+ audit._extract_principal("Basic YWxpY2U6ZHVtbXk="), "alice"
+ )
+
+ def test_basic_auth_invalid_base64_returns_anonymous(self):
+ self.assertEqual(
+ audit._extract_principal("Basic not-valid-base64!!"), "anonymous"
+ )
+
+ def test_unknown_scheme_returns_anonymous(self):
+ self.assertEqual(
+ audit._extract_principal("Negotiate abc123"), "anonymous"
+ )
+
+
+class TestAuditMiddlewareIntegration(unittest.TestCase):
+ """Integration tests: AuditMiddleware emits records via the full MCP tool
path."""
+
+ def setUp(self):
+ logging.disable(logging.NOTSET)
+ self.log_records = []
+ self.handler = _CapturingHandler(self.log_records)
+ audit_logger = logging.getLogger("gravitino.mcp.audit")
+ audit_logger.addHandler(self.handler)
+ audit_logger.setLevel(logging.INFO)
+ audit_logger.propagate = False
+
+ RESTClientFactory.set_rest_client(MockOperation)
+ server = GravitinoMCPServer(Setting("mock_metalake"))
+ self.mcp = server.mcp
+
+ def tearDown(self):
+ logging.disable(logging.INFO)
+ audit_logger = logging.getLogger("gravitino.mcp.audit")
+ audit_logger.removeHandler(self.handler)
+ audit_logger.propagate = True
+ # Restore original REST client so other tests are not affected.
+ RESTClientFactory.set_rest_client(PlainRESTClientOperation)
+
+ def test_successful_tool_call_emits_allow_record(self):
+ """A successful tool call produces an audit record with
outcome=allow."""
+
+ async def _run():
+ async with Client(self.mcp) as client:
+ await client.call_tool("get_list_of_catalogs")
+
+ asyncio.run(_run())
+
+ self.assertEqual(len(self.log_records), 1)
+ record = json.loads(self.log_records[0])
+ self.assertEqual(record["tool"], "get_list_of_catalogs")
+ self.assertEqual(record["outcome"], "allow")
+ self.assertEqual(record["principal"], "anonymous")
+
+ def test_principal_falls_back_to_startup_token(self):
+ """With no request header, the audit principal uses the startup
--token."""
+ RESTClientFactory.set_rest_client(MockOperation)
+ server = GravitinoMCPServer(
+ Setting("mock_metalake", token="abcdef123456")
+ )
+
+ async def _run():
+ async with Client(server.mcp) as client:
+ await client.call_tool("get_list_of_catalogs")
+
+ asyncio.run(_run())
+
+ record = json.loads(self.log_records[0])
+ self.assertEqual(record["principal"], "bearer:abcdef12")
+
+ def test_failed_tool_call_emits_deny_record(self):
+ """A tool call that raises an exception produces an audit record with
outcome=deny."""
+
+ class FailingOperation(MockOperation):
+ def as_catalog_operation(self):
+ return _FailingCatalogOperation()
+
+ RESTClientFactory.set_rest_client(FailingOperation)
+ server = GravitinoMCPServer(Setting("mock_metalake"))
+
+ async def _run():
+ async with Client(server.mcp) as client:
+ try:
+ await client.call_tool("get_list_of_catalogs")
+ except Exception: # pylint: disable=broad-exception-caught
+ pass
+
+ asyncio.run(_run())
+
+ deny_records = [
+ json.loads(r) for r in self.log_records if '"deny"' in r
+ ]
+ self.assertTrue(len(deny_records) >= 1)
+ self.assertEqual(deny_records[0]["tool"], "get_list_of_catalogs")
+ self.assertEqual(deny_records[0]["outcome"], "deny")
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+class _CapturingHandler(logging.Handler):
+ """Logging handler that stores formatted messages in a list."""
+
+ def __init__(self, records: list):
+ super().__init__()
+ self._records = records
+
+ def emit(self, record: logging.LogRecord) -> None:
+ self._records.append(self.format(record))
+
+
+class _FailingCatalogOperation:
+ async def get_list_of_catalogs(self) -> str:
+ raise GravitinoException("Error code: 1003, Error type: FORBIDDEN")
diff --git a/mcp-server/tests/unit/test_auth_flow.py
b/mcp-server/tests/unit/test_auth_flow.py
new file mode 100644
index 0000000000..b596654d84
--- /dev/null
+++ b/mcp-server/tests/unit/test_auth_flow.py
@@ -0,0 +1,192 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import asyncio
+import sys
+import unittest
+from unittest import mock
+
+from mcp_server.client.factory import RESTClientFactory
+from mcp_server.client.plain.plain_rest_client_operation import (
+ PlainRESTClientOperation,
+)
+from mcp_server.core.context import GravitinoContext
+from mcp_server.core.setting import Setting
+from mcp_server.main import _parse_args
+
+
+class _RealFactoryTestCase(unittest.TestCase):
+ """Base for tests that inspect the real PlainRESTClientOperation.
+
+ Other test modules globally swap RESTClientFactory to MockOperation without
+ restoring it, so pin the real client here to stay order-independent.
+ """
+
+ def setUp(self):
+ RESTClientFactory.set_rest_client(PlainRESTClientOperation)
+
+
+def _shared_rest_client(operation: PlainRESTClientOperation):
+ # pylint: disable=protected-access
+ return operation._catalog_operation.rest_client
+
+
+def _headers_of(operation: PlainRESTClientOperation):
+ return _shared_rest_client(operation).headers
+
+
+def _close(operation: PlainRESTClientOperation):
+ asyncio.run(_shared_rest_client(operation).aclose())
+
+
+class TestAuthorizationInjection(_RealFactoryTestCase):
+ """Verify the Authorization header is forwarded verbatim to the httpx
client."""
+
+ def test_bearer_authorization_header(self):
+ """A Bearer authorization value is forwarded unchanged."""
+ client = PlainRESTClientOperation(
+ "my_metalake",
+ "http://localhost:8090";,
+ authorization="Bearer my-secret-token",
+ )
+ try:
+ self.assertEqual(
+ _headers_of(client).get("Authorization"),
+ "Bearer my-secret-token",
+ )
+ finally:
+ _close(client)
+
+ def test_basic_authorization_header(self):
+ """A Basic authorization value (simple auth) is forwarded unchanged."""
+ client = PlainRESTClientOperation(
+ "my_metalake",
+ "http://localhost:8090";,
+ authorization="Basic YWxpY2U6ZHVtbXk=",
+ )
+ try:
+ self.assertEqual(
+ _headers_of(client).get("Authorization"),
+ "Basic YWxpY2U6ZHVtbXk=",
+ )
+ finally:
+ _close(client)
+
+ def test_empty_authorization_no_header(self):
+ """When authorization is empty, no Authorization header is added."""
+ client = PlainRESTClientOperation(
+ "my_metalake", "http://localhost:8090";, authorization=""
+ )
+ try:
+ self.assertIsNone(_headers_of(client).get("Authorization"))
+ finally:
+ _close(client)
+
+ def test_no_authorization_argument_no_header(self):
+ """When authorization argument is omitted, no Authorization header is
added."""
+ client = PlainRESTClientOperation(
+ "my_metalake", "http://localhost:8090";
+ )
+ try:
+ self.assertIsNone(_headers_of(client).get("Authorization"))
+ finally:
+ _close(client)
+
+
+class TestSettingTokenMasking(unittest.TestCase):
+ """Verify that the token is not exposed in Setting string
representations."""
+
+ def test_token_masked_in_str(self):
+ """Token value must not appear in Setting.__str__."""
+ setting = Setting(metalake="ml", token="super-secret-token-value")
+ self.assertNotIn("super-secret-token-value", str(setting))
+ self.assertIn("***", str(setting))
+
+ def test_token_not_in_repr(self):
+ """Token value must not appear in Setting.__repr__ either."""
+ setting = Setting(metalake="ml", token="super-secret-token-value")
+ self.assertNotIn("super-secret-token-value", repr(setting))
+
+ def test_empty_token_shows_empty_in_str(self):
+ """When no token is set, __str__ shows empty placeholder."""
+ setting = Setting(metalake="ml", token="")
+ self.assertNotIn("***", str(setting))
+
+
+class TestTokenArgParsing(unittest.TestCase):
+ """Verify --token CLI argument and GRAVITINO_TOKEN env var precedence."""
+
+ def test_env_var_used_when_token_omitted(self):
+ """GRAVITINO_TOKEN is used when --token is not passed."""
+ with mock.patch.dict(
+ "os.environ", {"GRAVITINO_TOKEN": "env-token"}
+ ), mock.patch.object(sys, "argv", ["prog", "--metalake", "ml"]):
+ args = _parse_args()
+ self.assertEqual(args.token, "env-token")
+
+ def test_cli_token_overrides_env_var(self):
+ """--token takes precedence over GRAVITINO_TOKEN when both are set."""
+ with mock.patch.dict(
+ "os.environ", {"GRAVITINO_TOKEN": "env-token"}
+ ), mock.patch.object(
+ sys, "argv", ["prog", "--metalake", "ml", "--token", "cli-token"]
+ ):
+ args = _parse_args()
+ self.assertEqual(args.token, "cli-token")
+
+ def test_no_token_anywhere_defaults_to_empty(self):
+ """Without --token and GRAVITINO_TOKEN, token defaults to empty
string."""
+ with mock.patch.dict("os.environ", {}, clear=True), mock.patch.object(
+ sys, "argv", ["prog", "--metalake", "ml"]
+ ):
+ args = _parse_args()
+ self.assertEqual(args.token, "")
+
+
+class TestGravitinoContextTokenPropagation(_RealFactoryTestCase):
+ """Verify GravitinoContext passes token from Setting to the REST client."""
+
+ def test_context_propagates_token(self):
+ """Token from Setting reaches the httpx client Authorization header."""
+ setting = Setting(
+ metalake="ml",
+ gravitino_uri="http://localhost:8090";,
+ token="ctx-token-xyz",
+ )
+ ctx = GravitinoContext(setting)
+ rest_client = ctx.rest_client()
+ try:
+ self.assertEqual(
+ _headers_of(rest_client).get("Authorization"),
+ "Bearer ctx-token-xyz",
+ )
+ finally:
+ _close(rest_client)
+
+ def test_context_anonymous_when_no_token(self):
+ """Empty token in Setting → no Authorization header in REST calls."""
+ setting = Setting(
+ metalake="ml",
+ gravitino_uri="http://localhost:8090";,
+ token="",
+ )
+ ctx = GravitinoContext(setting)
+ rest_client = ctx.rest_client()
+ try:
+ self.assertIsNone(_headers_of(rest_client).get("Authorization"))
+ finally:
+ _close(rest_client)
diff --git a/mcp-server/tests/unit/test_per_request_token.py
b/mcp-server/tests/unit/test_per_request_token.py
new file mode 100644
index 0000000000..db9b4ad92c
--- /dev/null
+++ b/mcp-server/tests/unit/test_per_request_token.py
@@ -0,0 +1,242 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for per-request identity isolation (Task 6).
+
+GravitinoContext.rest_client() must forward the current HTTP request's raw
+Authorization header (any scheme) to Gravitino, not the shared startup token,
+so concurrent multi-principal sessions stay fully isolated in HTTP mode.
+"""
+
+import asyncio
+import unittest
+from unittest.mock import MagicMock, patch
+
+from mcp_server.client.factory import RESTClientFactory
+from mcp_server.client.plain.plain_rest_client_operation import (
+ PlainRESTClientOperation,
+)
+from mcp_server.core import context as context_module
+from mcp_server.core.context import (
+ GravitinoContext,
+ _get_request_authorization,
+)
+from mcp_server.core.setting import Setting
+
+# Tests intentionally exercise context/client internals (e.g. _default_client,
+# _catalog_operation) to assert per-request isolation; protected access is
expected.
+# pylint: disable=protected-access
+
+
+class TestGetRequestAuthorization(unittest.TestCase):
+ """Unit tests for _get_request_authorization() (HTTP context
extraction)."""
+
+ def test_returns_raw_bearer_header(self):
+ mock_request = MagicMock()
+ mock_request.headers.get.return_value = "Bearer request-token-xyz"
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=mock_request,
+ ):
+ authorization = _get_request_authorization()
+
+ self.assertEqual(authorization, "Bearer request-token-xyz")
+
+ def test_returns_raw_basic_header_verbatim(self):
+ """Basic (simple auth) headers must pass through unchanged, not be
dropped."""
+ mock_request = MagicMock()
+ mock_request.headers.get.return_value = "Basic YWxpY2U6ZHVtbXk="
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=mock_request,
+ ):
+ authorization = _get_request_authorization()
+
+ self.assertEqual(authorization, "Basic YWxpY2U6ZHVtbXk=")
+
+ def test_returns_empty_when_no_http_context(self):
+ """Simulates stdio mode where get_http_request raises LookupError."""
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ side_effect=LookupError("no request context"),
+ ):
+ authorization = _get_request_authorization()
+
+ self.assertEqual(authorization, "")
+
+ def test_returns_empty_when_no_authorization_header(self):
+ mock_request = MagicMock()
+ mock_request.headers.get.return_value = ""
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=mock_request,
+ ):
+ authorization = _get_request_authorization()
+
+ self.assertEqual(authorization, "")
+
+
+class TestGravitinoContextPerRequestAuthorization(unittest.TestCase):
+ """GravitinoContext.rest_client() isolates per-request identities."""
+
+ def setUp(self):
+ # These tests inspect the real PlainRESTClientOperation; other test
+ # modules swap the factory to MockOperation without restoring it, so
pin
+ # the real client here to stay order-independent.
+ RESTClientFactory.set_rest_client(PlainRESTClientOperation)
+
+ def _make_context(self, startup_token: str = "") -> GravitinoContext:
+ return GravitinoContext(
+ Setting(
+ metalake="ml",
+ gravitino_uri="http://localhost:8090";,
+ token=startup_token,
+ )
+ )
+
+ def test_per_request_header_overrides_startup_token(self):
+ """An HTTP request's Authorization header takes priority over the
startup token."""
+ ctx = self._make_context(startup_token="startup-token")
+
+ mock_request = MagicMock()
+ mock_request.headers.get.return_value = "Basic YWxpY2U6ZHVtbXk="
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=mock_request,
+ ):
+ client = ctx.rest_client()
+
+ headers = dict(client._catalog_operation.rest_client.headers)
+ self.assertEqual(headers.get("authorization"), "Basic
YWxpY2U6ZHVtbXk=")
+
+ def test_falls_back_to_default_client_when_no_request_header(self):
+ """With no per-request header, the shared default client (startup
token) is used."""
+ ctx = self._make_context(startup_token="startup-token")
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ side_effect=LookupError,
+ ):
+ client = ctx.rest_client()
+
+ # Must be the exact same object as the cached default client.
+ self.assertIs(client, ctx._default_client)
+
+ def test_two_concurrent_requests_get_different_clients(self):
+ """Different request identities must produce different client
instances."""
+ ctx = self._make_context()
+
+ def make_mock(authorization: str) -> MagicMock:
+ m = MagicMock()
+ m.headers.get.return_value = authorization
+ return m
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=make_mock("Basic YWxpY2U6ZHVtbXk="),
+ ):
+ client_alice = ctx.rest_client()
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=make_mock("Basic Ym9iOmR1bW15"),
+ ):
+ client_bob = ctx.rest_client()
+
+ alice_headers = dict(
+ client_alice._catalog_operation.rest_client.headers
+ )
+ bob_headers = dict(client_bob._catalog_operation.rest_client.headers)
+ self.assertEqual(
+ alice_headers.get("authorization"), "Basic YWxpY2U6ZHVtbXk="
+ )
+ self.assertEqual(bob_headers.get("authorization"), "Basic
Ym9iOmR1bW15")
+ self.assertIsNot(client_alice, client_bob)
+
+ def test_same_principal_reuses_cached_client(self):
+ """Repeated calls from the same principal reuse one cached client."""
+ ctx = self._make_context()
+
+ mock_request = MagicMock()
+ mock_request.headers.get.return_value = "Basic YWxpY2U6ZHVtbXk="
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=mock_request,
+ ):
+ first = ctx.rest_client()
+ second = ctx.rest_client()
+
+ # Same Authorization header -> same client object (connection pool
reused).
+ self.assertIs(first, second)
+
+ def test_client_cache_is_bounded(self):
+ """The per-principal client cache evicts the oldest entries past its
cap."""
+ ctx = self._make_context()
+ cap = context_module._MAX_CACHED_CLIENTS
+
+ for i in range(cap + 5):
+ mock_request = MagicMock()
+ mock_request.headers.get.return_value = f"Bearer token-{i}"
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=mock_request,
+ ):
+ ctx.rest_client()
+
+ self.assertLessEqual(len(ctx._clients_by_auth), cap)
+
+ def test_evicted_client_is_closed(self):
+ """On eviction (with a running loop) the evicted client's pool is
closed."""
+ closed = []
+
+ class _ClosableClient:
+ def __init__(self, *_args, **_kwargs):
+ pass
+
+ async def close(self):
+ closed.append(self)
+
+ RESTClientFactory.set_rest_client(_ClosableClient)
+ try:
+ ctx = self._make_context()
+ cap = context_module._MAX_CACHED_CLIENTS
+
+ async def _drive():
+ for i in range(cap + 1):
+ mock_request = MagicMock()
+ mock_request.headers.get.return_value = f"Bearer t-{i}"
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=mock_request,
+ ):
+ ctx.rest_client()
+ # Let the scheduled close task run to completion.
+ await asyncio.sleep(0)
+ await asyncio.gather(*ctx._pending_closes)
+
+ asyncio.run(_drive())
+
+ # Exactly one client (the oldest) was evicted and closed.
+ self.assertEqual(len(closed), 1)
+ self.assertEqual(len(ctx._pending_closes), 0)
+ finally:
+ RESTClientFactory.set_rest_client(PlainRESTClientOperation)
diff --git a/mcp-server/tests/unit/test_transport_tls.py
b/mcp-server/tests/unit/test_transport_tls.py
new file mode 100644
index 0000000000..a7cba980ba
--- /dev/null
+++ b/mcp-server/tests/unit/test_transport_tls.py
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Tests for HTTP transport URL parsing, the streamable-http alias, and TLS
wiring."""
+
+import unittest
+from unittest.mock import patch
+
+from mcp_server.client.factory import RESTClientFactory
+from mcp_server.client.plain.plain_rest_client_operation import (
+ PlainRESTClientOperation,
+)
+from mcp_server.core.setting import Setting
+from mcp_server.server import GravitinoMCPServer, _parse_mcp_url
+from tests.unit.tools import MockOperation
+
+
+class TestParseMcpUrl(unittest.TestCase):
+ """_parse_mcp_url accepts http and https and rejects other schemes."""
+
+ def test_http_url(self):
+ self.assertEqual(
+ _parse_mcp_url("http://127.0.0.1:8000/mcp";),
+ ("127.0.0.1", 8000, "/mcp"),
+ )
+
+ def test_https_url(self):
+ self.assertEqual(
+ _parse_mcp_url("https://mcphost:9443/mcp";),
+ ("mcphost", 9443, "/mcp"),
+ )
+
+ def test_https_default_port(self):
+ _, port, _ = _parse_mcp_url("https://mcphost/mcp";)
+ self.assertEqual(port, 443)
+
+ def test_http_default_port(self):
+ _, port, _ = _parse_mcp_url("http://mcphost/mcp";)
+ self.assertEqual(port, 80)
+
+ def test_unsupported_scheme_rejected(self):
+ with self.assertRaises(ValueError):
+ _parse_mcp_url("ftp://mcphost/mcp";)
+
+
+class TestRunHttpTransport(unittest.TestCase):
+ """GravitinoMCPServer.run() wires transport name and TLS config
correctly."""
+
+ def setUp(self):
+ RESTClientFactory.set_rest_client(MockOperation)
+
+ def tearDown(self):
+ # Restore the default so this global mutation can't leak into other
tests.
+ RESTClientFactory.set_rest_client(PlainRESTClientOperation)
+
+ def _run_and_capture(self, setting: Setting) -> dict:
+ """Run the server with run_async patched; return the kwargs it was
called with."""
+ server = GravitinoMCPServer(setting)
+ captured = {}
+
+ async def fake_run_async(**kwargs):
+ captured.update(kwargs)
+
+ with patch.object(server.mcp, "run_async", side_effect=fake_run_async):
+ server.run()
+ return captured
+
+ def test_http_transport(self):
+ setting = Setting(
+ metalake="ml",
+ transport="http",
+ mcp_url="http://127.0.0.1:8000/mcp";,
+ )
+ kwargs = self._run_and_capture(setting)
+ self.assertEqual(kwargs["transport"], "http")
+ self.assertEqual(kwargs["host"], "127.0.0.1")
+ self.assertEqual(kwargs["port"], 8000)
+ self.assertEqual(kwargs["path"], "/mcp")
+ self.assertNotIn("uvicorn_config", kwargs)
+
+ def test_streamable_http_alias(self):
+ setting = Setting(
+ metalake="ml",
+ transport="streamable-http",
+ mcp_url="http://127.0.0.1:8000/mcp";,
+ )
+ kwargs = self._run_and_capture(setting)
+ self.assertEqual(kwargs["transport"], "streamable-http")
+
+ def test_tls_config_wired_when_cert_and_key_set(self):
+ setting = Setting(
+ metalake="ml",
+ transport="streamable-http",
+ mcp_url="https://127.0.0.1:8443/mcp";,
+ tls_cert="/path/to/cert.pem",
+ tls_key="/path/to/key.pem",
+ )
+ kwargs = self._run_and_capture(setting)
+ self.assertEqual(
+ kwargs["uvicorn_config"],
+ {
+ "ssl_certfile": "/path/to/cert.pem",
+ "ssl_keyfile": "/path/to/key.pem",
+ },
+ )
+
+ def test_lone_cert_rejected(self):
+ """TLS requires both cert and key; a lone cert is rejected."""
+ setting = Setting(
+ metalake="ml",
+ transport="http",
+ mcp_url="https://127.0.0.1:8443/mcp";,
+ tls_cert="/path/to/cert.pem",
+ tls_key="",
+ )
+ with self.assertRaises(ValueError):
+ self._run_and_capture(setting)
+
+ def test_https_url_without_tls_rejected(self):
+ """An https URL without cert/key must not silently serve plain HTTP."""
+ setting = Setting(
+ metalake="ml",
+ transport="http",
+ mcp_url="https://127.0.0.1:8443/mcp";,
+ )
+ with self.assertRaises(ValueError):
+ self._run_and_capture(setting)
+
+ def test_http_url_with_tls_rejected(self):
+ """TLS configured behind an http URL is a misconfiguration and
rejected."""
+ setting = Setting(
+ metalake="ml",
+ transport="http",
+ mcp_url="http://127.0.0.1:8000/mcp";,
+ tls_cert="/path/to/cert.pem",
+ tls_key="/path/to/key.pem",
+ )
+ with self.assertRaises(ValueError):
+ self._run_and_capture(setting)
diff --git a/mcp-server/tests/unit/tools/mock_operation.py
b/mcp-server/tests/unit/tools/mock_operation.py
index d76ba94b60..e1fc540442 100644
--- a/mcp-server/tests/unit/tools/mock_operation.py
+++ b/mcp-server/tests/unit/tools/mock_operation.py
@@ -31,7 +31,7 @@ from mcp_server.client.statistic_operation import
StatisticOperation
class MockOperation(GravitinoOperation):
- def __init__(self, metalake, uri):
+ def __init__(self, metalake, uri, authorization=""):
pass
def as_table_operation(self) -> TableOperation:
