schema.prismaを拡張してフィールドに指定した型が付くようにする
Commune Advent Calendar 2024 シリーズB 1日目の記事です。
みなさん、zod-prisma-typesはご存知ですか?
schema.prisma
上にコメントでZodスキーマっぽいアノテーションをつけることで、そこから実際のZodスキーマ(TypeScriptファイル)を生成してくれるライブラリです。やりたいことはすごく分かるし、ハマればとても便利そうなのですが、私にはちょっとヘビーウェイトで、あらゆるユースケースに合わせるのは難しそうな印象を持ちました。もっと気軽かつカスタマイズ容易な方法がないか考えて、真似して自分で作ればいいじゃんって気がついたので、試しに作ってみました。
環境
実現したいこと
例えばこんな感じでschema.prisma
に@schema
アノテーションを付けて、
model User {
/// @schema(UserId)
id Int @id @default(autoincrement())
/// @schema(Email)
email String @unique
/// @schema(ShortText)
name String?
community Community @relation(fields: [communityId], references: [id])
/// @schema(CommunityId)
communityId Int
createdAt DateTime @default(now())
updatedAt DateTime @updatedAt
}
prisma generate
すると、自動で以下のような(prismaClient
を拡張する)関数を出力する。
export function extendPrismaClient(prismaClient: PrismaClient) {
return prismaClient.$extends({
result: {
user: {
id: {
needs: { id: true },
compute(user) {
return UserId.of(user.id)
},
},
email: {
needs: { email: true },
compute(user) {
return Email.of(user.email)
},
},
name: {
needs: { name: true },
compute(user) {
if (user.name == null) return null
return ShortText.of(user.name)
},
},
communityId: {
needs: { communityId: true },
compute(user) {
return CommunityId.of(user.communityId)
},
},
},
},
})
}
これで拡張したprismaClient
を使えば、データ取得時に各フィールドに型が付いている状態を実現できる。
やったこと
Prismaで任意のジェネレータを作るのはそれほど難しくありません。
schema.prisma
に以下のように書いて、
generator resultExtender {
provider = "bun ./prisma/generator/result-extender/index.ts"
}
生成コマンドでprisma generate --generator resultExtender
のように引数でgeneratorを指定すれば良いです。すると、生成時に./prisma/generator/result-extender/index.ts
を実行してくれます。
実行するファイルはこんな感じ👇️
import { generatorHandler } from "@prisma/generator-helper"
import { z } from "zod"
import { generate } from "./generate"
const outputSchema = z.object({
fromEnvVar: z.string().nullable(),
value: z.string({ required_error: "No output path specified" }),
})
generatorHandler({
onManifest: () => {
return {
defaultOutput: "./generated/result-extender",
prettyName: "Result Extender",
}
},
onGenerate: async (generatorOptions) => {
const output = outputSchema.parse(generatorOptions.generator.output)
await generate({ outDir: output.value, dmmf: generatorOptions.dmmf })
},
})
schema.prisma
に書いたコメントを読み取って、ts-morphでコードを生成し、Biomeで整形します。
import * as path from "node:path"
import { Biome, Distribution } from "@biomejs/js-api"
import type { DMMF } from "@prisma/generator-helper/dist/dmmf"
import { parse as parseJsonc } from "jsonc-parser"
import { Project, type SourceFile } from "ts-morph"
import type { Context } from "./types/context"
import { uncapitalize } from "./utils/utils"
const schemaDocumentRegex = /^@schema\((\w+)\)$/m
export async function generate({
outDir,
dmmf,
}: { outDir: string; dmmf: DMMF.Document }) {
const ctx: Context = {
project: new Project({
tsConfigFilePath: "./tsconfig.json",
}),
}
const outSourceFile = ctx.project.createSourceFile(
path.join(outDir, "extend-prisma-client.ts"),
"",
{ overwrite: true },
)
outSourceFile.addFunction({
name: "extendPrismaClient",
parameters: [{ name: "prismaClient", type: "PrismaClient" }],
isExported: true,
statements: (writer) => {
writer
.writeLine("return prismaClient.$extends(")
.block(() => {
writer.writeLine("result:").block(() => {
dmmf.datamodel.models.forEach((model) => {
const uncapitalizedModelName = uncapitalize(model.name)
writer
.writeLine(`${uncapitalizedModelName}:`)
.block(() => {
model.fields.forEach((field) => {
const doc = field.documentation
if (doc == null) return
const match = doc.match(schemaDocumentRegex)
if (match == null || match[1] == null) return
const schemaName = match[1]
writer
.writeLine(`${field.name}:`)
.block(() => {
writer.write(`needs: { ${field.name}: true },`)
writer
.write(`compute(${uncapitalizedModelName})`)
.block(() => {
if (!field.isRequired) {
writer.writeLine(
`if (${uncapitalizedModelName}.${field.name} == null) return null`,
)
}
writer.write(
`return ${schemaName}.of(${uncapitalizedModelName}.${field.name})`,
)
})
})
.write(",")
})
})
.write(",")
})
})
})
.write(")")
},
})
outSourceFile.fixMissingImports()
outSourceFile.organizeImports()
const formattedContent = await getFixedContentWithBiome(outSourceFile)
outSourceFile.replaceWithText(formattedContent)
await outSourceFile.save()
}
async function getFixedContentWithBiome(outSourceFile: SourceFile) {
const biome = await Biome.create({
distribution: Distribution.NODE,
})
const biomeConfigFile = Bun.file("./biome.jsonc")
biome.applyConfiguration(parseJsonc(await biomeConfigFile.text()))
const filePath = outSourceFile.getFilePath()
const formatResult = biome.formatContent(outSourceFile.getText(), {
filePath,
})
const lintResult = biome.lintContent(formatResult.content, {
filePath,
fixFileMode: "SafeAndUnsafeFixes",
})
return lintResult.content
}
これだけで、実現したいことに書いたextend-prisma-client.ts
を生成できるようになりました!意外と簡単!
もちろんミニマム実装なので、実際のプロダクトで使いたかったら改善の余地は色々ありそうですが、自分でカスタマイズできる実装が手に入ったのは嬉しい。
さいごに
拡張性の高いライブラリとts-morphを組み合わせたら、他にもできることすごく色々ありそうですね!こんな感じで、これからも手強いプロダクト開発と戦っていきたいです。
Discussion