""" 数据库恢复脚本 从SQL备份文件恢复数据库 """ import os import sys import subprocess import pymysql from pathlib import Path from dotenv import load_dotenv import gzip # 加载环境变量 load_dotenv() class DatabaseRestore: """数据库恢复类""" def __init__(self): """初始化数据库配置""" self.db_config = { 'host': os.getenv('DB_HOST', '152.136.177.240'), 'port': int(os.getenv('DB_PORT', 5012)), 'user': os.getenv('DB_USER', 'finyx'), 'password': os.getenv('DB_PASSWORD', '6QsGK6MpePZDE57Z'), 'database': os.getenv('DB_NAME', 'finyx'), 'charset': 'utf8mb4' } def restore_with_mysql(self, backup_file, drop_database=False): """ 使用mysql命令恢复数据库（推荐方式） Args: backup_file: 备份文件路径 drop_database: 是否先删除数据库（危险操作） Returns: 是否成功 """ backup_file = Path(backup_file) if not backup_file.exists(): raise FileNotFoundError(f"备份文件不存在: {backup_file}") # 如果是压缩文件，先解压 sql_file = backup_file temp_file = None if backup_file.suffix == '.gz': print(f"检测到压缩文件，正在解压...") temp_file = backup_file.with_suffix('') with gzip.open(backup_file, 'rb') as f_in: with open(temp_file, 'wb') as f_out: f_out.write(f_in.read()) sql_file = temp_file print(f"解压完成: {sql_file}") try: print(f"开始恢复数据库 {self.db_config['database']}...") print(f"备份文件: {backup_file}") # 如果指定删除数据库 if drop_database: print("警告: 将删除现有数据库！") confirm = input("确认继续? (yes/no): ") if confirm.lower() != 'yes': print("已取消恢复操作") return False # 删除数据库 self._drop_database() # 构建mysql命令 cmd = [ 'mysql', f"--host={self.db_config['host']}", f"--port={self.db_config['port']}", f"--user={self.db_config['user']}", f"--password={self.db_config['password']}", '--default-character-set=utf8mb4', self.db_config['database'] ] # 执行恢复命令 with open(sql_file, 'r', encoding='utf-8') as f: result = subprocess.run( cmd, stdin=f, stderr=subprocess.PIPE, text=True ) if result.returncode != 0: error_msg = result.stderr.decode('utf-8') if result.stderr else '未知错误' raise Exception(f"mysql执行失败: {error_msg}") print("恢复完成！") return True except FileNotFoundError: print("错误: 未找到mysql命令，请确保MySQL客户端已安装并在PATH中") print("尝试使用Python方式恢复...") return self.restore_with_python(backup_file, drop_database) except Exception as e: print(f"恢复失败: {str(e)}") raise finally: # 清理临时解压文件 if temp_file and temp_file.exists(): temp_file.unlink() def restore_with_python(self, backup_file, drop_database=False): """ 使用Python直接连接数据库恢复（备用方式） Args: backup_file: 备份文件路径 drop_database: 是否先删除数据库（危险操作） Returns: 是否成功 """ backup_file = Path(backup_file) if not backup_file.exists(): raise FileNotFoundError(f"备份文件不存在: {backup_file}") # 如果是压缩文件，先解压 sql_file = backup_file temp_file = None if backup_file.suffix == '.gz': print(f"检测到压缩文件，正在解压...") temp_file = backup_file.with_suffix('') with gzip.open(backup_file, 'rb') as f_in: with open(temp_file, 'wb') as f_out: f_out.write(f_in.read()) sql_file = temp_file print(f"解压完成: {sql_file}") try: print(f"开始使用Python方式恢复数据库 {self.db_config['database']}...") print(f"备份文件: {backup_file}") # 如果指定删除数据库 if drop_database: print("警告: 将删除现有数据库！") confirm = input("确认继续? (yes/no): ") if confirm.lower() != 'yes': print("已取消恢复操作") return False # 删除数据库 self._drop_database() # 连接数据库 connection = pymysql.connect(**self.db_config) cursor = connection.cursor() # 读取SQL文件 print("读取SQL文件...") with open(sql_file, 'r', encoding='utf-8') as f: sql_content = f.read() # 分割SQL语句（按分号分割，但要注意字符串中的分号） print("执行SQL语句...") statements = self._split_sql_statements(sql_content) total = len(statements) print(f"共 {total} 条SQL语句") # 执行每条SQL语句 for i, statement in enumerate(statements, 1): statement = statement.strip() if not statement or statement.startswith('--'): continue try: cursor.execute(statement) if i % 100 == 0: print(f"进度: {i}/{total} ({i*100//total}%)") except Exception as e: # 某些错误可以忽略（如表已存在等） error_msg = str(e).lower() if 'already exists' in error_msg or 'duplicate' in error_msg: continue print(f"警告: 执行SQL语句时出错 (第{i}条): {str(e)}") print(f"SQL: {statement[:100]}...") # 提交事务 connection.commit() cursor.close() connection.close() print("恢复完成！") return True except Exception as e: print(f"恢复失败: {str(e)}") raise finally: # 清理临时解压文件 if temp_file and temp_file.exists(): temp_file.unlink() def _split_sql_statements(self, sql_content): """ 分割SQL语句（处理字符串中的分号） Args: sql_content: SQL内容 Returns: SQL语句列表 """ statements = [] current_statement = [] in_string = False string_char = None i = 0 while i < len(sql_content): char = sql_content[i] # 检测字符串开始/结束 if char in ("'", '"', '`') and (i == 0 or sql_content[i-1] != '\\'): if not in_string: in_string = True string_char = char elif char == string_char: in_string = False string_char = None current_statement.append(char) # 如果不在字符串中且遇到分号，分割语句 if not in_string and char == ';': statement = ''.join(current_statement).strip() if statement: statements.append(statement) current_statement = [] i += 1 # 添加最后一条语句 if current_statement: statement = ''.join(current_statement).strip() if statement: statements.append(statement) return statements def _drop_database(self): """删除数据库（危险操作）""" try: # 连接到MySQL服务器（不指定数据库） config = self.db_config.copy() config.pop('database') connection = pymysql.connect(**config) cursor = connection.cursor() cursor.execute(f"DROP DATABASE IF EXISTS `{self.db_config['database']}`") cursor.execute(f"CREATE DATABASE `{self.db_config['database']}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci") connection.commit() cursor.close() connection.close() print(f"数据库 {self.db_config['database']} 已删除并重新创建") except Exception as e: raise Exception(f"删除数据库失败: {str(e)}") def test_connection(self): """测试数据库连接""" try: connection = pymysql.connect(**self.db_config) cursor = connection.cursor() cursor.execute("SELECT VERSION()") version = cursor.fetchone()[0] cursor.close() connection.close() print(f"数据库连接成功！MySQL版本: {version}") return True except Exception as e: print(f"数据库连接失败: {str(e)}") return False def main(): """主函数""" import argparse parser = argparse.ArgumentParser(description='数据库恢复工具') parser.add_argument('backup_file', help='备份文件路径') parser.add_argument('--method', choices=['mysql', 'python', 'auto'], default='auto', help='恢复方法 (默认: auto)') parser.add_argument('--drop-db', action='store_true', help='恢复前删除现有数据库（危险操作）') parser.add_argument('--test', action='store_true', help='仅测试数据库连接') args = parser.parse_args() restore = DatabaseRestore() # 测试连接 if args.test: restore.test_connection() return # 执行恢复 try: if args.method == 'mysql': success = restore.restore_with_mysql(args.backup_file, args.drop_db) elif args.method == 'python': success = restore.restore_with_python(args.backup_file, args.drop_db) else: # auto try: success = restore.restore_with_mysql(args.backup_file, args.drop_db) except: print("\nmysql方式失败，切换到Python方式...") success = restore.restore_with_python(args.backup_file, args.drop_db) if success: print("\n恢复成功！") else: print("\n恢复失败！") sys.exit(1) except Exception as e: print(f"\n恢复失败: {str(e)}") sys.exit(1) if __name__ == '__main__': main()