feat(testcontainers): add task manager db backup and restore

Signed-off-by: Fatih Acar <[email protected]>

Author: fatih-acar

python_testcontainers/infrahub_testcontainers/container.py

@@ -517,3 +517,111 @@ def database_restore_backup(self, backup_file: Path) -> None:  # noqa: PLR0915
                 )
             self.start()
             print("Database restored successfully")
+
+    def task_manager_create_backup(self, backup_name: str = "prefect.dump", dest_dir: Path | None = None) -> Path:
+        """Create a backup of the task manager PostgreSQL database using ``pg_dump``.
+
+        Args:
+            backup_name: Name of the archive file to create. Defaults to ``prefect.dump``.
+            dest_dir: Optional host directory where the backup should be copied after it is
+                produced. When omitted, the backup remains in ``external_backup_dir``.
+
+        Returns:
+            Path to the backup archive on the host filesystem.
+
+        Raises:
+            FileNotFoundError: If the pg_dump command completes but no archive is produced.
+        """
+
+        service_name = "task-manager-db"
+
+        try:
+            self.get_container(service_name=service_name)
+        except ContainerIsNotRunning:
+            self.start_container(service_name=service_name)
+
+        self.external_backup_dir.mkdir(parents=True, exist_ok=True)
+
+        internal_backup_path = self.internal_backup_dir / backup_name
+        dump_command = [
+            "pg_dump",
+            "--format=custom",
+            "--blobs",
+            "--no-owner",
+            "--no-privileges",
+            "--dbname=postgresql://postgres:postgres@localhost:5432/prefect",
+            f"--file={internal_backup_path}",
+        ]
+        self.exec_in_container(command=dump_command, service_name=service_name)
+
+        source_path = self.external_backup_dir / backup_name
+        if not source_path.exists():
+            raise FileNotFoundError(f"Backup file {source_path} was not created")
+
+        final_path = source_path
+        if dest_dir:
+            dest_dir.mkdir(parents=True, exist_ok=True)
+            if dest_dir.resolve() != self.external_backup_dir.resolve():
+                final_path = dest_dir / backup_name
+                shutil.copy(source_path, final_path)
+
+        return final_path
+
+    def task_manager_restore_backup(self, backup_file: Path) -> None:
+        """Restore the task manager PostgreSQL database from a ``pg_restore`` archive.
+
+        Args:
+            backup_file: Path to the backup archive on the host filesystem.
+
+        Raises:
+            FileNotFoundError: If the provided backup archive does not exist.
+        """
+
+        if not backup_file.exists():
+            raise FileNotFoundError(f"Backup file {backup_file} does not exist")
+
+        service_name = "task-manager-db"
+
+        try:
+            self.get_container(service_name=service_name)
+        except ContainerIsNotRunning:
+            self.start_container(service_name=service_name)
+
+        self.external_backup_dir.mkdir(parents=True, exist_ok=True)
+        target_path = self.external_backup_dir / backup_file.name
+        shutil.copy(backup_file, target_path)
+
+        admin_dsn = "postgresql://postgres:postgres@localhost:5432/postgres"
+        prefect_dsn = "postgresql://postgres:postgres@localhost:5432/prefect"
+        internal_backup_path = self.internal_backup_dir / backup_file.name
+
+        terminate_sessions_command = [
+            "psql",
+            f"--dbname={admin_dsn}",
+            "--command",
+            "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = 'prefect';",
+        ]
+        drop_database_command = [
+            "psql",
+            f"--dbname={admin_dsn}",
+            "--command",
+            "DROP DATABASE IF EXISTS prefect WITH (FORCE);",
+        ]
+        create_database_command = [
+            "psql",
+            f"--dbname={admin_dsn}",
+            "--command",
+            "CREATE DATABASE prefect OWNER postgres;",
+        ]
+        restore_command = [
+            "pg_restore",
+            "--no-owner",
+            "--role=postgres",
+            f"--dbname={prefect_dsn}",
+            str(internal_backup_path),
+        ]
+
+        self.exec_in_container(command=terminate_sessions_command, service_name=service_name)
+        self.exec_in_container(command=drop_database_command, service_name=service_name)
+        self.exec_in_container(command=create_database_command, service_name=service_name)
+        self.exec_in_container(command=restore_command, service_name=service_name)

python_testcontainers/infrahub_testcontainers/docker-compose-cluster.test.yml

@@ -248,6 +248,7 @@ services:
       - POSTGRES_DB=prefect
     volumes:
       - workflow_db:/var/lib/postgresql/data
+      - "./${INFRAHUB_TESTING_LOCAL_DB_BACKUP_DIRECTORY}:${INFRAHUB_TESTING_INTERNAL_DB_BACKUP_DIRECTORY}"
     healthcheck:
       test: ["CMD-SHELL", "pg_isready"]
       interval: 10s

python_testcontainers/infrahub_testcontainers/docker-compose.test.yml

@@ -159,6 +159,7 @@ services:
       - POSTGRES_DB=prefect
     volumes:
       - workflow_db:/var/lib/postgresql/data
+      - "./${INFRAHUB_TESTING_LOCAL_DB_BACKUP_DIRECTORY}:${INFRAHUB_TESTING_INTERNAL_DB_BACKUP_DIRECTORY}"
     healthcheck:
       test: ["CMD-SHELL", "pg_isready"]
       interval: 10s

python_testcontainers/tests/test_container_postgres_backup.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Callable
+
+import pytest
+from infrahub_testcontainers.container import InfrahubDockerCompose
+from testcontainers.core.exceptions import ContainerIsNotRunning
+
+
+@pytest.fixture(name="compose_factory")
+def fixture_compose_factory(tmp_path: Path) -> Callable[[], InfrahubDockerCompose]:
+    """Provide a factory that yields an isolated ``InfrahubDockerCompose`` test double."""
+
+    def _factory() -> InfrahubDockerCompose:
+        instance = InfrahubDockerCompose.__new__(InfrahubDockerCompose)
+        instance.context = tmp_path
+        instance.env_vars = {
+            "INFRAHUB_TESTING_LOCAL_DB_BACKUP_DIRECTORY": "backups",
+            "INFRAHUB_TESTING_INTERNAL_DB_BACKUP_DIRECTORY": "/backups",
+            "NEO4J_DOCKER_IMAGE": "neo4j:2025.03.0-enterprise",
+        }
+        instance.services = []
+        instance.pull = False
+        instance.wait = False
+        instance.compose_file_name = None
+        instance.docker_command_path = None
+        instance.project_name = None
+        instance.env_file = None
+
+        def _noop_service(service_name: str) -> None:
+            """Consume a service name without performing any action."""
+
+            _ = service_name
+
+        instance.get_container = _noop_service  # type: ignore[assignment]
+        instance.start_container = _noop_service  # type: ignore[assignment]
+        instance.exec_calls: list[list[str]] = []
+
+        def _exec(command: list[str], service_name: str) -> tuple[str, str, int]:
+            _ = service_name
+            instance.exec_calls.append(command)
+            if command and command[0] == "pg_dump":
+                file_arg = next(arg for arg in command if arg.startswith("--file="))
+                container_path = Path(file_arg.split("=", 1)[1])
+                host_path = instance.external_backup_dir / container_path.name
+                host_path.parent.mkdir(parents=True, exist_ok=True)
+                host_path.write_bytes(b"backup-bytes")
+            return "", "", 0
+
+        instance.exec_in_container = _exec  # type: ignore[assignment]
+        return instance
+
+    return _factory
+
+
+def test_task_manager_db_create_backup_creates_archive(compose_factory: Callable[[], InfrahubDockerCompose]) -> None:
+    compose = compose_factory()
+
+    result = compose.task_manager_db_create_backup()
+
+    assert result.name == "prefect.dump"
+    assert result.exists()
+    assert compose.exec_calls[0][0] == "pg_dump"
+
+
+def test_task_manager_db_create_backup_respects_destination(
+    compose_factory: Callable[[], InfrahubDockerCompose], tmp_path: Path
+) -> None:
+    compose = compose_factory()
+    destination = tmp_path / "exports"
+
+    result = compose.task_manager_db_create_backup(dest_dir=destination)
+
+    assert result.parent == destination
+    assert result.exists()
+    assert (compose.external_backup_dir / "prefect.dump").exists()
+
+
+def test_task_manager_db_restore_backup_runs_expected_commands(
+    compose_factory: Callable[[], InfrahubDockerCompose], tmp_path: Path
+) -> None:
+    compose = compose_factory()
+
+    def _raise(service_name: str) -> None:
+        _ = service_name
+        raise ContainerIsNotRunning("task-manager-db")
+
+    compose.get_container = _raise  # type: ignore[assignment]
+    started_services: list[str] = []
+
+    def _record_start(service_name: str) -> None:
+        started_services.append(service_name)
+
+    compose.start_container = _record_start  # type: ignore[assignment]
+    recorded_commands: list[list[str]] = []
+
+    def _exec(command: list[str], service_name: str) -> tuple[str, str, int]:
+        _ = service_name
+        recorded_commands.append(command)
+        return "", "", 0
+
+    compose.exec_in_container = _exec  # type: ignore[assignment]
+
+    backup_file = tmp_path / "prefect.dump"
+    backup_file.write_bytes(b"backup")
+
+    compose.task_manager_db_restore_backup(backup_file)
+
+    assert started_services == ["task-manager-db"]
+    assert (compose.external_backup_dir / backup_file.name).exists()
+    assert [command[0] for command in recorded_commands] == [
+        "psql",
+        "psql",
+        "psql",
+        "pg_restore",
+    ]