Skip to content

DX-115085: MCP server hardening — auth, validation, annotations, and smoketests#86

Merged
aniket-s-kulkarni merged 1 commit into
dremio:mainfrom
aniket-s-kulkarni:dx-115085
Feb 23, 2026
Merged

DX-115085: MCP server hardening — auth, validation, annotations, and smoketests#86
aniket-s-kulkarni merged 1 commit into
dremio:mainfrom
aniket-s-kulkarni:dx-115085

Conversation

@aniket-s-kulkarni
Copy link
Copy Markdown
Contributor

@aniket-s-kulkarni aniket-s-kulkarni commented Feb 23, 2026

Summary

  • Fix GetDescriptionOfTableOrSchema auth: add missing @secured and @with_metrics decorators so OAuth tokens are injected
  • Add input validation to GetSchemaOfTable: empty string/list inputs now return a descriptive error instead of crashing with IndexError
  • Add MCP ToolAnnotations: all tools now declare readOnlyHint/destructiveHint; RunSqlQuery.destructiveHint is conditional on the allow_dml setting
  • Sanitize error messages: DremioAsyncHttpClient no longer leaks URI/PAT in errors; GetTableOrViewLineage returns a sanitized error dict on failure; RunSqlQuery returns a clean error for DML rejection
  • Rename RunSqlQuery parameter squery for clarity
  • Improve GetSchemaOfTable docstring with explicit format examples (dot-separated string vs list of path components)
  • Expand GetUsefulSystemTableNames from 1 entry to 6 (added jobs_recent, engines, users, COLUMNS, VIEWS)
  • Add 11 smoketest cases to stremable_http_cli.py covering all fixes, with --token, --local, --check-annotations, --check-new-contract flags
  • Fix ContextVar propagation in local MCP server: settings are now propagated to the server thread using model_copy().with_overrides() so the correct Dremio URI is used

Test plan

  • All 208 unit tests pass
  • Run remote smoketest: uv run python tests/stremable_http_cli.py cli test --token <TOKEN> --url <MCP_URL> --no-check-annotations --no-check-new-contract
  • Run local smoketest: uv run python tests/stremable_http_cli.py cli test --token <TOKEN> --url <MCP_URL> --local

…smoketests

   Critical fixes:
   - Add @secured and @with_metrics decorators to GetDescriptionOfTableOrSchema
     so OAuth tokens are properly injected
   - Add input validation to GetSchemaOfTable for empty string/list inputs
   - Add MCP ToolAnnotations (readOnlyHint/destructiveHint) to all tools;
     RunSqlQuery destructiveHint is conditional on allow_dml setting

   Medium fixes:
   - Sanitize error message in DremioAsyncHttpClient to not leak URI/PAT
   - Wrap GetTableOrViewLineage in try/except with sanitized error response
   - Return clean error dict for DML rejection in RunSqlQuery
   - Rename RunSqlQuery parameter from 's' to 'query' for clarity
   - Update GetSchemaOfTable docstring with explicit format examples
   - Expand GetUsefulSystemTableNames from 1 to 6 entries (jobs_recent,
     engines, users, COLUMNS, VIEWS)

   Smoketest improvements:
   - Add 11 smoketest cases to stremable_http_cli.py covering all fixes
   - Add --token, --local, --check-annotations, --check-new-contract flags
   - Add _local_mcp_server context manager with ContextVar propagation
     to server thread via settings.with_overrides()
@aniket-s-kulkarni aniket-s-kulkarni enabled auto-merge (squash) February 23, 2026 15:34
@aniket-s-kulkarni
Copy link
Copy Markdown
Contributor Author

Doing a full test run against Dremio cloud -

$ uv run python tests/stremable_http_cli.py test --local --url https://mcp.dremio.cloud/mcp/c63ebad3-..... --token "$TOKEN"
Starting local MCP server (dremio.uri=https://api.dremio.cloud, project_id=DREMIO_DYNAMIC)..
StreamableHTTP session manager started
Local server ready at http://127.0.0.1:10552/mcp/c63ebad3-...../
Connecting to server..
Created new transport with session ID: 6d3072a43d5642cfbbb65c0c68761cd9
HTTP Request: POST http://127.0.0.1:10552/mcp/c63ebad3-...../ "HTTP/1.1 200 OK"
Received session ID: 6d3072a43d5642cfbbb65c0c68761cd9
Negotiated protocol version: 2025-06-18
Processing request of type ListToolsRequest
HTTP Request: POST http://127.0.0.1:10552/mcp/c63ebad3-...../ "HTTP/1.1 200 OK"
['RunSqlQuery', 'GetUsefulSystemTableNames', 'GetSchemaOfTable', 'GetTableOrViewLineage', 'SearchTableAndViews', 'GetDescriptionOfTableOrSchema']
Checking tool annotations..
  RunSqlQuery: readOnly=True, destructive=False
OK
Checking RunSqlQuery parameter name.. OK
Checking RunSqlQuery SELECT.. 
Processing request of type CallToolRequest
HTTP Request: POST http://127.0.0.1:10552/mcp/c63ebad3-...../ "HTTP/1.1 200 OK"
[{'n': 1771859448}]
OK
Checking RunSqlQuery DML rejection.. 
Processing request of type CallToolRequest
HTTP Request: POST http://127.0.0.1:10552/mcp/c63ebad3-...../ "HTTP/1.1 200 OK"
OK
Checking RunSqlQuery job tracking.. 
Processing request of type CallToolRequest
HTTP Request: POST http://127.0.0.1:10552/mcp/c63ebad3-...../ "HTTP/1.1 200 OK"
[{'query': '/* dremioai: submitter=RunSqlQuery */\nSELECT 1771859448 as n'}]
OK
Checking GetUsefulSystemTableNames.. 
Processing request of type CallToolRequest
HTTP Request: POST http://127.0.0.1:10552/mcp/c63ebad3-...../ "HTTP/1.1 200 OK"
6 tables returned
OK
Checking GetSchemaOfTable empty input validation.. 
Processing request of type CallToolRequest
HTTP Request: POST http://127.0.0.1:10552/mcp/c63ebad3-...../ "HTTP/1.1 200 OK"
OK
Checking GetSchemaOfTable with valid table.. 
Processing request of type CallToolRequest
HTTP Request: POST http://127.0.0.1:10552/mcp/c63ebad3-...../ "HTTP/1.1 200 OK"
4 fields returned
OK
Checking GetSchemaOfTable parameter docs.. OK
Checking GetDescriptionOfTableOrSchema auth.. 
Processing request of type CallToolRequest
HTTP Request: POST http://127.0.0.1:10552/mcp/c63ebad3-...../ "HTTP/1.1 200 OK"
OK
Checking GetTableOrViewLineage error sanitization..
Processing request of type CallToolRequest
HTTP Request: POST http://127.0.0.1:10552/mcp/c63ebad3-...../ "HTTP/1.1 200 OK"
OK
Terminating session: 6d3072a43d5642cfbbb65c0c68761cd9
HTTP Request: DELETE http://127.0.0.1:10552/mcp/c63ebad3-...../ "HTTP/1.1 200 OK"

@aniket-s-kulkarni aniket-s-kulkarni requested review from tevin831 and removed request for tevin831 February 23, 2026 17:15
@aniket-s-kulkarni aniket-s-kulkarni merged commit f585af5 into dremio:main Feb 23, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Development

Successfully merging this pull request may close these issues.

3 participants