diff --git a/docs/modules/postgresql.md b/docs/modules/postgresql.md index ad9274588..ebd8ae14a 100644 --- a/docs/modules/postgresql.md +++ b/docs/modules/postgresql.md @@ -25,3 +25,18 @@ npm install @testcontainers/postgresql --save-dev [Set username:](../../packages/modules/postgresql/src/postgresql-container.test.ts) inside_block:setUsername + +### 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"`. + + +[Test with a reusable Postgres container](../../packages/modules/postgresql/src/postgresql-container.test.ts) inside_block:createAndRestoreFromSnapshot + \ No newline at end of file diff --git a/packages/modules/postgresql/src/postgresql-container.test.ts b/packages/modules/postgresql/src/postgresql-container.test.ts index 9dbffaf02..52ff5ede9 100644 --- a/packages/modules/postgresql/src/postgresql-container.test.ts +++ b/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(); + }); +}); diff --git a/packages/modules/postgresql/src/postgresql-container.ts b/packages/modules/postgresql/src/postgresql-container.ts index 3de192482..b841b672d 100755 --- a/packages/modules/postgresql/src/postgresql-container.ts +++ b/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 { + 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 { + 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 { + 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"); + } + } }