fix(tsql): make JSON functions work on the output and error columns (#4221)

## Summary

A Query page (TRQL) query that pulls fields out of a run's `output` with
JSON functions (`JSONExtractString`, `JSONExtractInt`, `JSONHas`, and
the rest of the family) failed with "The first argument of function ...
should be a string containing JSON, illegal type: JSON". Those queries
now work.

## Root cause and fix

`output` is a native ClickHouse `JSON` column, but `JSONExtract*`,
`JSONHas`, `JSONLength`, and `JSONType` all expect a String containing
JSON text. The compiler already swaps in the column's String companion
(`output_text`) when a JSON column is selected or compared, but not
inside function-call arguments, so it emitted `JSONExtractInt(output,
'x')` against the native column.

The fix prints the companion column for the first argument of these
functions when it resolves to a bare JSON field, keeping the table alias
when qualified (so it works in JOINs):

JSONExtractInt(output, 'x') -> JSONExtractInt(output_text, 'x')
JSONExtractArrayRaw(assumeNotNull(output), 'y') ->
JSONExtractArrayRaw(assumeNotNull(output_text), 'y')

It also reaches through value-preserving passthrough wrappers like
`assumeNotNull(...)`, while leaving value-changing wrappers like
`toJSONString(output)` on the native column (that argument is already a
String). The swap is also semantically correct, not just a type fix:
`output_text` is the unwrapped data JSON that the TRQL `output` model
already represents, so field paths line up.

Covered by printer unit tests and a ClickHouse integration test that
runs the whole family (plus the wrapped and `toJSONString` cases)
against a real native-JSON column. Both new cases fail with the exact
"illegal type: JSON" error without the fix.
This commit is contained in:
Eric Allam
2026-07-10 12:44:52 +01:00
committed by GitHub
parent de536622c8
commit 02cf9c81ad
4 changed files with 206 additions and 22 deletions
@@ -0,0 +1,6 @@
---
area: webapp
type: fix
---
Query page: extracting fields from a run's output with JSON functions (such as JSONExtractString or JSONExtractInt) no longer fails with an "illegal type: JSON" error.
@@ -28,6 +28,13 @@ const taskRunsSchema: TableSchema = {
completed_at: { name: "completed_at", ...column("Nullable(DateTime64)") },
is_test: { name: "is_test", ...column("UInt8") },
tags: { name: "tags", ...column("Array(String)") },
output: {
name: "output",
...column("JSON"),
nullValue: "'{}'",
textColumn: "output_text",
dataPrefix: "data",
},
usage_duration_ms: { name: "usage_duration_ms", ...column("UInt32") },
cost_in_cents: { name: "cost_in_cents", ...column("Float64") },
attempt: { name: "attempt", ...column("UInt8") },
@@ -65,7 +72,7 @@ const defaultTaskRun = {
started_at: Date.now() - 5000,
completed_at: Date.now(),
tags: ["tag-a", "tag-b"],
output: null,
output: { data: { count: 42, label: "ok", ratio: 1.5, enabled: true, items: [1, 2, 3] } },
error: null,
usage_duration_ms: 4500,
cost_in_cents: 1.5,
@@ -572,6 +579,33 @@ describe("TSQL Function Smoke Tests", () => {
],
["JSONExtractKeys", `SELECT JSONExtractKeys('{"a": 1, "b": 2}') AS r FROM task_runs`],
["toJSONString", "SELECT toJSONString(map('a', 1)) AS r FROM task_runs"],
// The `output` column is a native JSON type, so JSON functions must read its
// String companion (output_text). Without the printer fix these fail with
// "should be a string containing JSON, illegal type: JSON".
["JSONHas(output)", "SELECT JSONHas(output, 'count') AS r FROM task_runs"],
["JSONLength(output)", "SELECT JSONLength(output) AS r FROM task_runs"],
["JSONType(output)", "SELECT JSONType(output, 'count') AS r FROM task_runs"],
["JSONExtractInt(output)", "SELECT JSONExtractInt(output, 'count') AS r FROM task_runs"],
["JSONExtractUInt(output)", "SELECT JSONExtractUInt(output, 'count') AS r FROM task_runs"],
["JSONExtractFloat(output)", "SELECT JSONExtractFloat(output, 'ratio') AS r FROM task_runs"],
["JSONExtractBool(output)", "SELECT JSONExtractBool(output, 'enabled') AS r FROM task_runs"],
[
"JSONExtractString(output)",
"SELECT JSONExtractString(output, 'label') AS r FROM task_runs",
],
["JSONExtractRaw(output)", "SELECT JSONExtractRaw(output, 'count') AS r FROM task_runs"],
["JSONExtractKeys(output)", "SELECT JSONExtractKeys(output) AS r FROM task_runs"],
// The JSON field can be wrapped in a passthrough like assumeNotNull; the swap
// still has to reach the native column underneath.
[
"JSONExtractArrayRaw(assumeNotNull(output))",
"SELECT JSONExtractArrayRaw(assumeNotNull(output), 'items') AS r FROM task_runs",
],
// toJSONString(output) is already a String, so it stays on the native column.
[
"JSONExtractString(toJSONString(output))",
"SELECT JSONExtractString(toJSONString(output), 'data', 'label') AS r FROM task_runs",
],
]);
});
@@ -823,6 +823,89 @@ describe("ClickHousePrinter", () => {
expect(sql).toContain("output_text AS result");
});
it("should use text column as the first arg of JSONExtract* on a bare JSON column", () => {
const ctx = createTextColumnContext();
const { sql } = printQuery("SELECT JSONExtractInt(output, 'count') AS r FROM runs", ctx);
// JSONExtract* needs a String containing JSON, not the native JSON column.
expect(sql).toContain("JSONExtractInt(output_text,");
expect(sql).not.toMatch(/JSONExtractInt\(output,/);
});
it("should use text column for the whole JSONExtract* family and JSONHas/Length/Type", () => {
const ctx = createTextColumnContext();
const cases: Array<[string, string]> = [
["JSONExtract", "SELECT JSONExtract(output, 'x', 'Int64') AS r FROM runs"],
["JSONExtractUInt", "SELECT JSONExtractUInt(output, 'x') AS r FROM runs"],
["JSONExtractFloat", "SELECT JSONExtractFloat(output, 'x') AS r FROM runs"],
["JSONExtractBool", "SELECT JSONExtractBool(output, 'x') AS r FROM runs"],
["JSONExtractString", "SELECT JSONExtractString(output, 'x') AS r FROM runs"],
["JSONExtractRaw", "SELECT JSONExtractRaw(output, 'x') AS r FROM runs"],
["JSONExtractArrayRaw", "SELECT JSONExtractArrayRaw(output, 'x') AS r FROM runs"],
[
"JSONExtractKeysAndValues",
"SELECT JSONExtractKeysAndValues(output, 'String') AS r FROM runs",
],
["JSONExtractKeys", "SELECT JSONExtractKeys(output) AS r FROM runs"],
["JSONHas", "SELECT JSONHas(output, 'x') AS r FROM runs"],
["JSONLength", "SELECT JSONLength(output) AS r FROM runs"],
["JSONType", "SELECT JSONType(output, 'x') AS r FROM runs"],
];
for (const [fn, query] of cases) {
const { sql } = printQuery(query, ctx);
expect(sql, `${fn} should read from output_text`).toContain(`${fn}(output_text`);
expect(sql, `${fn} should not read from the native JSON column`).not.toMatch(
new RegExp(`${fn}\\(output,`)
);
}
});
it("should NOT rewrite a JSONExtract* first arg that is a String literal", () => {
const ctx = createTextColumnContext();
const { sql } = printQuery(`SELECT JSONExtractInt('{"a": 1}', 'a') AS r FROM runs`, ctx);
// String literals are parameterized, never swapped for the text column.
expect(sql).toMatch(/JSONExtractInt\(\{tsql_val_\d+: String\}/);
expect(sql).not.toContain("output_text");
});
it("should qualify the text column with the table alias inside JSONExtract*", () => {
const ctx = createTextColumnContext();
const { sql } = printQuery(
"SELECT JSONExtractInt(runs.output, 'count') AS r FROM runs",
ctx
);
expect(sql).toContain("JSONExtractInt(runs.output_text,");
});
it("should reach through assumeNotNull() to swap the JSON field for its text column", () => {
const ctx = createTextColumnContext();
const { sql } = printQuery(
"SELECT JSONExtractArrayRaw(assumeNotNull(output), 'losers') AS r FROM runs",
ctx
);
// The JSON field is wrapped in a value-preserving passthrough, so the swap
// has to descend into it: assumeNotNull(output) -> assumeNotNull(output_text).
expect(sql).toContain("JSONExtractArrayRaw(assumeNotNull(output_text),");
expect(sql).not.toContain("assumeNotNull(output)");
});
it("should NOT rewrite a JSON field consumed by toJSONString inside JSONExtract*", () => {
const ctx = createTextColumnContext();
const { sql } = printQuery(
"SELECT JSONExtractString(toJSONString(output), 'x') AS r FROM runs",
ctx
);
// toJSONString(output) is already a String and changes the value, so it must
// stay on the native column.
expect(sql).toContain("toJSONString(output)");
expect(sql).not.toContain("output_text");
});
it("should use text column for table-qualified JSON columns in SELECT", () => {
const ctx = createTextColumnContext();
const { sql } = printQuery("SELECT runs.output FROM runs", ctx);
+82 -21
View File
@@ -1979,28 +1979,10 @@ export class ClickHousePrinter {
CompareOperationOp.NotILike,
];
const useTextColumn = textColumnOps.includes(node.op);
const leftTextColumn = useTextColumn ? this.getTextColumnForExpression(node.left) : null;
const leftTextColumn = useTextColumn ? this.printTextColumnReference(node.left) : null;
// Build the left side, qualifying the text column with table alias if present
let left: string;
if (leftTextColumn) {
// Check if the field is qualified with a table alias (e.g., r.output)
// and prepend that alias to the text column to avoid ambiguity in JOINs
const fieldNode = node.left as Field;
if (fieldNode.expression_type === "field" && fieldNode.chain.length >= 2) {
const firstPart = fieldNode.chain[0];
if (typeof firstPart === "string" && this.tableContexts.has(firstPart)) {
// The field is qualified with a table alias, prepend it to the text column
left = this.printIdentifier(firstPart) + "." + this.printIdentifier(leftTextColumn);
} else {
left = this.printIdentifier(leftTextColumn);
}
} else {
left = this.printIdentifier(leftTextColumn);
}
} else {
left = this.visit(node.left);
}
// Build the left side, using the (alias-qualified) text column when available
const left = leftTextColumn ?? this.visit(node.left);
const right = this.visit(transformedRight);
switch (node.op) {
@@ -2529,6 +2511,49 @@ export class ClickHousePrinter {
return this.getTextColumnForField((expr as Field).chain);
}
/**
* Print a bare JSON field as its String companion (textColumn), keeping the table alias
* when qualified (r.output -> r.output_text). Returns null if there is no textColumn.
*/
private printTextColumnReference(expr: Expression): string | null {
const textColumn = this.getTextColumnForExpression(expr);
if (!textColumn) return null;
const fieldNode = expr as Field;
if (fieldNode.expression_type === "field" && fieldNode.chain.length >= 2) {
const firstPart = fieldNode.chain[0];
if (typeof firstPart === "string" && this.tableContexts.has(firstPart)) {
return this.printIdentifier(firstPart) + "." + this.printIdentifier(textColumn);
}
}
return this.printIdentifier(textColumn);
}
/**
* Print a JSON-string function argument, swapping a bare JSON field for its textColumn.
* Descends through value-preserving passthrough wrappers (e.g. assumeNotNull(output))
* so the swap still applies inside them. Returns null if no textColumn swap applies.
*/
private substituteJsonTextArg(expr: Expression): string | null {
const direct = this.printTextColumnReference(expr);
if (direct) return direct;
const call = expr as Call;
if (
call.expression_type === "call" &&
ClickHousePrinter.JSON_PASSTHROUGH_WRAPPERS.has(call.name.toLowerCase()) &&
call.args.length > 0
) {
const inner = this.substituteJsonTextArg(call.args[0]);
if (inner) {
const clickhouseName = findTSQLFunction(call.name)?.clickhouseName ?? call.name;
const rest = call.args.slice(1).map((arg) => this.visit(arg));
return `${clickhouseName}(${[inner, ...rest].join(", ")})`;
}
}
return null;
}
/**
* Get the dataPrefix for a field chain if the root column has one defined.
* Returns null if the column doesn't have a dataPrefix or if this isn't a subfield access.
@@ -2981,6 +3006,35 @@ export class ClickHousePrinter {
"date_diff",
]);
/**
* ClickHouse JSON functions whose first argument must be a String containing JSON text.
* Passing a native JSON column fails with "should be a string containing JSON, illegal
* type: JSON", so we print the column's String companion (textColumn) for that argument.
*/
private static readonly JSON_TEXT_ARG_FUNCTIONS = new Set([
"jsonhas",
"jsonlength",
"jsontype",
"jsonextract",
"jsonextractuint",
"jsonextractint",
"jsonextractfloat",
"jsonextractbool",
"jsonextractstring",
"jsonextractraw",
"jsonextractarrayraw",
"jsonextractkeysandvalues",
"jsonextractkeys",
]);
/**
* Value-preserving wrappers a user may put between a JSON column and a JSON function,
* e.g. `JSONExtractString(assumeNotNull(output), 'x')`. The text-column swap descends
* through these. Functions that transform the value (e.g. toJSONString) are excluded so
* their argument keeps reading the native column.
*/
private static readonly JSON_PASSTHROUGH_WRAPPERS = new Set(["assumenotnull"]);
/**
* Visit function call arguments, handling date functions that require an interval unit
* keyword as their first argument. For these functions, the first arg is output as a
@@ -2998,6 +3052,13 @@ export class ClickHousePrinter {
}
}
if (ClickHousePrinter.JSON_TEXT_ARG_FUNCTIONS.has(lowerName) && args.length > 0) {
const textColumn = this.substituteJsonTextArg(args[0]);
if (textColumn) {
return [textColumn, ...args.slice(1).map((arg) => this.visit(arg))];
}
}
return args.map((arg) => this.visit(arg));
}