Support custom callback on acquireConnection()

[giorgiogross]

trafficstars

As recently mentioned in https://github.com/kysely-org/kysely/issues/330#issuecomment-2154474432 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 https://github.com/kysely-org/kysely-postgres-js/blob/48af0e0684a34e037027958ff9a4fe0768c17900/src/driver.ts#L18-L22 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.