MCP Server for DQX

[souravg-db2]

Changes

Adds an MCP (Model Context Protocol) server for DQX, exposing DQX's data-quality
capabilities as tools that any MCP-compatible AI agent (Claude, Genie Code, Cursor, Mosaic AI)
can discover and orchestrate. It runs as a Databricks App with on-behalf-of (OBO)
authentication, so all data access is governed by the calling user's Unity Catalog permissions.

Architecture

• Databricks App + FastMCP server, exposed over Streamable HTTP. The server itself carries
  no PySpark/DQX dependency.
• OBO governance: for any data access the server creates a temporary definer's-rights view
  over the source table using the user's forwarded token (X-Forwarded-Access-Token), so the
  service principal reads data as the user, never directly.
• Async job pattern: long-running operations are submitted to a pre-deployed runner notebook
  job and return a run_id immediately; the client polls get_run_result. get_run_status does
  a single non-blocking poll (the client drives cadence) — this matches MCP's canonical
  long-running-tool model and avoids holding the HTTP connection / saturating the worker pool.
• Genie Code compatibility: stateless_http + json_response + CORS preflight scoped to
  Databricks domains.

Tools

Tool	Purpose
get_workflow	Recommended sequence of tool calls (call first)
get_table_schema	Table columns/types (direct, OBO)
profile_table	Profile data → summary stats + profiles
generate_rules	Generate checks from profiler output
generate_rules_from_contract	Generate checks from an ODCS data contract (deterministic)
validate_checks	Validate check definitions
run_checks	Execute checks → counts + per-rule summary + failing sample
apply_checks_and_save_to_table	Operationalized run: write valid/quarantine rows to Delta tables
save_checks / load_checks	Persist / retrieve a rule set (table, UC volume, or workspace file)
list_available_checks	Discover all built-in DQX check functions
get_run_result	Poll an async job for status/results

Temp-view lifecycle (stateless, restart/replica-safe)

• The runner job drops its own input view in a finally (it runs as the SP, which owns the
  temp schema — see setup), so cleanup happens in the guaranteed job execution regardless of
  whether/where the user polls.
• A throttled TTL sweeper (timestamped v_<epoch>_<uuid> view names) reaps any orphans whose
  job never ran. No per-request server state is kept, so app restarts / multiple replicas don't
  leak views or lose context.

Security & governance

• Data-access boundary enforced by OBO definer's-rights views; the app SP owns only the scratch
  tmp schema (object lifecycle), not the underlying data.
• CORS restricted to Databricks workspace/app domains with credentials; SQL identifiers validated
  and backtick-quoted; log values sanitized (CWE-117); catalog name sourced from a Databricks secret.

Deployment

• Databricks Asset Bundle (databricks.yml): runner job, one-time setup job (UC grants + schema
  ownership), and the app. requirements.txt (the App's runtime manifest) and pyproject.toml are
  kept in sync; the runner installs databricks-labs-dqx[datacontract].
• New make targets and fmt/CI wiring for the sub-project.

Linked issues

Resolves #1045

Tests

• added unit tests

Layered MCP-server test suite (62 passing, deterministic, no workspace needed for CI):

• Handler unit tests (test_tools.py)
• Protocol tests via FastMCP's in-memory Client over the real MCP protocol (test_mcp_protocol.py)
• HTTP integration over the real ASGI app — health, initialize/tools-list, and OBO token
  propagation (test_app_http.py); CORS policy (test_cors.py)
• Agent-in-the-loop integration (gated/skipped without a deployed server + workspace LLM
  endpoint): a tool-calling model is handed the tool schemas and an instruction, and we assert it
  discovers and invokes the right tools (test_integration_agent.py)

Documentation and Demos

• added/updated docs (docs/dqx/docs/guide/dqx_mcp_server.mdx)

---

This description was written by Isaac.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 92.75%. Comparing base (67138ec) to head (808a430).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1252      +/-   ##
==========================================
+ Coverage   92.56%   92.75%   +0.18%     
==========================================
  Files         102      102              
  Lines       10429    10429              
==========================================
+ Hits         9654     9673      +19     
+ Misses        775      756      -19     

Flag	Coverage Δ
anomaly	54.02% <ø> (ø)
anomaly-serverless	54.03% <ø> (ø)
integration	50.26% <ø> (+0.01%)	⬆️
integration-serverless	50.78% <ø> (+0.16%)	⬆️
unit	58.83% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[github-actions]

✅ 799/799 passed, 42 skipped, 5h15m1s total

_{Running from acceptance #5156}

[github-actions]

✅ 194/194 passed, 2 skipped, 6h54m24s total

_{Running from anomaly #1270}

[mwojtyczka]

Review: MCP Server for DQX

Strong PR — the OBO + definer's-rights-view + SP-job architecture is the right pattern, and it neatly avoids needing a Jobs scope at the user auth level. Nice touches: the 4.5 MB notebook.exit() guard, pure-ASGI OBO middleware, idempotent setup grants, and a fixed OPERATIONS dispatch (no arbitrary code exec).

A few things to address before merge — inline comments below. Two summary-level points:

Rebase on main (stale base). The root pyproject.toml mypy-comment change reintroduces a reference to apx dev check, but PR #1223 (already merged) removed apx and replaced it with first-party scripts (bun run tsc -b + basedpyright). The branch is behind main; please rebase so this drift is resolved.

Add an integration test with Genie (see inline on the tests). app.py is explicitly tuned for "Genie Code compatibility" (stateless_http, json_response, CORS preflight) — but nothing exercises that path. An end-to-end test that drives the MCP server through Databricks Genie (Genie Code as the MCP client) over the deployed app would protect exactly the behaviour these settings exist for.

Priorities: the rebase are blockers; CORS and the in-memory run-state are important follow-ups.

[mwojtyczka]

Going in the right direction, left some comments to address

[mwojtyczka]

Re-reviewed the whole PR from scratch

[mwojtyczka]

Fresh review of the current head (medium effort). The earlier security items are verified fixed. Three net-new findings below — one is a genuine cross-user data-disclosure that defeats the OBO governance model, so worth blocking on.

[mwojtyczka]

Re-review of the current head (bdc8f4b) after the fix commits. All three findings from the prior review are verified fixed and covered by regression tests — I've resolved those threads.

• IDOR / cross-user result disclosure — submit_job_async now stamps requesting_user (from the OBO X-Forwarded-Email) as a declared job parameter, and get_run_status reads it back and denies the read (not_found, no result download) unless caller == submitter, with an explicit empty-either-side guard. Verified against the SDK (Run.job_parameters round-trip) and the guard is correctly positioned before the result download / RUNNING branch. Tests: test_denies_run_submitted_by_another_user, test_denies_when_caller_has_no_identity, test_denies_when_run_has_no_submitter, test_allows_run_submitted_by_same_user (all assert files.download is not called on deny). Because submit and get_run_result both traverse the same OBO middleware, the guard is symmetric — no legitimate-caller lockout.
• Temp-view leak on submit failure — new _submit_or_drop_view wrapper drops the OBO view on failure across all view-backed tools; file-backed load_checks correctly skips it. Test: test_drops_temp_view_if_submit_fails.
• execute_sql 30s-timeout misreport — now polls get_statement to completion (2s/120s) and hardens the None-status deref; get_statement returns the same StatementResponse (manifest+result populated on SUCCEEDED), so the chunk-follow stays valid. Test: test_none_status_raises_clean_error_not_attributeerror.

Proactive hardening in the same commits also checks out: leading-digit identifier rejection (unquoted-FQN safety), partial-grant reporting (access_granted_to null unless every created table was granted), and the removed dead to_local_fuse_path. The changed deploy/CI/docs (requesting_user declared at job level, ci_deploy.sh threads the runner SP, mcp.yml stays pull_request + fork-gated, docs anchors resolve) are all clean.

No new findings on this head. Remaining open threads are the two I intentionally left for the authors (the docs-placement suggestion) — the type-gating thread was resolved separately once mcp-check landed in CI.

[mwojtyczka]

LGTM - pending manual testing, and code coverage