Parameterize mixed batches that start with a raw query

packages/client-engine-runtime/src/parameterization/parameterize-tests/batch.test.ts

@@ -406,4 +406,48 @@ describe('parameterizeBatch', () => {
       }
     `)
   })
+
+  it('parameterizes a model query that follows a raw statement', () => {
+    const batch: JsonBatchQuery = {
+      batch: [
+        {
+          action: 'executeRaw',
+          query: {
+            arguments: {
+              query: "SELECT set_config('app.user_id', $1, true)",
+              parameters: '["1"]',
+            },
+            selection: { $scalars: true },
+          },
+        },
+        {
+          modelName: 'User',
+          action: 'findUnique',
+          query: {
+            arguments: { where: { id: 1 } },
+            selection: { $scalars: true },
+          },
+        },
+      ],
+    }
+
+    const result = parameterizeBatch(batch, paramGraph)
+
+    expect(result.parameterizedBatch.batch[0].query.arguments).toEqual({
+      query: "SELECT set_config('app.user_id', $1, true)",
+      parameters: '["1"]',
+    })
+    expect(result.parameterizedBatch.batch[1].query.arguments).toEqual({
+      where: {
+        id: {
+          $type: 'Param',
+          value: {
+            name: '%1',
+            type: 'Int',
+          },
+        },
+      },
+    })
+    expect(result.placeholderValues).toEqual({ '%1': 1 })
+  })
 })

packages/client/src/runtime/core/engines/client/ClientEngine.ts

@@ -35,6 +35,7 @@ import { InteractiveTransactionInfo } from '../common/types/Transaction'
 import { getBatchRequestPayload } from '../common/utils/getBatchRequestPayload'
 import { getErrorMessageWithLink as genericGetErrorMessageWithLink } from '../common/utils/getErrorMessageWithLink'
 import type { Executor } from './Executor'
+import { isAllRawBatch } from './is-all-raw-batch'
 import { LocalExecutor } from './LocalExecutor'
 import { QueryPlanCache } from './query-plan-cache'
 import { RemoteExecutor } from './RemoteExecutor'
@@ -549,7 +550,7 @@ export class ClientEngine implements Engine {
       throw this.#transformRequestError(err, request)
     })
 
-    const hasRawQueries = firstModelName === undefined
+    const hasRawQueries = isAllRawBatch(queries)
     let batchResponse: BatchResponse
     let placeholderValues: Record<string, unknown> = {}
     let queryInfoQueries = queries.map((query) => query.query)

packages/client/src/runtime/core/engines/client/is-all-raw-batch.test.ts

@@ -0,0 +1,19 @@
+import { isAllRawBatch } from './is-all-raw-batch'
+
+describe('isAllRawBatch', () => {
+  it('is true for an empty batch', () => {
+    expect(isAllRawBatch([])).toBe(true)
+  })
+
+  it('is true when every item is raw', () => {
+    expect(isAllRawBatch([{ action: 'executeRaw' }, { action: 'queryRaw' }])).toBe(true)
+  })
+
+  it('is false when a raw item is followed by a model query', () => {
+    expect(isAllRawBatch([{ action: 'executeRaw' }, { action: 'updateMany', modelName: 'Item' }])).toBe(false)
+  })
+
+  it('is false when a model query is first', () => {
+    expect(isAllRawBatch([{ action: 'updateMany', modelName: 'Item' }, { action: 'executeRaw' }])).toBe(false)
+  })
+})

packages/client/src/runtime/core/engines/client/is-all-raw-batch.ts

@@ -0,0 +1,9 @@
+/**
+ * Whether every query in a batch is raw (no `modelName`).
+ *
+ * Mixed batches such as `$executeRaw` followed by a model query must return
+ * false so the model query can still be parameterized and plan-cached.
+ */
+export function isAllRawBatch(queries: Array<{ action: string; modelName?: string }>): boolean {
+  return queries.every((query) => query.modelName === undefined)
+}