diff --git a/apps/ai/src/mcp/__evals__/fake-warehouse.ts b/apps/ai/src/mcp/__evals__/fake-warehouse.ts index 2147e2de33..5c95bd6ef1 100644 --- a/apps/ai/src/mcp/__evals__/fake-warehouse.ts +++ b/apps/ai/src/mcp/__evals__/fake-warehouse.ts @@ -4,12 +4,19 @@ import { __testables } from "@maple/backend/services/warehouse/WarehouseQuerySer import { makeLargeTraceSpans, makeTraceLogs } from "./fixtures" import { FIXTURES } from "./utils" -export interface FixtureRule { +export type FixtureRule = { /** Match against the compiled SQL the WarehouseQueryService would execute. */ readonly match: (sql: string) => boolean - /** Fixed rows, or rows computed from the SQL (to honor a service filter, say). */ - readonly rows: ReadonlyArray | ((sql: string) => ReadonlyArray) -} +} & ( + | { + /** Fixed rows, or rows computed from the SQL (to honor a service filter, say). */ + readonly rows: ReadonlyArray | ((sql: string) => ReadonlyArray) + } + | { + /** Fail the statement as ClickHouse would, with this message and error type. */ + readonly failure: { readonly message: string; readonly type: string } + } +) /** A read of known traces' spans: `traces` seeked through the `trace_index` subquery. */ export const readsTraceSpans = (sql: string): boolean => @@ -75,6 +82,16 @@ export const installFakeWarehouse = ( }), ) } + if ("failure" in rule) { + return Effect.fail( + new WarehouseDriverError({ + reason: "server", + status: 400, + message: rule.failure.message, + type: rule.failure.type, + }), + ) + } const rows = typeof rule.rows === "function" ? rule.rows(sql) : rule.rows return Effect.succeed({ data: rows as ReadonlyArray> }) }, diff --git a/apps/ai/src/mcp/__evals__/run-sql-recovery.test.ts b/apps/ai/src/mcp/__evals__/run-sql-recovery.test.ts new file mode 100644 index 0000000000..2bd89349b3 --- /dev/null +++ b/apps/ai/src/mcp/__evals__/run-sql-recovery.test.ts @@ -0,0 +1,94 @@ +/** + * `run_sql` failures an agent hit in production, replayed against the real tool: the column + * names it fixes by itself and the internal SQL it keeps out of the error. + */ +import { afterAll, beforeAll, describe, expect, it } from "vitest" +import { Schema } from "effect" +import { RunSqlOutput } from "@maple/domain/mcp-outputs" +import { makeEvalRuntime, runToolDirect, markdown, type EvalRuntime } from "./eval-runtime" +import { installFakeWarehouse, restoreWarehouse } from "./fake-warehouse" + +const unknown = (name: string) => ({ + message: `[Error] Unknown expression or function identifier \`${name}\` in scope SELECT ${name} FROM traces`, + type: "UNKNOWN_IDENTIFIER", +}) + +const GATEWAY_WRAPPED = { + message: + "[Error] Unknown expression identifier `Level` in scope SELECT Level FROM (SELECT * FROM (SELECT * FROM " + + "maple_us.alert_checks WHERE OrgId = 'org_x') AS maple_us.alert_checks WHERE OrgId = 'org_x') AS alert_checks " + + "WHERE (OrgId = 'org_x') ) AS maple_raw_sql_limited LIMIT 1001 SETTINGS max_execution_time=10", + type: "UNKNOWN_IDENTIFIER", +} + +let rt: EvalRuntime + +beforeAll(() => { + installFakeWarehouse([ + { match: (sql) => sql.includes("service.name"), failure: unknown("service.name") }, + { match: (sql) => sql.includes("FROM alert_checks"), failure: GATEWAY_WRAPPED }, + { match: (sql) => /\bbogus\b/.test(sql), failure: unknown("bogus") }, + // Every pass fixes one name and the warehouse names the next one. + { match: (sql) => sql.includes("trace_id"), failure: unknown("trace_id") }, + { match: (sql) => sql.includes("span_id"), failure: unknown("span_id") }, + { match: (sql) => sql.includes("span_name"), failure: unknown("span_name") }, + { match: (sql) => sql.includes("status_code"), failure: unknown("status_code") }, + { match: (sql) => sql.includes("ServiceName"), rows: [{ ServiceName: "checkout", c: 3 }] }, + ]) + rt = makeEvalRuntime() +}) + +afterAll(async () => { + restoreWarehouse() + await rt.dispose() +}) + +const WINDOW = { start_time: "2026-09-24 00:00:00", end_time: "2026-09-24 06:00:00" } + +describe("run_sql recovery", () => { + it("reruns a misspelt column with its one right name and says so", async () => { + const result = await runToolDirect(rt, "run_sql", { + ...WINDOW, + sql: "SELECT service.name, count() AS c FROM traces WHERE $__orgFilter AND service.name != '' GROUP BY service.name", + }) + expect(result.isError).not.toBe(true) + const output = Schema.decodeUnknownSync(RunSqlOutput)(result.structuredContent) + expect(output.rows).toEqual([{ ServiceName: "checkout", c: 3 }]) + expect(output.rewrites).toEqual([{ from: "service.name", to: "ServiceName" }]) + expect(markdown(result)).toContain("Ran with `ServiceName` for `service.name`") + }) + + it("says which rewrite the error belongs to when the rerun fails on another name", async () => { + const result = await runToolDirect(rt, "run_sql", { + ...WINDOW, + sql: "SELECT service.name, bogus FROM traces WHERE $__orgFilter", + }) + expect(result.isError).toBe(true) + const text = markdown(result) + expect(text).toContain("`bogus`") + expect(text).toContain("This error is for the query rerun with `ServiceName` for `service.name`.") + }) + + it("stops rewriting after three passes", async () => { + const result = await runToolDirect(rt, "run_sql", { + ...WINDOW, + sql: "SELECT trace_id, span_id, span_name, status_code FROM traces WHERE $__orgFilter", + }) + expect(result.isError).toBe(true) + const text = markdown(result) + expect(text).toContain("`status_code`") + expect(text).toContain("`TraceId` for `trace_id`, `SpanId` for `span_id`, `SpanName` for `span_name`") + }) + + it("keeps the gateway's internal database names out of the error", async () => { + const result = await runToolDirect(rt, "run_sql", { + ...WINDOW, + sql: "SELECT Level FROM alert_checks WHERE $__orgFilter", + }) + const text = markdown(result) + expect(text).toContain("FROM alert_checks") + expect(text).not.toContain("maple_us") + expect(text).not.toContain("maple_raw_sql_limited") + expect(text).not.toContain("org_x") + }) +}) diff --git a/apps/ai/src/mcp/lib/raw-sql-error-text.test.ts b/apps/ai/src/mcp/lib/raw-sql-error-text.test.ts new file mode 100644 index 0000000000..e01367063d --- /dev/null +++ b/apps/ai/src/mcp/lib/raw-sql-error-text.test.ts @@ -0,0 +1,59 @@ +import { describe, expect, it } from "vitest" +import { rawSqlErrorText } from "./raw-sql-error-text" + +describe("rawSqlErrorText", () => { + it("collapses the gateway's per-org subquery back to the table name", () => { + const message = + "[Error] Unknown identifier in scope SELECT x FROM (SELECT * FROM (SELECT * FROM maple_us.traces WHERE " + + "OrgId = 'org_a') AS maple_us.traces WHERE OrgId = 'org_a') AS traces WHERE (OrgId = 'org_a')" + expect(rawSqlErrorText(message)).toBe( + "[Error] Unknown identifier in scope SELECT x FROM traces WHERE $__orgFilter", + ) + }) + + // A preview or staging warehouse uses its own database name. + it("collapses the wrapper under any database name, aliased or as a bare subquery", () => { + const wrap = (db: string, table: string) => + `(SELECT * FROM (SELECT * FROM ${db}.${table} WHERE OrgId = 'org_a') AS ${db}.${table} WHERE OrgId = 'org_a')` + expect( + rawSqlErrorText( + `In scope SELECT l.x FROM ${wrap("pr_1961", "logs")} AS l WHERE (l.OrgId = 'org_a')`, + ), + ).toBe("In scope SELECT l.x FROM logs AS l WHERE $__orgFilter(l)") + expect( + rawSqlErrorText( + "Actual SELECT * FROM (SELECT * FROM pr_1961.traces WHERE OrgId = 'org_a') AS pr_1961.traces WHERE OrgId = 'org_a'.", + ), + ).toBe("Actual traces.") + }) + + it("keeps parentheses balanced around a wrapper with no alias", () => { + expect( + rawSqlErrorText( + "x IN (SELECT * FROM (SELECT * FROM pr_1.traces WHERE OrgId = 'org_a') AS pr_1.traces WHERE OrgId = 'org_a')", + ), + ).toBe("x IN (traces)") + }) + + it("never leaves an org id or a database-qualified table behind", () => { + expect(rawSqlErrorText("FROM pr_1961.traces WHERE OrgId = 'org_3KkfPhr7'")).toBe( + "FROM traces WHERE OrgId = ''", + ) + }) + + it("drops the row-cap wrapper", () => { + expect( + rawSqlErrorText( + "in scope SELECT * FROM ( SELECT n FROM logs ) AS maple_raw_sql_limited LIMIT 1001 SETTINGS max_execution_time=10, max_threads=2", + ), + ).toBe("in scope SELECT n FROM logs") + }) + + it("turns a datasource permission error into a plain refusal", () => { + expect( + rawSqlErrorText( + "Not enough permissions for datasource 'audit_log', token needs DATASOURCES:READ:audit_log scope. curl -X PUT ...", + ), + ).toBe("`audit_log` cannot be read from run_sql.") + }) +}) diff --git a/apps/ai/src/mcp/lib/raw-sql-error-text.ts b/apps/ai/src/mcp/lib/raw-sql-error-text.ts new file mode 100644 index 0000000000..35b326db2b --- /dev/null +++ b/apps/ai/src/mcp/lib/raw-sql-error-text.ts @@ -0,0 +1,33 @@ +import { listWarehouseTables } from "@maple/backend/services/warehouse/warehouse-catalog" + +const DB = "[A-Za-z_]\\w*" + +const INNER = `SELECT \\* FROM \\(SELECT \\* FROM ${DB}\\.(\\w+) WHERE OrgId = '[^']*'\\) AS ${DB}\\.\\w+ WHERE OrgId = '[^']*'` + +/** The managed gateway's per-org subquery around a table, as `(...) AS alias` or bare. */ +const ALIASED_WRAPPER = new RegExp(`\\(${INNER}\\) AS (\\w+)`, "g") +const BARE_WRAPPER = new RegExp(INNER, "g") + +/** + * A warehouse error rewritten in the caller's own terms. The managed gateway wraps every + * table in a per-org subquery under its internal database name and our row cap wraps the + * statement; agents copied both back into their SQL (`FROM maple_us.alert_checks`). + */ +export const rawSqlErrorText = (message: string): string => { + const denied = message.match(/Not enough permissions for datasource '([^']+)'/) + if (denied !== null) return `\`${denied[1]}\` cannot be read from run_sql.` + const tables = listWarehouseTables() + .map((t) => t.name) + .join("|") + return message + .replace(ALIASED_WRAPPER, (_whole, table: string, alias: string) => + alias === table ? table : `${table} AS ${alias}`, + ) + .replace(BARE_WRAPPER, "$1") + .replace(new RegExp(`\\b${DB}\\.(${tables})\\b`, "g"), "$1") + .replace(/(scope )SELECT \* FROM \(\s*/g, "$1") + .replace(/\s*\) AS maple_raw_sql_limited\s+LIMIT \d+(?:\s+SETTINGS [\w=, ]+)?(?:\s+FORMAT \w+)?/g, "") + .replace(/\((\w+)\.OrgId = 'org_\w+'\)/g, "$__orgFilter($1)") + .replace(/\(OrgId = 'org_\w+'\)/g, "$__orgFilter") + .replace(/'org_[A-Za-z0-9]+'/g, "''") +} diff --git a/apps/ai/src/mcp/lib/sql-identifiers.test.ts b/apps/ai/src/mcp/lib/sql-identifiers.test.ts new file mode 100644 index 0000000000..95f5122336 --- /dev/null +++ b/apps/ai/src/mcp/lib/sql-identifiers.test.ts @@ -0,0 +1,185 @@ +import { describe, expect, it } from "vitest" +import { fixUnknownIdentifier } from "./sql-identifiers" + +const unknown = (name: string) => `Unknown expression or function identifier \`${name}\` in scope SELECT` + +describe("fixUnknownIdentifier", () => { + it.each([ + [ + "SELECT service.name FROM traces WHERE $__orgFilter", + "service.name", + "SELECT ServiceName FROM traces WHERE $__orgFilter", + ], + [ + "SELECT Events.Name FROM traces WHERE $__orgFilter", + "Events.Name", + "SELECT EventsName FROM traces WHERE $__orgFilter", + ], + [ + "SELECT `Events.Name` FROM traces WHERE $__orgFilter", + "Events.Name", + "SELECT EventsName FROM traces WHERE $__orgFilter", + ], + [ + "SELECT trace_id FROM logs WHERE $__orgFilter", + "trace_id", + "SELECT TraceId FROM logs WHERE $__orgFilter", + ], + [ + "SELECT ServiceVersion FROM traces WHERE $__orgFilter", + "ServiceVersion", + "SELECT ResourceAttributes['service.version'] FROM traces WHERE $__orgFilter", + ], + [ + "SELECT t.trace_id FROM traces t WHERE $__orgFilter(t)", + "t.trace_id", + "SELECT t.TraceId FROM traces t WHERE $__orgFilter(t)", + ], + ])("rewrites %s", (sql, name, expected) => { + expect(fixUnknownIdentifier(sql, unknown(name))?.sql).toBe(expected) + }) + + it("leaves string literals and comments alone", () => { + const sql = + "SELECT service.name FROM traces WHERE $__orgFilter AND Body LIKE '%service.name%' -- service.name" + expect(fixUnknownIdentifier(sql, unknown("service.name"))?.sql).toBe( + "SELECT ServiceName FROM traces WHERE $__orgFilter AND Body LIKE '%service.name%' -- service.name", + ) + }) + + // Two candidate columns, or none on the tables read: no guess. + it.each([ + ["SELECT Severity FROM logs WHERE $__orgFilter", "Severity"], + ["SELECT CommitSha FROM service_operations_hourly WHERE $__orgFilter", "CommitSha"], + ["SELECT http.response.status_code FROM traces WHERE $__orgFilter", "http.response.status_code"], + ["SELECT Count FROM traces WHERE $__orgFilter", "Count"], + ])("does not guess for %s", (sql, name) => { + expect(fixUnknownIdentifier(sql, unknown(name))).toBeUndefined() + }) + + // The `'` inside the backticks must not open a string literal for the masker. + it("keeps a literal intact next to a backticked name that contains a quote", () => { + const sql = + "SELECT `it's`, count() FROM traces WHERE $__orgFilter AND SpanName = 'service.name' AND service.name = 'a' GROUP BY 1" + expect(fixUnknownIdentifier(sql, unknown("service.name"))?.sql).toBe( + "SELECT `it's`, count() FROM traces WHERE $__orgFilter AND SpanName = 'service.name' AND ServiceName = 'a' GROUP BY 1", + ) + }) + + it("fixes a column of a backtick-quoted table", () => { + expect( + fixUnknownIdentifier( + "SELECT service.name FROM `traces` WHERE $__orgFilter", + unknown("service.name"), + )?.sql, + ).toBe("SELECT ServiceName FROM `traces` WHERE $__orgFilter") + }) + + // Only FROM/JOIN/comma positions name tables: not a literal, not an alias. + it.each([ + [ + "SELECT service.name FROM logs WHERE $__orgFilter AND Body = '`traces`'", + "SELECT ServiceName FROM logs WHERE $__orgFilter AND Body = '`traces`'", + ], + [ + "SELECT service.name FROM logs AS traces WHERE $__orgFilter", + "SELECT ServiceName FROM logs AS traces WHERE $__orgFilter", + ], + ])("fixes %s against the one table it reads", (sql, expected) => { + expect(fixUnknownIdentifier(sql, unknown("service.name"))?.sql).toBe(expected) + }) + + it("fixes a column qualified by the alias of a backtick-quoted table", () => { + expect( + fixUnknownIdentifier( + "SELECT t.trace_id FROM `traces` AS t WHERE $__orgFilter(t)", + unknown("t.trace_id"), + )?.sql, + ).toBe("SELECT t.TraceId FROM `traces` AS t WHERE $__orgFilter(t)") + }) + + it("fixes a qualified column only from the table the qualifier names", () => { + const sql = + "SELECT t.body FROM logs l JOIN traces t ON l.TraceId = t.TraceId WHERE $__orgFilter(l) AND $__orgFilter(t)" + expect(fixUnknownIdentifier(sql, unknown("t.body"))).toBeUndefined() + expect(fixUnknownIdentifier(sql.replace("t.body", "l.body"), unknown("l.body"))?.sql).toContain( + "SELECT l.Body FROM", + ) + }) + + it("never renames a table alias declared without AS", () => { + const sql = "SELECT service FROM traces service WHERE $__orgFilter(service)" + expect(fixUnknownIdentifier(sql, unknown("service"))).toBeUndefined() + }) + + it("never renames a derived table's alias declared without AS", () => { + const sql = "SELECT service FROM (SELECT ServiceName FROM traces WHERE $__orgFilter) service" + expect(fixUnknownIdentifier(sql, unknown("service"))).toBeUndefined() + }) + + it("never renames a quoted, dotted derived-table alias declared without AS", () => { + const sql = + "SELECT service.name FROM (SELECT ServiceName FROM traces WHERE $__orgFilter) `service.name`" + expect(fixUnknownIdentifier(sql, unknown("service.name"))).toBeUndefined() + }) + + it("never renames a quoted, dotted table alias declared without AS", () => { + const sql = "SELECT service.name FROM traces `service.name` WHERE $__orgFilter(`service.name`)" + expect(fixUnknownIdentifier(sql, unknown("service.name"))).toBeUndefined() + }) + + it("never renames the alias of a CTE reference", () => { + const sql = + "WITH s AS (SELECT ServiceName FROM traces) SELECT service FROM s service WHERE $__orgFilter(service)" + expect(fixUnknownIdentifier(sql, unknown("service"))).toBeUndefined() + }) + + it("reads the alias of a table joined after an unaliased one", () => { + const sql = + "SELECT l.body FROM traces JOIN logs l ON traces.TraceId = l.TraceId WHERE $__orgFilter(l)" + expect(fixUnknownIdentifier(sql, unknown("l.body"))?.sql).toContain("SELECT l.Body FROM") + }) + + it("declines a qualified fix when the qualifier names relations in two scopes", () => { + const sql = + "WITH t AS (SELECT ServiceName AS span_name FROM traces WHERE $__orgFilter) SELECT t.span_name FROM t " + + "WHERE t.span_name IN (SELECT t.span_name FROM traces t WHERE $__orgFilter(t))" + expect(fixUnknownIdentifier(sql, unknown("t.span_name"))).toBeUndefined() + }) + + it("ignores a table named in a # comment", () => { + expect( + fixUnknownIdentifier( + "SELECT service.name FROM logs WHERE $__orgFilter #, traces", + unknown("service.name"), + )?.sql, + ).toBe("SELECT ServiceName FROM logs WHERE $__orgFilter #, traces") + }) + + it("fixes a column qualified by a later item of a comma-separated FROM", () => { + const sql = "SELECT t.span_name FROM logs l, traces t WHERE $__orgFilter(l) AND $__orgFilter(t)" + expect(fixUnknownIdentifier(sql, unknown("t.span_name"))?.sql).toContain("SELECT t.SpanName FROM") + }) + + it("never rewrites a backtick-quoted alias the query defines", () => { + const sql = + "SELECT ServiceName AS `service.name` FROM traces WHERE $__orgFilter ORDER BY `service.name`" + expect(fixUnknownIdentifier(sql, unknown("service.name"))).toBeUndefined() + }) + + it("counts every table of a comma-separated FROM list", () => { + const sql = "SELECT service_name FROM traces t, logs l WHERE $__orgFilter(t) AND $__orgFilter(l)" + expect(fixUnknownIdentifier(sql, unknown("service_name"))).toBeUndefined() + }) + + it("does not fix an unqualified name both joined tables have", () => { + const sql = + "SELECT service_name FROM traces t JOIN logs l ON t.TraceId = l.TraceId WHERE $__orgFilter(t)" + expect(fixUnknownIdentifier(sql, unknown("service_name"))).toBeUndefined() + }) + + it("never rewrites a name the query defines itself", () => { + const sql = "SELECT ServiceName AS service FROM traces WHERE $__orgFilter ORDER BY service" + expect(fixUnknownIdentifier(sql, unknown("service"))).toBeUndefined() + }) +}) diff --git a/apps/ai/src/mcp/lib/sql-identifiers.ts b/apps/ai/src/mcp/lib/sql-identifiers.ts new file mode 100644 index 0000000000..c0d889e59f --- /dev/null +++ b/apps/ai/src/mcp/lib/sql-identifiers.ts @@ -0,0 +1,203 @@ +import type { RunSqlRewrite } from "@maple/domain/mcp-outputs" +import { + describeWarehouseTable, + listWarehouseTables, + type TableInfo, +} from "@maple/backend/services/warehouse/warehouse-catalog" + +/** The column an unknown-identifier error names, as written. */ +export const unknownIdentifierName = (message: string): string | undefined => { + const match = message.match( + /identifier [`']([^`']+)[`']|Missing columns: '([^']+)'|no column [`']([^`']+)[`']/i, + ) + return match?.[1] ?? match?.[2] ?? match?.[3] +} + +// String literals (with '' and \' escapes), `--`/`#`/`/* */` comments, and quoted identifiers. +const SQL_TOKEN = + /'(?:[^'\\]|\\[\s\S]|'')*'?|--[^\n]*|#[^\n]*|\/\*[\s\S]*?(?:\*\/|$)|`(?:[^`]|``)*`?|"(?:[^"]|"")*"?/g + +/** + * The SQL with literals and comments blanked and quoted identifiers unquoted, keeping every + * position, so `FROM \`traces\`` reads as a table while `'traces'` does not. + */ +export const identifierText = (sql: string): string => + sql.replace(SQL_TOKEN, (token) => { + if (token[0] !== "`" && token[0] !== '"') return token.replace(/[^\n]/g, " ") + const closed = token.length > 1 && token.endsWith(token[0]) + return ` ${token.slice(1, closed ? -1 : undefined)}${closed ? " " : ""}` + }) + +/** Names of the catalog tables the SQL reads (after FROM, JOIN or a comma), in first-mention order. */ +export const referencedTableNames = (sql: string): ReadonlyArray => { + const text = identifierText(sql) + return listWarehouseTables() + .map((t) => ({ + name: t.name, + at: text.search( + new RegExp(`(?:\\bFROM\\s+|\\bJOIN\\s+|,\\s*)(?:[A-Za-z_]\\w*\\.)?${t.name}\\b`, "i"), + ), + })) + .filter((t) => t.at >= 0) + .sort((a, b) => a.at - b.at) + .map((t) => t.name) +} + +/** The catalog tables the SQL reads, with their columns, in first-mention order. */ +export const referencedTables = (sql: string): ReadonlyArray => + referencedTableNames(sql).flatMap((name) => { + const info = describeWarehouseTable(name) + return info === null ? [] : [info] + }) + +/** `service.name`, `trace_id` and `TRACEID` all fold to `traceid`-style keys. */ +const fold = (name: string) => name.replace(/[._]/g, "").toLowerCase() + +/** Names agents write for a column or map key that have only one meaning, by folded name. */ +const KNOWN_ALIASES: ReadonlyMap = new Map([ + ["service", { column: "ServiceName", expr: "ServiceName" }], + ["logbody", { column: "Body", expr: "Body" }], + ["spanstatusmessage", { column: "StatusMessage", expr: "StatusMessage" }], + ["statusdescription", { column: "StatusMessage", expr: "StatusMessage" }], + ["serviceversion", { column: "ResourceAttributes", expr: "ResourceAttributes['service.version']" }], + ["servicenamespace", { column: "ResourceAttributes", expr: "ResourceAttributes['service.namespace']" }], + [ + "serviceinstanceid", + { column: "ResourceAttributes", expr: "ResourceAttributes['service.instance.id']" }, + ], + ["errortype", { column: "SpanAttributes", expr: "SpanAttributes['error.type']" }], +]) + +/** + * The real expression an unknown identifier was meant to be, only when exactly one fits + * the tables the query reads: a spelling of a real column (`service.name` → `ServiceName`, + * `Events.Name` → `EventsName`, `trace_id` → `TraceId`) or a known alias (`ServiceVersion`). + * In a join the column must belong to one table, or the fix would be ambiguous. + */ +const identifierFix = ( + written: string, + tables: ReadonlyArray, + qualifiers: ReadonlyMap, +): string | undefined => { + const owners = (column: string) => tables.filter((t) => t.columns.some((c) => c.name === column)).length + if (owners(written) > 0) return undefined + const byFold = (name: string, unqualified: boolean) => { + const columns = new Set(tables.flatMap((t) => t.columns.map((c) => c.name))) + const hits = [...columns].filter((c) => fold(c) === fold(name)) + return hits.length === 1 && (!unqualified || owners(hits[0]) === 1) ? hits[0] : undefined + } + const whole = byFold(written, true) + if (whole !== undefined) return whole + const alias = KNOWN_ALIASES.get(fold(written)) + if (alias !== undefined && owners(alias.column) === 1) return alias.expr + // `t.trace_id`: keep the qualifier and fix the column from the table it names. + const dot = written.lastIndexOf(".") + const qualified = + dot <= 0 ? undefined : tables.find((t) => t.name === qualifiers.get(written.slice(0, dot))) + if (qualified === undefined) return undefined + const hits = qualified.columns.filter((c) => fold(c.name) === fold(written.slice(dot + 1))) + return hits.length === 1 ? `${written.slice(0, dot + 1)}${hits[0].name}` : undefined +} + +const escapeRegExp = (text: string) => text.replace(/[.*+?^${}()|[\]\\]/g, "\\$&") + +/** + * Every use of `written` as an identifier, bare or quoted, replaced by `replacement`. + * Literals and comments are left alone, as are longer names it is part of and function calls. + */ +const replaceIdentifier = (sql: string, written: string, replacement: string): string => { + const pattern = new RegExp(`(? { + const quote = sql[m.index - 1] + const quoted = (quote === "`" || quote === '"') && sql[m.index + written.length] === quote + return quoted + ? { start: m.index - 1, end: m.index + written.length + 1 } + : { start: m.index, end: m.index + written.length } + }) + return spans.reduceRight( + (text, { start, end }) => `${text.slice(0, start)}${replacement}${text.slice(end)}`, + sql, + ) +} + +/** SQL words that are never an alias, wherever an alias could stand. */ +const KEYWORDS = new Set( + ( + "and or not in is like ilike between as on using where prewhere having qualify group order limit " + + "offset by asc desc nulls first last union except intersect settings format window over final " + + "sample array join left right inner outer full cross semi anti any all asof global then else end " + + "when from with into values interval" + ).split(" "), +) + +/** A relation after FROM, JOIN or a comma, and its alias unless that word is a keyword. */ +const RELATION = new RegExp( + `(?:\\bFROM\\s+|\\bJOIN\\s+|,\\s*)([A-Za-z_][\\w.]*)(?:\\s+(?:AS\\s+)?(?!(?:${[...KEYWORDS].join("|")})\\b)([A-Za-z_][\\w.]*))?`, + "gi", +) + +const relations = (sql: string) => + [...identifierText(sql).matchAll(RELATION)].map((m) => ({ name: m[1], alias: m[2] })) + +/** + * Names the SQL defines itself, which are never column typos: `AS x`, `x AS (`, a relation's + * alias (`FROM s service`, CTE or table), and an alias after `)` (`FROM (...) t`, `count() c`). + */ +export const definedNames = (sql: string): ReadonlySet => { + const text = identifierText(sql) + return new Set([ + ...[...text.matchAll(/\bAS\s+([A-Za-z_][\w.]*)|\b([A-Za-z_]\w*)\s+AS\s*\(/gi)].map( + (m) => m[1] ?? m[2], + ), + ...[...text.matchAll(/\)\s*([A-Za-z_][\w.]*)/g)] + .map((m) => m[1]) + .filter((name) => !KEYWORDS.has(name.toLowerCase())), + ...relations(sql).flatMap((r) => (r.alias === undefined ? [] : [r.alias])), + ]) +} + +/** The catalog table each qualifier names: `FROM traces t`, `JOIN logs AS l`, `FROM a, b`. */ +const tableQualifiers = (sql: string): ReadonlyMap => { + const catalog = new Set(listWarehouseTables().map((t) => t.name)) + return new Map( + relations(sql).flatMap(({ name, alias }) => { + const table = name.split(".").at(-1) ?? name + if (!catalog.has(table)) return [] + return [name, table, ...(alias === undefined ? [] : [alias])].map((q) => [q, table] as const) + }), + ) +} + +/** + * The qualifier of `written` names more than one relation (a CTE in one scope, a table in + * another), so a rewrite of every `q.col` would also change the scope where it was valid. + */ +const qualifierReused = (sql: string, written: string): boolean => { + const dot = written.lastIndexOf(".") + if (dot <= 0) return false + const qualifier = written.slice(0, dot).toLowerCase() + const binds = relations(sql).filter( + ({ name, alias }) => + alias?.toLowerCase() === qualifier || + (name.split(".").at(-1) ?? name).toLowerCase() === qualifier, + ) + return binds.length > 1 +} + +export type IdentifierRewrite = typeof RunSqlRewrite.Type + +/** The SQL with the identifier an error named replaced, or undefined when no single fix fits. */ +export const fixUnknownIdentifier = ( + sql: string, + message: string, +): { readonly sql: string; readonly rewrite: IdentifierRewrite } | undefined => { + const written = unknownIdentifierName(message) + if (written !== undefined && qualifierReused(sql, written)) return undefined + const qualifiers = tableQualifiers(sql) + // A table alias is a name the query defines, `AS` or not; renaming it breaks the FROM. + if (written === undefined || definedNames(sql).has(written) || qualifiers.has(written)) return undefined + const to = identifierFix(written, referencedTables(sql), qualifiers) + if (to === undefined) return undefined + const fixed = replaceIdentifier(sql, written, to) + return fixed === sql ? undefined : { sql: fixed, rewrite: { from: written, to } } +} diff --git a/apps/ai/src/mcp/tools/run-sql.ts b/apps/ai/src/mcp/tools/run-sql.ts index e9bc0a9e58..e0a1fe6ec6 100644 --- a/apps/ai/src/mcp/tools/run-sql.ts +++ b/apps/ai/src/mcp/tools/run-sql.ts @@ -7,12 +7,18 @@ import { truncate } from "../lib/format" import { toMcpQueryError } from "../lib/map-warehouse-error" import * as P from "../lib/params" import { MAX_RAW_SQL_RESULT_ROWS } from "@maple/domain/raw-sql" +import { rawSqlErrorText } from "../lib/raw-sql-error-text" import { doc } from "../lib/tool-doc" import { - describeWarehouseTable, listWarehouseTables, suggestWarehouseTables, } from "@maple/backend/services/warehouse/warehouse-catalog" +import { + fixUnknownIdentifier, + type IdentifierRewrite, + referencedTables, + unknownIdentifierName, +} from "../lib/sql-identifiers" // Rows returned to the model are capped so a wide/long result doesn't blow the // context. The full count is always reported via meta.rowCount. @@ -77,14 +83,6 @@ export const withTableListOnUnknownTable = ( const UNKNOWN_COLUMN = /unknown (?:expression (?:or function )?)?identifier|missing columns|there(?:'s| is) no column|no such column/i -/** The column an unknown-column error names, as written. */ -const missingColumnName = (message: string): string | undefined => { - const match = message.match( - /identifier [`']([^`']+)[`']|Missing columns: '([^']+)'|no column [`']([^`']+)[`']/i, - ) - return match?.[1] ?? match?.[2] ?? match?.[3] -} - const NO_FROM_MESSAGE = "Raw SQL must read from a Maple table: `$__orgFilter` filters a table's `OrgId`, and this query has no FROM. " + "Select constants from a table instead: `SELECT now() FROM traces WHERE $__orgFilter LIMIT 1`." @@ -92,14 +90,6 @@ const NO_FROM_MESSAGE = /** Cap on how many tables one enrichment describes. */ const MAX_DESCRIBED_TABLES = 3 -/** Catalog tables the submitted SQL actually names, in first-mention order. */ -const referencedTables = (sql: string): ReadonlyArray => - listWarehouseTables() - .map((t) => ({ name: t.name, at: sql.search(new RegExp(`\\b${t.name}\\b`)) })) - .filter((t) => t.at >= 0) - .sort((a, b) => a.at - b.at) - .map((t) => t.name) - /** * The column counterpart of `withTableListOnUnknownTable`, and it exists for the * same reason: the warehouse names the column the agent invented but never the @@ -114,7 +104,7 @@ export const withColumnListOnUnknownColumn = if (error._tag !== "@maple/mcp/errors/McpQueryError" || !UNKNOWN_COLUMN.test(error.message)) return error - const written = missingColumnName(error.message) + const written = unknownIdentifierName(error.message) const column = written?.split(".").at(-1) // A table-less SELECT fails on the org filter itself; say that, not "unknown OrgId". if (column === "OrgId" && !/\bFROM\b/i.test(sql)) { @@ -127,10 +117,7 @@ export const withColumnListOnUnknownColumn = // Only the tables the query names; a schema dump of all 38 would bury the // message it is attached to. - const described = referencedTables(sql) - .slice(0, MAX_DESCRIBED_TABLES) - .map((name) => describeWarehouseTable(name)) - .filter((info) => info !== null) + const described = referencedTables(sql).slice(0, MAX_DESCRIBED_TABLES) if (described.length === 0) return error // A column Maple's schema has but the warehouse lacks only happens on a cluster @@ -254,8 +241,9 @@ const hasUnsafeInteger = (rows: ReadonlyArray>): boo const BUDGET_ROLLUP_HINT = "Pre-aggregated tables answer most per-service questions without scanning raw spans: service_overview_minutely / " + "service_overview_hourly (requests, errors, latency per service), service_operations_minutely / service_operations_hourly " + - "(per span name), traces_aggregates_hourly, logs_aggregates_hourly, error_events_by_time, trace_list_mv (one row per trace). " + - "Otherwise narrow start_time/end_time or filter on the sorting key (ServiceName first on traces and logs)." + "(per span name), service_endpoint_details_hourly (per HTTP route and status), ai_trace_index (agent spans), " + + "traces_aggregates_hourly, logs_aggregates_hourly, error_events_by_time, trace_list_mv (one row per trace). " + + "Otherwise narrow start_time/end_time or filter on ServiceName with a Timestamp window." /** Point a query that ran out of time or memory at the rollups that answer it cheaply. */ export const withRollupHintOnBudget = ( @@ -282,6 +270,36 @@ export const callerSqlAsInvalidInput = ( ? new McpInvalidInputError({ message: error.message, parameter: "sql" }) : error +/** Reruns with a corrected column name before the error goes back to the agent. */ +const MAX_IDENTIFIER_FIXES = 3 + +/** Says which query a failure belongs to when the server already rewrote the agent's SQL. */ +const withRewriteNote = (message: string, rewrites: ReadonlyArray): string => + rewrites.length === 0 + ? message + : `${message}\n\nThis error is for the query rerun with ${rewrites.map((r) => `\`${r.to}\` for \`${r.from}\``).join(", ")}.` + +/** The warehouse message in the caller's terms: no gateway database names or row-cap wrapper. */ +const withCallerErrorText = + (rewrites: ReadonlyArray) => + (error: McpQueryError | McpQueryBudgetError): McpQueryError | McpQueryBudgetError => + error._tag === "@maple/mcp/errors/McpQueryError" + ? new McpQueryError({ + message: withRewriteNote(rawSqlErrorText(error.message), rewrites), + pipeName: error.pipeName, + cause: error.cause, + }) + : error + +type RawSqlRun = ReturnType + +interface FixFailure { + readonly error: Effect.Error + /** The statement that failed, after any rewrites. */ + readonly sql: string + readonly rewrites: ReadonlyArray +} + export function registerRunSqlTool(server: McpToolRegistrar) { server.define({ name: "run_sql", @@ -296,18 +314,42 @@ export function registerRunSqlTool(server: McpToolRegistrar) { const { st, et } = yield* WINDOW.resolve(params, "run_sql") const granularitySeconds = params.granularity_seconds ?? autoBucketSeconds(st, et) - const { rows, columns, rowCount, expandedSql } = yield* runRawSql({ - tenant, - sql: params.sql, - startTime: st, - endTime: et, - granularitySeconds, - }).pipe( - Effect.mapError((error) => + const execute = (sql: string) => + runRawSql({ tenant, sql, startTime: st, endTime: et, granularitySeconds }) + + // A misspelt column with exactly one right answer is fixed and rerun rather than + // costing the agent a round trip; each pass fixes the name the warehouse reported. + const runFixing = ( + sql: string, + rewrites: ReadonlyArray, + ): Effect.Effect< + Effect.Success & { readonly rewrites: ReadonlyArray }, + FixFailure, + Effect.Services + > => + execute(sql).pipe( + Effect.map((result) => ({ ...result, rewrites })), + Effect.catch((error) => { + const fix = + error._tag === "@maple/http/errors/RawSqlValidationError" || + rewrites.length >= MAX_IDENTIFIER_FIXES + ? undefined + : fixUnknownIdentifier(sql, error.message) + return fix === undefined + ? Effect.fail({ error, sql, rewrites }) + : runFixing(fix.sql, [...rewrites, fix.rewrite]) + }), + ) + + const { rows, columns, rowCount, expandedSql, rewrites } = yield* runFixing(params.sql, []).pipe( + Effect.mapError(({ error, sql, rewrites }) => // Macro/safety failures (a missing `$__orgFilter` above all) are caller-fixable. error._tag === "@maple/http/errors/RawSqlValidationError" ? new McpInvalidInputError({ - message: `SQL rejected (${error.code}): ${error.message}`, + message: withRewriteNote( + `SQL rejected (${error.code}): ${error.message}`, + rewrites, + ), parameter: "sql", example: RUN_SQL_EXAMPLE, }) @@ -317,10 +359,10 @@ export function registerRunSqlTool(server: McpToolRegistrar) { // plus the real table (or column) names when the one it named does not exist: // the warehouse message never says what they are. withRollupHintOnBudget( - withColumnListOnUnknownColumn(params.sql)( + withColumnListOnUnknownColumn(sql)( withTableListOnUnknownTable( - toMcpQueryError("run_sql")(error), - params.sql, + withCallerErrorText(rewrites)(toMcpQueryError("run_sql")(error)), + sql, ), ), ), @@ -343,51 +385,61 @@ export function registerRunSqlTool(server: McpToolRegistrar) { timeRange: { start: st, end: et }, maxCellChars, windowApplied: usesWindow(expandedSql, st, et), + ...(rewrites.length === 0 ? undefined : { rewrites: [...rewrites] }), ...(params.show_sql === true ? { expandedSql } : undefined), } }), - render: (output) => ({ - title: "SQL result", - scope: [ - [ - "Time range", - output.windowApplied !== false - ? `${output.timeRange.start} to ${output.timeRange.end}` - : "set by the SQL (no $__timeFilter/$__startTime/$__endTime macro)", + render: (output) => { + const notices = [ + ...(output.rewrites ?? []).map( + (r) => + `Ran with \`${r.to}\` for \`${r.from}\`, which is not a column here; use \`${r.to}\` directly.`, + ), + ...(hasUnsafeInteger(output.rows) + ? [ + "Some integers exceed 2^53 and lost their last digits in transit; select them as toString(col) for exact values.", + ] + : []), + ] + return { + title: "SQL result", + scope: [ + [ + "Time range", + output.windowApplied !== false + ? `${output.timeRange.start} to ${output.timeRange.end}` + : "set by the SQL (no $__timeFilter/$__startTime/$__endTime macro)", + ], + ["Rows", String(output.rowCount)], + ["Columns", String(output.columns.length)], ], - ["Rows", String(output.rowCount)], - ["Columns", String(output.columns.length)], - ], - ...(output.rowCount === 0 ? { empty: { message: "No rows returned." } } : undefined), - blocks: [ - ...(output.rowCount === 0 - ? [] - : [ - doc.table( - [...output.columns], - output.rows.map((row) => - output.columns.map((col) => cellToString(row[col], output.maxCellChars)), + ...(output.rowCount === 0 ? { empty: { message: "No rows returned." } } : undefined), + blocks: [ + ...(output.rowCount === 0 + ? [] + : [ + doc.table( + [...output.columns], + output.rows.map((row) => + output.columns.map((col) => + cellToString(row[col], output.maxCellChars), + ), + ), ), - ), - ]), - ...(output.expandedSql === undefined ? [] : [doc.code("sql", output.expandedSql)]), - ], - ...(hasUnsafeInteger(output.rows) - ? { - notices: [ - "Some integers exceed 2^53 and lost their last digits in transit; select them as toString(col) for exact values.", - ], - } - : undefined), - ...(output.truncated - ? { - truncation: { - shown: output.rows.length, - total: output.rowCount, - noun: "rows (refine with LIMIT or filters)", - }, - } - : undefined), - }), + ]), + ...(output.expandedSql === undefined ? [] : [doc.code("sql", output.expandedSql)]), + ], + ...(notices.length === 0 ? undefined : { notices }), + ...(output.truncated + ? { + truncation: { + shown: output.rows.length, + total: output.rowCount, + noun: "rows (refine with LIMIT or filters)", + }, + } + : undefined), + } + }, }) } diff --git a/packages/backend/src/services/warehouse/warehouse-catalog.ts b/packages/backend/src/services/warehouse/warehouse-catalog.ts index a67059e75b..bba5b787ca 100644 --- a/packages/backend/src/services/warehouse/warehouse-catalog.ts +++ b/packages/backend/src/services/warehouse/warehouse-catalog.ts @@ -23,8 +23,7 @@ const TABLE_NOTES: Record> = { "`SeverityText` casing varies by SDK ('Error' from Effect services, 'ERROR' from OTel SDKs), so filter on the OTel level number instead: `SeverityNumber BETWEEN 17 AND 20` is ERROR, 13-16 WARN, 9-12 INFO, 5-8 DEBUG, 1-4 TRACE, 21-24 FATAL. Where you must use text, use `upper(SeverityText) = 'ERROR'`.", "`SeverityNumber` follows OTel: 1-4 Trace, 5-8 Debug, 9-12 Info, 13-16 Warn, 17-20 Error, 21-24 Fatal.", "`ResourceAttributes` and `LogAttributes` are `Map(LowCardinality(String), String)` — access with `LogAttributes['key']`; missing keys return '' (empty string), not NULL.", - "Use `TimestampTime` (DateTime) for `$__timeFilter(TimestampTime)` if you want sort-key-prefix-friendly filtering; `Timestamp` is DateTime64 (nanosecond precision).", - "Sorting key: `(OrgId, ServiceName, TimestampTime, Timestamp)` — adding `ServiceName = '…'` to the WHERE clause speeds queries dramatically.", + "Sorting key: `(OrgId, toStartOfFiveMinutes(Timestamp), ServiceName, Timestamp)`: filter on a `Timestamp` window and `ServiceName = '…'`. Partitions are daily on `TimestampTime` (DateTime), so `$__timeFilter(TimestampTime)` also skips whole days. `Timestamp` is DateTime64 (nanosecond precision).", ], traces: [ "`StatusCode` values are Title Case: 'Ok', 'Error', 'Unset'. Filter with `StatusCode = 'Error'` (NOT `'ERROR'`).", @@ -34,7 +33,7 @@ const TABLE_NOTES: Record> = { "`SampleRate` defaults to 1.0; multiply counts by `SampleRate` for unbiased throughput estimates.", "Sorting key: `(OrgId, ServiceName, toStartOfFiveMinutes(Timestamp), TraceId)` — filter on `ServiceName` and a `Timestamp` window first.", "To read whole traces by id, seek them through `trace_index`: `(ServiceName, toStartOfFiveMinutes(Timestamp), TraceId) IN (SELECT ServiceName, Bucket, TraceId FROM trace_index WHERE OrgId = … AND TraceId IN (…))`. A bare `TraceId = …` probes every daily partition.", - "For service-level metrics (per-service throughput, latency), prefer `service_overview_spans` — it's pre-filtered to entry-point spans and ~10× smaller.", + "Before scanning raw spans, check the rollups: per-service requests, errors and latency are in `service_overview_minutely` / `service_overview_hourly`, per span name in `service_operations_minutely` / `service_operations_hourly`, per HTTP route and status in `service_endpoint_details_hourly`, agent spans in `ai_trace_index`. Raw `traces` reads of the attribute maps run out of time past about a day without a `ServiceName` filter.", ], ai_trace_index: [ "GenAI agent spans ONLY (every row carries a non-empty `VendorId`), with the `maple_ai.*` identity pre-extracted to plain columns. ALWAYS prefer this over `traces` + `mapContains(SpanAttributes, 'maple_ai.…')` for finding agent traces/sessions — the raw-traces scan reads the full attribute Map per span and times out on day-plus windows.", @@ -67,10 +66,25 @@ const TABLE_NOTES: Record> = { "Keeps 365 days. Sorting key: `(OrgId, Timestamp, TraceId)`.", ], service_endpoint_details_hourly: [ + "`HttpStatusCode` is a String: compare it as text (`HttpStatusCode = '500'`) or convert it (`toUInt16OrZero(HttpStatusCode) >= 500`).", "HTTP Server spans with a response status, one row per (service, hour, env, normalized `SpanName` = `METHOD /route`, `HttpStatusCode`, `ErrorType` = `error.type`). ALWAYS prefer this over `traces` + `SpanAttributes['http.response.status_code']` for per-endpoint status, error-type, latency, body-size and user-agent breakdowns.", "Aggregate state columns: merge `DurationQuantiles` with `quantilesTDigestMerge(0.5, 0.95, 0.99)` (NANOSECONDS), `BodySizeQuantiles` with `quantilesTDigestIfMerge(0.5, 0.95, 0.99)` (bytes), `UserAgents` with `approx_top_kMerge(10)`, and `SlowestTraceId`/`SlowestTimestamp`/`SlowestBodySize` with `argMaxMerge` (`SlowestBodySize` -1 = no body size). Sum `SpanCount`/`EstimatedSpanCount`/`BodyCount`.", "Keeps 90 days at hour grain. Sorting key: `(OrgId, ServiceName, Hour, DeploymentEnv, SpanName, HttpStatusCode, ErrorType)`.", ], + service_operations_hourly: [ + "One row per (service, env, hour, `SpanName`). Sum `SpanCount` / `ErrorCount`; p50 / p95 in NANOSECONDS are `quantilesTDigestMerge(0.5, 0.95)(DurationQuantiles)[1]` / `[2]`.", + "No `CommitSha`, `ServiceNamespace`, `SpanKind` or attributes here: per-commit numbers are in `service_overview_hourly`.", + ], + service_operations_minutely: [ + "`service_operations_hourly` at minute grain (`Minute`), 90-day TTL. Same columns and the same `quantilesTDigestMerge(0.5, 0.95)(DurationQuantiles)` merge; no `CommitSha`.", + ], + service_overview_hourly: [ + "One row per (service, hour, env, namespace, `CommitSha`) of entry-point spans. Sum `SpanCount` / `ErrorCount`; p50 / p95 / p99 in NANOSECONDS are `quantilesTDigestMerge(0.5, 0.95, 0.99)(DurationQuantiles)[1]` / `[2]` / `[3]`.", + "No `SpanName`: per-operation numbers are in `service_operations_hourly`. `CommitSha` is `vcs.ref.head.revision`.", + ], + service_overview_minutely: [ + "`service_overview_hourly` at minute grain (`Minute`), 90-day TTL, with the same columns and `quantilesTDigestMerge(0.5, 0.95, 0.99)(DurationQuantiles)` merge.", + ], service_overview_spans: [ "Pre-materialized projection of entry-point spans only (Server/Consumer kinds + root spans). Use for per-service request count, error rate, p50/p95/p99 latency.", "`Duration` is NANOSECONDS — divide by 1e6 for ms.", diff --git a/packages/domain/src/mcp-outputs/services.ts b/packages/domain/src/mcp-outputs/services.ts index a61178aae3..ce52f4dbbb 100644 --- a/packages/domain/src/mcp-outputs/services.ts +++ b/packages/domain/src/mcp-outputs/services.ts @@ -274,6 +274,9 @@ export const QueryDataOutput = Schema.Struct({ // run_sql +/** A column name the warehouse rejected and the one `run_sql` reran with. */ +export const RunSqlRewrite = Schema.Struct({ from: Schema.String, to: Schema.String }) + export const RunSqlOutput = Schema.Struct({ rowCount: Schema.Number, columns: Schema.Array(Schema.String), @@ -288,6 +291,8 @@ export const RunSqlOutput = Schema.Struct({ windowApplied: Schema.optionalKey(Schema.Boolean), /** The fully macro-expanded SQL that was executed, only when asked for (`show_sql`). */ expandedSql: Schema.optionalKey(Schema.String), + /** Column names the warehouse rejected and the server replaced before rerunning. */ + rewrites: Schema.optionalKey(Schema.Array(RunSqlRewrite)), }) // describe_warehouse_tables diff --git a/packages/domain/src/raw-sql.test.ts b/packages/domain/src/raw-sql.test.ts index ae0085517d..191b6a75e2 100644 --- a/packages/domain/src/raw-sql.test.ts +++ b/packages/domain/src/raw-sql.test.ts @@ -277,6 +277,9 @@ describe("rawSqlIssue org filter placement", () => { "WITH f AS (SELECT 'org_abc' AS OrgId) SELECT * FROM traces t JOIN f ON 1 = 1 WHERE $__orgFilter(f)", "SELECT * FROM traces t JOIN numbers(1) n ON 1 = 1 WHERE $__orgFilter(n)", "SELECT * FROM traces t WHERE $__orgFilter(x)", + // An unaliased first table must not let a joined subquery's alias bind. + "SELECT * FROM traces JOIN (SELECT 'org_abc' AS OrgId) f ON 1 = 1 WHERE $__orgFilter(f)", + "SELECT * FROM traces CROSS JOIN (SELECT 'org_abc' AS OrgId) f WHERE $__orgFilter(f)", ])("rejects an alias that is not a table read in FROM/JOIN: %s", (sql) => { expect(rawSqlIssue(sql)?.message).toContain("must name a table read in FROM or JOIN") }) @@ -286,6 +289,8 @@ describe("rawSqlIssue org filter placement", () => { "SELECT 1 FROM traces WHERE $__orgFilter(traces)", "SELECT 1 FROM maple.traces t FINAL WHERE $__orgFilter(t)", "SELECT 1 FROM traces t LEFT JOIN logs AS l ON t.TraceId = l.TraceId WHERE $__orgFilter(t) AND $__orgFilter(l)", + "SELECT 1 FROM traces JOIN logs l ON traces.TraceId = l.TraceId WHERE $__orgFilter(traces) AND $__orgFilter(l)", + "SELECT 1 FROM traces FINAL JOIN logs l ON traces.TraceId = l.TraceId WHERE $__orgFilter(traces) AND $__orgFilter(l)", ])("accepts an alias bound to a table: %s", (sql) => { expect(rawSqlIssue(sql)).toBeNull() }) diff --git a/packages/domain/src/raw-sql.ts b/packages/domain/src/raw-sql.ts index cd426090b7..c72b5a6b4c 100644 --- a/packages/domain/src/raw-sql.ts +++ b/packages/domain/src/raw-sql.ts @@ -1,6 +1,9 @@ import { Schema } from "effect" import { maskLiteralsAndComments } from "@maple-dev/effect-orm/sql" +/** Blanks literals and comments, keeping positions, so text searches only see SQL. */ +export { maskLiteralsAndComments } + // The static half of raw-SQL validation: everything that can be decided from the // query text alone, without an org, a time window, or a granularity. // @@ -537,6 +540,9 @@ const orgFilterEscapes = (masked: string, start: number, end: number): boolean = } const IDENT = "[A-Za-z_][A-Za-z0-9_]*" +// A clause word after a table is never its alias: `FROM traces JOIN logs l` must still bind `l`. +const NOT_ALIAS = + "(?!(?:WHERE|PREWHERE|FINAL|SAMPLE|ARRAY|GLOBAL|ANY|ALL|INNER|LEFT|RIGHT|FULL|CROSS|OUTER|SEMI|ANTI|ASOF|JOIN|ON|USING|GROUP|ORDER|LIMIT|UNION|SETTINGS|FORMAT)\\b)" const escapeRe = (s: string) => s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&") /** @@ -552,12 +558,12 @@ const aliasBindsTable = (masked: string, alias: string): boolean => { ) const a = escapeRe(alias) const binding = new RegExp( - `\\b(?:FROM|JOIN)\\s+(?:${IDENT}\\.)?(${IDENT})(?:\\s+(?:AS\\s+)?(${IDENT}))?(?=[\\s,)]|$)`, + `\\b(?:FROM|JOIN)\\s+(?:${IDENT}\\.)?(${IDENT})(?:\\s+(?:AS\\s+)?${NOT_ALIAS}(${IDENT}))?(?=[\\s,)]|$)`, "gi", ) // `FROM a x, b y` binds every comma-separated table; a subquery or table function item ends the run. const commaItem = new RegExp( - `\\s*,\\s*(?:${IDENT}\\.)?(${IDENT})(?:\\s+(?:AS\\s+)?(${IDENT}))?(?=[\\s,)]|$)`, + `\\s*,\\s*(?:${IDENT}\\.)?(${IDENT})(?:\\s+(?:AS\\s+)?${NOT_ALIAS}(${IDENT}))?(?=[\\s,)]|$)`, "iy", ) const bindings = [...masked.matchAll(binding)].flatMap((m) => { @@ -570,13 +576,7 @@ const aliasBindsTable = (masked: string, alias: string): boolean => { return items }) for (const [table, bound] of bindings) { - const names = - bound === undefined || - /^(?:WHERE|PREWHERE|FINAL|SAMPLE|ARRAY|GLOBAL|ANY|ALL|INNER|LEFT|RIGHT|FULL|CROSS|JOIN|ON|USING|GROUP|ORDER|LIMIT|UNION|SETTINGS|FORMAT)$/i.test( - bound, - ) - ? [table] - : [table, bound] + const names = bound === undefined ? [table] : [table, bound] if (names.some((n) => new RegExp(`^${a}$`, "i").test(n)) && !ctes.has(table.toLowerCase())) { return true }