New structure working

Author: codetheweb

src/index.ts

@@ -12,6 +12,7 @@ import type {
   GetTestPostgresDatabase,
   GetTestPostgresDatabaseFactoryOptions,
   GetTestPostgresDatabaseOptions,
+  GetTestPostgresDatabaseResult,
 } from "./public-types"
 import { Pool } from "pg"
 import type { Jsonifiable } from "type-fest"
@@ -104,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,
@@ -157,41 +158,28 @@ 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 =
+              manuallyBuildAdditionalTemplate: async () => {
+                const connection =
                   mapWorkerConnectionDetailsToConnectionDetails(
-                    connectionDetails
+                    await rpc.createEmptyDatabase()
                   )
 
-                createdNestedConnections.push(mappedConnection)
-
                 return {
-                  ...mappedConnection,
-                  beforeTemplateIsBakedResult,
+                  connection,
+                  finish: async () => {
+                    await teardownConnection(connection)
+                    return rpc.convertDatabaseToTemplate(connection.database)
+                  },
                 }
               },
             })
 
-            await Promise.all(
-              createdNestedConnections.map(async (connection) => {
-                await teardownConnection(connection)
-                await rpc.dropDatabase(connection.database)
-              })
-            )
-
             await teardownConnection(connectionDetails)
 
             if (hookResult && !isSerializable(hookResult)) {
@@ -226,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,
@@ -251,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,5 +30,11 @@ export interface SharedWorkerFunctions {
     beforeTemplateIsBakedResult: unknown
   }>
   execCommandInContainer: (command: string[]) => Promise<ExecResult>
-  dropDatabase: (databaseName: string) => Promise<void>
+  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
@@ -92,11 +92,10 @@ export interface GetTestPostgresDatabaseFactoryOptions<
      * })
      * ```
      */
-    beforeTemplateIsBaked: (
-      options: {
-        params: Params
-      } & Pick<GetTestPostgresDatabaseOptions, "databaseDedupeKey">
-    ) => Promise<GetTestPostgresDatabaseResult>
+    manuallyBuildAdditionalTemplate: () => Promise<{
+      connection: ConnectionDetails
+      finish: () => Promise<{ templateName: string }>
+    }>
   }) => Promise<any>
 }
 
@@ -155,14 +154,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

@@ -148,46 +148,39 @@ 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>({
+  const getTestDatabase = getTestPostgresDatabaseFactory({
     postgresVersion: process.env.POSTGRES_VERSION,
     workerDedupeKey: "beforeTemplateIsBakedHookNestedDatabase",
     beforeTemplateIsBaked: async ({
-      params,
       connection: { pool },
-      beforeTemplateIsBaked,
+      manuallyBuildAdditionalTemplate,
     }) => {
-      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
+      const fooTemplateBuilder = await manuallyBuildAdditionalTemplate()
+      await fooTemplateBuilder.connection.pool.query(
+        `CREATE TABLE "foo" ("id" SERIAL PRIMARY KEY)`
+      )
+      const { templateName: fooTemplateName } =
+        await fooTemplateBuilder.finish()
 
-      await t.notThrowsAsync(async () => {
-        await fooDatabase.pool.query(`INSERT INTO "foo" DEFAULT VALUES`)
-      })
-
-      return { createdBar: true }
+      return { fooTemplateName }
     },
   })
 
-  const database = await getTestServer(t, { type: "bar" })
-  t.deepEqual(database.beforeTemplateIsBakedResult, { createdBar: true })
+  const barDatabase = await getTestDatabase(t)
+  t.truthy(barDatabase.beforeTemplateIsBakedResult.fooTemplateName)
 
-  t.false(
-    await doesDatabaseExist(database.pool, nestedDatabaseName!),
-    "Nested database should have been cleaned up after the parent hook completed"
+  const fooDatabase = await getTestDatabase.fromTemplate(
+    t,
+    barDatabase.beforeTemplateIsBakedResult.fooTemplateName
   )
+
+  await t.notThrowsAsync(async () => {
+    await fooDatabase.pool.query('SELECT * FROM "foo"')
+  })
+
+  await t.throwsAsync(async () => {
+    await fooDatabase.pool.query('SELECT * FROM "bar"')
+  })
 })

src/worker.ts

@@ -73,9 +73,31 @@ export class Worker {
           const container = (await this.startContainerPromise).container
           return container.exec(command)
         },
-        dropDatabase: async (databaseName) => {
+        createEmptyDatabase: async () => {
           const { postgresClient } = await this.startContainerPromise
-          await postgresClient.query(`DROP DATABASE ${databaseName}`)
+          const databaseName = getRandomDatabaseName()
+          await postgresClient.query(`CREATE DATABASE ${databaseName}`)
+          return this.getConnectionDetails(databaseName)
+        },
+        createDatabaseFromTemplate: async (templateName) => {
+          const { postgresClient } = await this.startContainerPromise
+          const databaseName = getRandomDatabaseName()
+          await postgresClient.query(
+            `CREATE DATABASE ${databaseName} WITH TEMPLATE ${templateName};`
+          )
+
+          testWorker.teardown(async () => {
+            await this.teardownDatabase(databaseName)
+          })
+
+          return this.getConnectionDetails(databaseName)
+        },
+        convertDatabaseToTemplate: async (databaseName) => {
+          const { postgresClient } = await this.startContainerPromise
+          await postgresClient.query(
+            `ALTER DATABASE ${databaseName} WITH is_template TRUE;`
+          )
+          return { templateName: databaseName }
         },
       },
       rpcChannel
@@ -152,17 +174,7 @@ export class Worker {
         return
       }
 
-      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
-      }
+      await this.teardownDatabase(databaseName!)
     })
 
     return {
@@ -171,6 +183,22 @@ export class Worker {
     }
   }
 
+  private async teardownDatabase(databaseName: string) {
+    const { postgresClient } = await this.startContainerPromise
+
+    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
+    }
+  }
+
   private async createTemplate(rpc: WorkerRpc, params?: Jsonifiable) {
     const databaseName = getRandomDatabaseName()