如何优雅地提取和管理Prisma客户端扩展类型

// 假设 _prismaClient 是原始的 PrismaClient 实例
const prismaClient = _prismaClient.$extends({
  query: {
    company: {
      update: async ({ args, query }) => {
        // 业务逻辑：如果公司状态为DECLINED，则锁定相关用户账户
        if (args.data?.status === CompanyStatus.DECLINED) {
          args.data.user = {
            update: {
              accountLocked: AccountLockedReason.COMPANY_DECLINED,
            },
          };
        }
        return query(args);
      },
    },
  },
});

// prismaClient.ts
import { companyExtensions } from './companyExtensions';

const prismaClient = _prismaClient.$extends({
  query: {
    company: companyExtensions, // 需要 companyExtensions 具有正确的类型
  },
});

// companyExtensions.ts
// export const companyExtensions: NeedsType = { ... }; // 这里的 NeedsType 是挑战

type ExtensionArgs = Extract<
  Parameters<typeof _prismaClient.$extends>[0],
  { name?: string }
>;

// src/types/prisma.ts (或任何你喜欢的路径)
import { PrismaClient } from '@prisma/client';

// 假设 _prismaClient 是你原始的 PrismaClient 实例。
// 注意：为了在类型文件中引用它，你可能需要一个“虚拟”实例或一个类型别名。
// 最简单的方法是直接从 @prisma/client 导入 PrismaClient 类型。
// 如果你有一个自定义的基类，可能需要调整。
// 这里我们假设 _prismaClient 是 PrismaClient 的一个实例。
// 如果你想避免实际导入 _prismaClient，可以这样定义一个类型：
// type BasePrismaClient = InstanceType<typeof PrismaClient>;
// 然后在 ExtensionArgs 中使用 BasePrismaClient['$extends']
// 但为了简化，我们直接使用 _prismaClient 的类型。

// 假设你有一个 _prismaClient 实例，或者你可以直接使用 typeof PrismaClient
// 这里的 _prismaClient 应该指向你实际使用的 PrismaClient 实例或其类型
// 示例：如果你有一个 const _prismaClient = new PrismaClient();
// 那么 typeof _prismaClient 就是正确的。
// 如果你只是想定义类型，可以这样模拟：
declare const _prismaClient: PrismaClient; // 这是一个类型声明，不会生成运行时代码

export type ExtensionArgs = Extract<
  Parameters<typeof _prismaClient.$extends>[0],
  { name?: string } // 使用 name?: string 作为模式匹配，匹配Prisma扩展的顶层结构
>;

// src/extensions/companyExtensions.ts
import { CompanyStatus, AccountLockedReason } from '@prisma/client'; // 假设这些枚举已定义
import { ExtensionArgs } from '../types/prisma'; // 导入定义的类型

// 精确指定 company 模型的 query.update 扩展类型
// 我们可以通过 ExtensionArgs 进一步推断出具体模型的扩展类型
// 实际操作中，Prisma的类型推断通常足够智能，
// 但为了明确和分离，我们可以在这里直接应用更通用的 ExtensionArgs
// 然后让TypeScript在组合时进行验证。

// 这里的 companyExtensions 必须符合 ExtensionArgs 的部分结构
// 更精确的做法是只导出 query.company 的部分
export const companyExtensions: ExtensionArgs['query']['company'] = {
  update: async ({ args, query }) => {
    if (args.data?.status === CompanyStatus.DECLINED) {
      args.data.user = {
        update: {
          accountLocked: AccountLockedReason.COMPANY_DECLINED,
        },
      };
    }
    return query(args);
  },
};

// src/prismaClient.ts
import { PrismaClient } from '@prisma/client';
import { companyExtensions } from './extensions/companyExtensions';

const _prismaClient = new PrismaClient();

export const prismaClient = _prismaClient.$extends({
  query: {
    company: companyExtensions, // 类型现在可以正确推断和验证
  },
  // 可以继续添加其他模型的扩展
});

export type ExtendedPrismaClient = typeof prismaClient;