# BigQuery RAG MCP Server

[![Python](https://img.shields.io/badge/Python-3.12-3776AB?logo=python&logoColor=white)](https://www.python.org/)
[![BigQuery](https://img.shields.io/badge/BigQuery-Vector%20Search-669DF6?logo=googlebigquery&logoColor=white)](https://cloud.google.com/bigquery)
[![Cloud Run](https://img.shields.io/badge/Cloud%20Run-Deployed-4285F4?logo=googlecloud&logoColor=white)](https://cloud.google.com/run)
[![MCP](https://img.shields.io/badge/MCP-Streamable%20HTTP-111111)](https://modelcontextprotocol.io/)

A private Model Context Protocol (MCP) service that converts a natural-language question into an embedding, performs semantic retrieval over document chunks stored in BigQuery, and returns structured passages with source and page metadata.

This repository owns the **retrieval layer** of a larger document-grounded chatbot. The companion application repository owns Google ADK orchestration, Gemini answer generation, Firebase Authentication, the `/chat` API, and the React interface.

> Companion application: [shrprabh/atomic-habits-adk-rag](https://github.com/shrprabh/atomic-habits-adk-rag)

## Live deployment

| Resource | Value |
| --- | --- |
| Cloud Run service | `bigquery-rag-mcp` |
| Region | `us-central1` |
| Base URL | `https://bigquery-rag-mcp-nfp4nl2vna-uc.a.run.app` |
| MCP endpoint | `/mcp` |
| Health endpoint | `/health` |
| Access | Private; Cloud Run IAM authentication required |

The service URL is intentionally not browser-public. A caller must have `roles/run.invoker` on the service and send a Google-signed identity token whose audience is the MCP base URL.

## End-to-end architecture
![Secure BigQuery RAG and Google ADK Architecture](assets/secure-bigquery-rag-adk.png)
```text
React application on Firebase Hosting
        │ Firebase ID token
        ▼
ADK Agent API on Cloud Run
        │ Google service identity token
        ▼
Private MCP service on Cloud Run  ◀── this repository
        │ parameterized BigQuery SQL
        ▼
AI.GENERATE_EMBEDDING
        │ 1,536-dimensional query vector
        ▼
BigQuery VECTOR_SEARCH (COSINE)
        │
        ▼
Top document passages + page metadata
```

## What this service does

- Exposes a read-only MCP tool named `semantic_search`.
- Validates `query` and `top_k` using Pydantic-generated MCP schemas.
- Creates a query embedding with `AI.GENERATE_EMBEDDING` using `RETRIEVAL_QUERY`.
- Performs cosine-distance `VECTOR_SEARCH` over stored document embeddings.
- Uses a parameterized query value instead of inserting user input into SQL.
- Returns structured source, page, chapter, section, distance, and similarity fields.
- Runs as a stateless Streamable HTTP MCP server.
- Keeps the retrieval service private with Cloud Run IAM.
- Does not call Gemini to compose an answer; generation belongs to the companion ADK service.

## BigQuery resources used by this project

| Setting | Value |
| --- | --- |
| Google Cloud project | `bigquery-semantic-search` |
| BigQuery location | `us-central1` |
| Dataset | `atomic_habits_rag` |
| Cloud resource connection | `vertex_ai_connection` |
| Remote embedding model | `atomic_habits_rag.embedding_model` |
| Embedding table | `atomic_habits_rag.article_embeddings` |
| Current rows | 1,222 |
| Embedding dimension | 1,536 |
| Distance type | Cosine |
| Search mode | Exact brute-force search |

The current table is small, so this implementation deliberately uses brute-force vector search. A vector index becomes useful after the corpus grows enough to justify approximate nearest-neighbor search and index maintenance.

## MCP tool contract

### `semantic_search`

Input:

```json
{
  "query": "What is the two-minute rule?",
  "top_k": 5
}
```

Validation:

| Field | Rules |
| --- | --- |
| `query` | String, 2–500 characters |
| `top_k` | Integer, 1–10; default `5` |

Simplified output:

```json
{
  "query": "What is the two-minute rule?",
  "result_count": 5,
  "results": [
    {
      "chunk_id": 480,
      "document_id": "atomic_habits",
      "content": "Retrieved passage text...",
      "title": "Atomic Habits",
      "author": "James Clear",
      "source": "atomic-habits.pdf",
      "page_start": 96,
      "page_end": 96,
      "chapter": "...",
      "section": "...",
      "distance": 0.18,
      "similarity": 0.82
    }
  ]
}
```

## Repository layout

```text
bigquery-rag-mcp/
├── server.py                 # MCP tool, BigQuery query, health route
├── test_mcp.py               # In-process MCP regression test
├── test_deployed_mcp.py      # Authenticated test against Cloud Run
├── rag_client.py             # Local in-process RAG reference client
├── requirements.txt
├── Dockerfile
├── .env.example
└── .gitignore
```

`rag_client.py` imports `mcp` from `server.py`, so it executes the tool in the same Python process. It is useful as a local reference or regression client, but it is **not part of the deployed production request path**. The companion ADK application calls this service remotely over `/mcp`.

## Prerequisites

- Python 3.12+
- Google Cloud CLI
- A Google Cloud project with billing enabled
- BigQuery, BigQuery Connection, Vertex AI, Cloud Run, Cloud Build, and Artifact Registry APIs
- Existing BigQuery dataset, embedding model, and embedding table matching the configured schema
- Permission to create service accounts and manage Cloud Run and BigQuery IAM

## 1. Clone and install

```bash
git clone https://github.com/shrprabh/bigquery-rag-mcp.git
cd bigquery-rag-mcp

python3 -m venv .venv
source .venv/bin/activate

python -m pip install --upgrade pip
pip install -r requirements.txt
```

For local development outside Cloud Shell:

```bash
gcloud auth login
gcloud auth application-default login
gcloud config set project bigquery-semantic-search
```

Never commit Application Default Credentials or service-account key files.

## 2. Configure the environment

```bash
cp .env.example .env
```

Expected values:

```dotenv
GOOGLE_CLOUD_PROJECT=bigquery-semantic-search
BQ_DATASET=atomic_habits_rag
BQ_LOCATION=us-central1
EMBEDDING_DIM=1536
```

`server.py` reads these values from the process environment. The checked-in `.env.example` is documentation only; use explicit exports locally or Cloud Run environment variables in deployment.

## 3. Verify the BigQuery assets

Run in the BigQuery editor:

```sql
SELECT
  ARRAY_LENGTH(embedding) AS dimensions,
  COUNT(*) AS row_count
FROM `bigquery-semantic-search.atomic_habits_rag.article_embeddings`
GROUP BY dimensions;
```

Expected for the current dataset:

```text
dimensions  row_count
1536        1222
```

Confirm that the model exists:

```sql
SELECT
  model_name,
  model_type
FROM `bigquery-semantic-search.atomic_habits_rag.INFORMATION_SCHEMA.MODELS`
WHERE model_name = 'embedding_model';
```

## 4. Configure runtime IAM

Set variables:

```bash
export PROJECT_ID="bigquery-semantic-search"
export REGION="us-central1"
export CONNECTION_ID="vertex_ai_connection"
export MCP_SERVICE="bigquery-rag-mcp"
export MCP_SA_NAME="bigquery-rag-mcp-sa"
export MCP_SA="${MCP_SA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com"

gcloud config set project "$PROJECT_ID"
```

Enable APIs:

```bash
gcloud services enable \
  bigquery.googleapis.com \
  bigqueryconnection.googleapis.com \
  aiplatform.googleapis.com \
  run.googleapis.com \
  cloudbuild.googleapis.com \
  artifactregistry.googleapis.com \
  --project="$PROJECT_ID"
```

Create the runtime service account if it does not already exist:

```bash
gcloud iam service-accounts describe "$MCP_SA" \
  --project="$PROJECT_ID" >/dev/null 2>&1 || \
gcloud iam service-accounts create "$MCP_SA_NAME" \
  --project="$PROJECT_ID" \
  --display-name="BigQuery RAG MCP Server"
```

Grant the service account permission to run queries and read the dataset/model:

```bash
gcloud projects add-iam-policy-binding "$PROJECT_ID" \
  --member="serviceAccount:${MCP_SA}" \
  --role="roles/bigquery.jobUser"

gcloud projects add-iam-policy-binding "$PROJECT_ID" \
  --member="serviceAccount:${MCP_SA}" \
  --role="roles/bigquery.dataViewer"
```

### Required connection permission

Because `AI.GENERATE_EMBEDDING` uses the BigQuery Cloud resource connection, the MCP runtime identity must also be allowed to use `vertex_ai_connection`.

In the Google Cloud console:

1. Open **BigQuery → your project → Connections**.
2. Select `vertex_ai_connection` in `us-central1`.
3. Select **Share**.
4. Add `bigquery-rag-mcp-sa@bigquery-semantic-search.iam.gserviceaccount.com`.
5. Grant **BigQuery Connection User** (`roles/bigquery.connectionUser`).

Do not use `bq add-iam-policy-binding --connection_type=...`; that flag does not share a connection and is rejected by current `bq` versions. The Cloud console or BigQuery Connections API should be used for connection-level sharing.

The connection itself has a Google-managed service account. That connection service account must have the appropriate Vertex AI/Agent Platform user role in the project so the remote embedding model can call its endpoint.

Without these connection permissions, the MCP log contains an error similar to:

```text
403 Access Denied: User does not have bigquery.connections.use permission
```

## 5. Test locally

Compile the files:

```bash
python -m py_compile server.py test_mcp.py rag_client.py
```

Run the direct MCP regression test:

```bash
python test_mcp.py
```

Start the HTTP server:

```bash
python server.py
```

Endpoints:

```text
http://localhost:8000/health
http://localhost:8000/mcp
```

From another terminal:

```bash
curl http://localhost:8000/health
```

Expected:

```json
{"status":"healthy"}
```

Optionally run the local grounded-generation reference client:

```bash
python rag_client.py
```

## 6. Deploy the private MCP service

```bash
gcloud run deploy "$MCP_SERVICE" \
  --source=. \
  --project="$PROJECT_ID" \
  --region="$REGION" \
  --service-account="$MCP_SA" \
  --no-allow-unauthenticated \
  --memory="1Gi" \
  --timeout="300" \
  --set-env-vars="GOOGLE_CLOUD_PROJECT=$PROJECT_ID,BQ_DATASET=atomic_habits_rag,BQ_LOCATION=$REGION,EMBEDDING_DIM=1536"
```

Cloud Run supplies `PORT`; `server.py` binds to `0.0.0.0` and uses that port.

Get the canonical service URL:

```bash
export MCP_URL="$(
  gcloud run services describe "$MCP_SERVICE" \
    --project="$PROJECT_ID" \
    --region="$REGION" \
    --format='value(status.url)'
)"

echo "$MCP_URL"
```

Test the authenticated health route:

```bash
curl -i \
  -H "Authorization: Bearer $(gcloud auth print-identity-token)" \
  "$MCP_URL/health"
```

Expected: HTTP `200` and `{"status":"healthy"}`.

## 7. Test the deployed MCP tool

```bash
export MCP_URL="https://bigquery-rag-mcp-nfp4nl2vna-uc.a.run.app"
python test_deployed_mcp.py
```

Ask:

```text
What is the two-minute rule?
```

The test client should initialize an MCP session, call `semantic_search`, and print structured retrieval results. A successful HTTP status alone is insufficient; verify that `result_count` is greater than zero and the result contains page metadata.

## 8. Authorize the companion ADK service

After creating the agent service account in the companion repository, allow it to invoke this private service:

```bash
export AGENT_SA="bigquery-rag-agent-sa@bigquery-semantic-search.iam.gserviceaccount.com"

gcloud run services add-iam-policy-binding "$MCP_SERVICE" \
  --project="$PROJECT_ID" \
  --region="$REGION" \
  --member="serviceAccount:${AGENT_SA}" \
  --role="roles/run.invoker"
```

Verify:

```bash
gcloud run services get-iam-policy "$MCP_SERVICE" \
  --project="$PROJECT_ID" \
  --region="$REGION" \
  --flatten="bindings[].members" \
  --filter="bindings.members:serviceAccount:${AGENT_SA}" \
  --format="table(bindings.role,bindings.members)"
```

The ADK service account needs `run.invoker` on this service. It does not need the MCP service's BigQuery roles because each Cloud Run service has its own identity and responsibility.

Continue with [the ADK + React deployment guide](https://github.com/shrprabh/atomic-habits-adk-rag#readme).

## Observability

Read recent logs:

```bash
gcloud run services logs read "$MCP_SERVICE" \
  --project="$PROJECT_ID" \
  --region="$REGION" \
  --limit=100
```

Useful successful log message:

```text
Running semantic search with top_k=5
```

Cloud Run metrics are available under:

```text
Google Cloud Console → Cloud Run → bigquery-rag-mcp → Metrics
```

BigQuery query history and bytes processed are available in BigQuery job history or `INFORMATION_SCHEMA.JOBS_BY_PROJECT`.

## Troubleshooting

| Symptom | Cause | Resolution |
| --- | --- | --- |
| `/health` returns `403` | Private Cloud Run request has no valid identity token | Send an ID token and ensure the caller has `roles/run.invoker` |
| Tool result says semantic search could not complete | Inspect MCP logs for the underlying BigQuery exception | Run the logs command above |
| `bigquery.connections.use` denied | MCP runtime SA cannot use `vertex_ai_connection` | Share the connection with the runtime SA as BigQuery Connection User |
| Vertex/remote-model permission denied | Connection-managed SA cannot invoke the embedding endpoint | Grant the documented Vertex AI/Agent Platform user role to the connection SA |
| `streamable_http_client()` rejects `headers` or `auth` | Client code does not match the installed MCP SDK version | Use the committed `test_deployed_mcp.py` and keep `mcp` dependency versions aligned |
| `structured_content` is `null` | The SDK may expose an error through content blocks | Inspect the complete tool result rather than only `structured_content` |
| Origin/DNS-rebinding error behind Cloud Run | Transport security treats proxy host headers as untrusted | The server disables DNS-rebinding protection only when `K_SERVICE` confirms Cloud Run |
| No rows returned | Model/table location, dimension, or query status mismatch | Verify the model, table, connection, location, and 1,536 dimension |

## Security and data handling

- The MCP Cloud Run service remains private.
- No service-account JSON keys are deployed or committed.
- Cloud Run service identity and short-lived Google ID tokens are used.
- User query text is passed to BigQuery as a parameter.
- The tool is marked read-only and returns retrieval evidence only.
- `.env`, ADC files, PDFs, JSONL chunks, logs, and local databases are ignored by Git.
- The source document and extracted chunks are not redistributed in this repository.
- Do not expose authentication tokens in screenshots or logs.

## Current limitations

- The corpus contains 1,222 chunks from one document.
- Search is brute force and has no vector index.
- There is no reranker or retrieval evaluation suite yet.
- The MCP tool returns passages; answer quality and citations depend on the companion agent.
- The current public portfolio implementation is document-specific rather than a multi-tenant ingestion platform.

## Recommended next improvements

1. Add retrieval evaluation with a question/expected-source dataset.
2. Add similarity thresholds and abstention tests.
3. Support document ingestion and metadata validation as a separate pipeline.
4. Add tenant/document filters before retrieval.
5. Add a vector index after the dataset becomes large enough.
6. Add structured Cloud Logging fields for latency and result count without logging passage contents.
7. Add unit tests that mock BigQuery and integration tests for the deployed MCP service.

## GitHub publication

```bash
git add README.md
git commit -m "Add end-to-end MCP deployment documentation"

git remote add origin https://github.com/shrprabh/bigquery-rag-mcp.git
git push -u origin main
```

If `origin` already exists, do not add it again. Check with `git remote -v`, then run only `git push`.

## Official references

- [MCP Python SDK](https://github.com/modelcontextprotocol/python-sdk)
- [BigQuery `AI.GENERATE_EMBEDDING`](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-ai-generate-embedding)
- [BigQuery vector search](https://cloud.google.com/bigquery/docs/vector-search)
- [BigQuery connection access](https://cloud.google.com/bigquery/docs/working-with-connections)
- [Cloud Run service-to-service authentication](https://cloud.google.com/run/docs/authenticating/service-to-service)
- [Cloud Run service identity](https://cloud.google.com/run/docs/configuring/services/service-identity)

## Author

**Shreyas Prabhakar**

- GitHub: [@shrprabh](https://github.com/shrprabh)
- LinkedIn: [linkedin.com/in/shreyasprabhakar](https://www.linkedin.com/in/shreyasprabhakar)
- Medium: [@pshreyasgowda1997](https://medium.com/@pshreyasgowda1997)