Improve access errors to columns not in the schema

Author: matt-aitken

internal-packages/tsql/src/query/printer.test.ts

@@ -429,15 +429,15 @@ describe("ClickHousePrinter", () => {
     });
 
     it("should handle IS NULL syntax", () => {
-      const { sql } = printQuery("SELECT * FROM task_runs WHERE error IS NULL");
+      const { sql } = printQuery("SELECT * FROM task_runs WHERE started_at IS NULL");
 
-      expect(sql).toContain("isNull(error)");
+      expect(sql).toContain("isNull(started_at)");
     });
 
     it("should handle IS NOT NULL syntax", () => {
-      const { sql } = printQuery("SELECT * FROM task_runs WHERE error IS NOT NULL");
+      const { sql } = printQuery("SELECT * FROM task_runs WHERE started_at IS NOT NULL");
 
-      expect(sql).toContain("isNotNull(error)");
+      expect(sql).toContain("isNotNull(started_at)");
     });
   });
 
@@ -2172,6 +2172,138 @@ describe("Column metadata", () => {
   });
 });
 
+describe("Unknown column blocking", () => {
+  /**
+   * These tests verify that unknown columns are blocked at compile time,
+   * preventing access to internal ClickHouse columns that aren't exposed in the schema.
+   */
+
+  describe("should block unknown columns in SELECT", () => {
+    it("should throw error for unknown column in SELECT list", () => {
+      expect(() => {
+        printQuery("SELECT id, unknown_column FROM task_runs");
+      }).toThrow(QueryError);
+      expect(() => {
+        printQuery("SELECT id, unknown_column FROM task_runs");
+      }).toThrow(/unknown.*column.*unknown_column/i);
+    });
+
+    it("should throw error for internal ClickHouse column name not in schema", () => {
+      // The schema exposes 'created' but the internal ClickHouse column is 'created_at'
+      // Using the internal name directly should be blocked
+      const schema = createSchemaRegistry([runsSchema]);
+      const ctx = createPrinterContext({
+        organizationId: "org_test",
+        projectId: "proj_test",
+        environmentId: "env_test",
+        schema,
+      });
+
+      // 'created_at' is not in runsSchema - only 'created' which maps to 'created_at'
+      expect(() => {
+        printQuery("SELECT id, created_at FROM runs", ctx);
+      }).toThrow(QueryError);
+    });
+
+    it("should throw error for unknown qualified column (table.column)", () => {
+      expect(() => {
+        printQuery("SELECT task_runs.unknown_column FROM task_runs");
+      }).toThrow(QueryError);
+    });
+  });
+
+  describe("should block unknown columns in WHERE", () => {
+    it("should throw error for unknown column in WHERE clause", () => {
+      expect(() => {
+        printQuery("SELECT id FROM task_runs WHERE unknown_column = 'test'");
+      }).toThrow(QueryError);
+    });
+
+    it("should throw error for unknown column in complex WHERE", () => {
+      expect(() => {
+        printQuery(
+          "SELECT id FROM task_runs WHERE status = 'completed' AND unknown_column > 10"
+        );
+      }).toThrow(QueryError);
+    });
+  });
+
+  describe("should block unknown columns in ORDER BY", () => {
+    it("should throw error for unknown column in ORDER BY", () => {
+      expect(() => {
+        printQuery("SELECT id, status FROM task_runs ORDER BY unknown_column DESC");
+      }).toThrow(QueryError);
+    });
+  });
+
+  describe("should block unknown columns in GROUP BY", () => {
+    it("should throw error for unknown column in GROUP BY", () => {
+      expect(() => {
+        printQuery("SELECT count(*) FROM task_runs GROUP BY unknown_column");
+      }).toThrow(QueryError);
+    });
+  });
+
+  describe("should allow SELECT aliases", () => {
+    it("should allow ORDER BY to reference aliased columns", () => {
+      // This should NOT throw - 'cnt' is a valid alias from SELECT
+      const { sql } = printQuery(
+        "SELECT status, count(*) AS cnt FROM task_runs GROUP BY status ORDER BY cnt DESC"
+      );
+      expect(sql).toContain("ORDER BY cnt DESC");
+    });
+
+    it("should allow HAVING to reference aliased columns", () => {
+      const { sql } = printQuery(
+        "SELECT status, count(*) AS cnt FROM task_runs GROUP BY status HAVING cnt > 10"
+      );
+      expect(sql).toContain("HAVING");
+    });
+
+    it("should allow ORDER BY to reference implicit aggregation names", () => {
+      // COUNT() without alias gets implicit name 'count'
+      const { sql } = printQuery(
+        "SELECT status, count() FROM task_runs GROUP BY status ORDER BY count DESC"
+      );
+      expect(sql).toContain("ORDER BY count DESC");
+    });
+  });
+
+  // Note: CTE support is limited - CTEs are not added to the table context,
+  // so CTE column references are treated as unknown. This is a pre-existing limitation.
+  // The tests below verify that unknown column blocking doesn't break existing behavior.
+
+  describe("should allow subquery alias references", () => {
+    it("should allow referencing columns from subquery in FROM clause", () => {
+      const { sql } = printQuery(`
+        SELECT sub.status_name, sub.total
+        FROM (
+          SELECT status AS status_name, count(*) AS total
+          FROM task_runs
+          GROUP BY status
+        ) AS sub
+        ORDER BY sub.total DESC
+      `);
+      expect(sql).toContain("status_name");
+      expect(sql).toContain("total");
+    });
+
+    it("should allow unqualified references to subquery columns", () => {
+      const { sql } = printQuery(`
+        SELECT status_name, total
+        FROM (
+          SELECT status AS status_name, count(*) AS total
+          FROM task_runs
+          GROUP BY status
+        )
+        WHERE total > 10
+      `);
+      expect(sql).toContain("status_name");
+      expect(sql).toContain("total");
+    });
+  });
+});
+
 describe("Field Mapping Value Transformation", () => {
   // Test schema with a field mapping column
   const fieldMappingSchema: TableSchema = {

internal-packages/tsql/src/query/printer.ts

@@ -116,6 +116,16 @@ export class ClickHousePrinter {
   private inGroupByContext = false;
   /** Columns hidden when SELECT * is expanded to core columns only */
   private hiddenColumns: string[] = [];
+  /**
+   * Set of column aliases defined in the current SELECT clause.
+   * Used to allow ORDER BY/HAVING to reference aliased columns.
+   */
+  private selectAliases: Set<string> = new Set();
+  /**
+   * Set of internal ClickHouse column names that are allowed (e.g., tenant columns).
+   * These are populated from tableSchema.tenantColumns when processing joins.
+   */
+  private allowedInternalColumns: Set<string> = new Set();
 
   constructor(
     private context: PrinterContext,
@@ -310,9 +320,18 @@ export class ClickHousePrinter {
     const isTopLevelQuery =
       this.stack.length <= 1 || (this.stack.length === 2 && partOfSelectUnion);
 
-    // Clear table contexts for top-level queries (subqueries inherit parent context)
+    // Save and clear table contexts
+    // Top-level queries clear contexts; subqueries save parent context and create fresh context
+    const savedTableContexts = new Map(this.tableContexts);
+    const savedInternalColumns = new Set(this.allowedInternalColumns);
     if (isTopLevelQuery) {
       this.tableContexts.clear();
+      this.allowedInternalColumns.clear();
+    } else {
+      // Subqueries get fresh contexts - they don't inherit parent tables
+      // (the parent will restore its context after the subquery is processed)
+      this.tableContexts = new Map();
+      this.allowedInternalColumns = new Set();
     }
 
     // Build WHERE clause starting with any existing where
@@ -349,6 +368,16 @@ export class ClickHousePrinter {
       nextJoin = nextJoin.next_join;
     }
 
+    // Extract SELECT column aliases BEFORE visiting columns
+    // This allows ORDER BY/HAVING to reference aliased columns
+    const savedAliases = this.selectAliases;
+    this.selectAliases = new Set();
+    if (node.select) {
+      for (const col of node.select) {
+        this.extractSelectAlias(col);
+      }
+    }
+
     // Process SELECT columns and collect metadata
     // Using flatMap because asterisk expansion can return multiple columns
     let columns: string[];
@@ -478,9 +507,56 @@ export class ClickHousePrinter {
       response = this.pretty ? `(${response.trim()})` : `(${response})`;
     }
 
+    // Restore saved contexts (for nested queries)
+    this.selectAliases = savedAliases;
+    this.tableContexts = savedTableContexts;
+    this.allowedInternalColumns = savedInternalColumns;
+
     return response;
   }
 
+  /**
+   * Extract column aliases from a SELECT expression.
+   * Handles explicit aliases (AS name) and implicit names from aggregations/functions.
+   *
+   * NOTE: We intentionally do NOT add field names as aliases here.
+   * Field names (e.g., SELECT status) are columns from the table, not aliases.
+   * Only explicit aliases (SELECT x AS name) and implicit function names
+   * (SELECT COUNT() → 'count') should be added.
+   */
+  private extractSelectAlias(expr: Expression): void {
+    // Handle explicit Alias: SELECT ... AS name
+    if ((expr as Alias).expression_type === "alias") {
+      this.selectAliases.add((expr as Alias).alias);
+      return;
+    }
+
+    // Handle implicit names from function calls (e.g., COUNT() → 'count')
+    if ((expr as Call).expression_type === "call") {
+      const call = expr as Call;
+      // Aggregations and functions get implicit lowercase names
+      this.selectAliases.add(call.name.toLowerCase());
+      return;
+    }
+
+    // Handle implicit names from arithmetic operations (e.g., a + b → 'plus')
+    if ((expr as ArithmeticOperation).expression_type === "arithmetic_operation") {
+      const op = expr as ArithmeticOperation;
+      const opNames: Record<ArithmeticOperationOp, string> = {
+        [ArithmeticOperationOp.Add]: "plus",
+        [ArithmeticOperationOp.Sub]: "minus",
+        [ArithmeticOperationOp.Mult]: "multiply",
+        [ArithmeticOperationOp.Div]: "divide",
+        [ArithmeticOperationOp.Mod]: "modulo",
+      };
+      this.selectAliases.add(opNames[op.op]);
+      return;
+    }
+
+    // Field references (e.g., SELECT status) are NOT aliases - they're columns
+    // that will be validated against the table schema. Don't add them here.
+  }
+
   /**
    * Visit a SELECT column expression with metadata collection
    *
@@ -1338,6 +1414,15 @@ export class ClickHousePrinter {
         // Register this table context for column name resolution
         this.tableContexts.set(effectiveAlias, tableSchema);
 
+        // Register tenant columns as allowed internal columns
+        // These are ClickHouse column names used for tenant isolation
+        if (tableSchema.tenantColumns) {
+          const { organizationId, projectId, environmentId } = tableSchema.tenantColumns;
+          if (organizationId) this.allowedInternalColumns.add(organizationId);
+          if (projectId) this.allowedInternalColumns.add(projectId);
+          if (environmentId) this.allowedInternalColumns.add(environmentId);
+        }
+
         // Add tenant isolation guard
         extraWhere = this.createTenantGuard(tableSchema, effectiveAlias);
       } else if (
@@ -2081,6 +2166,8 @@ export class ClickHousePrinter {
   /**
    * Resolve field chain to use ClickHouse column names where applicable
    * Handles both qualified (table.column) and unqualified (column) references
+   *
+   * @throws QueryError if the column is not found in any table schema and is not a SELECT alias
    */
   private resolveFieldChain(chain: Array<string | number>): Array<string | number> {
     if (chain.length === 0) {
@@ -2101,11 +2188,11 @@ export class ClickHousePrinter {
         // This is a table.column reference
         const columnName = chain[1];
         if (typeof columnName === "string") {
-          const resolvedColumn = this.resolveColumnName(tableSchema, columnName);
+          const resolvedColumn = this.resolveColumnName(tableSchema, columnName, tableAlias);
           return [tableAlias, resolvedColumn, ...chain.slice(2)];
         }
       }
-      // Not a known table alias, might be a nested field - return as-is
+      // Not a known table alias, might be a nested field or JSON path - return as-is
       return chain;
     }
 
@@ -2119,20 +2206,59 @@ export class ClickHousePrinter {
       }
     }
 
-    // Column not found in any table context - return as-is (might be a function, subquery alias, etc.)
+    // Check if it's a SELECT alias (e.g., from COUNT() or explicit AS)
+    if (this.selectAliases.has(columnName)) {
+      return chain; // Valid alias reference
+    }
+
+    // Check if this is an allowed internal column (e.g., tenant columns)
+    // These are ClickHouse column names that are allowed for internal use only
+    if (this.allowedInternalColumns.has(columnName)) {
+      return chain;
+    }
+
+    // Column not found in any table context and not a SELECT alias
+    // This is a security issue - block access to unknown columns
+    if (this.tableContexts.size > 0) {
+      // Only throw if we have tables in context (otherwise might be subquery)
+      const availableColumns = this.getAvailableColumnNames();
+      throw new QueryError(
+        `Unknown column "${columnName}". Available columns: ${availableColumns.join(", ")}`
+      );
+    }
+
+    // No tables in context (might be FROM subquery without alias) - return as-is
     return chain;
   }
 
+  /**
+   * Get list of available column names from all tables in context
+   */
+  private getAvailableColumnNames(): string[] {
+    const columns: string[] = [];
+    for (const tableSchema of this.tableContexts.values()) {
+      columns.push(...Object.keys(tableSchema.columns));
+    }
+    return [...new Set(columns)].sort();
+  }
+
   /**
    * Resolve a column name to its ClickHouse name using the table schema
+   *
+   * @throws QueryError if the column is not found in the table schema
    */
-  private resolveColumnName(tableSchema: TableSchema, columnName: string): string {
+  private resolveColumnName(tableSchema: TableSchema, columnName: string, tableAlias?: string): string {
     const columnSchema = tableSchema.columns[columnName];
     if (columnSchema) {
       return columnSchema.clickhouseName || columnSchema.name;
     }
-    // Column not in schema - return as-is (might be a computed column, etc.)
-    return columnName;
+
+    // Column not in schema - this is a security issue, block access
+    const availableColumns = Object.keys(tableSchema.columns).sort().join(", ");
+    const tableName = tableAlias || tableSchema.name;
+    throw new QueryError(
+      `Unknown column "${columnName}" on table "${tableName}". Available columns: ${availableColumns}`
+    );
   }
 
   private visitPlaceholder(node: Placeholder): string {