fix: improve type inference for array of enums (#437)

* fix: improve type inference for array of enums

Enhances type inference for arrays of enums in the TypeScript generation logic, ensuring that it correctly expands to the enum's members instead of defaulting to `unknown[]`. Additionally, adds testing guidelines for running tests in the project.

* lint

Author: Newbie012

.changeset/fancy-chefs-cough.md

@@ -0,0 +1,6 @@
+---
+"@ts-safeql/eslint-plugin": patch
+"@ts-safeql/generate": patch
+---
+
+Fix inference for array of enums (it now expands to the enum's members instead of `unknown[]`).

.cursor/rules/testing.mdc

@@ -0,0 +1,29 @@
+---
+globs: *.test.ts
+alwaysApply: false
+---
+# Testing Rules
+
+When running tests in this project, you must follow these steps:
+
+1. **Navigate to the package directory**: `cd` into the specific package directory where the test is located (e.g., `cd packages/eslint-plugin`).
+2. **Run the test using pnpm and vitest**: Use the following command format:
+   ```bash
+   pnpm vitest run [filename] -t "[testname]"
+   ```
+
+Replace `[filename]` with the path to the test file relative to the package root (e.g., `src/rules/check-sql.test.ts`), and `[testname]` with the specific test name or a pattern matching the test name you want to run.
+
+## Guidelines
+- Always `cd` to the package root before running `pnpm vitest`.
+- Use the `run` command for a single execution.
+- Use the `-t` flag to isolate specific tests, especially in large test files.
+
+## Example
+
+To run a specific test in `packages/eslint-plugin/src/rules/check-sql.test.ts`:
+
+```bash
+cd packages/eslint-plugin
+pnpm vitest run src/rules/check-sql.test.ts -t "my test name"
+```

packages/generate/src/ast-describe.ts

@@ -114,13 +114,14 @@ export function getASTDescription(params: ASTDescriptionOptions): {
         };
       }
 
+      const pgType = params.pgTypes.get(p.oid);
       const typeByOid = getTypeByOid(p.oid);
 
       if (typeByOid.override) {
         const baseType: ASTDescribedColumnType = {
           kind: "type",
           value: typeByOid.value,
-          type: params.pgTypes.get(p.oid)?.name ?? "unknown",
+          type: pgType?.name ?? "unknown",
         };
         return typeByOid.isArray ? { kind: "array", value: baseType } : baseType;
       }
@@ -136,25 +137,54 @@ export function getASTDescription(params: ASTDescriptionOptions): {
         return typeByBaseOid.isArray ? { kind: "array", value: baseType } : baseType;
       }
 
-      const enumValue = "oid" in p ? params.pgEnums.get(p.oid) : undefined;
+      const getEnumByOid = (oid: number): ASTDescribedColumnType | undefined => {
+        const pgEnum = params.pgEnums.get(oid);
+
+        if (pgEnum === undefined) {
+          return undefined;
+        }
 
-      if (enumValue !== undefined) {
         return {
           kind: "union",
-          value: enumValue.values.map((value) => ({
+          value: pgEnum.values.map((value) => ({
             kind: "type",
             value: `'${value}'`,
-            type: enumValue.name,
+            type: pgEnum.name,
           })),
         };
+      };
+
+      const valueAsEnum = (() => {
+        const enumType = getEnumByOid(p.oid);
+
+        if (enumType !== undefined) {
+          return enumType;
+        }
+
+        if (pgType?.typelem !== undefined && pgType.typelem !== 0) {
+          const arrayEnumValue = getEnumByOid(pgType.typelem);
+
+          if (arrayEnumValue !== undefined) {
+            return {
+              kind: "array",
+              value: arrayEnumValue,
+            } satisfies ASTDescribedColumnType;
+          }
+        }
+
+        return undefined;
+      })();
+
+      if (valueAsEnum !== undefined) {
+        return valueAsEnum;
       }
 
       const { isArray, value } = typeByBaseOid ?? typeByOid;
 
       const type: ASTDescribedColumnType = {
         kind: "type",
         value: value,
-        type: params.pgTypes.get(p.oid)?.name ?? "unknown",
+        type: pgType?.name ?? "unknown",
       };
 
       if (p.baseOid !== null) {

packages/generate/src/generate.test.ts

@@ -2482,3 +2482,29 @@ test("regression: wrong inference of nullable in aggregation", async () => {
     ],
   });
 });
+
+test("select array of enums", async () => {
+  await testQuery({
+    schema: `
+      CREATE TYPE my_enum AS ENUM ('A', 'B', 'C');
+      CREATE TABLE test_array_enum (col my_enum[] NOT NULL);
+    `,
+    query: `SELECT col FROM test_array_enum`,
+    expected: [
+      [
+        "col",
+        {
+          kind: "array",
+          value: {
+            kind: "union",
+            value: [
+              { kind: "type", value: "'A'", type: "my_enum" },
+              { kind: "type", value: "'B'", type: "my_enum" },
+              { kind: "type", value: "'C'", type: "my_enum" },
+            ],
+          },
+        },
+      ],
+    ],
+  });
+});

packages/generate/src/generate.ts

@@ -2,7 +2,6 @@ import {
   assertNever,
   defaultTypesMap,
   DuplicateColumnsError,
-  fmap,
   getOrSetFromMapWithEnabled,
   groupBy,
   IdentiferCase,
@@ -13,11 +12,11 @@ import {
 } from "@ts-safeql/shared";
 import * as LibPgQueryAST from "@ts-safeql/sql-ast";
 import { either } from "fp-ts";
+import * as parser from "libpg-query";
 import postgres from "postgres";
 import { ASTDescribedColumn, getASTDescription } from "./ast-describe";
 import { ColType } from "./utils/colTypes";
 import { FlattenedRelationWithJoins } from "./utils/get-relations-with-joins";
-import * as parser from "libpg-query";
 import { isParsedInsertResult, validateInsertResult } from "./utils/validate-insert";
 
 type JSToPostgresTypeMap = Record<string, unknown>;
@@ -458,14 +457,44 @@ function getResolvedTargetEntry(params: {
   }
 
   const pgTypeOid = params.col.introspected?.colBaseTypeOid ?? params.col.described.type;
+  const pgType = params.context.pgTypes.get(pgTypeOid);
+
+  const getEnumByOid = (oid: number): ResolvedTarget | undefined => {
+    const pgEnum = params.context.pgEnums.get(oid);
+
+    if (pgEnum === undefined) {
+      return undefined;
+    }
 
-  const valueAsEnum = fmap(
-    params.context.pgEnums.get(pgTypeOid),
-    ({ values }): ResolvedTarget => ({
+    return {
       kind: "union",
-      value: values.map((x): ResolvedTarget => ({ kind: "type", value: `'${x}'`, type: "text" })),
-    }),
-  );
+      value: pgEnum.values.map(
+        (value): ResolvedTarget => ({
+          kind: "type",
+          value: `'${value}'`,
+          type: pgEnum.name,
+        }),
+      ),
+    };
+  };
+
+  const valueAsEnum = (() => {
+    const enumTarget = getEnumByOid(pgTypeOid);
+
+    if (enumTarget !== undefined) {
+      return enumTarget;
+    }
+
+    if (pgType?.typelem !== undefined && pgType.typelem !== 0) {
+      const elemEnumTarget = getEnumByOid(pgType.typelem);
+
+      if (elemEnumTarget !== undefined) {
+        return { kind: "array", value: elemEnumTarget } satisfies ResolvedTarget;
+      }
+    }
+
+    return undefined;
+  })();
 
   const valueAsType = getTsTypeFromPgTypeOid({
     pgTypeOid: pgTypeOid,
@@ -599,13 +628,14 @@ async function getPgEnums(sql: Sql): Promise<PgEnumsMaps> {
 interface PgTypeRow {
   oid: number;
   name: ColType;
+  typelem: number;
 }
 
 export type PgTypesMap = Map<number, PgTypeRow>;
 
 async function getPgTypes(sql: Sql): Promise<PgTypesMap> {
   const rows = await sql<PgTypeRow[]>`
-        SELECT oid, typname as name FROM pg_type
+        SELECT oid, typname as name, typelem FROM pg_type
     `;
 
   const map = new Map<number, PgTypeRow>();