From 9cea5f8f297310e4a108098692e8f08e2d2c4508 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 01:19:55 +0100 Subject: [PATCH 01/13] fix(ai): run_sql reruns misspelt column names and keeps internal SQL out of errors An unknown identifier with exactly one right name on the tables the query reads (service.name, Events.Name, trace_id, ServiceVersion) is replaced and the query rerun, up to three passes; the output reports each rewrite. Warehouse errors drop the gateway's maple_us/maple_eu subquery wrapper, the row-cap wrapper and the org id. Catalog notes: correct the logs sorting key, add rollup merge idioms and dimensions, HttpStatusCode is a String. --- apps/ai/src/mcp/__evals__/fake-warehouse.ts | 25 ++- .../mcp/__evals__/run-sql-recovery.test.ts | 94 ++++++++ .../ai/src/mcp/lib/raw-sql-error-text.test.ts | 29 +++ apps/ai/src/mcp/lib/raw-sql-error-text.ts | 23 ++ apps/ai/src/mcp/lib/sql-identifiers.test.ts | 79 +++++++ apps/ai/src/mcp/lib/sql-identifiers.ts | 134 ++++++++++++ apps/ai/src/mcp/tools/run-sql.ts | 206 +++++++++++------- .../services/warehouse/warehouse-catalog.ts | 20 +- packages/domain/src/mcp-outputs/services.ts | 5 + packages/domain/src/raw-sql.ts | 3 + 10 files changed, 534 insertions(+), 84 deletions(-) create mode 100644 apps/ai/src/mcp/__evals__/run-sql-recovery.test.ts create mode 100644 apps/ai/src/mcp/lib/raw-sql-error-text.test.ts create mode 100644 apps/ai/src/mcp/lib/raw-sql-error-text.ts create mode 100644 apps/ai/src/mcp/lib/sql-identifiers.test.ts create mode 100644 apps/ai/src/mcp/lib/sql-identifiers.ts 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..dc0fee2815 --- /dev/null +++ b/apps/ai/src/mcp/lib/raw-sql-error-text.test.ts @@ -0,0 +1,29 @@ +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", + ) + }) + + 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..9e44df06af --- /dev/null +++ b/apps/ai/src/mcp/lib/raw-sql-error-text.ts @@ -0,0 +1,23 @@ +/** + * 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.` + return message + .replace( + /\(SELECT \* FROM \(SELECT \* FROM maple_[a-z]+\.(\w+) WHERE OrgId = '[^']*'\) AS maple_[a-z]+\.\w+ WHERE OrgId = '[^']*'\) AS \w+/g, + "$1", + ) + .replace( + /SELECT \* FROM \(SELECT \* FROM maple_[a-z]+\.(\w+) WHERE OrgId = '[^']*'\) AS maple_[a-z]+\.\w+ WHERE OrgId = '[^']*'/g, + "$1", + ) + .replace(/\bmaple_(?:us|eu)\./g, "") + .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") +} 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..c6461d7cba --- /dev/null +++ b/apps/ai/src/mcp/lib/sql-identifiers.test.ts @@ -0,0 +1,79 @@ +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("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..ae7aa0f1c1 --- /dev/null +++ b/apps/ai/src/mcp/lib/sql-identifiers.ts @@ -0,0 +1,134 @@ +import type { RunSqlRewrite } from "@maple/domain/mcp-outputs" +import { maskLiteralsAndComments } from "@maple/domain/raw-sql" +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] +} + +/** Names of the catalog tables the SQL reads, in first-mention order. */ +export const referencedTableNames = (sql: string): ReadonlyArray => { + const masked = maskLiteralsAndComments(sql) + return listWarehouseTables() + .map((t) => ({ name: t.name, at: masked.search(new RegExp(`\\b${t.name}\\b`)) })) + .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: ReadonlySet, +): 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, fix the column, but only for a table or alias in scope. + const dot = written.lastIndexOf(".") + if (dot <= 0 || !qualifiers.has(written.slice(0, dot))) return undefined + const column = byFold(written.slice(dot + 1), false) + return column === undefined ? undefined : `${written.slice(0, dot + 1)}${column}` +} + +const escapeRegExp = (text: string) => text.replace(/[.*+?^${}()|[\]\\]/g, "\\$&") + +/** + * Every use of `written` as an identifier, bare or backticked, 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 => { + // Unquote only this name's backticks (same length) so other quoting still masks as before. + const quoted = `\`${written}\`` + const masked = maskLiteralsAndComments(sql.replaceAll(quoted, ` ${written} `)) + const pattern = new RegExp(`(? + sql.startsWith(quoted, m.index - 1) + ? { 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, + ) +} + +/** Names the SQL defines itself (`AS x`, `x AS (`), which are never column typos. */ +export const definedNames = (sql: string): ReadonlySet => + new Set( + [...maskLiteralsAndComments(sql).matchAll(/\bAS\s+([A-Za-z_]\w*)|\b([A-Za-z_]\w*)\s+AS\s*\(/gi)].map( + (m) => m[1] ?? m[2], + ), + ) + +/** Tables and aliases a column may be qualified with: `FROM traces t`, `JOIN logs AS l`. */ +const tableQualifiers = (sql: string): ReadonlySet => + new Set( + [ + ...maskLiteralsAndComments(sql).matchAll( + /\b(?:FROM|JOIN)\s+([A-Za-z_][\w.]*)(?:\s+(?:AS\s+)?([A-Za-z_]\w*))?/gi, + ), + ].flatMap((m) => [m[1], m[1].split(".").at(-1) ?? m[1], ...(m[2] === undefined ? [] : [m[2]])]), + ) + +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 || definedNames(sql).has(written)) return undefined + const to = identifierFix(written, referencedTables(sql), tableQualifiers(sql)) + 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.ts b/packages/domain/src/raw-sql.ts index cd426090b7..f85d75ff18 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. // From 026c578141f91566259ef5dfbeea8a1ed7a5a9a3 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 02:05:48 +0100 Subject: [PATCH 02/13] fix(ai): count backtick-quoted tables as read by run_sql SQL Masking blanks backtick-quoted names, so FROM `traces` read as no table and an unknown column there got neither the column list nor an automatic fix. --- apps/ai/src/mcp/lib/sql-identifiers.test.ts | 9 +++++++++ apps/ai/src/mcp/lib/sql-identifiers.ts | 6 +++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/apps/ai/src/mcp/lib/sql-identifiers.test.ts b/apps/ai/src/mcp/lib/sql-identifiers.test.ts index c6461d7cba..b1bdf80a5e 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.test.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.test.ts @@ -66,6 +66,15 @@ describe("fixUnknownIdentifier", () => { ) }) + 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") + }) + 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)" diff --git a/apps/ai/src/mcp/lib/sql-identifiers.ts b/apps/ai/src/mcp/lib/sql-identifiers.ts index ae7aa0f1c1..d6f99834da 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.ts @@ -18,7 +18,11 @@ export const unknownIdentifierName = (message: string): string | undefined => { export const referencedTableNames = (sql: string): ReadonlyArray => { const masked = maskLiteralsAndComments(sql) return listWarehouseTables() - .map((t) => ({ name: t.name, at: masked.search(new RegExp(`\\b${t.name}\\b`)) })) + .map((t) => { + // Masking blanks backticks too, so a quoted ``FROM `traces` `` is found in the raw text. + const at = masked.search(new RegExp(`\\b${t.name}\\b`)) + return { name: t.name, at: at >= 0 ? at : sql.indexOf(`\`${t.name}\``) } + }) .filter((t) => t.at >= 0) .sort((a, b) => a.at - b.at) .map((t) => t.name) From 92a6440ccc4e0ccabff8586fab9f4fd96a99dae1 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 02:22:11 +0100 Subject: [PATCH 03/13] fix(ai): scrub the gateway wrapper from run_sql errors under any database name The preview test showed pr_1961.traces and the org id in errors: the scrubber only knew the maple_us/maple_eu databases. It now matches the wrapper's shape for any database, keeps the caller's table alias, strips a database prefix in front of any catalog table and blanks any remaining org id literal. --- .../ai/src/mcp/lib/raw-sql-error-text.test.ts | 22 ++++++++++++++++ apps/ai/src/mcp/lib/raw-sql-error-text.ts | 25 +++++++++++++------ 2 files changed, 39 insertions(+), 8 deletions(-) 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 index dc0fee2815..2c93ad8deb 100644 --- a/apps/ai/src/mcp/lib/raw-sql-error-text.test.ts +++ b/apps/ai/src/mcp/lib/raw-sql-error-text.test.ts @@ -11,6 +11,28 @@ describe("rawSqlErrorText", () => { ) }) + // 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("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( diff --git a/apps/ai/src/mcp/lib/raw-sql-error-text.ts b/apps/ai/src/mcp/lib/raw-sql-error-text.ts index 9e44df06af..91562670d7 100644 --- a/apps/ai/src/mcp/lib/raw-sql-error-text.ts +++ b/apps/ai/src/mcp/lib/raw-sql-error-text.ts @@ -1,3 +1,13 @@ +import { listWarehouseTables } from "@maple/backend/services/warehouse/warehouse-catalog" + +const DB = "[A-Za-z_]\\w*" + +/** The managed gateway's per-org subquery around each table, under any database name. */ +const GATEWAY_WRAPPER = new RegExp( + `(\\()?SELECT \\* FROM \\(SELECT \\* FROM ${DB}\\.(\\w+) WHERE OrgId = '[^']*'\\) AS ${DB}\\.\\w+ WHERE OrgId = '[^']*'(?:\\) AS (\\w+))?`, + "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 @@ -6,18 +16,17 @@ 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( - /\(SELECT \* FROM \(SELECT \* FROM maple_[a-z]+\.(\w+) WHERE OrgId = '[^']*'\) AS maple_[a-z]+\.\w+ WHERE OrgId = '[^']*'\) AS \w+/g, - "$1", + .replace(GATEWAY_WRAPPER, (_whole, _open, table: string, alias: string | undefined) => + alias === undefined || alias === table ? table : `${table} AS ${alias}`, ) - .replace( - /SELECT \* FROM \(SELECT \* FROM maple_[a-z]+\.(\w+) WHERE OrgId = '[^']*'\) AS maple_[a-z]+\.\w+ WHERE OrgId = '[^']*'/g, - "$1", - ) - .replace(/\bmaple_(?:us|eu)\./g, "") + .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, "''") } From b06b31b025939a430b5006ef83665fd3bdfcaba9 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 02:31:15 +0100 Subject: [PATCH 04/13] fix(ai): keep parentheses balanced when scrubbing an unaliased gateway wrapper The wrapper pattern took an optional opening paren without its closing one, so a wrapper inside parentheses with no alias left a stray ')' behind. --- apps/ai/src/mcp/lib/raw-sql-error-text.test.ts | 8 ++++++++ apps/ai/src/mcp/lib/raw-sql-error-text.ts | 15 ++++++++------- 2 files changed, 16 insertions(+), 7 deletions(-) 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 index 2c93ad8deb..e01367063d 100644 --- a/apps/ai/src/mcp/lib/raw-sql-error-text.test.ts +++ b/apps/ai/src/mcp/lib/raw-sql-error-text.test.ts @@ -27,6 +27,14 @@ describe("rawSqlErrorText", () => { ).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 = ''", diff --git a/apps/ai/src/mcp/lib/raw-sql-error-text.ts b/apps/ai/src/mcp/lib/raw-sql-error-text.ts index 91562670d7..35b326db2b 100644 --- a/apps/ai/src/mcp/lib/raw-sql-error-text.ts +++ b/apps/ai/src/mcp/lib/raw-sql-error-text.ts @@ -2,11 +2,11 @@ import { listWarehouseTables } from "@maple/backend/services/warehouse/warehouse const DB = "[A-Za-z_]\\w*" -/** The managed gateway's per-org subquery around each table, under any database name. */ -const GATEWAY_WRAPPER = new RegExp( - `(\\()?SELECT \\* FROM \\(SELECT \\* FROM ${DB}\\.(\\w+) WHERE OrgId = '[^']*'\\) AS ${DB}\\.\\w+ WHERE OrgId = '[^']*'(?:\\) AS (\\w+))?`, - "g", -) +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 @@ -20,9 +20,10 @@ export const rawSqlErrorText = (message: string): string => { .map((t) => t.name) .join("|") return message - .replace(GATEWAY_WRAPPER, (_whole, _open, table: string, alias: string | undefined) => - alias === undefined || alias === table ? table : `${table} AS ${alias}`, + .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, "") From 517593419127d007a527b922bad67543685f85cc Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 02:39:35 +0100 Subject: [PATCH 05/13] fix(ai): read quoted identifiers and only FROM/JOIN positions as tables in run_sql fixes Masking blanked backtick-quoted names with the literals, so a quoted table, its alias or a quoted defined name was invisible, while a table name inside a literal or used as an alias counted as a table read. A scanner now blanks only literals and comments and unquotes identifiers in place, and tables count only after FROM, JOIN or a comma. --- apps/ai/src/mcp/lib/sql-identifiers.test.ts | 34 ++++++++++++++ apps/ai/src/mcp/lib/sql-identifiers.ts | 50 +++++++++++++-------- 2 files changed, 66 insertions(+), 18 deletions(-) diff --git a/apps/ai/src/mcp/lib/sql-identifiers.test.ts b/apps/ai/src/mcp/lib/sql-identifiers.test.ts index b1bdf80a5e..46013e396b 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.test.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.test.ts @@ -75,6 +75,40 @@ describe("fixUnknownIdentifier", () => { ).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("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)" diff --git a/apps/ai/src/mcp/lib/sql-identifiers.ts b/apps/ai/src/mcp/lib/sql-identifiers.ts index d6f99834da..c257cbf6c8 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.ts @@ -1,5 +1,4 @@ import type { RunSqlRewrite } from "@maple/domain/mcp-outputs" -import { maskLiteralsAndComments } from "@maple/domain/raw-sql" import { describeWarehouseTable, listWarehouseTables, @@ -14,15 +13,31 @@ export const unknownIdentifierName = (message: string): string | undefined => { return match?.[1] ?? match?.[2] ?? match?.[3] } -/** Names of the catalog tables the SQL reads, in first-mention order. */ +// String literals (with '' and \' escapes), comments, and backtick or double-quoted identifiers. +const SQL_TOKEN = + /'(?:[^'\\]|\\[\s\S]|'')*'?|--[^\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 masked = maskLiteralsAndComments(sql) + const text = identifierText(sql) return listWarehouseTables() - .map((t) => { - // Masking blanks backticks too, so a quoted ``FROM `traces` `` is found in the raw text. - const at = masked.search(new RegExp(`\\b${t.name}\\b`)) - return { name: t.name, at: at >= 0 ? at : sql.indexOf(`\`${t.name}\``) } - }) + .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) @@ -85,19 +100,18 @@ const identifierFix = ( const escapeRegExp = (text: string) => text.replace(/[.*+?^${}()|[\]\\]/g, "\\$&") /** - * Every use of `written` as an identifier, bare or backticked, replaced by `replacement`. + * 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 => { - // Unquote only this name's backticks (same length) so other quoting still masks as before. - const quoted = `\`${written}\`` - const masked = maskLiteralsAndComments(sql.replaceAll(quoted, ` ${written} `)) const pattern = new RegExp(`(? - sql.startsWith(quoted, m.index - 1) + const spans = [...identifierText(sql).matchAll(pattern)].map((m) => { + 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 }, - ) + : { start: m.index, end: m.index + written.length } + }) return spans.reduceRight( (text, { start, end }) => `${text.slice(0, start)}${replacement}${text.slice(end)}`, sql, @@ -107,7 +121,7 @@ const replaceIdentifier = (sql: string, written: string, replacement: string): s /** Names the SQL defines itself (`AS x`, `x AS (`), which are never column typos. */ export const definedNames = (sql: string): ReadonlySet => new Set( - [...maskLiteralsAndComments(sql).matchAll(/\bAS\s+([A-Za-z_]\w*)|\b([A-Za-z_]\w*)\s+AS\s*\(/gi)].map( + [...identifierText(sql).matchAll(/\bAS\s+([A-Za-z_][\w.]*)|\b([A-Za-z_]\w*)\s+AS\s*\(/gi)].map( (m) => m[1] ?? m[2], ), ) @@ -116,7 +130,7 @@ export const definedNames = (sql: string): ReadonlySet => const tableQualifiers = (sql: string): ReadonlySet => new Set( [ - ...maskLiteralsAndComments(sql).matchAll( + ...identifierText(sql).matchAll( /\b(?:FROM|JOIN)\s+([A-Za-z_][\w.]*)(?:\s+(?:AS\s+)?([A-Za-z_]\w*))?/gi, ), ].flatMap((m) => [m[1], m[1].split(".").at(-1) ?? m[1], ...(m[2] === undefined ? [] : [m[2]])]), From 10791f41e84eea1eff1aa4c620232c7428952960 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 02:45:53 +0100 Subject: [PATCH 06/13] fix(ai): fix a qualified column only from the table its qualifier names t.body in a logs/traces join was rewritten to t.Body because logs has Body, though traces (t) does not. Qualifiers now map to their table, and the fix must come from that table's columns. --- apps/ai/src/mcp/lib/sql-identifiers.test.ts | 9 +++++++ apps/ai/src/mcp/lib/sql-identifiers.ts | 28 +++++++++++++-------- 2 files changed, 27 insertions(+), 10 deletions(-) diff --git a/apps/ai/src/mcp/lib/sql-identifiers.test.ts b/apps/ai/src/mcp/lib/sql-identifiers.test.ts index 46013e396b..5ed6413a4d 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.test.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.test.ts @@ -98,6 +98,15 @@ describe("fixUnknownIdentifier", () => { ).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 rewrites a backtick-quoted alias the query defines", () => { const sql = "SELECT ServiceName AS `service.name` FROM traces WHERE $__orgFilter ORDER BY `service.name`" diff --git a/apps/ai/src/mcp/lib/sql-identifiers.ts b/apps/ai/src/mcp/lib/sql-identifiers.ts index c257cbf6c8..7f1b37635a 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.ts @@ -77,7 +77,7 @@ const KNOWN_ALIASES: ReadonlyMap, - qualifiers: ReadonlySet, + 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 @@ -90,11 +90,13 @@ const identifierFix = ( 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, fix the column, but only for a table or alias in scope. + // `t.trace_id`: keep the qualifier and fix the column from the table it names. const dot = written.lastIndexOf(".") - if (dot <= 0 || !qualifiers.has(written.slice(0, dot))) return undefined - const column = byFold(written.slice(dot + 1), false) - return column === undefined ? undefined : `${written.slice(0, dot + 1)}${column}` + 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, "\\$&") @@ -126,15 +128,21 @@ export const definedNames = (sql: string): ReadonlySet => ), ) -/** Tables and aliases a column may be qualified with: `FROM traces t`, `JOIN logs AS l`. */ -const tableQualifiers = (sql: string): ReadonlySet => - new Set( +/** 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( [ ...identifierText(sql).matchAll( - /\b(?:FROM|JOIN)\s+([A-Za-z_][\w.]*)(?:\s+(?:AS\s+)?([A-Za-z_]\w*))?/gi, + /(?:\bFROM\s+|\bJOIN\s+|,\s*)([A-Za-z_][\w.]*)(?:\s+(?:AS\s+)?([A-Za-z_]\w*))?/gi, ), - ].flatMap((m) => [m[1], m[1].split(".").at(-1) ?? m[1], ...(m[2] === undefined ? [] : [m[2]])]), + ].flatMap((m) => { + const table = m[1].split(".").at(-1) ?? m[1] + if (!catalog.has(table)) return [] + return [m[1], table, ...(m[2] === undefined ? [] : [m[2]])].map((q) => [q, table] as const) + }), ) +} export type IdentifierRewrite = typeof RunSqlRewrite.Type From c44b1627825d58229cb32d4b5d4e28f043cab53f Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 02:51:27 +0100 Subject: [PATCH 07/13] fix(ai): never rename an AS-less table alias, and mask # comments in run_sql fixes FROM traces service with an unknown service renamed the alias too, breaking the FROM and its $__orgFilter(service). A name the query uses as a table qualifier is now never rewritten. # comments are blanked like -- comments. --- apps/ai/src/mcp/lib/sql-identifiers.test.ts | 19 +++++++++++++++++++ apps/ai/src/mcp/lib/sql-identifiers.ts | 10 ++++++---- 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/apps/ai/src/mcp/lib/sql-identifiers.test.ts b/apps/ai/src/mcp/lib/sql-identifiers.test.ts index 5ed6413a4d..0575982fda 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.test.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.test.ts @@ -107,6 +107,25 @@ describe("fixUnknownIdentifier", () => { ) }) + 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("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`" diff --git a/apps/ai/src/mcp/lib/sql-identifiers.ts b/apps/ai/src/mcp/lib/sql-identifiers.ts index 7f1b37635a..91eea002fb 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.ts @@ -13,9 +13,9 @@ export const unknownIdentifierName = (message: string): string | undefined => { return match?.[1] ?? match?.[2] ?? match?.[3] } -// String literals (with '' and \' escapes), comments, and backtick or double-quoted identifiers. +// String literals (with '' and \' escapes), `--`/`#`/`/* */` comments, and quoted identifiers. const SQL_TOKEN = - /'(?:[^'\\]|\\[\s\S]|'')*'?|--[^\n]*|\/\*[\s\S]*?(?:\*\/|$)|`(?:[^`]|``)*`?|"(?:[^"]|"")*"?/g + /'(?:[^'\\]|\\[\s\S]|'')*'?|--[^\n]*|#[^\n]*|\/\*[\s\S]*?(?:\*\/|$)|`(?:[^`]|``)*`?|"(?:[^"]|"")*"?/g /** * The SQL with literals and comments blanked and quoted identifiers unquoted, keeping every @@ -152,8 +152,10 @@ export const fixUnknownIdentifier = ( message: string, ): { readonly sql: string; readonly rewrite: IdentifierRewrite } | undefined => { const written = unknownIdentifierName(message) - if (written === undefined || definedNames(sql).has(written)) return undefined - const to = identifierFix(written, referencedTables(sql), tableQualifiers(sql)) + 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 } } From db55d921f8fd6bde42202b831f5142732fc131a5 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 02:55:50 +0100 Subject: [PATCH 08/13] fix(ai): never rename an alias that follows a closing parenthesis FROM (...) service with an unknown service renamed the derived table's alias. A name straight after ')' that is not a keyword is now a name the query defines, which also covers select aliases like count() c. --- apps/ai/src/mcp/lib/sql-identifiers.test.ts | 5 ++++ apps/ai/src/mcp/lib/sql-identifiers.ts | 28 +++++++++++++++++---- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/apps/ai/src/mcp/lib/sql-identifiers.test.ts b/apps/ai/src/mcp/lib/sql-identifiers.test.ts index 0575982fda..aa3eb06beb 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.test.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.test.ts @@ -112,6 +112,11 @@ describe("fixUnknownIdentifier", () => { 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("ignores a table named in a # comment", () => { expect( fixUnknownIdentifier( diff --git a/apps/ai/src/mcp/lib/sql-identifiers.ts b/apps/ai/src/mcp/lib/sql-identifiers.ts index 91eea002fb..f919e59a7b 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.ts @@ -120,13 +120,31 @@ const replaceIdentifier = (sql: string, written: string, replacement: string): s ) } -/** Names the SQL defines itself (`AS x`, `x AS (`), which are never column typos. */ -export const definedNames = (sql: string): ReadonlySet => - new Set( - [...identifierText(sql).matchAll(/\bAS\s+([A-Za-z_][\w.]*)|\b([A-Za-z_]\w*)\s+AS\s*\(/gi)].map( +/** Words that can follow a closing parenthesis without being an alias. */ +const KEYWORDS_AFTER_PAREN = 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(" "), +) + +/** + * Names the SQL defines itself, which are never column typos: `AS x`, `x AS (`, and an + * alias straight after a closing parenthesis (`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_AFTER_PAREN.has(name.toLowerCase())), + ]) +} /** The catalog table each qualifier names: `FROM traces t`, `JOIN logs AS l`, `FROM a, b`. */ const tableQualifiers = (sql: string): ReadonlyMap => { From 9da86aa342dbb0ad54a5a82cbbfa4bf2db76a7b7 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 02:57:26 +0100 Subject: [PATCH 09/13] fix(ai): count a dotted alias after a closing parenthesis as defined --- apps/ai/src/mcp/lib/sql-identifiers.test.ts | 6 ++++++ apps/ai/src/mcp/lib/sql-identifiers.ts | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/apps/ai/src/mcp/lib/sql-identifiers.test.ts b/apps/ai/src/mcp/lib/sql-identifiers.test.ts index aa3eb06beb..cabc4ac28a 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.test.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.test.ts @@ -117,6 +117,12 @@ describe("fixUnknownIdentifier", () => { 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("ignores a table named in a # comment", () => { expect( fixUnknownIdentifier( diff --git a/apps/ai/src/mcp/lib/sql-identifiers.ts b/apps/ai/src/mcp/lib/sql-identifiers.ts index f919e59a7b..83f4bf84fd 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.ts @@ -140,7 +140,7 @@ export const definedNames = (sql: string): ReadonlySet => { ...[...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)] + ...[...text.matchAll(/\)\s*([A-Za-z_][\w.]*)/g)] .map((m) => m[1]) .filter((name) => !KEYWORDS_AFTER_PAREN.has(name.toLowerCase())), ]) From 8fd3173d439710b576a8d95f1160363d77947af3 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 03:01:31 +0100 Subject: [PATCH 10/13] fix(ai): treat a dotted table alias declared without AS as a qualifier --- apps/ai/src/mcp/lib/sql-identifiers.test.ts | 5 +++++ apps/ai/src/mcp/lib/sql-identifiers.ts | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/apps/ai/src/mcp/lib/sql-identifiers.test.ts b/apps/ai/src/mcp/lib/sql-identifiers.test.ts index cabc4ac28a..f4845ab737 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.test.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.test.ts @@ -123,6 +123,11 @@ describe("fixUnknownIdentifier", () => { 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("ignores a table named in a # comment", () => { expect( fixUnknownIdentifier( diff --git a/apps/ai/src/mcp/lib/sql-identifiers.ts b/apps/ai/src/mcp/lib/sql-identifiers.ts index 83f4bf84fd..58eeb5c67d 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.ts @@ -152,7 +152,7 @@ const tableQualifiers = (sql: string): ReadonlyMap => { return new Map( [ ...identifierText(sql).matchAll( - /(?:\bFROM\s+|\bJOIN\s+|,\s*)([A-Za-z_][\w.]*)(?:\s+(?:AS\s+)?([A-Za-z_]\w*))?/gi, + /(?:\bFROM\s+|\bJOIN\s+|,\s*)([A-Za-z_][\w.]*)(?:\s+(?:AS\s+)?([A-Za-z_][\w.]*))?/gi, ), ].flatMap((m) => { const table = m[1].split(".").at(-1) ?? m[1] From 2e6bedf9976567b80cdc0dd0eff2b84c1da0e356 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 03:04:23 +0100 Subject: [PATCH 11/13] fix(ai): one relation scan for run_sql aliases, keywords never an alias FROM traces JOIN logs l took JOIN as the first table's alias and missed l, and the alias of a CTE or other non-catalog relation (FROM s service) was not a defined name, so it could be renamed. Every FROM/JOIN/comma relation's alias is now a defined name, clause keywords are never captured as one, and the qualifier map keeps only catalog tables. --- apps/ai/src/mcp/lib/sql-identifiers.test.ts | 12 +++++++++ apps/ai/src/mcp/lib/sql-identifiers.ts | 30 ++++++++++++--------- 2 files changed, 30 insertions(+), 12 deletions(-) diff --git a/apps/ai/src/mcp/lib/sql-identifiers.test.ts b/apps/ai/src/mcp/lib/sql-identifiers.test.ts index f4845ab737..34a7f815ef 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.test.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.test.ts @@ -128,6 +128,18 @@ describe("fixUnknownIdentifier", () => { 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("ignores a table named in a # comment", () => { expect( fixUnknownIdentifier( diff --git a/apps/ai/src/mcp/lib/sql-identifiers.ts b/apps/ai/src/mcp/lib/sql-identifiers.ts index 58eeb5c67d..efe442fbd5 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.ts @@ -120,8 +120,8 @@ const replaceIdentifier = (sql: string, written: string, replacement: string): s ) } -/** Words that can follow a closing parenthesis without being an alias. */ -const KEYWORDS_AFTER_PAREN = new Set( +/** 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 " + @@ -130,9 +130,18 @@ const KEYWORDS_AFTER_PAREN = new Set( ).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 (`, and an - * alias straight after a closing parenthesis (`FROM (...) t`, `count() c`). + * 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) @@ -142,7 +151,8 @@ export const definedNames = (sql: string): ReadonlySet => { ), ...[...text.matchAll(/\)\s*([A-Za-z_][\w.]*)/g)] .map((m) => m[1]) - .filter((name) => !KEYWORDS_AFTER_PAREN.has(name.toLowerCase())), + .filter((name) => !KEYWORDS.has(name.toLowerCase())), + ...relations(sql).flatMap((r) => (r.alias === undefined ? [] : [r.alias])), ]) } @@ -150,14 +160,10 @@ export const definedNames = (sql: string): ReadonlySet => { const tableQualifiers = (sql: string): ReadonlyMap => { const catalog = new Set(listWarehouseTables().map((t) => t.name)) return new Map( - [ - ...identifierText(sql).matchAll( - /(?:\bFROM\s+|\bJOIN\s+|,\s*)([A-Za-z_][\w.]*)(?:\s+(?:AS\s+)?([A-Za-z_][\w.]*))?/gi, - ), - ].flatMap((m) => { - const table = m[1].split(".").at(-1) ?? m[1] + relations(sql).flatMap(({ name, alias }) => { + const table = name.split(".").at(-1) ?? name if (!catalog.has(table)) return [] - return [m[1], table, ...(m[2] === undefined ? [] : [m[2]])].map((q) => [q, table] as const) + return [name, table, ...(alias === undefined ? [] : [alias])].map((q) => [q, table] as const) }), ) } From 7db50416254ab390f139ecfefa3fc58a46ac5d15 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 03:07:25 +0100 Subject: [PATCH 12/13] fix(domain): bind a joined table's alias after an unaliased first table The raw-SQL gate read FROM traces JOIN logs l as traces aliased JOIN, so l never bound and $__orgFilter(l) was rejected on valid SQL. A clause keyword is never taken as an alias; a joined subquery's alias still does not bind. --- packages/domain/src/raw-sql.test.ts | 5 +++++ packages/domain/src/raw-sql.ts | 15 ++++++--------- 2 files changed, 11 insertions(+), 9 deletions(-) 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 f85d75ff18..c72b5a6b4c 100644 --- a/packages/domain/src/raw-sql.ts +++ b/packages/domain/src/raw-sql.ts @@ -540,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, "\\$&") /** @@ -555,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) => { @@ -573,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 } From 414492e98e417404fbedfcd637725231aec4e390 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Fri, 30 Oct 2026 03:24:36 +0100 Subject: [PATCH 13/13] fix(ai): decline a qualified fix when its qualifier names two relations With t bound to a CTE in one scope and to traces in another, fixing the inner t.span_name rewrote the outer valid t.span_name too, so the query could run and read the wrong column. --- apps/ai/src/mcp/lib/sql-identifiers.test.ts | 7 +++++++ apps/ai/src/mcp/lib/sql-identifiers.ts | 17 +++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/apps/ai/src/mcp/lib/sql-identifiers.test.ts b/apps/ai/src/mcp/lib/sql-identifiers.test.ts index 34a7f815ef..95f5122336 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.test.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.test.ts @@ -140,6 +140,13 @@ describe("fixUnknownIdentifier", () => { 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( diff --git a/apps/ai/src/mcp/lib/sql-identifiers.ts b/apps/ai/src/mcp/lib/sql-identifiers.ts index efe442fbd5..c0d889e59f 100644 --- a/apps/ai/src/mcp/lib/sql-identifiers.ts +++ b/apps/ai/src/mcp/lib/sql-identifiers.ts @@ -168,6 +168,22 @@ const tableQualifiers = (sql: string): ReadonlyMap => { ) } +/** + * 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. */ @@ -176,6 +192,7 @@ export const fixUnknownIdentifier = ( 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