feat: improve A_Expr inference (#387)

Author: Newbie012

.changeset/gentle-lizards-jump.md

@@ -0,0 +1,6 @@
+---
+"@ts-safeql/generate": patch
+"@ts-safeql/shared": patch
+---
+
+improved A_Expr inference

packages/generate/src/ast-describe.ts

@@ -1,4 +1,4 @@
-import { defaultTypeExprMapping, fmap, normalizeIndent } from "@ts-safeql/shared";
+import { fmap, normalizeIndent } from "@ts-safeql/shared";
 import * as LibPgQueryAST from "@ts-safeql/sql-ast";
 import {
   isColumnStarRef,
@@ -16,6 +16,7 @@ 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[]>>;
@@ -274,61 +275,119 @@ function getDescribedAExpr({
 
     if (column === undefined) return null;
 
-    if (column.type.kind === "array") {
-      return { value: "array", nullable: false };
-    }
+    const getFromType = (
+      type: ASTDescribedColumnType,
+    ): { value: string; array: boolean; nullable: boolean } | null => {
+      switch (true) {
+        case type.kind === "type":
+          return { value: type.base ?? type.type, array: false, nullable: false };
 
-    if (column.type.kind === "type") {
-      return { value: column.type.base ?? column.type.type, nullable: false };
-    }
+        case type.kind === "literal" && type.base.kind === "type":
+          return { value: type.base.type, array: false, nullable: false };
 
-    if (column.type.kind === "literal" && column.type.base.kind === "type") {
-      return { value: column.type.base.type, nullable: false };
-    }
+        case type.kind === "union" && type.value.every((x) => x.kind === "literal"): {
+          const resolved = getFromType(type.value[0].base);
+
+          if (resolved === null) return null;
+
+          return { value: resolved.value, nullable: false, array: false };
+        }
+
+        case type.kind === "union" && isTuple(type.value): {
+          let nullable = false;
+          let value: string | undefined = undefined;
+
+          for (const valueType of type.value) {
+            if (valueType.kind !== "type") return null;
+            if (valueType.value === "null") nullable = true;
+            if (valueType.value !== "null") value = valueType.type;
+          }
+
+          if (value === undefined) return null;
 
-    if (column.type.kind === "union" && isTuple(column.type.value)) {
-      let nullable = false;
-      let value: string | undefined = undefined;
+          return { value, nullable, array: false };
+        }
 
-      for (const type of column.type.value) {
-        if (type.kind !== "type") return null;
-        if (type.value === "null") nullable = true;
-        if (type.value !== "null") value = type.type;
+        default:
+          return null;
       }
+    };
+
+    if (column.type.kind === "array") {
+      const resolved = getFromType(column.type.value);
 
-      if (value === undefined) return null;
+      if (!resolved) return null;
 
-      return { value, nullable };
+      return { value: resolved.value, nullable: resolved.nullable, array: true };
     }
 
-    return null;
+    return getFromType(column.type);
   };
 
   const lnode = getResolvedNullableValueOrNull(node.lexpr);
   const rnode = getResolvedNullableValueOrNull(node.rexpr);
+  const operator = concatStringNodes(node.name);
 
   if (lnode === null || rnode === null) {
     return [];
   }
 
-  const operator = concatStringNodes(node.name);
-  const resolved: string | undefined =
-    defaultTypeExprMapping[`${lnode.value} ${operator} ${rnode.value}`];
+  const downcast = () => {
+    const left = lnode.array ? `_${lnode.value}` : lnode.value;
+    const right = rnode.array ? `_${rnode.value}` : rnode.value;
+
+    const overrides: Record<string, [string, string, string]> = {
+      "int4 ^ int4": ["float8", "^", "float8"],
+    };
+
+    if (overrides[`${left} ${operator} ${right}`]) {
+      return overrides[`${left} ${operator} ${right}`];
+    }
+
+    const adjust = (value: string) => (value === "varchar" ? "text" : value);
+
+    return [adjust(left), operator, adjust(right)];
+  };
 
-  if (resolved === undefined) {
+  const getType = (): ASTDescribedColumnType | undefined => {
+    const nullable = !context.nonNullableColumns.has(name) && (lnode.nullable || rnode.nullable);
+    const [dleft, doperator, dright] = downcast();
+
+    const type =
+      context.typeExprMap.get(dleft)?.get(doperator)?.get(dright) ??
+      context.typeExprMap.get("anycompatiblearray")?.get(operator)?.get("anycompatiblearray") ??
+      context.typeExprMap.get("anyarray")?.get(operator)?.get("anyarray") ??
+      context.typeExprMap.get(lnode.value)?.get(operator)?.values().next().value;
+
+    if (type === undefined) {
+      return;
+    }
+
+    if (type === "anycompatiblearray") {
+      return {
+        kind: "array",
+        value: resolveType({
+          context,
+          nullable,
+          type: context.toTypeScriptType({ name: lnode.value }),
+        }),
+      };
+    }
+
+    return resolveType({
+      context,
+      nullable,
+      type: context.toTypeScriptType({ name: type }),
+    });
+  };
+
+  const type = getType();
+
+  if (type === undefined) {
     return [];
   }
 
-  return [
-    {
-      name: name,
-      type: resolveType({
-        context: context,
-        nullable: !context.nonNullableColumns.has(name) && (lnode.nullable || rnode.nullable),
-        type: context.toTypeScriptType({ name: resolved }),
-      }),
-    },
-  ];
+  return [{ name, type }];
 }
 
 function getDescribedNullTest({
@@ -815,14 +874,14 @@ function getColumnRefOrigins({
       // lookup in cte
       context.select.withClause?.ctes
         .find((cte) => cte.CommonTableExpr?.ctename === source)
-        ?.CommonTableExpr?.ctequery?.SelectStmt?.targetList?.map((x) => x.ResTarget)
-        .find((x) => x?.name === column)?.val ??
+        ?.CommonTableExpr?.ctequery?.SelectStmt?.targetList?.find(
+          (x) => x.ResTarget?.name === column,
+        ) ??
       // lookup in subselect
       context.select.fromClause
         ?.map((from) => from.RangeSubselect)
         .find((subselect) => subselect?.alias?.aliasname === source)
-        ?.subquery?.SelectStmt?.targetList?.map((x) => x.ResTarget)
-        .find((x) => x?.name === column)?.val;
+        ?.subquery?.SelectStmt?.targetList?.find((x) => x.ResTarget?.name === column);
 
     if (!origin) return undefined;
 

packages/generate/src/generate.test.ts

@@ -220,7 +220,7 @@ const testQuery = async (params: {
         ({ output, unknownColumns }) => {
           assert.deepEqual(output?.value ?? null, params.expected);
 
-          if (unknownColumns.length > 0) {
+          if (unknownColumns.length > 0 || params.unknownColumns) {
             assert.deepEqual(unknownColumns, params.unknownColumns);
           }
         },
@@ -737,7 +737,6 @@ test("select from subselect with an alias", async () => {
   await testQuery({
     query: `SELECT subselect.id FROM (SELECT * FROM caregiver) AS subselect`,
     expected: [["id", { kind: "type", value: "number", type: "int4" }]],
-    unknownColumns: ["id"],
   });
 });
 
@@ -1971,46 +1970,48 @@ test("ARRAY[2, 3] <@ ARRAY[1, 2, 3] => boolean", async () => {
 test("ARRAY[1, 2] || ARRAY[3, 4] => array", async () => {
   await testQuery({
     query: `SELECT ARRAY[1, 2] || ARRAY[3, 4]`,
-    expected: [["?column?", { kind: "type", value: "array", type: "array" }]],
+    expected: [
+      ["?column?", { kind: "array", value: { kind: "type", type: "int4", value: "number" } }],
+    ],
   });
 });
 
-test("'{\"key\": \"value\"}'::jsonb ? 'key' => boolean", async () => {
+test(`'{"key": "value"}'::jsonb ? 'key' => boolean`, async () => {
   await testQuery({
     query: `SELECT '{"key": "value"}'::jsonb ? 'key'`,
     expected: [["?column?", { kind: "type", value: "boolean", type: "bool" }]],
   });
 });
 
-test("'{\"a\": 1, \"b\": 2}'::jsonb ?| array['a', 'c'] => boolean", async () => {
+test(`'{"a": 1, "b": 2}'::jsonb ?| array['a', 'c'] => boolean`, async () => {
   await testQuery({
     query: `SELECT '{"a": 1, "b": 2}'::jsonb ?| array['a', 'c']`,
     expected: [["?column?", { kind: "type", value: "boolean", type: "bool" }]],
   });
 });
 
-test("'{\"a\": 1, \"b\": 2}'::jsonb ?& array['a', 'b'] => boolean", async () => {
+test(`'{"a": 1, "b": 2}'::jsonb ?& array['a', 'b'] => boolean`, async () => {
   await testQuery({
     query: `SELECT '{"a": 1, "b": 2}'::jsonb ?& array['a', 'b']`,
     expected: [["?column?", { kind: "type", value: "boolean", type: "bool" }]],
   });
 });
 
-test("'{\"a\": {\"b\": 1}}'::jsonb -> 'a' => jsonb", async () => {
+test(`'{"a": {"b": 1}}'::jsonb -> 'a' => jsonb`, async () => {
   await testQuery({
     query: `SELECT '{"a": {"b": 1}}'::jsonb -> 'a'`,
     expected: [["?column?", { kind: "type", type: "jsonb", value: "any" }]],
   });
 });
 
-test("'{\"a\": {\"b\": 1}}'::jsonb ->> 'a' => string", async () => {
+test(`'{"a": {"b": 1}}'::jsonb ->> 'a' => string`, async () => {
   await testQuery({
     query: `SELECT '{"a": {"b": 1}}'::jsonb ->> 'a'`,
     expected: [["?column?", { kind: "type", value: "string", type: "text" }]],
   });
 });
 
-test("'{\"a\": 1, \"b\": 2}'::jsonb #- '{a}' => jsonb", async () => {
+test(`'{"a": 1, "b": 2}'::jsonb #- '{a}' => jsonb`, async () => {
   await testQuery({
     query: `SELECT '{"a": 1, "b": 2}'::jsonb #- '{a}'`,
     expected: [["?column?", { kind: "type", value: "any", type: "jsonb" }]],
@@ -2126,7 +2127,6 @@ test("with select from inner join and left join", async () => {
         INNER JOIN caregiver_agency ON x.id = caregiver_agency.caregiver_id
         LEFT JOIN agency ON caregiver_agency.agency_id = agency.id
     `,
-    unknownColumns: ["id"],
     expected: [
       ["id", { kind: "type", type: "int4", value: "number" }],
       [
@@ -2170,7 +2170,6 @@ test("select colref and const from left joined using col", async () => {
         },
       ],
     ],
-    unknownColumns: ["value"], // TODO: `ast-get-source` needs to be refactored to handle this case
   });
 });
 
@@ -2276,3 +2275,19 @@ test("select col.tbl from cte with array agg and col filter", async () => {
     ],
   });
 });
+
+test("varchar not like expr", async () => {
+  await testQuery({
+    schema: `CREATE TABLE tbl (email varchar(80) NOT NULL)`,
+    query: `SELECT jsonb_build_object('key', tbl.email NOT LIKE '%@example.com') AS col FROM tbl`,
+    expected: [
+      [
+        "col",
+        {
+          kind: "object",
+          value: [["key", { kind: "type", type: "bool", value: "boolean" }]],
+        },
+      ],
+    ],
+  });
+});

packages/generate/src/generate.ts

@@ -72,11 +72,12 @@ type Cache = {
     CacheKey,
     {
       pgTypes: PgTypesMap;
-      pgCols: PgColRow[];
+      pgCols: postgres.RowList<PgColRow[]>;
       pgEnums: PgEnumsMaps;
       pgColsByTableOidCache: Map<number, PgColRow[]>;
       pgColsBySchemaAndTableName: Map<string, Map<string, PgColRow[]>>;
       pgFnsByName: Map<string, PgFnRow[]>;
+      pgTypeExprMap: Map<string, Map<string, Map<string, string>>>;
     }
   >;
   overrides: {
@@ -128,13 +129,19 @@ async function generate(
 ): Promise<either.Either<GenerateError, GenerateResult>> {
   const { sql, query, cacheKey, cacheMetadata = true } = params;
 
-  const { pgColsByTableOidCache, pgColsBySchemaAndTableName, pgTypes, pgEnums, pgFnsByName } =
-    await getOrSetFromMapWithEnabled({
-      shouldCache: cacheMetadata,
-      map: cache.base,
-      key: cacheKey,
-      value: () => getDatabaseMetadata(sql),
-    });
+  const {
+    pgColsByTableOidCache,
+    pgColsBySchemaAndTableName,
+    pgTypes,
+    pgEnums,
+    pgFnsByName,
+    pgTypeExprMap,
+  } = await getOrSetFromMapWithEnabled({
+    shouldCache: cacheMetadata,
+    map: cache.base,
+    key: cacheKey,
+    value: () => getDatabaseMetadata(sql),
+  });
 
   const typesMap = await getOrSetFromMapWithEnabled({
     shouldCache: cacheMetadata,
@@ -266,6 +273,7 @@ async function generate(
       pgTypes: pgTypes,
       pgEnums: pgEnums,
       pgFns: functionsMap,
+      typeExprMap: pgTypeExprMap,
     });
 
     const columns = result.columns.map((col, position): ColumnAnalysisResult => {
@@ -329,6 +337,7 @@ async function getDatabaseMetadata(sql: Sql) {
   const pgColsByTableOidCache = groupBy(pgCols, "tableOid");
   const pgColsBySchemaAndTableName = groupBy(pgCols, "schemaName", "tableName");
   const pgFnsByName = groupBy(pgFns, "name");
+  const pgTypeExprMap = await getPgTypeExprMap(sql);
 
   return {
     pgTypes,
@@ -337,6 +346,7 @@ async function getDatabaseMetadata(sql: Sql) {
     pgColsByTableOidCache,
     pgColsBySchemaAndTableName,
     pgFnsByName,
+    pgTypeExprMap,
   };
 }
 
@@ -696,3 +706,30 @@ async function getPgFunctions(sql: Sql) {
     returnType: row.returnType,
   }));
 }
+
+async function getPgTypeExprMap(sql: Sql): Promise<Map<string, Map<string, Map<string, string>>>> {
+  const rows = await sql<{ left: string; operator: string; right: string; result: string }[]>`
+      SELECT
+        l.typname as left,
+        o.oprname as operator,
+        r.typname as right,
+        ret.typname AS result
+      FROM pg_operator o
+        JOIN pg_type l ON l.oid = o.oprleft
+        JOIN pg_type r ON r.oid = o.oprright
+        JOIN pg_type ret ON ret.oid = o.oprresult
+      WHERE o.oprleft <> 0 AND o.oprright <> 0
+  `;
+
+  const map = new Map<string, Map<string, Map<string, string>>>();
+
+  for (const row of rows) {
+    const leftMap = map.get(row.left) ?? new Map<string, Map<string, string>>();
+    const rightMap = leftMap.get(row.operator) ?? new Map<string, string>();
+    rightMap.set(row.right, row.result);
+    leftMap.set(row.operator, rightMap);
+    map.set(row.left, leftMap);
+  }
+
+  return map;
+}