205 lines
8.3 KiB
Python
205 lines
8.3 KiB
Python
"""
|
||
Query关键词导入模块
|
||
从Excel文件读取关键词，批量导入到baidu_keyword表
|
||
"""
|
||
import os
|
||
import glob
|
||
import time
|
||
import pandas as pd
|
||
from datetime import datetime
|
||
from loguru import logger
|
||
from config import Config
|
||
from db_manager import QueryKeywordManager, QueryImportLogManager
|
||
|
||
|
||
class QueryKeywordImporter:
|
||
"""Query关键词导入器"""
|
||
|
||
# 支持的query列名
|
||
QUERY_COLUMNS = ['query', 'Query', '查询词', 'keyword', '关键词']
|
||
# 支持的科室列名
|
||
DEPT_COLUMNS = ['科室', 'department', 'Department', '部门']
|
||
# 批量插入大小
|
||
BATCH_SIZE = 500
|
||
|
||
def __init__(self):
|
||
self.keyword_mgr = QueryKeywordManager()
|
||
self.log_mgr = QueryImportLogManager()
|
||
|
||
def _detect_column(self, df_columns, candidates):
|
||
"""从DataFrame列中检测匹配的列名"""
|
||
for col in candidates:
|
||
if col in df_columns:
|
||
return col
|
||
return None
|
||
|
||
def import_file(self, filepath: str, log_id: int = None, import_mode: str = 'query_only') -> dict:
|
||
"""
|
||
导入单个Excel文件到baidu_keyword表（批量模式）
|
||
|
||
Args:
|
||
filepath: Excel文件路径
|
||
log_id: 导入日志ID（可选，如果已创建）
|
||
import_mode: 导入模式
|
||
- 'query_only': 仅导入query列，query_status=draft
|
||
- 'full_import': 三列（科室/关键字/query），query_status=manual_review
|
||
|
||
Returns:
|
||
{'success': bool, 'stats': {...}, 'error': str}
|
||
"""
|
||
filename = os.path.basename(filepath)
|
||
stats = {'total': 0, 'success': 0, 'skip': 0, 'fail': 0}
|
||
|
||
# 根据模式确定 query_status
|
||
query_status = 'draft' if import_mode == 'query_only' else 'manual_review'
|
||
|
||
# 创建或获取日志记录
|
||
if log_id is None:
|
||
log_id = self.log_mgr.create_log(filename, filepath)
|
||
|
||
try:
|
||
# 更新状态为运行中
|
||
if log_id:
|
||
self.log_mgr.update_status(log_id, 'running')
|
||
|
||
logger.info(f"[Query导入] 开始导入文件: {filename}, 模式: {import_mode}, query_status: {query_status}")
|
||
|
||
# 1. 读取Excel
|
||
if not os.path.exists(filepath):
|
||
raise FileNotFoundError(f"文件不存在: {filepath}")
|
||
|
||
df = pd.read_excel(filepath)
|
||
logger.info(f"[Query导入] 文件 {filename} 包含 {len(df)} 行, 列名: {df.columns.tolist()}")
|
||
|
||
# 2. 根据模式确定列映射
|
||
if import_mode == 'query_only':
|
||
# 模式1：仅一列query
|
||
query_col = df.columns[0]
|
||
dept_col = None
|
||
keyword_col = None
|
||
logger.info(f"[Query导入] 仅导入Query模式，使用列: {query_col}")
|
||
else:
|
||
# 模式2：三列 - 科室/关键字/query（按位置）
|
||
dept_col = df.columns[0]
|
||
keyword_col = df.columns[1]
|
||
query_col = df.columns[2]
|
||
logger.info(f"[Query导入] 完整导入模式，科室列: {dept_col}, 关键字列: {keyword_col}, query列: {query_col}")
|
||
|
||
# 3. 预处理数据：提取所有有效关键词
|
||
keyword_list = []
|
||
for idx, row in df.iterrows():
|
||
query_val = str(row[query_col]).strip() if pd.notna(row[query_col]) else ''
|
||
if not query_val or query_val == 'nan':
|
||
stats['fail'] += 1
|
||
continue
|
||
|
||
item = {'keyword': query_val, 'department': ''}
|
||
|
||
if import_mode == 'full_import':
|
||
# 提取科室
|
||
if dept_col and pd.notna(row[dept_col]):
|
||
dept_val = str(row[dept_col]).strip()
|
||
if dept_val != 'nan':
|
||
item['department'] = dept_val
|
||
|
||
keyword_list.append(item)
|
||
|
||
stats['total'] = len(df)
|
||
total_valid = len(keyword_list)
|
||
logger.info(f"[Query导入] 有效关键词: {total_valid} / {stats['total']}")
|
||
|
||
# 5. 分批插入
|
||
batch_count = (total_valid + self.BATCH_SIZE - 1) // self.BATCH_SIZE
|
||
processed = 0
|
||
|
||
for batch_idx in range(batch_count):
|
||
start = batch_idx * self.BATCH_SIZE
|
||
end = min(start + self.BATCH_SIZE, total_valid)
|
||
batch = keyword_list[start:end]
|
||
|
||
batch_stats = self.keyword_mgr.batch_insert_keywords(batch, query_status=query_status)
|
||
stats['success'] += batch_stats['success']
|
||
stats['skip'] += batch_stats['skip']
|
||
stats['fail'] += batch_stats['fail']
|
||
processed += len(batch)
|
||
|
||
progress = (processed / total_valid) * 100 if total_valid > 0 else 100
|
||
logger.info(
|
||
f"[Query导入] [{filename}] 批次 {batch_idx+1}/{batch_count} | "
|
||
f"进度: {processed}/{total_valid} ({progress:.1f}%) | "
|
||
f"成功: {stats['success']} | 跳过: {stats['skip']} | 失败: {stats['fail']}"
|
||
)
|
||
|
||
# 6. 更新日志
|
||
if log_id:
|
||
self.log_mgr.update_status(
|
||
log_id, 'completed',
|
||
total_count=stats['total'],
|
||
success_count=stats['success'],
|
||
skip_count=stats['skip'],
|
||
fail_count=stats['fail']
|
||
)
|
||
|
||
logger.info(
|
||
f"[Query导入] 文件 {filename} 导入完成 | "
|
||
f"总数: {stats['total']} | 成功: {stats['success']} | "
|
||
f"跳过: {stats['skip']} | 失败: {stats['fail']}"
|
||
)
|
||
|
||
return {'success': True, 'stats': stats}
|
||
|
||
except Exception as e:
|
||
error_msg = str(e)
|
||
logger.error(f"[Query导入] 文件 {filename} 导入失败: {error_msg}")
|
||
|
||
if log_id:
|
||
self.log_mgr.update_status(
|
||
log_id, 'failed',
|
||
total_count=stats['total'],
|
||
success_count=stats['success'],
|
||
skip_count=stats['skip'],
|
||
fail_count=stats['fail'],
|
||
error_message=error_msg
|
||
)
|
||
|
||
return {'success': False, 'stats': stats, 'error': error_msg}
|
||
|
||
def scan_and_import(self):
|
||
"""
|
||
扫描上传目录，导入待处理文件
|
||
|
||
1. 处理数据库中 status='pending' 的记录
|
||
2. 扫描目录中未登记的Excel文件并导入
|
||
"""
|
||
upload_dir = Config.QUERY_UPLOAD_DIR
|
||
|
||
if not os.path.exists(upload_dir):
|
||
os.makedirs(upload_dir, exist_ok=True)
|
||
return
|
||
|
||
# 1. 处理数据库中 pending 状态的记录
|
||
pending_logs = self.log_mgr.get_pending_logs()
|
||
if pending_logs:
|
||
logger.info(f"[Query扫描] 发现 {len(pending_logs)} 个待处理导入任务")
|
||
for log in pending_logs:
|
||
filepath = log['filepath']
|
||
if os.path.exists(filepath):
|
||
logger.info(f"[Query扫描] 处理待导入文件: {log['filename']}")
|
||
self.import_file(filepath, log_id=log['id'])
|
||
else:
|
||
logger.warning(f"[Query扫描] 文件不存在，标记为失败: {filepath}")
|
||
self.log_mgr.update_status(log['id'], 'failed', error_message='文件不存在')
|
||
|
||
# 2. 扫描目录中未登记的文件
|
||
excel_files = []
|
||
for pattern in ['*.xlsx', '*.xls']:
|
||
excel_files.extend(glob.glob(os.path.join(upload_dir, pattern)))
|
||
|
||
for filepath in excel_files:
|
||
filepath = os.path.abspath(filepath)
|
||
if not self.log_mgr.is_file_logged(filepath):
|
||
filename = os.path.basename(filepath)
|
||
logger.info(f"[Query扫描] 发现未登记文件: {filename}")
|
||
log_id = self.log_mgr.create_log(filename, filepath)
|
||
self.import_file(filepath, log_id=log_id)
|