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')}"
 )