Better enum handling

tylergannon · tylergannon · commit 85cd6a877d58 · 2026-01-20T15:59:47.000-06:00
diff --git a/src/app.ts b/src/app.ts
@@ -11,7 +11,7 @@ import {
   factory,
 } from "typescript";
 
-import { GenerateRequest, GenerateResponse, File } from "./gen/plugin/codegen_pb";
+import { GenerateRequest, GenerateResponse, File, Enum } from "./gen/plugin/codegen_pb";
 
 import { argName, colName } from "./drivers/utils";
 import { assertUniqueNames } from "./validate";
@@ -24,8 +24,64 @@ const result = codegen(input);
 // Write the result to stdout
 writeOutput(result);
 
+/**
+ * Build a map of enum names to their values from the catalog.
+ * This allows us to recognize enum types and generate appropriate TypeScript types.
+ */
+function buildEnumMap(input: GenerateRequest): Map<string, Enum> {
+  const enumMap = new Map<string, Enum>();
+  const defaultSchema = input.catalog?.defaultSchema ?? "public";
+
+  for (const schema of input.catalog?.schemas ?? []) {
+    if (schema.name === "pg_catalog" || schema.name === "information_schema") {
+      continue;
+    }
+
+    for (const enumDef of schema.enums) {
+      // Store with both qualified and unqualified names
+      enumMap.set(enumDef.name, enumDef);
+      if (schema.name !== defaultSchema) {
+        enumMap.set(`${schema.name}.${enumDef.name}`, enumDef);
+      }
+    }
+  }
+
+  return enumMap;
+}
+
+/**
+ * Generate TypeScript union type for an enum.
+ * e.g., type EventSource = 'user' | 'runner' | 'system';
+ */
+function enumTypeDecl(name: string, enumDef: Enum): Node {
+  const unionType = factory.createUnionTypeNode(
+    enumDef.vals.map((val) => factory.createLiteralTypeNode(factory.createStringLiteral(val))),
+  );
+
+  return factory.createTypeAliasDeclaration(
+    [factory.createToken(SyntaxKind.ExportKeyword)],
+    factory.createIdentifier(pascalCase(name)),
+    undefined,
+    unionType,
+  );
+}
+
+/**
+ * Convert snake_case to PascalCase
+ */
+function pascalCase(str: string): string {
+  return str
+    .split("_")
+    .map((part) => part.charAt(0).toUpperCase() + part.slice(1).toLowerCase())
+    .join("");
+}
+
 function codegen(input: GenerateRequest): GenerateResponse {
-  let files = [];
+  const files = [];
+  const enumMap = buildEnumMap(input);
+
+  // Set the enum map in the postgres driver so columnType can use it
+  postgres.setEnumMap(enumMap);
 
   const querymap = new Map<string, typeof input.queries>();
 
@@ -37,9 +93,15 @@ function codegen(input: GenerateRequest): GenerateResponse {
     qs?.push(query);
   }
 
+  // Track which enums are used across all files
+  const usedEnums = new Set<string>();
+
   for (const [filename, queries] of querymap.entries()) {
     const nodes: Node[] = [...postgres.preamble()];
 
+    // Track enums used in this file
+    const fileEnums = new Set<string>();
+
     for (const query of queries) {
       const lowerName = query.name[0].toLowerCase() + query.name.slice(1);
 
@@ -56,22 +118,37 @@ function codegen(input: GenerateRequest): GenerateResponse {
           names,
         });
 
-        nodes.push(
-          factory.createInterfaceDeclaration(
-            [factory.createToken(SyntaxKind.ExportKeyword)],
-            factory.createIdentifier(argIface),
-            undefined,
-            undefined,
-            query.params.map((param, i) =>
-              factory.createPropertySignature(
-                undefined,
-                factory.createIdentifier(argName(i, param.column)),
-                undefined,
-                postgres.columnType(param.column),
+        // Check for enum usage in params
+        for (const param of query.params) {
+          const enumName = postgres.getEnumName(param.column);
+          if (enumName) {
+            fileEnums.add(enumName);
+            usedEnums.add(enumName);
+          }
+        }
+
+        try {
+          nodes.push(
+            factory.createInterfaceDeclaration(
+              [factory.createToken(SyntaxKind.ExportKeyword)],
+              factory.createIdentifier(argIface),
+              undefined,
+              undefined,
+              query.params.map((param, i) =>
+                factory.createPropertySignature(
+                  undefined,
+                  factory.createIdentifier(argName(i, param.column)),
+                  undefined,
+                  postgres.columnType(param.column),
+                ),
               ),
             ),
-          ),
-        );
+          );
+        } catch (err) {
+          throw new Error(
+            `Error in query "${query.name}" (${filename}): ${err instanceof Error ? err.message : String(err)}`,
+          );
+        }
       }
 
       if (query.columns.length > 0) {
@@ -84,22 +161,37 @@ function codegen(input: GenerateRequest): GenerateResponse {
           names,
         });
 
-        nodes.push(
-          factory.createInterfaceDeclaration(
-            [factory.createToken(SyntaxKind.ExportKeyword)],
-            factory.createIdentifier(returnIface),
-            undefined,
-            undefined,
-            query.columns.map((column, i) =>
-              factory.createPropertySignature(
-                undefined,
-                factory.createIdentifier(colName(i, column)),
-                undefined,
-                postgres.columnType(column),
+        // Check for enum usage in columns
+        for (const col of query.columns) {
+          const enumName = postgres.getEnumName(col);
+          if (enumName) {
+            fileEnums.add(enumName);
+            usedEnums.add(enumName);
+          }
+        }
+
+        try {
+          nodes.push(
+            factory.createInterfaceDeclaration(
+              [factory.createToken(SyntaxKind.ExportKeyword)],
+              factory.createIdentifier(returnIface),
+              undefined,
+              undefined,
+              query.columns.map((column, i) =>
+                factory.createPropertySignature(
+                  undefined,
+                  factory.createIdentifier(colName(i, column)),
+                  undefined,
+                  postgres.columnType(column),
+                ),
               ),
             ),
-          ),
-        );
+          );
+        } catch (err) {
+          throw new Error(
+            `Error in query "${query.name}" (${filename}): ${err instanceof Error ? err.message : String(err)}`,
+          );
+        }
       }
 
       switch (query.cmd) {
@@ -140,6 +232,19 @@ function codegen(input: GenerateRequest): GenerateResponse {
       }
     }
 
+    // Add enum type declarations at the beginning of the file (after imports)
+    const enumNodes: Node[] = [];
+    for (const enumName of fileEnums) {
+      const enumDef = enumMap.get(enumName);
+      if (enumDef) {
+        enumNodes.push(enumTypeDecl(enumName, enumDef));
+      }
+    }
+
+    // Insert enum declarations after the preamble (imports)
+    const preambleLength = postgres.preamble().length;
+    nodes.splice(preambleLength, 0, ...enumNodes);
+
     files.push(
       new File({
         name: `${filename.replace(".", "_")}.ts`,
@@ -169,7 +274,7 @@ function printNode(nodes: Node[]): string {
   );
   const printer = createPrinter({ newLine: NewLineKind.LineFeed });
   let output = "// Code generated by sqlc. DO NOT EDIT.\n\n";
-  for (let node of nodes) {
+  for (const node of nodes) {
     output += printer.printNode(EmitHint.Unspecified, node, resultFile);
     output += "\n\n";
   }
diff --git a/src/drivers/postgres.ts b/src/drivers/postgres.ts
@@ -9,37 +9,102 @@
 
 import { SyntaxKind, NodeFlags, TypeNode, factory, FunctionDeclaration } from "typescript";
 
-import { Parameter, Column } from "../gen/plugin/codegen_pb";
+import { Parameter, Column, Enum } from "../gen/plugin/codegen_pb";
 import { argName } from "./utils";
 
+// Map of enum names to their definitions, set by app.ts
+let enumMap: Map<string, Enum> = new Map();
+
+/**
+ * Set the enum map from the catalog. Called by app.ts before generating code.
+ */
+export function setEnumMap(map: Map<string, Enum>): void {
+  enumMap = map;
+}
+
+/**
+ * Check if a column type is an enum and return the enum name if so.
+ */
+export function getEnumName(column?: Column): string | null {
+  if (column === undefined || column.type === undefined) {
+    return null;
+  }
+  const typeName = column.type.name.toLowerCase();
+  if (enumMap.has(typeName)) {
+    return typeName;
+  }
+  return null;
+}
+
+/**
+ * Convert snake_case to PascalCase for enum type names
+ */
+function pascalCase(str: string): string {
+  return str
+    .split("_")
+    .map((part) => part.charAt(0).toUpperCase() + part.slice(1).toLowerCase())
+    .join("");
+}
+
 export function columnType(column?: Column): TypeNode {
   if (column === undefined || column.type === undefined) {
     return factory.createKeywordTypeNode(SyntaxKind.AnyKeyword);
   }
-  let typeName = column.type.name;
+  const originalTypeName = column.type.name;
+  let typeName = originalTypeName;
   const pgCatalog = "pg_catalog.";
   if (typeName.startsWith(pgCatalog)) {
     typeName = typeName.slice(pgCatalog.length);
   }
 
-  typeName = typeName.toLowerCase();
+  const lowerTypeName = typeName.toLowerCase();
 
-  let typ: TypeNode = factory.createKeywordTypeNode(SyntaxKind.StringKeyword);
-  switch (typeName) {
+  // Check if it's an enum type
+  if (enumMap.has(lowerTypeName)) {
+    const typ = factory.createTypeReferenceNode(
+      factory.createIdentifier(pascalCase(lowerTypeName)),
+      undefined,
+    );
+    if (column.isArray || column.arrayDims > 0) {
+      let arrayType: TypeNode = typ;
+      const dims = Math.max(column.arrayDims || 1);
+      for (let i = 0; i < dims; i++) {
+        arrayType = factory.createArrayTypeNode(arrayType);
+      }
+      if (column.notNull) {
+        return arrayType;
+      }
+      return factory.createUnionTypeNode([
+        arrayType,
+        factory.createLiteralTypeNode(factory.createNull()),
+      ]);
+    }
+    if (column.notNull) {
+      return typ;
+    }
+    return factory.createUnionTypeNode([typ, factory.createLiteralTypeNode(factory.createNull())]);
+  }
+
+  let typ: TypeNode;
+  switch (lowerTypeName) {
+    // Boolean types
     case "bool":
     case "boolean":
       typ = factory.createKeywordTypeNode(SyntaxKind.BooleanKeyword);
       break;
+    // Binary types
     case "bytea":
       typ = factory.createTypeReferenceNode(factory.createIdentifier("Buffer"), undefined);
       break;
+    // Date/time types
     case "date":
     case "timestamp":
     case "timestamp without time zone":
     case "timestamptz":
     case "timestamp with time zone":
       typ = factory.createTypeReferenceNode(factory.createIdentifier("Date"), undefined);
       break;
+    // Numeric types
     case "float4":
     case "real":
     case "float8":
@@ -61,11 +126,64 @@ export function columnType(column?: Column): TypeNode {
     case "oid":
       typ = factory.createKeywordTypeNode(SyntaxKind.NumberKeyword);
       break;
+    // JSON types - any to allow flexible object access
     case "json":
     case "jsonb":
       typ = factory.createKeywordTypeNode(SyntaxKind.AnyKeyword);
       break;
-    // All other types default to string (uuid, text, varchar, etc.)
+    // Void type (from functions like pg_advisory_xact_lock)
+    case "void":
+      typ = factory.createKeywordTypeNode(SyntaxKind.VoidKeyword);
+      break;
+    // String types - explicitly listed (unambiguously representable as string)
+    case "text":
+    case "varchar":
+    case "character varying":
+    case "char":
+    case "character":
+    case "bpchar":
+    case "name":
+    case "uuid":
+    case "citext":
+    case "inet":
+    case "cidr":
+    case "macaddr":
+    case "macaddr8":
+    case "money":
+    case "numeric":
+    case "decimal":
+    case "xml":
+    case "bit":
+    case "varbit":
+    case "bit varying":
+    case "interval":
+    case "time":
+    case "time without time zone":
+    case "timetz":
+    case "time with time zone":
+    case "tsvector":
+    case "tsquery":
+      typ = factory.createKeywordTypeNode(SyntaxKind.StringKeyword);
+      break;
+    // Geometric types - postgres.js returns these as objects, not strings
+    case "point":
+    case "line":
+    case "lseg":
+    case "box":
+    case "path":
+    case "polygon":
+    case "circle":
+      throw new Error(
+        `Unrecognized PostgreSQL type: "${originalTypeName}". ` +
+          `Please add support for this type in sqlc-gen-typescript/src/drivers/postgres.ts`,
+      );
+    default:
+      throw new Error(
+        `Unrecognized PostgreSQL type: "${originalTypeName}" for column "${column.name || "unknown"}". ` +
+          `This usually means sqlc couldn't infer the type. ` +
+          `Try adding an explicit cast like "sqlc.arg(${column.name})::text" or "sqlc.narg('${column.name}')" in your query. ` +
+          `If this is a valid PostgreSQL type that needs support, please add it to sqlc-gen-typescript/src/drivers/postgres.ts`,
+      );
   }
 
   if (column.isArray || column.arrayDims > 0) {
