fix(tsql): use text column for JSON functions on native JSON columns

JSONExtract*, JSONHas, JSONLength, and JSONType expect a String
containing JSON, but the compiler printed them against native JSON
columns (like a run's output), which ClickHouse rejects with
"illegal type: JSON". Print the column's String companion (its
textColumn) for the first argument of these functions when it is a
bare JSON field, mirroring the existing SELECT and WHERE behavior.

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 } },
   error: null,
   usage_duration_ms: 4500,
   cost_in_cents: 1.5,
@@ -572,6 +579,22 @@ 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"],
     ]);
   });
 

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

@@ -823,6 +823,63 @@ 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 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,24 @@ 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);
+  }
+
   /**
    * 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 +2981,27 @@ 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",
+  ]);
+
   /**
    * 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 +3019,13 @@ export class ClickHousePrinter {
       }
     }
 
+    if (ClickHousePrinter.JSON_TEXT_ARG_FUNCTIONS.has(lowerName) && args.length > 0) {
+      const textColumn = this.printTextColumnReference(args[0]);
+      if (textColumn) {
+        return [textColumn, ...args.slice(1).map((arg) => this.visit(arg))];
+      }
+    }
+
     return args.map((arg) => this.visit(arg));
   }