import sys
import click
from src.common.config import parse_crdb_uri, set_crdb_config_from_cli
from src.common.server import mcp
import src.tools.cluster_monitoring
import src.tools.database_operations
import src.tools.table_management
import src.tools.query_engine
class CockroachMCPServer:
def __init__(self):
print("Starting the CockroachDB MCP Server", file=sys.stderr)
def run(self, transport: str, http_host: str | None, http_port: int | None,
http_path: str | None, stateless_http: bool):
if not mcp:
return
if transport == "http":
import uvicorn
# Get the Starlette app for streamable HTTP
app = mcp.streamable_http_app()
print(f"Starting HTTP server on {http_host or '0.0.0.0'}:{http_port or 8000}", file=sys.stderr)
uvicorn.run(
app,
host=http_host or "0.0.0.0",
port=http_port or 8000,
)
else:
mcp.run()
@click.command()
@click.option('--url', help='CockroachDB connection URI (cockroach://<username>:<password>@<host>:<port>/<database> or postgresql://<username>:<password>@<host>:<port>/<database>)')
@click.option('--host', help='CockroachDB host')
@click.option('--port', default=26257, type=int, help='CockroachDB port')
@click.option('--db', default='defaultdb', help='Cockroach database name')
@click.option('--username', default='root', help='Username')
@click.option('--password', help='Password')
@click.option('--ssl-mode', default='disable', help='SSL mode for CockroachDB connection. Possible values: disable, allow, prefer, require, verify-ca or verify-full. Default is disable.')
@click.option('--ssl-key', help='Path to SSL Client key file')
@click.option('--ssl-cert', help='Path to SSL Client certificate file')
@click.option('--ssl-ca-cert', help='Path to CA (Root) certificate file')
@click.option('--transport', type=click.Choice(['stdio', 'http']), default='stdio', show_default=True, help='MCP transport to use.')
@click.option('--http-host', help='HTTP host to bind for streamable HTTP transport.')
@click.option('--http-port', type=int, help='HTTP port to bind for streamable HTTP transport.')
@click.option('--http-path', help='HTTP path for streamable HTTP transport (e.g., /mcp).')
@click.option('--stateless-http/--stateful-http', default=False, show_default=True, help='Enable stateless HTTP mode for horizontal scaling.')
@click.option('--use-env/--no-use-env', default=False, show_default=True, help='Use environment variables for CockroachDB configuration.')
def cli(url, host, port, db, username, password,
ssl_mode, ssl_key, ssl_cert, ssl_ca_cert,
transport, http_host, http_port, http_path, stateless_http, use_env):
"""CockroachDB MCP Server - Model Context Protocol server for CockroachDB."""
# Handle CockroachDB URI if provided
if url:
try:
uri_config = parse_crdb_uri(url)
set_crdb_config_from_cli(uri_config)
except ValueError as e:
click.echo(f"Error parsing CockroachDB URI: {e}", err=True)
sys.exit(1)
elif host:
# Set individual parameters
cfg = {
'host': host,
'port': port,
'username': username,
'database': db,
}
if password:
cfg['password'] = password
if ssl_mode:
cfg['ssl_mode'] = ssl_mode
if ssl_key:
cfg['ssl_key'] = ssl_key
if ssl_cert:
cfg['ssl_cert'] = ssl_cert
if ssl_ca_cert:
cfg['ssl_ca_cert'] = ssl_ca_cert
set_crdb_config_from_cli(cfg)
else:
if not use_env:
print(f"You are in CLI mode. You must fill in at least one of the two parameters --url or --host to launch the MCP server", file=sys.stderr)
sys.exit(1)
# Start the server
server = CockroachMCPServer()
server.run(transport, http_host, http_port, http_path, stateless_http)
def main():
"""Legacy main function for backward compatibility."""
server = CockroachMCPServer()
server.run("stdio", None, None, None, False)
if __name__ == "__main__":
main()
