Repository: prisma/prisma
File: packages/client/tests/functional/batching/tests.ts (TypeScript)

Code before:

    function resetQueryCounts() {
      queriesExecuted = 0
      beginCount = 0
      commitCount = 0
      rollbackCount = 0
      savepointCount = 0
      releaseSavepointCount = 0
      rollbackToSavepointCount = 0
    }

    function resetEngineCounts() {
      engineRequestCount = 0
      engineRequestBatchCount = 0
      engineRequestBatchSizes.length = 0
    }

    beforeAll(async () => {
      prisma = newPrismaClient({
        log: [{ emit: 'event', level: 'query' }],
      })

      // Count batching at the transport layer: this is what actually verifies that Prisma Client
      // grouped multiple operations into a single engine roundtrip.

      const engine = (prisma as any)._engine as any
      if (!engine || typeof engine.request !== 'function' || typeof engine.requestBatch !== 'function') {
        throw new Error('Expected PrismaClient to have an engine with request/requestBatch methods')
      }

      const originalRequest = engine.request.bind(engine) as (...args: any[]) => Promise<unknown>
      const originalRequestBatch = engine.requestBatch.bind(engine) as (...args: any[]) => Promise<unknown>

      engine.request = async (...args: any[]) => {
        engineRequestCount++
        return originalRequest(...args)
      }

      engine.requestBatch = async (...args: any[]) => {
        engineRequestBatchCount++
        const queries = args[0]
        engineRequestBatchSizes.push(Array.isArray(queries) ? queries.length : 0)
        return originalRequestBatch(...args)
      }

      await prisma.user.create({ data: { posts: { create: {} }, ...user1 } })
      await prisma.user.create({ data: user2 })

      // @ts-expect-error - client not typed for log opts for cross generator compatibility - can be improved once we drop the prisma-client-js generator
      prisma.$on('query', ({ query }: Prisma.QueryEvent) => {
        const normalized = query.trim().toUpperCase()

Review comment (suggestion):
To make it less brittle and more type safe, we could try to type it against `getPrismaClient`:

```suggestion
      const engine = (client as InstanceType<ReturnType<typeof getPrismaClient>>)._engine
```

Code after:
    const engineRequestBatchSizes: number[] = []

    function resetQueryCounts() {
      queriesExecuted = 0
      beginCount = 0
      commitCount = 0
      rollbackCount = 0
      savepointCount = 0
      releaseSavepointCount = 0
      rollbackToSavepointCount = 0
    }

    function resetEngineCounts() {
      engineRequestCount = 0
      engineRequestBatchCount = 0
      engineRequestBatchSizes.length = 0
    }

    beforeAll(async () => {
      prisma = newPrismaClient({
        log: [{ emit: 'event', level: 'query' }],
      })

      // Count batching at the transport layer: this is what actually verifies that Prisma Client
      // grouped multiple operations into a single engine roundtrip.
      const engine = (prisma as unknown as InstanceType<ReturnType<typeof getPrismaClient>>)._engine
      const originalRequest = engine.request.bind(engine)
      const originalRequestBatch = engine.requestBatch.bind(engine)

      engine.request = (query, options) => {
        engineRequestCount++
        return originalRequest(query, options)
      }

      engine.requestBatch = (queries, options) => {
        engineRequestBatchCount++
        engineRequestBatchSizes.push(Array.isArray(queries) ? queries.length : 0)
        return originalRequestBatch(queries, options)
      }

      await prisma.user.create({ data: { posts: { create: {} }, ...user1 } })
      await prisma.user.create({ data: user2 })

      // @ts-expect-error - client not typed for log opts for cross generator compatibility - can be improved once we drop the prisma-client-js generator
      prisma.$on('query', ({ query }: Prisma.QueryEvent) => {
        const normalized = query.trim().toUpperCase()

        // Track transaction boundary statements separately from "real" queries.
        // In qpe=remote mode, BEGIN isn't consistently surfaced in query logs, but COMMIT often is.
        if (normalized.startsWith('BEGIN')) beginCount++
        if (normalized.startsWith('COMMIT')) commitCount++