feat: pgbouncer container support (#26)

Co-authored-by: Severin Ibarluzea <seve@seve-squid.local>

Author: seveibar

src/index.ts

@@ -58,6 +58,7 @@ export const getTestPostgresDatabaseFactory = <
   const initialData: InitialWorkerData = {
     postgresVersion: options?.postgresVersion ?? "14",
     containerOptions: options?.container,
+    pgbouncerOptions: options?.pgbouncer,
   }
 
   const workerPromise = getWorker(initialData, options as any)
@@ -74,9 +75,18 @@ export const getTestPostgresDatabaseFactory = <
         connectionString: connectionDetailsFromWorker.connectionString,
       })
 
+      let pgbouncerPool: Pool | undefined
+      if (connectionDetailsFromWorker.pgbouncerConnectionString) {
+        pgbouncerPool = new Pool({
+          connectionString:
+            connectionDetailsFromWorker.pgbouncerConnectionString,
+        })
+      }
+
       t.teardown(async () => {
         try {
           await pool.end()
+          await pgbouncerPool?.end()
         } catch (error) {
           if (
             (error as Error).message.includes(
@@ -93,6 +103,7 @@ export const getTestPostgresDatabaseFactory = <
       return {
         ...connectionDetailsFromWorker,
         pool,
+        pgbouncerPool,
       }
     }
 

src/internal-types.ts

@@ -8,6 +8,7 @@ import type {
 export interface InitialWorkerData {
   postgresVersion: string
   containerOptions?: GetTestPostgresDatabaseFactoryOptions<any>["container"]
+  pgbouncerOptions?: GetTestPostgresDatabaseFactoryOptions<any>["pgbouncer"]
 }
 
 export type ConnectionDetailsFromWorker = Omit<ConnectionDetails, "pool">

src/public-types.ts

@@ -7,6 +7,8 @@ import { BindMode } from "testcontainers/build/types"
 export interface ConnectionDetails {
   connectionString: string
   connectionStringDocker: string
+  pgbouncerConnectionString?: string
+  pgbouncerConnectionStringDocker?: string
   dockerNetworkId: string
 
   host: string
@@ -16,6 +18,9 @@ export interface ConnectionDetails {
   database: string
 
   pool: Pool
+
+  // TODO if pgbouncer is enabled, this is defined, otherwise undefined
+  pgbouncerPool?: Pool
 }
 
 export interface GetTestPostgresDatabaseFactoryOptions<
@@ -32,6 +37,16 @@ export interface GetTestPostgresDatabaseFactoryOptions<
       mode?: BindMode
     }[]
   }
+
+  /**
+   * Pgbouncer container settings, disabled by default.
+   */
+  pgbouncer?: {
+    enabled: boolean
+    version?: string
+    poolMode?: "session" | "transaction" | "statement"
+  }
+
   /**
    * Test workers will be de-duped by this key. You probably don't need to set this.
    */

src/tests/pgbouncer.test.ts

@@ -0,0 +1,32 @@
+import test from "ava"
+import { QueryResult } from "pg"
+import { getTestPostgresDatabaseFactory } from "~/index"
+
+test("pgbouncer", async (t) => {
+  const getPostgres13 = getTestPostgresDatabaseFactory({
+    postgresVersion: "13.5",
+    pgbouncer: {
+      enabled: true,
+      version: "1.22.0",
+      poolMode: "statement",
+    },
+  })
+
+  const postgres13 = await getPostgres13(t)
+
+  t.truthy(postgres13.pgbouncerConnectionString)
+  const result = await postgres13.pool.query("SELECT 1 as result")
+
+  t.is(result.rows[0].result, 1)
+
+  // can't use a transaction with statement pool mode
+  const err = await t.throwsAsync(
+    postgres13.pgbouncerPool!.query(`BEGIN TRANSACTION`)
+  )
+
+  t.truthy(
+    err!
+      .toString()
+      .includes("transaction blocks not allowed in statement pooling mode")
+  )
+})

src/worker.ts

@@ -178,7 +178,9 @@ export class Worker {
     const databaseName = getRandomDatabaseName()
 
     // Create database
-    const { postgresClient, container } = await this.startContainerPromise
+    const { postgresClient, container, pgbouncerContainer } = await this
+      .startContainerPromise
+
     await postgresClient.query(`CREATE DATABASE ${databaseName};`)
 
     const msg = message.reply({
@@ -246,16 +248,29 @@ export class Worker {
   }
 
   private async getConnectionDetails(databaseName: string) {
-    const { container, network } = await this.startContainerPromise
+    const { container, network, pgbouncerContainer } = await this
+      .startContainerPromise
     const externalDatabaseUrl = `postgresql://postgres:@${container.getHost()}:${container.getMappedPort(
       5432
     )}/${databaseName}`
 
+    let pgbouncerConnectionString, pgbouncerConnectionStringDocker
+    if (pgbouncerContainer) {
+      pgbouncerConnectionString = `postgresql://postgres:@${pgbouncerContainer.getHost()}:${pgbouncerContainer.getMappedPort(
+        6432
+      )}/${databaseName}`
+      pgbouncerConnectionStringDocker = `postgresql://postgres:@${pgbouncerContainer
+        .getName()
+        .replace("/", "")}:5432/${databaseName}`
+    }
+
     return {
       connectionString: externalDatabaseUrl,
       connectionStringDocker: `postgresql://postgres:@${container
         .getName()
         .replace("/", "")}:5432/${databaseName}`,
+      pgbouncerConnectionString,
+      pgbouncerConnectionStringDocker,
       dockerNetworkId: network.getId(),
       host: container.getHost(),
       port: container.getMappedPort(5432),
@@ -300,13 +315,35 @@ export class Worker {
       )
     }
 
+    const connectionString = `postgresql://postgres:@${startedContainer.getHost()}:${startedContainer.getMappedPort(
+      5432
+    )}/postgres`
+
+    let startedPgbouncerContainer
+    if (this.initialData.pgbouncerOptions?.enabled) {
+      const pgbouncerContainer = new GenericContainer("edoburu/pgbouncer")
+        .withExposedPorts(6432)
+        .withName(getRandomDatabaseName())
+        .withEnvironment({
+          DB_HOST: startedContainer.getName().replace("/", ""),
+          DB_USER: "postgres",
+          DB_NAME: "*",
+          POOL_MODE:
+            this.initialData.pgbouncerOptions?.poolMode ?? "transaction",
+          LISTEN_PORT: "6432",
+          AUTH_TYPE: "trust",
+        })
+        .withStartupTimeout(120_000)
+        .withNetwork(network)
+      startedPgbouncerContainer = await pgbouncerContainer.start()
+    }
+
     return {
       container: startedContainer,
+      pgbouncerContainer: startedPgbouncerContainer,
       network,
       postgresClient: new pg.Pool({
-        connectionString: `postgresql://postgres:@${startedContainer.getHost()}:${startedContainer.getMappedPort(
-          5432
-        )}/postgres`,
+        connectionString,
       }),
     }
   }