406 lines
18 KiB
JavaScript
406 lines
18 KiB
JavaScript
#!/usr/bin/env node
|
|
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
|
|
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
|
|
import { createRequire } from "node:module";
|
|
import { z } from "zod";
|
|
import {
|
|
buildSchemaContext,
|
|
createBackend,
|
|
evaluateRedisCommandSafety,
|
|
evaluateMongoAggregateSafety,
|
|
evaluateSqlSafety,
|
|
formatCell,
|
|
formatSchemaContext,
|
|
isMainModule,
|
|
mdTable,
|
|
notifyReload,
|
|
parseMongoAggregateCommand,
|
|
postBridge,
|
|
sqlSafetyFromEnv,
|
|
splitSqlStatements,
|
|
type Backend,
|
|
type ConnectionConfig,
|
|
type QueryResult,
|
|
type RedisCommandResult,
|
|
} from "@dbx-app/node-core";
|
|
|
|
const require = createRequire(import.meta.url);
|
|
const packageJson = require("../package.json") as { version?: string };
|
|
export const DBX_MCP_PACKAGE_VERSION = packageJson.version ?? "0.0.0";
|
|
|
|
function text(s: string) {
|
|
return { content: [{ type: "text" as const, text: s }] };
|
|
}
|
|
|
|
function toolError(code: string, message: string) {
|
|
return { ...text(`${code}: ${message}`), isError: true };
|
|
}
|
|
|
|
function withDatabase(config: ConnectionConfig, database?: string): ConnectionConfig {
|
|
return database === undefined ? config : { ...config, database };
|
|
}
|
|
|
|
function formatQueryToolResult(result: QueryResult, title?: string) {
|
|
const prefix = title ? `${title}\n` : "";
|
|
if (result.columns.length === 0) return text(`${prefix}Query executed. ${result.row_count} row(s) affected.`);
|
|
const rows = result.rows.map((r) => result.columns.map((c) => formatCell(r[c])));
|
|
return text(`${prefix}${mdTable(result.columns, rows)}\n\n${result.row_count} row(s)`);
|
|
}
|
|
|
|
function redisDbFromValue(value?: string): number | undefined {
|
|
const trimmed = value?.trim();
|
|
if (!trimmed) return undefined;
|
|
const db = Number(trimmed);
|
|
return Number.isInteger(db) && db >= 0 ? db : undefined;
|
|
}
|
|
|
|
function defaultRedisDb(config: ConnectionConfig, scope: McpScope, db?: number): number {
|
|
return db ?? redisDbFromValue(scope.database) ?? redisDbFromValue(config.database) ?? 0;
|
|
}
|
|
|
|
function formatRedisCommandValue(value: unknown): string {
|
|
if (typeof value === "string") return value;
|
|
return JSON.stringify(value, null, 2) ?? String(value);
|
|
}
|
|
|
|
function formatRedisCommandToolResult(result: RedisCommandResult) {
|
|
return text(`Command: ${result.command}\nSafety: ${result.safety}\n\n${formatRedisCommandValue(result.value)}`);
|
|
}
|
|
|
|
export const DBX_CONNECTION_TYPE_DESCRIPTION =
|
|
"Database type: postgres, mysql, sqlite, rqlite, redis, duckdb, clickhouse, sqlserver, mongodb, oracle, elasticsearch, etcd, doris, starrocks, manticoresearch, milvus, qdrant, weaviate, chromadb, redshift, dameng, kingbase, highgo, vastbase, goldendb, databend, gaussdb, kwdb, yashandb, databricks, saphana, teradata, vertica, firebird, exasol, opengauss, oceanbase-oracle, questdb, gbase, h2, snowflake, trino, prestosql, hive, db2, informix, influxdb, iris, neo4j, cassandra, bigquery, kylin, sundb, oscar, tdengine, iotdb, xugu, zookeeper, jdbc, access, mq";
|
|
const FILE_CAPABLE_CONNECTION_TYPES = new Set(["sqlite", "duckdb", "access", "h2"]);
|
|
|
|
interface McpScope {
|
|
connectionId?: string;
|
|
connectionName?: string;
|
|
database?: string;
|
|
}
|
|
|
|
function scopedValue(value: string | undefined): string | undefined {
|
|
const trimmed = value?.trim();
|
|
return trimmed ? trimmed : undefined;
|
|
}
|
|
|
|
function mcpScopeFromEnv(): McpScope {
|
|
return {
|
|
connectionId: scopedValue(process.env.DBX_MCP_SCOPE_CONNECTION_ID),
|
|
connectionName: scopedValue(process.env.DBX_MCP_SCOPE_CONNECTION_NAME),
|
|
database: scopedValue(process.env.DBX_MCP_SCOPE_DATABASE),
|
|
};
|
|
}
|
|
|
|
function scopeEnabled(scope: McpScope): boolean {
|
|
return !!(scope.connectionId || scope.connectionName);
|
|
}
|
|
|
|
function connectionMatchesScope(config: ConnectionConfig, scope: McpScope): boolean {
|
|
return (!!scope.connectionId && config.id === scope.connectionId) || (!!scope.connectionName && config.name === scope.connectionName);
|
|
}
|
|
|
|
async function loadScopedConnections(backend: Backend, scope: McpScope): Promise<ConnectionConfig[]> {
|
|
const connections = await backend.loadConnections();
|
|
if (!scopeEnabled(scope)) return connections;
|
|
return connections.filter((config) => connectionMatchesScope(config, scope));
|
|
}
|
|
|
|
async function resolveConnection(backend: Backend, scope: McpScope, requestedName?: string): Promise<{ config?: ConnectionConfig; error?: ReturnType<typeof toolError> }> {
|
|
if (!scopeEnabled(scope)) {
|
|
if (!requestedName?.trim()) return { error: toolError("CONNECTION_NOT_FOUND", "Connection name is required.") };
|
|
const config = await backend.findConnection(requestedName);
|
|
return config ? { config } : { error: toolError("CONNECTION_NOT_FOUND", `Connection "${requestedName}" not found.`) };
|
|
}
|
|
|
|
const [scopedConfig] = await loadScopedConnections(backend, scope);
|
|
if (!scopedConfig) return { error: toolError("CONNECTION_NOT_FOUND", "Scoped DBX connection was not found.") };
|
|
if (requestedName?.trim() && requestedName !== scopedConfig.name && requestedName !== scopedConfig.id) {
|
|
return { error: toolError("CONNECTION_OUT_OF_SCOPE", `Connection "${requestedName}" is outside this DBX AI session scope.`) };
|
|
}
|
|
return { config: scopedConfig };
|
|
}
|
|
|
|
export function createDbxMcpServer(backend: Backend, options: { isWebMode?: boolean } = {}): McpServer {
|
|
const isWebMode = options.isWebMode ?? !!process.env.DBX_WEB_URL;
|
|
const scope = mcpScopeFromEnv();
|
|
const scoped = scopeEnabled(scope);
|
|
const server = new McpServer({
|
|
name: "dbx",
|
|
version: DBX_MCP_PACKAGE_VERSION,
|
|
});
|
|
|
|
server.tool("dbx_list_connections", "List all database connections configured in DBX", {}, async () => {
|
|
const connections = await loadScopedConnections(backend, scope);
|
|
if (connections.length === 0) return text("No connections configured in DBX.");
|
|
const rows = connections.map((c) => [c.name, c.db_type, c.host, String(c.port), c.database || ""]);
|
|
return text(mdTable(["Name", "Type", "Host", "Port", "Database"], rows));
|
|
});
|
|
|
|
server.tool(
|
|
"dbx_list_tables",
|
|
"List tables and views for a database connection",
|
|
{
|
|
connection_name: z.string().optional().describe("Name of the DBX connection"),
|
|
database: z.string().optional().describe("Database name"),
|
|
schema: z.string().optional().describe("Schema name (default: public for PostgreSQL)"),
|
|
},
|
|
async ({ connection_name, database, schema }) => {
|
|
const { config, error } = await resolveConnection(backend, scope, connection_name);
|
|
if (error) return error;
|
|
const tables = await backend.listTables(withDatabase(config!, database ?? scope.database), schema);
|
|
if (tables.length === 0) return text("No tables found.");
|
|
const rows = tables.map((t) => [t.name, t.type]);
|
|
return text(mdTable(["Table", "Type"], rows));
|
|
},
|
|
);
|
|
|
|
server.tool(
|
|
"dbx_describe_table",
|
|
"Get column definitions for a table",
|
|
{
|
|
connection_name: z.string().optional().describe("Name of the DBX connection"),
|
|
table: z.string().describe("Table name"),
|
|
database: z.string().optional().describe("Database name"),
|
|
schema: z.string().optional().describe("Schema name (default: public for PostgreSQL)"),
|
|
},
|
|
async ({ connection_name, table, database, schema }) => {
|
|
const { config, error } = await resolveConnection(backend, scope, connection_name);
|
|
if (error) return error;
|
|
const columns = await backend.describeTable(withDatabase(config!, database ?? scope.database), table, schema);
|
|
if (columns.length === 0) return text("No columns found.");
|
|
const rows = columns.map((c) => [c.is_primary_key ? `${c.name} (PK)` : c.name, c.data_type, c.is_nullable ? "YES" : "NO", c.column_default ?? "", c.comment ?? ""]);
|
|
return text(mdTable(["Column", "Type", "Nullable", "Default", "Comment"], rows));
|
|
},
|
|
);
|
|
|
|
server.tool(
|
|
"dbx_execute_query",
|
|
"Execute a SQL query on a database connection (max 100 rows returned)",
|
|
{
|
|
connection_name: z.string().optional().describe("Name of the DBX connection"),
|
|
database: z.string().optional().describe("Database name"),
|
|
sql: z.string().describe("SQL query to execute"),
|
|
},
|
|
async ({ connection_name, database, sql }) => {
|
|
const { config, error } = await resolveConnection(backend, scope, connection_name);
|
|
if (error) return error;
|
|
const scopedConfig = config!;
|
|
if (scopedConfig.db_type === "redis") {
|
|
return toolError("REDIS_COMMAND_REQUIRED", "Redis connections do not accept SQL through dbx_execute_query. Use dbx_execute_redis_command with a Redis command such as GET key or INFO.");
|
|
}
|
|
if (scopedConfig.db_type !== "mongodb") {
|
|
const safety = evaluateSqlSafety(sql, { ...sqlSafetyFromEnv(), allowMultipleStatements: true });
|
|
if (!safety.allowed) return toolError("SQL_BLOCKED", safety.reason ?? "SQL blocked.");
|
|
}
|
|
// MongoDB shell commands don't fit the SQL safety evaluator; the backend
|
|
// (node-core executeQuery) applies command-aware read/write gating.
|
|
try {
|
|
const statements = scopedConfig.db_type === "mongodb" ? [sql] : splitSqlStatements(sql);
|
|
const results = [];
|
|
for (const statement of statements) {
|
|
results.push(await backend.executeQuery(withDatabase(scopedConfig, database ?? scope.database), statement));
|
|
}
|
|
if (results.length === 1) return formatQueryToolResult(results[0]);
|
|
return text(results.map((result, index) => formatQueryToolResult(result, `Statement ${index + 1}`).content[0].text).join("\n\n"));
|
|
} catch (e: unknown) {
|
|
const msg = e instanceof Error ? e.message : String(e);
|
|
return toolError("QUERY_ERROR", msg);
|
|
}
|
|
},
|
|
);
|
|
|
|
server.tool(
|
|
"dbx_execute_redis_command",
|
|
"Execute a Redis command on a Redis connection",
|
|
{
|
|
connection_name: z.string().optional().describe("Name of the DBX Redis connection"),
|
|
db: z.number().int().min(0).optional().describe("Redis logical database number (default: scoped/default database or 0)"),
|
|
command: z.string().describe("Redis command to execute, for example: GET mykey, INFO, or DBSIZE"),
|
|
},
|
|
async ({ connection_name, db, command }) => {
|
|
const { config, error } = await resolveConnection(backend, scope, connection_name);
|
|
if (error) return error;
|
|
const scopedConfig = config!;
|
|
if (scopedConfig.db_type !== "redis") {
|
|
return toolError("INVALID_CONNECTION_TYPE", `Connection "${scopedConfig.name}" is ${scopedConfig.db_type}, not Redis.`);
|
|
}
|
|
if (!backend.executeRedisCommand) {
|
|
return toolError("UNSUPPORTED_BACKEND", "This DBX backend does not support Redis command execution.");
|
|
}
|
|
const safety = evaluateRedisCommandSafety(command, sqlSafetyFromEnv());
|
|
if (!safety.allowed) return toolError("REDIS_COMMAND_BLOCKED", safety.reason ?? "Redis command blocked.");
|
|
try {
|
|
const result = await backend.executeRedisCommand(scopedConfig, defaultRedisDb(scopedConfig, scope, db), command, {
|
|
skipSafetyCheck: safety.skipSafetyCheck,
|
|
});
|
|
return formatRedisCommandToolResult(result);
|
|
} catch (e: unknown) {
|
|
const msg = e instanceof Error ? e.message : String(e);
|
|
return toolError("REDIS_COMMAND_ERROR", msg);
|
|
}
|
|
},
|
|
);
|
|
|
|
server.tool(
|
|
"dbx_get_schema_context",
|
|
"Get compact table and column context for writing SQL",
|
|
{
|
|
connection_name: z.string().optional().describe("Name of the DBX connection"),
|
|
database: z.string().optional().describe("Database name"),
|
|
schema: z.string().optional().describe("Schema name (default: public for PostgreSQL)"),
|
|
tables: z.array(z.string()).optional().describe("Specific table names to include"),
|
|
max_tables: z.number().int().min(1).max(20).default(8).describe("Maximum number of tables to include"),
|
|
},
|
|
async ({ connection_name, database, schema, tables, max_tables }) => {
|
|
const { config, error } = await resolveConnection(backend, scope, connection_name);
|
|
if (error) return error;
|
|
const context = await buildSchemaContext(backend, withDatabase(config!, database ?? scope.database), {
|
|
schema,
|
|
tables,
|
|
maxTables: max_tables,
|
|
});
|
|
if (context.tables.length === 0) return text("No matching tables found.");
|
|
return text(formatSchemaContext(context));
|
|
},
|
|
);
|
|
|
|
if (!scoped) {
|
|
server.tool(
|
|
"dbx_add_connection",
|
|
"Add a new database connection to DBX",
|
|
{
|
|
name: z.string().describe("Connection name"),
|
|
db_type: z.string().describe(DBX_CONNECTION_TYPE_DESCRIPTION),
|
|
host: z.string().describe("Database host"),
|
|
port: z.number().optional().describe("Database port (TDengine defaults to 6041, IoTDB defaults to 6667, XuguDB defaults to 5138)"),
|
|
username: z.string().default("").describe("Username"),
|
|
password: z.string().default("").describe("Password"),
|
|
database: z.string().optional().describe("Default database name"),
|
|
ssl: z.boolean().default(false).describe("Enable SSL"),
|
|
driver_profile: z.string().optional().describe("Driver profile (e.g. 'gbase8a', 'gbase8s')"),
|
|
},
|
|
async ({ name, db_type, host, port, username, password, database, ssl, driver_profile }) => {
|
|
const existing = await backend.findConnection(name);
|
|
if (existing) return text(`Connection "${name}" already exists.`);
|
|
const DEFAULT_PORTS: Record<string, number> = {
|
|
kwdb: 26257,
|
|
rqlite: 4001,
|
|
tdengine: 6041,
|
|
oscar: 2003,
|
|
iotdb: 6667,
|
|
xugu: 5138,
|
|
};
|
|
const resolvedPort = port ?? DEFAULT_PORTS[db_type] ?? (FILE_CAPABLE_CONNECTION_TYPES.has(db_type) ? 0 : undefined);
|
|
if (resolvedPort === undefined) return text("Port is required for this database type.");
|
|
const config = await backend.addConnection({
|
|
name,
|
|
db_type,
|
|
host,
|
|
port: resolvedPort,
|
|
username,
|
|
password,
|
|
database,
|
|
ssl,
|
|
driver_profile,
|
|
ssh_enabled: false,
|
|
} as Omit<ConnectionConfig, "id">);
|
|
await notifyReload();
|
|
return text(`Connection "${config.name}" added (id: ${config.id}).`);
|
|
},
|
|
);
|
|
|
|
server.tool(
|
|
"dbx_remove_connection",
|
|
"Remove a database connection from DBX",
|
|
{
|
|
connection_name: z.string().describe("Name of the connection to remove"),
|
|
},
|
|
async ({ connection_name }) => {
|
|
const removed = await backend.removeConnection(connection_name);
|
|
if (!removed) return toolError("CONNECTION_NOT_FOUND", `Connection "${connection_name}" not found.`);
|
|
await notifyReload();
|
|
return text(`Connection "${connection_name}" removed.`);
|
|
},
|
|
);
|
|
}
|
|
|
|
// Desktop-only tools: open table and execute-and-show require the Tauri bridge
|
|
if (!isWebMode && !scoped) {
|
|
server.tool(
|
|
"dbx_open_table",
|
|
"Open a table in DBX desktop app UI. Requires DBX to be running.",
|
|
{
|
|
connection_name: z.string().describe("Name of the DBX connection"),
|
|
table: z.string().describe("Table name to open"),
|
|
database: z.string().optional().describe("Database name"),
|
|
schema: z.string().optional().describe("Schema name"),
|
|
},
|
|
async ({ connection_name, table, database, schema }) => {
|
|
const config = await backend.findConnection(connection_name);
|
|
if (!config) return toolError("CONNECTION_NOT_FOUND", `Connection "${connection_name}" not found.`);
|
|
return bridgeRequest("/open-table", { connection_name, table, database, schema }, `Opened ${table} in DBX`);
|
|
},
|
|
);
|
|
|
|
server.tool(
|
|
"dbx_execute_and_show",
|
|
"Execute a SQL query in DBX desktop app UI and show results there. Requires DBX to be running.",
|
|
{
|
|
connection_name: z.string().describe("Name of the DBX connection"),
|
|
sql: z.string().describe("SQL query to execute"),
|
|
database: z.string().optional().describe("Database name"),
|
|
},
|
|
async ({ connection_name, sql, database }) => {
|
|
const config = await backend.findConnection(connection_name);
|
|
if (!config) return toolError("CONNECTION_NOT_FOUND", `Connection "${connection_name}" not found.`);
|
|
const safetyOptions = sqlSafetyFromEnv();
|
|
if (config?.db_type === "mongodb") {
|
|
const aggregate = parseMongoAggregateCommand(sql);
|
|
if (aggregate) {
|
|
const safety = evaluateMongoAggregateSafety(aggregate, safetyOptions);
|
|
if (!safety.allowed) return toolError("SQL_BLOCKED", safety.reason ?? "Query blocked.");
|
|
}
|
|
} else {
|
|
const safety = evaluateSqlSafety(sql, { ...safetyOptions, allowMultipleStatements: true });
|
|
if (!safety.allowed) return toolError("SQL_BLOCKED", safety.reason ?? "SQL blocked.");
|
|
}
|
|
// MongoDB shell commands bypass the SQL safety evaluator; pass MCP
|
|
// safety flags to the desktop executor for command-aware gating.
|
|
return bridgeRequest(
|
|
"/execute-query",
|
|
{
|
|
connection_name,
|
|
sql,
|
|
database,
|
|
allow_writes: safetyOptions.allowWrites,
|
|
allow_dangerous: safetyOptions.allowDangerous,
|
|
},
|
|
"Query sent to DBX",
|
|
);
|
|
},
|
|
);
|
|
}
|
|
|
|
return server;
|
|
}
|
|
|
|
async function bridgeRequest(path: string, body: Record<string, unknown>, successMsg: string) {
|
|
const res = await postBridge(path, body);
|
|
if (res.ok) return text(successMsg);
|
|
const message = res.text.startsWith("DBX is not running") ? res.text : `Failed: ${res.text}`;
|
|
return toolError("DBX_NOT_RUNNING", message);
|
|
}
|
|
|
|
async function main() {
|
|
const backend = await createBackend();
|
|
const server = createDbxMcpServer(backend);
|
|
const transport = new StdioServerTransport();
|
|
await server.connect(transport);
|
|
}
|
|
|
|
if (isMainModule(import.meta.url, process.argv[1])) {
|
|
main().catch((e) => {
|
|
console.error("MCP Server failed to start:", e);
|
|
process.exit(1);
|
|
});
|
|
}
|