fix: correct type inference for CTEs with aggregations and joins (#430)

Improves the type inference logic for Common Table Expressions (CTEs) involving aggregations and joins. Fixes issues where nullable types were incorrectly inferred, especially in cases with LEFT JOINs and aggregated columns. Adds a helper function to adjust context when resolving column references from CTE sources, ensuring accurate type descriptions.

Author: Newbie012

.changeset/fast-planes-work.md

@@ -0,0 +1,5 @@
+---
+"@ts-safeql/generate": patch
+---
+
+fix: correct type inference for CTEs with aggregations and joins

packages/generate/src/ast-describe.ts

@@ -990,17 +990,43 @@ function getColumnRefOrigins({
   }
 }
 
+function getContextForColumnRef(
+  context: ASTDescriptionContext,
+  node: LibPgQueryAST.ColumnRef,
+): ASTDescriptionContext {
+  if (isColumnTableColumnRef(node.fields) || isColumnTableStarRef(node.fields)) {
+    const sourceName = node.fields[0].String.sval;
+    const source = context.resolver.sources.get(sourceName);
+    if (source?.kind === "cte") {
+      return { ...context, resolver: source.sources };
+    }
+  }
+
+  return context;
+}
+
 function getDescribedColumnRef({
   alias,
   context,
   node,
 }: GetDescribedParamsOf<LibPgQueryAST.ColumnRef>): ASTDescribedColumn[] {
-  const origins = getColumnRefOrigins({ alias, context, node })
-    ?.map((origin) => getDescribedNode({ alias, node: origin, context }))
-    .flat();
+  const definitionNodes = getColumnRefOrigins({ alias, context, node });
 
-  if (origins) return origins;
+  if (definitionNodes) {
+    const defContext = getContextForColumnRef(context, node);
+    return definitionNodes.flatMap((node) =>
+      getDescribedNode({ alias, node, context: defContext }),
+    );
+  }
 
+  return getDescribedColumnRefFromSchema({ alias, context, node });
+}
+
+function getDescribedColumnRefFromSchema({
+  alias,
+  context,
+  node,
+}: GetDescribedParamsOf<LibPgQueryAST.ColumnRef>): ASTDescribedColumn[] {
   // select *
   if (isColumnStarRef(node.fields)) {
     return getDescribedColumnByResolvedColumns({

packages/generate/src/generate.test.ts

@@ -56,7 +56,7 @@ function runMigrations(sql: SQL) {
       id INTEGER PRIMARY KEY GENERATED ALWAYS AS IDENTITY,
       nullable_col TEXT
     );
-    
+
     CREATE TYPE overriden_enum AS ENUM ('foo', 'bar');
     
     CREATE TABLE test_overriden_enum (
@@ -2443,3 +2443,42 @@ test("nullable columns in INNER JOIN subselect should remain nullable", async ()
     ],
   });
 });
+
+test("regression: wrong inference of nullable in aggregation", async () => {
+  await testQuery({
+    query: `
+      with subquery as (
+        select a_id, array_agg(b_id) as list
+        from b
+        group by a_id
+      )
+      select subquery.list
+      from a
+      left join subquery on (subquery.a_id = a.id);
+    `,
+    schema: `
+      CREATE TABLE a (
+        id int primary key,
+        name text not null default ''
+      );
+ 
+      CREATE TABLE b (
+        a_id int not null,
+        b_id int not null,
+        primary key (a_id, b_id)
+      );
+    `,
+    expected: [
+      [
+        "list",
+        {
+          kind: "union",
+          value: [
+            { kind: "array", value: { kind: "type", value: "number", type: "int4" } },
+            { kind: "type", value: "null", type: "null" },
+          ],
+        },
+      ],
+    ],
+  });
+});