准备创建mysqlv1分支的提交
This commit is contained in:
huangzhenpc
327
init_database.py
Normal file
327
init_database.py
Normal file
@@ -0,0 +1,327 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
数据库初始化脚本
|
||||
用于自动创建MySQL数据库、用户和表结构，并配置Redis
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import getpass
|
||||
import asyncio
|
||||
import yaml
|
||||
import pymysql
|
||||
import redis
|
||||
|
||||
from loguru import logger
|
||||
|
||||
|
||||
# 默认配置
|
||||
DEFAULT_CONFIG = {
|
||||
"database": {
|
||||
"host": "localhost",
|
||||
"port": 3306,
|
||||
"username": "auto_cursor_reg",
|
||||
"password": "auto_cursor_pass",
|
||||
"database": "auto_cursor_reg",
|
||||
"pool_size": 10,
|
||||
"use_redis": True
|
||||
},
|
||||
"redis": {
|
||||
"host": "127.0.0.1",
|
||||
"port": 6379,
|
||||
"password": "",
|
||||
"db": 0
|
||||
}
|
||||
}
|
||||
|
||||
# 数据库表结构
|
||||
EMAIL_ACCOUNTS_TABLE = '''
|
||||
CREATE TABLE IF NOT EXISTS email_accounts (
|
||||
id INT AUTO_INCREMENT PRIMARY KEY,
|
||||
email VARCHAR(255) UNIQUE NOT NULL,
|
||||
password VARCHAR(255) NOT NULL,
|
||||
client_id VARCHAR(255) NOT NULL,
|
||||
refresh_token TEXT NOT NULL,
|
||||
in_use BOOLEAN DEFAULT 0,
|
||||
cursor_password VARCHAR(255),
|
||||
cursor_cookie TEXT,
|
||||
cursor_token TEXT,
|
||||
sold BOOLEAN DEFAULT 0,
|
||||
status VARCHAR(20) DEFAULT 'pending',
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
|
||||
INDEX idx_status_inuse_sold (status, in_use, sold)
|
||||
)
|
||||
'''
|
||||
|
||||
|
||||
def setup_logger():
|
||||
"""设置日志"""
|
||||
logger.remove()
|
||||
logger.add(sys.stderr, level="INFO")
|
||||
logger.add("init_database.log", rotation="1 MB", level="DEBUG")
|
||||
|
||||
|
||||
def get_mysql_root_credentials():
|
||||
"""获取MySQL root用户凭据"""
|
||||
print("\n===== MySQL配置 =====")
|
||||
print("请输入MySQL root用户凭据以创建数据库和用户")
|
||||
|
||||
mysql_host = input("MySQL 主机地址 [localhost]: ") or "localhost"
|
||||
mysql_port = input("MySQL 端口 [3306]: ") or "3306"
|
||||
try:
|
||||
mysql_port = int(mysql_port)
|
||||
except ValueError:
|
||||
mysql_port = 3306
|
||||
print("端口无效，使用默认端口 3306")
|
||||
|
||||
mysql_root = input("MySQL root用户名 [root]: ") or "root"
|
||||
mysql_root_password = getpass.getpass("MySQL root密码: ")
|
||||
|
||||
return {
|
||||
"host": mysql_host,
|
||||
"port": mysql_port,
|
||||
"user": mysql_root,
|
||||
"password": mysql_root_password
|
||||
}
|
||||
|
||||
|
||||
def get_app_database_config():
|
||||
"""获取应用数据库配置"""
|
||||
print("\n请设置应用程序的数据库配置:")
|
||||
|
||||
db_name = input("数据库名称 [auto_cursor_reg]: ") or "auto_cursor_reg"
|
||||
db_user = input("数据库用户名 [auto_cursor_reg]: ") or "auto_cursor_reg"
|
||||
db_pass = getpass.getpass(f"为用户 {db_user} 设置密码 [auto_cursor_pass]: ")
|
||||
if not db_pass:
|
||||
db_pass = "auto_cursor_pass"
|
||||
|
||||
return {
|
||||
"database": db_name,
|
||||
"username": db_user,
|
||||
"password": db_pass
|
||||
}
|
||||
|
||||
|
||||
def get_redis_config():
|
||||
"""获取Redis配置"""
|
||||
use_redis = input("\n是否使用Redis缓存? (y/n) [y]: ").lower() != "n"
|
||||
|
||||
if not use_redis:
|
||||
return {"use_redis": False}
|
||||
|
||||
print("\n===== Redis配置 =====")
|
||||
redis_host = input("Redis 主机地址 [127.0.0.1]: ") or "127.0.0.1"
|
||||
redis_port = input("Redis 端口 [6379]: ") or "6379"
|
||||
try:
|
||||
redis_port = int(redis_port)
|
||||
except ValueError:
|
||||
redis_port = 6379
|
||||
print("端口无效，使用默认端口 6379")
|
||||
|
||||
redis_password = getpass.getpass("Redis 密码 (如无密码请直接回车): ")
|
||||
redis_db = input("Redis 数据库索引 [0]: ") or "0"
|
||||
try:
|
||||
redis_db = int(redis_db)
|
||||
except ValueError:
|
||||
redis_db = 0
|
||||
print("数据库索引无效，使用默认值 0")
|
||||
|
||||
return {
|
||||
"use_redis": True,
|
||||
"redis": {
|
||||
"host": redis_host,
|
||||
"port": redis_port,
|
||||
"password": redis_password,
|
||||
"db": redis_db
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_mysql_connection(config):
|
||||
"""测试MySQL连接"""
|
||||
try:
|
||||
conn = pymysql.connect(**config)
|
||||
conn.close()
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"MySQL连接失败: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def test_redis_connection(config):
|
||||
"""测试Redis连接"""
|
||||
try:
|
||||
r = redis.Redis(
|
||||
host=config["host"],
|
||||
port=config["port"],
|
||||
password=config["password"] if config["password"] else None,
|
||||
db=config["db"]
|
||||
)
|
||||
r.ping()
|
||||
r.close()
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Redis连接失败: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def setup_mysql(root_config, app_config):
|
||||
"""设置MySQL数据库和用户"""
|
||||
logger.info("开始设置MySQL数据库...")
|
||||
|
||||
try:
|
||||
# 连接到MySQL
|
||||
conn = pymysql.connect(**root_config)
|
||||
cursor = conn.cursor()
|
||||
|
||||
db_name = app_config["database"]
|
||||
db_user = app_config["username"]
|
||||
db_pass = app_config["password"]
|
||||
|
||||
# 创建数据库
|
||||
logger.info(f"创建数据库: {db_name}")
|
||||
cursor.execute(f"CREATE DATABASE IF NOT EXISTS `{db_name}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
|
||||
|
||||
# 创建用户并授权
|
||||
logger.info(f"创建用户: {db_user}")
|
||||
|
||||
# 检查用户是否已存在
|
||||
cursor.execute(f"SELECT User FROM mysql.user WHERE User = '{db_user}'")
|
||||
user_exists = cursor.fetchone()
|
||||
|
||||
if user_exists:
|
||||
logger.info(f"用户 {db_user} 已存在，更新密码")
|
||||
cursor.execute(f"ALTER USER '{db_user}'@'localhost' IDENTIFIED BY '{db_pass}'")
|
||||
cursor.execute(f"ALTER USER '{db_user}'@'%' IDENTIFIED BY '{db_pass}'")
|
||||
else:
|
||||
# 创建用户 (同时创建本地和远程连接权限)
|
||||
try:
|
||||
cursor.execute(f"CREATE USER '{db_user}'@'localhost' IDENTIFIED BY '{db_pass}'")
|
||||
cursor.execute(f"CREATE USER '{db_user}'@'%' IDENTIFIED BY '{db_pass}'")
|
||||
except pymysql.err.MySQLError as e:
|
||||
logger.warning(f"创建用户时出现警告 (可能用户已存在): {str(e)}")
|
||||
|
||||
# 授权
|
||||
cursor.execute(f"GRANT ALL PRIVILEGES ON `{db_name}`.* TO '{db_user}'@'localhost'")
|
||||
cursor.execute(f"GRANT ALL PRIVILEGES ON `{db_name}`.* TO '{db_user}'@'%'")
|
||||
cursor.execute("FLUSH PRIVILEGES")
|
||||
|
||||
# 切换到新创建的数据库
|
||||
cursor.execute(f"USE `{db_name}`")
|
||||
|
||||
# 创建表
|
||||
logger.info("创建数据表: email_accounts")
|
||||
cursor.execute(EMAIL_ACCOUNTS_TABLE)
|
||||
|
||||
conn.commit()
|
||||
cursor.close()
|
||||
conn.close()
|
||||
|
||||
logger.success("MySQL数据库设置成功")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"设置MySQL失败: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def update_config_file(db_config, redis_config=None):
|
||||
"""更新配置文件"""
|
||||
logger.info("更新配置文件...")
|
||||
|
||||
config_file = "config.yaml"
|
||||
|
||||
try:
|
||||
# 读取现有配置
|
||||
with open(config_file, "r", encoding="utf-8") as f:
|
||||
config = yaml.safe_load(f)
|
||||
|
||||
# 更新数据库配置
|
||||
config["database"].update({
|
||||
"host": db_config.get("host", "localhost"),
|
||||
"port": db_config.get("port", 3306),
|
||||
"username": db_config["username"],
|
||||
"password": db_config["password"],
|
||||
"database": db_config["database"],
|
||||
"use_redis": db_config.get("use_redis", False)
|
||||
})
|
||||
|
||||
# 如果启用Redis，更新Redis配置
|
||||
if redis_config and db_config.get("use_redis"):
|
||||
config["redis"] = redis_config
|
||||
|
||||
# 备份原配置文件
|
||||
if os.path.exists(config_file):
|
||||
os.rename(config_file, f"{config_file}.bak")
|
||||
logger.info(f"已备份原配置文件为 {config_file}.bak")
|
||||
|
||||
# 写入新配置
|
||||
with open(config_file, "w", encoding="utf-8") as f:
|
||||
yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
|
||||
|
||||
logger.success(f"配置文件已更新: {config_file}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"更新配置文件失败: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
setup_logger()
|
||||
logger.info("开始初始化数据库")
|
||||
|
||||
# 获取MySQL root凭据
|
||||
root_config = get_mysql_root_credentials()
|
||||
|
||||
# 测试MySQL连接
|
||||
if not test_mysql_connection(root_config):
|
||||
logger.error("无法连接到MySQL，请检查凭据和服务状态")
|
||||
return
|
||||
|
||||
# 获取应用数据库配置
|
||||
db_config = get_app_database_config()
|
||||
|
||||
# 获取Redis配置
|
||||
redis_info = get_redis_config()
|
||||
db_config["use_redis"] = redis_info["use_redis"]
|
||||
|
||||
# 如果启用了Redis，测试连接
|
||||
if redis_info["use_redis"]:
|
||||
redis_config = redis_info["redis"]
|
||||
if not test_redis_connection(redis_config):
|
||||
use_anyway = input("Redis连接测试失败，是否继续? (y/n) [n]: ").lower() == "y"
|
||||
if not use_anyway:
|
||||
logger.warning("用户取消了初始化")
|
||||
return
|
||||
else:
|
||||
redis_config = None
|
||||
|
||||
# 设置MySQL
|
||||
if not setup_mysql(root_config, db_config):
|
||||
return
|
||||
|
||||
# 合并配置
|
||||
final_db_config = {
|
||||
"host": root_config["host"],
|
||||
"port": root_config["port"],
|
||||
"username": db_config["username"],
|
||||
"password": db_config["password"],
|
||||
"database": db_config["database"],
|
||||
"use_redis": db_config["use_redis"]
|
||||
}
|
||||
|
||||
# 更新配置文件
|
||||
update_config_file(final_db_config, redis_config)
|
||||
|
||||
logger.success("数据库初始化完成!")
|
||||
print("\n===== 初始化完成 =====")
|
||||
print("您现在可以运行以下命令导入邮箱账号:")
|
||||
print(" python import_emails.py")
|
||||
print("\n然后运行主程序开始注册:")
|
||||
print(" python main.py")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user