Fixed pgSchema enum types (fixes #4421)

Author: Sukairo-02

changelogs/drizzle-orm/0.43.1.md

@@ -0,0 +1,3 @@
+## Fixes
+
+- [Fixed incorrect types of schema enums in PostgreSQL](https://github.com/drizzle-team/drizzle-orm/issues/4421)

drizzle-orm/src/pg-core/schema.ts

@@ -1,6 +1,7 @@
 import { entityKind, is } from '~/entity.ts';
 import { SQL, sql, type SQLWrapper } from '~/sql/sql.ts';
-import { pgEnumObjectWithSchema, pgEnumWithSchema } from './columns/enum.ts';
+import type { NonArray, Writable } from '~/utils.ts';
+import { type PgEnum, type PgEnumObject, pgEnumObjectWithSchema, pgEnumWithSchema } from './columns/enum.ts';
 import { type pgSequence, pgSequenceWithSchema } from './sequence.ts';
 import { type PgTableFn, pgTableWithSchema } from './table.ts';
 import { type pgMaterializedView, pgMaterializedViewWithSchema, type pgView, pgViewWithSchema } from './view.ts';
@@ -23,15 +24,25 @@ export class PgSchema<TName extends string = string> implements SQLWrapper {
 		return pgMaterializedViewWithSchema(name, columns, this.schemaName);
 	}) as typeof pgMaterializedView;
 
-	public enum = ((enumName: any, input: any) => {
+	public enum<U extends string, T extends Readonly<[U, ...U[]]>>(
+		enumName: string,
+		values: T | Writable<T>,
+	): PgEnum<Writable<T>>;
+
+	public enum<E extends Record<string, string>>(
+		enumName: string,
+		enumObj: NonArray<E>,
+	): PgEnumObject<E>;
+
+	public enum(enumName: any, input: any): any {
 		return Array.isArray(input)
 			? pgEnumWithSchema(
 				enumName,
 				[...input] as [string, ...string[]],
 				this.schemaName,
 			)
 			: pgEnumObjectWithSchema(enumName, input, this.schemaName);
-	}) as any;
+	}
 
 	sequence: typeof pgSequence = ((name, options) => {
 		return pgSequenceWithSchema(name, options, this.schemaName);

drizzle-orm/type-tests/pg/tables.ts

@@ -1160,8 +1160,12 @@ await db.refreshMaterializedView(newYorkers2).withNoData().concurrently();
 }
 
 {
+	const testSchema = pgSchema('test');
+
 	const e1 = pgEnum('test', ['a', 'b', 'c']);
 	const e2 = pgEnum('test', ['a', 'b', 'c'] as const);
+	const e3 = testSchema.enum('test', ['a', 'b', 'c']);
+	const e4 = testSchema.enum('test', ['a', 'b', 'c'] as const);
 
 	const test = pgTable('test', {
 		col1: char('col1', { enum: ['a', 'b', 'c'] as const }),
@@ -1175,6 +1179,8 @@ await db.refreshMaterializedView(newYorkers2).withNoData().concurrently();
 		col9: varchar('col9', { enum: ['a', 'b', 'c'] as const }),
 		col10: varchar('col10', { enum: ['a', 'b', 'c'] }),
 		col11: varchar('col11'),
+		col12: e3('col4'),
+		col13: e4('col5'),
 	});
 
 	Expect<Equal<['a', 'b', 'c'], typeof test.col1.enumValues>>;
@@ -1188,11 +1194,17 @@ await db.refreshMaterializedView(newYorkers2).withNoData().concurrently();
 	Expect<Equal<['a', 'b', 'c'], typeof test.col9.enumValues>>;
 	Expect<Equal<['a', 'b', 'c'], typeof test.col10.enumValues>>;
 	Expect<Equal<[string, ...string[]], typeof test.col11.enumValues>>;
+	Expect<Equal<['a', 'b', 'c'], typeof test.col12.enumValues>>;
+	Expect<Equal<['a', 'b', 'c'], typeof test.col13.enumValues>>;
 }
 
 {
+	const testSchema = pgSchema('test');
+
 	const e1 = pgEnum('test', ['a', 'b', 'c']);
 	const e2 = pgEnum('test', ['a', 'b', 'c'] as const);
+	const e3 = testSchema.enum('test', ['a', 'b', 'c']);
+	const e4 = testSchema.enum('test', ['a', 'b', 'c'] as const);
 
 	const test = pgTable('test', {
 		col1: char('col1', { enum: ['a', 'b', 'c'] as const }).generatedAlwaysAs(sql``),
@@ -1206,6 +1218,8 @@ await db.refreshMaterializedView(newYorkers2).withNoData().concurrently();
 		col9: varchar('col9', { enum: ['a', 'b', 'c'] as const }).generatedAlwaysAs(sql``),
 		col10: varchar('col10', { enum: ['a', 'b', 'c'] }).generatedAlwaysAs(sql``),
 		col11: varchar('col11').generatedAlwaysAs(sql``),
+		col12: e3('col4').generatedAlwaysAs(sql``),
+		col13: e4('col5').generatedAlwaysAs(sql``),
 	});
 
 	Expect<Equal<['a', 'b', 'c'], typeof test.col1.enumValues>>;
@@ -1219,6 +1233,8 @@ await db.refreshMaterializedView(newYorkers2).withNoData().concurrently();
 	Expect<Equal<['a', 'b', 'c'], typeof test.col9.enumValues>>;
 	Expect<Equal<['a', 'b', 'c'], typeof test.col10.enumValues>>;
 	Expect<Equal<[string, ...string[]], typeof test.col11.enumValues>>;
+	Expect<Equal<['a', 'b', 'c'], typeof test.col12.enumValues>>;
+	Expect<Equal<['a', 'b', 'c'], typeof test.col13.enumValues>>;
 }
 
 {
@@ -1439,4 +1455,22 @@ await db.refreshMaterializedView(newYorkers2).withNoData().concurrently();
 	const res = await db.select().from(table);
 
 	Expect<Equal<{ enum: Role | null }[], typeof res>>;
+
+	const mySchema = pgSchema('my_schema');
+
+	const schemaRole = mySchema.enum('role', Role);
+
+	// @ts-expect-error
+	mySchema.enum('role', RoleNonString);
+
+	// @ts-expect-error
+	mySchema.enum('role', RolePartiallyString);
+
+	const schemaTable = mySchema.table('table', {
+		enum: schemaRole('enum'),
+	});
+
+	const schemaRes = await db.select().from(schemaTable);
+
+	Expect<Equal<{ enum: Role | null }[], typeof schemaRes>>;
 }