diff --git a/packages/mcp-server/tests/server.test.ts b/packages/mcp-server/tests/server.test.ts index d64f01220..617a65e2c 100644 --- a/packages/mcp-server/tests/server.test.ts +++ b/packages/mcp-server/tests/server.test.ts @@ -109,6 +109,25 @@ test("execute query runs safe multi-statement SQL one statement at a time", asyn assert.match(result.content[0].text, /Statement 2/); }); +test("execute query preserves string literals and PostgreSQL dollar quotes", async () => { + const executed: string[] = []; + const scopedBackend: Backend = { + ...backend, + executeQuery: async (_config, sql) => { + executed.push(sql); + return { columns: ["value"], rows: [{ value: 1 }], row_count: 1 }; + }, + }; + const server = createDbxMcpServer(scopedBackend, { isWebMode: true }); + + await (server as any)._registeredTools.dbx_execute_query.handler({ + connection_name: "local", + sql: "SELECT 'a;b' AS first; SELECT $tag$c;d$tag$ AS second;", + }); + + assert.deepEqual(executed, ["SELECT 'a;b' AS first", "SELECT $tag$c;d$tag$ AS second"]); +}); + test("execute query reports the blocked statement number for unsafe multi-statement SQL", async () => { const server = createDbxMcpServer(backend, { isWebMode: true }); diff --git a/packages/node-core/src/sql-safety.ts b/packages/node-core/src/sql-safety.ts index 782b6ccc1..49c16b08e 100644 --- a/packages/node-core/src/sql-safety.ts +++ b/packages/node-core/src/sql-safety.ts @@ -81,5 +81,106 @@ export function sqlSafetyFromEnv(env: NodeJS.ProcessEnv = process.env): SqlSafet } export function splitSqlStatements(sql: string): string[] { - return splitSqlStatementsForSafety(sql); + const statements: string[] = []; + let statementStart = 0; + let index = 0; + let state: "none" | "single" | "double" | "backtick" | "bracket" | "lineComment" | "blockComment" | "dollar" = "none"; + let dollarTag = ""; + + const pushStatement = (end: number) => { + const statement = sql.slice(statementStart, end).trim(); + if (statement) statements.push(statement); + }; + + while (index < sql.length) { + const char = sql[index] ?? ""; + const next = sql[index + 1] ?? ""; + + if (state === "lineComment") { + if (char === "\n" || char === "\r") state = "none"; + index += 1; + continue; + } + if (state === "blockComment") { + if (char === "*" && next === "/") { + state = "none"; + index += 2; + } else { + index += 1; + } + continue; + } + if (state === "dollar") { + if (sql.startsWith(dollarTag, index)) { + index += dollarTag.length; + state = "none"; + } else { + index += 1; + } + continue; + } + if (state === "single" || state === "double" || state === "backtick") { + const quote = state === "single" ? "'" : state === "double" ? '"' : "`"; + if (char === quote) { + if (next === quote) { + index += 2; + continue; + } + state = "none"; + } else if (char === "\\" && next) { + // Preserve dialects that accept backslash escapes without letting an escaped quote end the literal. + index += 2; + continue; + } + index += 1; + continue; + } + if (state === "bracket") { + if (char === "]") { + if (next === "]") { + index += 2; + continue; + } + state = "none"; + } + index += 1; + continue; + } + + if (char === "-" && next === "-") { + state = "lineComment"; + index += 2; + continue; + } + if (char === "#") { + state = "lineComment"; + index += 1; + continue; + } + if (char === "/" && next === "*") { + state = "blockComment"; + index += 2; + continue; + } + if (char === "'") state = "single"; + else if (char === '"') state = "double"; + else if (char === "`") state = "backtick"; + else if (char === "[") state = "bracket"; + else if (char === "$") { + const match = /^\$[A-Za-z_0-9]*\$/.exec(sql.slice(index)); + if (match) { + dollarTag = match[0]; + state = "dollar"; + index += dollarTag.length; + continue; + } + } else if (char === ";") { + pushStatement(index); + statementStart = index + 1; + } + index += 1; + } + + pushStatement(sql.length); + return statements; } diff --git a/packages/node-core/tests/sql-safety.test.ts b/packages/node-core/tests/sql-safety.test.ts index fe2375a5b..0ff4f5a98 100644 --- a/packages/node-core/tests/sql-safety.test.ts +++ b/packages/node-core/tests/sql-safety.test.ts @@ -1,6 +1,6 @@ import assert from "node:assert/strict"; import { test } from "vitest"; -import { evaluateSqlSafety, sqlSafetyFromEnv } from "../src/sql-safety.js"; +import { evaluateSqlSafety, splitSqlStatements, sqlSafetyFromEnv } from "../src/sql-safety.js"; test("allows read-only SQL by default", () => { const decision = evaluateSqlSafety("select * from users limit 5"); @@ -79,6 +79,24 @@ test("checks every statement in a multi-statement SQL string", () => { assert.match(decision.reason ?? "", /WHERE/i); }); +test("splits statements without altering SQL literals or comments", () => { + const sql = "SELECT 'a;b' AS value, ''abc'' AS quoted; -- keep comment\nSELECT $$c;d$$ AS dollar;"; + + assert.deepEqual(splitSqlStatements(sql), [ + "SELECT 'a;b' AS value, ''abc'' AS quoted", + "-- keep comment\nSELECT $$c;d$$ AS dollar", + ]); +}); + +test("keeps tagged dollar quotes and quoted identifiers intact", () => { + const sql = 'SELECT $body$begin; end$body$ AS body, "semi;colon" AS "quoted;column"; SELECT 2;'; + + assert.deepEqual(splitSqlStatements(sql), [ + 'SELECT $body$begin; end$body$ AS body, "semi;colon" AS "quoted;column"', + "SELECT 2", + ]); +}); + test("sqlSafetyFromEnv allows writes by default but keeps dangerous SQL blocked", () => { const options = sqlSafetyFromEnv({});