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.

Author: ericallam

.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.

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

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);

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