🐳

schema.prismaを拡張してフィールドに指定した型が付くようにする

2024/12/01に公開

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を拡張する)関数を出力する。

extend-prisma-client.ts
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を実行してくれます。

実行するファイルはこんな感じ👇️

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で整形します。

prisma/generator/result-extender/generate.ts
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