feat: override column types (#213)

.changeset/little-knives-flash.md

@@ -0,0 +1,20 @@
+---
+"@ts-safeql/eslint-plugin": minor
+"@ts-safeql/generate": minor
+---
+
+You can now modify the expected type for each table column. This could be useful when dealing with
+dynamic types such as JSONB or when you want to enforce a specific type for a column.
+
+```json
+{
+  // ...,
+  "overrides": {
+    "columns": {
+      "table_name.column_name": "CustomType"
+    }
+  }
+}
+```
+
+You can read more about it in the [documentation](https://safeql.dev/api/#connections-overrides-columns-optional)

docs/.vitepress/config.ts

@@ -73,6 +73,10 @@ export default defineConfig({
                 text: "overrides.types",
                 link: "/api/index.md#connections-overrides-types-optional",
               },
+              {
+                text: "overrides.columns",
+                link: "/api/index.md#connections-overrides-columns-optional",
+              },
             ],
           },
         ],

docs/api/index.md

@@ -393,3 +393,23 @@ In this case, you can use the following syntax:
   }
 }
 ```
+
+### `connections.overrides.columns` (Optional)
+
+::: info
+
+While SafeQL generally succeeds in inferring column types, it sometimes cannot, particularly with `JSONB` columns where the type is not explicitly defined. Additionally, there may be scenarios where you want to enforce a specific type for a column.
+
+:::
+
+```json
+{
+  "connections": {
+    "overrides": {
+      "columns": {
+        "table_name.column_name": "CustomType"
+      }
+    }
+  }
+}
+```

packages/eslint-plugin/src/rules/check-sql.rule.ts

@@ -123,6 +123,7 @@ const zBaseSchema = z.object({
         z.record(z.enum(objectKeysNonEmpty(defaultTypeMapping)), zOverrideTypeResolver),
         z.record(z.string(), zOverrideTypeResolver),
       ]),
+      columns: z.record(z.string(), z.string()),
     })
     .partial()
     .optional(),
@@ -359,14 +360,18 @@ function reportCheck(params: {
         }
 
         const reservedTypes = memoize({
-          key: `reserved-types:${JSON.stringify(connection.overrides?.types)}`,
+          key: `reserved-types:${JSON.stringify(connection.overrides)}`,
           value: () => {
             const types = new Set<string>();
 
             for (const value of Object.values(connection.overrides?.types ?? {})) {
               types.add(typeof value === "string" ? value : value.return);
             }
 
+            for (const columnType of Object.values(connection.overrides?.columns ?? {})) {
+              types.add(columnType);
+            }
+
             return types;
           },
         });

packages/eslint-plugin/src/rules/check-sql.test.ts

@@ -71,6 +71,12 @@ const runMigrations1 = <TTypes extends Record<string, unknown>>(sql: Sql<TTypes>
       jsonb_col JSONB NOT NULL
     );
 
+    CREATE TABLE test_override_column_type (
+      jsonb_col JSONB NOT NULL,
+      jsonb_col_nullable JSONB,
+      jsonb_col_not_overriden JSONB
+    );
+
     CREATE TABLE all_types (
       id SERIAL PRIMARY KEY NOT NULL,
       text_column TEXT NOT NULL,
@@ -1163,6 +1169,99 @@ RuleTester.describe("check-sql", () => {
     ],
   });
 
+  ruleTester.run("connection with overrides.columns", rules["check-sql"], {
+    valid: [
+      {
+        filename,
+        options: withConnection(connections.withTag, {
+          overrides: {
+            columns: {
+              "test_override_column_type.jsonb_col": "JsonbColType",
+              "test_override_column_type.jsonb_col_nullable": "JsonbColType",
+            },
+          },
+        }),
+        name: "overriden-column: valid type annotation",
+        code: `
+          type JsonbColType = { foo: string; };
+
+          sql<{
+            jsonb_col: JsonbColType;
+            jsonb_col_nullable: JsonbColType | null;
+            jsonb_col_not_overriden: any | null
+          }>\`
+            select
+              jsonb_col,
+              jsonb_col_nullable,
+              jsonb_col_not_overriden
+            from
+              test_override_column_type
+          \`
+        `,
+      },
+    ],
+    invalid: [
+      {
+        filename,
+        options: withConnection(connections.withTag, {
+          overrides: {
+            columns: {
+              "test_override_column_type.jsonb_col": "JsonbColType",
+              "test_override_column_type.jsonb_col_nullable": "JsonbColType",
+            },
+          },
+        }),
+        name: "overriden-column: invalid missing type annotation",
+        code: `
+          type JsonbColType = { foo: string; };
+
+          sql\`
+            select
+              jsonb_col,
+              jsonb_col_nullable,
+              jsonb_col_not_overriden
+            from
+              test_override_column_type
+          \`
+        `,
+        output: `
+          type JsonbColType = { foo: string; };
+
+          sql<{ jsonb_col: JsonbColType; jsonb_col_nullable: JsonbColType | null; jsonb_col_not_overriden: any | null }>\`
+            select
+              jsonb_col,
+              jsonb_col_nullable,
+              jsonb_col_not_overriden
+            from
+              test_override_column_type
+          \`
+        `,
+        errors: [{ messageId: "missingTypeAnnotations" }],
+      },
+      {
+        filename,
+        options: withConnection(connections.withTag, {
+          overrides: {
+            columns: {
+              "invalid-config": "JsonbColType",
+            },
+          },
+        }),
+        name: "overriden-column: invalid configuration",
+        code: "sql`select 1`",
+        errors: [
+          {
+            messageId: "error",
+            data: {
+              error:
+                "Internal error: Invalid override column key: invalid-config. Expected format: table.column",
+            },
+          },
+        ],
+      },
+    ],
+  });
+
   ruleTester.run("connection with fieldTransform", rules["check-sql"], {
     valid: [
       {

packages/generate/src/ast-describe.ts

@@ -14,6 +14,7 @@ type ASTDescriptionOptions = {
   parsed: LibPgQueryAST.ParseResult;
   relations: FlattenedRelationWithJoins[];
   typesMap: Map<string, { override: boolean; value: string }>;
+  overridenColumnTypesMap: Map<string, Map<string, string>>;
   nonNullableColumns: Set<string>;
   pgColsByTableName: Map<string, PgColRow[]>;
   pgTypes: PgTypesMap;
@@ -604,15 +605,27 @@ function getDescribedColumnByResolvedColumns(params: {
   context: ASTDescriptionContext;
 }) {
   return params.resolved.map(({ column, isNotNull }) => {
+    const getType = (): ASTDescribedColumnType => {
+      const overridenType = params.context.overridenColumnTypesMap
+        .get(column.tableName)
+        ?.get(column.colName);
+
+      if (overridenType !== undefined) {
+        return { kind: "type", value: overridenType };
+      }
+
+      return params.context.toTypeScriptType({
+        oid: column.colTypeOid,
+        baseOid: column.colBaseTypeOid,
+      });
+    };
+
     return {
       name: params.alias ?? column.colName,
       type: resolveType({
         context: params.context,
         nullable: !isNotNull,
-        type: params.context.toTypeScriptType({
-          oid: column.colTypeOid,
-          baseOid: column.colBaseTypeOid,
-        }),
+        type: getType(),
       }),
     };
   });

packages/generate/src/generate.ts

@@ -44,6 +44,7 @@ type CacheKey = string;
 type OverrideValue = string | { parameter: string | { regex: string }; return: string };
 type Overrides = {
   types: TypesMap;
+  columns: Map<string, Map<string, string>>;
 };
 
 export interface GenerateParams {
@@ -55,6 +56,7 @@ export interface GenerateParams {
   fieldTransform: IdentiferCase | undefined;
   overrides?: Partial<{
     types: Record<string, OverrideValue>;
+    columns: Record<string, string>;
   }>;
 }
 
@@ -74,14 +76,20 @@ type Cache = {
       pgFnsByName: Map<string, PgFnRow[]>;
     }
   >;
-  types: Map<string, TypesMap>;
+  overrides: {
+    types: Map<string, TypesMap>;
+    columns: Map<string, Map<string, Map<string, string>>>;
+  };
   functions: Map<string, FunctionsMap>;
 };
 
 function createEmptyCache(): Cache {
   return {
     base: new Map(),
-    types: new Map(),
+    overrides: {
+      types: new Map(),
+      columns: new Map(),
+    },
     functions: new Map(),
   };
 }
@@ -94,7 +102,8 @@ export function createGenerator() {
     dropCacheKey: (cacheKey: CacheKey) => cache.base.delete(cacheKey),
     clearCache: () => {
       cache.base.clear();
-      cache.types.clear();
+      cache.overrides.types.clear();
+      cache.overrides.columns.clear();
       cache.functions.clear();
     },
   };
@@ -126,7 +135,7 @@ async function generate(
 
   const typesMap = await getOrSetFromMapWithEnabled({
     shouldCache: cacheMetadata,
-    map: cache.types,
+    map: cache.overrides.types,
     key: JSON.stringify(params.overrides?.types),
     value: () => {
       const map: TypesMap = new Map();
@@ -143,6 +152,29 @@ async function generate(
     },
   });
 
+  const overridenColumnTypesMap = await getOrSetFromMapWithEnabled({
+    shouldCache: cacheMetadata,
+    map: cache.overrides.columns,
+    key: JSON.stringify(params.overrides?.columns),
+    value: () => {
+      const map: Map<string, Map<string, string>> = new Map();
+
+      for (const [colPath, type] of Object.entries(params.overrides?.columns ?? {})) {
+        const [table, column] = colPath.split(".");
+
+        if (table === undefined || column === undefined) {
+          throw new Error(`Invalid override column key: ${colPath}. Expected format: table.column`);
+        }
+
+        map.has(table)
+          ? map.get(table)?.set(column, type)
+          : map.set(table, new Map([[column, type]]));
+      }
+
+      return map;
+    },
+  });
+
   function byReturnType(a: PgFnRow, b: PgFnRow) {
     const priority = ["numeric", "int8"];
     return priority.indexOf(a.returnType) - priority.indexOf(b.returnType);
@@ -205,6 +237,7 @@ async function generate(
       parsed: params.pgParsed,
       relations: relationsWithJoins,
       typesMap: typesMap,
+      overridenColumnTypesMap: overridenColumnTypesMap,
       nonNullableColumns: nonNullableColumnsBasedOnAST,
       pgColsByTableName: pgColsByTableName,
       pgTypes: pgTypes,
@@ -232,7 +265,10 @@ async function generate(
       pgTypes,
       pgEnums,
       relationsWithJoins,
-      overrides: { types: typesMap },
+      overrides: {
+        types: typesMap,
+        columns: overridenColumnTypesMap,
+      },
       pgColsByTableName,
       fieldTransform: params.fieldTransform,
     };