Linter now works with AS queries

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");
+    });
+  });
+});
+

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<string>;
 }
 
 /**
@@ -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()) {