Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 79 additions & 0 deletions packages/adapter-pg/src/__tests__/pg.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -153,3 +153,82 @@ describe('PrismaPgAdapterFactory', () => {
await adapter.dispose()
})
})

describe('query serialization', () => {
const query = (sql: string): SqlQuery => ({ sql, args: [], argTypes: [] })
const emptyResult = { rows: [], fields: [], rowCount: 0 }

function trackingQueryMock() {
let inFlight = 0
let maxInFlight = 0
const started: string[] = []
const mock = vi.fn(async ({ text }: { text: string }) => {
started.push(text)
maxInFlight = Math.max(maxInFlight, ++inFlight)
await new Promise((resolve) => setImmediate(resolve))
inFlight--
return emptyResult
})
return { mock, started, maxInFlight: () => maxInFlight }
}

async function connectedAdapter() {
const factory = new PrismaPgAdapterFactory('postgresql://test:test@localhost:5432/test')
return await factory.connect()
}

it('serializes concurrent queries on a transaction connection', async () => {
const adapter = await connectedAdapter()
const { mock, maxInFlight } = trackingQueryMock()
const mockConnection = { on: vi.fn(), removeListener: vi.fn(), query: mock, release: vi.fn() }
adapter['client'].connect = vi.fn().mockResolvedValue(mockConnection)

const transaction = await adapter.startTransaction()
await Promise.all([
transaction.queryRaw(query('SELECT 1')),
transaction.queryRaw(query('SELECT 2')),
transaction.queryRaw(query('SELECT 3')),
])

// A pg.PoolClient is a single connection: queries must never overlap.
expect(maxInFlight()).toBe(1)
await transaction.commit()
await adapter.dispose()
})

it('does not serialize queries on the pool', async () => {
const adapter = await connectedAdapter()
const { mock, maxInFlight } = trackingQueryMock()
adapter['client'].query = mock

await Promise.all([
adapter.queryRaw(query('SELECT 1')),
adapter.queryRaw(query('SELECT 2')),
adapter.queryRaw(query('SELECT 3')),
])

// The pool handles concurrency itself; serializing here would limit the
// whole application to one query at a time.
expect(maxInFlight()).toBe(3)
await adapter.dispose()
})

it('keeps serializing after a failed query', async () => {
const adapter = await connectedAdapter()
const mock = vi
.fn()
.mockResolvedValueOnce(emptyResult) // BEGIN
.mockRejectedValueOnce(new Error('boom'))
.mockResolvedValue(emptyResult)
const mockConnection = { on: vi.fn(), removeListener: vi.fn(), query: mock, release: vi.fn() }
adapter['client'].connect = vi.fn().mockResolvedValue(mockConnection)

const tx = await adapter.startTransaction()
const failing = tx.queryRaw(query('SELECT 1'))
const following = tx.queryRaw(query('SELECT 2'))

await expect(failing).rejects.toThrow()
await expect(following).resolves.toBeDefined()
await adapter.dispose()
})
})
21 changes: 21 additions & 0 deletions packages/adapter-pg/src/pg.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,13 @@ class PgQueryable<ClientT extends StdClient | TransactionClient> implements SqlQ
readonly provider = 'postgres'
readonly adapterName = packageName

// `pg.Client` and `pg.PoolClient` are single connections and don't support
// concurrent queries (deprecated in pg@8, an error in pg@9), so queries must
// be serialized. `pg.Pool` handles concurrency itself and must not be
// serialized, or the whole pool would be limited to one query at a time.
protected readonly serializeQueries: boolean = true
#queryLock: Promise<unknown> = Promise.resolve()

constructor(
protected readonly client: ClientT,
protected readonly pgOptions?: PrismaPgOptions,
Expand Down Expand Up @@ -99,6 +106,18 @@ class PgQueryable<ClientT extends StdClient | TransactionClient> implements SqlQ
* marked as unhealthy.
*/
private async performIO(query: SqlQuery): Promise<pg.QueryArrayResult<any>> {
if (!this.serializeQueries) {
return this.#performIO(query)
}
const previous = this.#queryLock
const current = previous.then(() => this.#performIO(query))
// Keep the lock chain alive even if the query fails; the failure still
// propagates to the caller through `current`.
this.#queryLock = current.catch(() => {})
return current
}

async #performIO(query: SqlQuery): Promise<pg.QueryArrayResult<any>> {
const { sql, args } = query
const values = args.map((arg, i) => mapArg(arg, query.argTypes[i]))

Expand Down Expand Up @@ -197,6 +216,8 @@ export type UserDefinedTypeParser = (oid: number, value: unknown, adapter: SqlQu
export type StatementNameGenerator = (query: SqlQuery) => string

export class PrismaPgAdapter extends PgQueryable<StdClient> implements SqlDriverAdapter {
protected override readonly serializeQueries = false

constructor(
client: StdClient,
protected readonly pgOptions?: PrismaPgOptions,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,47 @@ test('merges chunked query results without overflowing the stack', async () => {
expect(result).toHaveLength(rowsPerLaterChunk)
})

// Loading sibling relations concurrently is intentional: adapters whose connection
// cannot run queries concurrently (e.g. a single pg connection) are responsible for
// serializing them in `performIO` (see https://github.com/prisma/prisma/issues/29407).
// This pins the interpreter side of that contract so join loading stays parallel.
test('loads join children in parallel', async () => {
let inFlight = 0
let maxInFlight = 0
const queryable: SqlQueryable = {
provider: 'postgres',
adapterName: 'test',
queryRaw: async () => {
maxInFlight = Math.max(maxInFlight, ++inFlight)
await new Promise((resolve) => setImmediate(resolve))
inFlight--
return userResultSet(1, 'Alice')
},
executeRaw: () => Promise.resolve(0),
}

const joinChild = (parentField: string) => ({
child: queryNode(`SELECT * FROM ${parentField}`),
on: [['id', 'id']] as [string, string][],
parentField,
isRelationUnique: true,
})

const queryPlan: QueryPlanNode = {
type: 'join',
args: {
parent: queryNode('SELECT * FROM users'),
children: [joinChild('posts'), joinChild('profile'), joinChild('settings')],
canAssumeStrictEquality: true,
},
}

const interpreter = QueryInterpreter.forSql({ tracingHelper: noopTracingHelper })
await interpreter.run(queryPlan, { queryable, transactionManager: { enabled: false }, scope: {} })

expect(maxInFlight).toBe(3)
})

class MockTransactionAdapter implements SqlDriverAdapter {
adapterName = 'mock-adapter'
provider = 'postgres' as const
Expand Down