Fix effect-schema $type output inference

drizzle-orm/src/effect-schema/column.types.ts

@@ -2,7 +2,7 @@ import type { Schema as s } from 'effect';
 import type { Array$, Literal, NullOr, optional, Schema, Struct, Tuple, Tuple2, UndefinedOr } from 'effect/Schema';
 import type { ColumnTypeData, ExtractColumnTypeData } from '~/column-builder.ts';
 import type { Column } from '~/column.ts';
-import type { Assume } from '~/utils.ts';
+import type { Assume, Equal } from '~/utils.ts';
 import type { bigintStringModeSchema, bufferSchema, jsonSchema, unsignedBigintStringModeSchema } from './column.ts';
 
 type GetArrayDepth<T, Depth extends number = 0> = Depth extends 5 ? 5
@@ -17,18 +17,29 @@ type WrapInEffectSchemaArray<TSchema extends Schema.Any, TDepth extends number>
 	: TDepth extends 5 ? Array$<Array$<Array$<Array$<Array$<TSchema>>>>>
 	: Array$<typeof s.Any>;
 
-type IsPgArrayColumn<TColumn extends Column<any>, TType extends ColumnTypeData> = TType['type'] extends 'array' ? false // Already handled as explicit array type
-	: GetArrayDepth<TColumn['_']['data']> extends 0 ? false
+type IsPgArrayColumn<TData, TType extends ColumnTypeData> = TType['type'] extends 'array' ? false // Already handled as explicit array type
+	: GetArrayDepth<TData> extends 0 ? false
 	: true;
 
+type RebuildEffectSchema<TSchema extends Schema.Any, TData> =
+	& Omit<TSchema, keyof Schema.Any>
+	& Schema<TData, s.Schema.Encoded<TSchema>, s.Schema.Context<TSchema>>;
+
+type ApplyTypeOverride<TSchema extends Schema.Any, TData> = Equal<s.Schema.Type<TSchema>, TData> extends true ? TSchema
+	: RebuildEffectSchema<TSchema, TData>;
+
 export type GetEffectSchemaType<
 	TColumn extends Column<any>,
+	TData = TColumn['_']['data'],
 	TType extends ColumnTypeData = ExtractColumnTypeData<TColumn['_']['dataType']>,
-> = IsPgArrayColumn<TColumn, TType> extends true ? WrapInEffectSchemaArray<
-		GetBaseEffectSchemaType<TColumn, TType>,
-		GetArrayDepth<TColumn['_']['data']>
-	>
-	: GetBaseEffectSchemaType<TColumn, TType>;
+> = ApplyTypeOverride<
+	IsPgArrayColumn<TData, TType> extends true ? WrapInEffectSchemaArray<
+			GetBaseEffectSchemaType<TColumn, TType>,
+			GetArrayDepth<TData>
+		>
+		: GetBaseEffectSchemaType<TColumn, TType>,
+	TData
+>;
 
 type GetBaseEffectSchemaType<
 	TColumn extends Column<any>,

integration-tests/tests/validators/effect-schema/pg.test.ts

@@ -15,7 +15,7 @@ import {
 	text,
 } from 'drizzle-orm/pg-core';
 import { CONSTANTS } from 'drizzle-orm/utils';
-import { Schema as s } from 'effect';
+import { type Brand, Schema as s } from 'effect';
 import type { TopLevelCondition } from 'json-rules-engine';
 import { test } from 'vitest';
 import { Equal, Expect } from '~/utils';
@@ -73,6 +73,41 @@ test('table in schema - select', (tc) => {
 	Expect<Equal<typeof result, typeof expected>>();
 });
 
+test('$type override preserves effect schema output type', (t) => {
+	type MyBrandedId = string & Brand.Brand<'MyBrandedId'>;
+
+	const table = pgTable('test', ({ uuid }) => ({
+		id: uuid().$type<MyBrandedId>().notNull(),
+	}));
+
+	const result = createSelectSchema(table);
+	const runtimeExpected = s.Struct({ id: s.UUID });
+
+	expectSchemaShape(t, runtimeExpected).from(result);
+	Expect<Equal<s.Schema.Type<typeof result>['id'], MyBrandedId>>();
+	Expect<Equal<s.Schema.Type<typeof result.fields.id>, MyBrandedId>>();
+	Expect<Equal<s.Schema.Context<typeof result>, never>>();
+	Expect<Equal<s.Schema.Context<typeof result.fields.id>, never>>();
+});
+
+test('$type literal override preserves effect schema output type', (t) => {
+	const fooLiteral = s.Literal('Foo');
+	type Foo = s.Schema.Type<typeof fooLiteral>;
+
+	const table = pgTable('test', ({ varchar }) => ({
+		value: varchar().$type<Foo>().notNull(),
+	}));
+
+	const result = createSelectSchema(table);
+	const runtimeExpected = s.Struct({ value: s.String });
+
+	expectSchemaShape(t, runtimeExpected).from(result);
+	Expect<Equal<s.Schema.Type<typeof result>['value'], Foo>>();
+	Expect<Equal<s.Schema.Type<typeof result.fields.value>, Foo>>();
+	Expect<Equal<s.Schema.Context<typeof result>, never>>();
+	Expect<Equal<s.Schema.Context<typeof result.fields.value>, never>>();
+});
+
 test('table - insert', (t) => {
 	const table = pgTable('test', {
 		id: integer().generatedAlwaysAsIdentity().primaryKey(),