Add snapshot support to postgres module

docs/modules/postgresql.md

@@ -25,3 +25,18 @@ npm install @testcontainers/postgresql --save-dev
 <!--codeinclude-->
 [Set username:](../../packages/modules/postgresql/src/postgresql-container.test.ts) inside_block:setUsername
 <!--/codeinclude-->
+
+### Using Snapshots
+
+This example shows the usage of the postgres module's Snapshot feature to give each test a clean database without having
+to recreate the database container on every test or run heavy scripts to clean your database. This makes the individual
+tests very modular, since they always run on a brand-new database.
+
+!!!tip
+    You should never pass the `"postgres"` system database as the container database name if you want to use snapshots. 
+    The Snapshot logic requires dropping the connected database and using the system database to run commands, which will
+    not work if the database for the container is set to `"postgres"`.
+
+<!--codeinclude-->
+[Test with a reusable Postgres container](../../packages/modules/postgresql/src/postgresql-container.test.ts) inside_block:createAndRestoreFromSnapshot
+<!--/codeinclude-->

packages/modules/postgresql/src/postgresql-container.test.ts

@@ -113,3 +113,209 @@ describe("PostgreSqlContainer", { timeout: 180_000 }, () => {
     await expect(() => container.start()).rejects.toThrow();
   });
 });
+
+describe("PostgreSqlContainer snapshot and restore", { timeout: 180_000 }, () => {
+  // createAndRestoreFromSnapshot {
+  it("should create and restore from snapshot", async () => {
+    const container = await new PostgreSqlContainer().start();
+
+    // Connect to the database
+    let client = new Client({
+      connectionString: container.getConnectionUri(),
+    });
+    await client.connect();
+
+    // Create some test data
+    await client.query("CREATE TABLE test_table (id SERIAL PRIMARY KEY, name TEXT)");
+    await client.query("INSERT INTO test_table (name) VALUES ('initial data')");
+
+    // Close connection before snapshot (otherwise we'll get an error because user is already connected)
+    await client.end();
+
+    // Take a snapshot
+    await container.snapshot();
+
+    // Reconnect to database
+    client = new Client({
+      connectionString: container.getConnectionUri(),
+    });
+    await client.connect();
+
+    // Modify the database
+    await client.query("INSERT INTO test_table (name) VALUES ('data after snapshot')");
+
+    // Verify both records exist
+    let result = await client.query("SELECT * FROM test_table ORDER BY id");
+    expect(result.rows).toHaveLength(2);
+    expect(result.rows[0].name).toEqual("initial data");
+    expect(result.rows[1].name).toEqual("data after snapshot");
+
+    // Close connection before restore (same reason as above)
+    await client.end();
+
+    // Restore to the snapshot
+    await container.restoreSnapshot();
+
+    // Reconnect to database
+    client = new Client({
+      connectionString: container.getConnectionUri(),
+    });
+    await client.connect();
+
+    // Verify only the initial data exists after restore
+    result = await client.query("SELECT * FROM test_table ORDER BY id");
+    expect(result.rows).toHaveLength(1);
+    expect(result.rows[0].name).toEqual("initial data");
+
+    await client.end();
+    await container.stop();
+  });
+  // }
+
+  it("should use custom snapshot name", async () => {
+    const container = await new PostgreSqlContainer().start();
+    const customSnapshotName = "my_custom_snapshot";
+
+    // Connect to the database
+    let client = new Client({
+      connectionString: container.getConnectionUri(),
+    });
+    await client.connect();
+
+    // Create a test table and insert data
+    await client.query("CREATE TABLE test_table (id SERIAL PRIMARY KEY, name TEXT)");
+    await client.query("INSERT INTO test_table (name) VALUES ('initial data')");
+
+    // Close connection before snapshot
+    await client.end();
+
+    // Take a snapshot with custom name
+    await container.snapshot(customSnapshotName);
+
+    // Reconnect to database
+    client = new Client({
+      connectionString: container.getConnectionUri(),
+    });
+    await client.connect();
+
+    // Modify the database
+    await client.query("INSERT INTO test_table (name) VALUES ('data after snapshot')");
+
+    // Close connection before restore
+    await client.end();
+
+    // Restore using the custom snapshot name
+    await container.restoreSnapshot(customSnapshotName);
+
+    // Reconnect to database
+    client = new Client({
+      connectionString: container.getConnectionUri(),
+    });
+    await client.connect();
+
+    // Verify only the initial data exists after restore
+    const result = await client.query("SELECT * FROM test_table ORDER BY id");
+    expect(result.rows).toHaveLength(1);
+    expect(result.rows[0].name).toEqual("initial data");
+
+    await client.end();
+    await container.stop();
+  });
+
+  it("should handle multiple snapshots", async () => {
+    const container = await new PostgreSqlContainer().start();
+
+    // Connect to the database
+    let client = new Client({
+      connectionString: container.getConnectionUri(),
+    });
+    await client.connect();
+
+    // Create a test table
+    await client.query("CREATE TABLE test_table (id SERIAL PRIMARY KEY, name TEXT)");
+
+    // Close connection before snapshot
+    await client.end();
+
+    // Take first snapshot with empty table
+    await container.snapshot("snapshot1");
+
+    // Reconnect to database
+    client = new Client({
+      connectionString: container.getConnectionUri(),
+    });
+    await client.connect();
+
+    // Add first record
+    await client.query("INSERT INTO test_table (name) VALUES ('data for snapshot 2')");
+
+    // Close connection before snapshot
+    await client.end();
+
+    // Take second snapshot with one record
+    await container.snapshot("snapshot2");
+
+    // Reconnect to database
+    client = new Client({
+      connectionString: container.getConnectionUri(),
+    });
+    await client.connect();
+
+    // Add second record
+    await client.query("INSERT INTO test_table (name) VALUES ('data after snapshots')");
+
+    // Verify we have two records
+    let result = await client.query("SELECT COUNT(*) as count FROM test_table");
+    expect(result.rows[0].count).toEqual("2");
+
+    // Close connection before restore
+    await client.end();
+
+    // Restore to first snapshot (empty table)
+    await container.restoreSnapshot("snapshot1");
+
+    // Reconnect to database
+    client = new Client({
+      connectionString: container.getConnectionUri(),
+    });
+    await client.connect();
+
+    // Verify table is empty
+    result = await client.query("SELECT COUNT(*) as count FROM test_table");
+    expect(result.rows[0].count).toEqual("0");
+
+    // Close connection before restore
+    await client.end();
+
+    // Restore to second snapshot (one record)
+    await container.restoreSnapshot("snapshot2");
+
+    // Reconnect to database
+    client = new Client({
+      connectionString: container.getConnectionUri(),
+    });
+    await client.connect();
+
+    // Verify we have one record
+    result = await client.query("SELECT * FROM test_table");
+    expect(result.rows).toHaveLength(1);
+    expect(result.rows[0].name).toEqual("data for snapshot 2");
+
+    await client.end();
+    await container.stop();
+  });
+
+  it("should throw an error when trying to snapshot postgres system database", async () => {
+    const container = await new PostgreSqlContainer().withDatabase("postgres").start();
+
+    await expect(container.snapshot()).rejects.toThrow(
+      "Snapshot feature is not supported when using the postgres system database"
+    );
+
+    await expect(container.restoreSnapshot()).rejects.toThrow(
+      "Snapshot feature is not supported when using the postgres system database"
+    );
+
+    await container.stop();
+  });
+});

packages/modules/postgresql/src/postgresql-container.ts

@@ -48,6 +48,7 @@ export class PostgreSqlContainer extends GenericContainer {
 }
 
 export class StartedPostgreSqlContainer extends AbstractStartedContainer {
+  private snapshotName: string = "migrated_template";
   constructor(
     startedTestContainer: StartedTestContainer,
     private readonly database: string,
@@ -85,4 +86,104 @@ export class StartedPostgreSqlContainer extends AbstractStartedContainer {
     url.password = this.getPassword();
     return url.toString();
   }
+
+  /**
+   * Sets the name to be used for database snapshots.
+   * This name will be used as the default for snapshot() and restore() methods.
+   *
+   * @param snapshotName The name to use for snapshots (default is "migrated_template" if this method is not called)
+   * @returns this (for method chaining)
+   */
+  public withSnapshotName(snapshotName: string): this {
+    this.snapshotName = snapshotName;
+    return this;
+  }
+
+  /**
+   * Takes a snapshot of the current state of the database as a template, which can then be restored using
+   * the restore method.
+   *
+   * @param snapshotName Name for the snapshot, defaults to the value set by withSnapshotName() or "migrated_template" if not specified
+   * @returns Promise resolving when snapshot is complete
+   * @throws Error if attempting to snapshot the postgres system database or if using the same name as the database
+   */
+  public async snapshot(snapshotName = this.snapshotName): Promise<void> {
+    this.snapshotSanityCheck(snapshotName);
+
+    // Execute the commands to create the snapshot, in order
+    await this.execCommandsSQL([
+      // Update pg_database to remove the template flag, then drop the database if it exists.
+      // This is needed because dropping a template database will fail.
+      `UPDATE pg_database SET datistemplate = FALSE WHERE datname = '${snapshotName}'`,
+      `DROP DATABASE IF EXISTS "${snapshotName}"`,
+      // Create a copy of the database to another database to use as a template now that it was fully migrated
+      `CREATE DATABASE "${snapshotName}" WITH TEMPLATE "${this.getDatabase()}" OWNER "${this.getUsername()}"`,
+      // Snapshot the template database so we can restore it onto our original database going forward
+      `ALTER DATABASE "${snapshotName}" WITH is_template = TRUE`,
+    ]);
+  }
+
+  /**
+   * Restores the database to a specific snapshot.
+   *
+   * @param snapshotName Name of the snapshot to restore from, defaults to the value set by withSnapshotName() or "migrated_template" if not specified
+   * @returns Promise resolving when restore is complete
+   * @throws Error if attempting to restore the postgres system database or if using the same name as the database
+   */
+  public async restoreSnapshot(snapshotName = this.snapshotName): Promise<void> {
+    this.snapshotSanityCheck(snapshotName);
+
+    // Execute the commands to restore the snapshot, in order
+    await this.execCommandsSQL([
+      // Drop the entire database by connecting to the postgres global database
+      `DROP DATABASE "${this.getDatabase()}" WITH (FORCE)`,
+      // Then restore the previous snapshot
+      `CREATE DATABASE "${this.getDatabase()}" WITH TEMPLATE "${snapshotName}" OWNER "${this.getUsername()}"`,
+    ]);
+  }
+
+  /**
+   * Executes a series of SQL commands against the Postgres database
+   *
+   * @param commands Array of SQL commands to execute in sequence
+   * @throws Error if any command fails to execute with details of the failure
+   */
+  private async execCommandsSQL(commands: string[]): Promise<void> {
+    for (const command of commands) {
+      try {
+        const result = await this.exec([
+          "psql",
+          "-v",
+          "ON_ERROR_STOP=1",
+          "-U",
+          this.getUsername(),
+          "-d",
+          "postgres",
+          "-c",
+          command,
+        ]);
+
+        if (result.exitCode !== 0) {
+          throw new Error(`Command failed with exit code ${result.exitCode}: ${result.output}`);
+        }
+      } catch (error) {
+        console.error(`Failed to execute command: ${command}`, error);
+        throw error;
+      }
+    }
+  }
+
+  /**
+   * Checks if the snapshot name is valid and if the database is not the postgres system database
+   * @param snapshotName The name of the snapshot to check
+   */
+  private snapshotSanityCheck(snapshotName: string): void {
+    if (this.getDatabase() === "postgres") {
+      throw new Error("Snapshot feature is not supported when using the postgres system database");
+    }
+
+    if (this.getDatabase() === snapshotName) {
+      throw new Error("Snapshot name cannot be the same as the database name");
+    }
+  }
 }