Merge pull request #68 from fdarian/drizzle-custom-mode

feat: `drizzle.custom.field.mode` for customizing `bigint` mode
fdarian · Jul 5, 2024 · fa88c56 · fa88c56
2 parents 1458235 + 36c20d1
commit fa88c56
Show file tree

Hide file tree

Showing 6 changed files with 22 additions and 7 deletions.
diff --git a/packages/generator/src/lib/adapter/fields/createField.ts b/packages/generator/src/lib/adapter/fields/createField.ts
@@ -16,7 +16,9 @@ export type DefineImport = {
 export interface CreateFieldInput {
 	field: DMMF.Field
 	imports?: ImportValue[]
-	func: string
+	func:
+		| string
+		| (<T extends Partial<Record<string, unknown>>>(opts: T) => string)
 	onDefault?: (
 		field: FieldWithDefault
 	) => { code: string; imports?: ImportValue[] } | undefined
@@ -30,8 +32,6 @@ export function createField(input: CreateFieldInput) {
 
 	let imports = input.imports ?? []
 
-	let func = `${input.func}`
-
 	const custom = getCustomDirective(field)
 	if (custom?.imports) {
 		imports = imports.concat(
@@ -43,6 +43,8 @@ export function createField(input: CreateFieldInput) {
 		)
 	}
 
+	let func = `${typeof input.func === 'string' ? input.func : input.func(custom?.field ?? {})}`
+
 	// .type<...>()
 	if (custom?.$type) {
 		func += `.$type<${custom.$type}>()`

diff --git a/packages/generator/src/lib/adapter/fields/directives/custom.ts b/packages/generator/src/lib/adapter/fields/directives/custom.ts
@@ -43,10 +43,18 @@ const ImportSchema = v.object({
 	type: v.optional(v.boolean()),
 })
 
+const BigIntMode = v.picklist(['number', 'bigint'])
+export type BigIntMode = v.InferOutput<typeof BigIntMode>
+
+const FieldSchema = v.object({
+	mode: v.optional(BigIntMode),
+})
+
 const DirectiveSchema = v.object({
 	imports: v.optional(v.array(ImportSchema)),
 	$type: v.optional(v.string()),
 	default: v.optional(v.string()),
+	field: v.optional(FieldSchema),
 })
 
 class InvalidDirectiveShapeError extends Error {

diff --git a/packages/generator/src/lib/adapter/providers/mysql.ts b/packages/generator/src/lib/adapter/providers/mysql.ts
@@ -5,6 +5,7 @@ import { createModule } from '~/lib/syntaxes/module'
 import { getDateMode } from '~/shared/date-mode'
 import { createAdapter } from '../adapter'
 import { createField, hasDefault, isDefaultFunc } from '../fields/createField'
+import type { BigIntMode } from '../fields/directives/custom'
 
 const coreModule = 'drizzle-orm/mysql-core'
 
@@ -63,7 +64,8 @@ export const mysqlAdapter = createAdapter({
 			return createField({
 				field,
 				imports: [namedImport(['bigint'], coreModule)],
-				func: `bigint('${getDbName(field)}', { mode: 'bigint' })`,
+				func: (opts: { mode?: BigIntMode }) =>
+					`bigint('${getDbName(field)}', { mode: '${opts?.mode ?? 'bigint'}' })`,
 				onDefault(field) {
 					if (
 						field.isId &&

diff --git a/packages/generator/src/lib/adapter/providers/postgres.ts b/packages/generator/src/lib/adapter/providers/postgres.ts
@@ -10,6 +10,7 @@ import {
 	hasDefault,
 	isDefaultFunc,
 } from '../fields/createField'
+import type { BigIntMode } from '../fields/directives/custom'
 
 const coreModule = 'drizzle-orm/pg-core'
 
@@ -81,7 +82,8 @@ export const postgresAdapter = createAdapter({
 			return createField({
 				field,
 				imports: [namedImport([func], coreModule)],
-				func: `${func}('${getDbName(field)}', { mode: 'bigint' })`,
+				func: (opts: { mode?: BigIntMode }) =>
+					`${func}('${getDbName(field)}', { mode: '${opts?.mode ?? 'bigint'}' })`,
 			})
 		},
 		// https://orm.drizzle.team/docs/column-types/pg/#boolean

diff --git a/packages/usage/prisma/schema.prisma b/packages/usage/prisma/schema.prisma
@@ -206,7 +206,8 @@ model FieldCustomization {
   /// drizzle.custom {
   ///   "imports": [{ "name": ["SomeBigInt"], "module": "~/tests/shared/testFieldCustomization", "type": true }],
   ///   "$type": "SomeBigInt",
-  ///   "default": "() => 1n"
+  ///   "default": "() => 1n",
+  ///   "field": { "mode": "number" }
   /// }
   allFields BigInt
 }
diff --git a/packages/usage/tests/shared/testFieldCustomization.ts b/packages/usage/tests/shared/testFieldCustomization.ts
@@ -22,7 +22,7 @@ export function testFieldCustomization({ db, schema, provider }: TestContext) {
 
 		test('should contain correct field definition', () => {
 			expect(content).include(
-				"allFields: bigint('allFields', { mode: 'bigint' }).$type<SomeBigInt>().$defaultFn(() => 1n)"
+				"allFields: bigint('allFields', { mode: 'number' }).$type<SomeBigInt>().$defaultFn(() => 1n)"
 			)
 		})
 	})