fix: scalar subquery in select list should infer correct type (#412)

Author: Newbie012

.changeset/wet-ears-cross.md

@@ -0,0 +1,5 @@
+---
+"@ts-safeql/generate": patch
+---
+
+Fixed type inference for scalar subqueries in select lists.

packages/generate/src/ast-describe.ts

@@ -10,15 +10,18 @@ import {
 } from "./ast-decribe.utils";
 import { ResolvedColumn, SourcesResolver, getSources } from "./ast-get-sources";
 import { PgColRow, PgEnumsMaps, PgTypesMap } from "./generate";
-import { FlattenedRelationWithJoins } from "./utils/get-relations-with-joins";
+import { getNonNullableColumns } from "./utils/get-nonnullable-columns";
+import {
+  FlattenedRelationWithJoins,
+  flattenRelationsWithJoinsMap,
+  getRelationsWithJoins,
+} from "./utils/get-relations-with-joins";
 
 type ASTDescriptionOptions = {
   parsed: LibPgQueryAST.ParseResult;
-  relations: FlattenedRelationWithJoins[];
   typesMap: Map<string, { override: boolean; value: string }>;
   typeExprMap: Map<string, Map<string, Map<string, string>>>;
   overridenColumnTypesMap: Map<string, Map<string, string>>;
-  nonNullableColumns: Set<string>;
   pgColsBySchemaAndTableName: Map<string, Map<string, PgColRow[]>>;
   pgTypes: PgTypesMap;
   pgEnums: PgEnumsMaps;
@@ -29,6 +32,8 @@ type ASTDescriptionContext = ASTDescriptionOptions & {
   select: LibPgQueryAST.SelectStmt;
   resolver: SourcesResolver;
   resolved: WeakMap<LibPgQueryAST.Node, string>;
+  nonNullableColumns: Set<string>;
+  relations: FlattenedRelationWithJoins[];
   toTypeScriptType: (
     params: { oid: number; baseOid: number | null } | { name: string },
   ) => ASTDescribedColumnType;
@@ -43,15 +48,28 @@ export type ASTDescribedColumnType =
   | { kind: "type"; value: string; type: string; base?: string }
   | { kind: "literal"; value: string; base: ASTDescribedColumnType };
 
-export function getASTDescription(
-  params: ASTDescriptionOptions,
-): Map<number, ASTDescribedColumn | undefined> {
+export function getASTDescription(params: ASTDescriptionOptions): {
+  map: Map<number, ASTDescribedColumn | undefined>;
+  meta: {
+    relations: FlattenedRelationWithJoins[];
+    nonNullableColumns: Set<string>;
+  };
+} {
   const select = params.parsed.stmts[0]?.stmt?.SelectStmt;
 
   if (select === undefined) {
-    return new Map();
+    return {
+      map: new Map(),
+      meta: {
+        relations: [],
+        nonNullableColumns: new Set(),
+      },
+    };
   }
 
+  const nonNullableColumns = getNonNullableColumns(params.parsed);
+  const relations = flattenRelationsWithJoinsMap(getRelationsWithJoins(params.parsed));
+
   function getTypeByOid(oid: number) {
     const name = params.pgTypes.get(oid)?.name;
 
@@ -74,10 +92,12 @@ export function getASTDescription(
 
   const context: ASTDescriptionContext = {
     ...params,
+    nonNullableColumns,
+    relations,
     resolver: getSources({
-      relations: params.relations,
+      relations: relations,
       select: select,
-      nonNullableColumns: params.nonNullableColumns,
+      nonNullableColumns: nonNullableColumns,
       pgColsBySchemaAndTableName: params.pgColsBySchemaAndTableName,
     }),
     select: select,
@@ -176,7 +196,13 @@ export function getASTDescription(
     final.set(i, result);
   }
 
-  return final;
+  return {
+    map: final,
+    meta: {
+      relations,
+      nonNullableColumns,
+    },
+  };
 }
 
 function mergeColumns(columns: (ASTDescribedColumn | undefined)[]): ASTDescribedColumn | undefined {
@@ -247,6 +273,10 @@ function getDescribedNode(params: {
     return getDescribedAExpr({ alias: alias, node: node.A_Expr, context });
   }
 
+  if (node.SelectStmt !== undefined) {
+    return getDescribedSelectStmt({ alias: alias, node: node.SelectStmt, context });
+  }
+
   return [];
 }
 
@@ -488,24 +518,74 @@ function getDescribedSubLink({
   context,
   node,
 }: GetDescribedParamsOf<LibPgQueryAST.SubLink>): ASTDescribedColumn[] {
+  const getSubLinkType = (): ASTDescribedColumnType => {
+    if (node.subLinkType === LibPgQueryAST.SubLinkType.EXISTS_SUBLINK) {
+      return context.toTypeScriptType({ name: "bool" });
+    }
+
+    if (node.subLinkType === LibPgQueryAST.SubLinkType.EXPR_SUBLINK) {
+      const described = node.subselect?.SelectStmt
+        ? getDescribedNode({
+            alias: undefined,
+            node: { SelectStmt: node.subselect.SelectStmt },
+            context,
+          })
+        : [];
+
+      return described.length > 0
+        ? described[0].type
+        : context.toTypeScriptType({ name: "unknown" });
+    }
+
+    return context.toTypeScriptType({ name: "unknown" });
+  };
+
   return [
     {
       name: alias ?? "exists",
       type: resolveType({
         context: context,
         nullable: false,
-        type: (() => {
-          if (node.subLinkType === LibPgQueryAST.SubLinkType.EXISTS_SUBLINK) {
-            return context.toTypeScriptType({ name: "bool" });
-          }
-
-          return context.toTypeScriptType({ name: "unknown" });
-        })(),
+        type: getSubLinkType(),
       }),
     },
   ];
 }
 
+function getDescribedSelectStmt({
+  alias,
+  context,
+  node,
+}: GetDescribedParamsOf<LibPgQueryAST.SelectStmt>): ASTDescribedColumn[] {
+  const subParsed: LibPgQueryAST.ParseResult = {
+    version: 0,
+    stmts: [{ stmt: { SelectStmt: node }, stmtLocation: 0, stmtLen: 0 }],
+  };
+
+  const subDescription = getASTDescription({
+    parsed: subParsed,
+    typesMap: context.typesMap,
+    typeExprMap: context.typeExprMap,
+    overridenColumnTypesMap: context.overridenColumnTypesMap,
+    pgColsBySchemaAndTableName: context.pgColsBySchemaAndTableName,
+    pgTypes: context.pgTypes,
+    pgEnums: context.pgEnums,
+    pgFns: context.pgFns,
+  });
+
+  const firstColumn = subDescription.map.get(0);
+  if (firstColumn) {
+    return [
+      {
+        name: alias ?? firstColumn.name,
+        type: firstColumn.type,
+      },
+    ];
+  }
+
+  return [];
+}
+
 function getDescribedCoalesceExpr({
   alias,
   context,

packages/generate/src/generate.test.ts

@@ -2056,7 +2056,6 @@ test(`jsonb subselect ->> key => string | null`, async () => {
         },
       ],
     ],
-    unknownColumns: ["extracted_value"],
   });
 });
 
@@ -2383,3 +2382,40 @@ test("jsonb ->> operator should return string | null", async () => {
     ],
   });
 });
+
+test("scalar subquery in select list should infer correct type", async () => {
+  await testQuery({
+    schema: `
+      CREATE TABLE tbl (
+        id INTEGER PRIMARY KEY,
+        col TEXT
+      );
+    `,
+    query: `SELECT (SELECT col FROM tbl LIMIT 1) AS col`,
+    expected: [
+      [
+        "col",
+        {
+          kind: "union",
+          value: [
+            { kind: "type", value: "string", type: "text" },
+            { kind: "type", value: "null", type: "null" },
+          ],
+        },
+      ],
+    ],
+  });
+});
+
+test("scalar subquery with WHERE should infer non-nullable type", async () => {
+  await testQuery({
+    schema: `
+      CREATE TABLE tbl (
+        id INTEGER PRIMARY KEY,
+        col TEXT
+      );
+    `,
+    query: `SELECT (SELECT col FROM tbl WHERE col IS NOT NULL LIMIT 1) AS col`,
+    expected: [["col", { kind: "type", value: "string", type: "text" }]],
+  });
+});

packages/generate/src/generate.ts

@@ -16,12 +16,7 @@ import { either } from "fp-ts";
 import postgres from "postgres";
 import { ASTDescribedColumn, getASTDescription } from "./ast-describe";
 import { ColType } from "./utils/colTypes";
-import { getNonNullableColumns } from "./utils/get-nonnullable-columns";
-import {
-  FlattenedRelationWithJoins,
-  flattenRelationsWithJoinsMap,
-  getRelationsWithJoins,
-} from "./utils/get-relations-with-joins";
+import { FlattenedRelationWithJoins } from "./utils/get-relations-with-joins";
 import * as parser from "libpg-query";
 
 type JSToPostgresTypeMap = Record<string, unknown>;
@@ -262,15 +257,11 @@ async function generate(
     }
 
     const parsed = await parser.parse(query.text);
-    const relationsWithJoins = flattenRelationsWithJoinsMap(getRelationsWithJoins(parsed));
-    const nonNullableColumnsBasedOnAST = getNonNullableColumns(parsed);
 
     const astQueryDescription = getASTDescription({
       parsed: parsed,
-      relations: relationsWithJoins,
       typesMap: typesMap,
       overridenColumnTypesMap: overridenColumnTypesMap,
-      nonNullableColumns: nonNullableColumnsBasedOnAST,
       pgColsBySchemaAndTableName: pgColsBySchemaAndTableName,
       pgTypes: pgTypes,
       pgEnums: pgEnums,
@@ -283,21 +274,21 @@ async function generate(
         .get(col.table)
         ?.find((x) => x.colNum === col.number);
 
-      const astDescribed = astQueryDescription.get(position);
+      const astDescribed = astQueryDescription.map.get(position);
 
       return {
         described: col,
         astDescribed: astDescribed,
         introspected: introspected,
-        isNonNullableBasedOnAST: nonNullableColumnsBasedOnAST.has(col.name),
+        isNonNullableBasedOnAST: astQueryDescription.meta.nonNullableColumns.has(col.name),
       };
     });
 
     const context: GenerateContext = {
       columns,
       pgTypes,
       pgEnums,
-      relationsWithJoins,
+      relationsWithJoins: astQueryDescription.meta.relations,
       overrides: {
         types: typesMap,
         columns: overridenColumnTypesMap,