Files
banban/talkingq-url/scripts/restore_db.py
2026-03-24 15:04:36 +08:00

113 lines
3.2 KiB
Python

import os
import sys
import asyncio
import subprocess
import gzip
import tempfile
import shutil
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
from config import settings
from scripts.backup_db import list_backups, BACKUP_DIR
async def restore_database(backup_file):
"""从备份文件恢复数据库"""
file_path = backup_file
if not os.path.isabs(backup_file):
file_path = os.path.join(BACKUP_DIR, backup_file)
if not os.path.exists(file_path):
print(f"错误: 备份文件 '{file_path}' 不存在")
return False
confirm = input(f"警告: 这将覆盖当前数据库的所有数据。确定要继续吗? (y/N): ")
if confirm.lower() != 'y':
print("操作已取消")
return False
temp_file = None
sql_file = file_path
try:
if file_path.endswith('.gz'):
print("解压备份文件...")
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.sql')
temp_file.close()
with gzip.open(file_path, 'rb') as f_in:
with open(temp_file.name, 'wb') as f_out:
shutil.copyfileobj(f_in, f_out)
sql_file = temp_file.name
cmd = [
"mysql",
"-h", settings.db_host,
"-P", str(settings.db_port),
"-u", settings.db_user,
f"-p{settings.db_password}",
settings.db_name
]
print(f"开始恢复数据库...")
with open(sql_file, 'r') as f:
process = subprocess.Popen(
cmd,
stdin=f,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
)
stdout, stderr = process.communicate()
if process.returncode != 0:
print(f"恢复失败: {stderr.decode()}")
return False
print("数据库恢复成功!")
return True
except Exception as e:
print(f"恢复过程中出错: {e}")
return False
finally:
if temp_file and os.path.exists(temp_file.name):
os.unlink(temp_file.name)
async def interactive_restore():
"""交互式选择要恢复的备份文件"""
backups = await list_backups()
if not backups:
print("没有找到可用的备份文件")
return False
print("\n可用的备份文件:")
for i, backup in enumerate(backups):
file_path = os.path.join(BACKUP_DIR, backup)
size = os.path.getsize(file_path) / (1024 * 1024) # 转换为MB
print(f"{i+1}. {backup} ({size:.2f} MB)")
try:
choice = int(input("\n请选择要恢复的备份文件 (输入序号): "))
if 1 <= choice <= len(backups):
selected_backup = backups[choice-1]
return await restore_database(selected_backup)
else:
print("无效的选择")
return False
except ValueError:
print("无效的输入")
return False
async def main():
if len(sys.argv) > 1:
await restore_database(sys.argv[1])
else:
await interactive_restore()
if __name__ == "__main__":
asyncio.run(main())