diff --git a/.server-changes/trql-json-extract-on-output.md b/.server-changes/trql-json-extract-on-output.md new file mode 100644 index 00000000000..4a2dcfd62bf --- /dev/null +++ b/.server-changes/trql-json-extract-on-output.md @@ -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. diff --git a/internal-packages/clickhouse/src/tsqlFunctions.test.ts b/internal-packages/clickhouse/src/tsqlFunctions.test.ts index 6959338256a..8c1f9ec4221 100644 --- a/internal-packages/clickhouse/src/tsqlFunctions.test.ts +++ b/internal-packages/clickhouse/src/tsqlFunctions.test.ts @@ -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", + ], ]); }); diff --git a/internal-packages/tsql/src/query/printer.test.ts b/internal-packages/tsql/src/query/printer.test.ts index 0efa0d34fc4..30d7b3ef82c 100644 --- a/internal-packages/tsql/src/query/printer.test.ts +++ b/internal-packages/tsql/src/query/printer.test.ts @@ -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); diff --git a/internal-packages/tsql/src/query/printer.ts b/internal-packages/tsql/src/query/printer.ts index 82d97f5491b..6f1af07d71c 100644 --- a/internal-packages/tsql/src/query/printer.ts +++ b/internal-packages/tsql/src/query/printer.ts @@ -1979,28 +1979,10 @@ export class ClickHousePrinter { CompareOperationOp.NotILike, ]; const useTextColumn = textColumnOps.includes(node.op); - const leftTextColumn = useTextColumn ? this.getTextColumnForExpression(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); - } + const leftTextColumn = useTextColumn ? this.printTextColumnReference(node.left) : null; + + // 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)); }