feat(mcp): improve schema context and query safety
This commit is contained in:
t8y2
parent
f4566221b1
commit
9c4ac66baf
|
|
@ -7,9 +7,10 @@ MCP server for [DBX](https://github.com/t8y2/dbx) — lets AI agents (Claude Cod
|
|||
## Features
|
||||
|
||||
- **Zero config** — Automatically reads your DBX connections (including passwords from system keyring)
|
||||
- **7 tools** — List/add/remove connections, list tables, describe table, execute SQL, open table in DBX UI
|
||||
- **8 tools** — List/add/remove connections, list tables, describe table, get schema context, execute SQL, open table in DBX UI
|
||||
- **Connection pooling** — Reuses database connections across queries
|
||||
- **PostgreSQL & MySQL** — Supports PostgreSQL, MySQL, and compatible databases (Doris, StarRocks, etc.)
|
||||
- **Direct execution** — PostgreSQL, MySQL, SQLite, and compatible databases (Doris, StarRocks, etc.) can run without opening DBX
|
||||
- **Read-only by default** — SQL execution blocks write and dangerous statements unless explicitly enabled
|
||||
- **DBX UI integration** — Open tables directly in the DBX desktop app from your AI agent
|
||||
|
||||
## Quick Start
|
||||
|
|
@ -73,9 +74,24 @@ In Claude Code, just ask:
|
|||
| `dbx_remove_connection` | Remove a database connection |
|
||||
| `dbx_list_tables` | List tables and views for a connection |
|
||||
| `dbx_describe_table` | Get column definitions for a table |
|
||||
| `dbx_get_schema_context` | Get compact table and column context for writing SQL |
|
||||
| `dbx_execute_query` | Execute a SQL query (max 100 rows) |
|
||||
| `dbx_open_table` | Open a table in DBX desktop app UI |
|
||||
|
||||
## SQL Safety
|
||||
|
||||
`dbx_execute_query` is read-only by default. To allow write statements such as `INSERT` or `UPDATE`, set:
|
||||
|
||||
```bash
|
||||
DBX_MCP_ALLOW_WRITES=1
|
||||
```
|
||||
|
||||
Dangerous statements such as `DROP`, `TRUNCATE`, and `ALTER` remain blocked unless you also set:
|
||||
|
||||
```bash
|
||||
DBX_MCP_ALLOW_DANGEROUS_SQL=1
|
||||
```
|
||||
|
||||
## How It Works
|
||||
|
||||
```
|
||||
|
|
@ -94,6 +110,8 @@ The MCP server reads your database connections from DBX's SQLite database:
|
|||
|
||||
The `dbx_open_table` tool communicates with the running DBX app to open tables directly in the UI. This requires DBX to be running. If DBX is not running, the tool will return an error message.
|
||||
|
||||
PostgreSQL, MySQL, SQLite, Doris, StarRocks, and Redshift queries run directly from the MCP server. Other database types still use the DBX desktop bridge for query, table, and column operations unless `DBX_WEB_URL` is configured.
|
||||
|
||||
## Requirements
|
||||
|
||||
- [DBX](https://github.com/t8y2/dbx) installed with at least one connection configured
|
||||
|
|
@ -112,9 +130,10 @@ MIT
|
|||
### 特性
|
||||
|
||||
- **零配置** — 自动读取 DBX 的连接配置
|
||||
- **7 个工具** — 列出/添加/删除连接、列出表、查看表结构、执行 SQL、在 DBX 中打开表
|
||||
- **8 个工具** — 列出/添加/删除连接、列出表、查看表结构、获取 Schema 上下文、执行 SQL、在 DBX 中打开表
|
||||
- **连接池** — 跨查询复用数据库连接
|
||||
- **PostgreSQL 和 MySQL** — 支持 PostgreSQL、MySQL 及兼容数据库（Doris、StarRocks 等）
|
||||
- **直接执行** — PostgreSQL、MySQL、SQLite 及兼容数据库（Doris、StarRocks 等）无需打开 DBX 即可查询
|
||||
- **默认只读** — SQL 执行默认拦截写操作和危险语句
|
||||
- **DBX UI 联动** — 从 AI 助手直接在 DBX 桌面端打开表
|
||||
|
||||
### 快速开始
|
||||
|
|
@ -164,9 +183,24 @@ npx @dbx-app/mcp-server
|
|||
| `dbx_remove_connection` | 删除数据库连接 |
|
||||
| `dbx_list_tables` | 列出指定连接的表和视图 |
|
||||
| `dbx_describe_table` | 获取表的列定义 |
|
||||
| `dbx_get_schema_context` | 获取适合 AI 写 SQL 的紧凑表结构上下文 |
|
||||
| `dbx_execute_query` | 执行 SQL 查询（最多返回 100 行） |
|
||||
| `dbx_open_table` | 在 DBX 桌面端打开指定表 |
|
||||
|
||||
### SQL 安全
|
||||
|
||||
`dbx_execute_query` 默认只读。若要允许 `INSERT`、`UPDATE` 等写操作，设置：
|
||||
|
||||
```bash
|
||||
DBX_MCP_ALLOW_WRITES=1
|
||||
```
|
||||
|
||||
`DROP`、`TRUNCATE`、`ALTER` 等危险语句仍会被拦截，除非额外设置：
|
||||
|
||||
```bash
|
||||
DBX_MCP_ALLOW_DANGEROUS_SQL=1
|
||||
```
|
||||
|
||||
### 工作原理
|
||||
|
||||
MCP Server 从 DBX 的 SQLite 数据库读取连接信息：
|
||||
|
|
@ -179,6 +213,8 @@ MCP Server 从 DBX 的 SQLite 数据库读取连接信息：
|
|||
|
||||
`dbx_open_table` 工具通过本地 HTTP 接口与运行中的 DBX 应用通信，直接在 UI 中打开表。需要 DBX 正在运行。
|
||||
|
||||
PostgreSQL、MySQL、SQLite、Doris、StarRocks、Redshift 查询可由 MCP Server 直接执行。其他数据库类型的查询、表列表、字段读取仍会走 DBX 桌面端 bridge，除非配置了 `DBX_WEB_URL` 使用 Web 后端。
|
||||
|
||||
### 系统要求
|
||||
|
||||
- 已安装 [DBX](https://github.com/t8y2/dbx) 并配置了至少一个数据库连接
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@
|
|||
],
|
||||
"scripts": {
|
||||
"start": "tsx src/index.ts",
|
||||
"test": "tsx --test tests/*.test.ts",
|
||||
"build": "tsc",
|
||||
"prepublishOnly": "tsc"
|
||||
},
|
||||
|
|
@ -52,4 +53,4 @@
|
|||
"esbuild"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@ import { createServer, connect as netConnect, type Server, type Socket } from "n
|
|||
import { readFile } from "node:fs/promises";
|
||||
import { join } from "node:path";
|
||||
import { homedir, platform } from "node:os";
|
||||
import Database from "better-sqlite3";
|
||||
import { sqlSafetyFromEnv } from "./sql-safety.js";
|
||||
|
||||
export interface TableInfo {
|
||||
name: string;
|
||||
|
|
@ -235,6 +237,7 @@ function isDirectType(dbType: string): boolean {
|
|||
case "mysql":
|
||||
case "doris":
|
||||
case "starrocks":
|
||||
case "sqlite":
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
|
|
@ -355,10 +358,40 @@ async function mysqlQuery(config: ConnectionConfig, sql: string, params?: unknow
|
|||
}
|
||||
|
||||
async function query(config: ConnectionConfig, sql: string, params?: unknown[]): Promise<QueryResult> {
|
||||
if (config.db_type === "sqlite") return sqliteQuery(config, sql);
|
||||
if (isMysqlType(config.db_type)) return mysqlQuery(config, sql, params);
|
||||
return pgQuery(config, sql, params);
|
||||
}
|
||||
|
||||
function sqlitePath(config: ConnectionConfig): string {
|
||||
return expandTilde(config.host || config.database || "");
|
||||
}
|
||||
|
||||
function expandTilde(path: string): string {
|
||||
if (path === "~") return homedir();
|
||||
if (path.startsWith("~/")) return join(homedir(), path.slice(2));
|
||||
return path;
|
||||
}
|
||||
|
||||
function quoteSqliteIdentifier(identifier: string): string {
|
||||
return `"${identifier.replace(/"/g, '""')}"`;
|
||||
}
|
||||
|
||||
function sqliteQuery(config: ConnectionConfig, sql: string): QueryResult {
|
||||
const db = new Database(sqlitePath(config), { readonly: !sqlSafetyFromEnv().allowWrites });
|
||||
try {
|
||||
const stmt = db.prepare(sql);
|
||||
if (stmt.reader) {
|
||||
const rows = stmt.all().slice(0, MAX_ROWS) as Record<string, unknown>[];
|
||||
return { columns: stmt.columns().map((column) => column.name), rows, row_count: rows.length };
|
||||
}
|
||||
const result = stmt.run();
|
||||
return { columns: [], rows: [], row_count: result.changes };
|
||||
} finally {
|
||||
db.close();
|
||||
}
|
||||
}
|
||||
|
||||
export async function executeQuery(config: ConnectionConfig, sql: string): Promise<QueryResult> {
|
||||
if (isDirectType(config.db_type)) {
|
||||
return query(config, sql);
|
||||
|
|
@ -372,6 +405,13 @@ export async function executeQuery(config: ConnectionConfig, sql: string): Promi
|
|||
}
|
||||
|
||||
export async function listTables(config: ConnectionConfig, schema?: string): Promise<TableInfo[]> {
|
||||
if (config.db_type === "sqlite") {
|
||||
const result = await query(
|
||||
config,
|
||||
`SELECT name, type FROM sqlite_master WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%' ORDER BY name`,
|
||||
);
|
||||
return result.rows.map((r) => ({ name: String(r.name || ""), type: String(r.type || "table") }));
|
||||
}
|
||||
if (!isDirectType(config.db_type)) {
|
||||
const tables = await bridgeDataRequest<BridgeTableInfo[]>("/data/list-tables", {
|
||||
connection_name: config.name,
|
||||
|
|
@ -394,6 +434,17 @@ export async function listTables(config: ConnectionConfig, schema?: string): Pro
|
|||
}
|
||||
|
||||
export async function describeTable(config: ConnectionConfig, table: string, schema?: string): Promise<ColumnInfo[]> {
|
||||
if (config.db_type === "sqlite") {
|
||||
const result = await query(config, `PRAGMA table_info(${quoteSqliteIdentifier(table)})`);
|
||||
return result.rows.map((r) => ({
|
||||
name: String(r.name || ""),
|
||||
data_type: String(r.type || ""),
|
||||
is_nullable: Number(r.notnull || 0) === 0,
|
||||
column_default: r.dflt_value != null ? String(r.dflt_value) : null,
|
||||
is_primary_key: Number(r.pk || 0) > 0,
|
||||
comment: null,
|
||||
}));
|
||||
}
|
||||
if (!isDirectType(config.db_type)) {
|
||||
const columns = await bridgeDataRequest<BridgeColumnInfo[]>("/data/describe-table", {
|
||||
connection_name: config.name,
|
||||
|
|
|
|||
|
|
@ -18,6 +18,8 @@ import {
|
|||
} from "./database.js";
|
||||
import type { ConnectionConfig } from "./connections.js";
|
||||
import type { TableInfo, ColumnInfo, QueryResult } from "./database.js";
|
||||
import { buildSchemaContext, formatSchemaContext } from "./schema-context.js";
|
||||
import { evaluateSqlSafety, sqlSafetyFromEnv } from "./sql-safety.js";
|
||||
|
||||
const isWebMode = !!process.env.DBX_WEB_URL;
|
||||
|
||||
|
|
@ -122,6 +124,8 @@ server.tool(
|
|||
async ({ connection_name, sql }) => {
|
||||
const config = await backend.findConnection(connection_name);
|
||||
if (!config) return text(`Connection "${connection_name}" not found`);
|
||||
const safety = evaluateSqlSafety(sql, sqlSafetyFromEnv());
|
||||
if (!safety.allowed) return text(`Query blocked: ${safety.reason}`);
|
||||
try {
|
||||
const result = await backend.executeQuery(config, sql);
|
||||
if (result.columns.length === 0) return text(`Query executed. ${result.row_count} row(s) affected.`);
|
||||
|
|
@ -134,6 +138,24 @@ server.tool(
|
|||
},
|
||||
);
|
||||
|
||||
server.tool(
|
||||
"dbx_get_schema_context",
|
||||
"Get compact table and column context for writing SQL",
|
||||
{
|
||||
connection_name: z.string().describe("Name of the DBX connection"),
|
||||
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, schema, tables, max_tables }) => {
|
||||
const config = await backend.findConnection(connection_name);
|
||||
if (!config) return text(`Connection "${connection_name}" not found`);
|
||||
const context = await buildSchemaContext(backend, config, { schema, tables, maxTables: max_tables });
|
||||
if (context.tables.length === 0) return text("No matching tables found.");
|
||||
return text(formatSchemaContext(context));
|
||||
},
|
||||
);
|
||||
|
||||
server.tool(
|
||||
"dbx_add_connection",
|
||||
"Add a new database connection to DBX",
|
||||
|
|
@ -142,7 +164,7 @@ server.tool(
|
|||
db_type: z
|
||||
.string()
|
||||
.describe(
|
||||
"Database type: postgres, mysql, sqlite, redis, duckdb, clickhouse, sqlserver, mongodb, oracle, elasticsearch, doris, starrocks, redshift, dameng, kingbase, highgo, vastbase, goldendb, gaussdb, h2, snowflake, trino, hive, db2, informix, neo4j, cassandra, bigquery, kylin, sundb, tdengine, jdbc",
|
||||
"Database type: postgres, mysql, sqlite, redis, duckdb, clickhouse, sqlserver, mongodb, oracle, elasticsearch, doris, starrocks, redshift, dameng, kingbase, highgo, vastbase, goldendb, gaussdb, h2, snowflake, trino, hive, db2, informix, neo4j, cassandra, bigquery, kylin, sundb, tdengine, jdbc, access",
|
||||
),
|
||||
host: z.string().describe("Database host"),
|
||||
port: z.number().optional().describe("Database port (TDengine defaults to 6041)"),
|
||||
|
|
@ -154,8 +176,9 @@ server.tool(
|
|||
async ({ name, db_type, host, port, username, password, database, ssl }) => {
|
||||
const existing = await backend.findConnection(name);
|
||||
if (existing) return text(`Connection "${name}" already exists.`);
|
||||
const resolvedPort = port ?? (db_type === "tdengine" ? 6041 : undefined);
|
||||
if (!resolvedPort) return text("Port is required for this database type.");
|
||||
const FILE_BASED_TYPES = new Set(["sqlite", "duckdb", "access"]);
|
||||
const resolvedPort = port ?? (db_type === "tdengine" ? 6041 : FILE_BASED_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,
|
||||
|
|
@ -242,6 +265,8 @@ if (!isWebMode) {
|
|||
database: z.string().optional().describe("Database name"),
|
||||
},
|
||||
async ({ connection_name, sql, database }) => {
|
||||
const safety = evaluateSqlSafety(sql, sqlSafetyFromEnv());
|
||||
if (!safety.allowed) return text(`Query blocked: ${safety.reason}`);
|
||||
return bridgeRequest("/execute-query", { connection_name, sql, database }, "Query sent to DBX");
|
||||
},
|
||||
);
|
||||
|
|
|
|||
|
|
@ -0,0 +1,82 @@
|
|||
import type { ConnectionConfig } from "./connections.js";
|
||||
import type { ColumnInfo, TableInfo } from "./database.js";
|
||||
|
||||
export interface SchemaContextBackend {
|
||||
listTables(config: ConnectionConfig, schema?: string): Promise<TableInfo[]>;
|
||||
describeTable(config: ConnectionConfig, table: string, schema?: string): Promise<ColumnInfo[]>;
|
||||
}
|
||||
|
||||
export interface SchemaContextOptions {
|
||||
schema?: string;
|
||||
tables?: string[];
|
||||
maxTables?: number;
|
||||
}
|
||||
|
||||
export interface SchemaContextTable {
|
||||
name: string;
|
||||
type: string;
|
||||
columns: ColumnInfo[];
|
||||
}
|
||||
|
||||
export interface SchemaContext {
|
||||
connection: string;
|
||||
database: string;
|
||||
schema: string;
|
||||
truncated: boolean;
|
||||
tables: SchemaContextTable[];
|
||||
}
|
||||
|
||||
const DEFAULT_MAX_TABLES = 8;
|
||||
|
||||
export async function buildSchemaContext(
|
||||
backend: SchemaContextBackend,
|
||||
config: ConnectionConfig,
|
||||
options: SchemaContextOptions = {},
|
||||
): Promise<SchemaContext> {
|
||||
const maxTables = Math.max(1, Math.min(options.maxTables ?? DEFAULT_MAX_TABLES, 20));
|
||||
const availableTables = await backend.listTables(config, options.schema);
|
||||
const requested = new Set((options.tables ?? []).map((table) => table.toLowerCase()));
|
||||
const selected = requested.size
|
||||
? availableTables.filter((table) => requested.has(table.name.toLowerCase()))
|
||||
: availableTables.slice(0, maxTables);
|
||||
|
||||
const limited = selected.slice(0, maxTables);
|
||||
const tables = await Promise.all(
|
||||
limited.map(async (table) => ({
|
||||
name: table.name,
|
||||
type: table.type,
|
||||
columns: await backend.describeTable(config, table.name, options.schema),
|
||||
})),
|
||||
);
|
||||
|
||||
return {
|
||||
connection: config.name,
|
||||
database: config.database || "",
|
||||
schema: options.schema || "",
|
||||
truncated: selected.length > limited.length || (!requested.size && availableTables.length > limited.length),
|
||||
tables,
|
||||
};
|
||||
}
|
||||
|
||||
export function formatSchemaContext(context: SchemaContext): string {
|
||||
const header = [
|
||||
`Connection: ${context.connection}`,
|
||||
context.database ? `Database: ${context.database}` : "",
|
||||
context.schema ? `Schema: ${context.schema}` : "",
|
||||
].filter(Boolean);
|
||||
const sections = context.tables.map((table) => {
|
||||
const lines = table.columns.map((column) => {
|
||||
const parts = [
|
||||
column.name,
|
||||
column.data_type,
|
||||
column.is_nullable ? "NULL" : "NOT NULL",
|
||||
column.is_primary_key ? "PK" : "",
|
||||
].filter(Boolean);
|
||||
return `- ${parts.join(" ")}${column.comment ? ` -- ${column.comment}` : ""}`;
|
||||
});
|
||||
return [`## ${table.name}`, `Type: ${table.type}`, ...lines].join("\n");
|
||||
});
|
||||
const suffix = context.truncated ? "\n\nNote: table list was truncated; request specific table names for more context." : "";
|
||||
return `${header.join("\n")}\n\n${sections.join("\n\n")}${suffix}`;
|
||||
}
|
||||
|
||||
|
|
@ -0,0 +1,115 @@
|
|||
export interface SqlSafetyOptions {
|
||||
allowWrites?: boolean;
|
||||
allowDangerous?: boolean;
|
||||
}
|
||||
|
||||
export interface SqlSafetyDecision {
|
||||
allowed: boolean;
|
||||
reason?: string;
|
||||
}
|
||||
|
||||
const READ_KEYWORDS = new Set(["select", "with", "show", "describe", "desc", "explain"]);
|
||||
const DANGEROUS_KEYWORDS = new Set(["drop", "truncate", "alter"]);
|
||||
|
||||
export function evaluateSqlSafety(sql: string, options: SqlSafetyOptions = {}): SqlSafetyDecision {
|
||||
const statements = splitSqlStatements(sql);
|
||||
if (statements.length === 0) return { allowed: false, reason: "SQL is empty." };
|
||||
if (statements.length > 1) return { allowed: false, reason: "Only one SQL statement is allowed per MCP query." };
|
||||
|
||||
const normalized = stripSqlCommentsAndStrings(statements[0]).trim();
|
||||
const firstKeyword = normalized.match(/^[a-zA-Z_]+/)?.[0]?.toLowerCase();
|
||||
if (!firstKeyword) return { allowed: false, reason: "SQL statement is not recognized." };
|
||||
|
||||
const tokens: string[] = normalized.toLowerCase().match(/[a-z_]+/g) ?? [];
|
||||
const dangerous = tokens.find((token) => DANGEROUS_KEYWORDS.has(token));
|
||||
if (dangerous && !options.allowDangerous) {
|
||||
return { allowed: false, reason: `Dangerous SQL keyword "${dangerous.toUpperCase()}" is blocked.` };
|
||||
}
|
||||
|
||||
if (!options.allowWrites && !READ_KEYWORDS.has(firstKeyword)) {
|
||||
return {
|
||||
allowed: false,
|
||||
reason: "MCP SQL execution is read-only by default. Set DBX_MCP_ALLOW_WRITES=1 to allow write statements.",
|
||||
};
|
||||
}
|
||||
|
||||
if (options.allowWrites && !options.allowDangerous) {
|
||||
if (firstKeyword === "update" && !tokens.includes("where")) {
|
||||
return { allowed: false, reason: "UPDATE statements must include a WHERE clause." };
|
||||
}
|
||||
if (firstKeyword === "delete" && !tokens.includes("where")) {
|
||||
return { allowed: false, reason: "DELETE statements must include a WHERE clause." };
|
||||
}
|
||||
}
|
||||
|
||||
return { allowed: true };
|
||||
}
|
||||
|
||||
export function sqlSafetyFromEnv(env: NodeJS.ProcessEnv = process.env): SqlSafetyOptions {
|
||||
return {
|
||||
allowWrites: env.DBX_MCP_ALLOW_WRITES === "1" || env.DBX_MCP_ALLOW_WRITES === "true",
|
||||
allowDangerous: env.DBX_MCP_ALLOW_DANGEROUS_SQL === "1" || env.DBX_MCP_ALLOW_DANGEROUS_SQL === "true",
|
||||
};
|
||||
}
|
||||
|
||||
function splitSqlStatements(sql: string): string[] {
|
||||
const statements: string[] = [];
|
||||
let current = "";
|
||||
let quote: "'" | '"' | "`" | null = null;
|
||||
let inLineComment = false;
|
||||
let inBlockComment = false;
|
||||
|
||||
for (let i = 0; i < sql.length; i++) {
|
||||
const char = sql[i];
|
||||
const next = sql[i + 1];
|
||||
|
||||
if (inLineComment) {
|
||||
current += char;
|
||||
if (char === "\n") inLineComment = false;
|
||||
continue;
|
||||
}
|
||||
if (inBlockComment) {
|
||||
current += char;
|
||||
if (char === "*" && next === "/") {
|
||||
current += next;
|
||||
i++;
|
||||
inBlockComment = false;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (quote) {
|
||||
current += char;
|
||||
if (char === quote) {
|
||||
if (next === quote) {
|
||||
current += next;
|
||||
i++;
|
||||
} else {
|
||||
quote = null;
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (char === "-" && next === "-") inLineComment = true;
|
||||
if (char === "/" && next === "*") inBlockComment = true;
|
||||
if (char === "'" || char === '"' || char === "`") quote = char;
|
||||
|
||||
if (char === ";") {
|
||||
if (current.trim()) statements.push(current.trim());
|
||||
current = "";
|
||||
} else {
|
||||
current += char;
|
||||
}
|
||||
}
|
||||
if (current.trim()) statements.push(current.trim());
|
||||
return statements;
|
||||
}
|
||||
|
||||
function stripSqlCommentsAndStrings(sql: string): string {
|
||||
return sql
|
||||
.replace(/--.*$/gm, " ")
|
||||
.replace(/\/\*[\s\S]*?\*\//g, " ")
|
||||
.replace(/'([^']|'')*'/g, "''")
|
||||
.replace(/"([^"]|"")*"/g, '""')
|
||||
.replace(/`([^`]|``)*`/g, "``");
|
||||
}
|
||||
|
|
@ -0,0 +1,88 @@
|
|||
import assert from "node:assert/strict";
|
||||
import test from "node:test";
|
||||
import type { ConnectionConfig } from "../src/connections.js";
|
||||
import type { ColumnInfo, TableInfo } from "../src/database.js";
|
||||
import { buildSchemaContext, formatSchemaContext } from "../src/schema-context.js";
|
||||
|
||||
const config: ConnectionConfig = {
|
||||
id: "pg",
|
||||
name: "analytics",
|
||||
db_type: "postgres",
|
||||
host: "localhost",
|
||||
port: 5432,
|
||||
username: "app",
|
||||
password: "",
|
||||
database: "warehouse",
|
||||
ssh_enabled: false,
|
||||
ssl: false,
|
||||
};
|
||||
|
||||
const tableRows: TableInfo[] = [
|
||||
{ name: "users", type: "BASE TABLE" },
|
||||
{ name: "orders", type: "BASE TABLE" },
|
||||
{ name: "events", type: "BASE TABLE" },
|
||||
];
|
||||
|
||||
const columns: Record<string, ColumnInfo[]> = {
|
||||
users: [
|
||||
{ name: "id", data_type: "integer", is_nullable: false, column_default: null, is_primary_key: true, comment: null },
|
||||
{ name: "email", data_type: "text", is_nullable: false, column_default: null, is_primary_key: false, comment: "Login email" },
|
||||
],
|
||||
orders: [
|
||||
{ name: "id", data_type: "integer", is_nullable: false, column_default: null, is_primary_key: true, comment: null },
|
||||
{ name: "user_id", data_type: "integer", is_nullable: false, column_default: null, is_primary_key: false, comment: null },
|
||||
],
|
||||
};
|
||||
|
||||
test("builds schema context for requested tables", async () => {
|
||||
const context = await buildSchemaContext(
|
||||
{
|
||||
listTables: async () => tableRows,
|
||||
describeTable: async (_config, table) => columns[table] ?? [],
|
||||
},
|
||||
config,
|
||||
{ tables: ["users", "orders"], maxTables: 5 },
|
||||
);
|
||||
|
||||
assert.equal(context.connection, "analytics");
|
||||
assert.deepEqual(
|
||||
context.tables.map((table) => table.name),
|
||||
["users", "orders"],
|
||||
);
|
||||
assert.equal(context.tables[0].columns[0].is_primary_key, true);
|
||||
});
|
||||
|
||||
test("limits schema context when no table list is provided", async () => {
|
||||
const context = await buildSchemaContext(
|
||||
{
|
||||
listTables: async () => tableRows,
|
||||
describeTable: async (_config, table) => columns[table] ?? [],
|
||||
},
|
||||
config,
|
||||
{ maxTables: 2 },
|
||||
);
|
||||
|
||||
assert.deepEqual(
|
||||
context.tables.map((table) => table.name),
|
||||
["users", "orders"],
|
||||
);
|
||||
});
|
||||
|
||||
test("formats schema context as compact markdown", async () => {
|
||||
const context = await buildSchemaContext(
|
||||
{
|
||||
listTables: async () => tableRows,
|
||||
describeTable: async (_config, table) => columns[table] ?? [],
|
||||
},
|
||||
config,
|
||||
{ tables: ["users"] },
|
||||
);
|
||||
|
||||
const markdown = formatSchemaContext(context);
|
||||
|
||||
assert.match(markdown, /Connection: analytics/);
|
||||
assert.match(markdown, /## users/);
|
||||
assert.match(markdown, /id integer NOT NULL PK/);
|
||||
assert.match(markdown, /email text NOT NULL -- Login email/);
|
||||
});
|
||||
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
import assert from "node:assert/strict";
|
||||
import test from "node:test";
|
||||
import { evaluateSqlSafety } from "../src/sql-safety.js";
|
||||
|
||||
test("allows read-only SQL by default", () => {
|
||||
const decision = evaluateSqlSafety("select * from users limit 5");
|
||||
|
||||
assert.equal(decision.allowed, true);
|
||||
});
|
||||
|
||||
test("blocks write SQL by default", () => {
|
||||
const decision = evaluateSqlSafety("update users set role = 'admin' where id = 1");
|
||||
|
||||
assert.equal(decision.allowed, false);
|
||||
assert.match(decision.reason ?? "", /read-only/i);
|
||||
});
|
||||
|
||||
test("blocks dangerous SQL even when writes are enabled", () => {
|
||||
const decision = evaluateSqlSafety("drop table users", { allowWrites: true });
|
||||
|
||||
assert.equal(decision.allowed, false);
|
||||
assert.match(decision.reason ?? "", /dangerous/i);
|
||||
});
|
||||
|
||||
test("blocks update without where when writes are enabled", () => {
|
||||
const decision = evaluateSqlSafety("update users set disabled = true", { allowWrites: true });
|
||||
|
||||
assert.equal(decision.allowed, false);
|
||||
assert.match(decision.reason ?? "", /WHERE/i);
|
||||
});
|
||||
|
||||
|
|
@ -0,0 +1,70 @@
|
|||
import assert from "node:assert/strict";
|
||||
import { mkdtempSync, rmSync } from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import test from "node:test";
|
||||
import Database from "better-sqlite3";
|
||||
import type { ConnectionConfig } from "../src/connections.js";
|
||||
import { describeTable, executeQuery, listTables } from "../src/database.js";
|
||||
|
||||
function sqliteConfig(path: string): ConnectionConfig {
|
||||
return {
|
||||
id: "sqlite-test",
|
||||
name: "local-sqlite",
|
||||
db_type: "sqlite",
|
||||
host: path,
|
||||
port: 0,
|
||||
username: "",
|
||||
password: "",
|
||||
ssh_enabled: false,
|
||||
ssl: false,
|
||||
};
|
||||
}
|
||||
|
||||
test("queries SQLite connections without the DBX bridge", async () => {
|
||||
const dir = mkdtempSync(join(tmpdir(), "dbx-mcp-sqlite-"));
|
||||
const path = join(dir, "app.db");
|
||||
const db = new Database(path);
|
||||
db.exec("create table users (id integer primary key, name text not null); insert into users (name) values ('Ada');");
|
||||
db.close();
|
||||
|
||||
try {
|
||||
const result = await executeQuery(sqliteConfig(path), "select id, name from users");
|
||||
|
||||
assert.deepEqual(result.columns, ["id", "name"]);
|
||||
assert.deepEqual(result.rows, [{ id: 1, name: "Ada" }]);
|
||||
assert.equal(result.row_count, 1);
|
||||
} finally {
|
||||
rmSync(dir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
test("lists and describes SQLite tables without the DBX bridge", async () => {
|
||||
const dir = mkdtempSync(join(tmpdir(), "dbx-mcp-sqlite-"));
|
||||
const path = join(dir, "app.db");
|
||||
const db = new Database(path);
|
||||
db.exec("create table users (id integer primary key, name text not null);");
|
||||
db.close();
|
||||
|
||||
try {
|
||||
const tables = await listTables(sqliteConfig(path));
|
||||
const columns = await describeTable(sqliteConfig(path), "users");
|
||||
|
||||
assert.deepEqual(tables, [{ name: "users", type: "table" }]);
|
||||
assert.deepEqual(
|
||||
columns.map((column) => ({
|
||||
name: column.name,
|
||||
data_type: column.data_type,
|
||||
is_nullable: column.is_nullable,
|
||||
is_primary_key: column.is_primary_key,
|
||||
})),
|
||||
[
|
||||
{ name: "id", data_type: "INTEGER", is_nullable: true, is_primary_key: true },
|
||||
{ name: "name", data_type: "TEXT", is_nullable: false, is_primary_key: false },
|
||||
],
|
||||
);
|
||||
} finally {
|
||||
rmSync(dir, { recursive: true, force: true });
|
||||
}
|
||||
});
|
||||
|
||||
Loading…
Reference in New Issue