diff --git a/internal-packages/tsql/src/query/validator.test.ts b/internal-packages/tsql/src/query/validator.test.ts new file mode 100644 index 000000000..7d0b42fec --- /dev/null +++ b/internal-packages/tsql/src/query/validator.test.ts @@ -0,0 +1,99 @@ +import { describe, it, expect } from "vitest"; +import { validateQuery } from "./validator.js"; +import { parseTSQLSelect } from "../index.js"; +import { column, type TableSchema } from "./schema.js"; + +const runsSchema: TableSchema = { + name: "runs", + clickhouseName: "trigger_dev.task_runs_v2", + columns: { + id: { name: "id", ...column("String") }, + status: { + name: "status", + ...column("String", { + allowedValues: ["PENDING", "COMPLETED", "FAILED"], + }), + }, + task_id: { name: "task_id", ...column("String") }, + created_at: { name: "created_at", ...column("DateTime64") }, + }, + tenantColumns: { + organizationId: "organization_id", + }, +}; + +function validateSQL(query: string, schema: TableSchema[] = [runsSchema]) { + const ast = parseTSQLSelect(query); + return validateQuery(ast, schema); +} + +describe("validateQuery", () => { + describe("SELECT aliases", () => { + it("should allow ORDER BY to reference aliased columns", () => { + const result = validateSQL( + "SELECT status, count(*) as count FROM runs GROUP BY status ORDER BY count DESC" + ); + expect(result.valid).toBe(true); + expect(result.issues).toHaveLength(0); + }); + + it("should allow ORDER BY to reference multiple aliased columns", () => { + const result = validateSQL( + "SELECT status, count(*) as total, avg(created_at) as avg_time FROM runs GROUP BY status ORDER BY total DESC, avg_time ASC" + ); + expect(result.valid).toBe(true); + expect(result.issues).toHaveLength(0); + }); + + it("should still report unknown columns that are not aliases", () => { + const result = validateSQL( + "SELECT status, count(*) as count FROM runs GROUP BY status ORDER BY unknown_col DESC" + ); + expect(result.valid).toBe(true); // unknown column is a warning, not error + expect(result.issues).toHaveLength(1); + expect(result.issues[0].type).toBe("unknown_column"); + expect(result.issues[0].columnName).toBe("unknown_col"); + }); + + it("should allow ORDER BY to reference both aliases and real columns", () => { + const result = validateSQL( + "SELECT status, count(*) as count FROM runs GROUP BY status ORDER BY status ASC, count DESC" + ); + expect(result.valid).toBe(true); + expect(result.issues).toHaveLength(0); + }); + }); + + describe("column validation", () => { + it("should validate known columns", () => { + const result = validateSQL("SELECT id, status FROM runs LIMIT 10"); + expect(result.valid).toBe(true); + expect(result.issues).toHaveLength(0); + }); + + it("should warn about unknown columns", () => { + const result = validateSQL("SELECT id, unknown_column FROM runs LIMIT 10"); + expect(result.valid).toBe(true); // warnings don't affect validity + expect(result.issues).toHaveLength(1); + expect(result.issues[0].type).toBe("unknown_column"); + expect(result.issues[0].columnName).toBe("unknown_column"); + }); + }); + + describe("enum validation", () => { + it("should validate enum values", () => { + const result = validateSQL("SELECT * FROM runs WHERE status = 'COMPLETED' LIMIT 10"); + expect(result.valid).toBe(true); + expect(result.issues).toHaveLength(0); + }); + + it("should error on invalid enum values", () => { + const result = validateSQL("SELECT * FROM runs WHERE status = 'INVALID_STATUS' LIMIT 10"); + expect(result.valid).toBe(false); + expect(result.issues).toHaveLength(1); + expect(result.issues[0].type).toBe("invalid_enum_value"); + expect(result.issues[0].invalidValue).toBe("INVALID_STATUS"); + }); + }); +}); + diff --git a/internal-packages/tsql/src/query/validator.ts b/internal-packages/tsql/src/query/validator.ts index 6a9ede319..fe4e414de 100644 --- a/internal-packages/tsql/src/query/validator.ts +++ b/internal-packages/tsql/src/query/validator.ts @@ -67,6 +67,8 @@ interface ValidationContext { schema: TableSchema[]; /** Accumulated issues */ issues: ValidationIssue[]; + /** Set of column aliases defined in the SELECT clause */ + selectAliases: Set; } /** @@ -84,6 +86,7 @@ export function validateQuery( tables: new Map(), schema, issues: [], + selectAliases: new Set(), }; if (ast.expression_type === "select_set_query") { @@ -121,11 +124,25 @@ function validateSelectSetQuery(node: SelectSetQuery, context: ValidationContext * Validate a SELECT query */ function validateSelectQuery(node: SelectQuery, context: ValidationContext): void { + // Save parent aliases and create fresh set for this query + const parentAliases = context.selectAliases; + context.selectAliases = new Set(); + // First, extract tables from FROM clause to build context if (node.select_from) { extractTablesFromJoin(node.select_from, context); } + // Extract column aliases from SELECT clause before validation + // This allows ORDER BY to reference aliased columns + if (node.select) { + for (const expr of node.select) { + if ((expr as Alias).expression_type === "alias") { + context.selectAliases.add((expr as Alias).alias); + } + } + } + // Validate SELECT columns if (node.select) { for (const expr of node.select) { @@ -156,6 +173,9 @@ function validateSelectQuery(node: SelectQuery, context: ValidationContext): voi validateExpression(expr, context); } } + + // Restore parent aliases + context.selectAliases = parentAliases; } /** @@ -316,8 +336,14 @@ function validateField(field: Field, context: ValidationContext): void { return; } - // Case 2: Unqualified reference - try to find in any table + // Case 2: Unqualified reference - try to find in any table or SELECT alias const columnName = firstPart; + + // Check if it's a SELECT alias (e.g., from "count(*) as count") + if (context.selectAliases.has(columnName)) { + return; + } + let found = false; for (const tableSchema of context.tables.values()) {