Citus MCP

🚀 Citus MCP Server

An AI-powered MCP server for managing Citus distributed PostgreSQL clusters

This server enables AI assistants like GitHub Copilot to interact with your Citus distributed PostgreSQL cluster, offering features such as read - only inspection, intelligent advisors, guarded operations, and real - time monitoring.

Quick Start • Features • Installation • Configuration • Tools Reference • Examples

🚀 Quick Start

Prerequisites

• Go 1.23+ (for building from source)
• Citus 12.x–14.x cluster with coordinator access
• GitHub Copilot with MCP support (VS Code or CLI)

1. Build the Server

git clone https://github.com/citusdata/citus-mcp.git
cd citus-mcp
make build
# Binary created at ./bin/citus-mcp

Or using Go directly:

go build -o bin/citus-mcp ./cmd/citus-mcp

2. Configure Your Connection

Create a configuration file at ~/.config/citus-mcp/config.yaml:

# Minimum required configuration
coordinator_dsn: postgres://username:password@localhost:5432/mydb?sslmode=disable

Or set the environment variable:

export CITUS_MCP_COORDINATOR_DSN="postgres://username:password@localhost:5432/mydb?sslmode=disable"

3. Set Up VS Code

Create .vscode/mcp.json in your workspace (or mcp.json at the project root):

{
  "mcpServers": {
    "citus-mcp": {
      "command": "/path/to/citus-mcp/bin/citus-mcp",
      "args": [],
      "env": {
        "CITUS_MCP_COORDINATOR_DSN": "postgres://username:password@localhost:5432/mydb?sslmode=disable"
      }
    }
  }
}

4. Test the Connection

In VS Code Copilot Chat, type:

@citus-mcp ping

You should see a "pong" response confirming the connection works.

✨ Features

🔍 Cluster Inspection (Read-Only)

📊 Monitoring & Analysis

🤖 Intelligent Advisors

⚡ Execute Operations (Requires Approval)

📦 Installation

Option 1: Build from Source

# Clone the repository
git clone https://github.com/citusdata/citus-mcp.git
cd citus-mcp

# Build using Make
make build

# Or build directly with Go
go build -o bin/citus-mcp ./cmd/citus-mcp

# (Optional) Install to your PATH
sudo cp bin/citus-mcp /usr/local/bin/

Option 2: Go Install

go install github.com/citusdata/citus-mcp/cmd/citus-mcp@latest

Verify Installation

citus-mcp --help

⚙️ Configuration

Connection String (DSN)

The most important configuration is the PostgreSQL connection string to your Citus coordinator:

postgres://[user]:[password]@[host]:[port]/[database]?sslmode=[mode]

Examples:

# Local development (no SSL)
postgres://postgres:secret@localhost:5432/mydb?sslmode=disable

# Production with SSL
postgres://admin:secret@citus-coord.example.com:5432/production?sslmode=require

# With specific schema
postgres://user:pass@host:5432/db?sslmode=require&search_path=myschema

Configuration Methods

Configuration can be provided via (in order of precedence):

1. Command-line flags
2. Environment variables
3. Configuration file

Method 1: Environment Variables

# Required
export CITUS_MCP_COORDINATOR_DSN="postgres://user:pass@localhost:5432/mydb?sslmode=disable"

# Optional
export CITUS_MCP_MODE="read_only"           # read_only (default) or admin
export CITUS_MCP_ALLOW_EXECUTE="false"      # Enable execute operations
export CITUS_MCP_APPROVAL_SECRET="secret"   # Required if allow_execute=true
export CITUS_MCP_LOG_LEVEL="info"           # debug, info, warn, error

Method 2: Configuration File

Create ~/.config/citus-mcp/config.yaml:

# ===========================================
# Citus MCP Server Configuration
# ===========================================

# Database Connection (REQUIRED)
# -----------------------------
coordinator_dsn: postgres://user:password@localhost:5432/mydb?sslmode=disable

# Optional: Override credentials from DSN
# coordinator_user: myuser
# coordinator_password: mypassword

# Optional: Direct worker connections (comma-separated)
# worker_dsns: postgres://user:pass@worker1:5432/db,postgres://user:pass@worker2:5432/db

# Server Mode
# -----------
# read_only: Only inspection tools available (default, safest)
# admin: All tools available including execute operations
mode: read_only

# Execute Operations (only if mode=admin)
# ---------------------------------------
allow_execute: false
# approval_secret: your-secret-key  # Required if allow_execute=true

# Performance Settings
# --------------------
cache_ttl_seconds: 5          # Cache duration for metadata queries
enable_caching: true          # Set to false to disable caching
max_rows: 200                 # Maximum rows returned per query
max_text_bytes: 200000        # Maximum text size in responses

# Timeouts
# --------
connect_timeout_seconds: 10   # Connection timeout
statement_timeout_ms: 30000   # Query timeout (30 seconds)

# Logging
# -------
log_level: info               # debug, info, warn, error

# Transport (NEW)
# ---------------
# stdio: Standard input/output (default, for VS Code/CLI integration)
# sse: Server-Sent Events over HTTP (for remote/network access)
# streamable: Streamable HTTP transport (for remote/network access)
transport: stdio

# HTTP Settings (only used when transport is sse or streamable)
# http_addr: 127.0.0.1        # Listen address (use 0.0.0.0 for all interfaces)
# http_port: 8080             # Listen port
# http_path: /mcp             # Endpoint path
# sse_keepalive_seconds: 30   # SSE keepalive interval

Method 3: Command-Line Flags

# Using flags (note: use underscores in flag names)
bin/citus-mcp --coordinator_dsn "postgres://..." --mode read_only

# Using positional argument for DSN
bin/citus-mcp "postgres://user:pass@localhost:5432/mydb?sslmode=disable"

# Specify config file
bin/citus-mcp --config /path/to/config.yaml

# Start with SSE transport
bin/citus-mcp --transport sse --http_port 8080 --coordinator_dsn "postgres://..."

Configuration File Locations

The server searches for configuration files in this order:

1. --config / -c flag
2. CITUS_MCP_CONFIG environment variable
3. $XDG_CONFIG_HOME/citus-mcp/config.yaml
4. ~/.config/citus-mcp/config.yaml
5. ./citus-mcp.yaml (current directory)

Supported formats: YAML, JSON, TOML

🌐 Transport Options

Citus MCP supports three transport modes for different deployment scenarios:

1. Stdio Transport (Default)

Standard input/output transport — the server communicates via stdin/stdout. This is the default and is used for direct integration with VS Code and GitHub Copilot CLI.

# Default - stdio transport
bin/citus-mcp --coordinator_dsn "postgres://..."

# Explicit
bin/citus-mcp --transport stdio --coordinator_dsn "postgres://..."

Use cases:

• VS Code Copilot Chat integration
• GitHub Copilot CLI
• Local development

2. SSE Transport (Server-Sent Events)

HTTP-based transport using Server-Sent Events. The server runs as an HTTP daemon that clients can connect to remotely.

# Start server on HTTP with SSE
bin/citus-mcp --transport sse --http_addr 0.0.0.0 --http_port 8080 --coordinator_dsn "postgres://..."

# Or via environment variables
export CITUS_MCP_TRANSPORT=sse
export CITUS_MCP_HTTP_ADDR=0.0.0.0
export CITUS_MCP_HTTP_PORT=8080
export CITUS_MCP_COORDINATOR_DSN="postgres://..."
bin/citus-mcp

Endpoints:

• GET /mcp - Establish SSE connection
• POST /mcp/session/{id} - Send messages to session
• GET /health - Health check

Use cases:

• Remote MCP server deployment
• Docker/Kubernetes deployments
• Shared server for multiple clients
• Network-accessible MCP services

3. Streamable HTTP Transport

Modern HTTP transport with streaming support. Recommended for new deployments.

# Start server with streamable HTTP transport
bin/citus-mcp --transport streamable --http_addr 0.0.0.0 --http_port 8080 --coordinator_dsn "postgres://..."

Endpoints:

• POST /mcp - Handle MCP requests with streaming responses
• GET /health - Health check

Use cases:

• Same as SSE, with better streaming support
• Environments where SSE is not ideal

Docker Deployment Example

FROM golang:1.22-alpine AS builder
WORKDIR /app
COPY . .
RUN go build -o citus-mcp ./cmd/citus-mcp

FROM alpine:latest
COPY --from=builder /app/citus-mcp /usr/local/bin/
EXPOSE 8080
CMD ["citus-mcp", "--transport", "sse", "--http-addr", "0.0.0.0", "--http-port", "8080"]

# docker-compose.yml
version: '3.8'
services:
  citus-mcp:
    build: .
    ports:
      - "8080:8080"
    environment:
      CITUS_MCP_TRANSPORT: sse
      CITUS_MCP_HTTP_ADDR: 0.0.0.0
      CITUS_MCP_HTTP_PORT: 8080
      CITUS_MCP_COORDINATOR_DSN: postgres://user:pass@citus-coordinator:5432/mydb?sslmode=disable

Connecting to Remote Server

For SSE/Streamable transports, configure your MCP client to connect via HTTP:

{
  "mcpServers": {
    "citus-mcp": {
      "type": "sse",
      "url": "http://citus-mcp-server:8080/mcp"
    }
  }
}

🔌 Setting Up with GitHub Copilot

VS Code Setup

1. Install Prerequisites
   • VS Code with GitHub Copilot extension
   • MCP support enabled in Copilot settings
2. Create MCP Configuration Create .vscode/mcp.json in your workspace:

   {
     "mcpServers": {
       "citus-mcp": {
         "command": "/absolute/path/to/bin/citus-mcp",
         "args": [],
         "env": {
           "CITUS_MCP_COORDINATOR_DSN": "postgres://user:pass@localhost:5432/mydb?sslmode=disable"
         }
       }
     }
   }
   

   Or for development (using go run):

   {
     "mcpServers": {
       "citus-mcp": {
         "command": "go",
         "args": ["run", "./cmd/citus-mcp"],
         "cwd": "/path/to/citus-mcp",
         "env": {
           "CITUS_MCP_COORDINATOR_DSN": "postgres://user:pass@localhost:5432/mydb?sslmode=disable"
         }
       }
     }
   }
   

3. Reload VS Code and open Copilot Chat
4. Verify Connection

   @citus-mcp ping
   

GitHub Copilot CLI Setup

1. Create Global MCP Config Create ~/.config/github-copilot/mcp.json:

   {
     "mcpServers": {
       "citus-mcp": {
         "command": "/usr/local/bin/citus-mcp",
         "args": [],
         "env": {
           "CITUS_MCP_COORDINATOR_DSN": "postgres://user:pass@localhost:5432/mydb?sslmode=disable"
         }
       }
     }
   }
   

2. Verify Setup

   copilot mcp list
   copilot mcp test citus-mcp
   

3. Use in CLI

   copilot -p "Show me the cluster summary"
   

💻 Usage Examples

Basic Cluster Inspection

@citus-mcp Show me the cluster summary

@citus-mcp List all distributed tables

@citus-mcp Inspect the public.users table including shards and indexes

Monitoring

@citus-mcp Show current cluster activity

@citus-mcp Are there any lock waits in the cluster?

@citus-mcp Show background job progress

Analysis

@citus-mcp Analyze shard skew for the orders table

@citus-mcp Show me the shard heatmap grouped by node

@citus-mcp Explain this query: SELECT * FROM orders WHERE customer_id = 123

Advisor

@citus-mcp Run the advisor with focus on skew

@citus-mcp Check operational health (long queries, locks, jobs)

@citus-mcp Suggest the best source node for snapshot-based scaling

@citus-mcp Check metadata health with deep validation across nodes

Configuration Analysis

@citus-mcp Analyze cluster configuration and recommend improvements

@citus-mcp Run config advisor with focus on memory settings

Colocation Analysis

@citus-mcp Show all colocation groups

@citus-mcp Which tables are colocated with the orders table?

Node Addition

@citus-mcp Run pre-flight checks for adding node at postgres://user:pass@newworker:5432/db

📚 Tools Reference

Inspection Tools

Monitoring Tools

Advisor Tools

Execute Tools (Require Approval)

📋 Built-in Prompts

Use these prompts in Copilot Chat for guided workflows:

🔐 Security

Read-Only Mode (Default)

By default, citus-mcp runs in read-only mode. This means:

• ✅ All inspection and monitoring tools work
• ✅ Advisors provide recommendations
• ❌ Execute operations are disabled
• ❌ No data can be modified

Admin Mode with Approval Tokens

To enable execute operations:

1. Set admin mode in configuration:

   mode: admin
   allow_execute: true
   approval_secret: your-secret-key-here
   

2. Request an approval token before executing:

   @citus-mcp Request approval token for rebalance
   

3. Use the token in the execute command:

   @citus-mcp Execute rebalance with token: <token>
   

Tokens are time-limited and action-specific (HMAC-signed).

🔧 Troubleshooting

Connection Issues

Error: connection refused

• Verify the coordinator host and port are correct
• Check that PostgreSQL is running and accepting connections
• Ensure firewall rules allow the connection

Error: authentication failed

• Verify username and password in DSN
• Check that the user has permissions on the database
• For SSL issues, try sslmode=disable for local testing

MCP Issues

Copilot doesn't see citus-mcp

• Ensure mcp.json is in the correct location
• Check that the command path is absolute
• Reload VS Code after changing configuration

Tools return errors

• Check logs: CITUS_MCP_LOG_LEVEL=debug bin/citus-mcp
• Verify Citus extension is installed: SELECT * FROM pg_extension WHERE extname = 'citus'

Testing Connection

# Test directly
CITUS_MCP_COORDINATOR_DSN="postgres://..." bin/citus-mcp

# Then send a ping via stdin
echo '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}' | bin/citus-mcp

🛠️ Development

Running Tests

# Unit tests
make test

# With verbose output
go test -v ./...

# Integration tests (requires Docker)
make docker-up
make integration
make docker-down

Linting

make lint

Project Structure

citus-mcp/
├── cmd/citus-mcp/       # Main entry point
├── internal/
│   ├── mcpserver/       # MCP server implementation
│   │   ├── tools/       # Tool implementations (30+ tools)
│   │   ├── prompts/     # Prompt templates
│   │   └── resources/   # Static resources
│   ├── db/              # Database layer and worker management
│   ├── citus/           # Citus-specific logic and queries
│   │   ├── advisor/     # Advisor implementations
│   │   └── guc/         # GUC (configuration) analysis
│   ├── cache/           # Query result caching
│   ├── config/          # Configuration management
│   ├── errors/          # Error types and codes
│   ├── fanout/          # Parallel query execution
│   ├── logging/         # Structured logging
│   └── safety/          # Guardrails and approval tokens
├── docker/              # Docker Compose setup for testing
├── docs/                # Additional documentation
└── tests/               # Integration tests

🤝 Contributing

Contributions are welcome! Please see CONTRIBUTING.md for guidelines.

📄 License

MIT License — see LICENSE for details.

⬆ Back to Top

Made with ❤️ for the Citus community