From 8305ef964dcdf072a0cf738901c2f5f5487063ae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lukas=20L=C3=B6sche?= Date: Wed, 17 Jul 2024 11:13:38 +0200 Subject: [PATCH] Add PostgreSQL support --- Dockerfile | 1 + fixbackup/__init__.py | 2 +- fixbackup/backup/__init__.py | 16 ++++- fixbackup/backup/postgresql.py | 109 +++++++++++++++++++++++++++++++++ pyproject.toml | 2 +- 5 files changed, 127 insertions(+), 3 deletions(-) create mode 100644 fixbackup/backup/postgresql.py diff --git a/Dockerfile b/Dockerfile index ff99036..477e95e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -68,6 +68,7 @@ RUN groupadd -g "${PGID:-0}" -o fix \ python3-pip \ redis-tools \ mysql-client \ + postgresql \ && dpkg -i /usr/local/tmp/arangodb3-client_*.deb \ && ln -s /usr/bin/busybox /usr/local/bin/vi \ && ln -s /usr/bin/busybox /usr/local/bin/wget \ diff --git a/fixbackup/__init__.py b/fixbackup/__init__.py index 16dbf2b..448516a 100644 --- a/fixbackup/__init__.py +++ b/fixbackup/__init__.py @@ -11,4 +11,4 @@ __author__ = "Some Engineering Inc." __license__ = "Apache 2.0" __copyright__ = "Copyright © 2023 Some Engineering Inc." -__version__ = "0.0.8" +__version__ = "0.0.9" diff --git a/fixbackup/backup/__init__.py b/fixbackup/backup/__init__.py index 427cd75..0be2ba5 100644 --- a/fixbackup/backup/__init__.py +++ b/fixbackup/backup/__init__.py @@ -4,10 +4,11 @@ from typing import List, Tuple from .redis import backup as redis_backup, add_args as redis_add_args from .mysql import backup as mysql_backup, add_args as mysql_add_args +from .postgresql import backup as pg_backup, add_args as postgresql_add_args from .arangodb import backup as arangodb_backup, add_args as arangodb_add_args from ..utils import valid_hostname, valid_ip, valid_dbname -add_args = [redis_add_args, mysql_add_args, arangodb_add_args] +add_args = [redis_add_args, mysql_add_args, arangodb_add_args, postgresql_add_args] def backup(args: Namespace, backup_directory: Path) -> Tuple[List[Path], bool]: @@ -39,6 +40,19 @@ def backup(args: Namespace, backup_directory: Path) -> Tuple[List[Path], bool]: else: all_success = False + if args.pg_host and (valid_hostname(args.pg_host) or valid_ip(args.pg_host)): + if args.pg_database: + db = str(args.pg_database) + if not valid_dbname(db): + raise ValueError(f"Invalid database name: {db}") + else: + db = "all" + pg_backup_file = backup_directory / f"{environment}-{date_prefix}-mysql-{args.pg_host}-{db}.sql.gz" + if pg_backup(args, pg_backup_file): + result.append(pg_backup_file) + else: + all_success = False + if args.arangodb_host and (valid_hostname(args.arangodb_host)): if args.arangodb_database: db = str(args.arangodb_database) diff --git a/fixbackup/backup/postgresql.py b/fixbackup/backup/postgresql.py new file mode 100644 index 0000000..879905a --- /dev/null +++ b/fixbackup/backup/postgresql.py @@ -0,0 +1,109 @@ +import os +import subprocess +from pathlib import Path +from argparse import ArgumentParser, Namespace +from ..utils import BackupFile +from ..logger import log + + +def add_args(arg_parser: ArgumentParser) -> None: + arg_parser.add_argument( + "--pg-host", + help="PostgreSQL host", + dest="pg_host", + type=str, + default=os.getenv("PG_HOST"), + ) + + arg_parser.add_argument( + "--pg-port", + help="PostgreSQL port", + dest="pg_port", + type=int, + default=os.getenv("PG_PORT", 5432), + ) + + arg_parser.add_argument( + "--pg-user", + help="PostgreSQL user", + dest="gp_user", + type=str, + default=os.getenv("PG_USER", "root"), + ) + + arg_parser.add_argument( + "--pg-password", + help="PostgreSQL password", + dest="pg_password", + type=str, + default=os.getenv("PG_PASSWORD"), + ) + + arg_parser.add_argument( + "--pg-database", + help="PostgreSQL database", + dest="pg_database", + type=str, + default=os.getenv("PG_DATABASE"), + ) + + arg_parser.add_argument( + "--pg-dump-args", + help="Extra arguments to pass to pg_dump", + dest="pg_dump_args", + action="append", + default=[], + ) + + +def backup(args: Namespace, backup_file_path: Path, timeout: int = 900, compress: bool = True) -> bool: + log.info("Starting PostgreSQL backup...") + + if not args.pg_host: + return False + + env = os.environ.copy() + command = [ + "pg_dump", + "-w", + "--if-exists", + "--inserts", + "-h", + str(args.pg_host), + "-p", + str(args.pg_port), + "-u", + str(args.pg_user), + *args.pg_dump_args, + ] + if args.pg_database: + command.append("-d") + command.append(args.pg_database) + else: + command[0] = "pg_dumpall" + + if args.pg_password: + env["PGPASSWORD"] = args.pg_password + + log.debug(f"Running command: {' '.join(command)}") + + try: + with BackupFile(backup_file_path, compress) as backup_fd: + process = subprocess.Popen(command, stdout=backup_fd, stderr=subprocess.PIPE, env=env) + _, stderr = process.communicate(timeout=timeout) + + if process.returncode == 0: + log.info(f"PostgreSQL backup completed successfully. Saved to {backup_file_path}") + if stderr: + log.debug(stderr.decode().strip()) + return True + else: + log.error(f"PostgreSQL backup failed with return code: {process.returncode}") + if stderr: + log.error(stderr.decode().strip()) + except subprocess.TimeoutExpired: + log.error(f"PostgreSQL backup failed with timeout after {timeout} seconds") + process.kill() + process.communicate() + + return False diff --git a/pyproject.toml b/pyproject.toml index fdeac9f..9010efe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "fixbackup" -version = "0.0.8" +version = "0.0.9" authors = [{name="Some Engineering Inc."}] description = "FIX Database Backup System" license = {file="LICENSE"}