Feature Request: Add snapshot support in Postgres module

[Nikola-Milovic]

Hey!

The test containers for go have this awesome snapshot support which makes migrating the db and cleaning it up before every test a breeze, and it would help out in node environments as well.

I created my own wrapper for the time being, I can look into creating a PR.

export class PostgresContainerWrapper {
  private container: PostgreSqlContainer;
  private snapshotName = 'migrated_template';
  private startedContainer: StartedPostgreSqlContainer | null = null;

  constructor(image?: string) {
    this.container = new PostgreSqlContainer(image);
  }

  /**
   * Forwards method calls to the underlying container
   */
  public withDatabase(database: string): this {
    this.container.withDatabase(database);
    return this;
  }

  /**
   * Forwards method calls to the underlying container
   */
  public withUsername(username: string): this {
    this.container.withUsername(username);
    return this;
  }

  /**
   * Forwards method calls to the underlying container
   */
  public withPassword(password: string): this {
    this.container.withPassword(password);
    return this;
  }

  /**
   * Forwards method calls to the underlying container
   */
  public withExposedPorts(...ports: number[]): this {
    this.container.withExposedPorts(...ports);
    return this;
  }

  /**
   * Forwards method calls to the underlying container
   */
  public withEnvironment(environment: Record<string, string>): this {
    this.container.withEnvironment(environment);
    return this;
  }

  /**
   * Starts the container if not already started
   */
  public async start(): Promise<StartedPostgreSqlContainer> {
    if (!this.startedContainer) {
      this.startedContainer = await this.container.start();
    }
    return this.startedContainer;
  }

  /**
   * Stops the container if it's running
   */
  public async stop(): Promise<void> {
    if (this.startedContainer) {
      await this.startedContainer.stop();
      this.startedContainer = null;
    }
  }

  /**
   * Takes a snapshot of the current state of the database as a template, which can then be restored using
   * the restore method. By default, the snapshot will be created under a database called migrated_template.
   *
   * If a snapshot already exists under the given/default name, it will be overwritten with the new snapshot.
   *
   * @param opts Optional snapshot configuration options
   * @returns Promise resolving when snapshot is complete
   */
  public async snapshot(opts: SnapshotOption[] = []): Promise<void> {
    const startedContainer = await this.start();
    const snapshotName = this.checkSnapshotConfig(opts);

    // 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 "${startedContainer.getDatabase()}" OWNER "${startedContainer.getUsername()}"`,
      // Snapshot the template database so we can restore it onto our original database going forward
      `ALTER DATABASE "${snapshotName}" WITH is_template = TRUE`,
    ]);

    this.snapshotName = snapshotName;
  }

  /**
   * Restores the database to a specific snapshot. By default, it will restore the last snapshot taken on the
   * database by the snapshot method. If a snapshot name is provided, it will instead try to restore the snapshot by name.
   *
   * @param opts Optional snapshot configuration options
   * @returns Promise resolving when restore is complete
   */
  public async restore(opts: SnapshotOption[] = []): Promise<void> {
    const startedContainer = await this.start();
    const snapshotName = this.checkSnapshotConfig(opts);

    // 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 "${startedContainer.getDatabase()}" WITH (FORCE)`,
      // Then restore the previous snapshot
      `CREATE DATABASE "${startedContainer.getDatabase()}" WITH TEMPLATE "${snapshotName}" OWNER "${startedContainer.getUsername()}"`,
    ]);
  }

  /**
   * Checks the snapshot configuration and returns the snapshot name to use
   */
  private checkSnapshotConfig(opts: SnapshotOption[]): string {
    const config: SnapshotConfig = {};
    for (const opt of opts) {
      Object.assign(config, opt(config));
    }

    const snapshotName = config.snapshotName || this.snapshotName;

    // Make sure we have started container before checking database name
    if (!this.startedContainer) {
      throw new Error('Container must be started before taking or restoring snapshots');
    }

    if (this.startedContainer.getDatabase() === 'postgres') {
      throw new Error('Cannot restore the postgres system database as it cannot be dropped to be restored');
    }

    return snapshotName;
  }

  /**
   * Executes a series of SQL commands against the Postgres database
   * @param commands SQL commands to execute
   */
  private async execCommandsSQL(commands: string[]): Promise<void> {
    if (!this.startedContainer) {
      throw new Error('Container must be started before executing SQL commands');
    }

    try {
      // Connect to the postgres database (not the app database)
      const connectionString = this.startedContainer.getConnectionUri().replace(
        `/${this.startedContainer.getDatabase()}`,
        '/postgres'
      );

      // Use the pg client from the container's dependencies
      const { Client } = await import('pg');
      const client = new Client({
        connectionString,
        ssl: false,
      });

      await client.connect();

      try {
        // Execute each command
        for (const command of commands) {
          await client.query(command);
        }
      } finally {
        await client.end();
      }
    } catch (error) {
      // Fall back to using the docker exec approach if direct connection fails
      console.log('Direct database connection failed, falling back to docker exec:', error);
      await this.execCommandsFallback(commands);
    }
  }

  /**
   * Fallback method to execute SQL commands using docker exec if direct connection fails
   * @param commands SQL commands to execute
   */
  private async execCommandsFallback(commands: string[]): Promise<void> {
    if (!this.startedContainer) {
      throw new Error('Container must be started before executing fallback commands');
    }

    for (const command of commands) {
      try {
        const result = await this.startedContainer.exec([
          'psql',
          '-v',
          'ON_ERROR_STOP=1',
          '-U',
          this.startedContainer.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;
      }
    }
  }
}

[cristianrgreco]

Hi, thanks for sharing. What's the benefit in this approach over using container.commit, which works for all use cases?

[Nikola-Milovic]

Hey, honestly not familiar with that feature, I can only explain the workflow that I was looking for

pseudo-ish code

beforeAll(() => {
    container = await setupPostgresContainer()
    const dbClient = new Pool({conn: container.getConnectionURL})
    await migrate(dbClient)
   // await seedDB(dbClient)
    await container.snapshot("migrated_db")
})

beforeEach(() => {
   await container.restore("migrated_db")
})

This way I only have to wait for the PG container to startup once, and then having a fresh DB before each test run takes miliseconds.

How would I achieve this with the container.commit @cristianrgreco ?

[cristianrgreco]

Hi @Nikola-Milovic, I see, this snapshot feature lets you reset the DB without having to create a new container.

I assume however that this does not allow you to run tests in parallel, and does not help if applying the schema is slow?

[Nikola-Milovic]

I see, this snapshot feature lets you reset the DB without having to create a new container.

@cristianrgreco Yeah it's quite simple.

I assume however that this does not allow you to run tests in parallel, and does not help if applying the schema is slow?

It does not allow parallel tests, no, but I don't think the commit or any method for that matter does. You still need separate instances for concurrent access if you don't want race conditions. You can restore the snapshot into another database so in theory this method scales the best, since instead of spinning up X instances, you can have a single instance and X databases

Yeah applying the schema is a bottleneck but it happens once on startup and the way I develop apps I usually have schema defined in code so its rather quick to apply. In general I think this is the smoothest experience for me, if anything, combining snapshot feature with a commited, pre-migrated container could be an amazing workflow. Since these two features don't really collide as much as I initially thought, but rather would work very well together

[cristianrgreco]

Thanks for the info! I agree I think it would be a nice feature to have. PR is welcome if you have the time.

[Nikola-Milovic]

Sweet. I think I'll take it and create a PR today or tomorrow.