Skip to content

Commit 17ae9b5

Browse files
Nikola-MilovicNikola-Milovic
authored
Add snapshot support to postgres module (#961)
1 parent 7f2aff2 commit 17ae9b5

File tree

3 files changed

+322
-0
lines changed

3 files changed

+322
-0
lines changed

‎docs/modules/postgresql.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,3 +25,18 @@ npm install @testcontainers/postgresql --save-dev
2525
<!--codeinclude-->
2626
[Set username:](../../packages/modules/postgresql/src/postgresql-container.test.ts) inside_block:setUsername
2727
<!--/codeinclude-->
28+
29+
### Using Snapshots
30+
31+
This example shows the usage of the postgres module's Snapshot feature to give each test a clean database without having
32+
to recreate the database container on every test or run heavy scripts to clean your database. This makes the individual
33+
tests very modular, since they always run on a brand-new database.
34+
35+
!!!tip
36+
You should never pass the `"postgres"` system database as the container database name if you want to use snapshots.
37+
The Snapshot logic requires dropping the connected database and using the system database to run commands, which will
38+
not work if the database for the container is set to `"postgres"`.
39+
40+
<!--codeinclude-->
41+
[Test with a reusable Postgres container](../../packages/modules/postgresql/src/postgresql-container.test.ts) inside_block:createAndRestoreFromSnapshot
42+
<!--/codeinclude-->

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

Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,3 +113,209 @@ describe("PostgreSqlContainer", { timeout: 180_000 }, () => {
113113
await expect(() => container.start()).rejects.toThrow();
114114
});
115115
});
116+
117+
describe("PostgreSqlContainer snapshot and restore", { timeout: 180_000 }, () => {
118+
// createAndRestoreFromSnapshot {
119+
it("should create and restore from snapshot", async () => {
120+
const container = await new PostgreSqlContainer().start();
121+
122+
// Connect to the database
123+
let client = new Client({
124+
connectionString: container.getConnectionUri(),
125+
});
126+
await client.connect();
127+
128+
// Create some test data
129+
await client.query("CREATE TABLE test_table (id SERIAL PRIMARY KEY, name TEXT)");
130+
await client.query("INSERT INTO test_table (name) VALUES ('initial data')");
131+
132+
// Close connection before snapshot (otherwise we'll get an error because user is already connected)
133+
await client.end();
134+
135+
// Take a snapshot
136+
await container.snapshot();
137+
138+
// Reconnect to database
139+
client = new Client({
140+
connectionString: container.getConnectionUri(),
141+
});
142+
await client.connect();
143+
144+
// Modify the database
145+
await client.query("INSERT INTO test_table (name) VALUES ('data after snapshot')");
146+
147+
// Verify both records exist
148+
let result = await client.query("SELECT * FROM test_table ORDER BY id");
149+
expect(result.rows).toHaveLength(2);
150+
expect(result.rows[0].name).toEqual("initial data");
151+
expect(result.rows[1].name).toEqual("data after snapshot");
152+
153+
// Close connection before restore (same reason as above)
154+
await client.end();
155+
156+
// Restore to the snapshot
157+
await container.restoreSnapshot();
158+
159+
// Reconnect to database
160+
client = new Client({
161+
connectionString: container.getConnectionUri(),
162+
});
163+
await client.connect();
164+
165+
// Verify only the initial data exists after restore
166+
result = await client.query("SELECT * FROM test_table ORDER BY id");
167+
expect(result.rows).toHaveLength(1);
168+
expect(result.rows[0].name).toEqual("initial data");
169+
170+
await client.end();
171+
await container.stop();
172+
});
173+
// }
174+
175+
it("should use custom snapshot name", async () => {
176+
const container = await new PostgreSqlContainer().start();
177+
const customSnapshotName = "my_custom_snapshot";
178+
179+
// Connect to the database
180+
let client = new Client({
181+
connectionString: container.getConnectionUri(),
182+
});
183+
await client.connect();
184+
185+
// Create a test table and insert data
186+
await client.query("CREATE TABLE test_table (id SERIAL PRIMARY KEY, name TEXT)");
187+
await client.query("INSERT INTO test_table (name) VALUES ('initial data')");
188+
189+
// Close connection before snapshot
190+
await client.end();
191+
192+
// Take a snapshot with custom name
193+
await container.snapshot(customSnapshotName);
194+
195+
// Reconnect to database
196+
client = new Client({
197+
connectionString: container.getConnectionUri(),
198+
});
199+
await client.connect();
200+
201+
// Modify the database
202+
await client.query("INSERT INTO test_table (name) VALUES ('data after snapshot')");
203+
204+
// Close connection before restore
205+
await client.end();
206+
207+
// Restore using the custom snapshot name
208+
await container.restoreSnapshot(customSnapshotName);
209+
210+
// Reconnect to database
211+
client = new Client({
212+
connectionString: container.getConnectionUri(),
213+
});
214+
await client.connect();
215+
216+
// Verify only the initial data exists after restore
217+
const result = await client.query("SELECT * FROM test_table ORDER BY id");
218+
expect(result.rows).toHaveLength(1);
219+
expect(result.rows[0].name).toEqual("initial data");
220+
221+
await client.end();
222+
await container.stop();
223+
});
224+
225+
it("should handle multiple snapshots", async () => {
226+
const container = await new PostgreSqlContainer().start();
227+
228+
// Connect to the database
229+
let client = new Client({
230+
connectionString: container.getConnectionUri(),
231+
});
232+
await client.connect();
233+
234+
// Create a test table
235+
await client.query("CREATE TABLE test_table (id SERIAL PRIMARY KEY, name TEXT)");
236+
237+
// Close connection before snapshot
238+
await client.end();
239+
240+
// Take first snapshot with empty table
241+
await container.snapshot("snapshot1");
242+
243+
// Reconnect to database
244+
client = new Client({
245+
connectionString: container.getConnectionUri(),
246+
});
247+
await client.connect();
248+
249+
// Add first record
250+
await client.query("INSERT INTO test_table (name) VALUES ('data for snapshot 2')");
251+
252+
// Close connection before snapshot
253+
await client.end();
254+
255+
// Take second snapshot with one record
256+
await container.snapshot("snapshot2");
257+
258+
// Reconnect to database
259+
client = new Client({
260+
connectionString: container.getConnectionUri(),
261+
});
262+
await client.connect();
263+
264+
// Add second record
265+
await client.query("INSERT INTO test_table (name) VALUES ('data after snapshots')");
266+
267+
// Verify we have two records
268+
let result = await client.query("SELECT COUNT(*) as count FROM test_table");
269+
expect(result.rows[0].count).toEqual("2");
270+
271+
// Close connection before restore
272+
await client.end();
273+
274+
// Restore to first snapshot (empty table)
275+
await container.restoreSnapshot("snapshot1");
276+
277+
// Reconnect to database
278+
client = new Client({
279+
connectionString: container.getConnectionUri(),
280+
});
281+
await client.connect();
282+
283+
// Verify table is empty
284+
result = await client.query("SELECT COUNT(*) as count FROM test_table");
285+
expect(result.rows[0].count).toEqual("0");
286+
287+
// Close connection before restore
288+
await client.end();
289+
290+
// Restore to second snapshot (one record)
291+
await container.restoreSnapshot("snapshot2");
292+
293+
// Reconnect to database
294+
client = new Client({
295+
connectionString: container.getConnectionUri(),
296+
});
297+
await client.connect();
298+
299+
// Verify we have one record
300+
result = await client.query("SELECT * FROM test_table");
301+
expect(result.rows).toHaveLength(1);
302+
expect(result.rows[0].name).toEqual("data for snapshot 2");
303+
304+
await client.end();
305+
await container.stop();
306+
});
307+
308+
it("should throw an error when trying to snapshot postgres system database", async () => {
309+
const container = await new PostgreSqlContainer().withDatabase("postgres").start();
310+
311+
await expect(container.snapshot()).rejects.toThrow(
312+
"Snapshot feature is not supported when using the postgres system database"
313+
);
314+
315+
await expect(container.restoreSnapshot()).rejects.toThrow(
316+
"Snapshot feature is not supported when using the postgres system database"
317+
);
318+
319+
await container.stop();
320+
});
321+
});

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

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ export class PostgreSqlContainer extends GenericContainer {
4848
}
4949

5050
export class StartedPostgreSqlContainer extends AbstractStartedContainer {
51+
private snapshotName: string = "migrated_template";
5152
constructor(
5253
startedTestContainer: StartedTestContainer,
5354
private readonly database: string,
@@ -85,4 +86,104 @@ export class StartedPostgreSqlContainer extends AbstractStartedContainer {
8586
url.password = this.getPassword();
8687
return url.toString();
8788
}
89+
90+
/**
91+
* Sets the name to be used for database snapshots.
92+
* This name will be used as the default for snapshot() and restore() methods.
93+
*
94+
* @param snapshotName The name to use for snapshots (default is "migrated_template" if this method is not called)
95+
* @returns this (for method chaining)
96+
*/
97+
public withSnapshotName(snapshotName: string): this {
98+
this.snapshotName = snapshotName;
99+
return this;
100+
}
101+
102+
/**
103+
* Takes a snapshot of the current state of the database as a template, which can then be restored using
104+
* the restore method.
105+
*
106+
* @param snapshotName Name for the snapshot, defaults to the value set by withSnapshotName() or "migrated_template" if not specified
107+
* @returns Promise resolving when snapshot is complete
108+
* @throws Error if attempting to snapshot the postgres system database or if using the same name as the database
109+
*/
110+
public async snapshot(snapshotName = this.snapshotName): Promise<void> {
111+
this.snapshotSanityCheck(snapshotName);
112+
113+
// Execute the commands to create the snapshot, in order
114+
await this.execCommandsSQL([
115+
// Update pg_database to remove the template flag, then drop the database if it exists.
116+
// This is needed because dropping a template database will fail.
117+
`UPDATE pg_database SET datistemplate = FALSE WHERE datname = '${snapshotName}'`,
118+
`DROP DATABASE IF EXISTS "${snapshotName}"`,
119+
// Create a copy of the database to another database to use as a template now that it was fully migrated
120+
`CREATE DATABASE "${snapshotName}" WITH TEMPLATE "${this.getDatabase()}" OWNER "${this.getUsername()}"`,
121+
// Snapshot the template database so we can restore it onto our original database going forward
122+
`ALTER DATABASE "${snapshotName}" WITH is_template = TRUE`,
123+
]);
124+
}
125+
126+
/**
127+
* Restores the database to a specific snapshot.
128+
*
129+
* @param snapshotName Name of the snapshot to restore from, defaults to the value set by withSnapshotName() or "migrated_template" if not specified
130+
* @returns Promise resolving when restore is complete
131+
* @throws Error if attempting to restore the postgres system database or if using the same name as the database
132+
*/
133+
public async restoreSnapshot(snapshotName = this.snapshotName): Promise<void> {
134+
this.snapshotSanityCheck(snapshotName);
135+
136+
// Execute the commands to restore the snapshot, in order
137+
await this.execCommandsSQL([
138+
// Drop the entire database by connecting to the postgres global database
139+
`DROP DATABASE "${this.getDatabase()}" WITH (FORCE)`,
140+
// Then restore the previous snapshot
141+
`CREATE DATABASE "${this.getDatabase()}" WITH TEMPLATE "${snapshotName}" OWNER "${this.getUsername()}"`,
142+
]);
143+
}
144+
145+
/**
146+
* Executes a series of SQL commands against the Postgres database
147+
*
148+
* @param commands Array of SQL commands to execute in sequence
149+
* @throws Error if any command fails to execute with details of the failure
150+
*/
151+
private async execCommandsSQL(commands: string[]): Promise<void> {
152+
for (const command of commands) {
153+
try {
154+
const result = await this.exec([
155+
"psql",
156+
"-v",
157+
"ON_ERROR_STOP=1",
158+
"-U",
159+
this.getUsername(),
160+
"-d",
161+
"postgres",
162+
"-c",
163+
command,
164+
]);
165+
166+
if (result.exitCode !== 0) {
167+
throw new Error(`Command failed with exit code ${result.exitCode}: ${result.output}`);
168+
}
169+
} catch (error) {
170+
console.error(`Failed to execute command: ${command}`, error);
171+
throw error;
172+
}
173+
}
174+
}
175+
176+
/**
177+
* Checks if the snapshot name is valid and if the database is not the postgres system database
178+
* @param snapshotName The name of the snapshot to check
179+
*/
180+
private snapshotSanityCheck(snapshotName: string): void {
181+
if (this.getDatabase() === "postgres") {
182+
throw new Error("Snapshot feature is not supported when using the postgres system database");
183+
}
184+
185+
if (this.getDatabase() === snapshotName) {
186+
throw new Error("Snapshot name cannot be the same as the database name");
187+
}
188+
}
88189
}

0 commit comments

Comments
 (0)