backup-docker-to-local/src/baudolo/backup/db.py

from __future__ import annotations

import logging
import pathlib
import re
from typing import TYPE_CHECKING

from baudolo.databases import CLUSTER_ROW, validate_database
from baudolo.generation import CLUSTER_SUFFIX, DUMP_SUFFIX, SQL_DIR

from .docker import docker_exec_argv
from .shell import BackupError, execute_to_file

if TYPE_CHECKING:
    import pandas as pd

log = logging.getLogger(__name__)


def get_instance(container: str, database_containers: list[str]) -> str | None:
    """The databases.csv instance a container serves, or None for no database.

    A declared container is its own instance. Every other name is normalised by
    stripping a database suffix token, which maps both `<app>-database` from
    compose and `<app>_database.1.<task>` from swarm onto the same instance.

    Args:
        container: the running container's name.
        database_containers: names passed via --database-containers, taken as
            declared engines whatever they are called.

    Returns:
        The instance name, or None when the name carries no database token: an
        application container is not an engine, even when it ships the client
        tools that would let a dump command start.
    """
    if container in database_containers:
        return container
    parts = re.split(r"(_|-)(database|db|postgres)", container)
    return parts[0] if len(parts) > 1 else None


def fallback_pg_dumpall(
    container: str, username: str, password: str, out_file: str
) -> None:
    """
    Perform a full Postgres cluster dump using pg_dumpall.
    """
    execute_to_file(
        docker_exec_argv(
            container,
            ["pg_dumpall", "-U", username, "-h", "localhost"],
            interactive=True,
            forward_env=["PGPASSWORD"],
        ),
        out_file,
        env={"PGPASSWORD": password},
    )


def backup_database(
    *,
    container: str,
    volume_dir: str,
    db_type: str,
    dump_tool: str,
    databases_df: pd.DataFrame,
    database_containers: list[str],
) -> bool:
    """
    Backup databases for a given DB container.

    Args:
        dump_tool: the MariaDB client found in the container, so an image
            that ships only mysqldump is dumped with the tool it has.

    Returns True if at least one dump was produced.
    """
    instance_name = get_instance(container, database_containers)
    if instance_name is None:
        log.debug("Container '%s' carries no database token", container)
        return False

    entries = databases_df[databases_df["instance"] == instance_name]
    if entries.empty:
        log.debug("No database entries for instance '%s'", instance_name)
        return False

    out_dir = pathlib.Path(volume_dir) / SQL_DIR
    out_dir.mkdir(parents=True, exist_ok=True)

    produced = False

    for row in entries.itertuples(index=False):
        raw_db = getattr(row, "database", "")
        user = (getattr(row, "username", "") or "").strip()
        password = (getattr(row, "password", "") or "").strip()

        db_value = validate_database(raw_db, instance=instance_name)

        if db_value == CLUSTER_ROW:
            if db_type != "postgres":
                raise ValueError(
                    f"databases.csv entry for instance '{instance_name}': "
                    f"'{CLUSTER_ROW}' is currently only supported for Postgres."
                )

            cluster_file = str(out_dir / f"{instance_name}{CLUSTER_SUFFIX}")
            fallback_pg_dumpall(container, user, password, cluster_file)
            produced = True
            continue

        db_name = db_value
        dump_file = str(out_dir / f"{db_name}{DUMP_SUFFIX}")

        if db_type == "mariadb":
            # Force TCP so auth matches '<user>'@'%' instead of socket -> 'localhost'.
            execute_to_file(
                docker_exec_argv(
                    container,
                    [
                        dump_tool,
                        "-h",
                        "127.0.0.1",
                        "--protocol=tcp",
                        "-u",
                        user,
                        f"-p{password}",
                        db_name,
                    ],
                ),
                dump_file,
            )
            produced = True
            continue

        if db_type == "postgres":
            try:
                execute_to_file(
                    docker_exec_argv(
                        container,
                        [
                            "pg_dump",
                            "-U",
                            user,
                            "-d",
                            db_name,
                            "-h",
                            "localhost",
                            "--no-owner",
                            "--no-privileges",
                        ],
                        interactive=True,
                        forward_env=["PGPASSWORD"],
                    ),
                    dump_file,
                    env={"PGPASSWORD": password},
                )
                produced = True
            except BackupError as e:
                raise BackupError(
                    f"Postgres dump failed for instance '{instance_name}', "
                    f"database '{db_name}'. This database was explicitly configured "
                    "and therefore must succeed.\n"
                    f"{e}"
                ) from e
            continue

    return produced