Enable nested beforeTemplateIsBaked calls

README.md

@@ -148,6 +148,55 @@ This works across the entire test suite.
 
 Note that if unique parameters are passed to the `beforeTemplateIsBaked` (`null` in the above example), separate databases will still be created.
 
+### Manual template creation
+
+In some cases, if you do extensive setup in your `beforeTemplateIsBaked` hook, you might want to obtain a separate, additional database within it if your application uses several databases for different purposes. This is possible by using the `manuallyBuildAdditionalTemplate()` function passed to your hook callback:
+
+```ts
+import test from "ava"
+
+const getTestDatabase = getTestPostgresDatabaseFactory<DatabaseParams>({
+  beforeTemplateIsBaked: async ({
+    params,
+    connection: { pool },
+    manuallyBuildAdditionalTemplate,
+  }) => {
+    await pool.query(`CREATE TABLE "bar" ("id" SERIAL PRIMARY KEY)`)
+
+    const fooTemplateBuilder = await manuallyBuildAdditionalTemplate()
+    await fooTemplateBuilder.connection.pool.query(
+      `CREATE TABLE "foo" ("id" SERIAL PRIMARY KEY)`
+    )
+    const { templateName: fooTemplateName } = await fooTemplateBuilder.finish()
+
+    return { fooTemplateName }
+  },
+})
+
+test("foo", async (t) => {
+  const barDatabase = await getTestDatabase({ type: "bar" })
+
+  // the "bar" database has the "bar" table...
+  await t.notThrowsAsync(async () => {
+    await barDatabase.pool.query(`SELECT * FROM "bar"`)
+  })
+
+  // ...but not the "foo" table...
+  await t.throwsAsync(async () => {
+    await barDatabase.pool.query(`SELECT * FROM "foo"`)
+  })
+
+  // ...and we can obtain a separate database with the "foo" table
+  const fooDatabase = await getTestDatabase.fromTemplate(
+    t,
+    barDatabase.beforeTemplateIsBakedResult.fooTemplateName
+  )
+  await t.notThrowsAsync(async () => {
+    await fooDatabase.pool.query(`SELECT * FROM "foo"`)
+  })
+})
+```
+
 ### Bind mounts & `exec`ing in the container
 
 `ava-postgres` uses [testcontainers](https://www.npmjs.com/package/testcontainers) under the hood to manage the Postgres container.

src/index.ts

@@ -12,12 +12,12 @@ import type {
   GetTestPostgresDatabase,
   GetTestPostgresDatabaseFactoryOptions,
   GetTestPostgresDatabaseOptions,
+  GetTestPostgresDatabaseResult,
 } from "./public-types"
 import { Pool } from "pg"
 import type { Jsonifiable } from "type-fest"
 import type { ExecutionContext } from "ava"
-import { once } from "node:events"
-import { createBirpc } from "birpc"
+import { BirpcReturn, createBirpc } from "birpc"
 import { ExecResult } from "testcontainers"
 import isPlainObject from "lodash/isPlainObject"
 
@@ -105,7 +105,7 @@ export const getTestPostgresDatabaseFactory = <
   Params extends Jsonifiable = never
 >(
   options?: GetTestPostgresDatabaseFactoryOptions<Params>
-) => {
+): GetTestPostgresDatabase<Params> => {
   const initialData: InitialWorkerData = {
     postgresVersion: options?.postgresVersion ?? "14",
     containerOptions: options?.container,
@@ -136,57 +136,73 @@ export const getTestPostgresDatabaseFactory = <
   }
 
   let rpcCallback: (data: any) => void
-  const rpc = createBirpc<SharedWorkerFunctions, TestWorkerFunctions>(
-    {
-      runBeforeTemplateIsBakedHook: async (connection, params) => {
-        if (options?.beforeTemplateIsBaked) {
-          const connectionDetails =
-            mapWorkerConnectionDetailsToConnectionDetails(connection)
-
-          // Ignore if the pool is terminated by the shared worker
-          // (This happens in CI for some reason even though we drain the pool first.)
-          connectionDetails.pool.on("error", (error) => {
-            if (
-              error.message.includes(
-                "terminating connection due to administrator command"
-              )
-            ) {
-              return
-            }
+  const rpc: BirpcReturn<SharedWorkerFunctions, TestWorkerFunctions> =
+    createBirpc<SharedWorkerFunctions, TestWorkerFunctions>(
+      {
+        runBeforeTemplateIsBakedHook: async (connection, params) => {
+          if (options?.beforeTemplateIsBaked) {
+            const connectionDetails =
+              mapWorkerConnectionDetailsToConnectionDetails(connection)
 
-            throw error
-          })
+            // Ignore if the pool is terminated by the shared worker
+            // (This happens in CI for some reason even though we drain the pool first.)
+            connectionDetails.pool.on("error", (error) => {
+              if (
+                error.message.includes(
+                  "terminating connection due to administrator command"
+                )
+              ) {
+                return
+              }
 
-          const hookResult = await options.beforeTemplateIsBaked({
-            params: params as any,
-            connection: connectionDetails,
-            containerExec: async (command): Promise<ExecResult> =>
-              rpc.execCommandInContainer(command),
-          })
+              throw error
+            })
 
-          await teardownConnection(connectionDetails)
+            const hookResult = await options.beforeTemplateIsBaked({
+              params: params as any,
+              connection: connectionDetails,
+              containerExec: async (command): Promise<ExecResult> =>
+                rpc.execCommandInContainer(command),
+              // This is what allows a consumer to get a "nested" database from within their beforeTemplateIsBaked hook
+              manuallyBuildAdditionalTemplate: async () => {
+                const connection =
+                  mapWorkerConnectionDetailsToConnectionDetails(
+                    await rpc.createEmptyDatabase()
+                  )
 
-          if (hookResult && !isSerializable(hookResult)) {
-            throw new TypeError(
-              "Return value of beforeTemplateIsBaked() hook could not be serialized. Make sure it returns only JSON-serializable values."
-            )
-          }
+                return {
+                  connection,
+                  finish: async () => {
+                    await teardownConnection(connection)
+                    return rpc.convertDatabaseToTemplate(connection.database)
+                  },
+                }
+              },
+            })
 
-          return hookResult
-        }
-      },
-    },
-    {
-      post: async (data) => {
-        const worker = await workerPromise
-        await worker.available
-        worker.publish(data)
-      },
-      on: (data) => {
-        rpcCallback = data
+            await teardownConnection(connectionDetails)
+
+            if (hookResult && !isSerializable(hookResult)) {
+              throw new TypeError(
+                "Return value of beforeTemplateIsBaked() hook could not be serialized. Make sure it returns only JSON-serializable values."
+              )
+            }
+
+            return hookResult
+          }
+        },
       },
-    }
-  )
+      {
+        post: async (data) => {
+          const worker = await workerPromise
+          await worker.available
+          worker.publish(data)
+        },
+        on: (data) => {
+          rpcCallback = data
+        },
+      }
+    )
 
   // Automatically cleaned up by AVA since each test file runs in a separate worker
   const _messageHandlerPromise = (async () => {
@@ -198,11 +214,11 @@ export const getTestPostgresDatabaseFactory = <
     }
   })()
 
-  const getTestPostgresDatabase: GetTestPostgresDatabase<Params> = async (
+  const getTestPostgresDatabase = async (
     t: ExecutionContext,
     params: any,
     getTestDatabaseOptions?: GetTestPostgresDatabaseOptions
-  ) => {
+  ): Promise<GetTestPostgresDatabaseResult> => {
     const testDatabaseConnection = await rpc.getTestDatabase({
       databaseDedupeKey: getTestDatabaseOptions?.databaseDedupeKey,
       params,
@@ -223,7 +239,22 @@ export const getTestPostgresDatabaseFactory = <
     }
   }
 
-  return getTestPostgresDatabase
+  getTestPostgresDatabase.fromTemplate = async (
+    t: ExecutionContext,
+    templateName: string
+  ) => {
+    const connection = mapWorkerConnectionDetailsToConnectionDetails(
+      await rpc.createDatabaseFromTemplate(templateName)
+    )
+
+    t.teardown(async () => {
+      await teardownConnection(connection)
+    })
+
+    return connection
+  }
+
+  return getTestPostgresDatabase as any
 }
 
 export * from "./public-types"

src/internal-types.ts

@@ -30,4 +30,11 @@ export interface SharedWorkerFunctions {
     beforeTemplateIsBakedResult: unknown
   }>
   execCommandInContainer: (command: string[]) => Promise<ExecResult>
+  createEmptyDatabase: () => Promise<ConnectionDetailsFromWorker>
+  createDatabaseFromTemplate: (
+    templateName: string
+  ) => Promise<ConnectionDetailsFromWorker>
+  convertDatabaseToTemplate: (
+    databaseName: string
+  ) => Promise<{ templateName: string }>
 }

src/public-types.ts

@@ -1,8 +1,8 @@
 import type { Pool } from "pg"
 import type { Jsonifiable } from "type-fest"
-import { ExecutionContext } from "ava"
-import { ExecResult } from "testcontainers"
-import { BindMode } from "testcontainers/build/types"
+import type { ExecutionContext } from "ava"
+import type { ExecResult } from "testcontainers"
+import type { BindMode } from "testcontainers/build/types"
 
 export interface ConnectionDetails {
   connectionString: string
@@ -58,6 +58,59 @@ export interface GetTestPostgresDatabaseFactoryOptions<
     connection: ConnectionDetails
     params: Params
     containerExec: (command: string[]) => Promise<ExecResult>
+    /**
+     * In some cases, if you do extensive setup in your `beforeTemplateIsBaked` hook, you might want to obtain a separate, additional database within it if your application uses several databases for different purposes.
+     *
+     * @example
+     * ```ts
+     * import test from "ava"
+     *
+     * const getTestDatabase = getTestPostgresDatabaseFactory<DatabaseParams>({
+     *   beforeTemplateIsBaked: async ({
+     *     params,
+     *     connection: { pool },
+     *     manuallyBuildAdditionalTemplate,
+     *   }) => {
+     *     await pool.query(`CREATE TABLE "bar" ("id" SERIAL PRIMARY KEY)`)
+     *
+     *     const fooTemplateBuilder = await manuallyBuildAdditionalTemplate()
+     *     await fooTemplateBuilder.connection.pool.query(
+     *       `CREATE TABLE "foo" ("id" SERIAL PRIMARY KEY)`
+     *     )
+     *     const { templateName: fooTemplateName } = await fooTemplateBuilder.finish()
+     *
+     *     return { fooTemplateName }
+     *   },
+     * })
+     *
+     * test("foo", async (t) => {
+     *   const barDatabase = await getTestDatabase({ type: "bar" })
+     *
+     *   // the "bar" database has the "bar" table...
+     *   await t.notThrowsAsync(async () => {
+     *     await barDatabase.pool.query(`SELECT * FROM "bar"`)
+     *   })
+     *
+     *   // ...but not the "foo" table...
+     *   await t.throwsAsync(async () => {
+     *     await barDatabase.pool.query(`SELECT * FROM "foo"`)
+     *   })
+     *
+     *   // ...and we can obtain a separate database with the "foo" table
+     *   const fooDatabase = await getTestDatabase.fromTemplate(
+     *     t,
+     *     barDatabase.beforeTemplateIsBakedResult.fooTemplateName
+     *   )
+     *   await t.notThrowsAsync(async () => {
+     *     await fooDatabase.pool.query(`SELECT * FROM "foo"`)
+     *   })
+     * })
+     * ```
+     */
+    manuallyBuildAdditionalTemplate: () => Promise<{
+      connection: ConnectionDetails
+      finish: () => Promise<{ templateName: string }>
+    }>
   }) => Promise<any>
 }
 
@@ -94,14 +147,23 @@ export type GetTestPostgresDatabaseOptions = {
 // https://github.com/microsoft/TypeScript/issues/23182#issuecomment-379091887
 type IsNeverType<T> = [T] extends [never] ? true : false
 
+interface BaseGetTestPostgresDatabase {
+  fromTemplate(
+    t: ExecutionContext,
+    templateName: string
+  ): Promise<ConnectionDetails>
+}
+
 export type GetTestPostgresDatabase<Params> = IsNeverType<Params> extends true
-  ? (
+  ? ((
       t: ExecutionContext,
       args?: null,
       options?: GetTestPostgresDatabaseOptions
-    ) => Promise<GetTestPostgresDatabaseResult>
-  : (
+    ) => Promise<GetTestPostgresDatabaseResult>) &
+      BaseGetTestPostgresDatabase
+  : ((
       t: ExecutionContext,
       args: Params,
       options?: GetTestPostgresDatabaseOptions
-    ) => Promise<GetTestPostgresDatabaseResult>
+    ) => Promise<GetTestPostgresDatabaseResult>) &
+      BaseGetTestPostgresDatabase

src/tests/hooks.test.ts

@@ -145,3 +145,41 @@ test("beforeTemplateIsBaked (result isn't serializable)", async (t) => {
     }
   )
 })
+
+test("beforeTemplateIsBaked with manual template build", async (t) => {
+  const getTestDatabase = getTestPostgresDatabaseFactory({
+    postgresVersion: process.env.POSTGRES_VERSION,
+    workerDedupeKey: "beforeTemplateIsBakedHookManualTemplateBuild",
+    beforeTemplateIsBaked: async ({
+      connection: { pool },
+      manuallyBuildAdditionalTemplate,
+    }) => {
+      await pool.query(`CREATE TABLE "bar" ("id" SERIAL PRIMARY KEY)`)
+
+      const fooTemplateBuilder = await manuallyBuildAdditionalTemplate()
+      await fooTemplateBuilder.connection.pool.query(
+        `CREATE TABLE "foo" ("id" SERIAL PRIMARY KEY)`
+      )
+      const { templateName: fooTemplateName } =
+        await fooTemplateBuilder.finish()
+
+      return { fooTemplateName }
+    },
+  })
+
+  const barDatabase = await getTestDatabase(t)
+  t.truthy(barDatabase.beforeTemplateIsBakedResult.fooTemplateName)
+
+  const fooDatabase = await getTestDatabase.fromTemplate(
+    t,
+    barDatabase.beforeTemplateIsBakedResult.fooTemplateName
+  )
+
+  await t.notThrowsAsync(async () => {
+    await fooDatabase.pool.query('SELECT * FROM "foo"')
+  }, "foo table should exist on database manually created from template")
+
+  await t.throwsAsync(async () => {
+    await fooDatabase.pool.query('SELECT * FROM "bar"')
+  })
+})