124 lines
3.9 KiB
Python
124 lines
3.9 KiB
Python
|
import argparse
|
|||
|
|
import sys
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
from app.engines.reporter import Reporter
|
|||
|
|
|
|||
|
|
from .config import load_config, AppConfig
|
|||
|
|
from .models.base import connect_db, create_database
|
|||
|
|
|
|||
|
|
from loguru import logger
|
|||
|
|
import sqlalchemy.exc
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MainApp:
|
|||
|
|
"""主应用"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self.args = None
|
|||
|
|
self.config: AppConfig = None
|
|||
|
|
self.db_engine = None
|
|||
|
|
|
|||
|
|
def parse_args(self):
|
|||
|
|
"""解析命令行参数"""
|
|||
|
|
parser = argparse.ArgumentParser(description="Baidu Reporter")
|
|||
|
|
|
|||
|
|
# 指定配置文件
|
|||
|
|
parser.add_argument(
|
|||
|
|
"-c",
|
|||
|
|
"--config",
|
|||
|
|
default="./config.local.toml",
|
|||
|
|
help="指定配置文件路径，默认为 ./config.local.toml",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 添加输入文件参数
|
|||
|
|
parser.add_argument(
|
|||
|
|
"-f",
|
|||
|
|
"--file",
|
|||
|
|
default="./urls.txt",
|
|||
|
|
help="指定输入文件路径，默认为 ./urls.txt",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 添加运行模式参数
|
|||
|
|
parser.add_argument(
|
|||
|
|
"-m",
|
|||
|
|
"--mode",
|
|||
|
|
help="指定运行模式：pc/site/wap，不指定则运行所有模式，多个模式使用英文逗号分隔",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 添加 web 服务器参数
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--web", action="store_true", help="启动 web 服务器，启动后将忽略其他选项"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 如果没有传入任何参数，显示帮助信息
|
|||
|
|
if len(sys.argv) == 1:
|
|||
|
|
parser.print_help()
|
|||
|
|
sys.exit(0)
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
# 处理模式参数
|
|||
|
|
if args.mode:
|
|||
|
|
modes = [m.strip() for m in args.mode.split(",")]
|
|||
|
|
valid_modes = ["pc", "site", "wap"]
|
|||
|
|
invalid_modes = [m for m in modes if m not in valid_modes]
|
|||
|
|
if invalid_modes:
|
|||
|
|
parser.error(f'无效的运行模式: {", ".join(invalid_modes)}')
|
|||
|
|
args.mode = modes
|
|||
|
|
else:
|
|||
|
|
args.mode = ["pc", "site", "wap"]
|
|||
|
|
|
|||
|
|
# 检查输入的文件是否存在
|
|||
|
|
if not os.path.exists(args.file):
|
|||
|
|
parser.error(f"输入的文件不存在: {args.file}")
|
|||
|
|
|
|||
|
|
# 检查配置文件是否存在
|
|||
|
|
if not os.path.exists(args.config):
|
|||
|
|
parser.error(f"配置文件不存在: {args.config}")
|
|||
|
|
|
|||
|
|
self.args = args
|
|||
|
|
|
|||
|
|
def start_cli(self):
|
|||
|
|
"""开启 CLI 模式"""
|
|||
|
|
reporter = Reporter(self.args.file, self.args.mode, self.db_engine)
|
|||
|
|
reporter.run()
|
|||
|
|
|
|||
|
|
def start_web(self):
|
|||
|
|
"""开启 Web 模式"""
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
def run(self):
|
|||
|
|
"""运行应用"""
|
|||
|
|
|
|||
|
|
# 解析命令行参数
|
|||
|
|
self.parse_args()
|
|||
|
|
|
|||
|
|
# 加载配置文件
|
|||
|
|
self.config = load_config(self.args.config)
|
|||
|
|
logger.info(f"加载配置文件 {self.args.config} 成功")
|
|||
|
|
|
|||
|
|
# 连接数据库
|
|||
|
|
try:
|
|||
|
|
self.db_engine = connect_db(self.config)
|
|||
|
|
logger.info(f"连接数据库 {self.config.database.database} 成功")
|
|||
|
|
except sqlalchemy.exc.OperationalError as e:
|
|||
|
|
# 如果错误类型是数据库不存在，询问用户是否执行初始化操作
|
|||
|
|
if "1049" in str(e):
|
|||
|
|
logger.info("数据库不存在，尝试初始化数据库")
|
|||
|
|
create_database(self.config)
|
|||
|
|
logger.info("数据库初始化成功，尝试连接数据库")
|
|||
|
|
self.db_engine = connect_db(self.config)
|
|||
|
|
logger.info(f"连接数据库 {self.config.database.database} 成功")
|
|||
|
|
else:
|
|||
|
|
logger.error(f"连接数据库失败，请检查配置文件或数据库服务是否正常: {e}")
|
|||
|
|
sys.exit(1)
|
|||
|
|
|
|||
|
|
# 如果指定了 --web 参数，启动 web 服务器，忽略其他选项
|
|||
|
|
if self.args.web:
|
|||
|
|
logger.info("启动 Web 模式")
|
|||
|
|
return self.start_web()
|
|||
|
|
else:
|
|||
|
|
logger.info("启动 CLI 模式")
|
|||
|
|
return self.start_cli()
|