Snapshot feature for postgres?

[Nikola-Milovic]

The go version comes with an amazing support for postgres and snapshotting which makes testing and having a clean environment super easy, is something like this possible in python?

[alexanderankin]

doesn't seem too difficult to do yourself, not sure if i think it needs to be part of tc necessarily but i wouldn't turn down a PR i guess https://github.com/testcontainers/testcontainers-go/blob/8b5211f88a6d4ce21164209ab004aa09f6bfde3f/modules/postgres/postgres.go#L253 https://www.postgresql.org/docs/current/manage-ag-templatedbs.html

…

On Thu, Jan 23, 2025 at 3:25 PM Nikola Milovic ***@***.***> wrote: The go version <https://golang.testcontainers.org/modules/postgres/#examples> comes with an amazing support for postgres and snapshotting which makes testing and having a clean environment super easy, is something like this possible in python? — Reply to this email directly, view it on GitHub <#760>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ACECGJDDI4Q5JEWT5PBV2AT2MFF3BAVCNFSM6AAAAABVYGI5NWVHI2DSMVQWIX3LMV43ERDJONRXK43TNFXW4OZXHA3DGNJVGA> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[Nikola-Milovic]

I managed to get something working, haven't tested it thoroughly, but seems good. I can contribute it once I get some free time

from testcontainers.postgres import PostgresContainer


class PostgresContainerHelper:
    def __init__(self, container: PostgresContainer):
        self.container = container
        self.snapshot_name: str | None = "migrated_template"
        self.db_name = container.dbname
        self.user = container.username

    async def snapshot(self, snapshot_name: str = "migrated_template") -> None:
        """Create a template database snapshot"""
        if snapshot_name:
            self.snapshot_name = snapshot_name

        commands = [
            # Allow dropping the template database if it exists
            f"UPDATE pg_database SET datistemplate = FALSE WHERE datname = '{snapshot_name}'",
            f"DROP DATABASE IF EXISTS {snapshot_name}",
            # Create new template from current database
            f"CREATE DATABASE {snapshot_name} WITH TEMPLATE {self.db_name} OWNER {self.user}",
            # Mark as template
            f"ALTER DATABASE {snapshot_name} WITH is_template = TRUE",
        ]

        print(f"Creating snapshot: {snapshot_name}")

        await self._execute_system_commands(commands)

        print(f"Created snapshot: {snapshot_name}")

    async def restore(self, snapshot_name: str | None = None) -> None:
        """Restore database from snapshot"""
        if not snapshot_name:
            snapshot_name = self.snapshot_name

        if not snapshot_name:
            raise ValueError("No snapshot name provided and no default exists")

        commands = [
            f"DROP DATABASE {self.db_name} WITH (FORCE)",
            f"CREATE DATABASE {self.db_name} WITH TEMPLATE {snapshot_name} OWNER {self.user}",
        ]

        print(f"Restoring snapshot: {snapshot_name}")

        await self._execute_system_commands(commands)

        print(f"Restored snapshot: {snapshot_name}")

    async def _execute_system_commands(self, commands: list[str]) -> None:
        for cmd in commands:
            (exit_code, output) = self.container.exec(
                f"psql -U {self.container.username} -d {self.container.dbname} -c '{cmd}'"
            )

            if exit_code != 0:
                print(
                    f"\nExit code: {exit_code} for Command: {cmd} \nOutput: {output.decode('utf-8')}"
                )