2195299055 / pg数据库

#!/usr/bin/env python3 """ PostgreSQL 压缩备份到 S3 一体化脚本 支持自动生成配置模板、环境变量覆盖、错误重试、进度显示 """ import os import subprocess import boto3 from botocore.exceptions import ClientError from datetime import datetime, timedelta import logging import gzip import shutil import yaml from dataclasses import dataclass from typing import List, NoReturn import argparse from tqdm import tqdm from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type from boto3.s3.transfer import TransferConfig # ------------------------- # 配置数据类与默认值 # ------------------------- @dataclass class BackupConfig: db_name: str = "your_database_name" db_user: str = "your_db_user" db_password: str = "your_db_password" s3_endpoint: str = "https://s3.your-provider.com" s3_access_key: str = "your_access_key" s3_secret_key: str = "your_secret_key" s3_bucket: str = "your-bucket-name" backup_dir: str = "/var/lib/pg_backup" compress_level: int = 6 keep_days: int = 7 log_file: str = "/var/log/pg_backup.log" pg_host: str = "localhost" pg_port: int = 5432 use_ssl: bool = False # ------------------------- # 配置文件管理 # ------------------------- def create_env_template(env_path: str) -> None: """生成.env环境变量模板文件""" template = """# PostgreSQL 备份到 S3 环境变量配置 # 将敏感信息放在这里，不要提交到版本控制 PG_DB_PASSWORD=your_db_password # 数据库密码 S3_ACCESS_KEY=your_access_key # S3 Access Key S3_SECRET_KEY=your_secret_key # S3 Secret Key """ with open(env_path, "w") as f: f.write(template) print(f"已生成环境变量模板到 {env_path}，请修改敏感信息") # 添加文件权限设置 if os.name != 'nt': # 非Windows系统 os.chmod(env_path, 0o600) # 仅用户可读写 def create_config_template(config_path: str) -> None: """生成配置模板文件""" default_config = BackupConfig() template = f"""# PostgreSQL 备份到 S3 配置文件 # 敏感信息建议通过环境变量设置（优先级高于文件配置） # 环境变量：PG_DB_PASSWORD, S3_ACCESS_KEY, S3_SECRET_KEY db_name: {default_config.db_name} # 数据库名称 db_user: {default_config.db_user} # 数据库用户 # db_password: {default_config.db_password} # 数据库密码（推荐通过环境变量设置） s3_endpoint: {default_config.s3_endpoint} # S3 端点（例如：https://s3.example.com） s3_access_key: {default_config.s3_access_key} # S3 Access Key # s3_secret_key: {default_config.s3_secret_key} # S3 Secret Key（推荐通过环境变量设置） s3_bucket: {default_config.s3_bucket} # S3 存储桶名称 backup_dir: {default_config.backup_dir} # 本地备份存储目录（需可写） keep_days: {default_config.keep_days} # 保留天数（删除超过天数的备份） pg_host: {default_config.pg_host} # 数据库主机（默认localhost） pg_port: {default_config.pg_port} # 数据库端口（默认5432） use_ssl: {default_config.use_ssl} # 是否启用SSL连接（默认false） log_file: {default_config.log_file} # 日志文件路径 compress_level: {default_config.compress_level} # 压缩级别（0-9，默认6） """ with open(config_path, "w") as f: f.write(template) print(f"已生成配置模板到 {config_path}，请修改后重新运行") # 同时生成.env文件模板 env_path = os.path.join(os.path.dirname(config_path), ".env") if not os.path.exists(env_path): create_env_template(env_path) def load_or_create_config(config_path: str) -> BackupConfig: """加载配置文件，不存在则生成模板""" if not os.path.exists(config_path): create_config_template(config_path) raise SystemExit(0) # 加载.env文件（如果存在） env_path = os.path.join(os.path.dirname(config_path), ".env") if os.path.exists(env_path): from dotenv import load_dotenv load_dotenv(env_path) with open(config_path, "r") as f: cfg = yaml.safe_load(f) # 环境变量覆盖敏感信息 env_override = { "db_password": os.getenv("PG_DB_PASSWORD"), "s3_access_key": os.getenv("S3_ACCESS_KEY"), "s3_secret_key": os.getenv("S3_SECRET_KEY") } for key, value in env_override.items(): if value: cfg[key] = value return BackupConfig(**cfg) # ------------------------- # 日志初始化 # ------------------------- def setup_logger(log_file: str) -> logging.Logger: logger = logging.getLogger("PGBackup") logger.setLevel(logging.DEBUG) formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') # 控制台和文件日志 console_handler = logging.StreamHandler() console_handler.setLevel(logging.INFO) console_handler.setFormatter(formatter) file_handler = logging.FileHandler(log_file) file_handler.setLevel(logging.DEBUG) file_handler.setFormatter(formatter) logger.addHandler(console_handler) logger.addHandler(file_handler) return logger # ------------------------- # 核心功能类 # ------------------------- class BackupManager: def __init__(self, config: BackupConfig): self.config = config self.logger = setup_logger(config.log_file) self.s3_client = self._init_s3_client() def _init_s3_client(self): return boto3.client( 's3', endpoint_url=self.config.s3_endpoint, aws_access_key_id=self.config.s3_access_key, aws_secret_access_key=self.config.s3_secret_key, region_name='cn-sy1', config=boto3.session.Config(signature_version='s3v4') ) def check_prerequisites(self) -> bool: try: os.makedirs(self.config.backup_dir, exist_ok=True) # 添加对备份目录可写性的检查 test_file = os.path.join(self.config.backup_dir, ".test") with open(test_file, "w") as f: f.write("test") os.remove(test_file) subprocess.run(["pg_dump", "--version"], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) return True except Exception as e: self.logger.error(f"前置检查失败: {str(e)}") return False def create_compressed_backup(self) -> str: timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") base_name = f"{self.config.db_name}_backup_{timestamp}" sql_path = os.path.join(self.config.backup_dir, base_name + ".sql") gz_path = sql_path + ".gz" try: # 检查关键参数是否为空 if not all([self.config.db_name, self.config.db_user, self.config.db_password, self.config.pg_host, str(self.config.pg_port)]): raise ValueError("数据库连接参数不完整") env = os.environ.copy() env["PGPASSWORD"] = self.config.db_password pg_cmd = [ "pg_dump", f"--dbname={self.config.db_name}", f"--host={self.config.pg_host}", f"--port={self.config.pg_port}", f"--username={self.config.db_user}", f"--file={sql_path}" ] if self.config.use_ssl: pg_cmd.append("--ssl-mode=require") # 添加详细的日志记录 self.logger.debug(f"执行pg_dump命令: {' '.join(pg_cmd)}") result = subprocess.run(pg_cmd, env=env, capture_output=True, text=True) if result.returncode != 0: error_msg = result.stderr.strip() self.logger.error(f"pg_dump详细错误: {error_msg}") raise RuntimeError(f"pg_dump失败: {error_msg}") with open(sql_path, "rb") as f_in, gzip.open(gz_path, "wb", self.config.compress_level) as f_out: shutil.copyfileobj(f_in, f_out) os.remove(sql_path) return gz_path except Exception as e: self._cleanup_temp_files([sql_path, gz_path]) self.logger.error(f"备份创建失败: {str(e)}") raise @retry(stop=stop_after_attempt(3), wait=wait_exponential(min=2, max=10), retry=retry_if_exception_type(ClientError)) def upload_to_s3(self, file_path: str) -> None: file_name = os.path.basename(file_path) transfer_config = TransferConfig(multipart_threshold=10*1024**2, max_concurrency=10) with open(file_path, "rb") as f, tqdm( total=os.path.getsize(file_path), unit="B", unit_scale=True, desc=f"上传 {file_name}", leave=False ) as pbar: def progress(bytes_transferred): pbar.update(bytes_transferred) self.s3_client.upload_fileobj(f, self.config.s3_bucket, file_name, Config=transfer_config, Callback=progress) def clean_old_backups(self) -> None: cutoff = datetime.now() - timedelta(days=self.config.keep_days) for file in os.listdir(self.config.backup_dir): if file.endswith(".gz"): path = os.path.join(self.config.backup_dir, file) mtime = datetime.fromtimestamp(os.path.getmtime(path)) if mtime < cutoff: try: os.remove(path) self.logger.info(f"删除过期备份: {file}") except Exception as e: self.logger.warning(f"删除失败: {file} - {str(e)}") @staticmethod def _cleanup_temp_files(files: List[str]) -> None: for f in files: if os.path.exists(f): try: os.remove(f) except: pass # ------------------------- # 命令行接口 # ------------------------- def parse_arguments() -> argparse.Namespace: parser = argparse.ArgumentParser(description="PostgreSQL 备份到 S3 工具") parser.add_argument("-c", "--config", default="pg_backup.yaml", help="配置文件路径（默认: pg_backup.yaml）") return parser.parse_args() # ------------------------- # 主流程 # ------------------------- def check_dependencies(): """检查所有依赖是否已安装""" try: import boto3 import yaml import tqdm import tenacity import dotenv except ImportError as e: print(f"[!] 缺少依赖: {str(e)}") print("请运行: pip install -r requirements.txt") raise SystemExit(1) # 在main()函数开头调用 def main(): check_dependencies() args = parse_arguments() config_path = args.config try: config = load_or_create_config(config_path) manager = BackupManager(config) print("[*] 开始备份流程") if not manager.check_prerequisites(): raise SystemExit("[!] 前置检查失败，流程终止") backup_path = manager.create_compressed_backup() manager.upload_to_s3(backup_path) manager.clean_old_backups() os.remove(backup_path) # 清理本地文件 print("\n[✓] 备份流程成功完成") except SystemExit: raise except Exception as e: manager.logger.error(f"备份失败: {str(e)}", exc_info=True) print(f"\n[!] 备份失败: {str(e)}") raise SystemExit(1) if __name__ == "__main__": main()

#!/usr/bin/env python3 """ PostgreSQL 恢复脚本（优化版） 此脚本用于恢复由 pg_backup_s3.py 备份到 S3 的数据库 支持从配置文件读取信息、错误重试、进度显示和详细日志记录 """ import os import boto3 import gzip import subprocess import shutil from datetime import datetime import logging import yaml from botocore.exceptions import ClientError from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type from tqdm import tqdm import argparse # 配置文件默认路径 DEFAULT_CONFIG_PATH = "pg_backup.yaml" # 日志设置 logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S', handlers=[ logging.StreamHandler(), logging.FileHandler('/var/log/pg_restore.log') ] ) logger = logging.getLogger('PG_Restore') def print_step(message): print(f"→ {message}") def load_config(config_path): """ 从 YAML 配置文件中加载配置信息 """ if not os.path.exists(config_path): logger.error(f"配置文件 {config_path} 不存在，请检查路径。") raise FileNotFoundError(f"配置文件 {config_path} 不存在。") with open(config_path, 'r') as f: config = yaml.safe_load(f) return config def get_s3_client(config): """ 创建 S3 客户端 """ return boto3.client( 's3', endpoint_url=config['s3_endpoint'], aws_access_key_id=config['s3_access_key'], aws_secret_access_key=config['s3_secret_key'], region_name='cn-sy1', config=boto3.session.Config(signature_version='s3v4') ) @retry(stop=stop_after_attempt(3), wait=wait_exponential(min=2, max=10), retry=retry_if_exception_type(ClientError)) def list_backup_files(config): """ 列出 S3 中的备份文件，并按时间倒序排列 """ try: s3 = get_s3_client(config) response = s3.list_objects_v2(Bucket=config['s3_bucket']) if 'Contents' not in response: print_step("S3 桶中没有找到备份文件") return [] files = [obj['Key'] for obj in response['Contents'] if obj['Key'].endswith('.gz')] files.sort(reverse=True) # 按时间倒序排列 if not files: print_step("没有找到 .gz 格式的备份文件") return [] return files except Exception as e: logger.error(f"获取备份列表失败: {str(e)}") raise class DownloadProgressPercentage: """ 下载进度显示 """ def __init__(self, filename, total_size): self._filename = filename self._size = total_size self._seen_so_far = 0 self._pbar = tqdm(total=total_size, unit='B', unit_scale=True, desc=f"下载 {filename}", leave=False) def __call__(self, bytes_amount): self._seen_so_far += bytes_amount self._pbar.update(bytes_amount) @retry(stop=stop_after_attempt(3), wait=wait_exponential(min=2, max=10), retry=retry_if_exception_type(ClientError)) def download_from_s3(config, file_name): """ 从 S3 下载备份文件 """ try: restore_dir = config.get('restore_dir', '/tmp/pg_restores') os.makedirs(restore_dir, exist_ok=True) local_path = os.path.join(restore_dir, file_name) s3 = get_s3_client(config) print_step(f"正在下载 {file_name}...") # 获取文件大小用于进度显示 file_size = s3.head_object(Bucket=config['s3_bucket'], Key=file_name)['ContentLength'] s3.download_file( Bucket=config['s3_bucket'], Key=file_name, Filename=local_path, Callback=DownloadProgressPercentage(file_name, file_size) ) print() # 换行 return local_path except Exception as e: logger.error(f"下载备份文件失败: {str(e)}") raise def decompress_file(compressed_path): """ 解压备份文件 """ try: print_step("正在解压备份文件...") decompressed_path = compressed_path[:-3] # 去掉 .gz 后缀 with gzip.open(compressed_path, 'rb') as f_in: with open(decompressed_path, 'wb') as f_out: shutil.copyfileobj(f_in, f_out) return decompressed_path except Exception as e: logger.error(f"解压备份文件失败: {str(e)}") raise def restore_database(config, sql_file): """ 执行数据库恢复 """ try: # 让用户选择恢复模式 print("\n请选择恢复模式:") print("1. 完全恢复 (先清空数据库，再恢复)") print("2. 追加恢复 (保留现有数据，只添加备份数据)") while True: try: mode = int(input("请输入选择(1或2): ")) if mode in [1, 2]: break print("输入无效，请输入 1 或 2") except ValueError: print("请输入有效的数字") env = os.environ.copy() env['PGPASSWORD'] = config['db_password'] # 完全恢复模式 if mode == 1: print_step("正在准备完全恢复...") temp_db = f"{config['db_name']}_temp" # 0. 先检查并删除已存在的临时数据库 print_step("正在清理可能存在的临时数据库...") drop_temp_cmd = [ 'sudo', '-u', 'postgres', 'psql', '-c', f"DROP DATABASE IF EXISTS {temp_db};" ] subprocess.run(drop_temp_cmd, check=True) # 1. 创建临时数据库 print_step("正在创建临时数据库...") create_temp_cmd = [ 'sudo', '-u', 'postgres', 'psql', '-c', f"CREATE DATABASE {temp_db} WITH OWNER {config['db_user']} ENCODING 'UTF8';" ] subprocess.run(create_temp_cmd, check=True) # 2. 将备份恢复到临时数据库 print_step("正在恢复数据到临时数据库...") restore_temp_cmd = [ 'psql', '-U', config['db_user'], '-h', 'localhost', '-d', temp_db, '-f', sql_file ] subprocess.run(restore_temp_cmd, env=env, check=True) # 3. 终止所有连接到原数据库的会话 print_step("正在终止原数据库连接...") terminate_cmd = [ 'sudo', '-u', 'postgres', 'psql', '-c', f"SELECT pg_terminate_backend(pg_stat_activity.pid) FROM pg_stat_activity WHERE pg_stat_activity.datname = '{config['db_name']}';" ] subprocess.run(terminate_cmd, check=True) # 4. 删除原数据库 print_step("正在清理原数据库...") drop_orig_cmd = [ 'sudo', '-u', 'postgres', 'psql', '-c', f"DROP DATABASE IF EXISTS {config['db_name']};" ] subprocess.run(drop_orig_cmd, check=True) # 5. 重命名临时数据库 print_step("正在完成恢复...") rename_cmd = [ 'sudo', '-u', 'postgres', 'psql', '-c', f"ALTER DATABASE {temp_db} RENAME TO {config['db_name']};" ] subprocess.run(rename_cmd, check=True) # 普通恢复操作 print_step("正在恢复数据库...") restore_cmd = [ 'psql', '-U', config['db_user'], '-h', 'localhost', '-d', config['db_name'], '-f', sql_file ] result = subprocess.run( restore_cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True ) if result.returncode != 0: raise Exception(f"恢复失败: {result.stderr.strip()}") print_step("数据库恢复成功") except Exception as e: logger.error(f"数据库恢复失败: {str(e)}") raise def cleanup(file_path): """ 清理临时文件 """ try: if os.path.exists(file_path): os.remove(file_path) except Exception as e: logger.warning(f"清理文件失败: {str(e)}") def main(): parser = argparse.ArgumentParser(description="PostgreSQL 恢复脚本") parser.add_argument("-c", "--config", default=DEFAULT_CONFIG_PATH, help="配置文件路径") args = parser.parse_args() print("\n" + "=" * 50) print("PostgreSQL 恢复脚本") print("=" * 50 + "\n") try: config = load_config(args.config) # 列出备份文件 backup_files = list_backup_files(config) if not backup_files: return # 显示备份文件列表 print("\n可用的备份文件:") for i, file in enumerate(backup_files, 1): print(f"{i}. {file}") # 选择要恢复的备份 while True: try: choice = int(input("\n请输入要恢复的备份编号: ")) if 1 <= choice <= len(backup_files): selected_file = backup_files[choice - 1] break print("输入无效，请重新输入") except ValueError: print("请输入有效的数字") # 下载并恢复 compressed_path = download_from_s3(config, selected_file) sql_path = decompress_file(compressed_path) restore_database(config, sql_path) except Exception as e: print_step(f"[错误] {str(e)}") finally: # 清理临时文件 if 'compressed_path' in locals(): cleanup(compressed_path) if 'sql_path' in locals(): cleanup(sql_path) print("\n[操作完成]") if __name__ == "__main__": main()

boto3>=1.26.0,<2.0.0 pyyaml>=6.0,<7.0 tenacity>=8.2.2,<9.0 tqdm>=4.65.0,<5.0 python-dotenv>=0.19.0,<1.0.0