Skip to content

Commit f4dc2b9

Browse files
1stvampTrigger.dev RepoOps
authored andcommitted
fix(tsql): print column CTEs as expression aliases
Queries using `WITH <expression> AS <name>` now run instead of failing with a syntax error, and the name can be used in the rest of the query. The name can't reuse a table alias used anywhere in the query. Mono-RevId: d22b6e0d58d14a11f286d15e62a74b6b1f021da3
1 parent bdc994c commit f4dc2b9

5 files changed

Lines changed: 210 additions & 3 deletions

File tree

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
---
2+
area: webapp
3+
type: fix
4+
---
5+
6+
Queries using `WITH <expression> AS <name>` now run instead of failing with a syntax error, and the name can be used in the rest of the query. The name can't reuse a table alias used anywhere in the query.

‎internal-packages/clickhouse/src/tsql.test.ts‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,6 +295,83 @@ describe("TSQL Integration Tests", () => {
295295
}
296296
);
297297

298+
clickhouseTest(
299+
"should not let a WITH alias replace a tenant guard in another scope",
300+
async ({ clickhouseContainer }) => {
301+
const client = new ClickhouseClient({
302+
name: "test",
303+
url: clickhouseContainer.getConnectionUrl(),
304+
});
305+
306+
const insert = insertTaskRuns(client, { async_insert: 0 });
307+
308+
await insert([
309+
createTaskRun({
310+
run_id: "run_tenant1_a",
311+
organization_id: "org_tenant1",
312+
project_id: "proj_tenant1",
313+
environment_id: "env_tenant1",
314+
}),
315+
createTaskRun({
316+
run_id: "run_tenant2_a",
317+
organization_id: "org_tenant2",
318+
project_id: "proj_tenant2",
319+
environment_id: "env_tenant2",
320+
}),
321+
]);
322+
323+
const enforcedWhereClause = {
324+
organization_id: { op: "eq" as const, value: "org_tenant1" },
325+
project_id: { op: "eq" as const, value: "proj_tenant1" },
326+
environment_id: { op: "eq" as const, value: "env_tenant1" },
327+
};
328+
const ownTenant = `JSONExtract('{"organization_id":"org_tenant1","project_id":"proj_tenant1","environment_id":"env_tenant1"}', 'Tuple(organization_id String, project_id String, environment_id String)')`;
329+
330+
// ClickHouse carries a WITH alias into every UNION branch, where it would replace the
331+
// branch's guards, so a WITH alias may not reuse a table alias anywhere in the query.
332+
for (const query of [
333+
`WITH ${ownTenant} AS r SELECT run_id FROM task_runs AS q UNION ALL SELECT run_id FROM task_runs AS r`,
334+
`WITH ${ownTenant} AS task_runs SELECT 'x' AS run_id UNION ALL SELECT run_id FROM task_runs`,
335+
`WITH ${ownTenant} AS r SELECT run_id FROM (SELECT run_id FROM task_runs AS r)`,
336+
`WITH ${ownTenant} AS task_runs SELECT run_id FROM (SELECT run_id FROM task_runs)`,
337+
`WITH (SELECT ${ownTenant}) AS task_runs SELECT run_id FROM task_runs AS q WHERE run_id IN (SELECT run_id FROM task_runs)`,
338+
]) {
339+
const [error, result] = await executeTSQL(client, {
340+
name: "test-outer-with-shadowing",
341+
query,
342+
schema: z.object({ run_id: z.string() }),
343+
enforcedWhereClause,
344+
tableSchema: [taskRunsSchema],
345+
});
346+
347+
expect(error?.message).toMatch(/conflicts with the table alias/);
348+
expect(result).toBeNull();
349+
}
350+
351+
const [withError, withResult] = await executeTSQL(client, {
352+
name: "test-with-expression-alias",
353+
query: "WITH 1 AS x SELECT x",
354+
schema: z.object({ x: z.coerce.number() }),
355+
enforcedWhereClause,
356+
tableSchema: [taskRunsSchema],
357+
});
358+
359+
expect(withError).toBeNull();
360+
expect(withResult?.rows).toEqual([{ x: 1 }]);
361+
362+
const [tableError, tableResult] = await executeTSQL(client, {
363+
name: "test-with-expression-alias-and-table",
364+
query: "WITH 1 AS x SELECT run_id, x FROM task_runs WHERE x = 1",
365+
schema: z.object({ run_id: z.string(), x: z.coerce.number() }),
366+
enforcedWhereClause,
367+
tableSchema: [taskRunsSchema],
368+
});
369+
370+
expect(tableError).toBeNull();
371+
expect(tableResult?.rows).toEqual([{ run_id: "run_tenant1_a", x: 1 }]);
372+
}
373+
);
374+
298375
clickhouseTest(
299376
"should not allow cross-tenant access even with malicious WHERE",
300377
async ({ clickhouseContainer }) => {

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

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1506,6 +1506,21 @@ describe("ClickHousePrinter", () => {
15061506
});
15071507
});
15081508

1509+
describe("WITH clauses", () => {
1510+
it("should print a column CTE as an expression alias", () => {
1511+
const { sql } = printQuery("WITH 1 AS x SELECT x");
1512+
1513+
expect(sql).toBe("WITH 1 AS x SELECT x LIMIT 10000");
1514+
});
1515+
1516+
it("should let a column CTE be selected alongside a table", () => {
1517+
const { sql } = printQuery("WITH 1 AS x SELECT id, x FROM task_runs WHERE x = 1");
1518+
1519+
expect(sql).toMatch(/^WITH 1 AS x SELECT id, x FROM /);
1520+
expect(sql).toContain("equals(x, 1)");
1521+
});
1522+
});
1523+
15091524
describe("UNION queries", () => {
15101525
it("should print UNION ALL", () => {
15111526
const { sql } = printQuery(`
@@ -3000,6 +3015,25 @@ describe("Unknown column blocking", () => {
30003015
}).toThrow(QueryError);
30013016
});
30023017

3018+
it("should resolve WITH expression aliases against the table schema", () => {
3019+
const schema = createSchemaRegistry([runsSchema]);
3020+
const ctx = createPrinterContext({
3021+
schema,
3022+
enforcedWhereClause: {
3023+
organization_id: { op: "eq", value: "org_test" },
3024+
project_id: { op: "eq", value: "proj_test" },
3025+
environment_id: { op: "eq", value: "env_test" },
3026+
},
3027+
});
3028+
3029+
expect(() => {
3030+
printQuery("WITH created_at AS x SELECT x FROM runs", ctx);
3031+
}).toThrow(QueryError);
3032+
expect(printQuery("WITH created AS x SELECT x FROM runs", ctx).sql).toMatch(
3033+
/^WITH created_at AS x SELECT x FROM /
3034+
);
3035+
});
3036+
30033037
it("should suggest TSQL column name when user types ClickHouse column name", () => {
30043038
// The schema exposes 'created' but the internal ClickHouse column is 'created_at'
30053039
// When user types 'created_at', we should suggest 'created'
@@ -3305,6 +3339,13 @@ describe("Internal-only column blocking", () => {
33053339
}).toThrow(/not available for querying/i);
33063340
});
33073341

3342+
it("should throw error when a WITH alias reads a hidden tenant column", () => {
3343+
const ctx = createHiddenTenantContext();
3344+
expect(() => {
3345+
printQuery("WITH organization_id AS x SELECT id, x FROM hidden_tenant_runs", ctx);
3346+
}).toThrow(/not available for querying/i);
3347+
});
3348+
33083349
it("should throw error when selecting hidden tenant column (project_id)", () => {
33093350
const ctx = createHiddenTenantContext();
33103351
expect(() => {

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

Lines changed: 70 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,7 @@ export class ClickHousePrinter {
156156
print(node: SelectQuery | SelectSetQuery): PrintResult {
157157
this.outputColumns = [];
158158
this.hiddenColumns = [];
159+
this.validateColumnCtesDoNotShadowTables(node);
159160
const sql = this.visit(node);
160161
const result: PrintResult = {
161162
sql,
@@ -359,11 +360,16 @@ export class ClickHousePrinter {
359360
// Build WHERE clause starting with any existing where
360361
let where: Expression | undefined = node.where;
361362

362-
// Process CTEs
363-
const cteStrings: string[] = [];
363+
// Process subquery CTEs; column CTEs are printed after the joins below
364+
const printedCtes = new Map<string, string>();
364365
if (node.ctes) {
365366
for (const [name, cte] of Object.entries(node.ctes)) {
366-
cteStrings.push(`${this.printIdentifier(name)} AS (${this.visit(cte.expr)})`);
367+
printedCtes.set(
368+
name,
369+
cte.cte_type === "column"
370+
? ""
371+
: `${this.printIdentifier(name)} AS (${this.visit(cte.expr)})`
372+
);
367373
}
368374
}
369375

@@ -403,6 +409,21 @@ export class ClickHousePrinter {
403409
}
404410
}
405411

412+
// Process column CTEs after the joins so they resolve fields against this scope's tables,
413+
// with the same internal-only column rules as SELECT
414+
if (node.ctes) {
415+
for (const cte of Object.values(node.ctes)) {
416+
if (cte.cte_type === "column") this.selectAliases.set(cte.name.toLowerCase(), cte.name);
417+
}
418+
for (const [name, cte] of Object.entries(node.ctes)) {
419+
if (cte.cte_type !== "column") continue;
420+
this.inProjectionContext = true;
421+
printedCtes.set(name, `${this.visit(cte.expr)} AS ${this.printIdentifier(name)}`);
422+
this.inProjectionContext = false;
423+
}
424+
}
425+
const cteStrings = [...printedCtes.values()];
426+
406427
// Track if query has GROUP BY for JSON subfield type hint decisions
407428
// (ClickHouse requires .:String for Dynamic types in GROUP BY, and SELECT must match)
408429
const savedQueryHasGroupBy = this.queryHasGroupBy;
@@ -1954,6 +1975,52 @@ export class ClickHousePrinter {
19541975
}
19551976
}
19561977

1978+
/**
1979+
* Reject a column CTE (`WITH expr AS name`) whose name, or dotted prefix, matches a table
1980+
* alias anywhere in the query. ClickHouse carries a WITH alias into every branch of a UNION,
1981+
* where it would replace that branch's tenant guards, so a per-scope check is not enough.
1982+
*/
1983+
private validateColumnCtesDoNotShadowTables(root: SelectQuery | SelectSetQuery): void {
1984+
const tableAliases = new Map<string, string>();
1985+
const cteNames: string[] = [];
1986+
const seen = new WeakSet<object>();
1987+
const walk = (expr: unknown): void => {
1988+
if (expr === null || typeof expr !== "object" || seen.has(expr)) return;
1989+
seen.add(expr);
1990+
if (Array.isArray(expr)) {
1991+
for (const item of expr) walk(item);
1992+
return;
1993+
}
1994+
const candidate = expr as Partial<JoinExpr> & Partial<CTE> & { expression_type?: string };
1995+
if (candidate.expression_type === "join_expr") {
1996+
const table = candidate.table as Field | undefined;
1997+
if (table?.expression_type === "field" && typeof table.chain[0] === "string") {
1998+
const alias = candidate.alias || table.chain[0];
1999+
tableAliases.set(alias.toLowerCase(), alias);
2000+
}
2001+
} else if (candidate.expression_type === "cte" && candidate.cte_type === "column") {
2002+
if (typeof candidate.name === "string") cteNames.push(candidate.name);
2003+
}
2004+
for (const [property, value] of Object.entries(expr)) {
2005+
if (property === "type" || property === "parent") continue;
2006+
walk(value);
2007+
}
2008+
};
2009+
walk(root);
2010+
2011+
for (const name of cteNames) {
2012+
const lowerName = name.toLowerCase();
2013+
for (const [lowerTableAlias, original] of tableAliases) {
2014+
if (lowerName === lowerTableAlias || lowerName.startsWith(`${lowerTableAlias}.`)) {
2015+
throw new QueryError(
2016+
`Alias '${name}' conflicts with the table alias '${original}'. ` +
2017+
`Choose a different alias.`
2018+
);
2019+
}
2020+
}
2021+
}
2022+
}
2023+
19572024
private findShadowingAlias(
19582025
expr: unknown,
19592026
tableAliases: Map<string, string>,

‎internal-packages/tsql/src/query/security.test.ts‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,22 @@ describe("Cross-Tenant Security", () => {
233233
"a joined table's alias",
234234
`SELECT r.id, ${ownTenant} AS e FROM task_runs AS r JOIN task_events AS e ON r.id = e.run_id`,
235235
],
236+
[
237+
"a WITH alias carried into a UNION branch",
238+
`WITH ${ownTenant} AS r SELECT id FROM task_runs AS q UNION ALL SELECT id FROM task_runs AS r`,
239+
],
240+
[
241+
"a WITH alias carried into a UNION branch's implicit table alias",
242+
`WITH ${ownTenant} AS task_runs SELECT 'x' AS id UNION ALL SELECT id FROM task_runs`,
243+
],
244+
[
245+
"a WITH alias around a FROM subquery",
246+
`WITH ${ownTenant} AS r SELECT id FROM (SELECT id FROM task_runs AS r)`,
247+
],
248+
[
249+
"a dotted WITH alias around an IN subquery",
250+
"WITH 'org_tenant1' AS `task_runs.organization_id` SELECT 1 AS x WHERE 1 IN (SELECT 1 FROM task_runs)",
251+
],
236252
])("should reject %s", (_, query) => {
237253
expect(() => compile(query)).toThrow(/conflicts with the table alias/);
238254
});

0 commit comments

Comments
 (0)