From 7fc74d93090d7248e0e00f44496e67fde4e3f788 Mon Sep 17 00:00:00 2001
From: Eliya Cohen <co.eliya2@gmail.com>
Date: Mon, 22 Dec 2025 11:10:31 +0200
Subject: [PATCH] FEAT feat: Support unnest function for array expansion in
 type inference

---
 .changeset/slimy-walls-prove.md        |  5 ++
 packages/generate/src/ast-describe.ts  | 45 ++++++++++++++++
 packages/generate/src/generate.test.ts | 75 ++++++++++++++++++++++++++
 3 files changed, 125 insertions(+)
 create mode 100644 .changeset/slimy-walls-prove.md

diff --git a/.changeset/slimy-walls-prove.md b/.changeset/slimy-walls-prove.md
new file mode 100644
index 00000000..81c60cff
--- /dev/null
+++ b/.changeset/slimy-walls-prove.md
@@ -0,0 +1,5 @@
+---
+"@ts-safeql/generate": minor
+---
+
+Added support for `unnest` function in type inference. This allows correct typing when expanding arrays into rows, including support for multidimensional arrays and enums.
diff --git a/packages/generate/src/ast-describe.ts b/packages/generate/src/ast-describe.ts
index 09fa8f1e..35c42b02 100644
--- a/packages/generate/src/ast-describe.ts
+++ b/packages/generate/src/ast-describe.ts
@@ -800,11 +800,56 @@ function getDescribedFuncCall(
       return getDescribedJsonAggFunCall(params);
     case functionName === "array_agg":
       return getDescribedArrayAggFunCall(params);
+    case functionName === "unnest":
+      return getDescribedUnnestFunCall(params);
     default:
       return getDescribedFuncCallByPgFn(params);
   }
 }
 
+function getDescribedUnnestFunCall({
+  alias,
+  context,
+  node,
+}: GetDescribedParamsOf<LibPgQueryAST.FuncCall>): ASTDescribedColumn[] {
+  const functionName = node.funcname.at(-1)?.String?.sval ?? "";
+  const name = alias ?? functionName;
+
+  const firstArg = node.args?.at(0);
+  if (firstArg === undefined) {
+    return [{ name, type: context.toTypeScriptType({ name: "unknown" }) }];
+  }
+
+  const described = getDescribedNode({ alias: undefined, node: firstArg, context }).at(0);
+  if (described === undefined) {
+    return [{ name, type: context.toTypeScriptType({ name: "unknown" }) }];
+  }
+
+  function unwrap(type: ASTDescribedColumnType): ASTDescribedColumnType {
+    if (type.kind === "array") {
+      return unwrap(type.value);
+    }
+
+    if (type.kind === "union") {
+      const hasArray = type.value.some((t) => t.kind === "array");
+
+      const unwrapped = type.value
+        .filter((t) => !hasArray || !(t.kind === "type" && t.value === "null"))
+        .map(unwrap);
+
+      return mergeDescribedColumnTypes(unwrapped);
+    }
+
+    if (type.kind === "type" && type.type.startsWith("_")) {
+      return { ...type, type: type.type.slice(1) };
+    }
+
+    return type;
+  }
+
+  return [{ name, type: unwrap(described.type) }];
+}
+
 function getDescribedFuncCallByPgFn({
   alias,
   context,
diff --git a/packages/generate/src/generate.test.ts b/packages/generate/src/generate.test.ts
index 38f6b1a0..474bbbfd 100644
--- a/packages/generate/src/generate.test.ts
+++ b/packages/generate/src/generate.test.ts
@@ -2508,3 +2508,78 @@ test("select array of enums", async () => {
     ],
   });
 });
+
+test("select unnest(array_text_column)", async () => {
+  await testQuery({
+    query: `SELECT unnest(array_text_column) as unnested_text FROM all_types`,
+    expected: [["unnested_text", { kind: "type", value: "string", type: "text" }]],
+  });
+});
+
+test("select unnest(array_enum_column)", async () => {
+  await testQuery({
+    schema: `
+      CREATE TYPE my_enum AS ENUM ('A', 'B', 'C');
+      CREATE TABLE test_unnest_enum (col my_enum[] NOT NULL);
+    `,
+    query: `SELECT unnest(col) as unnested_enum FROM test_unnest_enum`,
+    expected: [
+      [
+        "unnested_enum",
+        {
+          kind: "union",
+          value: [
+            { kind: "type", value: "'A'", type: "my_enum" },
+            { kind: "type", value: "'B'", type: "my_enum" },
+            { kind: "type", value: "'C'", type: "my_enum" },
+          ],
+        },
+      ],
+    ],
+  });
+});
+
+test("select unnest(multidimensional_array)", async () => {
+  await testQuery({
+    query: `SELECT unnest(ARRAY[[1,2],[3,4]]) as unnested_int`,
+    expected: [
+      [
+        "unnested_int",
+        {
+          kind: "union",
+          value: [
+            { kind: "literal", value: "1", base: { kind: "type", value: "number", type: "int4" } },
+            { kind: "literal", value: "2", base: { kind: "type", value: "number", type: "int4" } },
+            { kind: "literal", value: "3", base: { kind: "type", value: "number", type: "int4" } },
+            { kind: "literal", value: "4", base: { kind: "type", value: "number", type: "int4" } },
+          ],
+        },
+      ],
+    ],
+  });
+});
+
+test("select unnest(nullable_array_column)", async () => {
+  await testQuery({
+    query: `SELECT unnest(nullable_date_arr) as unnested_date FROM test_date_column`,
+    expected: [["unnested_date", { kind: "type", value: "Date", type: "date" }]],
+  });
+});
+
+test("select unnest(array_with_nulls)", async () => {
+  await testQuery({
+    query: `SELECT unnest(ARRAY[1, NULL]) as unnested_int`,
+    expected: [
+      [
+        "unnested_int",
+        {
+          kind: "union",
+          value: [
+            { kind: "literal", value: "1", base: { kind: "type", value: "number", type: "int4" } },
+            { kind: "type", value: "null", type: "null" },
+          ],
+        },
+      ],
+    ],
+  });
+});