fix(tsql): reach through passthrough wrappers for JSON functions on native JSON columns

A JSON function whose JSON argument is wrapped in a value-preserving
passthrough (e.g. JSONExtractArrayRaw(assumeNotNull(output), 'x')) still
printed against the native JSON column and failed with "illegal type:
JSON". Descend through those wrappers so the text-column swap applies
inside them, while leaving value-changing wrappers like toJSONString on
the native column.

Author: ericallam

internal-packages/clickhouse/src/tsqlFunctions.test.ts

@@ -72,7 +72,7 @@ const defaultTaskRun = {
   started_at: Date.now() - 5000,
   completed_at: Date.now(),
   tags: ["tag-a", "tag-b"],
-  output: { data: { count: 42, label: "ok", ratio: 1.5, enabled: true } },
+  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,
@@ -595,6 +595,17 @@ describe("TSQL Function Smoke Tests", () => {
       ],
       ["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

@@ -880,6 +880,32 @@ describe("ClickHousePrinter", () => {
         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

@@ -2529,6 +2529,31 @@ export class ClickHousePrinter {
     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.
@@ -3002,6 +3027,14 @@ export class ClickHousePrinter {
     "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
@@ -3020,7 +3053,7 @@ export class ClickHousePrinter {
     }
 
     if (ClickHousePrinter.JSON_TEXT_ARG_FUNCTIONS.has(lowerName) && args.length > 0) {
-      const textColumn = this.printTextColumnReference(args[0]);
+      const textColumn = this.substituteJsonTextArg(args[0]);
       if (textColumn) {
         return [textColumn, ...args.slice(1).map((arg) => this.visit(arg))];
       }