MCP Server for DQX

.github/workflows/mcp.yml

@@ -0,0 +1,98 @@
+# MCP server integration tests.
+#
+# Runs the tests/integration_mcp suite via the shared acceptance harness (same workspace,
+# auth, and Model Serving endpoint as the anomaly/integration suites). The suite deploys an
+# isolated MCP app from the bundle, drives it as an agent would, and tears it down — see
+# tests/integration_mcp/conftest.py. No MCP-specific secrets are needed; auth comes from the
+# acceptance action exactly like the other integration workflows.
+name: mcp
+
+on:
+  pull_request:
+    types: [opened, synchronize, ready_for_review]
+    paths:
+      - "mcp-server/**"
+      - "src/databricks/labs/dqx/**" # the runner installs the in-repo DQX wheel
+      - "tests/integration_mcp/**"
+      - ".github/workflows/mcp.yml"
+  merge_group:
+    types: [checks_requested]
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: false # let teardown finish so the isolated app/bundle is removed
+
+jobs:
+  # Gate downstream work: fork PRs (no access to the tool environment) and drafts do not run.
+  not-a-fork:
+    runs-on:
+      group: databrickslabs-protected-runner-group
+      labels: linux-ubuntu-latest
+    if: github.event_name == 'pull_request' && !github.event.pull_request.draft && !github.event.pull_request.head.repo.fork
+    steps:
+      - run: echo "Not a fork PR, proceeding"
+
+  mcp-tests:
+    needs: not-a-fork
+    environment: tool
+    runs-on:
+      group: larger-runners
+      labels: larger
+    permissions:
+      id-token: write
+      pull-requests: write
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+        with:
+          fetch-depth: 0
+
+      - name: Setup environment
+        uses: ./.github/actions/setup-env
+
+      - name: Pre-build DQX wheel
+        uses: ./.github/actions/prebuild-wheel
+
+      - name: Prepare code coverage configuration for MCP tests
+        run: |
+          cat > tests/integration_mcp/.coveragerc << EOF
+          [run]
+          source =
+              ../../mcp-server/server
+              ../../mcp-server/runner/src/dqx_mcp_runner
+          relative_files = true
+          parallel = true
+          EOF
+
+      # Required for the DAB (Databricks Asset Bundle) deploy the MCP suite runs (ci_deploy.sh).
+      # The CLI lands on PATH (via GITHUB_PATH) for the acceptance action's test subprocess.
+      - name: Install Databricks CLI
+        uses: databricks/setup-cli@596b0a354ba14aa59921aca1b02bd67c2b0a81a5 # v0.297.2
+
+      - name: Run MCP integration tests
+        uses: databrickslabs/sandbox/acceptance@83461e5dd7021feabb1a9ca3ee10d6f46b72092a # acceptance/v0.4.6
+        with:
+          vault_uri: ${{ secrets.VAULT_URI }}
+          timeout: 1h
+          codegen_path: tests/integration_mcp/.codegen.json
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          ARM_CLIENT_ID: ${{ secrets.ARM_CLIENT_ID }}
+          ARM_TENANT_ID: ${{ secrets.ARM_TENANT_ID }}
+          # The runner job's run_as identity: a workspace SP with workspace + consumer access to the
+          # CI workspace. ci_deploy.sh requires this (bundle run_as); without it the deploy aborts.
+          DQX_MCP_RUNNER_SERVICE_PRINCIPAL_ID: ${{ secrets.TOOLS_CLIENT_ID }}
+          COVERAGE_FILE: ${{ github.workspace }}/.coverage  # make sure the coverage report is preserved
+
+      - name: Merge coverage reports and convert them to XML
+        run: make combine-coverage
+
+      - name: Publish test coverage
+        uses: codecov/codecov-action@b9fd7d16f6d7d1b5d2bec1a2887e65ceed900238 # v4.6.0
+        with:
+          use_oidc: true
+          files: coverage-combined.xml
+          flags: mcp

.github/workflows/push.yml

@@ -158,3 +158,23 @@ jobs:
       - name: Run app backend tests
         run: make app-test
         working-directory: .
+
+  mcp:
+    needs: not-a-fork
+    name: Build and Check MCP Server
+    runs-on:
+      group: databrickslabs-protected-runner-group
+      labels: linux-ubuntu-latest
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+
+      - name: Setup environment
+        uses: ./.github/actions/setup-env
+
+      - name: Type-check MCP server
+        run: make mcp-check
+
+      - name: Run MCP server tests
+        run: make mcp-test

Makefile

@@ -68,6 +68,15 @@ perf: ## Run performance benchmarks (long timeout)
 anomaly: ## Run anomaly integration tests (long timeout, with reruns)
 	$(UV_RUN) pytest tests/integration_anomaly/ -v -n 10 --timeout 1200 --durations 20 --reruns 2 --reruns-delay 5
 
+# The integration test's deploy (ci_deploy.sh -> bundle deploy) builds the artifact wheels with
+# `uv build ../` from inside mcp-server/, so the global *relative* UV_BUILD_CONSTRAINT
+# (.build-constraints.txt) would resolve against the wrong directory and fail with
+# "File not found: .build-constraints.txt". Pin it to the absolute repo-root path (same fix as
+# mcp-deploy). CI doesn't hit this — it runs via the acceptance harness, not make.
+mcp-integration: export UV_BUILD_CONSTRAINT := $(CURDIR)/.build-constraints.txt
+mcp-integration: ## Run MCP server integration tests (deploys an isolated app; requires workspace auth + Databricks CLI)
+	$(UV_RUN) pytest tests/integration_mcp/ -v --timeout 1800 --durations 10
+
 coverage: ## Run all tests (excl. e2e/perf) and open HTML coverage report
 	$(UV_TEST) --ignore=tests/e2e --ignore=tests/perf --cov --cov-report=html tests/
 	open htmlcov/index.html
@@ -177,6 +186,60 @@ app-test: ## Run app backend pytest suite (K=<expr> filter, COV=1 for coverage)
 	  $(if $(K),-k "$(K)") \
 	  $(if $(COV),--cov=src/databricks_labs_dqx_app/backend --cov-report=term-missing --cov-report=xml:coverage-app.xml)
 
+# Run the MCP server's unit-test suite (pytest, no Databricks/Spark dependencies).
+# Usage:  make mcp-test           # run everything
+#         make mcp-test K=expr    # forward -k filter to pytest
+#
+# pytest is not declared in mcp-server's pyproject/uv.lock — the suite only
+# needs anyio, which ships transitively. With UV_FROZEN=1 we can't add pytest
+# to the lock from here, so inject it ephemerally with ``uv run --with pytest``
+# over the synced runtime env (anyio's pytest plugin handles the async tests).
+mcp-test: ## Run MCP server pytest suite (K=<expr> filter)
+	cd mcp-server && uv run --with pytest pytest tests/ \
+	  $(if $(K),-k "$(K)")
+
+# Static type-check the MCP server (parity with app-check). basedpyright is not in mcp-server's
+# pyproject/uv.lock (like pytest), so inject it ephemerally with ``uv run --with basedpyright`` over
+# the synced runtime env. Scoped to ``server`` (the app package): the runner sub-project's
+# pyspark/dqx deps aren't in this env — it type-checks separately as its own wheel. ``standard`` mode
+# + ``--level error`` config lives in mcp-server/pyproject.toml [tool.basedpyright].
+mcp-check: ## Type-check the MCP server with basedpyright (standard mode, errors only)
+	cd mcp-server && uv run --with basedpyright basedpyright --level error server
+
+# One-command MCP server deploy (parity with app-deploy). Ensures the runner-wheel volume
+# exists, deploys the bundle (app + runner job + setup job), runs the one-time setup job (UC
+# grants + temp-schema ownership), then deploys & starts the app via ``bundle run``. The catalog
+# is passed as CATALOG below and injected into the app as the DQX_CATALOG config value (no secret).
+#
+# Usage: make mcp-deploy PROFILE=my-profile CATALOG=main
+# CATALOG is REQUIRED: workspace.artifact_path is the UC volume /Volumes/<CATALOG>/tmp/dqx_artifacts
+# that hosts the runner wheels; it is resolved at deploy time and must pre-exist. It is also the
+# catalog the app + runner use for temp views and per-user output schemas.
+# TARGET defaults to the bundle's default target (dev); override with TARGET=<t>.
+# The bundle artifact build runs `uv build ../` from inside mcp-server/, so the global
+# relative UV_BUILD_CONSTRAINT (.build-constraints.txt) would resolve against the wrong
+# directory. Pin it to the absolute repo-root path so the wheel build finds it.
+mcp-deploy: export UV_BUILD_CONSTRAINT := $(CURDIR)/.build-constraints.txt
+mcp-deploy: ## Deploy the MCP server bundle, run setup, and (re)deploy + start the app
+	@test -n "$(PROFILE)" || (echo "Usage: make mcp-deploy PROFILE=<databricks-profile> CATALOG=<catalog> [TARGET=<bundle-target>] [BUNDLE_VARS=...]"; exit 1)
+	@test -n "$(CATALOG)" || (echo "Usage: make mcp-deploy PROFILE=<databricks-profile> CATALOG=<catalog> [TARGET=<bundle-target>] [BUNDLE_VARS=...]"; exit 1)
+	cd mcp-server && DQX_MCP_CATALOG=$(CATALOG) DATABRICKS_PROFILE=$(PROFILE) ./scripts/ensure_artifacts_volume.sh
+	cd mcp-server && databricks bundle deploy -p $(PROFILE) $(if $(TARGET),-t $(TARGET)) --var catalog_name=$(CATALOG) $(BUNDLE_VARS)
+	cd mcp-server && databricks bundle run dqx_setup -p $(PROFILE) $(if $(TARGET),-t $(TARGET)) --var catalog_name=$(CATALOG) $(BUNDLE_VARS)
+	cd mcp-server && databricks bundle run mcp-dqx -p $(PROFILE) $(if $(TARGET),-t $(TARGET)) --var catalog_name=$(CATALOG) $(BUNDLE_VARS)
+
+# One-command teardown (nothing in the MCP bundle is destroy-protected, so — unlike the Studio's
+# multi-step unbind-then-destroy uninstall — this is a single command). Removes the app + runner/
+# setup jobs, then drops the out-of-band runner-wheel volume. Leaves the <catalog>.tmp schema and any
+# dqx_mcp_<user> output schemas intact (they may hold user data) — drop those manually if you want a
+# fully clean wipe. runner_service_principal_id isn't needed to destroy, so its placeholder is fine.
+mcp-destroy: ## Tear down the MCP server (app + jobs + runner-wheel volume). Requires PROFILE + CATALOG
+	@test -n "$(PROFILE)" || (echo "Usage: make mcp-destroy PROFILE=<databricks-profile> CATALOG=<catalog> [TARGET=<bundle-target>]"; exit 1)
+	@test -n "$(CATALOG)" || (echo "Usage: make mcp-destroy PROFILE=<databricks-profile> CATALOG=<catalog> [TARGET=<bundle-target>]"; exit 1)
+	cd mcp-server && databricks bundle destroy -p $(PROFILE) $(if $(TARGET),-t $(TARGET)) --auto-approve --var catalog_name=$(CATALOG) $(BUNDLE_VARS)
+	cd mcp-server && databricks volumes delete $(CATALOG).tmp.dqx_artifacts -p $(PROFILE) 2>/dev/null || true
+	@echo "Destroyed the MCP app + jobs and the dqx_artifacts volume. The $(CATALOG).tmp schema and any $(CATALOG).dqx_mcp_<user> output schemas are left intact — drop them manually if you want a full wipe."
+
 ##@ App deploy (require PROFILE=<databricks-profile>; most also need TARGET=<bundle-target>)
 
 # Minimum Databricks CLI version required to deploy. The ``postgres_roles``
@@ -307,6 +370,14 @@ lock-app-dependencies: ## Regenerate app/uv.lock, app/yarn.lock, app/.build-cons
 	  uv pip compile --generate-hashes --universal --no-header - > app/build-constraints-new.txt
 	mv app/build-constraints-new.txt app/.build-constraints.txt
 
+# Regenerate the MCP server lockfile and scrub private-proxy URLs so the
+# committed file resolves against whatever registry the install environment
+# is configured for (JFrog in CI, public in fork PRs).
+lock-mcp-dependencies: export UV_FROZEN := 0
+lock-mcp-dependencies: ## Regenerate mcp-server/uv.lock
+	cd mcp-server && uv lock --exclude-newer "7 days"
+	perl -pi -e 's|registry = "https://[^"]*"|registry = "https://pypi.org/simple"|g' mcp-server/uv.lock
+
 lock-dependencies: export UV_FROZEN := 0
 lock-dependencies: ## Regenerate top-level uv.lock and .build-constraints.txt
 	uv lock --exclude-newer "7 days"