Enable nested beforeTemplateIsBaked calls

src/index.ts

@@ -16,7 +16,6 @@ import type {
 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 { ExecResult } from "testcontainers"
 import isPlainObject from "lodash/isPlainObject"
@@ -157,13 +156,39 @@ export const getTestPostgresDatabaseFactory = <
             throw error
           })
 
+          const createdNestedConnections: 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
+            beforeTemplateIsBaked: async (options) => {
+              const { connectionDetails, beforeTemplateIsBakedResult } =
+                await rpc.getTestDatabase({
+                  params: options.params,
+                  databaseDedupeKey: options.databaseDedupeKey,
+                })
+
+              const mappedConnection =
+                mapWorkerConnectionDetailsToConnectionDetails(connectionDetails)
+
+              createdNestedConnections.push(mappedConnection)
+
+              return {
+                ...mappedConnection,
+                beforeTemplateIsBakedResult,
+              }
+            },
           })
 
+          await Promise.all(
+            createdNestedConnections.map(async (connection) => {
+              await teardownConnection(connection)
+              await rpc.dropDatabase(connection.database)
+            })
+          )
+
           await teardownConnection(connectionDetails)
 
           if (hookResult && !isSerializable(hookResult)) {

src/internal-types.ts

@@ -30,4 +30,5 @@ export interface SharedWorkerFunctions {
     beforeTemplateIsBakedResult: unknown
   }>
   execCommandInContainer: (command: string[]) => Promise<ExecResult>
+  dropDatabase: (databaseName: string) => Promise<void>
 }

src/public-types.ts

@@ -55,6 +55,11 @@ export interface GetTestPostgresDatabaseFactoryOptions<
     connection: ConnectionDetails
     params: Params
     containerExec: (command: string[]) => Promise<ExecResult>
+    beforeTemplateIsBaked: (
+      options: {
+        params: Params
+      } & Pick<GetTestPostgresDatabaseOptions, "databaseDedupeKey">
+    ) => Promise<GetTestPostgresDatabaseResult>
   }) => Promise<any>
 }
 

src/tests/hooks.test.ts

@@ -1,6 +1,7 @@
 import test from "ava"
 import { getTestPostgresDatabaseFactory } from "~/index"
 import { countDatabaseTemplates } from "./utils/count-database-templates"
+import { doesDatabaseExist } from "./utils/does-database-exist"
 
 test("beforeTemplateIsBaked", async (t) => {
   let wasHookCalled = false
@@ -145,3 +146,48 @@ test("beforeTemplateIsBaked (result isn't serializable)", async (t) => {
     }
   )
 })
+
+test("beforeTemplateIsBaked, get nested database", async (t) => {
+  type DatabaseParams = {
+    type: "foo" | "bar"
+  }
+
+  let nestedDatabaseName: string | undefined = undefined
+
+  const getTestServer = getTestPostgresDatabaseFactory<DatabaseParams>({
+    postgresVersion: process.env.POSTGRES_VERSION,
+    workerDedupeKey: "beforeTemplateIsBakedHookNestedDatabase",
+    beforeTemplateIsBaked: async ({
+      params,
+      connection: { pool },
+      beforeTemplateIsBaked,
+    }) => {
+      if (params.type === "foo") {
+        await pool.query(`CREATE TABLE "foo" ("id" SERIAL PRIMARY KEY)`)
+        return { createdFoo: true }
+      }
+
+      await pool.query(`CREATE TABLE "bar" ("id" SERIAL PRIMARY KEY)`)
+      const fooDatabase = await beforeTemplateIsBaked({
+        params: { type: "foo" },
+      })
+      t.deepEqual(fooDatabase.beforeTemplateIsBakedResult, { createdFoo: true })
+
+      nestedDatabaseName = fooDatabase.database
+
+      await t.notThrowsAsync(async () => {
+        await fooDatabase.pool.query(`INSERT INTO "foo" DEFAULT VALUES`)
+      })
+
+      return { createdBar: true }
+    },
+  })
+
+  const database = await getTestServer(t, { type: "bar" })
+  t.deepEqual(database.beforeTemplateIsBakedResult, { createdBar: true })
+
+  t.false(
+    await doesDatabaseExist(database.pool, nestedDatabaseName!),
+    "Nested database should have been cleaned up after the parent hook completed"
+  )
+})

src/tests/utils/does-database-exist.ts

@@ -0,0 +1,12 @@
+import { Pool } from "pg"
+
+export const doesDatabaseExist = async (pool: Pool, databaseName: string) => {
+  const {
+    rows: [{ count }],
+  } = await pool.query(
+    'SELECT COUNT(*) FROM "pg_database" WHERE "datname" = $1',
+    [databaseName]
+  )
+
+  return count > 0
+}

src/worker.ts

@@ -73,6 +73,10 @@ export class Worker {
           const container = (await this.startContainerPromise).container
           return container.exec(command)
         },
+        dropDatabase: async (databaseName) => {
+          const { postgresClient } = await this.startContainerPromise
+          await postgresClient.query(`DROP DATABASE ${databaseName}`)
+        },
       },
       rpcChannel
     )
@@ -148,8 +152,17 @@ export class Worker {
         return
       }
 
-      await this.forceDisconnectClientsFrom(databaseName!)
-      await postgresClient.query(`DROP DATABASE ${databaseName}`)
+      try {
+        await this.forceDisconnectClientsFrom(databaseName!)
+        await postgresClient.query(`DROP DATABASE ${databaseName}`)
+      } catch (error) {
+        if ((error as Error)?.message?.includes("does not exist")) {
+          // Database was likely a nested database and manually dropped by the test worker, ignore
+          return
+        }
+
+        throw error
+      }
     })
 
     return {