fix jsonb operator nullability inference (#408)

Author: Newbie012

.changeset/clever-rockets-act.md

@@ -0,0 +1,5 @@
+---
+"@ts-safeql/generate": patch
+---
+
+Fixed jsonb `->>` and `#>>` operators to correctly infer nullable types when left expressions contain column references.

packages/generate/src/ast-describe.ts

@@ -308,6 +308,9 @@ function getDescribedAExpr({
           return { value, nullable, array: false };
         }
 
+        case type.kind === "object":
+          return { value: "jsonb", array: false, nullable: false };
+
         default:
           return null;
       }
@@ -349,8 +352,23 @@ function getDescribedAExpr({
     return [adjust(left), operator, adjust(right)];
   };
 
+  const getNullable = () => {
+    if (context.nonNullableColumns.has(name)) {
+      return false;
+    }
+
+    if (lnode.nullable || rnode.nullable) {
+      return true;
+    }
+
+    const operatorForcesNullable =
+      ["->>", "#>>"].includes(operator) && hasColumnReference(node.lexpr);
+
+    return operatorForcesNullable;
+  };
+
   const getType = (): ASTDescribedColumnType | undefined => {
-    const nullable = !context.nonNullableColumns.has(name) && (lnode.nullable || rnode.nullable);
+    const nullable = getNullable();
     const [dleft, doperator, dright] = downcast();
 
     const type =
@@ -1087,6 +1105,70 @@ function concatStringNodes(nodes: LibPgQueryAST.Node[] | undefined): string {
   );
 }
 
+function hasColumnReference(node: LibPgQueryAST.Node | undefined): boolean {
+  if (node === undefined) {
+    return false;
+  }
+
+  if (node.ColumnRef !== undefined) {
+    return true;
+  }
+
+  if (node.A_Const !== undefined) {
+    return false;
+  }
+
+  if (node.TypeCast !== undefined) {
+    return hasColumnReference(node.TypeCast.arg);
+  }
+
+  if (node.A_Expr !== undefined) {
+    return hasColumnReference(node.A_Expr.lexpr) || hasColumnReference(node.A_Expr.rexpr);
+  }
+
+  if (node.BoolExpr !== undefined) {
+    return (node.BoolExpr.args ?? []).some(hasColumnReference);
+  }
+
+  if (node.FuncCall !== undefined) {
+    return (node.FuncCall.args ?? []).some(hasColumnReference);
+  }
+
+  if (node.CoalesceExpr !== undefined) {
+    return node.CoalesceExpr.args.some(hasColumnReference);
+  }
+
+  if (node.CaseExpr !== undefined) {
+    if (node.CaseExpr.arg && hasColumnReference(node.CaseExpr.arg)) {
+      return true;
+    }
+
+    if (node.CaseExpr.defresult && hasColumnReference(node.CaseExpr.defresult)) {
+      return true;
+    }
+
+    return node.CaseExpr.args.some((caseWhen) => {
+      if (caseWhen.CaseWhen === undefined) {
+        return false;
+      }
+
+      return (
+        hasColumnReference(caseWhen.CaseWhen.expr) || hasColumnReference(caseWhen.CaseWhen.result)
+      );
+    });
+  }
+
+  if (node.A_ArrayExpr !== undefined) {
+    return (node.A_ArrayExpr.elements ?? []).some(hasColumnReference);
+  }
+
+  if (node.SubLink !== undefined) {
+    return true;
+  }
+
+  return false;
+}
+
 type GetDescribedParamsOf<T> = {
   alias: string | undefined;
   node: T;

packages/generate/src/generate.test.ts

@@ -2034,6 +2034,57 @@ test(`'{"a": {"b": 1}}'::jsonb ->> 'a' => string`, async () => {
   });
 });
 
+test(`'{"a": {"b": 1}}'::jsonb #>> '{a,b}' => string`, async () => {
+  await testQuery({
+    query: `SELECT '{"a": {"b": 1}}'::jsonb #>> '{a,b}'`,
+    expected: [["?column?", { kind: "type", value: "string", type: "text" }]],
+  });
+});
+
+test(`jsonb subselect ->> key => string | null`, async () => {
+  await testQuery({
+    query: `SELECT (SELECT data FROM employee LIMIT 1) ->> 'myKey' as extracted_value`,
+    expected: [
+      [
+        "extracted_value",
+        {
+          kind: "union",
+          value: [
+            { kind: "type", value: "string", type: "text" },
+            { kind: "type", value: "null", type: "null" },
+          ],
+        },
+      ],
+    ],
+    unknownColumns: ["extracted_value"],
+  });
+});
+
+test(`jsonb_build_object with column ->> key => string | null`, async () => {
+  await testQuery({
+    query: `SELECT jsonb_build_object('name', caregiver.last_name) ->> 'name' as extracted_value FROM caregiver`,
+    expected: [
+      [
+        "extracted_value",
+        {
+          kind: "union",
+          value: [
+            { kind: "type", value: "string", type: "text" },
+            { kind: "type", value: "null", type: "null" },
+          ],
+        },
+      ],
+    ],
+  });
+});
+
+test(`jsonb_build_object without column ->> key => string`, async () => {
+  await testQuery({
+    query: `SELECT jsonb_build_object('name', 'value') ->> 'name'`,
+    expected: [["?column?", { kind: "type", value: "string", type: "text" }]],
+  });
+});
+
 test(`'{"a": 1, "b": 2}'::jsonb #- '{a}' => jsonb`, async () => {
   await testQuery({
     query: `SELECT '{"a": 1, "b": 2}'::jsonb #- '{a}'`,
@@ -2314,3 +2365,21 @@ test("varchar not like expr", async () => {
     ],
   });
 });
+
+test("jsonb ->> operator should return string | null", async () => {
+  await testQuery({
+    query: `SELECT data->>'myKey' as extracted_value FROM employee`,
+    expected: [
+      [
+        "extracted_value",
+        {
+          kind: "union",
+          value: [
+            { kind: "type", value: "string", type: "text" },
+            { kind: "type", value: "null", type: "null" },
+          ],
+        },
+      ],
+    ],
+  });
+});