Files
Django/import_backup.py

146 lines
4.4 KiB
Python

#!/usr/bin/env python3
"""只导入用户迁移所需的表到临时数据库 xaio_old"""
import MySQLdb
import re
from app_secrets import DATABASE_HOST, DATABASE_PORT, DATABASE_USER, DATABASE_PASSWORD
# 需要导入的表(用户相关)
NEEDED_TABLES = {
'user_main', 'user_dashou', 'user_guanshi', 'user_kefu',
'user_shangjia', 'user_boss', 'admin_profile', 'user_zuzhang',
'user_shenheguan', 'user_role', 'role', 'permission',
'role_permission', 'account_permission',
}
def get_conn(db=None):
kwargs = dict(
host=DATABASE_HOST, port=int(DATABASE_PORT),
user=DATABASE_USER, password=DATABASE_PASSWORD,
charset='utf8mb4', connect_timeout=60,
read_timeout=300, write_timeout=300,
)
if db:
kwargs['database'] = db
return MySQLdb.connect(**kwargs)
# Drop and recreate
conn = get_conn()
cursor = conn.cursor()
cursor.execute('DROP DATABASE IF EXISTS xaio_old')
cursor.execute('CREATE DATABASE xaio_old CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci')
conn.commit()
conn.close()
print('临时数据库 xaio_old 已重建')
# Read the SQL file
with open('xaio_cheng_xu_backup.sql', 'r', encoding='utf-8') as f:
sql_content = f.read()
# Parse and filter statements
# We need CREATE TABLE and INSERT statements for NEEDED_TABLES only
# Also need DROP TABLE IF EXISTS for those tables
statements = []
current = []
current_table = None
in_create = False
for line in sql_content.split('\n'):
stripped = line.strip()
if not stripped or stripped.startswith('--') or stripped.startswith('/*'):
continue
# Detect table name from CREATE TABLE or DROP TABLE or INSERT INTO
create_match = re.match(r'CREATE\s+TABLE\s+`(\w+)`', stripped, re.IGNORECASE)
drop_match = re.match(r'DROP\s+TABLE\s+IF\s+EXISTS\s+`(\w+)`', stripped, re.IGNORECASE)
insert_match = re.match(r'INSERT\s+INTO\s+`(\w+)`', stripped, re.IGNORECASE)
if create_match:
current_table = create_match.group(1)
in_create = True
elif drop_match:
current_table = drop_match.group(1)
elif insert_match:
current_table = insert_match.group(1)
in_create = False
# Only include statements for needed tables
if current_table and current_table in NEEDED_TABLES:
current.append(line)
if stripped.endswith(';'):
stmt = '\n'.join(current).strip()
current = []
current_table = None
in_create = False
upper = stmt.upper()
if not upper.startswith('SET ') and not upper.startswith('USE '):
statements.append(stmt)
elif stripped.endswith(';'):
# Not a needed table, skip
current = []
current_table = None
in_create = False
elif current_table and current_table in NEEDED_TABLES:
current.append(line)
else:
# Not needed, skip but track state
if stripped.endswith(';'):
current = []
current_table = None
print(f'筛选出 {len(statements)} 条用户相关SQL语句')
# Execute
conn = get_conn(db='xaio_old')
cursor = conn.cursor()
cursor.execute('SET FOREIGN_KEY_CHECKS = 0')
executed = 0
errors = 0
for stmt in statements:
try:
cursor.execute(stmt)
executed += 1
if executed % 20 == 0:
conn.commit()
except MySQLdb.OperationalError as e:
if e.args[0] in (2006, 2013):
try:
conn.close()
except:
pass
conn = get_conn(db='xaio_old')
cursor = conn.cursor()
try:
cursor.execute(stmt)
executed += 1
except Exception as e2:
errors += 1
if errors <= 5:
print(f' ERR(retry): {str(e2)[:100]}')
else:
errors += 1
if errors <= 5:
print(f' ERR: {str(e)[:100]}')
except Exception as e:
errors += 1
if errors <= 5:
print(f' ERR: {str(e)[:100]}')
conn.commit()
cursor.execute('SET FOREIGN_KEY_CHECKS = 1')
conn.commit()
print(f'\n导入完成: 成功 {executed}, 错误 {errors}')
# Check row counts
print('\n关键表行数:')
for table in sorted(NEEDED_TABLES):
try:
cursor.execute('SELECT COUNT(*) FROM `%s`' % table)
count = cursor.fetchone()[0]
print(f' {table}: {count} rows')
except Exception:
print(f' {table}: NOT FOUND')
conn.close()