Support custom callback on acquireConnection()

[giorgiogross]

As recently mentioned in kysely-org/kysely#330 (comment) I'm using RLS and configure it while initiating the client connection. For various reasons I had to switch to postgresjs where currently there is no configuration param exposed for hooking into the db client lifecycle.

As a workaround, I'm manually overriding Dialect.createDriver() as such:

		const config = { ... };
		const dialect = new PostgresJSDialect(config);
		dialect.createDriver = (): Driver => {
			return new RlsDriver(config, this.vars);
		};

And am subclassing PostgresJSDriver as such:

class RlsDriver extends PostgresJSDriver {
	constructor(
		protected config: PostgresJSDialectConfig,
		protected vars: UserVariables
	) {
		super(config);
	}

        async acquireConnection(): Promise<PostgresJSConnection> {
		const reservedConnection = await this.config.postgres.reserve();

                // apply any configs for RLS
		await reservedConnection`SELECT set_config('request.jwt.claim.sub', ${this.vars.userId}, FALSE)`;
		await reservedConnection`SET ROLE authenticated`;
		
		return new PostgresJSConnection(reservedConnection);
         }
}

Feels a bit hacky bc I'm exposing that private #config prop, and it seems like the API wasn't designed for this use case.

Proposed solution

It would be helpful to have a config param on PostgresJSDialectConfig that would allow to specify a callback that may modify the reservedConnection at

kysely-postgres-js/src/driver.ts

Lines 18 to 22 in 48af0e0

	async acquireConnection(): Promise<PostgresJSConnection> {
	const reservedConnection = await this.#config.postgres.reserve()
	return new PostgresJSConnection(reservedConnection)
	}

For example:

const reservedConnection = await this.#config.postgres.reserve()
const modifiedConnection = await this.#config.onAcquireConnection?.(reservedConnection) ?? reservedConnection;
return new PostgresJSConnection(modifiedConnection)

For completeness onDestroyConnection might also be useful.. There might also be a better place to put that param or use a different approach, in any case the goal is to modify fresh connections before they are used.