fix(db): reject comments and unusual quoting in agent SQL and filter every agent table on the server

Agent SQL is checked by a validator and then run by ClickHouse, and the
two could read a query differently when it used syntax the validator
did not model. Agent SQL now rejects comments, dollar-quoted strings,
backslashes outside string literals, and quoted identifiers other than
plain names. The server-side tenant filter now covers every agent table
instead of only the tables the validator found in the query.

Author: izadoesdev

packages/ai/src/ai/tools/execute-sql-query.ts

@@ -1,7 +1,7 @@
 import {
 	AGENT_SQL_VALIDATION_ERROR,
+	AGENT_TENANT_COLUMN_BY_TABLE,
 	buildAdditionalTableFilters,
-	extractAllowlistedTables,
 	validateAgentSQL,
 } from "@databuddy/db/clickhouse";
 import { tool } from "ai";
@@ -46,9 +46,8 @@ export async function executeAgentSqlForWebsite({
 		throw new Error(validation.reason ?? AGENT_SQL_VALIDATION_ERROR);
 	}
 
-	const referencedTables = extractAllowlistedTables(sql);
 	const additional_table_filters = buildAdditionalTableFilters(
-		referencedTables,
+		Object.keys(AGENT_TENANT_COLUMN_BY_TABLE),
 		websiteId
 	);
 

packages/db/src/clickhouse/sql-validation.test.ts

@@ -156,6 +156,61 @@ describe("validateAgentSQL", () => {
 		expect(result).toEqual({ valid: true, reason: null });
 	});
 
+	it("rejects a # comment that hides a second query from the validator", () => {
+		const result = validateAgentSQL(
+			`SELECT 'filtered' AS src, count() AS c FROM analytics.events
+WHERE client_id = {websiteId:String} AND 1 = 1 # '
+UNION ALL SELECT 'all_revenue' AS src, count() AS c FROM analytics.revenue -- '`
+		);
+		expect(result.valid).toBe(false);
+		expect(result.reason).toContain("Comments are not allowed");
+	});
+
+	it("rejects every comment form ClickHouse accepts", () => {
+		for (const comment of ["# note", "#! note", "-- note", "/* note */"]) {
+			const result = validateAgentSQL(
+				`SELECT count() FROM analytics.events ${TENANT} ${comment}`
+			);
+			expect(result.valid).toBe(false);
+		}
+	});
+
+	it("rejects settings hidden behind a comment and quote", () => {
+		const result = validateAgentSQL(
+			`SELECT count() FROM analytics.events ${TENANT} # '
+SETTINGS additional_table_filters = {} -- '`
+		);
+		expect(result.valid).toBe(false);
+	});
+
+	it("rejects quoted identifiers that could hide a quote", () => {
+		for (const sql of [
+			`SELECT count() AS \`a'b\` FROM analytics.events ${TENANT}`,
+			`SELECT count() AS "a'b" FROM analytics.events ${TENANT}`,
+			`SELECT count() AS \`unclosed FROM analytics.events ${TENANT}`,
+		]) {
+			expect(validateAgentSQL(sql).valid).toBe(false);
+		}
+	});
+
+	it("rejects dollar-quoted strings and unterminated literals", () => {
+		expect(
+			validateAgentSQL(`SELECT $$x$$ FROM analytics.events ${TENANT}`).valid
+		).toBe(false);
+		expect(
+			validateAgentSQL(
+				`SELECT count() FROM analytics.events ${TENANT} AND path = '/a`
+			).valid
+		).toBe(false);
+	});
+
+	it("keeps comment and quote characters inside string literals", () => {
+		const result = validateAgentSQL(
+			`SELECT count() FROM analytics.events ${TENANT} AND path = '/a#b--c/*d$e"f''g\\'h'`
+		);
+		expect(result).toEqual({ valid: true, reason: null });
+	});
+
 	it("is case-insensitive for FROM/JOIN keywords", () => {
 		const result = validateAgentSQL(
 			"select count() from analytics.events where client_id = {websiteId:String}"
@@ -363,12 +418,12 @@ describe("validateAgentSQL", () => {
 			expect(result.valid).toBe(false);
 		});
 
-		it("ignores tenant markers inside comments", () => {
+		it("rejects tenant markers inside comments", () => {
 			const result = validateAgentSQL(
 				"SELECT * FROM analytics.events /* client_id = {websiteId:String} */ WHERE time > now()"
 			);
 			expect(result.valid).toBe(false);
-			expect(result.reason).toContain("client_id");
+			expect(result.reason).toContain("Comments are not allowed");
 		});
 
 		it("ignores tenant markers inside string literals", () => {

packages/db/src/clickhouse/sql-validation.ts

@@ -185,6 +185,63 @@ const WILDCARD_PROJECTION_PATTERN =
 	/(?:^|,)\s*(?:(?:DISTINCT|ALL)\s+)?(?:[a-zA-Z_][a-zA-Z0-9_]*\s*\.\s*)?\*\s*(?=,|$|\b(?:APPLY|EXCEPT|REPLACE)\b)/i;
 const SENSITIVE_PROJECTION_PATTERN = /\b(?:ip|properties|url|user_agent)\b/i;
 
+const PLAIN_QUOTED_IDENTIFIER_PATTERN = /^[A-Za-z0-9_.]+$/;
+
+function agentSqlSyntaxError(sql: string): string | null {
+	let index = 0;
+	while (index < sql.length) {
+		const char = sql[index];
+		const next = sql[index + 1];
+		if (char === "'") {
+			index += 1;
+			let closed = false;
+			while (index < sql.length) {
+				if (sql[index] === "\\") {
+					index += 2;
+					continue;
+				}
+				if (sql[index] === "'") {
+					if (sql[index + 1] === "'") {
+						index += 2;
+						continue;
+					}
+					index += 1;
+					closed = true;
+					break;
+				}
+				index += 1;
+			}
+			if (!closed) {
+				return "Unterminated string literal.";
+			}
+			continue;
+		}
+		if (char === "`" || char === '"') {
+			const end = sql.indexOf(char, index + 1);
+			if (
+				end < 0 ||
+				!PLAIN_QUOTED_IDENTIFIER_PATTERN.test(sql.slice(index + 1, end))
+			) {
+				return "Quoted identifiers may contain only letters, digits, underscores and dots.";
+			}
+			index = end + 1;
+			continue;
+		}
+		if (
+			char === "#" ||
+			(char === "-" && next === "-") ||
+			(char === "/" && next === "*")
+		) {
+			return "Comments are not allowed; remove them and resend the query.";
+		}
+		if (char === "$" || char === "\\") {
+			return `The character ${char} is not allowed outside string literals.`;
+		}
+		index += 1;
+	}
+	return null;
+}
+
 function maskCommentsAndStrings(sql: string): string {
 	let result = "";
 	let index = 0;
@@ -446,6 +503,10 @@ export function validateAgentSQL(sql: string): {
 	valid: boolean;
 	reason: string | null;
 } {
+	const syntaxError = agentSqlSyntaxError(sql);
+	if (syntaxError) {
+		return { valid: false, reason: syntaxError };
+	}
 	const sanitized = maskCommentsAndStrings(sql);
 
 	if (!SELECT_OR_WITH_PATTERN.test(sanitized)) {